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

要旨

報酬ガイダンスは、連続拡散モデルのテスト時適応において大きな成功を収めており、下流の報酬モデルからの勾配を用いて各ノイズ除去ステップを更新する。本研究では、離散拡散言語モデルに対する報酬ガイダンスに着目する。離散トークンがモデルの自然な出力であるため、これらのモデルでは出力を微分することができない。既存手法では、離散トークンを連続緩和で置き換えるか、Straight-through Estimatorなどの技術を採用している。本研究では、これらの両手法に欠点があることを示す。前者は、報酬モデルが連続入力を用いて訓練されていないため、勾配フィードバックが劣化する。後者は、離散トークンで評価された勾配を連続ロジットの更新に用いるため、不正確な最適化を伴う。我々の核心的な革新は、このトレードオフを超える新規メカニズム「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.