기술 요약: CForce: 일관성 강제(Consistency Forcing)를 통한 dLLM의 병렬 디코딩 성능 향상

1. 문제 정의

확산 대규모 언어 모델(Diffusion Large Language Models, dLLMs)은 단일 순방향 패스(forward pass)에서 여러 개의 마스크된 토큰을 예측함으로써 텍스트 생성을 가속화하고 높은 병렬성을 달성합니다. 그러나 공격적인 병렬 디코딩 전략(신뢰도 임계값에 따라 많은 토큰을 동시에 확정하는 방식)을 사용할 경우, 초기 디노이징 단계에서는 제한된 문맥으로 인해 예측의 신뢰도가 떨어지는 문제가 발생합니다. 이러한 초기 오류는 후속 단계로 전파되어 전체적인 생성 품질을 저하시킵니다.

기존의 가속화 방법들은 두 가지 주요 한계점에 직면해 있습니다:

1. 추론 시점 접근법(예: 캐싱, 투기적 디코딩)은 모델의 내재적 예측 신뢰도를 개선하기 위해 모델 파라미터를 수정하지 않습니다.
2. 학습 기반 접근법은 종종 구성된 교사(teacher) 생성물이나 특권 정보(privileged) 궤적에 의존합니다. 이러한 타겟은 학생 모델의 실제 추론 상태(낮은 문맥, 임계값에 의해 결정된 경로)와 다르기 때문에 학습-추론 간의 불일치가 발생합니다.

핵심 과제는 외부의 교사 없이, 또는 근본적인 디코딩 알고리즘을 변경하지 않고도 높은 병렬성 디코딩 환경에서 초기 단계 예측의 신뢰도를 개선하는 것입니다.

2. 방법론: 일관성 강제 (CForce)

CForce는 모델 자체의 디코딩 궤적 내에서 초기(낮은 문맥) 디노이징 단계의 예측을 더 정보가 풍부한 후기 단계의 예측과 정렬하도록 설계된 증류(distillation) 방법입니다.

2.1 단계별 궤적 구축 (Staged Trajectory Construction)

CForce는 모든 네이티브 디코딩 스텝에 대해 학습하는 대신, 사전 학습된 dLLM의 자기 실행(self-rollout)으로부터 "단계별" 궤적을 구축합니다:

• 자기 실행 (Self-Rollout): 모델은 신뢰도 기반 병렬 디코딩(임계값 $\tau$ 이상의 토큰을 드러냄)을 사용하여 궤적을 생성합니다.
• 단계 분할 (Stage Partitioning): 궤적은 네이티브 스텝 수가 아닌, 새롭게 드러난 누적 토큰 수(단계 크기 $S$)를 기준으로 단계별로 다운샘플링됩니다. 이는 인접한 단계들이 유의미한 문맥 변화를 나타내도록 보장합니다.
• 인접 쌍 (Adjacent Pairs): 학습 쌍은 이전 단계($\bar{s}_i$)와 이후 단계($\bar{s}_{i+1}$)로 구성됩니다. 이후 단계는 이전 단계의 stop-gradient 타겟 역할을 합니다.

2.2 학습 목적 함수 (Training Objective)

이 방법은 일관성을 강제하기 위해 복합 손실 함수를 사용합니다:

1. 신뢰도 적응형 KL 발산 (Confidence Adaptive KL Divergence, CAD):

   • 두 단계 모두에서 마스크 상태로 남아 있는 위치($U_i$)에 적용됩니다.
   • 후기 단계 예측의 신뢰도($c_j$)에 따라 Forward KL과 Reverse KL 사이를 동적으로 보간합니다.
   • Forward KL: 분포 드리프트(distributional drift)를 최소화하여 안정성을 유지합니다.
   • Reverse KL: 모드 추구(mode-seeking)를 촉진하여 예측을 날카롭게 만들지만, 불확실한 타겟에 대해 붕괴되는 것을 방지하기 위해 신뢰도에 의해 가중치가 조절됩니다.
   • 공식: $L_{CAD} = \frac{1}{|U_i|} \sum_{j \in U_i} [D_{KL}(p^{tar}_j \| p^{stu}_j) + c_j D_{KL}(p^{stu}_j \| p^{tar}_j)]$

2. 교차 엔트로피 (Cross-Entropy, CE) 앵커:

   • 단계 사이의 새롭게 드러난 위치($\Delta_i$)에 적용됩니다.
   • 학생의 예측을 토큰이 확정되기 직전의 깨끗한 타겟 토큰에 고정(anchor)시켜, 마스크 상태에서 드러난 상태로의 전환을 안정화합니다.

3. 커리큘럼 학습 (Curriculum Learning):

   • 훨씬 앞선 단계에서 드러날 토큰을 예측하는 어려움을 해결하기 위해, 모델은 단계 사이의 전체 간격에 점진적으로 노출됩니다. 타겟 상태에서 드러나는 토큰의 비율은 학습 단계에 따라 10%에서 100%까지 증가합니다.

2.3 편집 가능 모델로의 확장 (Extension to Edit-Capable Models)

토큰 간 수정(Token-to-Token, T2T) 기능이 있는 모델(예: LLaDA 2.1)의 경우, CForce는 후기 단계의 교정 정보를 활용합니다:

• 편집 신호 (Edit Signals): 후기 단계에서 토큰이 수정되는 위치($\bar{E}_i$)는 CAD 도메인에 포함됩니다.
• 감독 (Supervision): 이러한 T2T 교정은 초기 Mask-to-Token (M2T) 예측에 더 풍부한 감독 신호를 제공하여, 모델이 스스로의 자기 교정 능력을 통해 학습할 수 있게 합니다.

2.4 이론적 정당성

본 논문은 초기 단계의 예측 오차가 인접 단계 간의 분포 드리프트와 드러남 경계(reveal boundary)에서의 토큰 오차의 합에 의해 상한선이 결정된다는 이론적 경계(theoretical bound)를 제시합니다. CAD 항은 드리프트를 최소화하고, CE 앵커는 경계 오차를 최소화함으로써 초기 단계의 예측 오차를 근사적으로 최소화합니다.

3. 주요 기여

1. 가속화의 재정의: dLLM 가속의 초점을 단순히 샘플링 단계를 줄이는 것에서, 공격적인 임계값 디코딩 하에서의 초기 단계 예측 신뢰도를 개선하는 것으로 전환했습니다.
2. 일관성 강제 (Consistency Forcing): 고정된 교사 모델 없이, 모델 자신의 궤적 내에서 동일 모델의 stop-gradient 타겟을 사용하여 인접 단계를 정렬하는 증류 방법을 도입했습니다.
3. 새로운 목적 함수: 분포적 안정성과 모드 추구 사이의 균형을 맞추는 Confidence Adaptive KL Divergence (CAD)와 토큰 확정을 위한 CE 앵커를 제안했습니다.
4. 편집 가능 모델 통합: T2T 정교화 신호가 어떻게 초기 M2T 초안 작성을 감독할 수 있는지 입증하여, 편집 및 비편집 설정 모두에서 성능을 향상시켰습니다.

4. 실험 결과

LLaDA 2.0(비편집) 및 LLaDA 2.1(편집 가능) 변형 모델을 대상으로 GSM8K, MATH500, MBPP, HumanEval의 네 가지 벤치마크에서 실험을 수행했습니다.

• 편집 가능 모델 (LLaDA 2.1-mini):

  • CForce는 평균 TPF(Tokens Per Forward)를 6.94에서 9.08로 증가시켰습니다.
  • 평균 정확도가 85.57%에서 86.41%로 향상되었습니다.
  • 특히 코드 생성 벤치마크(MBPP, HumanEval)에서 큰 폭의 이득을 보였습니다.

• 비편집 모델 (LLaDA 2.0-mini):

  • CForce는 모델을 더 높은 병렬성 작동 지점으로 이동시켜 평균 TPF를 3.60에서 6.42로 높였습니다.
  • 평균 정확도는 기저 모델보다 약간 낮았으나(내재적인 속도-품질 트레이드오프 때문), 높은 병렬성 환경에서 CAP 베이스라인을 크게 앞질렀습니다.
  • 고정된 공격적 TPF 예산(예: TPF=8) 하에서, CForce는 모든 변형 모델 중 가장 높은 정확도를 달성했습니다.

• 처리량 (Throughput):

  • CForce는 LLaDA 2.0-mini 대비 평균 42.82%, LLaDA 2.1-mini 대비 평균 **22.84%**의 실시간 처리량(wall-clock throughput) 개선을 달성했습니다.

• 절제 연구 (Ablation Studies):

  • 단계 크기 (Stage Size): $S=8$이 병렬성과 정확도 사이에서 최적의 균형을 제공했습니다.
  • 타겟 유형 (Target Type): 동일 모델의 stop-gradient 타겟이 고정된 교사 베이스라인보다 우수한 성능을 보였습니다.
  • KL 형태 (KL Form): CAD는 순수 Forward KL(높은 병렬성)과 순수 Reverse KL(높은 정확도) 모두보다 뛰어난 성능을 보이며 최적의 트레이드오프를 달성했습니다.

5. 의의 및 주장

본 논문은 Consistency Forcing가 고병렬 dLLM 디코딩의 신뢰도 병목 현상을 효과적으로 해결한다고 주장합니다. 초기 단계의 예측이 모델 자신의 궤적에서 유도된 더 정보가 풍부한 후기 상태와 일치하도록 강제함으로써, CForce는 다음을 가능하게 합니다:

• 더 높은 병렬성: 치명적인 오류 전파 없이 더 많은 토큰을 단일 순방향 패스에서 확정할 수 있습니다.
• 개선된 속도-품질 트레이드오프: 이 방법은 운영 경계(operating frontier)를 확장하여, 특히 자기 교정 신호를 사용할 수 있는 편집 가능 시나리오에서 품질을 유지하거나 향상시키면서 더 빠른 생성을 가능하게 합니다.
• 범용적 적용 가능성: 이 접근 방식은 표준 마스크 확산 모델과 편집 가능 모델 모두에 효과적이며, 각각의 강점인 M2T 초안 작성과 T2T 정교화를 활용합니다.

저자들은 현재 이 방법이 고정된 초기화로부터 수집된 사전 수집된 궤적에 의존하고 있어, 학습이 진행됨에 따라 학생의 분포와 괴리될 수 있음을 언급했습니다. 향후 연구로는 forcing 목적 함수가 진화하는 모델에 적응할 수 있도록 온라인 궤적 수집을 탐색하는 것을 제안했습니다.