CDLM: 일관성 확산 언어 모델을 통한 추론 가속화

핵심 요약

확산 언어 모델(DLM)은 자동 회귀 모델의 대안으로 주목받고 있으나, 전체 양방향 어텐션 사용으로 인한 KV 캐싱 불가와 많은 정제 단계로 인해 추론 효율성이 낮았다. Together AI가 제안한 CDLM(Consistency Diffusion Language Models)은 블록 단위 인과적 마스크를 도입하여 정확한 KV 캐싱을 가능하게 하고, 일관성 학습을 통해 정제 단계를 획기적으로 줄였다. 이를 통해 수학 및 코딩 작업에서 품질 저하 없이 기존 DLM 대비 최대 14.5배의 지연 시간 단축과 높은 처리량을 달성했다. 이 기법은 사후 학습 레시피로서 다양한 블록 확산 모델에 범용적으로 적용될 수 있는 것이 특징이다.

배경

Diffusion Models, Autoregressive Decoding, KV Caching, Consistency Models

대상 독자

LLM 추론 최적화 엔지니어 및 확산 모델 기반 생성 기술 연구자

의미 / 영향

확산 언어 모델의 고질적인 문제였던 추론 속도를 자동 회귀 모델 수준 이상으로 끌어올림으로써, 텍스트 인필링이나 양방향 문맥 이해가 필요한 작업에서 DLM의 실용성을 크게 높였다.

섹션별 상세

확산 언어 모델(DLM)의 추론 효율성을 저해하는 두 가지 핵심 병목 현상을 해결했다. 기존 DLM은 양방향 어텐션을 사용하여 매 단계 전체 컨텍스트를 재계산해야 하므로 KV 캐싱이 불가능했으며, 고품질 생성을 위해 생성 길이와 맞먹는 많은 정제 단계가 필요했다. CDLM은 블록 단위 인과적 마스크(Block-wise Causal Mask)를 도입하여 이전 블록의 정보를 캐싱하고 현재 블록 내에서만 확산 과정을 수행함으로써 이 문제를 해결했다.

CDLM의 학습은 궤적 수집, 블록 인과적 학생 모델 구성, 세 가지 목적 함수 최적화의 과정을 거친다. 교사 모델로부터 고품질의 디코딩 궤적을 수집한 후, 학생 모델이 블록 내 다중 토큰을 한 번에 확정할 수 있도록 증류 손실(Distillation Loss)과 일관성 손실(Consistency Loss)을 적용한다. 여기에 일반적인 마스크 기반 복원 손실을 추가하여 모델의 추론 및 일반화 능력을 유지하도록 설계했다.

성능 평가 결과 Dream-7B-Instruct 모델 기반의 CDLM은 GSM8K-CoT에서 11.2배, MBPP-Instruct에서 14.5배의 지연 시간 단축을 기록했다. 특히 정제 단계를 4.1배에서 7.7배까지 줄이면서도 정확도 하락은 최소화했다. 이는 단순히 단계를 줄이는 것만으로는 정확도가 급락하는 기존 방식과 달리, 학습을 통해 다중 토큰 정제의 안정성을 확보했음을 증명한다.

시스템 수준 분석에서 CDLM은 자동 회귀(AR) 모델과 바닐라 DLM 사이의 최적의 균형점(Sweet Spot)을 제공한다. AR 모델은 작은 배치 사이즈에서 메모리 대역폭에 제한을 받지만, CDLM은 블록 내 병렬 처리를 통해 연산 집약도(Arithmetic Intensity)를 높여 하드웨어 활용도를 개선한다. 이는 특히 배치 사이즈가 작은 실시간 추론 환경에서 높은 효율성을 발휘하는 근거가 된다.