[Paper] CD4LM: 일관성 증류와 aDaptive 디코딩 for Diffusion Language Models

Source: arXiv - 2601.02236v1

개요

논문 “CD4LM: Consistency Distillation and aDaptive Decoding for Diffusion Language Models” 은 현대 언어 생성의 핵심 병목 현상인 자동회귀 디코딩으로 인한 지연 시간을 해결합니다. 확산 언어 모델(DLM)의 학습 및 디코딩 방식을 재고함으로써, 저자들은 텍스트를 고도로 병렬적으로 생성하면서도 최첨단 자동회귀 모델과 동등한 품질을 유지하는 프레임워크를 제시합니다.

주요 기여

• Discrete‑Space Consistency Distillation (DSCD): 새로운 학습 목표로, “학생” 확산 모델이 trajectory‑invariant가 되도록 강제합니다. 즉, 어떤 잡음이 섞인 중간 상태라도 직접 깨끗한 토큰 분포로 매핑할 수 있습니다.
• Confidence‑Adaptive Decoding (CAD): 토큰 수준의 신뢰도를 모니터링하고, 높은 확신을 가진 토큰에 대해 확산 단계를 동적으로 건너뛰는 추론 알고리즘으로, 함수 평가 횟수를 크게 줄입니다.
• Empirical Pareto‑frontier improvement: 수학, 코드, 추론 벤치마크(예: GSM8K, MBPP)에서 CD4LM은 강력한 베이스라인 대비 3–5배의 실제 시간 속도 향상을 달성하면서 정확도는 동등하거나 능가합니다.
• Open‑source implementation: 저자들은 코드와 사전 학습된 체크포인트를 공개하여 실무자들이 CD4LM을 기존 파이프라인에 쉽게 적용할 수 있도록 했습니다.

Source: …

방법론

1. 확산 언어 모델링 입문

   • 기존의 확산 모델은 무작위 노이즈에서 시작해 점차 깨끗한 문장으로 이동하면서 이산 토큰 시퀀스를 반복적으로 디노이징하여 텍스트를 생성합니다. 각 디노이징 단계는 별도의 신경망 호출을 필요로 하므로 추론 비용이 많이 듭니다.

2. 일관성 증류

   • 고정된 시간 단계에서 다음 토큰을 예측하도록 모델을 학습하는 일반적인 “지역” 손실 대신, DSCD는 학생 모델이 이미 수행된 확산 단계 수와 무관하게 동일한 출력을 생성하도록 훈련합니다.
   • 구체적으로, 교사는 전체 확산 궤적을 실행하고, 학생은 중간의 어느 노이즈 상태(예: 2단계, 5단계, 10단계 후)든 최종 깨끗한 분포로 바로 매핑하도록 요구됩니다. 이 “궤적 불변성”은 학생에게 건너뛴 단계에 대한 내재된 강인성을 부여합니다.

3. 적응형 디코딩

   • 생성 과정에서 CAD는 각 토큰에 대해 신뢰도 점수(예: 최대 소프트맥스 확률)를 계산합니다.
   • 신뢰도가 설정 가능한 임계값을 초과하는 토큰은 고정되어, 디코더가 해당 토큰에 대한 추가 정제를 중단하고 여러 확산 단계를 “점프”합니다.
   • 신뢰도가 낮은 토큰은 계속 정제되어, 문장의 어려운 부분이 충분한 계산을 받도록 보장합니다.

4. 병렬 생성

   • 확산 과정이 토큰별이 아니라 전체 시퀀스를 한 번에 처리하기 때문에, CAD는 GPU 배치 병렬성을 활용하여 실제 소요 시간을 크게 단축할 수 있습니다.

결과 및 발견

• 품질 보존: 높은 신뢰도의 토큰에 대해 최대 80 %의 diffusion 단계를 건너뛰었음에도 최종 출력은 전체 단계 기준과 통계적으로 구별되지 않는다.
• 효율성 프론티어: 정확도‑효율성 그래프에서 CD4LM은 기존 diffusion 기반 및 autoregressive 방법들을 모두 압도하며 새로운 Pareto‑optimal 영역을 형성한다.
• 소거 실험: DSCD를 제거하고(즉, 표준 diffusion 손실 사용) 하면 CAD가 몇 번의 건너뛰기 후 바로 붕괴되어, trajectory‑invariance가 안전한 가속에 필수적임을 확인한다.

실용적 시사점

• 저지연 AI 서비스: 챗봇, 코드 어시스턴트, 실시간 추론 도구가 이제 자동회귀 디코딩에서 흔히 발생하는 수초 수준의 지연 없이 확산 모델을 활용할 수 있습니다.
• 클라우드 GPU 비용 절감: 생성된 토큰당 전방 패스 횟수가 줄어들면 계산 비용이 직접적으로 감소하며, 특히 대량 처리 워크로드(예: 문서나 테스트 케이스의 배치 생성)에서 효과적입니다.
• 가변 컴퓨팅 예산에 대한 견고성: CAD의 신뢰도 임계값을 실시간으로 조정할 수 있어, 서비스가 트래픽 급증 시 품질을 약간만 희생하고 속도를 높일 수 있습니다.
• 배포 간소화: 모델이 여전히 단일 패스 신경망(외부 토큰 수준 스케줄러 없음)이라 기존 추론 스택(TensorRT, ONNX Runtime)과 최소한의 엔지니어링 노력으로 CD4LM을 통합할 수 있습니다.

제한 사항 및 향후 연구

• 이산 토큰 공간 가정: DSCD는 현재 토큰‑레벨 확산을 위해 설계되었습니다; 이를 서브워드 또는 문자‑레벨 공간으로 확장하려면 추가적인 트릭이 필요할 수 있습니다.
• 신뢰도 보정: 적응형 스키핑은 소프트맥스 확률에 의존하는데, 특정 도메인(예: 고도로 전문적인 용어)에서는 보정이 맞지 않을 수 있습니다. 더 나은 불확실성 추정기가 견고성을 향상시킬 수 있습니다.
• 대규모 모델로의 확장: 실험은 약 2 B 파라미터까지의 모델을 대상으로 수행되었습니다. DSCD와 CAD를 >10 B‑파라미터 규모로 확장하면 새로운 안정성 문제가 드러날 수 있습니다.
• 다양한 모달리티: 저자들은 일관성‑증류 원칙이 이미지나 오디오용 확산 모델에 도움이 될 수 있다고 제안하지만, 구체적인 실험은 향후 연구로 남겨두었습니다.

CD4LM을 직접 사용해보고 싶다면, 저자들은 바로 실행 가능한 Docker 이미지와 벤치마크 재현 스크립트를 제공합니다. 기존 생성 파이프라인에 연결하여 오늘부터 지연 시간 개선을 측정해 보세요!

저자

• Yihao Liang
• Ze Wang
• Hao Chen
• Ximeng Sun
• Jialian Wu
• Xiaodong Yu
• Jiang Liu
• Emad Barsoum
• Zicheng Liu
• Niraj K. Jha

논문 정보

• arXiv ID: 2601.02236v1
• Categories: cs.CL
• Published: 2026년 1월 5일
• PDF: PDF 다운로드