dOPSD: On-Policy Self-Distillation for Diffusion Language Models
dOPSD enhances reasoning in diffusion language models by self-distillation, excelling on Dream and LLaDA datasets.
Key Findings
Methodology
dOPSD derives privileged information from the student's own denoising trajectory, replacing external labels, for self-distillation. This method leverages the model's own decoding steps to provide dense, on-policy supervision signals.
Key Results
- On Dream, dOPSD improved GSM8K accuracy from 81.41% to 83.04% and MATH500 from 38.97% to 42.20%.
- On LLaDA, dOPSD increased MATH500 performance by 4.76 points, from 31.24% to 36.00%.
- dOPSD improved HumanEval performance by 4.17 points, from 52.54% to 56.71%.
Significance
dOPSD significantly enhances diffusion language models' performance in mathematical reasoning and code generation tasks, overcoming traditional methods' bottlenecks in reasoning capabilities, and offering new possibilities for language model applications.
Technical Contribution
dOPSD introduces a technique to extract privileged information from the model's own decoding trajectory, avoiding reliance on external labels, and providing a new self-distillation strategy that significantly enhances reasoning capabilities.
Novelty
dOPSD is the first method to utilize the model's own decoding trajectory for privileged information extraction in self-distillation, offering more effective reasoning capability improvements compared to traditional methods relying on external labels.
Limitations
- dOPSD relies on the model's own decoding trajectory, which may not provide sufficient privileged information for some complex tasks.
- The method requires verification of each decoding result's correctness, increasing computational overhead.
Future Work
Future work could explore how to effectively extract privileged information in more complex tasks and reduce computational overhead to enhance the method's practicality.
AI Executive Summary
Diffusion language models (dLLMs) have gained attention as a non-autoregressive text generation method due to their parallel decoding and bidirectional context advantages. However, enhancing their reasoning capabilities in post-training remains a challenge. Traditional supervised fine-tuning and reinforcement learning methods have limitations in this area.
dOPSD derives privileged information from the student's own denoising trajectory, replacing external labels, for self-distillation. This method leverages the model's own decoding steps to provide dense, on-policy supervision signals, significantly enhancing the model's performance in mathematical reasoning and code generation tasks.
Experimental results show that dOPSD surpasses existing supervised and self-distillation baseline methods on Dream and LLaDA, demonstrating its potential in improving reasoning capabilities. Nevertheless, the method's computational overhead and applicability in complex tasks require further investigation.
Deep Analysis
Background
In recent years, diffusion language models have gained widespread attention for their advantages in text generation tasks. Unlike traditional autoregressive models, diffusion models generate text through a denoising process, offering parallel decoding and bidirectional context. However, enhancing their reasoning capabilities in post-training remains a challenge. Traditional supervised fine-tuning methods are prone to exposure bias, while reinforcement learning methods are difficult to apply due to sparse rewards.
Core Problem
Enhancing the reasoning capabilities of diffusion language models has been a challenge. Traditional methods lack dense, on-policy supervision signals, leading to suboptimal performance in complex tasks. Effectively improving the model's reasoning capabilities in post-training is a key research focus.
Innovation
dOPSD derives privileged information from the student's own denoising trajectory, replacing external labels, for self-distillation. This method leverages the model's own decoding steps to provide dense, on-policy supervision signals, significantly enhancing the model's performance in mathematical reasoning and code generation tasks.
Methodology
- �� Extract privileged information from the student's denoising trajectory.
- �� Utilize the model's own decoding steps for self-distillation.
- �� Verify the correctness of each decoding result to ensure the reliability of the distillation signal.
Experiments
Experiments were conducted on Dream and LLaDA, using GSM8K and MATH500 datasets to evaluate mathematical reasoning capabilities, and HumanEval and MBPP for code generation capabilities. Baseline methods include supervised fine-tuning and traditional self-distillation methods.
Results
dOPSD improved GSM8K accuracy from 81.41% to 83.04% on Dream, and MATH500 from 38.97% to 42.20%. On LLaDA, dOPSD increased MATH500 performance by 4.76 points, from 31.24% to 36.00%.
Applications
dOPSD can be directly applied to language model tasks requiring high reasoning capabilities, such as mathematical reasoning and code generation. Its independence from external labels makes it more flexible in practical applications.
Limitations & Outlook
dOPSD relies on the model's own decoding trajectory, which may not provide sufficient privileged information for some complex tasks. Additionally, the method requires verification of each decoding result's correctness, increasing computational overhead.
Plain Language Accessible to non-experts
Imagine a factory where workers are on an assembly line. Traditional language models are like a single worker handling one part at a time, while diffusion language models are like a team handling multiple parts simultaneously. dOPSD is like giving this team a new way to work together more effectively, improving overall production efficiency. By observing the team's work process, dOPSD finds the most effective working methods, enhancing overall productivity.
ELI14 Explained like you're 14
Imagine you're playing a puzzle game. Traditional methods are like solving the puzzle step by step, while dOPSD is like giving you a new strategy to consider multiple clues at once. This way, you can find the answer faster! dOPSD helps you find the most effective way to solve puzzles by observing your solving process, boosting your game skills.
Glossary
Diffusion Language Model
A model that generates text through a denoising process, offering parallel decoding and bidirectional context.
Used for text generation and enhancing reasoning capabilities.
Self-Distillation
A method of training using signals generated by the model itself, avoiding reliance on external labels.
Used to enhance the model's reasoning capabilities.
Privileged Information
Additional information provided to the model during training to improve performance.
dOPSD extracts privileged information from the model's own decoding trajectory.
Denoising Trajectory
The process of progressively removing noise during text generation.
Used to extract privileged information and perform self-distillation.
Exposure Bias
A problem where model performance degrades due to inconsistencies between training and inference phases.
Traditional supervised fine-tuning methods are prone to exposure bias.
Open Questions Unanswered questions from this research
- 1 How to effectively extract privileged information in more complex tasks remains to be studied.
- 2 Reducing computational overhead to enhance the method's practicality is a future research direction.
Applications
Immediate Applications
Mathematical Reasoning Tasks
dOPSD can be used to improve model performance in mathematical reasoning tasks, especially in scenarios requiring high reasoning capabilities.
Long-term Vision
General Language Models
dOPSD has the potential to become a core technology for general language models, driving applications in more fields.
Abstract
Diffusion large language models (dLLMs) generate text by iteratively denoising a masked sequence, offering a parallel alternative to autoregressive models, but eliciting strong reasoning through post-training remains difficult: supervised fine-tuning is off-policy and suffers from exposure bias, while reinforcement learning gives only sparse, sequence-level rewards and is hard to apply without tractable sequence likelihoods. On-policy self-distillation (OPSD) offers a promising alternative, using one model as both student and teacher to provide dense, token-level, on-policy supervision, but its effectiveness hinges on giving the teacher privileged information (PI) - typically an instance-specific ground-truth reference unavailable at inference - so the student ends up distilling a weak PI-free consensus policy that yields little improvement on dLLM reasoning. We introduce dOPSD, which instead derives the teacher's privilege directly from the student's own denoising trajectory, evaluating masked positions using later, more-decoded steps of that same trajectory rather than an external label, so the teacher's advantage emerges from the model's own decoding process; on Dream and LLaDA, dOPSD improves both in-domain math reasoning and out-of-domain code generation, outperforming supervised and on-policy baselines.