EntRGi : Orientation par Récompense Consciente de l'Entropie pour les Modèles de Langage à Diffusion

EntRGi: Entropy Aware Reward Guidance for Diffusion Language Models

February 4, 2026

Auteurs: Atula Tejaswi, Litu Rout, Constantine Caramanis, Sanjay Shakkottai, Sujay Sanghavi

cs.AI

Résumé

Le guidage par récompense a été appliqué avec un grand succès pour l'adaptation en temps de test des modèles de diffusion continus ; il met à jour chaque étape de dé-bruitage en utilisant les gradients d'un modèle de récompense en aval. Nous étudions le guidage par récompense pour les modèles de langage à diffusion discrète, où il est impossible de différencier les sorties naturelles du modèle car elles sont constituées de tokens discrets. Les approches existantes remplacent soit ces tokens discrets par des relaxations continues, soit emploient des techniques comme l'estimateur straight-through. Dans ce travail, nous montrons les inconvénients de ces deux méthodes. La première dégrade la rétroaction du gradient car le modèle de récompense n'a jamais été entraîné avec des entrées continues. La seconde implique une optimisation incorrecte car le gradient évalué sur des tokens discrets est utilisé pour mettre à jour des logits continus. Notre innovation clé consiste à dépasser ce compromis en introduisant un nouveau mécanisme appelé EntRGi : un guidage par récompense sensible à l'entropie qui régule dynamiquement les gradients provenant du modèle de récompense. En modulant la relaxation continue en utilisant la confiance du modèle, notre approche améliore considérablement le guidage par récompense tout en fournissant des entrées fiables au modèle de récompense. Nous validons empiriquement notre approche sur un modèle de langage à diffusion de 7 milliards de paramètres, en utilisant 3 modèles de récompense diversifiés et 3 benchmarks multi-compétences, montrant des améliorations constantes par rapport aux méthodes de l'état de l'art.

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.