전망 경로 가능도 최적화로 확산형 LLM 성능 향상

초록

본 논문은 확산 대형 언어 모델(dLLM)의 비정형 마스킹 순서가 생성 품질에 미치는 영향을 분석하고, 경로 로그가능도(Path LL)를 새로운 품질 지표로 제안한다. 부분 디코딩 상태에서 향후 기대 Path LL을 예측하는 POKE(Value Estimator)를 설계하고, 이를 순차적 몬테카를로(Sequential Monte Carlo, SMC) 탐색에 결합한 POKE‑SMC 프레임워크를 제시한다. 6개의 추론·추리 벤치마크에서 기존 휴리스틱 기반 전략보다 2~3%p(평균) 높은 정확도를 달성하면서도 추론 비용을 크게 늘리지 않는다.

상세 분석

이 연구는 확산 기반 언어 모델이 “임의 순서 마스킹”이라는 특성을 갖는 점에 주목한다. 기존 작업에서는 마스크 해제 순서를 confidence, entropy, margin 등 로컬 불확실성 지표에 기반한 휴리스틱으로 정했지만, 이러한 방법은 작업마다 성능 편차가 크고 전역적인 품질을 보장하지 못한다는 한계가 있었다. 저자들은 이를 해결하기 위해 “경로 로그가능도(Path LL)”라는 새로운 전역 목표 함수를 정의한다. Path LL은 특정 마스킹·디코딩 순서 τ에 대해 전체 토큰 시퀀스 x의 로그가능도를 단계별 로그가능도의 합으로 표현한다(식 4). 이는 기존 ELBO와 달리 순서 의존성을 명시적으로 반영하므로, 실제 downstream 정확도와 높은 상관관계를 보인다(그림 1 우측).

하지만 Path LL을 직접 최적화하려면 모든 가능한 순열을 탐색해야 하는 조합 폭이 기하급수적으로 커진다. 따라서 저자들은 “부분 디코딩 상태 x_O”에서 남은 마스크 토큰 M에 대해 무작위 순열 σ를 사용한 롤아웃을 통해 기대 Path LL을 추정하는 가치 함수 V(x_O)를 정의한다(식 5). 이때 V는 실제 조인트 분포 pθ와 독립 마진 곱분포 qθ 사이의 차이인 Total Correlation(TC)을 포함한다(식 7). TC는 토큰 간 상호 의존성을 정량화하지만 직접 계산이 불가능하다.

핵심 기여는 TC를 상한으로 근사하는 “Optimistic Approximation”이다. 무작위 순열이 토큰 의존성을 크게 약화시킨다는 이론적 결과(Anari et al., 2025)를 이용해, 기대 TC를 각 마스크 토큰의 마진 엔트로피 평균으로 제한한다(식 10). 이를 바탕으로 K‑step(예: K=4~32)으로 나눈 단계별 롤아웃을 수행하고, 각 단계에서 마진 로그가능도와 마진 엔트로피를 합산해 POKE 추정기(식 11)를 만든다. POKE는 단일 전방 패스로 모든 마스크 토큰의 마진 확률을 얻고, 추가적인 K‑step 롤아웃을 통해 전체 Path LL을 높은 정확도로 예측한다. 실험 결과, POKE는 기존 “Product LL”가 크게 과소평가하던 값을 거의 복원한다(그림 2).

이후 POKE 값을 탐색 가이드로 활용해 Sequential Monte Carlo(SMC) 기반 탐색기인 POKE‑SMC를 설계한다. 입자(particle)들은 현재 부분 경로와 POKE 가치에 따라 가중치를 부여받으며, 재샘플링 단계에서 높은 가치의 입자를 집중시킨다. 이렇게 동적으로 최적 경로를 탐색함으로써, 정적 휴리스틱보다 전역적인 품질을 지속적으로 개선한다.

실험에서는 LLaDA 모델군(다양한 규모)과 6개의 추론·추리 데이터셋(HumanEval, GSM8K, Sudoku 등)을 사용했다. 동일한 연산 예산(예: 토큰당 48개의 샘플) 하에서 POKE‑SMC는 기존 “Semi‑AR”, “Margin”, “Confidence” 기반 전략보다 평균 23%p, 최고 5%p까지 정확도 향상을 기록했다. 또한, 연산량 대비 정확도 향상 곡선(accuracy‑compute Pareto frontier)에서도 기존 디코딩‑시간 스케일링 방법들을 앞선다.

이 논문의 의의는 세 가지로 정리할 수 있다. 첫째, dLLM에서 전역 품질을 정량화할 수 있는 Path LL이라는 새로운 지표를 제시함으로써, 모델 자체가 순서에 민감함을 명확히 밝혔다. 둘째, TC 기반의 낙관적 상한을 이용해 복잡한 조인트 확률을 효율적으로 근사하는 POKE 추정기를 설계해, 실시간 가치 추정이 가능하도록 했다. 셋째, 이러한 가치 추정을 SMC와 결합해 동적 경로 탐색을 수행함으로써, 기존 정적 휴리스틱의 한계를 넘어선 실용적인 디코딩 전략을 제공한다. 앞으로는 POKE‑SMC를 다른 생성형 모델(예: 이미지·음성 확산 모델)이나 더 큰 스케일의 LLM에 적용하는 연구가 기대된다.