Source-linked AI summary
Learning from the Self-future: On-policy Self-distillation for dLLMs
Yifu Luo, Zeyu Chen, Haoyu Wang, Xinhao Hu, Yuxuan Zhang, Zhizhou Sha, Shiwei Liu
TL;DR
Existing on-policy self-distillation methods are designed for autoregressive models, leaving their formulation for diffusion LLMs unresolved. d-OPSD uses self-generated answers as suffix conditioning and step-level divergence, outperforming RLVR and SFT across four reasoning tasks while requiring around 10% of RLVR’s optimization steps.
Problem
Existing OPSD formulations rely on autoregressive prefix conditioning and token-level divergence, which are incompatible with diffusion LLMs’ arbitrary-order generation.
Method
d-OPSD constructs the self-teacher with self-generated answers as suffix privileged information and applies step-level divergence aligned with iterative denoising.
Results
Around 10% of the optimization steps required by RLVR are sufficient for d-OPSD, which consistently outperforms or matches SFT and RLVR across four reasoning tasks.
Takeaways & Limitations
d-OPSD provides a dLLM-specific on-policy self-distillation approach for self-improvement with superior reasoning performance and sample efficiency.
Takeaways & Limitations
Like RLVR, OPSD can undergo policy collapse after reaching peak performance in some settings.
Abstract
from arXiv · showhide
On-policy self-distillation (OPSD) has proven effective for post-training large language models (LLMs), yet its application to diffusion LLMs (dLLMs) remains unexplored. Existing OPSD methods are inherently autoregressive-centric. They inject privileged information via left-to-right prefix conditioning with token-level divergence supervision, a design that fundamentally conflicts with the arbitraryorder generation of dLLMs. We introduce d-OPSD, the first OPSD framework tailored for dLLMs. Our approach makes two core contributions. First, we reframe self-teacher construction by using self-generated answers as suffix conditioning, enabling the student model to learn from "self future-experience" rather than privileged prefixes. Second, we shift supervision from token-level to step-level, aligning training with the iterative denoising process of dLLMs. Experiments across four reasoning benchmarks show that d-OPSD consistently outperforms RLVR and SFT baselines with superior sample efficiency, requiring only around 10% of the optimization steps by RLVR and opening a promising pathway for dLLM posttraining. The code is available at https://github.com/xingzhejun/d-OPSD.
1 Introduction
d-OPSD introduces the first OPSD framework for diffusion LLMs, replacing autoregressive prefix-based self-teacher construction and token-level supervision with suffix conditioning and step-level divergence. It improves dLLM reasoning performance and sample efficiency over RLVR.
- Contributions: d-OPSD is presented as the first application of OPSD to dLLMs and enables a single model to serve as both student and self-teacher.The framework is designed to drive self-improvement without relying on a stronger external teacher model.
- Motivation: Existing OPSD methods remain largely unexplored for dLLMs because their autoregressive design appends privileged information as a prefix and applies token-level divergence.dLLMs generate non-autoregressively, making these assumptions incompatible with arbitrary-order denoising.
- Results: d-OPSD achieves superior reasoning performance and sample efficiency compared with the RLVR baseline for dLLMs.The comparison is made against diffu-GRPO, identified as the RLVR baseline in Figure 1.
- Method: d-OPSD constructs the self-teacher by conditioning on self-generated answers as suffix privileged information, enabling learning from “self future-experience”.The suffix formulation exploits dLLMs’ ability to incorporate privileged information without left-to-right prefix constraints.
- Method: d-OPSD replaces token-level divergence with step-level divergence aligned with the iterative denoising process of dLLMs.Each denoising step is treated as an independent Markov transition, while dLLMs predict masked tokens simultaneously and remask some outputs.
2 Preliminaries
dLLMs train by progressively masking clean responses and infer by iterative denoising from fully masked sequences. OPSD uses the same model as teacher and student with different conditioning contexts, but its autoregressive formulation faces challenges under dLLMs’ arbitrary-order generation.
- dLLM paradigms: dLLM training gradually corrupts clean responses by replacing tokens with a special mask token.The forward process operates over sequence positions and diffusion steps t from 0 to T.
- dLLM paradigms: dLLM inference is a T-step iterative denoising process from a fully masked sequence yT = {mask}L to a clean response y0.At each step, the model predicts all tokens, samples and reveals the top-k most confident masked positions, and retains the others as mask.
- OPD and OPSD: Autoregressive OPSD supervises the student through token-level divergence between teacher and student next-token distributions along student-sampled trajectories.Reverse KL is used in the described formulation, although forward KL and generalized Jensen-Shannon divergence can also be employed.
- OPD and OPSD: OPSD instantiates teacher and student from the same model pθ but differentiates their conditioning contexts using privileged information appended to the prompt.For autoregressive models, the teacher-specific prompt is x* = x + y*, while the response is generated solely by the student.
- OPD and OPSD: Adapting OPSD to dLLMs is challenging because their arbitrary-order generation offers a different way to inject privileged information that better aligns with on-policy learning.The passage identifies this issue as a fundamental challenge for transferring the autoregressive formulation to dLLMs.
3 Methods
d-OPSD uses one dLLM as both student and teacher, with the teacher conditioned on selectively revealed tokens from the student’s own generated answer. It replaces autoregressive token-level supervision with step-level divergence over the masked tokens that drive each denoising transition.
- Self-teacher construction: The student samples an on-policy T-step trajectory, while the teacher uses selectively revealed tokens from the final response as self-generated future information.Both distributions are instantiated from the same model pθ, but their conditioning inputs differ.
- Step-level supervision: At each denoising step, d-OPSD supervises the student with divergence between teacher and student distributions rather than autoregressive token-level targets.This matches dLLMs’ next-step decoding process, where only selected tokens are revealed at each step.
- Step-level supervision: The step-level KL divergence is computed only over the top-k confident masked positions that determine the transition from t to t−1.The top-k subset can be selected from either distribution; the teacher-derived subset performs better in ablation.
- d-OPSD procedure: d-OPSD constructs student and teacher inputs independently at each denoising step before minimizing a step-level learning objective with the shared model pθ.The procedure begins from the student’s on-policy trajectory for prompt x and applies the two conditioning constructions at every step.
- Distinction from prior work: Unlike prior dLLM self-distillation methods, d-OPSD reveals answers generated on-policy by the student rather than ground-truth answers from static datasets.The paper identifies this self-experience distinction as a fundamental difference from d3llm and Cd4lm.
4 Experiments
Across four reasoning tasks, d-OPSD matches or surpasses SFT and RLVR while converging with substantially higher sample efficiency. Experiments also validate the self-teacher, explain the advantage of dLLM-specific construction, assess training choices, and identify policy collapse as a failure mode.
- 4.1 Toy Verification: The self-teacher can resume correct answers and guide high-quality distillation, validating the teacher construction.Higher ρteacher makes self-teacher performance nearly match its origin at Pass@8.
- 4.2 Main Results: d-OPSD consistently outperforms or matches SFT and RLVR across four reasoning tasks, achieving state-of-the-art performance in most settings.The tasks cover GSM8K, MATH500, 4x4 Sudoku, and Countdown (3 numbers), using configurations consistent with diffu-GRPO.
- 4.2 Main Results: Around 10% of the optimization steps are sufficient for d-OPSD to converge, demonstrating vastly superior sample efficiency relative to the RLVR baseline.Even k = 1 surpasses the RLVR baseline, which uses group k = 8 rollouts.
- 4.3 Teacher Construction: d-OPSD significantly outperforms the AR-style teacher construction because its self-teacher introduces more transferable new knowledge and thinking patterns.AR-style OPSD has Overlap Top-Kt nearly 1, whereas d-OPSD lies in a suitable range for transferring new knowledge.
- 4.4 Ablations: Reverse KL, teacher-derived Kt selection, pointwise clipping, and suitable retaining ratios improve training, while k = 1 retains an advantage over RLVR in sample efficiency.The no-clipping variant begins collapsing around step 150 and drops to 69.37 by step 500.
- 4.5 Failure Modes: Like RLVR, d-OPSD can undergo catastrophic policy collapse after reaching peak performance, potentially because model-seeking behavior becomes overly narrow.The collapse is observed in some settings and can degrade training catastrophically.
5 Related Works
Prior work established knowledge distillation and sequence-level off-policy distillation before extending distillation on-policy to address exposure bias by training on student-generated outputs.
- On-policy Distillation: Knowledge distillation transfers knowledge from a large teacher to a smaller student using the teacher’s soft output distributions.
- On-policy Distillation: Sequence-level distillation established dominant off-policy approaches in subsequent work.
- On-policy Distillation: On-policy distillation addresses exposure-bias mismatch by shifting the training distribution to the student’s own generations.
6 Conclusion … B Self-teacher Construction Illustrations
d-OPSD is presented as the first on-policy self-distillation framework for dLLMs, using self-generated suffixes as privileged information and step-level rather than token-level supervision. The illustrations show how this self-teacher is constructed from an on-policy trajectory and contrast it with appending a reference solution.
- 6 Conclusion: d-OPSD is the first on-policy self-distillation approach for dLLMs, aligning its design with their on-policy and diffusion characteristics.Its self-teacher uses self-generated answers as suffix conditioning, while dense divergence supervision moves from token level to step level.
- A.1 Additional Preliminaries: Block-diffusion partitions a response into B contiguous, non-overlapping blocks, each containing L′ = LThe supplied passage introduces the block structure and gives the block-length notation, but is truncated after the expression.
- A.1 Additional Preliminaries: Block-diffusion decodes autoregressively across blocks and with diffusion within each block, beginning the next block only after the previous block is fully decoded.This hybrid inference strategy is described as common in current dLLMs.
- A.2 Additional Related Works: Existing reinforcement-learning work for dLLMs commonly applies GRPO, estimating log-probabilities through one-step estimation or the ELBO.The passage also identifies heavy-computation challenges from RLVR, but the supplied text is truncated.
- B Self-teacher Construction Illustrations: The self-teacher construction illustration uses a question sampled from the GSM8K training set and omits some mask and end-of-text tokens for brevity.This example is presented to explain how the Section 3.1 construction works.
- B Self-teacher Construction Illustrations: The construction first samples an on-policy trajectory from the student model and obtains the final clean answer as the self-generated future.The supplied passage labels this answer the self-generated future.
- B Self-teacher Construction Illustrations: For comparison, the AR-style construction appends a reference solution to the prompt.The paper illustrates this alternative in Figure 8.
C Additional Implementation Details … D Additional Experiment Details
The implementation stabilizes dLLM training through pointwise clipping and memory-efficient trajectory batching, while the default objective is computed only on correct generations. The supplied details also note an AR-style teacher-construction illustration and compare correct-only with all-generation training.
- C.1 Per-Token pointwise clipping: Pointwise clipping stabilizes vocabulary-level divergence contributions and improves performance.Token-level divergence is highly skewed across vocabulary entries, motivating clipping.
- C.2 Inputs Concatenation: Full-attention gradients across all denoising steps can cause Equation (12) to exceed memory capacity.The gradients must otherwise remain stored until the trajectory is fully decoded.
- C.2 Inputs Concatenation: Figure 8 presents an AR-style teacher construction.The supplied figure caption identifies the construction but provides no further procedural details.
- C.2 Inputs Concatenation: Inputs from every trajectory step are concatenated into a batch tensor shaped (bsz×steps, seq-length).This engineering technique reorganizes all status tensors across trajectory steps into one batch.
- C.2 Inputs Concatenation: Because the concatenated inputs share one model, their gradients remain constant and need not be stored separately.This reduces the storage requirement described for the original stepwise computation.
- C.3 Compute only on Correct Generations: By default, Equation (12) is optimized only on correct generations.The paper states that this default setting achieves superior results, although computing on all generations also improves reasoning performance.
- C.3 Compute only on Correct Generations: Computing the objective on all generations improves reasoning performance but underperforms the correct-generation default.Detailed experimental results are provided in Section E.1.
D.1 Training Details
d-OPSD was implemented with TRL and LoRA, using specified optimization and hardware settings. Sudoku examples entered the loss only above a 0.25 score threshold, while diffu-GRPO was reproduced with its default settings on 8 GPUs.
- Implementation and optimization: d-OPSD used TRL with LoRA at rank r = 128 and scaling factor α = 64.Training used 4 NVIDIA GPUs, a learning rate of 5 × 10−6, accumulation steps of 1, and AdamW.
- Task-specific criterion: For Sudoku, generations were included in loss computation only when their score exceeded the 0.25 threshold.The threshold addresses Sudoku’s score range of [0, 1], where answers are not categorized as simply right or wrong.
- Baseline reproduction: The diffu-GRPO RLVR baseline was reproduced on 8 NVIDIA GPUs using its default settings.The reproduction corresponds to the baseline shown in Figure 1.
D.2 Toy Experiment Details and Examples · E Additional Experiment Results
The toy experiment uses partially revealed self-generated answers to construct self-teachers, while controlling generation length, unmasking, and block-level information leakage. No passage supplied here reports additional experiment results.
- D.2 Toy Experiment Details and Examples: Generation length is fixed at 256 for all tasks.After self-teacher construction, the remaining mask-token count is smaller than 256.
- D.2 Toy Experiment Details and Examples: When entering a new block, all previously unmasked tokens in that block are cleared.The new block is therefore filled entirely with mask tokens.
- D.2 Toy Experiment Details and Examples: This block reset prevents leakage of the final answer during self-teacher construction.The passage specifically warns against retaining the final answer between <answer><answer>.
- D.2 Toy Experiment Details and Examples: The examples include a question from the GSM8K training set and its self-generated answer.These are presented in Figures 9 and 10.
- D.2 Toy Experiment Details and Examples: Self-teachers are constructed by partially revealing the final generation.This construction is illustrated in Figure 11.
E.1 Additional Ablation Studies
The ablations show that fixing the teacher model improves d-OPSD’s reasoning performance, while leaving it unfixed still nearly matches RLVR baselines.
- Fixing the Teacher: Fixing the teacher model produces greater reasoning-performance gains for d-OPSD.The comparison is reported in Table 10.
- Fixing the Teacher: Without a fixed teacher, d-OPSD’s reasoning performance nearly matches the RLVR baselines.This result further supports d-OPSD’s effectiveness under the unfixed-teacher setting.
- Self-teacher Toy Experiment: The additional ablation studies include a toy experiment examining the self-teacher.This analysis is presented in Figure 11.
Compute only on Correct Generations · E.2 Qualitative Examples on GSM8k · E.3 Failure Mode
Computing loss on all trajectories slightly degrades performance but still surpasses the RLVR baseline. A GSM8k example shows the proposed approach correcting an RLVR error, while Figure 12 illustrates the reported collapse failure mode.
- Compute only on Correct Generations: Computing loss on all trajectories causes a slight performance degradation.This comparison is reported in Table 11.
- Compute only on Correct Generations: Despite this degradation, computing loss on all trajectories still outperforms the RLVR baseline.The result is presented as a reasoning performance comparison in Table 11.
- E.2 Qualitative Examples on GSM8k: A GSM8k testing example shows the RLVR model producing an incorrect answer while the proposed approach produces the correct one.The example is shown in Figure 13.
- E.2 Qualitative Examples on GSM8k: The qualitative GSM8k comparison demonstrates a case where the proposed approach succeeds despite RLVR failure.The passage specifically identifies this as an example from the GSM8k testing set.
- E.3 Failure Mode: Figure 12 presents the failure mode discussed in Section 4.5.The figure is captioned as the failure mode of collapse.
- E.3 Failure Mode: The reported failure mode is collapse.Figure 12 is titled “Failure Mode of collapse.”