On-Policy Diffusion Language Model

• The paper introduces OPDLM, where an autoregressive LM is transformed into a diffusion LM via on-policy reverse decoding trajectories.
• It addresses training-inference mismatches by aligning the student's iterative denoising states with its own generation process.
• Empirical results show that OPDLM achieves remarkable token efficiency and competitive performance across tasks such as MMLU, GSM8K, and MATH500.

On-Policy Diffusion LLM (OPDLM) denotes an on-policy ARLM-to-DLM conversion method in which a pretrained autoregressive LLM is transformed into a diffusion LLM by training a bidirectionally-attentive student on its own reverse decoding trajectories while distilling token-level targets from the original frozen ARLM (Su et al., 4 Jun 2026). The method is introduced to address two distribution shifts that arise in prior ARLM-to-DLM conversion: the shift from next-token prediction to a diffusion objective, and the shift between randomly masked training states and confidence-guided reverse unmasking at inference. In subsequent diffusion-LM work, the term also functions as a broader label for on-policy post-training regimes that optimize or distill over iterative denoising trajectories rather than left-to-right token chains (Zhan, 5 Oct 2025).

1. Problem setting and background

The OPDLM formulation begins from the observation that standard masked diffusion LLMs are trained on randomly corrupted sequences but decoded through reverse unmasking. In the paper’s notation, given a clean sequence $x_0 \sim p_{\mathrm{data}}$, each token is independently replaced by a mask token $m$ according to diffusion time $t \in [0,1]$: $$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$ The corresponding standard MDLM objective is

$$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$

where $\mathcal{M}(x)=\{i:x^i=m\}$ is the set of masked positions. For block diffusion, the sequence is partitioned into $B$ blocks and the model conditions on the clean prefix plus the corrupted active block (Su et al., 4 Jun 2026).

The motivation for OPDLM is that this standard recipe leaves two distinct mismatches unresolved. First, ARLMs are trained with next-token prediction, whereas DLMs are trained with a masked denoising or diffusion objective; the paper states that useful knowledge acquired by the ARLM during training can be weakened or lost if the objective is simply replaced. Second, standard DLMs are optimized on randomly masked states from the forward corruption process but deployed on reverse decoding trajectories induced by the model and sampler (Su et al., 4 Jun 2026).

This train-inference discrepancy is independently emphasized in later work on masked diffusion LLMs. "MDPO: Overcoming the Training-Inference Divide of Masked Diffusion LLMs" characterizes the divide as the contrast between random masking during training and a progressive, confidence-guided refinement schedule during inference, and argues that this mismatch is especially harmful for reasoning tasks because the model must progressively refine an answer (He et al., 18 Aug 2025). This suggests that OPDLM’s central intervention is not merely architectural conversion, but a redefinition of the training state distribution.

2. Trajectory formulation and the OPDLM objective

A central conceptual move in OPDLM is to recast diffusion training in terms of trajectories. Let $\Gamma$ be a distribution over trajectories $\tau$, where $\tau(t)$ is the state at time $m$0, and let $m$1 be a distribution over times along a trajectory. The paper defines a generalized block diffusion objective: $m$2 Here $m$3 specifies the trajectory distribution, $m$4 specifies which state on the trajectory is supervised, $m$5 is a time weighting, $m$6 is the target distribution at masked position $m$7, and $m$8 is the student prediction at that position (Su et al., 4 Jun 2026).

Standard MDLM training is recovered when $m$9 is the forward random masking trajectory distribution, $t \in [0,1]$0, $t \in [0,1]$1, and $t \in [0,1]$2. OPDLM changes the trajectory distribution from forward corruption to reverse decoding by sampling

$t \in [0,1]$3

that is, a reverse decoding trajectory generated by the student itself. The on-policy objective becomes

$t \in [0,1]$4

The remaining question is the choice of target distribution $t \in [0,1]$5. OPDLM answers this by introducing self-OPD. The student is an ARLM checkpoint converted into a DLM by changing the attention mask to blockwise or bidirectional attention; the teacher is the original frozen ARLM; and the student generates the reverse trajectories on which it is trained. For each masked token $t \in [0,1]$6 in block $t \in [0,1]$7, the teacher target is

$t \in [0,1]$8

where $t \in [0,1]$9 is the terminal generated sequence from the student’s reverse trajectory. Substituting this target into the on-policy objective yields the final OPDLM loss: $$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$0 The paper uses $$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$1 following prior block diffusion work (Su et al., 4 Jun 2026).

Conceptually, this makes OPDLM different from standard knowledge distillation on dataset states or teacher-generated states. Supervision is queried on student-generated reverse trajectories, so the training distribution is aligned with the same decoding process used at inference.

3. Training pipeline, rollout curriculum, and decoding

The OPDLM training algorithm is explicitly trajectory-based. The paper describes the procedure as follows: input the frozen ARLM teacher $$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$2, the student DLM $$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$3, and the sampler $$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$4; sample a reverse trajectory

$$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$5

set $$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$6; sample a non-terminal step $$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$7 uniformly; set $$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$8, $$q^{t\mid0}(x_t^i \mid x_0^i) = \alpha_t \delta_{x_0^i}(x_t^i) + (1-\alpha_t)\delta_{m}(x_t^i), \qquad \alpha_t = 1-t.$$9; and for each block $$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$0 and each masked position $$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$1, compute the teacher distribution $$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$2, the student distribution $$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$3, and accumulate $$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$4 (Su et al., 4 Jun 2026).

A practical complication is that reverse rollouts can be poor immediately after conversion. To stabilize early training, the paper introduces a curriculum on rollout length. Let $$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$5 and $$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$6 denote the minimum and maximum rollout lengths. At step $$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$7,

$$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$8

Training therefore begins with short rollouts and gradually increases the maximum sequence length over a warmup schedule (Su et al., 4 Jun 2026).

At inference, OPDLM uses reverse unmasking or confidence-guided decoding and can generate multiple tokens per step. The paper studies two controls on throughput: the confidence threshold $$\mathcal{L}_{\mathrm{MDLM}}(\theta) = - \mathbb{E}_{x_0 \sim p_{\mathrm{data}},\, t \sim \mathcal{U}(0,1),\, x_t \sim q^{t \mid 0}(\cdot\mid x_0)} \left[ \frac{1}{t} \sum_{i\in\mathcal{M}(x_t)} \log p_\theta(x_0^i\mid x_t) \right],$$9, where tokens whose confidence exceeds $\mathcal{M}(x)=\{i:x^i=m\}$0 can be committed earlier, and the block size, where larger blocks permit more parallel token updates. On MATH-500 with block size 4, static decoding with $\mathcal{M}(x)=\{i:x^i=m\}$1 gives 1 token per step, while lowering the threshold to $\mathcal{M}(x)=\{i:x^i=m\}$2 gives over 2 tokens per step with moderate accuracy cost (Su et al., 4 Jun 2026).

These details clarify that, in OPDLM, “on-policy” refers to the distribution of intermediate denoising states encountered under the student’s own sampler. The algorithm is therefore aligned with iterative unmasking rather than with a left-to-right causal rollout.

4. Empirical profile

The paper’s headline empirical claim is data efficiency. It reports that OPDLM requires about 15× to 7,000× fewer tokens than prior DLM baselines. A concrete example given is OPDLM-8B trained on 0.066B tokens with 4.2 × 10$\mathcal{M}(x)=\{i:x^i=m\}$3 FLOPs, compared with LLaDA-8B: 1500B tokens, Dream-7B: 580B tokens, SDAR-8B: 55B tokens, and Fast-dLLM-v2-7B: 1B token (Su et al., 4 Jun 2026).

The reported general benchmark results for OPDLM-8B are: MMLU 70.9, MMLU-Pro 53.7, GPQA-Diamond 36.1, GSM8K 87.1, MATH-500 71.2, AIME-24 14.7, AIME-25 12.4, HumanEval-base 59.8, and MBPP-base 48.7. The paper states that OPDLM is competitive across general knowledge, mathematics, code generation, and instruction following, and that its relative advantage grows with task difficulty, especially on GPQA-Diamond, AIME, and LiveCodeBench (Su et al., 4 Jun 2026).

The work also reports preservation of pretrained ARLM priors not explicitly trained on, including zero-shot extended thinking and zero-shot multilingual ability. This is presented as evidence that distillation from the original ARLM preserves broader latent capabilities (Su et al., 4 Jun 2026).

Task-specialized post-training is demonstrated through OPDLM-MATH and OPDLM-MATH-Thinking. The paper reports OPDLM-MATH-4B-Thinking: GSM8K 91.7, MATH500 90.2, AIME24 43.3 and OPDLM-MATH-8B-Thinking: GSM8K 93.8, MATH500 92.4, AIME24 50.0 (Su et al., 4 Jun 2026).

The ablation results sharpen the interpretation of why the method works. The paper compares an Off-Policy baseline, OPDLM$\mathcal{M}(x)=\{i:x^i=m\}$4, and OPDLM, and concludes that On-policy data is the primary driver of performance. It further reports that the exact choice of masking trajectory is less important once the data is on-policy, because OPDLM and OPDLM$\mathcal{M}(x)=\{i:x^i=m\}$5 are often close. A second comparison with supervised fine-tuning on ARLM-generated responses finds that SFT is a strong baseline under static decoding but degrades more strongly under dynamic sampling, whereas OPDLM is more stable because it is trained on its own on-policy generations. The appendix additionally reports that for 4B and 8B students, self-distillation from the same-sized teacher works best or matches a much larger teacher, while for 0.6B students, a larger teacher can help more (Su et al., 4 Jun 2026).

5. Relation to the broader on-policy diffusion-LM literature

Subsequent work broadens the meaning of on-policy diffusion language modeling beyond ARLM-to-DLM conversion. "Principled and Tractable RL for Reasoning with Diffusion LLMs" presents Amortized Group Relative Policy Optimization (AGRPO) as the first principled on-policy RL algorithm designed specifically for dLLMs or on-policy diffusion language modeling (OPDLM) settings. AGRPO reinterprets the GRPO objective as an expectation over denoising timesteps and estimates that expectation with Monte Carlo sampling, yielding what the paper describes as the first tractable, faithful adaptation of policy gradient methods to dLLMs. Starting from LLaDA-8B-Instruct, it reports +7.6% absolute on GSM8K, +5.2% absolute on MATH, and +29.6% absolute on Countdown over the diffusion baseline, and states that the model can match the baseline with roughly 4x fewer sampling steps on GSM8K (Zhan, 5 Oct 2025).

"MDPO: Overcoming the Training-Inference Divide of Masked Diffusion LLMs" addresses the same train-inference discrepancy from an RL perspective. It reframes denoising as sequential decision-making, treats each denoising iteration as an action, and uses a PPO-style clipped objective with importance sampling and group-relative advantage estimation. The paper states that MDPO matches the performance of the previous state-of-the-art (SOTA) method with 60x fewer gradient updates, while achieving average improvements of 9.6% on MATH500 and 54.2% on Countdown over SOTA under the same number of weight updates. It also introduces Running Confidence Remasking (RCR), a training-free inference replacement that tracks the running maximum confidence and yields additional gains when combined with MDPO (He et al., 18 Aug 2025).

"Diffusion Policy Optimization without Drifting Apart" studies instability in diffusion policy-gradient methods through the double-drift phenomenon, consisting of ELBO–likelihood drift and proxy-gradient drift. Its practical intervention, DiPOD, augments each diffusion policy-gradient update with an on-policy ELBO regularizer: $\mathcal{M}(x)=\{i:x^i=m\}$6 For diffusion LLM post-training on GSM8K, MATH500, Countdown, and Sudoku, the paper reports that DiPOD substantially stabilizes training and reaches higher rewards than previous methods. The language experiments use LLaDA-8B-Instruct, sequence length 256, decoding steps 128, and $\mathcal{M}(x)=\{i:x^i=m\}$7 (Jiang et al., 11 Jun 2026).

"Learning from the Self-future: On-policy Self-distillation for dLLMs" introduces d-OPSD, described as the first on-policy self-distillation framework tailored specifically for diffusion LLMs. Its two core changes are suffix conditioning from the model’s own final answer and step-level divergence aligned with denoising. The paper states that d-OPSD consistently outperforms RLVR and SFT baselines with superior sample efficiency, requiring only around 10% of the optimization steps by RLVR, with examples including GSM8K: 7700 vs 425, MATH500: 6600 vs 100, Countdown: 5000 vs 175, and Sudoku: 3800 vs 425 (Luo et al., 16 Jun 2026).

Outside language, "DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models" generalizes OPD from token sequences to continuous-state Markov processes and derives a closed-form per-step KL objective that unifies stochastic SDE and deterministic ODE refinement. The paper’s central claim is that on-policy distillation is fundamentally about matching local transition kernels along student-generated trajectories. Although the work targets multi-task diffusion models rather than dLLMs, it gives a unified perspective under which OPDLM can be understood as one instance of on-policy transition matching (Li et al., 14 May 2026).

6. Interpretation, misconceptions, and open questions

A common misconception is to treat OPDLM as ordinary supervised conversion from an ARLM to a DLM. The defining feature is not merely that a causal model is converted to bidirectional or blockwise attention; it is that the student is trained on its own reverse diffusion trajectories and supervised by the frozen ARLM on those states. In that sense, OPDLM is a form of post-training rather than fresh diffusion pretraining, and the paper explicitly argues that Diffusion LLM transformation should be viewed as a post-training problem, not as pretraining from scratch (Su et al., 4 Jun 2026).

A second misconception is to equate all on-policy dLLM methods with the same optimization principle. OPDLM uses forward KL distillation from an ARLM teacher on student-generated trajectories; AGRPO uses an unbiased Monte Carlo estimate of a GRPO-style on-policy objective over denoising timesteps; MDPO casts denoising trajectories as Markovian sequential decision-making and optimizes them with PPO-style clipping and group-relative advantages; DiPOD adds an on-policy ELBO regularizer to stabilize variational policy-gradient updates; and d-OPSD performs step-level self-distillation from the model’s own future outputs (Zhan, 5 Oct 2025, He et al., 18 Aug 2025, Jiang et al., 11 Jun 2026, Luo et al., 16 Jun 2026). This suggests that “on-policy” is a shared training-distribution principle, while the learning signal may be distillation, RL, or self-distillation.

The broader literature also identifies open stability questions. d-OPSD notes a failure mode of policy collapse after reaching peak performance (Luo et al., 16 Jun 2026). DiPOD attributes instability in diffusion policy-gradient training to double drift, namely the separation of the ELBO from the true likelihood and the resulting misalignment of the proxy policy gradient (Jiang et al., 11 Jun 2026). MDPO argues that the main bottleneck in MDLMs is not just model scale, but the mismatch between pretraining and inference dynamics, and highlights Answer Backslide, where an intermediate correct reasoning state is later refined into an incorrect final answer (He et al., 18 Aug 2025). These observations indicate that on-policy state matching alone does not eliminate all optimization pathologies.

Within this landscape, OPDLM occupies a specific position. It is an on-policy distillation method for converting pretrained ARLMs into DLMs efficiently, with the additional claim that the resulting model avoids the prohibitive cost of DLM pretraining. The paper’s broader implication is therefore both practical and conceptual: diffusion LLM transformation can be treated as a post-training problem in which a student learns directly on the reverse decoding states it will encounter at inference, while retaining the pretrained knowledge of the original autoregressive model (Su et al., 4 Jun 2026).