EntRGi: Guía de Recompensa Basada en Entropía para Modelos de Lenguaje de Difusión
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
Resumen
La guía por recompensa se ha aplicado con gran éxito en la adaptación durante la inferencia de modelos de difusión continuos; actualiza cada paso de eliminación de ruido utilizando los gradientes de un modelo de recompensa auxiliar. Estudiamos la guía por recompensa para modelos de lenguaje de difusión discreta, donde no es posible diferenciar a través de las salidas naturales del modelo porque son tokens discretos. Los enfoques existentes o bien reemplazan estos tokens discretos con relajaciones continuas, o emplean técnicas como el estimador de paso directo. En este trabajo, mostramos las desventajas de ambos métodos. El primero degrada la retroalimentación del gradiente porque el modelo de recompensa nunca ha sido entrenado con entradas continuas. El segundo implica una optimización incorrecta porque el gradiente evaluado en tokens discretos se utiliza para actualizar los logits continuos. Nuestra innovación clave es superar esta disyuntiva introduciendo un mecanismo novedoso llamado EntRGi: Guía por Recompensa con Conciencia de la Entropía, que regula dinámicamente los gradientes del modelo de recompensa. Al modular la relajación continua utilizando la confianza del modelo, nuestro enfoque mejora sustancialmente la guía por recompensa mientras proporciona entradas confiables al modelo de recompensa. Validamos empíricamente nuestro enfoque en un modelo de lenguaje de difusión de 7 mil millones de parámetros a través de 3 modelos de recompensa diversos y 3 benchmarks de habilidades múltiples, mostrando mejoras consistentes sobre los métodos más avanzados.
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.