EntRGi: Orientação por Recompensa com Consciência de Entropia para Modelos de Linguagem de Difusão
EntRGi: Entropy Aware Reward Guidance for Diffusion Language Models
February 4, 2026
Autores: Atula Tejaswi, Litu Rout, Constantine Caramanis, Sanjay Shakkottai, Sujay Sanghavi
cs.AI
Resumo
A orientação por recompensa tem sido aplicada com grande sucesso na adaptação em tempo de teste de modelos de difusão contínua; ela atualiza cada etapa de remoção de ruído usando os gradientes de um modelo de recompensa downstream. Estudamos a orientação por recompensa para modelos de linguagem de difusão discreta, onde não é possível diferenciar as saídas naturais do modelo porque elas são tokens discretos. As abordagens existentes ou substituem esses tokens discretos por relaxações contínuas, ou empregam técnicas como o estimador straight-through. Neste trabalho, mostramos as desvantagens de ambos os métodos. O primeiro degrada o feedback do gradiente porque o modelo de recompensa nunca foi treinado com entradas contínuas. O segundo envolve uma otimização incorreta porque o gradiente avaliado em tokens discretos é usado para atualizar logits contínuos. Nossa principal inovação é ir além desse tradeoff introduzindo um novo mecanismo chamado EntRGi: Orientação por Recompensa com Consciência de Entropia, que regula dinamicamente os gradientes do modelo de recompensa. Ao modular a relaxação contínua usando a confiança do modelo, nossa abordagem melhora substancialmente a orientação por recompensa enquanto fornece entradas confiáveis para o modelo de recompensa. Validamos empiricamente nossa abordagem em um modelo de linguagem de difusão de 7 bilhões de parâmetros, abrangendo 3 modelos de recompensa diversos e 3 benchmarks de múltiplas habilidades, mostrando melhorias consistentes em relação aos métodos state-of-the-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.