엔트로피 인식 보상 가이드가 디스크리트 확산 언어 모델을 혁신한다

초록

본 논문은 디스크리트 확산 언어 모델( dLLM )의 테스트‑타임 적응을 위해, 보상 모델의 그래디언트를 활용하는 새로운 방법인 EntRGi(Entropy‑aware Reward Guidance)를 제안한다. EntRGi는 모델의 예측 엔트로피를 이용해 연속 완화(soft embedding)와 하드 토큰(샘플링된 토큰) 사이를 동적으로 가중합함으로써, 기존의 연속 완화 방식이 초래하는 입력‑분포 불일치와 STE(직접 전파) 방식이 초래하는 로그‑그라디언트 불일치를 동시에 완화한다. 7B 규모의 확산 언어 모델에 3가지 보상 모델과 3가지 멀티스킬 벤치마크를 적용한 실험에서, EntRGi는 최신 방법들을 일관되게 능가한다.

상세 분석

EntRGi의 핵심 아이디어는 “엔트로피 기반 가중합”이다. 디스크리트 확산 과정에서 각 마스크된 위치 l에 대해 현재 로그잇 ψ_l 로부터 소프트맥스 확률 q_l 를 구하고, 이를 이용해 두 종류의 임베딩을 만든다. 첫 번째는 확률 q_l 로 가중합한 연속형 임베딩 (\bar e_l = \sum_i q_{l,i} E^R_i) 로, 이는 보상 모델이 훈련된 연속 입력 공간에 가깝다. 두 번째는 q_l 로 샘플링한 하드 토큰 x_l 에 대한 임베딩 (\tilde e_l = E^R_{x_l}) 로, 이는 실제 토큰 형태이지만 그래디언트 전파 시 STE를 사용하면 로그와 임베딩 사이에 불일치가 발생한다. EntRGi는 엔트로피 (H(q_l)) 를 정규화한 가중치 (w_l = H(q_l)/\log K) 를 정의하고, 최종 입력 (\hat e_l = \bar e_l + \text{sg}(w_l)(\tilde e_l - \bar e_l)) 로 구성한다. 여기서 sg는 stop‑gradient 연산으로, 그래디언트는 오직 (\bar e_l) 를 통해서만 흐른다. 따라서 엔트로피가 낮아 확률이 한 토큰에 집중될 때는 (w_l) 가 작아 연속 완화가 주도되고, 엔트로피가 높아 불확실성이 클 때는 (w_l) 가 커져 하드 토큰의 정보를 반영한다. 이 메커니즘은 두 기존 접근법의 장점을 동적으로 결합하면서, 보상 모델이 “알아볼 수 없는” 연속 입력에 노출되는 위험을 최소화한다.

수식적으로 보면, 그래디언트는 (\nabla_{\psi_l} R(\hat e) = \frac{\partial R}{\partial \hat e_l} E^{R\top} J_{\text{softmax}}(\psi_l)) 로 전파된다. 여기서 (\frac{\partial R}{\partial \hat e_l}) 은 보상 모델의 임베딩‑레벨 민감도를 나타내며, (J_{\text{softmax}}) 는 소프트맥스의 야코비안이다. stop‑gradient 덕분에 (\tilde e_l) 에 대한 직접적인 미분은 차단되므로, 보상 모델이 하드 토큰을 평가하더라도 그래디언트는 연속 임베딩 공간에 머문다. 이 설계는 “평가‑전파 불일치”를 수학적으로 제한하고, 근사 오차 (E_l = w_l | \tilde e_l - \bar e_l |) 를 엔트로피에 비례하도록 만든다. 엔트로피가 0에 가까워지면 (E_l) 도 0에 수렴해 연속 완화가 완벽히 정확해지고, 엔트로피가 높을 때는 STE가 보완 역할을 수행한다.

실험에서는 7B 파라미터 dLLM(Ye et al., 2025)을 사용했으며, 보상 모델로는 Skywork‑Reward, Qwen‑Reward, 그리고 자체 구축한 인간 피드백 기반 모델을 채택했다. 3가지 멀티스킬 벤치마크(코드 생성, 요약, 대화)에서 EntRGi는 기존의 APS(연속 완화 + STE)와 BoN(샘플링 기반) 방법보다 평균 4.2%~7.5% 포인트 높은 보상 점수를 기록했다. 특히 초기 디노이징 단계에서 엔트로피가 높은 경우에 EntRGi가 보상 신호를 효과적으로 활용해 토큰 선택을 조정하는 것이 두드러졌다. 추가적인 ablation 연구에서는 (w_l) 를 고정하거나 엔트로피 기반 가중합을 제거했을 때 성능이 급격히 하락함을 확인했으며, 이는 제안 기법이 엔트로피 정보를 핵심적으로 활용한다는 증거다.

요약하면, EntRGi는 (1) 엔트로피를 정량적 신뢰도 지표로 사용해 연속‑이산 간의 트레이드오프를 동적으로 조정하고, (2) 보상 모델 입력을 “신뢰할 수 있는” 형태로 유지하면서도 (3) 그래디언트 흐름을 연속 공간에 제한함으로써 최적화 오류를 최소화한다. 이러한 설계는 훈련‑프리 방식이면서도 대규모 dLLM에 적용 가능하다는 점에서 실용적 가치가 크다.