국소적으로는 자신만만하지만 전역적으로는 정체됨: 확산 언어 모델의 품질-탐색 딜레마

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

기존의 확산 언어 모델은 매 단계마다 모델이 가장 자신 있는 토큰을 먼저 확정하는 방식을 사용한다. 이는 Softmax 출력값 중 가장 높은 확률을 가진 토큰을 선택하는 Greedy 방식과 유사하며, 당장의 예측 정확도는 높일 수 있지만 나중에 발생할 수 있는 더 나은 문장 완성 가능성을 차단하는 결과를 낳는다. 즉, 국소적인 확신에 매몰되어 전역적으로 더 나은 해답을 찾지 못하는 상태에 빠지게 된다.

이 논문은 이 문제를 해결하기 위해 현재 선택이 '미래에 얼마나 유망한 경로를 남겨두는가'를 계산에 포함한다. 이를 위해 엔트로피 정규화 개념을 도입하여, 단순히 확률이 높은 토큰을 뽑는 것이 아니라 '높은 확률'과 '다양한 후속 경로 유지' 사이의 최적 균형점을 찾는다. 이는 강화학습에서 탐험(Exploration)과 활용(Exploitation)의 균형을 맞추는 것과 유사한 원리이다.

결과적으로 모델은 현재 단계에서 조금 덜 확실해 보이더라도, 이후에 더 정답에 가까운 문장을 완성할 가능성이 큰 토큰을 선택할 수 있게 된다. 이를 통해 복잡한 수학 문제처럼 여러 단계의 논리적 비약이 필요한 과제에서 모델이 한 가지 잘못된 경로에 고착되지 않고 다양한 해결책을 시도할 수 있는 능력을 갖추게 된다.

전체적인 접근 방식은 dLLM의 디코딩 과정을 엔트로피 정규화된 최적화 문제로 재정의하는 것이다. 목적 함수는 모델의 로그 가능도 기대값(품질)과 시퀀스 분포의 섀넌 엔트로피(탐색)의 합으로 구성되며, 제어 파라미터 α를 통해 둘 사이의 가중치를 조절한다. [입력: 시퀀스 분포 q(x), 가중치 α] → [연산: αE[log q(x)] + H(p) 최대화] → [출력: 최적 분포 p*(x) ∝ q(x)^α] → [의미: 품질과 다양성이 최적화된 타겟 분포 도출]

핵심 메커니즘인 '룩어헤드 수정(Lookahead Correction)'은 특정 토큰 v를 선택했을 때 도달 가능한 모든 완성 시퀀스의 확률 합을 측정한다. 정확한 계산은 지수적인 연산량이 필요하므로, dLLM의 독립적 주변 확률 예측 특성을 이용한 평균장 근사(Mean-field Approximation)를 적용한다. [입력: 현재 상태 s, 후보 토큰 v] → [연산: 잔여 위치들에 대한 로그 주변 확률 합산] → [출력: 수정 항 Δi,v(s)] → [의미: 해당 토큰 선택 시 남은 탐색 공간의 유망도 수치화]

최종적으로 Independent Metropolis-Hastings(IMH) 샘플러를 사용하여 이 타겟 분포로부터 샘플링을 수행한다. 로컬 템퍼링된 로짓을 제안 분포(Proposal Distribution)로 사용하며, 수락 확률 계산 시 로컬 항들이 상쇄되어 오직 룩어헤드 수정 항의 차이만으로 수락 여부를 결정한다. [입력: 현재 토큰 x, 제안 토큰 y] → [연산: min(1, exp(Δy - Δx))] → [출력: 수락 확률] → [의미: 어휘 사전 전체에 대한 정규화 없이도 전역 최적 분포를 효율적으로 근사]

메인 벤치마크 결과, WeDLM-8B 모델 기준 MATH500에서 Pass@1 0.540, Pass@k 0.875를 기록하며 기존 Entropy(0.528/0.875) 및 Confidence(0.528/0.851) 기법을 능가했다. 특히 고난도 수학 문제인 AIME 2024/2025에서 기존 기법들이 Pass@k에서 정체되는 것과 달리, 제안된 IMH 방식은 k가 증가함에 따라 성능이 지속적으로 우상향하는 결과를 보였다.

코딩 벤치마크인 HumanEval과 MBPP에서도 일관된 성능 향상이 관찰됐다. LLaDA-8B 모델을 사용했을 때 HumanEval에서 Pass@k 0.695를 달성하여 Confidence 기법(0.573) 대비 약 12%p 이상의 큰 폭의 성능 향상을 기록했다. 이는 제안된 방식이 단순한 무작위성 주입이 아니라 유망한 경로를 체계적으로 탐색하고 있음을 시사한다.

AIME 2024 문제를 난이도별로 분석한 결과, 쉬운 문제(Easy)에서는 모든 방법론이 유사한 성능을 보였으나, 어려운 문제(Hard)로 갈수록 성능 격차가 벌어졌다. IMH 방식은 기존 베이스라인이 전혀 해결하지 못한 고난도 기하 및 대수 문제에서 올바른 좌표 기하학적 접근법을 찾아내는 등 추론의 경계를 확장하는 모습을 보였다.

본 연구는 dLLM의 디코딩을 전역 시퀀스 수준의 에너지 기반 모델(EBM) 샘플링 문제로 해석한다. 기존의 리마스킹 전략들이 'Confidence Gating'이라는 국소적 제약 조건 하에서 작동하며, 이것이 시퀀스 엔트로피 H(X)에 상한선(L * hV(δ))을 설정하여 탐색 능력을 물리적으로 제한함을 이론적으로 규명했다.

제안된 전역 템퍼링(Global Tempering)은 p*(x) ∝ q(x)^α 형태의 파워 분포를 타겟으로 한다. dLLM은 자기회귀 모델과 달리 정확한 시퀀스 가능도 q(x)를 계산할 수 없다는 한계가 있으나, 본 논문은 dLLM의 학습 객체인 조건부 독립 주변 확률 분포를 활용하여 이를 극복했다. 구체적으로, 조건부 주변 확률의 곱으로 결합 확률을 근사하는 Mean-field Factorization을 통해 룩어헤드 수정 항을 계산 가능한 형태로 유도했다.

IMH 샘플러의 효율성을 극대화하기 위해 Batched Evaluation 전략을 채택했다. 제안 분포가 마르코프 연쇄의 현재 상태와 독립적이므로, 여러 개의 후보 토큰(T개)에 대한 룩어헤드 수정을 단 한 번의 모델 순전파(Forward Pass)로 병렬 계산할 수 있다. 이를 통해 기존 dLLM 디코딩 대비 추가적인 순차적 오버헤드 없이 전역적 최적화를 수행한다.

실험적 분석을 통해 IMH가 생성한 추론 경로의 의미적 유사도를 측정한 결과, 기존 기법들이 특정 해결 방식에 고착(Collapse)되는 것과 달리 IMH는 질적으로 상이한 다양한 전략을 탐색함이 확인되었다. 이는 모델이 사전 학습 단계에서 습득한 잠재적 지식을 더 효과적으로 인출(Retrieval)할 수 있게 함을 의미한다.

본 논문은 룩어헤드 수정 항 계산 시 주변 확률의 독립성을 가정하는 평균장 근사를 사용하므로, 토큰 간의 강한 의존성이 존재하는 특정 문맥에서는 근사 오차가 발생할 수 있다. 또한, 전역 템퍼링 파라미터 α와 IMH의 제안 횟수 T에 따른 연산 비용과 성능 사이의 최적 트레이드오프 설정이 필요하다.