PaperlayerView exact revision on arXiv

Source-linked AI summary

Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective

Jingyang Ou, Jiaqi Han, Minkai Xu, Shaoxuan Xu, Jianwen Xie, Stefano Ermon, Yi Wu, Chongxuan Li

arXiv:2512.03759v1cs.CLcs.AIcs.LG

TL;DR

Standard token-level RL relies on autoregressive likelihood factorization, which is ill-defined or costly for non-autoregressive dLLMs. ESPO treats sequence generation as one action, uses the ELBO as a tractable likelihood proxy, and stabilizes optimization with per-token normalization and robust KL estimation. It consistently outperforms prior dLLM-RL baselines across math, coding, and planning, including 20–40 absolute-point gains on Countdown.

  • Problem

    Autoregressive RL objectives require token-level likelihood factorization, but dLLMs generate non-autoregressively, making those conditionals ill-defined or computationally expensive.

  • Method

    ESPO treats an entire generated sequence as one action, uses the ELBO as a sequence-likelihood proxy, and adds per-token ratio normalization with robust KL regularization.

  • Results

    ESPO consistently outperforms existing dLLM-RL baselines across math, coding, and planning, with 20–40 absolute-point gains over the strongest baselines on Countdown.

  • Takeaways & Limitations

    Sequence-level optimization provides a principled and empirically effective paradigm for reinforcement learning in diffusion language models.

Abstract

from arXiv · show

Reinforcement Learning (RL) has proven highly effective for autoregressive language models, but adapting these methods to diffusion large language models (dLLMs) presents fundamental challenges. The core difficulty lies in likelihood approximation: while autoregressive models naturally provide token-level conditional probabilities essential for token-level RL objectives (e.g., GRPO), dLLMs generate sequences through iterative non-autoregressive denoising steps that lack this factorization. To address this fundamental mismatch, we propose ELBO-based Sequence-level Policy Optimization (ESPO), a principled RL framework that treats entire sequence generation as a single action and uses the ELBO as a tractable sequence-level likelihood proxy. Our method incorporates per-token normalization of importance ratios and robust KL-divergence estimation to ensure stable large-scale training. Extensive experiments on mathematical reasoning, coding, and planning tasks demonstrate that ESPO significantly outperforms token-level baselines, achieving dramatic improvements of 20-40 points on the Countdown task, while maintaining consistent gains on math and coding benchmarks. Our approach establishes sequence-level optimization as a principled and empirically effective paradigm for RL in dLLMs. Our code is available at https://github.com/ML-GSAI/ESPO.

1 INTRODUCTION

RL has improved autoregressive language models, motivating its extension to dLLMs, but non-autoregressive generation creates a likelihood mismatch for standard token-level objectives. ESPO addresses this mismatch with sequence-level optimization and achieves consistent gains across math, coding, and planning benchmarks.

  • 1 INTRODUCTION: Standard RL methods such as GRPO assume autoregressive token-level likelihoods that are ill-defined or expensive for non-autoregressive dLLMs.Prior solutions use heuristic mean-field or token-level ELBO approximations, or computationally heavy trajectory-level formulations.
  • 1 INTRODUCTION: ESPO treats an entire dLLM-generated sequence as one action and uses the ELBO as a tractable sequence-likelihood proxy.It also adds per-token ELBO-ratio normalization and robust KL regularization for stable, computationally efficient training.
  • 1 INTRODUCTION: Across math, coding, and planning tasks, ESPO consistently outperforms prior dLLM-RL baselines on both LLaDA and Dream.The strongest gains occur on planning tasks, while coding performance is comparable to LLaDA-1.5.

2 BACKGROUND

dLLMs generate sequences through masked diffusion, while RL methods such as GRPO use group-relative rewards, token-level importance ratios, and KL regularization. Because exact dLLM likelihoods are difficult to compute, ELBO-based approximations provide a practical alternative.

  • 2 BACKGROUND: Masked diffusion models gradually corrupt clean completions by replacing tokens with a mask token M.The forward process is defined over the prompt x, clean completion y, and corrupted sequence at time t.
  • 2 BACKGROUND: dLLM sequence likelihoods are typically approximated with an ELBO rather than computed exactly.The ELBO has an equivalent lower-variance formulation based on uniformly sampling a discrete number of masked tokens.
  • 2 BACKGROUND: The ELBO is a lower bound that has been shown empirically effective for likelihood evaluation and training.The supplied passage cites perplexity and DPO variants as examples.
  • 2 BACKGROUND: GRPO samples completion groups, computes each sample’s advantage relative to the group mean reward, and includes a KL penalty.Its standard objective uses token-level importance ratios between current and old autoregressive policies.

3 THE CHALLENGE OF THE TOKEN-LEVEL PERSPECTIVE IN DLLMS

Standard policy-gradient objectives rely on autoregressive token factorization, but dLLMs refine complete sequences through iterative denoising. Token-level substitutes therefore lack a formally valid conditional-likelihood interpretation, motivating a sequence-level formulation.

  • 3 THE CHALLENGE OF THE TOKEN-LEVEL PERSPECTIVE IN DLLMS: Autoregressive RL factorizes sequence likelihood into token conditionals, whereas dLLMs generate complete sequences non-autoregressively through iterative denoising.The autoregressive conditional probability is consequently ill-defined or difficult to compute for dLLMs.
  • 3 THE CHALLENGE OF THE TOKEN-LEVEL PERSPECTIVE IN DLLMS: Mean-field approximations ignore context from other sequence tokens, making them inaccurate proxies for autoregressive conditionals.Later token-level ELBO proxies better reflect masked-token prediction but do not resolve the underlying decomposition problem.
  • 3 THE CHALLENGE OF THE TOKEN-LEVEL PERSPECTIVE IN DLLMS: Individual token-level ELBO components are not conditional likelihoods, so inserting them into GRPO introduces an unknown inconsistency.The ELBO is formally a sequence-level lower bound, not a collection of independently valid token-level bounds.
  • 3 THE CHALLENGE OF THE TOKEN-LEVEL PERSPECTIVE IN DLLMS: The mismatch is structural: dLLM RL should respect holistic sequence generation rather than force the model into a token-level autoregressive framework.This observation motivates adapting the RL algorithm to the dLLM action space.

4 A PRINCIPLED SEQUENCE-LEVEL RL FRAMEWORK FOR DLLMS

ESPO aligns RL with dLLM generation by treating a sequence as one action and using an ELBO-based likelihood ratio, with normalization and KL choices that stabilize training. Sudoku ablations support the sequence-level ELBO design and the k2 KL estimator.

  • 4 A PRINCIPLED SEQUENCE-LEVEL RL FRAMEWORK FOR DLLMS: ESPO uses a sequence-level action space, an ELBO likelihood proxy, and stabilization techniques for policy gradients and KL regularization.This design is tailored to the non-autoregressive generation process of dLLMs.
  • 4.1 THE SEQUENCE-LEVEL POLICY OBJECTIVE WITH ELBO APPROXIMATION: Sequence-level ESPO avoids token-level summation by replacing the intractable sequence log-likelihood with an ELBO-based sequence importance ratio.The vanilla ratio is theoretically aligned but unstable because ELBO differences grow with sequence length and are exponentiated.
  • 4.1 THE SEQUENCE-LEVEL POLICY OBJECTIVE WITH ELBO APPROXIMATION: Per-token normalization of the ELBO log-ratio converts sequence-length-scaled differences to a stable scale and enables effective training.The stabilized ratio is incorporated into the sequence-level objective.
  • 4.1 THE SEQUENCE-LEVEL POLICY OBJECTIVE WITH ELBO APPROXIMATION: On Sudoku, sequence-level ELBO optimization is the only tested design that learns quickly, stably, and to the highest reward.Mean-field variants fail to learn effectively, while token-level ELBO is unstable and eventually collapses.
  • 4.2 STABLE KL-DIVERGENCE ESTIMATION: The k2 KL estimator avoids the exponential instability of k3 by using a quadratic ELBO-difference form with a stable gradient signal.In the Sudoku ablation, k2 learns stably and converges to the highest reward, whereas k3 stagnates and k1 is unstable.
  • 4.2 STABLE KL-DIVERGENCE ESTIMATION: Table 1 reports separate models for mathematics and planning tasks, with Δ measuring ESPO’s improvement over the corresponding non-RL base model.Countdown results marked † come from prior work, while other results are reproduced as described in the paper.

5 EXPERIMENT

ESPO is evaluated on mathematics, coding, and planning tasks using LLaDA and Dream models, with results showing consistent gains over prior dLLM-RL baselines and especially large improvements on planning tasks.

  • 5.2 BENCHMARK RESULTS: ESPO consistently and significantly outperforms original models and prior token-level dLLM-RL baselines, while remaining comparable to LLaDA-1.5 on coding tasks.Although trained only at sequence length 256, the improvements generalize to lengths 128, 256, and 512.
  • 5.2 BENCHMARK RESULTS: 20–40 absolute points on Countdown and up to over 60 points on Sudoku are the largest reported gains over strong baselines.These results are reported across sequence lengths and support the value of sequence-level optimization for planning tasks.
  • 5.2 BENCHMARK RESULTS: Mathematics and coding gains are more modest but consistently positive, surpassing previous token-level dLLM-RL methods averaged over sequence lengths.The base models’ pre-existing knowledge limits the maximum improvement available from RL fine-tuning alone.
  • 5.3 TRAINING DYNAMICS: On Sudoku, ESPO rapidly converges to a near-optimal policy while d1 and wd1 stagnate at low rewards; similar but less pronounced trends appear on Countdown.The comparison comes from training reward dynamics shown in Figure 3.
  • 5.4 ABLATION EXPERIMENTS: Increasing MC samples improves training stability, substantially accelerating Sudoku convergence from MC=1 to MC=2 or MC=4, while helping sparse-reward Countdown much less.The ablation compares MC sample counts of 1, 2, and 4.
  • 5.4 ABLATION EXPERIMENTS: ESPO remains robust across policy update values µ, reaching similarly high rewards on Countdown and Sudoku although smaller values such as 8 and 12 converge faster initially on Sudoku.The policy-update ablation varies µ while holding other settings fixed.
  • 5.5 DISCUSSION OF TRAINING EFFICIENCY: Theoretical FLOPs capture the overall trend of empirical wall-clock time because dLLM inference is compute-bound, with discrepancies attributed to hardware and communication factors.The cited factors include GPU utilization, memory bandwidth, and communication overhead.
  • 5.5 DISCUSSION OF TRAINING EFFICIENCY: Increasing MC from 1 to 4 raises ESPO’s total FLOPs by only about 47%, whereas VRPO training time scales almost linearly with MC.Generation costs are fixed and dominate dLLM computation, so additional MC samples primarily affect the less expensive policy-update term.

6 RELATED WORK

Prior work develops diffusion language models and applies reinforcement learning to language models, but RL methods designed for autoregressive likelihoods remain difficult to apply to dLLMs.

  • Diffusion language models: dLLMs use discrete or masked diffusion for sequence modeling rather than the left-to-right causal factorization central to autoregressive models.The related-work discussion frames this modeling distinction as foundational to subsequent RL challenges.
  • RL for language models: RL methods such as GRPO improve language-model reasoning with verifiable rewards, but applying them directly to dLLMs is challenging because likelihood computation is difficult.The difficulty follows from the mismatch between autoregressive RL objectives and diffusion generation.
  • RL for diffusion language models: Existing RL research for diffusion models includes score-matching, trajectory-balance, score-entropy, and masked-diffusion approaches, motivating more principled dLLM-specific optimization.The cited approaches span general discrete diffusion and masked diffusion models.

7 CONCLUSION

The paper identifies an incompatibility between autoregressive RL objectives and non-autoregressive dLLMs, then proposes ESPO as a sequence-level alternative based on ELBO likelihood estimation.

  • 7 CONCLUSION: ESPO treats sequence generation as a single action and uses the ELBO as a tractable proxy for sequence likelihood.It also introduces stabilized importance ratios and KL regularization for robust and efficient large-scale training.
  • 7 CONCLUSION: Across mathematics, coding, and planning benchmarks, ESPO consistently outperforms existing dLLM-RL baselines.The conclusion presents sequence-level optimization as a principled and empirically effective paradigm for RL in diffusion language models.

ETHICS STATEMENT

The paper frames dLLM reinforcement learning around a mismatch between autoregressive token-level objectives and non-autoregressive denoising, motivating sequence-level optimization with ELBO-based likelihood proxies.

  • B.1 GENERIC RL OBJECTIVE AND GRADIENT DERIVATIONS: The derivation begins from expected accumulated reward, uses the log-derivative trick and importance sampling, and identifies cumulative products as sources of exploding or vanishing gradients for long sequences.The advantage function is introduced for variance reduction.
  • C.1 ABLATION FOR ACTION SPACE AND LIKELIHOOD APPROXIMATION: Token-level RL methods decompose importance ratios into per-token terms, while sequence-level methods retain the sequence perspective and stabilize ratios through length normalization.The paper presents token-level approaches as practical approximations and sequence-level approaches as responses to their heuristic treatment of sequence-level rewards.
  • B.2 ADAPTING RL OBJECTIVES TO DIFFUSION LANGUAGE MODELS: dLLM trajectory-level RL requires T sequential network evaluations and backpropagation through them, creating a computational cost that grows linearly with the number of denoising steps.The passage contrasts this with autoregressive models, whose conditional probabilities can be computed in parallel through one forward pass.
  • B.2 ADAPTING RL OBJECTIVES TO DIFFUSION LANGUAGE MODELS: The ELBO preserves sequence-level validity because it is a variational lower bound on joint sequence log-likelihood, whereas token-decomposed ELBO terms do not approximate autoregressive conditionals.The mismatch arises because autoregressive conditionals use past-token context, while diffusion ELBO terms use bidirectional denoising context.
  • B.2 ADAPTING RL OBJECTIVES TO DIFFUSION LANGUAGE MODELS: ESPO treats complete sequence generation as one action and uses the ELBO as a tractable proxy for sequence likelihood.The framework also introduces stabilized importance-ratio and KL estimators for large-scale training.
  • C.1 ABLATION FOR ACTION SPACE AND LIKELIHOOD APPROXIMATION: The ablation compares token- and sequence-level objectives with mean-field and ELBO likelihood surrogates, with ESPO using sequence-level ELBO and normalized log-ratios.The four variants distinguish both the action granularity and the likelihood approximation.

C.2 ABLATION FOR KL DIVERGENCE ESTIMATOR

The KL-divergence estimator determines whether ESPO’s sequence-level optimization remains stable. Among k1, k2, and k3, k2 is the only estimator reported as both stable and theoretically sound.

  • Pitfalls in the k3 estimator: k3 destabilizes training because exponentiating ELBO differences magnifies estimation errors into rare, extreme gradient spikes.These spikes can derail the entire training process.
  • Pitfalls in the k3 estimator: Although k3 estimates KL without bias, its gradient instead estimates reverse-KL, leading to a different converged policy under bounded model capacity.This mismatch makes k3 unsuitable for the intended KL optimization objective.
  • Pitfalls in the k1 estimator: k1 provides no effective KL constraint because its gradient is zero, causing unconstrained policy drift and catastrophic collapse.The KL estimate becomes negative while the gradient norm grows uncontrollably.
  • The k2 estimator: k2 is the only robust estimator, producing stable KL estimates and gradient norms throughout sequence-level training.Its gradient is unbiased, unlike k3’s, and its quadratic form avoids exponential terms.
  • Variance reduction: Antithetic mask sharing and coupled sampling reduce ELBO-estimation variance by correlating policy samples and ensuring every token contributes to learning.Coupled sampling averages complementary masks, yielding a lower-variance and more stable objective.

E EXPERIMENT DETAILS

Experiments train task-specific or unified models across reasoning, planning, and coding benchmarks under fixed sequence-length and optimization protocols. The setup also distinguishes Dream and d1 Sudoku evaluations and reports coding-specific instability in baseline training.

  • Datasets and tasks: Training covers GSM8K, Math, Sudoku, Countdown, and four code benchmarks using separate task models or one AceCoder-87K model.Code evaluation uses HumanEval, HumanEval-Plus, MBPP, and MBPP-Plus.
  • Evaluation setup: All models train at maximum sequence length 256, while evaluation uses lengths 128, 256, and 512 to assess length generalization.The evaluation protocol therefore tests performance beyond the training length.
  • Denoising configuration: Planning tasks use 2-token unmasking per denoising step for Sudoku and Countdown, while other tasks use denoising steps equal to sequence length.For planning tasks with L = 256, this gives L/2 denoising steps.
  • Parameter-efficient fine-tuning: LoRA is used for GSM8K, Math, Countdown, and Sudoku, whereas coding uses full-parameter fine-tuning.Coding therefore uses a different adaptation regime from math and planning.
  • Sudoku settings: The main Sudoku results use the d1 setting because its longer generation lengths permit intermediate reasoning, while Dream restricts outputs to 24 tokens.Dream-setting results are additionally reported in Table 4.
  • Coding stability: Under coding full fine-tuning, wd1 is highly hyperparameter-sensitive and can reach gradient norms of 10^3–10^4 with immediate collapse, whereas ESPO remains stable.Careful tuning stabilizes wd1 but still leaves performance degradation on some metrics relative to the base model.

F DETAILS OF TRAINING FLOPS

The training-cost analysis focuses on vanilla dLLMs without KV cache and decomposes total FLOPs into generation and policy-update costs. Coupled sampling increases the policy-update term, but its effect is mild when sampling dominates.

  • Scope: The analysis excludes KV cache because partial KV-cache use can reduce downstream performance, so it studies the vanilla setting.This defines the scope of the computational-cost comparison.
  • Baseline FLOPs: Without coupled sampling, total per-sample cost is Ftotal = 2ND(K + 3µM), combining K forward passes with µM policy updates.The forward and backward approximations are Cforward = 2ND and Cbackward = 4ND.
  • Coupled sampling cost: With coupled sampling, total cost becomes Ftotal = 2ND(K + 6µM) because policy-update FLOPs double.The additional cost is confined to the policy-update term.
  • Cost implication: When K ≫ 6µM, increasing the Monte Carlo sample count raises overall training cost only mildly.Generation therefore dominates the cost under this condition.
  • Training dynamics: Sudoku training shows a phase transition near step 2500, where a sharp reward increase coincides with spikes in gradient norm and KL divergence.The signals indicate discovery of Sudoku’s strict logical rules and rapid model restructuring.

G.2 TRAINING CURVES ON MATHEMATICAL BENCHMARKS

Training rewards for ESPO and baselines are relatively stable and similar, but ESPO achieves superior held-out performance, indicating better generalization rather than merely fitting the training distribution.

  • Training dynamics: Similar training reward trends show that ESPO and baselines can fit the training distribution under their reward signals.The comparison is presented for MATH and GSM8K benchmarks on LLaDA-8B-Instruct.
  • Training dynamics: Weak correlation between final training reward and held-out accuracy can arise from overfitting.This limits the value of training reward alone as an indicator of benchmark performance.
  • Training dynamics: ESPO’s superior test performance demonstrates better generalization than the baselines despite similar training rewards.The reported comparison is based on the results in Table 1.
Loading 2512.03759v1…