Technical Summary: Learning from the Self-future (d-OPSD)

1. Problem Statement

While On-policy Self-distillation (OPSD) has proven effective for post-training autoregressive Large Language Models (AR-LLMs), its application to Diffusion Large Language Models (dLLMs) remains unexplored. Existing OPSD frameworks are inherently designed for AR models, creating two fundamental mismatches when applied to dLLMs:

1. Suboptimal Teacher Construction: AR-based OPSD injects privileged information (e.g., reference solutions) as a prefix to the prompt, relying on left-to-right conditioning $p(\text{suffix}|\text{prefix})$. This conflicts with dLLMs, which generate sequences non-autoregressively and support arbitrary-order generation, including the capability to condition on suffixes $p(\text{prefix}|\text{suffix})$. Furthermore, static reference solutions do not align with the "on-policy" nature of self-distillation.
2. Incompatible Supervision Level: AR models utilize token-level divergence supervision based on next-token prediction. In contrast, dLLMs operate via an iterative denoising process where multiple tokens are predicted simultaneously at each step, and only a subset is revealed while others remain masked. Token-level supervision is therefore incompatible with the step-wise state transitions of dLLMs.

Consequently, there is a need for a tailored OPSD framework that leverages the unique generative mechanics of dLLMs to enable efficient self-improvement without relying on external, stronger teacher models.

2. Methodology: d-OPSD

The authors propose d-OPSD (diffusion On-Policy Self-distillation), the first OPSD framework specifically designed for dLLMs. The method operates on a single model $p_\theta$ acting as both student and teacher, utilizing two core innovations:

A. Self-Teacher Construction via "Self-Future"

Instead of appending static reference solutions as a prefix, d-OPSD leverages the bidirectional conditioning capability of dLLMs.

• Mechanism: The student samples a full trajectory $Y = \{y_T, \dots, y_0\}$ from a fully masked state to a clean response $y_0$.
• Teacher Input: The teacher is constructed by partially revealing tokens from the student's own generated final response $y_0$ as a suffix condition. Specifically, at each denoising step $t$, the teacher input $y_{\text{teacher},t}$ consists of the current noisy sequence with a subset of indices $S_t$ (selected from the final answer) revealed.
• Concept: This allows the student to learn from its own "future experience" (the self-generated answer) conditioned as a suffix, adhering strictly to the on-policy paradigm where all data is generated by the student itself.

B. Step-Level Divergence Supervision

Recognizing that dLLMs decode via next-step prediction rather than next-token prediction, the authors shift the supervision objective from token-level to step-level.

• Mechanism: At each denoising step $t$, the model computes distributions for all tokens. The supervision focuses exclusively on the top-$k$ most confident tokens among the currently masked positions (the subset $K_t$ that dictates the state transition from $t$ to $t-1$).
• Objective: The loss function minimizes the Kullback-Leibler (KL) divergence between the student's distribution and the teacher's distribution over this specific subset $K_t$ at each step:
  $$L_t = \frac{1}{|K_t|} \sum_{i \in K_t} D_{KL}(P_{\text{student},t}^i \parallel P_{\text{teacher},t}^i)$$
  The total objective aggregates this step-level loss across the entire $T$-step trajectory.

3. Key Contributions

1. Identification of Suboptimality: The paper identifies that existing AR-centric OPSD formulations (prefix conditioning and token-level supervision) are fundamentally misaligned with dLLMs.
2. Novel Framework (d-OPSD): The authors introduce the first OPSD framework for dLLMs, featuring:
   • A suffix-based self-teacher construction that uses self-generated answers as privileged information, enabling learning from "self-future experience."
   • Step-level divergence supervision that aligns with the iterative denoising mechanics of diffusion models.
3. Empirical Validation: Extensive experiments demonstrate that d-OPSD enables a single model to drive its own self-improvement, outperforming both Reinforcement Learning with Verifiable Rewards (RLVR) and Supervised Fine-Tuning (SFT) baselines.

4. Experimental Results

Experiments were conducted on four reasoning benchmarks: GSM8K, MATH500 (mathematical reasoning), and Countdown, Sudoku (planning). The base model used was LLaDA-8B-Instruct.

• Reasoning Performance: d-OPSD consistently outperforms or matches RLVR (e.g., diffu-GRPO) and SFT baselines across all tasks and sequence lengths. For instance, on GSM8K (256 tokens), d-OPSD achieved 81.0% accuracy compared to 79.8% for diffu-GRPO and 78.8% for SFT.
• Sample Efficiency: d-OPSD demonstrates vastly superior sample efficiency, converging in approximately 10% of the optimization steps required by RLVR. For example, on GSM8K, d-OPSD required only 425 steps compared to 7,700 for diffu-GRPO.
• Comparison with AR-Style OPSD: An ablation study replacing the suffix-based teacher with an AR-style prefix-based teacher showed that the proposed suffix construction significantly outperforms the AR counterpart. Analysis of "Overlap Top-K" metrics revealed that the AR-style approach fails to introduce new knowledge (high overlap), whereas the d-OPSD suffix approach provides a suitable range of new thinking patterns for the student to learn.
• Ablation Findings:
  • Divergence Objective: Reverse KL outperforms Forward KL.
  • Retaining Ratio: A retaining ratio ($\rho_{\text{teacher}}$) of 0.25 yielded optimal results, though even 0.10 provided significant gains.
  • Top-k Selection: Deriving the top-$k$ subset from the teacher distribution yielded better performance than from the student distribution.
  • Stability: Pointwise KL clipping was found essential to prevent policy collapse after peak performance, a failure mode observed in both RLVR and unclipped d-OPSD.

5. Significance and Claims

The paper claims to open a promising pathway for dLLM post-training by establishing that self-distillation can be effectively adapted to diffusion models. The significance lies in:

• Efficiency: Providing a method that achieves state-of-the-art reasoning performance with drastically reduced computational cost (sample efficiency) compared to RLVR.
• Theoretical Alignment: Successfully bridging the gap between on-policy learning and diffusion mechanics by reframing privileged information as "self-future" suffixes and supervision as step-level divergence.
• Self-Improvement: Demonstrating that a single model can serve as its own teacher without requiring a stronger external model or static reference datasets, provided the generation trajectory is correctly utilized.

The authors note that while effective, the approach shares RLVR's susceptibility to policy collapse after reaching peak performance, suggesting a need for future work on stabilization techniques.