PaperlayerView exact revision on arXiv

Source-linked AI summary

Entropy Aware Reward Guidance for Diffusion Language Model Alignment

Atula Tejaswi, Litu Rout, Constantine Caramanis, Sanjay Shakkottai, Sujay Sanghavi

arXiv:2602.05000v2cs.LGcs.AIcs.CL

TL;DR

Discrete outputs make it difficult to apply reward-guidance gradient feedback to diffusion language models. EntRGi interpolates between continuous token relaxations and hard tokens according to predictive entropy. The paper reports that EntRGi outperforms APS in test-time adaptation and that RGRL outperforms diffu-GRPO by 70% relative improvement.

  • Problem

    Discrete-token outputs make it challenging to transfer reward-guidance gradient feedback from continuous diffusion to diffusion language models.

    Source: §1 · p. 2

  • Method

    EntRGi dynamically interpolates between continuous relaxations and hard token embeddings based on the diffusion model’s predictive entropy.

    Source: §5 · p. 12

  • Results

    RGRL significantly outperforms diffu-GRPO by 70% relative improvement, while EntRGi outperforms APS in test-time adaptation.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results suggest model uncertainty is a useful signal for regulating reward feedback in inference-time steering and post-training of dLLMs.

    Source: §5 · p. 13

    The paper identifies a mismatch between discrete outputs and continuous representations for gradient propagation, most pronounced during high-entropy early denoising steps.

    Source: §2 · p. 3

Abstract

from arXiv · show

Reward guidance, also known as posterior sampling, is a popular method for test-time adaptation and post-training in continuous diffusion models. In this paper, we study reward guidance for discrete diffusion language models; now, one cannot differentiate through the natural outputs of the model because they are discrete tokens. We introduce a novel mechanism called EntRGi (Entropy aware Reward Guidance) to address this issue. EntRGi dynamically interpolates between continuous token relaxations and sampled hard tokens, on a token-by-token basis, using the diffusion model's predictive entropy. We demonstrate that EntRGi maintains both reward model reliability and optimization accuracy, while existing approaches sacrifice one for the other. We empirically validate our approach on 7B-parameter diffusion language models across two settings: (1) test-time adaptation, and (2) RGRL (Reward Guided Reinforcement Learning), our recipe for post-training on reward-guided data, showing consistent improvements over state-of-the-art methods. Our code is available at https://atutej.github.io/entrgi-rgrl

1 Introduction

The introduction frames reward guidance for discrete diffusion language models as a way to use reward feedback during iterative denoising, but discrete tokens create a tension between gradient accuracy and reward-model reliability. EntRGi addresses this by adapting soft-versus-hard token inputs to per-token entropy, and RGRL uses reward-guided samples for post-training.

1 Introduction

  • Unlike autoregressive LLMs, dLLMs iteratively denoise fully masked sequences in parallel, enabling particle-filtering-style guidance that rejects noisy completions with low reward.The cited approaches include Dang et al. (2025), Ou et al. (2025), and Singhal et al. (2025).Source: §1 · pp. 1–2
  • Discrete tokens prevent direct gradient propagation, while soft out-of-vocabulary inputs can hurt reward-model performance and APS (Rout et al., 2025c) evaluates hard tokens but applies soft-input gradients.The resulting mismatch trades off gradient accuracy against reward-model reliability, limiting inference-time steering and dense reward-gradient post-training.Source: §1 · p. 2
  • EntRGi uses per-token predictive entropy to favor soft token embeddings when the dLLM is confident and sampled hard-token embeddings when it is uncertain.The interpolation aims to provide reliable optimization gradients while keeping reward-model inputs interpretable throughout denoising.Source: §1 · p. 2
  • EntRGi outperforms APS in test-time adaptation, while RGRL achieves a 70% relative improvement over diffu-GRPO, with higher gains from EntRGi guidance.RGRL is presented as a post-training recipe that fine-tunes on reward-gradient-guided samples rather than using standard RL methods that do not guide generation.Source: §1 · p. 2
  • RGRL generates reward-gradient-guided samples and fine-tunes on them, unlike standard RL methods that do not guide generation.The recipe works with both APS and EntRGi guidance, with EntRGi providing higher gains.Source: §1 · p. 2
  • The experiments use up to two 7B+ parameter models, five multi-skill datasets, and four reward models to analyze improvements over prior methods..Source: §1 · p. 2

2 Related Work

Related work applies reward guidance to discrete diffusion through gradient-free sampling and training or reward-gradient methods, but both continuous relaxations and straight-through estimation face a mismatch with discrete outputs. This work addresses the mismatch by using model uncertainty to regulate gradients and reward-guided completions for dense training feedback.

Discrete diffusion posterior sampling

  • Discrete diffusion generates predictive distributions over all tokens in parallel, enabling posterior sampling under external constraints such as reward models without retraining or task-specific fine-tuning.Prior work develops this setting for categorical sequences (Dang et al., 2025; Rout et al., 2025c).Source: §2 · p. 3

Reward-gradient-free methods

Reward-gradient-based methods

  • Reward-gradient methods are well studied for continuous diffusion but less developed for discrete diffusion, where relaxations query rewards out of distribution and STE evaluates gradients at sampled hard tokens.APS (Rout et al., 2025c) is prior state-of-the-art for inference-time dLLM steering, while DRAKES (Wang et al., 2025) uses STE for biological-sequence diffusion fine-tuning; relaxation methods include Murata et al. (2024) and Tae et al. (2025).Source: §2 · p. 3

Challenges and limitations

  • Both relaxation and STE methods mismatch discrete outputs with continuous gradient representations, especially early in denoising when token entropy is high; this work uses uncertainty to regulate gradients and completions for dense feedback.The paper describes uncertainty-based adaptive gradient regulation at inference as the first such approach in dLLMs, and uses reward-guided completions at training time beyond scalar-reward RL.Source: §2 · p. 3

3 Reward Guidance for Discrete Diffusion LLMs

This section formulates reward-guided masked diffusion generation and introduces EntRGi, which uses predictive entropy to balance reward-model input reliability against gradient accuracy. It also describes RGRL self-distillation and handling different tokenizers.

3 Reward Guidance for Discrete Diffusion LLMs

  • Masked diffusion generation starts from an all-mask sequence and repeatedly unmasks selected positions, with unmasked tokens committed and the lowest-entropy positions selected in the cited models.The described selection rule is used in Dream-v0-Instruct-7B and LLaDA-8B-Instruct.Source: §3 · p. 4

3.1 Entropy Aware Reward Guidance

  • Expectation methods (Murata et al., 2024) risk unreliable feedback off-token, while APS (Rout et al., 2025c) uses hard inputs but incurs entropy-growing gradient approximation error.The two approaches trade off reward reliability and gradient accuracy, motivating EntRGi’s interpolation.Source: §3.1 · p. 5
  • EntRGi guides masked-token logits using reward-model inputs that interpolate per token between expected soft embeddings and sampled hard-token embeddings, weighted by predictive entropy.The interpolation weight is H(q_l)/log K: it recovers the expectation method at zero weight and APS at weight one.Source: §3.1 · pp. 5–6
  • EntRGi strictly reduces APS approximation error at every entropy level, while keeping vocabulary error small at low entropy and driving it toward zero at high entropy to maintain gradient reliability.The high-entropy limit follows because the interpolation weight approaches one, making the reward-model input a hard token embedding.Source: §3.1 · p. 6
  • RGRL self-distills reward-guided completions by sampling N completions per prompt under current parameters and updating the diffusion model to increase their likelihood.Using EntRGi or APS for generation defines RGRL-EntRGi or RGRL-APS, respectively.Source: §3.1 · p. 7
  • For mismatched tokenizers, EntRGi assigns zero reward embeddings to diffusion-model tokens absent from the reward vocabulary, so only matched tokens contribute to gradients.The formulation allows the two vocabularies to overlap only partially.Source: §3.1 · p. 7

4 Experiments

Across test-time adaptation and reward-guided post-training, EntRGi improves reward-guided generation, while analysis shows how its entropy-aware token weighting reduces early denoising approximation error. Its gains also persist under tokenizer mismatch and vary with reward-model scale and optimization steps.

4.1 Test-Time Adaptation

  • EntRGi yields approximately 33% relative improvement over APS in reward-model-judged output quality, with LMUnit scores rising from 4.19 to 4.22 on RewardBench-2 and 4.01 to 4.06 on RM-Bench.EntRGi also achieves higher Top@1 reward across all tasks and its strongest results at τ=0.7, where APS noticeably degrades.Source: Table 1 · p. 8, §4.1 · p. 9
  • At low-entropy positions, Expectation provides consistent improvements, while STE is critical at high-entropy positions: removing it reduces EntRGi to the Expectation baseline, and EntRGi outperforms Expectation.EntRGi uses soft representations where appropriate and STE to provide hard-token inputs elsewhere, supporting reliable reward-model gradients.Source: §4.1 · p. 9
  • On LLaDA-8B-Instruct with tokenizer mismatch (45% with Llama, 55% with Qwen), gradient-based methods outperform BoN at test time, and EntRGi performs best.Non-overlapping tokens receive no gradients. Under mismatch, RGRL variants outperform their respective baselines, with gains up to +0.52 on WildChat-IF.Source: Table 2 · p. 9, §4.1 · pp. 9–10, §4.2 · p. 10

4.2 Reward Guided Post-Training

  • RGRL consistently beats diffu-GRPO at matched training steps or examples, with up to 70% (+0.90 absolute) improvement on WildChat-IF with Dream; RGRL-EntRGi has the highest sustained gains.Faster wall-clock convergence occurs in 1 out of 4 settings, while reward-model differentiation adds computational overhead.Source: Figure 2 · p. 10, §4.2 · p. 10
  • RGRL’s gains scale inversely with initial reward across Dream and LLaDA: they are largest on WildChat-IF, smaller on lmsys-chat-1m, and smallest on Magpie-Ultra.Initial rewards on Magpie-Ultra are already nearly positive.Source: §4.2 · p. 11

4.3 Analysis

  • EntRGi reduces early-step approximation error at moderate-to-high entropy (≈4–6) by weighting soft embeddings, whereas APS hard-token embeddings deviate substantially; both methods’ error converges to zero as denoising progresses.Initially, all tokens contribute to the error because the sequence is fully masked; as tokens become determined, fewer positions contribute.Source: Figure 3 · p. 11, §4.3 · p. 11
  • Larger reward models improve LMUnit scores across datasets, and EntRGi outperforms APS at each size: APS rises from 4.00 at 0.6B to 4.08 at 4B, while EntRGi rises from 4.04 to 4.12.Reward-model sizes range from 0.6B to 4B parameters.Source: §4.3 · p. 11, Figure 4 · p. 12
  • Increasing EntRGi optimization steps from M=1 to approximately 3–4 improves reward and LMUnit scores on JudgeBench and RM-Bench, but performance then degrades; LMUnit collapses beyond M=4.Reward scores on Reward-Bench-2 roughly improve up to M=5, and the optimal step count varies by dataset, with large M risking reward hacking or over-optimization.Source: Figure 5 · p. 12, §4.3 · p. 12

5 Conclusion · A Appendix

EntRGi uses predictive entropy to interpolate between continuous relaxations and hard token embeddings, balancing gradient accuracy with reward-model reliability. The conclusion reports that reward-guided rollouts improve dLLM post-training over scalar-reward RL, while the appendix points to implementation details and additional results.

5 Conclusion

  • EntRGi dynamically interpolates between continuous relaxations and hard token embeddings, trusting soft embeddings when confident and discrete tokens when uncertainty is high.The interpolation is based on the diffusion model’s predictive entropy and addresses the tension between gradient accuracy and reward-model reliability.Source: §5 · p. 12
  • RGRL is presented as the first study of dLLM post-training using reward gradients during rollouts, and it surpasses widely adopted scalar-reward RL.The conclusion also says EntRGi enables improved inference-time steering and RGRL extends these benefits to post-training.Source: §5 · pp. 12–13
  • EntRGi and RGRL require differentiable reward models and add back-propagation compute, which can offset wall-clock gains in some settings.These are limitations shared with gradient-based reward guidance methods.Source: §5 · p. 13
  • Future work includes selective denoising-step feedback, multi-objective rewards, principled tokenizer-mismatch treatments, and dense differentiable-reward feedback for RLVR-style methods.Selective feedback is proposed to better trade off sample efficiency against compute.Source: §5 · p. 13
  • Reward-guided systems risk reward hacking, proxy-reward misalignment with human preferences, and misuse of enhanced controllability to generate targeted harmful content.The paper recommends precautions and auxiliary quality checks when deploying these methods.Source: §5 · p. 13

A Appendix

  • Appendix B provides implementation details, including prompts, hyperparameters, and compute; Appendix C gives additional results and raw values underlying plots and figures.Source: A Appendix · p. 18

B Experimental Setup … C.3 Weighting Mechanism

The appendix specifies experimental protocols, evaluation inputs, and supplementary analyses of reward-model size, guidance iterations, and token-weighting mechanisms. It reports that entropy weighting outperforms the alternatives examined, while increasing guidance steps raises reward but can lead to reward hacking.

B.1 Implementation Details.

  • Test-time adaptation uses five-seed averages over 320 prompts per dataset, sequences up to 128 tokens, and one token per denoising step.The setup uses four H100 GPUs, η=0.5, M=3, N=4, and τ=0.7 unless stated otherwise; EOS is deprioritized for all methods.Source: §B.1 · p. 18
  • Training averages three runs and uses 500 steps, batch size 4, and four 128-token completions per prompt, decoding one token at a time.Training uses two GH200 GPUs, temperature 0.9, learning rate 5e −6, and LoRA with r = 32, α = 32, and dropout 0.1; diffu-GRPO uses KL-β = 0.04 and clipping ratio ϵ = 0.2, while RGRL uses M = 1 and η = 0.5. The codebase is adapted from Zhou et al. (2026)5, implementing diffu-GRPO (Zhao et al., 2025).Source: §B.1 · p. 18
  • LMUnit-Qwen2.5-72B scores responses from 1–5 using five unit tests, with the final score averaged across tests.Evaluation follows the official library protocol with greedy decoding and logprobs=20; tests cover relevance, correctness, coherence, and safety. The reward models include Skywork-Reward-V2-Qwen3-0.6B, 1.7B, and 4B, plus Skywork-Reward-V2-Llama-3.2-1B.Source: §B.1 · p. 18
  • Test-time adaptation uses Reward-Bench-2 (Malik et al., 2025), RM-Bench (Liu et al., 2024), and JudgeBench (Tan et al., 2025); post-training uses WildChat-IF (Bhaskar et al., 2025; Zhao et al., 2024), lmsys-chat-1m (Zheng et al., 2024), and Magpie-Ultra (Xu et al., 2024).The lmsys-chat-1m train split is filtered to English, and Magpie-Ultra's train set is filtered to quality “good” and above.Source: §B.1 · p. 19

B.2 Model Inputs

C.1 Scaling Reward Model Size

  • Results show similar trends across Skywork-Reward-v2 reward models of 0.6B, 1.7B, and 4B parameters.Table 3 reports results for the 0.6B and 4B models; the 1.7B results appear in Table 1 of the main paper.Source: §C.1 · p. 20, Table 3 · p. 20

C.2 Scaling Reward Model Iterations

  • Increasing guidance steps M increases reward but risks reward hacking beyond a point; the optimal M varies by dataset, and main experiments fix M = 3.The analysis varies M from 1 to 5 across Skywork-Reward-v2 models of 0.6B, 1.7B, and 4B parameters.Source: §C.2 · p. 21

C.3 Weighting Mechanism

  • EntRGi outperforms Inv-EntRGi, and L2-norm weighting improves on APS but does not match EntRGi.The authors suggest normalized token entropy offers a comparable signal across tokens and sequences, whereas L2 distances are unbounded and may require careful tuning.Source: §C.3 · p. 21

C.4 Timestep Ablation

EntRGi’s gradient-guidance benefits persist when denoising steps are reduced from 128 to 64, though the paper recommends using the highest number of available steps for best performance.

C.4 Timestep Ablation

  • EntRGi’s gradient-guidance benefits persist when denoising steps are reduced from 128 to 64.The results are reported in Table 8.Source: §C.4 · p. 21
  • For best performance, the paper recommends applying EntRGi at the highest number of denoising timesteps available.Source: §C.4 · p. 22

C.5 EntRGi Error Analysis

EntRGi’s approximation error grows approximately linearly with predictive entropy, unlike APS’s sharp increase. Its token-level entropy-aware balancing limits error at moderate entropy while preserving reward-model fidelity at high entropy, yielding more reliable gradients.

C.5 EntRGi Error Analysis

  • APS’s approximation error grows sharply with entropy, indicating a mismatch between discretized reward inputs and the continuous logits being updated.This steep error–entropy coupling produces unreliable gradient signals.Source: §C.5 · p. 22
  • EntRGi exhibits a controlled, approximately linear relationship between predictive entropy and approximation error across three datasets.Figure 9 shows this relationship in the bottom row.Source: §C.5 · p. 22
  • EntRGi adaptively reweights soft embeddings and hard tokens per token, limiting error at moderate entropy while preserving reward-model fidelity at high entropy.This balancing produces more stable, reliable reward gradients and translates into improved generation performance.Source: §C.5 · p. 22

C.6 Handling Tokenizer Mismatch

LLaDA’s custom tokenizer differs substantially from the tokenizers used by many reward models, so the paper handles non-overlapping tokens by setting their embeddings to zero. Despite this mismatch, EntRGi and RGRL-EntRGi achieve the best performance among the evaluated methods.

C.6 Handling Tokenizer Mismatch

  • The proposed approach handles tokenizer mismatch by setting the embeddings of all non-overlapping tokens to zero.This situation arises for LLaDA-8B-Instruct, which uses a custom tokenizer, while most existing reward models use autoregressive backbones adapted for classification and do not share its tokenizer.Source: §C.6 · pp. 22–23
  • Despite tokenizer mismatch, gradient-based methods outperform BoN, and EntRGi and RGRL-EntRGi achieve the best performance among all methods.BoN avoids the mismatch by decoding to text and re-encoding.Source: §C.6 · p. 23
  • 45%-55% of LLaDA’s tokenizer is mismatched with Qwen3/Llama-3, whereas Dream-v0-7B-Instruct has no such mismatch because it uses a Qwen2.5 backbone.The paper expects the discrepancy to diminish as more reward models use dLLMs such as LLaDA for initialization.Source: §C.6 · p. 23

C.7 Throughput Analysis

EntRGi preserves throughput relative to APS and BoN, while RGRL speed gains over diffu-GRPO vary by setting.

C.7 Throughput Analysis

  • EntRGi incurs no throughput overhead over APS and achieves throughput comparable to BoN (N=4).At similar throughput to BoN (N=4), EntRGi consistently outperforms it on Avg@N.Source: §C.7 · p. 23
  • EntRGi consistently outperforms BoN on Avg@N across both reward models and is competitive on Top@1; larger reward models slightly reduce throughput for APS and EntRGi.Source: §C.7 · p. 23
  • RGRL-EntRGi and RGRL-APS achieve an approximate 1.6× speedup over diffu-GRPO with Dream on WildChat-IF, but are slower in the other three settings.The authors attribute this tradeoff to reward-model differentiation exchanging computational efficiency for sample efficiency, and note that speedup may depend on the setting.Source: §C.7 · pp. 23–24, Figure 10 · p. 24

C.8 Qualitative Comparison

At τ = 0.1, the qualitative comparisons show EntRGi more explicitly reflecting requested themes and stylistic details than APS across several prompts.

C.8 Qualitative Comparison

  • For the robot-learning-to-love poem, EntRGi explicitly focuses on robotic themes, whereas APS’s poem is somewhat ambiguous.The comparison is shown in Figure 12.Source: §C.8 · pp. 24–25, Figure 12 · p. 26
  • In the sky explanation, EntRGi adds a personalized opening, “Well, honey,” while APS uses analogies such as ice cream.The prompt asks for an explanation suited to a five-year-old; the comparison is shown in Figure 13.Source: §C.8 · p. 25, Figure 13 · p. 26
  • For the cat-rulers story, EntRGi adds cat-related details such as toys, treats, and humans catering to them, while APS uses few cat analogies.The comparison is shown in Figure 16.Source: §C.8 · p. 25, Figure 16 · p. 28
  • For the clumsy-chimp detective story, EntRGi consistently incorporates the chimp’s clumsiness, unlike APS.The comparison is shown in Figure 17.Source: §C.8 · p. 25, Figure 17 · p. 28

Related papers

Loading 2602.05000v2…