DualDiffusion: A Speculative Decoding Strategy for Masked Diffusion Models

1. 문제 정의 (Problem Statement)

마스크드 확산 모델 (MDM, Masked Diffusion Models) 은 autoregressive(자기회귀) 언어 모델의 대안으로 주목받고 있습니다. MDM 은 병렬 토큰 생성과 양방향 컨텍스트 모델링을 가능하게 하여 생성 속도와 제어 가능성을 향상시킵니다. 그러나 MDM 은 다음과 같은 치명적인 계산 병목 현상에 직면해 있습니다.

• KV 캐싱 (Key-Value Caching) 의 부재: MDM 의 양방향 어텐션 메커니즘은 이전 토큰의 키 - 값 (KV) 쌍을 캐싱하여 재사용할 수 없게 만듭니다. AR 모델은 시퀀스 길이에 따라 선형 $O(L)$ 의 복잡도를 가지지만, MDM 은 매 생성 단계마다 전체 시퀀스에 대한 어텐션을 재계산해야 하므로 $O(N^2)$ 의 복잡도를 가집니다.
• 기존 가속화 방법의 한계: FastDLLM 이나 DkvCache 와 같은 최근 방법들은 어텐션 근사나 캐싱 전략을 통해 속도를 높였으나, 이는 토큰 표현을 고정하거나 어텐션을 제한함으로써 생성 품질 (Quality) 을 희생하는 대가를 치릅니다. 즉, MDM 의 핵심 강점인 '반복적 정제 (Iterative Refinement)' 능력이 저하됩니다.

이러한 상황에서 품질과 효율성 (Quality-Efficiency) 의 트레이드오프를 개선할 수 있는 새로운 접근법이 필요했습니다.

2. 방법론 (Methodology)

저자들은 DualDiffusion이라는 새로운 스펙큘레이티브 디코딩 (Speculative Decoding) 프레임워크를 제안합니다. 이는 빠른 '드래fter(Draft)' 모델과 정확한 '버리파이어 (Verifier)' 모델을 결합하여 MDM 의 추론 속도를 높이면서도 품질을 유지하는 것을 목표로 합니다.

핵심 프로세스

DualDiffusion 은 다음 두 단계를 반복합니다:

1. 드래프팅 단계 (Drafting Phase):
   • 빠른 드래fter 모델 (예: FastDLLM 또는 dKV-Cache 를 적용한 모델) 을 사용하여 마스크된 시퀀스에서 $K$개의 단계를 병렬로 수행합니다.
   • 이 단계에서는 속도를 위해 근사 기법을 사용하므로 토큰 예측에 오차가 발생할 수 있습니다.
2. 검증 단계 (Verification Phase):
   • 정확한 버리파이어 모델 (예: 전체 양방향 어텐션을 가진 LLaDA) 을 사용하여 드래fter 가 생성한 시퀀스를 한 번의 포워드 패스로 검증합니다.
   • 검증 알고리즘은 드래fter 의 예측과 버리파이어의 확률 분포를 비교하여 신뢰도가 낮은 토큰을 재마스크 (Remask) 합니다.
   • 수정된 시퀀스가 다음 드래프팅 단계의 입력으로 사용됩니다.

검증 알고리즘 (Verification Algorithms)

MDM 의 비자기회귀적 (Non-autoregressive) 특성과 양방향 의존성으로 인해 AR 모델용 검증 전략을 직접 적용할 수 없습니다. 이에 따라 저자들은 세 가지 새로운 검증 전략을 개발했습니다.

1. Trust Verifier: 드래fter 의 출력을 수정 없이 그대로 수락합니다. 검증 오버헤드가 없으나 오류 수정 기능이 없습니다.
2. KL Divergence-based Remasking: 드래fter 와 버리파이어의 예측 분포 간 KL 발산 (KL Divergence) 을 계산합니다. 발산이 임계값 ($\tau_{KL}$) 을 초과하거나 확률적으로 높은 경우 해당 토큰을 재마스크합니다.
3. Confidence-based Remasking: 버리파이어가 특정 토큰에 대해 가지는 자신감 (최대 확률) 을 기준으로 합니다. 버리파이어의 자신감이 낮은 토큰을 재마스크하여 추가 정제를 유도합니다.

3. 주요 기여 (Key Contributions)

1. DualDiffusion 프레임워크 제안: 마스크드 확산 언어 모델을 위한 최초의 스펙큘레이티브 디코딩 프레임워크를 제안하여, 빠른 근사 드래fter 와 정확한 버리파이어를 결합했습니다.
2. 새로운 검증 알고리즘 개발: 신뢰도 (Trust), KL 발산, 자신감 (Confidence) 기반의 세 가지 재마스크링 알고리즘을 도입하여 드래fter 의 오류를 선택적으로 수정하는 메커니즘을 정립했습니다.
3. 파레토 프론티어 (Pareto Frontier) 개선: 기존 방법들보다 생성 단계 수와 정확도 사이의 최적 균형을 달성했습니다. 즉, 적은 생성 단계로 높은 정확도를 유지합니다.
4. 실험적 검증: MMLU 와 GSM8K 벤치마크를 통해 LLaDA(버리파이어 단독) 및 FastDLLM(드래fter 단독) 대비 성능을 입증했습니다.

4. 실험 결과 (Results)

실험은 MMLU(논리적 추론) 와 GSM8K(수학적 추론) 에서 수행되었습니다.

• MMLU (논리적 추론):

  • 정확도: DualDiffusion 은 0.47 의 정확도를 기록하여, 완전한 LLaDA 모델 (0.48) 과 거의 동일한 수준을 유지했습니다.
  • 속도: LLaDA 대비 3.9 배의 속도 향상 (320.5 초 $\to$ 82.0 초) 을 달성했습니다.
  • 의미: KL 발산 기반 검증이 논리적 추론 작업에서 드래fter 의 오류를 효과적으로 보정하여, 속도와 정확도의 균형을 성공적으로 이뤘음을 보여줍니다.

• GSM8K (수학적 추론):

  • 정확도: DualDiffusion 은 0.25 의 정확도를 보였으며, 이는 LLaDA(0.57) 보다 현저히 낮습니다.
  • 원인: 수학적 추론은 단일 토큰 오류가 전체 해답 경로에 연쇄적으로 영향을 미치는 '다단계 정밀성'이 요구됩니다. 현재 단일 단계 검증 전략은 이러한 복잡한 오류 전파를 완전히 복구하지 못했습니다.
  • 속도: 정확도 하락에도 불구하고 LLaDA 대비 4.0 배의 속도 향상 (298.3 초 $\to$ 75.4 초) 을 기록하여 계산 효율성은 유지했습니다.

• 메모리 오버헤드: 드래fter 와 버리파이어 두 모델을 동시에 메모리에 로드해야 하므로 메모리 사용량이 약 2 배 증가했습니다 (약 31.7 GB).

5. 의의 및 결론 (Significance & Conclusion)

• 비자기회귀 모델의 가속화 가능성 입증: DualDiffusion 은 스펙큘레이티브 디코딩 패러다이가 autoregressive 모델뿐만 아니라 비자기회귀 (Non-autoregressive) 및 양방향 어텐션 기반 모델에도 적용 가능함을 처음으로 증명했습니다.
• 품질 - 효율성 트레이드오프 해소: 기존 가속화 방법들이 품질을 희생하며 속도를 얻었던 반면, DualDiffusion 은 드래fter 의 오류를 버리파이어가 교정하는 방식을 통해 높은 품질을 유지하면서 속도를 획기적으로 개선하는 새로운 파레토 프론티어를 제시했습니다.
• 향후 연구 방향:
  • 수학적 추론과 같은 고정밀 다단계 작업에서는 **적응형 검증 (Adaptive Verification)**이나 다단계 검증 (Multi-step Verification) 전략이 필요합니다.
  • 메모리 오버헤드를 줄이기 위해 드래fter 와 버리파이어 간의 파라미터 공유 또는 모델 증류 (Distillation) 기법이 필요하다는 점을 지적했습니다.

결론적으로, DualDiffusion 은 MDM 의 실용적 배포를 위한 핵심 가속화 전략으로, 특히 논리적 추론과 같은 작업에서 높은 정확도와 빠른 추론 속도를 동시에 달성할 수 있는 가능성을 열었습니다.