국소적 확신, 전역적 정체: 확산 언어 모델의 품질-탐색 딜레마

확산 언어 모델(dLLM)은 토큰 생성 순서가 자유로워 복잡한 추론에 유리하지만, 실제로는 확신이 높은 토큰만 먼저 생성하려는 경향 때문에 다양한 해결 경로를 탐색하지 못하는 한계가 있다. 이 논문은 이러한 '품질-탐색 딜레마'를 수학적으로 규명하고, 전역적인 관점에서 최적의 토큰을 선택하는 새로운 디코딩 전략을 통해 추론 성능(Pass@k)을 대폭 향상시켰다.

기존의 확산 언어 모델은 각 단계에서 가장 '확실해 보이는' 토큰을 먼저 확정하는 방식을 사용한다. 이는 딥러닝의 Softmax 출력값에서 확률이 가장 높은 클래스를 선택하는 것과 유사한데, 당장 눈앞의 정확도는 높일 수 있지만 나중에 더 나은 문장을 만들 수 있는 가능성을 미리 차단해버리는 '근시안적 최적화' 문제를 야기한다.

이 논문은 이 문제를 해결하기 위해 '글로벌 템퍼링'이라는 개념을 도입한다. 단순히 현재 위치에서 확률이 높은 토큰을 고르는 것이 아니라, 해당 토큰을 선택했을 때 앞으로 남은 빈칸들을 채워나갈 수 있는 '미래의 선택지'가 얼마나 풍부하고 유망한지를 함께 고려한다. 이는 체스에서 당장 상대 기물을 잡는 수보다, 장기적으로 유리한 형세를 만드는 수를 두는 것과 같다.

결과적으로 모델은 국소적인 확신에만 매몰되지 않고, 전체 문장의 완성도와 다양성을 동시에 확보할 수 있게 된다. 실험 결과, 특히 난이도가 높은 수학 문제에서 기존 방식들이 해결하지 못한 복잡한 논리 전개를 성공적으로 수행하는 모습을 보였다.

전체 시퀀스 분포 p에 대해 품질(Ex~p[log q(x)])과 탐색(H(p))의 가중 합을 최대화하는 엔트로피 정규화 목적 함수를 정의한다. 이 최적화 문제의 해는 전역 시퀀스 분포 q(x)에 온도 파라미터 α를 적용한 파워 분포(Power Distribution) 형태임이 확인됐다.

전역 분포 q(x)는 계산이 불가능하므로, 평균장 분해(Mean-field Factorization)를 통해 각 단계의 조건부 확률을 근사한다. 구체적으로 현재 토큰 v를 선택했을 때의 로그 확률 αℓi,v(s)에, 남은 위치 R(s')들에 대한 미래 기여도의 합인 룩어헤드 보정치 Δi,v(s)를 더한다. Δi,v(s)는 [남은 위치 j의 모든 가능한 토큰 u에 대한 확률 합 → 로그 변환 → 모든 j 위치에 대해 합산] 과정을 거쳐 계산되며, 이는 해당 선택 이후 도달 가능한 완성된 시퀀스들의 총 질량을 의미한다.

실제 디코딩 시에는 Independent Metropolis-Hastings(IMH) 알고리즘을 사용한다. [현재 로짓 기반의 제안 분포 ri에서 후보 토큰 y를 샘플링] → [룩어헤드 보정치 Δ의 차이를 이용해 수식 A(x→y) = min{1, exp(Δ_y - Δ_x)}에 따라 수락 여부 결정] 과정을 거친다. 이 과정은 병렬 연산이 가능하여 추가적인 순차적 오버헤드 없이 구현 가능하다.

WeDLM-8B 모델을 사용한 실험에서, 제안된 IMH 방식은 MATH500 데이터셋 기준 Pass@1 0.540, Pass@k 0.875를 기록하며 Entropy(0.528/0.875) 및 Confidence(0.528/0.851) 기반의 기존 전략들을 압도했다. 특히 고난도 수학 벤치마크인 AIME'24에서 기존 방식들이 Pass@1 0.05~0.09 수준에 머물 때 0.095를 기록하며 뛰어난 성능을 보였다.

LLaDA-8B 모델에서도 MATH500 Pass@1 0.360, Pass@k 0.700을 달성하여 Random(0.251/0.630) 및 Confidence(0.349/0.465) 전략보다 높은 효율성을 증명했다. 특히 난이도별 분석 결과, 쉬운 문제보다 어려운(Hard) 문제일수록 기존 방식과의 성능 격차가 더 벌어지는 것으로 나타나 전역적 탐색의 효과를 입증했다.

IMH 샘플러의 효율성 분석에서는 체인 길이 T=7 정도에서 성능이 수렴하는 빠른 수렴성을 보였으며, 평균 수락률(Acceptance Rate)이 97% 이상으로 유지되어 제안 분포가 타겟 분포를 매우 효과적으로 근사하고 있음을 확인했다.

본 연구는 dLLM의 디코딩 과정을 마르코프 연쇄 몬테카를로(MCMC) 프레임워크 내에서 재해석했다. 기존의 리마스킹 기법들이 특정 임계값 δ를 기준으로 토큰을 확정하는 'Confidence Gating' 방식임을 정의하고, 이것이 시퀀스 엔트로피 H(X)에 상한선(L · hV(δ))을 설정하여 탐색을 저해함을 수학적으로 증명했다.

제안된 글로벌 템퍼링은 로컬 로짓 템퍼링과 달리 전체 시퀀스 수준에서 에너지 기반 모델(EBM)과 유사한 구조를 가진다. 룩어헤드 보정치 Δ는 미래의 주변부 로그 가능도(Marginal Log-likelihood)의 합으로 해석될 수 있으며, 이는 강화학습의 가치 함수(Value Function)와 유사한 역할을 수행하여 현재의 선택이 미래의 보상(성공적인 시퀀스 완성)에 미치는 영향을 평가한다.

구현 측면에서는 KV 캐시의 Copy-on-Write(CoW) 메커니즘을 활용하여 여러 후보 시퀀스를 효율적으로 관리한다. 이를 통해 N개의 후보를 생성할 때 메모리 오버헤드를 O(N·L)에서 O(L + N·T)로 줄였으며, 8개의 후보를 동시에 처리하는 환경에서도 실시간 추론이 가능한 수준의 효율성을 확보했다.

룩어헤드 보정치를 계산할 때 평균장 가정을 사용하여 토큰 간의 독립성을 전제하므로, 토큰 간 의존성이 매우 강한 특정 문맥에서는 근사 오차가 발생할 수 있다. 또한, 제안된 방식은 확산 모델의 구조적 특성에 의존하므로 일반적인 자기회귀(AR) 모델에 직접 적용하기에는 구조적 차이가 존재한다.