KL 기반 안정성 샘플링으로 빠른 마스크 확산 모델 추론

초록

마스크 확산 모델의 반복적인 복원 과정은 추론 속도를 크게 제한한다. 본 논문은 토큰‑레벨 KL 발산과 confidence 점수를 활용해 안정적인 토큰을 조기에 언마스크함으로써, 추가 학습 없이도 한 번에 여러 토큰을 복원하는 KL‑Adaptive Stability Sampling(KLASS)을 제안한다. KLASS는 기존 greedy·Top‑k 방식에 비해 최대 2.78배 빠른 추론을 달성하면서, 수학·코드·이미지·분자 생성 등 다양한 베치마크에서 품질을 향상시킨다.

상세 분석

본 논문은 마스크 확산 모델(Masked Diffusion Models, MDM)의 핵심 병목인 “느린 샘플링 속도”를 해결하기 위해 토큰‑레벨 KL 발산(Kullback‑Leibler divergence)과 confidence 점수를 동시에 활용하는 새로운 샘플링 전략을 제시한다. 기존의 Top‑k, stochastic, 혹은 planner 기반 방법들은 고정된 스케줄에 따라 매 단계마다 제한된 수의 토큰만을 언마스크하거나, 외부 모델을 도입해 추가 연산 비용을 발생시킨다. 반면 KLASS는 모델 자체가 출력하는 확률 분포의 변화를 직접 측정한다. 구체적으로, 각 토큰 i에 대해 현재 시점 t의 예측 분포 pᵢₜ와 이전 시점 t+1의 분포 pᵢₜ₊₁ 사이의 KL(pᵢₜ‖pᵢₜ₊₁)를 계산하고, 이 값이 사전에 정의된 임계값 ε_KL 이하이며 동시에 confidence(confᵢₜ = max_v pᵢₜ(v))가 τ 이상인 경우 해당 토큰을 “stable”하다고 판단한다. 이렇게 선정된 토큰 집합 Sₜ는 한 번에 언마스크되며, Sₜ가 비어 있을 경우에는 fallback으로 상위 u개의 confidence가 높은 토큰을 언마스크한다.

이 접근법은 두 가지 중요한 장점을 제공한다. 첫째, KL 점수는 단순 confidence보다 더 강력한 신뢰 지표이다. 논문에서는 올바른 답변이 잘못된 답변에 비해 평균 KL 값이 현저히 낮다는 실험적 증거를 제시한다(그림 1b). 이는 모델이 올바른 토큰에 대해 예측 분포가 시간에 따라 안정적으로 유지된다는 가정을 뒷받침한다. 둘째, 토큰을 병렬로 언마스크함으로써 전체 샘플링 단계 수를 크게 감소시킨다. 실험 결과, KLASS는 기존 greedy 또는 Top‑k 방식 대비 평균 2배 이상의 스텝을 절감하면서도, 정확도는 오히려 상승한다.

이론적 근거로는 “stable token must have low per‑step KL”라는 명제를 제시한다. 모델이 잘 학습된 경우, 잘못된 토큰은 컨텍스트가 점진적으로 해소됨에 따라 예측 분포가 크게 변동하므로 KL이 크게 된다. 반면 올바른 토큰은 초기부터 안정적인 확률 질량을 유지한다. 따라서 KL 기준을 도입하면 잘못된 토큰이 조기에 언마스크되는 위험을 방지하고, 전체 생성 품질을 보장한다.

실험은 수학 문제(GSM8K, MA​TH), 코드 생성(HumanEval, MBPP), 이미지 합성, 분자 설계 등 네 가지 도메인에 걸쳐 수행되었다. 표 1에 따르면, LLaDA와 Dream 모델에 KLASS를 적용했을 때 정확도는 각각 35%p 상승하고, 샘플링 스텝은 3050% 감소하였다. 특히 reasoning 베치마크에서는 2.78배의 wall‑clock 속도 향상을 기록했다. 또한, 기존 confidence‑only 기반 샘플러(Fast‑dLLM, Prophet 등)와 비교했을 때, KLASS는 낮은 perplexity와 높은 정확도를 동시에 달성한다는 점에서 실용적 우수성을 입증한다.

한계점으로는 KL 임계값 ε_KL와 confidence 임계값 τ를 도메인별로 튜닝해야 한다는 점이며, 매우 긴 시퀀스에서는 KL 계산 비용이 누적될 수 있다. 그러나 이는 사전 계산된 캐시나 히스토리 길이 n을 조절함으로써 완화 가능하다. 전반적으로 KLASS는 추가 학습 없이도 기존 MDM의 내부 신호를 활용해 효율성과 품질을 동시에 끌어올리는 실용적인 샘플링 프레임워크라 할 수 있다.