PaperlayerView exact revision on arXiv

Source-linked AI summary

GDSD: Reinforcement Learning as Guided Denoiser Self-Distillation for Diffusion Language Models

Xiaohang Tang, Keyue Jiang, Che Liu, Qifang Zhao, Xiaoxiao Xu, Sangwoong Yoon, Ilija Bogunovic

arXiv:2605.29398v1cs.LGcs.AI

TL;DR

ELBO-based reinforcement learning for diffusion language models is biased by training–inference mismatch. GDSD instead distills an advantage-guided denoising teacher, consistently improving prior ELBO-based methods across benchmarks, with gains up to +19.6% test accuracy.

  • Problem

    ELBO-based reinforcement learning for dLLMs suffers training–inference mismatch because surrogate likelihoods and iterative decoding induce policy bias.

  • Method

    GDSD recasts reverse-KL-regularized reinforcement learning as off-policy self-distillation from an advantage-guided denoising teacher using normalization-free squared-logit matching.

  • Results

    +19.6% test-accuracy gains are reported on planning tasks, while GDSD consistently improves prior ELBO-based methods across planning, math, and coding benchmarks.

  • Takeaways & Limitations

    GDSD provides a stable alternative for reinforcement learning with dLLMs that avoids reliance on ELBO likelihood surrogates.

  • Takeaways & Limitations

    The approach assumes a consistent forward masking process across the reference, old, and current diffusion models.

Abstract

from arXiv · show

Reinforcement learning (RL) can be used to improve the policy (denoiser) of diffusion large language models (dLLMs), while being hindered by the intractability of the policy likelihood. A dominant and efficient family of methods replaces the likelihood in standard RL with its evidence lower bound (ELBO), estimated from randomly masked sequences. Despite being well aligned with pre-training, these approaches introduce bias through training--inference mismatch by using the ELBO as a likelihood surrogate, which can degrade performance. In this work, we propose Guided Denoiser Self-Distillation (GDSD) to directly distill the denoiser of dLLMs from an advantage-guided self-teacher, derived from the closed-form optimum of reverse-KL regularized RL. GDSD matches the dLLM's denoiser logits to the teacher's via a normalization-free objective, which reduces RL to likelihood-free self-distillation and thus bypasses the TIM biases. Recent ELBO-based methods emerge as instances of applying different distillation divergences, but with diagnosable pathologies that GDSD avoids. On planning, math, and coding benchmarks with LLaDA-8B and Dream-7B, GDSD consistently outperforms prior state-of-the-art ELBO-based methods with a more stable training reward dynamics, achieving test-accuracy improvements of up to $+19.6\%$. These results suggest that direct denoiser self-distillation, without relying on an ELBO likelihood surrogate, can provide a more stable and effective RL procedure for dLLMs. Code is available at https://github.com/GaryBall/GDSD.

1 Introduction

dLLMs offer efficient, non-left-to-right generation, but reinforcement learning is hindered by intractable policy likelihoods and ELBO-induced training-inference mismatch. GDSD addresses this by recasting RL as normalization-free denoiser self-distillation, outperforming prior ELBO-based methods with more stable rewards and gains up to +19.6%.

  • Motivation: dLLMs generate multiple tokens per decoding step without strict left-to-right ordering, improving inference efficiency and enabling broader token dependencies.This positions dLLMs as efficient alternatives to autoregressive models.
  • RL challenges: Intractable policy likelihoods motivate trajectory-likelihood and ELBO-surrogate RL, but trajectory estimates are costly while ELBO methods introduce training-inference mismatch bias.Trajectory likelihood is unbiased in principle but computationally expensive and misaligned with pre-training objectives.
  • Method: GDSD recasts dLLM reinforcement learning as direct off-policy denoiser self-distillation using an advantage-guided denoiser derived from reverse-KL regularized RL.The guided denoiser acts as a self-teacher at each RL step.
  • Method: GDSD uses squared-logit distillation with logit centralization to eliminate the partition function, enabling off-policy updates that bypass training-inference mismatch bias.The normalization-free reformulation is the mechanism used to avoid the partition function.
  • Results: +19.6% test-accuracy gains are achieved by GDSD with Dream-7B on planning tasks, alongside more stable training reward dynamics than prior ELBO-based methods.With LLaDA-8B, gains range from +0.6% to +5% across planning, math, and coding benchmarks.

2 Preliminaries

Masked diffusion models generate sequences through a continuous-time masking process and learn a reverse denoiser, while reinforcement learning optimizes rewards using policy objectives. Because sequence likelihoods are intractable, common ELBO-based methods introduce training–inference mismatch that can bias optimization and cause collapse.

  • Masked Diffusion Models: Masked diffusion models operate over a discrete vocabulary augmented with a dedicated mask token.The augmented state space is V′ = V ∪ {[M]}.
  • Masked Diffusion Models: A forward transition progressively masks clean sequences, while the learned reverse model predicts clean sequences from masked states and is called the denoiser.The forward process is defined by q(x_t|x_s) = Cat(x_t; Q(s,t)^⊤x_s), and pθ models the reverse process.
  • ELBO-Based Reinforcement Learning: ELBO-based RL methods estimate the sequence likelihood from randomly masked samples, leveraging computational efficiency and alignment with masked-sequence pre-training.RL can remain effective with a small Monte Carlo sample size, such as 4.
  • Training-Inference Mismatch: The ELBO surrogate has a non-negligible gap and estimation error, producing bias and variance in likelihood ratios that can degrade performance or cause training collapse.These issues are described as Training-Inference Mismatch (TIM).
  • Training-Inference Mismatch: The dLLM sampling process further differs from the likelihood represented by the ELBO surrogate, making likelihood-ratio correction ineffective and motivating likelihood-free off-policy self-distillation.The mismatch arises from discretizing the reverse process and selecting tokens to remask, including through policies such as low-confidence selection.

3 Method

GDSD converts reverse-KL-regularized RL for masked diffusion models into direct denoiser self-distillation from an advantage-guided teacher. By avoiding likelihood ratios and ELBO computation, it bypasses training–inference-mismatch biases while retaining efficient masked-sequence training.

  • Reverse-KL formulation: Reverse-KL regularization yields a closed-form optimal policy π* that GDSD approximates with a masked diffusion model under a consistent forward masking process.The same masking process is applied to the reference, old, and current models.
  • Teacher denoiser: The resulting teacher is an energy-guided denoiser built from the old and reference models, with reward-derived advantage providing the guidance.Its energy function is E(x0) = −ψA(x0), and the teacher uses the old model as its base with reference-model regularization.
  • GDSD objective: GDSD trains the parameterized denoiser by matching its logits to the guided teacher, reducing RL policy optimization to likelihood-free self-distillation.The sequence-level denoising distribution factorizes into token-level denoisers, while the advantage is estimated from rewards.
  • Advantages over ELBO-based RL: GDSD avoids both the ELBO–likelihood gap and the re-masking sampler–policy gap because its objective forms no likelihood ratio.It can train on off-policy samples produced through iterative self-refinement and requires only minor changes to existing dLLM RL pipelines.
  • Normalization-free reformulation: Normalization-free logit matching bypasses the teacher partition function by exploiting Softmax translation invariance, with token-level logit centralization preventing uncontrolled logit-scale shifts.The partition function would otherwise require additional completion samples and model inference, increasing online RL overhead.

4 Connection and Comparison to Related Work

GDSD connects logit-matching self-distillation to advantage-weighted ELBO and regularized policy-gradient objectives through alternative distillation divergences. Unlike exponential-weighting schemes and ELBO-based policy-gradient methods, it avoids their identified data-efficiency, instability, and TIM-bias limitations.

  • GDSD formulation: GDSD matches student and teacher denoiser logits with squared l2 loss using samples from an old/reference-policy mixture.Forward-KL and reverse-KL divergences are also eligible for the self-distillation objective.
  • Advantage-Weighted ELBO: Forward-KL distillation recovers advantage-weighted ELBO objectives used by wd1 (AW-DCE) and DMPO.Importance sampling makes the objective equivalent, up to θ-independent terms, to an ELBO-style denoising loss reweighted by exponential advantage.
  • Advantage-Weighted ELBO: Exponential weighting is data-inefficient because negative-advantage samples receive negligible weights, while auxiliary negative-sample penalties can cause training instability.The passage contrasts these limitations with GDSD’s stated natural incorporation of the sig...
  • Regularized Policy Gradient: Reverse-KL distillation recovers a regularized policy-gradient objective whose leading term maximizes expected advantage and whose remaining term regularizes the current denoiser.At(xt) serves as a state-dependent baseline, while the regularization is against the old/reference-policy mixture.
  • Regularized Policy Gradient: PG-based dLLM methods including SPG, UniGRPO, and ESPO approximate likelihood with ELBO and therefore incur inevitable TIM bias when importance sampling is required.Logit matching has also appeared in autoregressive post-training, dLLM preference learning through EMBR, and continuous diffusion RL.

5 Experiment

GDSD is evaluated across diffusion language models and benchmarks spanning mathematical reasoning, planning, and coding. It improves test accuracy over ELBO-based methods, stabilizes training rewards, and reveals distinct effects from token logit centralization and guidance strength.

  • Experimental Setup: GDSD is evaluated with LLaDA-8B-Instruct and Dream-v0-Instruct-7B across six benchmarks covering mathematical reasoning, planning, and coding.The benchmarks are GSM8K, MATH500, Countdown, Sudoku, HumanEval, and MBPP.
  • Dream-7B Results: +19.6% is the maximum reported improvement achieved by GDSD over ELBO-based methods on Dream-7B-Instruct.GDSD improves average test accuracy by +9.5% overall and +10% at the best-performing generation length.
  • LLaDA-8B Results: On LLaDA-8B-Instruct, GDSD shows stable reward convergence on planning and coding tasks and outperforms prior methods on almost all benchmarks.The evaluation reports training reward dynamics alongside testing accuracy across planning, math, and coding benchmarks.
  • Token Logit Centralization (TLC): Direct matching without TLC improves over prior ELBO-based methods, while TLC can degrade test performance despite being more faithful to the theory.The lack of consistent test-accuracy gains with TLC suggests a potential generalization gap caused by overfitting to training-specific signals.
  • Guidance Coefficient ψ: Increasing the guidance coefficient ψ generally raises training rewards by producing a teacher denoiser more biased toward high-advantage samples.This supports GDSD’s formulation as direct distillation of an energy-guided target distribution.

6 Conclusion

The conclusion identifies ELBO-based likelihood surrogates as a source of bias that can degrade dLLM RL performance and potentially cause training collapse. It presents GDSD as an off-policy self-distillation framework that performs reverse-KL-regularized RL without likelihood ratios, thereby bypassing this bias.

  • ELBO-based likelihood surrogates can bias RL for dLLMs, degrading performance and potentially causing training collapse.
  • GDSD is an off-policy self-distillation framework that equivalently performs reverse-KL-regularized RL without likelihood ratios.
  • Token-level Logit Centralization produces more stable training behavior, while increasing the energy-guidance coefficient ψ improves training reward.

A Proof … D.1 Histogram of Accuracy

The appendix proves equivalences underlying GDSD, relates its objective to prior diffusion-RL formulations, and details practical training choices. Experimental summaries report consistent accuracy gains over ELBO-based baselines and improved reward stability with token-logit centralization.

  • A.1 Lemma: Token-level Logit Centralization: Token-level logit centralization is equivalent to sequence-level centralization and preserves GDSD’s optimal probability distribution.The proof establishes the equivalence by counting vocabulary assignments across all possible sequences and applying denoising-distribution factorization.
  • A.2 Proof to Proposition 3.3: Centralizing the energy-guided teacher logits cancels sequence-independent normalizers and regroups policy terms with the centralized advantage.The resulting expression follows from substituting the log-form target distribution and using the token-level centralization lemma.
  • A.3 Proof of Proposition 4.1: The forward-KL gradient reduces to an advantage-weighted denoising objective identified with AW-DCE and DMPO when ELBO approximates likelihood.Terms involving the reference policy, advantage function, and partition function are constant with respect to model parameters; importance sampling introduces the advantage weight.
  • A.4 Proof of Reverse-KL Distillation to Policy Gradient: Reverse-KL distillation to the energy-guided teacher uses sampling equivalence between the current policy and the denoising forward process.The proof invokes x0 sampled from πθ and xt sampled from q(· | x0).
  • B Additional Related Work: Related work spans continuous diffusion-policy RL, autoregressive logit matching, exact trajectory-likelihood methods, and likelihood-free continuous diffusion objectives.The appendix contrasts these approaches with discrete-token denoising distributions and notes computational or formulation-specific limitations of trajectory-likelihood and flow-based methods.
  • C.1 Algorithm: GDSD samples multiple old-policy completions, computes within-prompt mean-centered advantages, and converts them into denoising-level supervision.The algorithm applies iterative re-masking decoding before reward scoring and advantage computation.
  • C.2 Additional Practical Design: The practical algorithm centralizes logits, constructs an advantage-guided teacher, and optionally separates reference regularization from the MSE objective.Efficiency and variance are addressed through single-inference batching, coupled complementary-mask sampling, and 1/t logit reweighting; the external regularizer approximates a k2 KL divergence.
  • D.1 Histogram of Accuracy: GDSD consistently outperforms ELBO-based baselines in test accuracy, with especially significant gains on Dream-7B base models.Figure 4 reports best accuracy across generation lengths 128, 256, and 512 for LLaDA-8B, while Figure 5 shows that token-logit centralization improves reward stability and reward level.

D.2 Coding Benchmarks

This section reports additional details and results for coding benchmarks. On these benchmarks, GDSD generally improves average performance over ELBO-based methods, with results assembled from reevaluated and extracted baselines.

  • The section presents additional details and results on coding benchmarks.
  • GDSD generally improves average performance over ELBO-based methods on coding benchmarks.
  • † results are reevaluated using pass@1 with lm-eval on specified checkpoints or models trained from scratch, while other baselines are extracted from Wang et al. [52].

D.3 Experimental Setups

The experiments fine-tune LLaDA-8B-Instruct and Dreamv0-Instruct-7B across six mathematical, logical, and coding benchmarks. Evaluation is unified across tasks, while rewards combine correctness with task-specific formatting or execution criteria.

  • Experimental Setup: Experiments use LLaDA-8B-Instruct and Dreamv0-Instruct-7B on GSM8K, MATH500, Countdown, Sudoku, HumanEval, and MBPP.The benchmarks cover mathematical reasoning, logical reasoning, and coding.
  • Evaluation Setup: Evaluation uses zero-shot generation lengths of 128, 256, and 512 for Sudoku and Countdown, with denoising steps set to half the sequence length.For other tasks, denoising steps equal the sequence length.
  • Evaluation Setup: The unified protocol addresses substantial variation in prior evaluations, including denoising steps of either 1× or 0.5× generation length.ESPO uses 1× generation length, whereas d1, wd1, and SPG use 0.5× generation length in the cited settings.
  • Reward Setup: Sudoku extracts and length-normalizes answers to 16 digits, then scores the fraction of originally empty cells matching the ground truth.Missing or invalid answers receive reward 0.0.
  • Reward Setup: Coding rewards valid Markdown code blocks and execute syntactically valid completions in a sandbox, with correctness based on the fraction of passed test cases.Completions failing format or syntax checks are not executed and receive execution reward 0.0.

D.4 Hyperparameter Settings and Implementation Details

The implementation follows ESPO for most hyperparameters, using LoRA adaptation and a specified diffusion rollout configuration. Training budgets vary by task, with evaluation every 100 steps and selection of the checkpoint with the highest average accuracy.

  • D.4 Hyperparameter Settings and Implementation Details: LoRA uses rank r = 128 and scaling factor α = 64 for training.The hyperparameter choices largely follow ESPO [32].
  • D.4 Hyperparameter Settings and Implementation Details: 256 tokens, 128 diffusion steps, block size 32, temperature 0.9, g = 6 completions, and m = 2 Monte Carlo samples define RL rollout.Confidence-based semi-autoregressive generation is used, and the sample count is limited by computational constraints.
  • D.4 Hyperparameter Settings and Implementation Details: 3K, 2K, 5K, and 10K gradient-update steps are used for GSM8K/MATH500, coding, Sudoku, and countdown, respectively.Models are evaluated every 100 steps, and the checkpoint with the highest average accuracy is reported.

D.5 Additional Results Analysis · E Others

Additional analyses show that Token Logits Centralization (TLC) improves reward stability and final performance across several tasks, while GDSD trains more effectively than prior ELBO-based methods. In Sudoku, TLC’s delayed reward jump suggests it can help escape local minima, supporting the proposed unbiased optimization account.

  • D.5 Additional Results Analysis: TLC consistently improves final attainable reward and optimization stability across GSM8K, Countdown, and Sudoku.With other hyperparameters fixed, TLC produces higher and more stable reward plateaus across these datasets.
  • D.5 Additional Results Analysis: In GSM8K, the model without TLC peaks early but later degrades, whereas TLC remains around a high reward level.This pattern suggests TLC stabilizes optimization after rewarding behaviors are discovered.
  • D.5 Additional Results Analysis: In Countdown, the reward dynamics show a similar stabilization pattern when TLC is used.The supplied analysis identifies Countdown as exhibiting the same qualitative contrast between TLC and non-TLC optimization.
  • D.5 Additional Results Analysis: In Sudoku, TLC initially stays near a moderate reward of roughly 0.8 before sharply improving after approximately 4k training steps and surpassing the non-TLC baseline.Checkpoint analysis confirms that this transition corresponds to a change in the model’s optimization behavior.
  • D.5 Additional Results Analysis: The empirical TLC results support the theoretical claim that TLC enables unbiased RL optimization and more stable, effective optimization than ELBO-based and likelihood-based methods.This is the authors’ stated interpretation of the ablation findings.
  • E Others: GDSD exhibits more stable and effective training than previous ELBO-based methods, with TLC converging slightly better than direct matching.SPG on Countdown and ESPO on coding show poor training performance in comparison.

E.1 Compute Resources · E.2 Broader Impact

Experiments use 8 GPUs, including AMD Instinct MI308X and A800. The method may make diffusion-language-model post-training more accessible, but stronger and cheaper training can amplify misuse and reward-specification risks, requiring broader evaluation and deployment safeguards.

  • E.1 Compute Resources: Experiments are conducted on 8 GPUs, including AMD Instinct MI308X and A800.
  • E.2 Broader Impact: The work studies efficient reinforcement learning for diffusion large language models.
  • E.2 Broader Impact: Avoiding likelihood estimation and reverse-chain policy gradients can reduce post-training computational cost and instability.
  • E.2 Broader Impact: These efficiency gains may make alignment and reasoning-oriented training more accessible.
  • E.2 Broader Impact: Stronger and cheaper post-training may amplify risks including misinformation, spam, academic dishonesty, automated persuasion, and harmful content generation.
  • E.2 Broader Impact: Misspecified reward signals may reinforce biases, unsafe behaviors, or undesirable shortcuts.
  • E.2 Broader Impact: Models should be evaluated for safety, robustness, bias, hallucination, and misuse beyond reasoning and coding benchmarks.
  • E.2 Broader Impact: Deployment should use safeguards such as data curation, red-teaming, monitoring, and access controls when appropriate.
Loading 2605.29398v1…