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 diffusion language models make reward guidance difficult because their token outputs block direct gradients, while soft relaxations can be unreliable inputs for reward models. EntRGi addresses this by entropy-aware interpolation between soft and hard token embeddings, and RGRL uses the resulting reward-guided completions for post-training. Across test-time adaptation and post-training, the paper reports improved performance over prior methods, including a 70% relative improvement over diffu-GRPO.

  • Problem

    Discrete token outputs prevent direct reward-gradient propagation, while continuous relaxations can provide reward models with out-of-distribution inputs and reduce feedback reliability.

  • Method

    EntRGi interpolates per-token between continuous and sampled hard embeddings according to predictive entropy, and RGRL fine-tunes dLLMs on reward-guided completions.

  • Results

    70% relative improvement over diffu-GRPO is reported, with consistent gains across test-time adaptation and reward-guided post-training settings.

  • Takeaways & Limitations

    Model uncertainty can regulate reward feedback, supporting improved inference-time steering and reward-gradient-based post-training for discrete diffusion language models.

  • Takeaways & Limitations

    The discrete–continuous mismatch is most pronounced during early denoising when per-token predictive entropy is high, and reward models are trained on hard-token embeddings.

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 paper studies reward guidance for discrete diffusion language models, where discrete tokens block direct gradient propagation and continuous relaxations can make reward-model feedback unreliable. It introduces EntRGi, which uses predictive entropy to interpolate between soft and hard token representations, and applies reward-guided samples to post-training.

  • Motivation: Discrete diffusion language models denoise masked sequences in parallel, enabling reward-guided generation without a fixed left-to-right order.A naive alternative is particle-filtering-style sampling that generates multiple noisy completions and rejects low-reward ones.
  • Motivation: Continuous diffusion uses differentiable reward feedback to iteratively refine generation, motivating analogous reward guidance for discrete diffusion language models.In dLLMs, outputs are discrete tokens and reward models are fine-tuned language models.
  • Challenge: Replacing discrete tokens with continuous embeddings enables gradients but can hurt reward-model performance because such soft inputs are out of distribution.APS instead evaluates sampled hard tokens while propagating gradients, forming the prior approach addressed by EntRGi.
  • Contributions: EntRGi interpolates between continuous and sampled hard token embeddings using per-token predictive entropy, favoring soft inputs when confident and hard inputs when uncertain.The mechanism aims to provide reliable optimization gradients while keeping reward-model inputs interpretable during denoising.
  • Contributions: RGRL post-trains dLLMs by generating reward-guided completions and fine-tuning the model to increase their likelihood.The recipe uses either EntRGi or APS guidance during generation.
  • Results: 70% relative improvement over diffu-GRPO is reported for the proposed reward-guided reinforcement-learning recipe.The recipe works with both APS and EntRGi guidance, with EntRGi providing higher gains.

2 Related Work

Prior work for discrete diffusion reward guidance includes gradient-free search and particle methods, continuous-relaxation approaches, and straight-through estimators. The paper identifies a shared mismatch between discrete outputs and gradient-compatible continuous representations, especially under high uncertainty.

  • Discrete diffusion posterior sampling: Discrete diffusion models generate predictive distributions over all tokens in parallel, supporting posterior sampling under external reward constraints without retraining.This non-autoregressive structure distinguishes them from autoregressive generation.
  • Reward-gradient-free methods: Reward-gradient-free methods use scalar reward queries with search, Best-of-N, or particle-based sampling rather than back-propagating through reward models.These approaches have been developed for continuous diffusion and extended to discrete diffusion.
  • Reward-gradient-based methods: Reward-gradient-based methods either feed continuous token relaxations to reward models or use straight-through estimation at sampled hard tokens.APS is the prior state of the art for inference-time dLLM steering, while DRAKES uses straight-through fine-tuning for biological sequence diffusion.
  • Challenges and limitations: Both relaxation and straight-through methods face a mismatch between discrete outputs and continuous gradient representations, most strongly during early denoising when token entropy is high.The paper uses uncertainty for adaptive gradient regulation and applies reward-guided completions as dense training feedback beyond scalar-reward RL.

3 Reward Guidance for Discrete Diffusion LLMs

Reward guidance for discrete diffusion LLMs must preserve differentiability without feeding reward models unfamiliar soft inputs. The paper develops EntRGi, which combines soft and hard token representations using entropy-aware interpolation while guiding denoising and post-training.

  • Discrete diffusion generation: Discrete diffusion LLMs iteratively unmask tokens in parallel from a fully masked sequence, with selected positions sampled according to predictive distributions.The described models use an unmask-and-commit generation mode, so unmasked tokens remain fixed in later steps.
  • The guidance tradeoff: Direct reward gradients are unavailable for discrete outputs, while soft embeddings enable gradients but can fall outside the reward model’s training distribution.Reward models trained on hard-token embeddings become less reliable as soft inputs drift from real token embeddings.
  • The guidance tradeoff: APS evaluates sampled hard-token embeddings while propagating gradients through soft embeddings, creating approximation error that grows as token distributions become more diffuse.The approximation error is zero for point-mass distributions and increases with predictive entropy.
  • EntRGi: EntRGi constructs reward-model inputs as a convex combination of soft expected embeddings and sampled hard-token embeddings, with interpolation weights chosen from per-token entropy.Soft inputs are favored for confident predictions, whereas hard inputs receive greater weight under uncertainty; the method unifies Expectation and APS as endpoint cases.
  • EntRGi: EntRGi theoretically reduces approximation error below APS while bounding vocabulary error, preserving gradient reliability through entropy-aware balancing.The supplied analysis states that high entropy drives the weight toward hard tokens, reducing vocabulary mismatch, while lower entropy permits more soft representation.
  • RGRL: RGRL self-distills reward-guided completions by generating samples with EntRGi or APS and updating the diffusion model to increase their likelihood.The resulting algorithms are named RGRL-EntRGi and RGRL-APS, and the approach uses reward-guided data for post-training.

4 Experiments

Experiments evaluate EntRGi and related guidance methods across test-time adaptation and reward-guided post-training, including mismatched-tokenizer settings. EntRGi generally outperforms APS and other baselines, while reward-model size and guidance-step choices affect performance.

  • 4 Experiments: Gradient-based methods consistently outperform Best-of-N across all evaluated benchmarks.The comparison covers Reward-Bench-2, JudgeBench, and RM-Bench.
  • 4 Experiments: EntRGi improves reward-model-judged output quality by approximately 33% relative to APS.It also improves LMUnit from 4.19 to 4.22 on RewardBench-2 and from 4.01 to 4.06 on RM-Bench.
  • 4.2 Reward Guided Post-Training: RGRL gains consistently over diffu-GRPO and reaches up to 70% relative improvement, or +0.90 absolute improvement, on WildChat-IF with Dream.RGRL-EntRGi yields the highest sustained gains across settings, while differentiating through the reward model adds computational overhead.
  • 4.2 Reward Guided Post-Training: RGRL variants continue to outperform their respective baselines under tokenizer mismatch, with gains reaching up to +0.52 on WildChat-IF.The absolute gains are smaller in this setting.
  • 4.3 Analysis: Approximation error decays as denoising progresses, because fewer token positions remain uncertain; both methods converge to zero error.At the initial fully masked step, all tokens contribute to the error.
  • 4.3 Analysis: Increasing reward-model size improves LMUnit scores, with APS rising from 4.00 to 4.08 and EntRGi from 4.04 to 4.12 between 0.6B and 4B models.EntRGi outperforms APS at every reward-model size across the three datasets.
  • 4.3 Analysis: Increasing guidance steps improves reward but risks over-optimization; M = 3–4 is a reliable range, while the optimum varies by dataset.LMUnit collapses beyond M = 4, and the main experiments use M = 3 for all datasets.

5 Conclusion

EntRGi uses predictive entropy to balance continuous relaxations and hard token embeddings for reward guidance, while RGRL applies reward gradients during dLLM rollouts. The authors report benefits for inference-time steering and post-training, alongside compute and deployment caveats.

  • EntRGi interpolates between continuous relaxations and hard token embeddings according to predictive entropy.
  • RGRL leverages reward gradients during dLLM rollouts and surpasses widely adopted scalar-reward reinforcement learning.
  • Together, EntRGi and RGRL extend reward-gradient guidance to inference-time steering and post-training of discrete diffusion language models.
  • EntRGi and RGRL require differentiable reward models and additional back-propagation compute, which can offset wall-clock gains in some settings.
  • Reward-guided systems retain risks including reward hacking, proxy misalignment with human preferences, and misuse for targeted harmful content.

B.1 Implementation Details.

The experiments evaluate test-time adaptation and reward-guided training with specified hardware, sampling, model, dataset, and evaluation configurations. Implementation details include LMUnit scoring, multiple reward-model sizes, and ablations over gradient steps and weighting schemes.

  • Test-Time Adaptation: Test-time adaptation uses 4 H100 GPUs, five seeds, 320 prompts per dataset, sequences up to 128 tokens, and one decoded token per denoising step.The default settings are η=0.5, M=3, N=4, and τ=0.7.
  • Training Experiments: Training experiments use 2 GH200 GPUs, three runs, 500 steps, batch size 4, and four 128-token completions per prompt.Sampling temperature is 0.9, and the implementation adapts diffu-GRPO code.
  • Training Experiments: All methods use learning rate 5e-6 and LoRA with r=32, α=32, and dropout 0.1.
  • LMUnit evaluation: LMUnit scores responses on a 1–5 continuous scale using greedy decoding with logprobs=20 and five unit tests.The tests cover relevance, correctness, coherence, and safety.
  • Results and Ablations: EntRGi records the strongest displayed values across the three benchmark groups, including 11.40±0.27, 8.60±0.12, and 13.67±0.15.These values are reported in the EntRGi row of the supplied results passage.
  • Reward Models and Datasets: The study evaluates Skywork reward models at 0.6B, 1.7B, and 4B scales across Reward-Bench-2, RM-Bench, and JudgeBench.Table 3 reports Dream-v0-7B-Instruct performance across these reward-model sizes and benchmarks.
  • Results and Ablations: Increasing gradient steps M raises reward but becomes prone to reward hacking after a dataset-dependent point; main experiments fix M=3.
  • Results and Ablations: Alternative weighting experiments compare inverse-entropy and L2-based token weighting against entropy-aware weighting.The supplied passage states that Inv-EntRGi consistently underperforms.

C.4 Timestep Ablation

EntRGi’s gradient-guidance benefits persist when denoising timesteps are reduced from 128 to 64. The authors nevertheless recommend using the highest available timestep count for best performance.

  • EntRGi’s gradient-guidance benefits persist when denoising timesteps decrease from 128 to 64.
  • The authors recommend applying EntRGi at the highest available number of denoising timesteps for best performance.

C.5 EntRGi Error Analysis

APS’s approximation error grows sharply with predictive entropy, whereas EntRGi maintains a controlled, approximately linear relationship. EntRGi thereby limits moderate-entropy error while preserving reward-model fidelity at high entropy.

  • APS exhibits sharply increasing approximation error with entropy, indicating mismatch between discretized reward inputs and updated continuous logits.
  • EntRGi produces a controlled, approximately linear error–entropy relationship through token-level reweighting of soft embeddings and hard tokens.
  • EntRGi limits approximation error in moderate-entropy regions while preserving reward-model fidelity at high entropy.The passage links this balancing to more stable and reliable reward gradients and improved generation performance.

C.6 Handling Tokenizer Mismatch

LLaDA’s tokenizer mismatches the reward-model tokenizer substantially, so EntRGi handles non-overlapping tokens by setting their embeddings to zero. Despite this mismatch, EntRGi and RGRL-EntRGi achieve the best reported performance.

  • EntRGi handles non-overlapping tokenizer entries by setting their embeddings to zero.This mismatch arises because LLaDA uses a custom tokenizer while many reward models use autoregressive backbones.
  • Inverse-entropy weighting shows no noticeable improvement, while L2-norm weighting improves over APS but remains below regular EntRGi.
  • Despite tokenizer mismatch, gradient-based methods outperform BoN, and EntRGi and RGRL-EntRGi achieve the best performance.BoN avoids the mismatch by decoding to text and re-encoding.

C.7 Throughput Analysis

EntRGi has throughput comparable to BoN and no overhead over APS, while gradient-guided RGRL can trade computational efficiency for sample efficiency. The observed speed advantage is setting dependent.

  • EntRGi incurs no overhead over APS and achieves throughput comparable to BoN with N=4.
  • At similar throughput to BoN with N=4, EntRGi consistently outperforms BoN on Avg@N across both reward models.EntRGi is also competitive on Top@1.
  • Approximately 1.6× speedup is observed for Dream on WildChat-IF when comparing RGRL-EntRGi and RGRL-APS with diffu-GRPO.
  • RGRL-EntRGi and RGRL-APS are slower in the other three settings because reward-model differentiation trades computational efficiency for sample efficiency.

C.8 Qualitative Comparison

The qualitative comparisons examine APS and EntRGi generations across several prompts using low-temperature decoding. EntRGi more consistently incorporates requested themes and stylistic traits in the reported examples.

  • The comparisons use low temperature τ = 0.1 to minimize randomness in final outputs.
  • For a robot-learning-to-love poem, EntRGi produces a more tailored output with explicit robotic themes than APS.
  • For a child-oriented sky explanation, EntRGi adds finer-grained stylistic personalization, including the opening “Well, honey,” compared with APS.
  • For a story about cats ruling the world, EntRGi includes richer cat-related details than APS, including toys, treats, and humans catering to cats.
  • For a clumsy-chimp detective story, EntRGi consistently incorporates the chimp’s clumsiness, whereas APS provides little indication of that trait.
Loading 2602.05000v2…