💬 NLP

Entropy Aware Reward Guidance for Diffusion Language Model Alignment

This paper introduces EntRGi, a novel method that dynamically interpolates between continuous token relaxations and hard tokens based on predictive entropy to enable effective reward guidance for discrete diffusion language models, demonstrating consistent improvements over state-of-the-art approaches in both test-time adaptation and post-training via Reward Guided Reinforcement Learning (RGRL).

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

Published 2026-05-14

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very talented but slightly confused artist named Diffusion. This artist is trying to paint a picture (or write a story) by starting with a canvas covered entirely in static noise (like a TV with no signal) and slowly clearing away the noise to reveal the image.

Usually, this artist works in a "continuous" world, like mixing paints where you can blend any shade of blue. But in this paper, the artist is working with discrete tokens—think of them as distinct, separate Lego bricks. You can't blend a "red" brick with a "blue" brick to make "purple"; you have to pick one or the other.

The problem the authors tackle is how to tell this artist, "No, not that red brick! Make it a better red brick," using a Reward Model (a critic who judges the final picture).

The Core Problem: The "Translator" Breakdown

In the old way of doing things (continuous models), the artist and the critic speak the same language. The critic can say, "Move your brush slightly to the left," and the artist understands exactly how to adjust.

But with these Lego-brick models, the critic only knows how to judge finished, hard bricks.

The "Expectation" Method: Some tried to give the critic a "blurred" version of the bricks (a mix of red and blue) so the artist could get smooth instructions. But the critic was trained only on sharp, real bricks. It's like asking a food critic to judge a smoothie made of blended carrots and celery; they don't know what to say because they've never tasted that before. The feedback is unreliable.
The "APS" Method (Previous Best): Others tried to show the critic a real, sharp brick to get a good score, but then told the artist to pretend the brick was still blurred so they could get smooth instructions. This is like showing a judge a perfect steak, but then telling the chef, "Imagine this steak is a liquid soup" while giving instructions. The instructions don't match the reality, leading to confusion.

The Solution: EntRGi (The "Confidence Meter")

The authors introduce EntRGi (Entropy-Aware Reward Guidance). Think of this as a smart Confidence Meter for the artist.

As the artist works, they are sometimes very sure about which brick to pick (low entropy/confidence is high), and sometimes they are totally guessing (high entropy/confidence is low).

When the Artist is Confident: The Confidence Meter says, "You know what you're doing! Show the critic the real, hard brick." This ensures the critic gives a reliable, accurate score because they are judging something they understand.
When the Artist is Guessing: The Confidence Meter says, "You're unsure! Let's show the critic a blurred mix of possibilities." This allows the artist to get smooth, continuous instructions on how to improve without breaking the flow of the creative process.

The Magic: EntRGi automatically switches between these two modes, brick by brick, based on how sure the artist is. It gets the best of both worlds: reliable feedback from the critic and smooth, accurate instructions for the artist.

The Result: Better Art and New Training

The paper shows that this method works better than previous attempts at two levels:

Test-Time Adaptation (The "Live Edit"): When the artist is painting a picture, EntRGi helps them steer the process to create a better final image without needing to retrain the artist from scratch. It's like having a director on set who knows exactly when to give a gentle nudge and when to let the actor improvise.
RGRL (Reward Guided Reinforcement Learning): The authors also created a new training recipe called RGRL. Instead of just showing the artist a finished picture and saying "Good job" or "Bad job" (which is vague), they use the smooth feedback from EntRGi to teach the artist how to get better. It's like a coach giving specific, detailed drills based on the player's current form, rather than just keeping score.

The Bottom Line

The paper claims that by using this "Confidence Meter" (EntRGi), they can guide these discrete Lego-brick models to produce higher-quality results than before. They tested this on large models (7 billion parameters) and found it consistently beats the previous best methods, both when just guiding the model during generation and when using it to train the model to be smarter.

They also noted that this method works even when the artist and the critic use slightly different "vocabularies" (different sets of Lego bricks), which is a common real-world problem.

In short: They figured out how to talk to a confused Lego artist in a way that keeps the critic happy and the artist learning, resulting in better stories and images.

Technical Summary: Entropy Aware Reward Guidance for Diffusion Language Model Alignment

Problem Statement

Reward guidance (or posterior sampling) has proven highly effective for continuous diffusion models, where differentiable feedback from a reward model allows for iterative refinement of generation steps. However, applying this paradigm to Discrete Diffusion Large Language Models (dLLMs) presents a fundamental challenge: dLLMs generate discrete tokens, which are non-differentiable.

Existing approaches attempt to bridge this gap but face a critical trade-off between reward model reliability and gradient accuracy:

Continuous Relaxation (Expectation): Feeding the reward model "soft" embeddings (expected values of token probabilities) enables gradient flow. However, reward models are typically trained on hard tokens; thus, soft, out-of-vocabulary inputs lead to unreliable reward signals and degraded performance.
Straight-Through Estimator (STE): Methods like APS (Rout et al., 2025c) evaluate the reward on sampled hard tokens (ensuring reliability) but propagate gradients as if the input were soft. This creates a mismatch between where the reward is evaluated and where gradients are applied, leading to approximation errors, particularly when the model's predictive entropy is high (i.e., during early denoising steps).

Methodology: EntRGi

The authors introduce EntRGi (Entropy-aware Reward Guidance), a mechanism that dynamically interpolates between continuous token relaxations and sampled hard tokens on a token-by-token basis. The core innovation is using the diffusion model's own per-token predictive entropy to determine the interpolation weight.

The Mechanism

At each masked position $l$ during the denoising process, EntRGi constructs an input embedding $e_l^{EntRGi}$ for the reward model as a convex combination:
$e_l^{EntRGi} = (1 - w_l)e_l^{soft} + w_l e_l^{hard}$
Where:

$e_l^{soft}$ is the expected soft embedding derived from the model's probability distribution.
$e_l^{hard}$ is the embedding of a sampled hard token.
$w_l$ is an adaptive weight derived from the token's entropy $H(q_l)$ , specifically $w_l = H(q_l) / \log K$ .

Logic of Adaptation:

Low Entropy (High Confidence): When the model is confident (low entropy), $w_l \to 0$ . The input is primarily the soft embedding, allowing for clean gradient flow without significant approximation error because the soft and hard embeddings are close.
High Entropy (High Uncertainty): When the model is uncertain (high entropy), $w_l \to 1$ . The input shifts toward the hard token embedding. This ensures the reward model receives realistic, in-distribution inputs, preserving reward reliability even though the gradient approximation error is higher.

This approach resolves the tension by trusting soft embeddings when the model is confident and reverting to discrete tokens when uncertainty is high.

Post-Training: RGRL

The paper also proposes RGRL (Reward Guided Reinforcement Learning), a post-training recipe that leverages EntRGi (or APS) to generate reward-guided completions. These completions are used to fine-tune the dLLM via self-distillation, creating a dense reward-gradient feedback loop rather than relying solely on scalar rewards used in standard RL methods like diffu-GRPO.

Key Contributions

EntRGi: A novel entropy-aware reward guidance mechanism for discrete diffusion LLMs that adaptively balances gradient accuracy and reward model reliability.
Test-Time Adaptation Superiority: Empirical validation showing EntRGi outperforms the prior state-of-the-art (APS) and other baselines (Expectation, Best-of-N) across multiple benchmarks, particularly at higher sampling temperatures where entropy is higher.
RGRL Post-Training: A new training recipe that utilizes reward-guided samples for fine-tuning. The authors demonstrate that RGRL-EntRGi significantly outperforms diffu-GRPO (a widely adopted RL algorithm for dLLMs), achieving up to a 70% relative improvement in reward scores on specific datasets.
Robustness to Tokenizer Mismatch: The method remains effective even when the dLLM and reward model have mismatched tokenizers (e.g., LLaDA vs. Qwen/Llama), a common real-world scenario where existing gradient-based methods often struggle.

Experimental Results

The authors evaluated EntRGi on 7B-parameter models (Dream-v0-Instruct-7B and LLaDA-8B-Instruct) across five multi-skill datasets and four reward models.

Inference-Time Performance: EntRGi consistently achieved higher Top@1 and Avg@4 reward scores compared to APS and Expectation. Notably, while APS performance degraded significantly at higher temperatures ( $\tau=0.7$ ), EntRGi maintained or improved performance, demonstrating its ability to handle high-entropy states effectively.
Post-Training Efficiency: RGRL-EntRGi showed faster convergence and higher final reward scores compared to diffu-GRPO. The gains were most pronounced on datasets with lower initial rewards (e.g., WildChat-IF), suggesting the method is particularly effective for difficult alignment tasks.
Ablation Studies: Analysis confirmed that the entropy-based weighting is superior to alternative heuristics (like L2-norm distance) and that the method reduces approximation error during early denoising steps where entropy is highest.

Significance

The paper claims that EntRGi provides a principled solution to the "discrete vs. differentiable" conflict in diffusion language models. By leveraging model uncertainty as a signal to regulate feedback incorporation, EntRGi enables:

Stronger Inference-Time Steering: Allowing dLLMs to optimize task-specific objectives without retraining, with reliability comparable to continuous diffusion models.
Viable Post-Training Alternatives: Unlocking dense reward-gradient feedback as a more effective alternative to supervised fine-tuning and scalar-reward RL for dLLMs.

The authors position this work as a necessary step toward making reward-guided adaptation and post-training robust and scalable for the emerging class of discrete diffusion language models.