빠르고 정확한 장문 텍스트 생성을 위한 소스텝 디퓨전 언어 모델

초록

FS-DFM은 디스크리트 흐름 매칭 기반의 확산 언어 모델에 “샘플링 단계 수”를 명시적 제어 변수로 도입해, 몇 단계만으로도 1 024 토큰 길이의 텍스트를 고품질로 생성한다. 8단계 샘플링으로 1 024‑step 기준 모델과 동등한 퍼플렉시티를 달성하며, 최대 128배 빠른 추론 속도를 제공한다.

상세 분석

FS-DFM은 기존 디스크리트 흐름 매칭(DFM) 프레임워크를 확장해 두 가지 핵심 아이디어를 결합한다. 첫째, “스텝‑어웨어” 메커니즘을 도입해 모델 입력에 목표 샘플링 단계 수 h를 추가한다. 이를 통해 모델은 “큰 한 걸음”이 여러 작은 걸음의 합과 동일한 확률 경로를 따라가도록 학습한다. 학습 시에는 다양한 h값을 무작위로 샘플링해 다중 단계 예산에 대해 일관된 출력을 강제하는 self‑consistency 손실을 적용한다. 둘째, 업데이트 규칙으로 Runge‑Kutta(Heun (RK‑2)와 RK‑4) 기반의 “shortcut teacher”를 사용한다. 교사는 장시간(1 024 step) 흐름을 미리 실행해 얻은 고정밀 경로를 제공하고, 이를 통해 큰 단계에서도 과도한 오버슈팅 없이 안정적인 확률 이동을 가능하게 한다.

이론적 배경으로는 연속시간 마코프 체인(CTMC) 위에 정의된 infinitesimal generator uₜ를 학습한다. 기존 DFM은 토큰별 전이율을 독립적으로 예측해 |V|·L 차원의 출력을 만든다. FS-DFM은 여기서 h에 따라 스칼라 형태의 “cumulative scalar”를 도입해, uₜ를 h‑scaled 형태로 변환한다. 즉, pₜ→pₜ₊ₕ 전이 확률을 직접 계산하는 대신, g(t)=κ̇ₜ/(1‑κₜ)와 같은 스케일링 함수를 학습해 큰 h에서도 정확히 이동한다.

실험에서는 0.17 B 파라미터 규모의 모델을 포함해 Dream‑7B, LLaDA‑8B 등 다양한 규모의 베이스라인과 비교했다. 8 NFE(함수 평가) 설정에서 FS‑DFM은 1 024‑step DFM과 거의 동일한 퍼플렉시티(≈7.5)와 토큰 정확도를 기록했으며, 토큰당 지연시간은 1/128 수준으로 감소했다. 특히 장문(1 024 토큰) 무조건 생성 실험에서, FS‑DFM은 반복된 쉼표·공백 같은 아티팩트를 최소화하고, 텍스트 흐름과 의미 일관성을 유지했다.

또한, 교사 모델의 Runge‑Kutta 선택이 성능에 미치는 영향을 분석했으며, RK‑2가 대부분의 설정에서 좋은 안정성과 정확도 균형을 제공한다는 결론을 내렸다. 교사 distillation 과정에서는 KL‑divergence 기반의 Bregman 손실을 사용해, 목표 토큰 분포 p₁|ₜ를 정확히 복원하도록 유도한다.

전체적으로 FS‑DFM은 “few‑step”이라는 새로운 설계 목표를 명시적으로 반영함으로써, 확산 기반 텍스트 생성의 전통적인 “많은 단계‑높은 품질” 트레이드오프를 깨뜨렸다. 이는 실시간 대화, 대규모 문서 자동 생성 등 고속·고품질 텍스트 생성이 요구되는 실제 서비스에 바로 적용 가능한 기술적 진보라 할 수 있다.