EntRGi: Entropie-bewusste Belohnungssteuerung für Diffusions-Sprachmodelle
EntRGi: Entropy Aware Reward Guidance for Diffusion Language Models
February 4, 2026
Autoren: Atula Tejaswi, Litu Rout, Constantine Caramanis, Sanjay Shakkottai, Sujay Sanghavi
cs.AI
Zusammenfassung
Belohnungsgestützung wurde bei der Testzeit-Anpassung kontinuierlicher Diffusionsmodelle mit großem Erfolg eingesetzt; dabei wird jeder Denoising-Schritt mithilfe der Gradienten eines nachgeschalteten Belohnungsmodells aktualisiert. Wir untersuchen Belohnungsgestützung für diskrete Diffusions-Sprachmodelle, bei denen man nicht durch die natürlichen Ausgaben des Modells differenzieren kann, da es sich um diskrete Tokens handelt. Bestehende Ansätze ersetzen entweder diese diskreten Tokens durch kontinuierliche Relaxierungen oder verwenden Techniken wie den Straight-Through-Schätzer. In dieser Arbeit zeigen wir die Nachteile beider Methoden auf. Erstere verschlechtert die Gradientenrückmeldung, da das Belohnungsmodell nie mit kontinuierlichen Eingaben trainiert wurde. Letztere beinhaltet eine inkorrekte Optimierung, da der Gradient, der an diskreten Tokens ausgewertet wird, zur Aktualisierung kontinuierlicher Logits verwendet wird. Unsere zentrale Innovation besteht darin, über diesen Kompromiss hinauszugehen, indem wir einen neuartigen Mechanismus namens EntRGi einführen: Entropiebewusste Belohnungsgestützung, die die Gradienten aus dem Belohnungsmodell dynamisch reguliert. Durch die Modulation der kontinuierlichen Relaxierung mithilfe des Konfidenzniveaus des Modells verbessert unser Ansatz die Belohnungsgestützung erheblich und liefert gleichzeitig zuverlässige Eingaben für das Belohnungsmodell. Wir validieren unseren Ansatz empirisch an einem Diffusions-Sprachmodell mit 7B Parametern über 3 verschiedene Belohnungsmodelle und 3 Multi-Skill-Benchmarks hinweg und zeigen konsistente Verbesserungen gegenüber state-of-the-art Methoden.
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.