EntRGi: 확산 언어 모델을 위한 엔트로피 인식 보상 가이던스
EntRGi: Entropy Aware Reward Guidance for Diffusion Language Models
February 4, 2026
저자: Atula Tejaswi, Litu Rout, Constantine Caramanis, Sanjay Shakkottai, Sujay Sanghavi
cs.AI
초록
보안 가이던스는 연속 확산 모델의 테스트 타임 적응에 큰 성공을 거두며 적용되어 왔으며, 이는 다운스트림 보상 모델의 그래디언트를 사용하여 각 노이즈 제거 단계를 업데이트합니다. 본 연구는 이산 확산 언어 모델에 대한 보상 가이던스를 다루며, 여기서는 모델의 자연스러운 출력이 이산 토큰이기 때문에 이를 직접 미분할 수 없다는 문제가 있습니다. 기존 접근법은 이 이산 토큰을 연속 완화로 대체하거나, 스트레이트-스루 추정기와 같은 기법을 사용합니다. 본 연구에서는 이 두 방법의 단점을 보여줍니다. 전자의 경우 보상 모델이 연속 입력으로 훈련된 적이 없기 때문에 그래디언트 피드백이 저하됩니다. 후자의 경우 이산 토큰에서 평가된 그래디언트가 연속 로짓을 업데이트하는 데 사용되므로 부정확한 최적화가 발생합니다. 우리의 핵심 혁신은 이러한 절충점을 넘어서 EntRGi라는 새로운 메커니즘, 즉 엔트로피 인식 보안 가이던스를 도입하여 보상 모델의 그래디언트를 동적으로 조절한다는 점입니다. 모델의 신뢰도를 사용하여 연속 완화를 조절함으로써, 우리의 접근법은 보상 모델에 신뢰할 수 있는 입력을 제공하면서도 보안 가이던스를 크게 개선합니다. 우리는 7B 매개변수 확산 언어 모델을 대상으로 3가지 다양한 보상 모델과 3가지 다중 기술 벤치마크에 걸쳐 제안 방법을 실증적으로 검증하며, 최첨단 방법 대비 일관된 성능 향상을 보여줍니다.
English
Reward guidance has been applied to great success in the test-time adaptation of continuous diffusion models; it updates each denoising step using the gradients from a downstream reward model. We study reward guidance for discrete diffusion language models, where one cannot differentiate through the natural outputs of the model because they are discrete tokens. Existing approaches either replace these discrete tokens with continuous relaxations, or employ techniques like the straight-through estimator. In this work, we show the downsides of both these methods. The former degrades gradient feedback because the reward model has never been trained with continuous inputs. The latter involves incorrect optimization because the gradient evaluated at discrete tokens is used to update continuous logits. Our key innovation is to go beyond this tradeoff by introducing a novel mechanism called EntRGi: Entropy aware Reward Guidance that dynamically regulates the gradients from the reward model. By modulating the continuous relaxation using the model's confidence, our approach substantially improves reward guidance while providing reliable inputs to the reward model. We empirically validate our approach on a 7B-parameter diffusion language model across 3 diverse reward models and 3 multi-skill benchmarks, showing consistent improvements over state-of-the-art methods.