EntRGi: Guida alla Ricompensa Basata sull'Entropia per Modelli Linguistici di Diffusione
EntRGi: Entropy Aware Reward Guidance for Diffusion Language Models
February 4, 2026
Autori: Atula Tejaswi, Litu Rout, Constantine Caramanis, Sanjay Shakkottai, Sujay Sanghavi
cs.AI
Abstract
La guida basata su ricompensa è stata applicata con grande successo nell'adattamento al momento del test di modelli di diffusione continui; essa aggiorna ogni passo di denoising utilizzando i gradienti provenienti da un modello di ricompensa a valle. Studiamo la guida basata su ricompensa per modelli linguistici di diffusione discreta, dove non è possibile differenziare attraverso gli output naturali del modello poiché sono token discreti. Gli approcci esistenti sostituiscono questi token discreti con rilassamenti continui, oppure impiegano tecniche come lo straight-through estimator. In questo lavoro, mostriamo gli svantaggi di entrambi i metodi. Il primo degrada il feedback del gradiente perché il modello di ricompensa non è mai stato addestrato con input continui. Il secondo comporta un'ottimizzazione scorretta perché il gradiente valutato sui token discreti viene utilizzato per aggiornare i logit continui. La nostra innovazione chiave consiste nel superare questo compromesso introducendo un nuovo meccanismo chiamato EntRGi: Entropy aware Reward Guidance, che regola dinamicamente i gradienti provenienti dal modello di ricompensa. Modulando il rilassamento continuo utilizzando la confidenza del modello, il nostro approccio migliora sostanzialmente la guida basata su ricompensa fornendo al contempo input affidabili al modello di ricompensa. Convalidiamo empiricamente il nostro approccio su un modello linguistico di diffusione da 7 miliardi di parametri, utilizzando 3 diversi modelli di ricompensa e 3 benchmark multi-abilità, mostrando miglioramenti consistenti rispetto ai metodi allo stato dell'arte.
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.