기술 요약: 이산 확산 언어 모델을 위한 평균장 병렬 디코딩 (Mean-Field Parallel Decoding)

1. 문제 정의

이산 확산 언어 모델(dLLM)은 병렬 토큰 생성을 가능하게 함으로써 저지연 디코딩을 위한 경로를 제공하며, 이는 순차적인 자기회귀(AR) 모델과 대조를 이룹니다. 그러나 결정적인 병목 현상은 **결합 불일치(joint inconsistency)**에 존재합니다. 대부분의 dLLM은 마스크된 위치의 토큰 분포를 독립적으로 예측하는 분해된 역방향 정책(factorized reverse policy)을 채택합니다. 여러 토큰을 오직 개별적인 신뢰도(예: 가장 높은 확률 또는 마진)에만 기반하여 병렬로 확정(commit)할 경우, 결과 시퀀스는 종-종 "결합 미스매치(joint mismatch)" 문제를 겪게 됩니다. 개별적으로는 신뢰할 수 있어 보이는 토큰들이, 모델링되지 않은 구문적 및 의미적 의존성(예: 주어-동사 일치 또는 논리적 일관성)으로 인해 동시에 업데이트될 때 서로 호환되지 않는 구성을 형성할 수 있기 때문입니다.

이를 완화하기 위한 기존 솔루션으로는 안전한 확정을 스케줄링하는 의존성 인식 추론(예: DAWN)이나 블록 단위 준-자기회귀 디코딩 등이 있습니다. 그러나 이러한 접근 방식은 종종 보조 모델, 학습된 예측기, 그래프 구축 휴리스틱 또는 추가적인 학습 목적 함수를 필요로 하며, 이는 계산 오버헤드나 복잡성을 초래합니다. 따라서 생성 품질을 희생하지 않으면서 병렬 확정을 조절하기 위해 추론 중에 토큰 간 상호작용을 추정할 수 있는 가볍고 학습이 필요 없는 메커니즘이 필요합니다.

2. 방법론: 평균장 병렬 디코딩 (Mean-Field Parallel Decoding)

저자들은 토큰 확정을 구조화된 추론 문제로 취급하여 병렬 업데이트를 조정하는, 학습이 필요 없는 단일 패스(single-pass) 프레임워크인 평균장 병렬 디코딩을 제안합니다.

2.1 에너지 기반 정식화

확정 세트 $C_t$(언마스킹될 마스크된 위치들의 집합)의 선택은 이진 지표 $S = (s_1, \dots, s_m)$에 대한 쌍방향 마르코프 무작위 장(pairwise Markov Random Field, MRF)으로 모델링됩니다. 여기서 $s_i=1$은 확정을 나타냅니다. 구성의 확률은 다음과 같은 깁스 분포(Gibbs distribution)를 통해 정의됩니다:
$$P(S; U_t) \propto \exp\left( \sum_{i} \phi_i(s_i) + \sum_{i<j} \phi_{ij}(s_i, s_j) \right)$$

• 단항 포텐셜 (Unary Potentials, 국소 신뢰도): $\phi_i(s_i) = c_i s_i$, 여기서 $c_i$는 위치 $i$에서의 상위-2 예측 토큰 간의 로그 마진(log-margin)입니다. 이는 높은 국소 확실성을 가진 위치의 확정을 장려합니다.
• 쌍방향 포텐셜 (Pairwise Potentials, 예측 중첩): $\phi_{ij}(s_i, s_j) = -D_{ij} s_i s_j$. 이 항은 강하게 결합된 위치들의 동시 확정을 억제합니다. 상호작용 행렬 $D$는 마스크된 토큰들의 예측 분포 $\pi_{t,i}$와 $\pi_{t,j}$ 사이의 **정규화된 젠슨-샤논 발산(normalized Jensen-Shannon Divergence, JSD)**을 사용하여 구축됩니다:
  $$D_{ij} = \frac{1 - \text{JSD}(\pi_{t,i}, \pi_{t,j}) / \ln 2}{\max_{a,b} \tilde{D}_{ab}}$$
  높은 $D_{ij}$는 유사한 예측 분포를 의미하며, 이는 두 위치가 모두 확정될 경우 중복 또는 충돌의 위험이 높음을 시사합니다. 이는 유사한 위치들이 확정을 위해 경쟁하게 만드는 "반강자성(antiferromagnetic)" 결합을 생성합니다.

2.2 변분 평균장 추론 (Variational Mean-Field Inference)

최대 사후 확률(MAP) 할당을 찾는 것은 NP-hard 문제입니다. 저자들은 완전 분해된 분포 $q(S) = \prod q_i^{s_i}(1-q_i)^{1-s_i}$로 결합 분포를 근사하는 **변분 평균장 완화(variational mean-field relaxation)**를 사용합니다. 여기서 $q_i$는 확정 강도(확정의 주변 확률)를 나타냅니다.

증거 하한(ELBO)을 최대화하면 다음과 같은 고정점 업데이트 방정식을 얻습니다:
$$q_i = \sigma\left( c_i - \sum_{j \neq i} D_{ij} q_j \right)$$
여기서 $\sigma$는 로지스틱 시그모이드 함수입니다. 이 업데이트는 높은 신뢰도를 가진 위치가 쌍방향 상호작용을 통해 다른 위치를 조절할 수 있게 합니다. 알고리즘은 이 업데이트를 $R$번(통상적으로 $R=2$) 반복하여 안정적인 확정 강도 세트에 수렴하며, 이후 이를 임계값 처리하여 최종 확정 세트 $C_t$를 결정합니다.

2.3 구현 세부 사항

• 학습 불필요 (Training-Free): 보조 모델, 재학습 또는 그래디언트 업데이트를 필요로 하지 않습니다. 이는 베이스 모델의 단일 순방향 패스에서 얻은 예측 분포만을 사용하여 작동합니다.
• 복잡도: 주요 오버헤드는 JSD 계산을 위한 $O(m^2|V|)$이며, 여기서 $m$은 마스크된 위치의 수, $|V|$는 어휘 사전의 크기입니다. 이는 GPU에서 매우 병렬화가 가능합니다.
• 수렴: 실험 결과, 확정 강도가 안정적인 고정점에 도달하는 데 $R=2$번의 반복만으로도 충분하다는 것을 보여줍니다.

3. 주요 기여

1. 정식화: dLLM에서의 병렬 확정 선택 문제를 이터레이티브 디노이징(iterative denoising)과 에너지 기반 모델을 잇는 구조화된 추론 작업으로 정식화했습니다.
2. 학습 불필요 지표: 단일 순방향 패스와 보조 구성 요소 없이도 작동하는, 분포 중첩(JSD)에 기반한 그래디언트 프리(gradient-free) 지표를 도입했습니다.
3. 효율적 추론: 국소적 확실성과 상호작용 인식을 모두 고려하여 적응형 병렬성을 구현하는 효율적인 평균장 반복법을 도출했습니다.
4. 실험적 검증: 제안된 방법이 추론 및 코드 생성 벤치마크 전반에서 품질-지연 시간 트레이드오프를 유의미하게 개선함을 입증했습니다.

4. 실험 결과

본 방법은 세 가지 dLLM 백본(LLaDA-8B-Instruct, LLaDA-1.5, Dream-v0-Instruct-7B)을 사용하여 네 가지 벤치마크(GSM8K, MATH, HumanEval, MBPP)에서 평가되었습니다.

• 처리량 vs 정확도: 제안된 방법은 비교 대상인 모든 베이스라인(Entropy, KLASS, LocalLeap, DAWN) 중에서 경쟁력 있는 작업 정확도를 유지하면서도 **가장 높은 초당 토큰 수(TPS)**를 지속적으로 달고했습니다.
  • LLaDA-8B-Instruct의 경우, Entropy 베이스라인 대비 4.07배에서 6.10배의 속도 향상을 달성했습니다.
  • LLaDA-1.5의 경우, MBPP에서 최대 8.62배의 속도 향상을 기록했습니다.
• 베이스라인과의 비교:
  • Entropy 베이스라인(보수적, 높은 정확도, 낮은 처리량)이나 KLASS(주변 신뢰도만 고려)와 달리, 평균장 디코딩은 정확도의 과도한 손실 없이 처리량을 높이는 방향으로 트레이드오프를 이동시킵니다.
  • DAWN 및 LocalLeap과 같은 의존성 인식 베이스라인과 비교했을 때, 제안된 방법은 보조 검증기나 추가 순방향 패스 없이도 일관된 처리량 우위를 제공하여 더 효율적인 확정 세트를 선택함을 보여주었습니다.
• 절제 연구 (Ablation Studies):
  • 쌍방향 상호작용을 제거하면 정확도가 크게 떨어졌으며, 이는 국소 신뢰도만으로는 공격적인 병렬 처리가 불충분함을 확인시켜 줍니다.
  • 균일한 상호작용 행렬을 사용했을 때 JSD 유도 행렬보다 성능이 낮게 나타났으며, 이는 특정 상호작용 구조의 중요성을 강조합니다.
  • 반복 횟수를 $R=1$로 줄이면 성능이 저하되어, 반복적 정교화의 필요성을 입증했습니다.
• 민감도: 이 방법은 확정 임계값 $\tau$와 블록 크이에 의해 제어되는 매끄럽고 해석 가능한 트레이드오프를 보여줍니다. 낮은 임계값은 처리량을 높이지만 정확도를 낮추고, 높은 임계값은 속도를 희생하여 정확도를 보존합니다.

5. 의의 및 주장

본 논문은 평균장 병렬 디코딩이 이산 확산 언어 모델의 "결합 불일치" 병목 현상을 완화하기 위한 원칙적이고 단일 패스적인 메커니즘을 제공한다고 주장합니다. 주변 확률의 독립적 업데이트를 국소적 신뢰도와 전역적 호환성 사이의 균형을 맞추는 구조적 프로세스로 대체함으로써, 디코더가 각 단계에서 더 많은 수의 토큰을 안전하게 언마스킹할 수 있게 합니다.

이 방법의 의의는 그 경량성에 있습니다. 보조 모델, 재학습 또는 복잡한 스케줄링 휴리스틱 없이도 처리량과 생성 일관성을 개선합니다. 이는 모델 자체의 예측 분포를 기반으로 한 협력적 선택이 추론 및 코드 생성과 같은 복잡한 텍스트 합성 작업에서 우수한 품질-지연 시간 트레이드오프를 달 수 있음을 보여줍니다. 저자들은 이를 기존의 확산 디코딩 파이프라인에 즉시 적용 가능한 실용적인 솔루션으로 제시합니다.