EntRGi: Entropiebewuste Beloningsgeleiding voor Diffusie Taalmodellen

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

Samenvatting

Beloningsgeleiding is met groot succes toegepast bij de testtijd-aanpassing van continue diffusiemodellen; hierbij wordt elke denoiseringsstap bijgewerkt met behulp van de gradienten van een downstream beloningsmodel. Wij bestuderen beloningsgeleiding voor discrete diffusie-taalmodellen, waarbij men niet kan differentiëren door de natuurlijke uitvoeren van het model omdat deze discrete tokens zijn. Bestaande benaderingen vervangen deze discrete tokens ofwel door continue relaxaties, of maken gebruik van technieken zoals de straight-through estimator. In dit werk tonen we de nadelen van beide methoden aan. De eerste degradeert de gradientfeedback omdat het beloningsmodel nooit is getraind met continue invoeren. De tweede houdt een incorrecte optimalisatie in omdat de gradient, geëvalueerd op discrete tokens, wordt gebruikt om continue logits bij te werken. Onze belangrijkste innovatie is om voorbij deze afweging te gaan door een nieuw mechanisme te introduceren genaamd EntRGi: Entropie-bewuste Beloningsgeleiding, dat de gradienten van het beloningsmodel dynamisch reguleert. Door de continue relaxatie te moduleren met behulp van het vertrouwen van het model, verbetert onze aanpak de beloningsgeleiding aanzienlijk en biedt het tegelijkertijd betrouwbare invoeren voor het beloningsmodel. We valideren onze aanpak empirisch op een diffusie-taalmodel met 7B parameters, over 3 diverse beloningsmodellen en 3 multi-vaardigheid benchmarks, en tonen consistente verbeteringen aan ten opzichte van 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.