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

섹션별 상세

01

확산 언어 모델은 전체 시퀀스를 반복적으로 정제하는 특성상 모든 단계에서 전체 문맥에 대한 어텐션을 재계산해야 하므로 표준적인 KV 캐싱을 사용할 수 없었다. 또한 고품질 생성을 위해 생성 길이와 맞먹는 수많은 정제 단계가 필요하여 실질적인 추론 속도가 자기회귀 모델에 비해 느린 경우가 많았다.

02

CDLM은 블록 단위 인과적 마스크(Block-wise Causal Mask)를 도입하여 이전 블록의 KV 캐시를 재사용할 수 있는 구조를 설계했다. 이는 전체 양방향 어텐션을 사용하는 기존 DLM과 달리, 현재 처리 중인 블록 내에서만 양방향 문맥을 유지하고 완료된 블록은 고정함으로써 효율적인 메모리 관리를 가능하게 한다.

교사 모델과 학생 모델의 어텐션 마스크 구조 비교도 — Diagram왼쪽의 교사 모델은 전체 양방향 어텐션을 사용하는 반면, 오른쪽의 학생 모델(CDLM)은 블록 단위 인과적 마스크를 사용함을 보여준다. 이 구조적 차이가 CDLM에서 정확한 블록 단위 KV 캐싱을 가능하게 하는 핵심 원리임을 설명한다.

03

학습 과정에서 증류 손실(Distillation Loss), 일관성 손실(Consistency Loss), 보조 DLM 손실의 세 가지 목적 함수를 사용한다. 증류 손실은 교사 모델의 예측 분포를 따르도록 유도하며, 일관성 손실은 블록 내의 시간적 일관성을 강제하여 적은 단계로도 안정적인 다중 토큰 확정이 가능하도록 만든다.

CDLM의 학습 목적 함수 및 토큰 확정 과정 시각화 — Diagram단계별로 마스크가 해제되는 과정과 함께 증류 손실 및 일관성 손실이 계산되는 지점을 도식화했다. 학생 모델이 중간 상태에서 최종 블록 완성 상태의 예측을 일치시키도록 학습되는 과정을 구체적으로 보여준다.

04

성능 평가 결과, CDLM-Dream 모델은 GSM8K-CoT 벤치마크에서 11.2배, MBPP-Instruct에서 14.5배의 지연 시간 단축을 기록했다. 특히 정제 단계를 4.1배에서 7.7배까지 줄이면서도 정확도 손실을 최소화하여 실용적인 추론 효율성을 입증했다.

Dream-7B-Instruct 모델 기반 CDLM 성능 벤치마크 결과표 — ChartGSM8K, HumanEval, MATH, MBPP 등 주요 벤치마크에서 CDLM이 기존 방식 대비 TPS는 최대 20.9배 향상시키고 지연 시간은 대폭 단축했음을 수치로 증명한다. 특히 정확도(Score)를 유지하면서도 총 단계(Total Steps)를 획기적으로 줄였음을 확인할 수 있다.

05

시스템 분석에 따르면 CDLM은 자기회귀 모델의 메모리 제한(Memory-bound) 특성과 일반 DLM의 연산 제한(Compute-bound) 특성 사이의 균형 잡힌 지점에 위치한다. 블록 내 병렬 처리를 통해 메모리 접근을 최적화하면서도 하드웨어 자원을 효율적으로 활용할 수 있는 최적의 운영 지점을 제공한다.

배치 크기에 따른 모델별 연산 강도(Arithmetic Intensity) 분석 그래프 — Chart자기회귀(AR) 모델, 일반 DLM, 그리고 블록 크기별 CDLM의 연산 강도를 비교 분석했다. CDLM이 AR 모델보다 높은 연산 강도를 유지하여 메모리 대역폭 병목 현상을 완화하고 하드웨어 효율성을 높이는 '스위트 스팟'에 위치함을 시각적으로 나타낸다.

용어 해설

확산 언어 모델(Diffusion Language Model): — 데이터의 노이즈를 점진적으로 제거하며 텍스트를 생성하는 모델로, 한 번에 하나의 토큰을 생성하는 자기회귀 모델과 달리 전체 시퀀스를 병렬로 정제할 수 있습니다.
일관성 모델(Consistency Model): — 확산 모델의 샘플링 궤적 상의 모든 지점을 동일한 최종 결과물로 매핑하도록 학습하여, 단 한 번 혹은 매우 적은 단계의 추론만으로 고품질 샘플을 생성하는 기법입니다.
KV 캐싱(KV Caching): — Transformer 모델의 추론 시 이전 토큰들의 Key와 Value 행렬을 저장해 두어 중복 계산을 방지함으로써 생성 속도를 비약적으로 높이는 최적화 기술입니다.
연산 강도(Arithmetic Intensity): — 메모리에서 읽어온 데이터 1바이트당 수행되는 부동 소수점 연산(FLOPs)의 비율로, 모델의 성능이 메모리 대역폭에 제한되는지 연산 능력에 제한되는지를 판단하는 지표입니다.

기술

CDLM
Dream-7B-Instruct
LLaDA-8B-Instruct
Block-wise Causal Attention

활용 사례

수학 문제 풀이 (GSM8K, MATH)
코드 생성 (HumanEval, MBPP)
텍스트 인필링 및 정제

언급된 리소스

논문Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models