LangFlow: 언어 모델링에서 이산 확산 모델에 필적하는 연속 확산 모델

이미지 생성에서 강력한 성능을 보였던 연속 확산 모델이 언어 모델링에서는 이산 모델에 뒤처졌으나, LangFlow는 이를 극복하고 동등하거나 더 우수한 성능을 증명했다. 이는 언어 생성 분야에서 연속 확산 모델이 새로운 표준이 될 수 있는 가능성을 열어준 연구이다.

기존의 언어용 확산 모델은 단어라는 이산적인 데이터를 다루기 위해 복잡한 상태 전이 행렬을 사용하거나, 연속적인 공간으로 변환하더라도 데이터가 희소하게 분포하는 한계가 있었다. LangFlow는 단어를 고차원 벡터 공간인 Embedding에 매핑한 뒤, 이 공간에서 가우시안 노이즈를 데이터 분포로 변환하는 속도장(Velocity Field)을 학습하는 Flow Matching 기법을 적용했다.

핵심은 언어 데이터가 이미지와 달리 매우 노이즈가 많은 상태에서도 정답 토큰을 쉽게 예측할 수 있다는 점에 착안한 것이다. 연구진은 노이즈 수준에 따른 손실 함수의 변화를 분석하여, 정보가 실제로 생성되는 특정 노이즈 구간에 학습 자원을 집중 배치하는 전략을 세웠다.

결과적으로 LangFlow는 단순한 노이즈 제거를 넘어, 데이터의 기하학적 구조에 맞춘 최적화된 경로를 따라 토큰을 생성한다. 이를 통해 연속 확산 모델임에도 불구하고 이산 모델보다 더 낮은 Perplexity를 달성하며 언어 생성의 효율성과 품질을 동시에 확보했다.

LangFlow는 텍스트 토큰을 임베딩 행렬 E를 통해 D차원 벡터로 변환한 후, Flow Matching 프레임워크 내에서 가우시안 분포로부터 데이터 분포로의 결정론적 궤적을 학습한다. 학습 목적 함수는 Bregman Divergence를 최소화하는 것과 Cross-Entropy 손실 함수를 최소화하는 것이 동치임을 증명하여, 토큰 공간에서의 직접적인 최적화를 수행한다.

노이즈 스케줄링을 위해 로그 노이즈 대 신호비(logNSR)인 γ를 도입했다. [시간 t → γ 변환 → 노이즈 수준 결정] 과정을 거치며, 정보 획득 속도 H'γ가 Gumbel 분포를 따른다는 관찰을 바탕으로 π(γ) ∝ H'γ가 되도록 스케줄러를 구성했다. 이는 정보 이득이 큰 구간에 샘플링 단계를 더 많이 할당하여 효율성을 극대화한다.

학습 시에는 이전 단계의 예측값 z_hat을 다시 입력으로 사용하는 Self-Conditioning을 적용한다. [현재 노이즈 상태 z_γ와 이전 예측 z_hat을 입력 → 데노이저 통과 → 새로운 예측 산출] 과정을 통해 모델이 자신의 이전 출력을 참고하여 더 정교한 임베딩을 생성하도록 유도하며, 이는 특히 연속 공간에서의 모드 붕괴를 방지하는 역할을 한다.

LangFlow는 LM1B 데이터셋에서 Perplexity(PPL) 30.0, OpenWebText에서 24.6을 기록하며 기존의 최첨단 이산 확산 모델인 MDLM(31.0)과 SEDD(32.0)를 능가했다. 생성 품질을 나타내는 Generative PPL에서도 LM1B 기준 92.2를 달성하여 비교 모델 중 가장 우수한 성능을 보였다.

제로샷 전이 성능 평가에서는 PTB, Wikitext, Lambada 등 7개 벤치마크 중 4개에서 자기회귀(Autoregressive) 모델인 Transformer 베이스라인을 앞질렀다. 특히 PTB 데이터셋에서는 81.20의 PPL을 기록하여 모든 비교 대상 확산 모델 중 1위를 차지했다.

Ablation Study 결과, 제안된 Gumbel 노이즈 스케줄러는 Generative PPL을 약 1000에서 154.2로 획기적으로 낮추었으며, Self-Conditioning은 PPL을 49.0에서 30.0으로 19포인트 가량 개선하는 결정적인 기여를 했음이 확인됐다.

LangFlow 아키텍처는 130M 파라미터 규모의 수정된 Diffusion Transformer(DiT)를 기반으로 하며, 12개의 레이어와 768의 숨겨진 차원을 갖는다. 시간 t 대신 logNSR인 γ를 조건부 입력으로 사용하며, AdaLN(Adaptive Layer Norm) 대신 임베딩 유닛 구체 상의 정규화를 적용하여 노이즈 분산과 데이터 분산을 일치시켰다.

이론적으로는 Bregman Divergence Flow Matching을 통해 토큰 예측 확률로부터 연속적인 데노이저 z_hat을 결정론적으로 도출한다. [예측된 토큰 확률 x_hat → 임베딩 행렬 E와의 곱 → 기대 임베딩 값 z_hat]의 과정을 거치며, 이는 ODE 샘플링 시 속도 벡터를 결정하는 핵심 요소가 된다.

기존 연구들이 MSE(Mean Squared Error) 손실 함수를 사용하여 임베딩 붕괴(Embedding Collapse) 문제를 겪었던 것과 달리, LangFlow는 Cross-Entropy를 사용하여 각 토큰 임베딩이 고유한 영역을 유지하도록 강제함으로써 모델의 표현력을 보존했다.

LangFlow는 강력한 PPL 성능에도 불구하고 샘플의 엔트로피가 일부 이산 확산 모델에 비해 낮게 나타나는 경향이 있다. 이는 생성된 텍스트에서 특정 단어가 과도하게 반복되는 현상으로 이어질 수 있으며, 더 큰 규모의 모델과 데이터셋에서 이러한 효과가 어떻게 변화하는지에 대한 추가 연구가 필요하다.