Matching Supervision to the Student's Learning Capacity: A Unified Framework for On-Policy Self-Distillation
Unified on-policy self-distillation (USD) dynamically adjusts token weights and privileged information via a single dual variable, boosting reasoning performance.
Key Findings
Methodology
This paper models the two key variables in OPSD—token selection and privileged information (PI) control—as a joint optimization problem. By introducing a capacity constraint and a single dual variable λ, the authors develop USD, an online primal-dual algorithm that adaptively adjusts λ and PI strength β during training. The core mechanism is that λ governs both token thresholding and PI adjustment, ensuring supervision aligns with the student's current learning capacity. Extensive experiments on Qwen3 models (1.7B, 4B, 8B) demonstrate that USD outperforms vanilla OPSD and related baselines, with an average performance gain of +2.3% on reasoning benchmarks.
Key Results
- Across three model scales, USD improves average accuracy from 56.4% (vanilla OPSD) to 58.7%, with significant gains on tasks like AIME and HMMT. For instance, in Qwen3-8B, HMMT25 accuracy rises from 46.9% to 49.7%. Ablation studies confirm that joint tuning of token weights and PI strength yields the best results, outperforming single-variable approaches.
- The algorithm's implementation is simple, requiring only real-time updates of λ and β without extra models or complex scheduling. Results indicate that the capacity-aware adjustment effectively balances supervision quality and difficulty, leading to more robust reasoning performance.
- The experimental setup includes training for 300 steps, hyperparameters such as ε=0.3, τ=0.1, and learning rates ηλ=0.1, ηβ=0.03. Evaluation uses 12 samples per problem, with the best checkpoint reported. The method consistently surpasses baselines like TIP and PAINT across all scales.
Significance
This work addresses the longstanding challenge of balancing supervision strength and relevance in self-distillation, especially for large models. By formalizing the coupling between token-level divergence and privileged information, the proposed framework enables models to adapt supervision dynamically, improving reasoning and generalization. The theoretical foundation and practical algorithm open new avenues for capacity-aware training, with broad implications for AI development in complex reasoning tasks. It bridges the gap between static supervision strategies and the need for adaptive, capacity-matched learning, fostering more efficient and capable models.
Technical Contribution
The paper introduces a capacity-constrained joint optimization framework for OPSD, leveraging a single dual variable λ to coordinate token selection and privileged information control. The approach employs an online primal-dual algorithm that updates λ and β in real time, guided by the residual of the capacity constraint. Theoretically, it proves that λ simultaneously sets token thresholds and PI adjustment directions, providing a rigorous basis for capacity matching. The method avoids non-differentiable operations via surrogate updates, ensuring practical deployability. Empirical results validate the effectiveness of this approach across multiple model scales and reasoning benchmarks, outperforming existing methods.
Novelty
This is the first work to unify token-level supervision and privileged information control under a single capacity-aware optimization framework. The key innovation lies in using a single dual variable to coordinate both axes, enabling dynamic, capacity-matched supervision. Unlike prior methods that optimize these variables independently, this approach ensures their joint evolution, leading to superior performance. The theoretical insights into the coupling mechanism and the practical primal-dual algorithm represent significant advances in self-distillation research, opening new directions for adaptive supervision strategies.
Limitations
- The method relies on accurate estimation of the student's current capacity, which may be challenging in highly dynamic or uncertain environments, potentially affecting the adjustment quality.
- In extremely large or complex tasks, the fixed capacity budget might be insufficient or overly conservative, requiring further tuning or adaptive schemes.
- Additional computational overhead from real-time λ and β updates, although minimal, could impact training efficiency in large-scale deployment scenarios.
Future Work
Future research could focus on developing more robust capacity estimation techniques, possibly integrating meta-learning or self-assessment modules. Extending the framework to multi-task and multi-modal settings would test its generality. Moreover, exploring reinforcement learning-based policies for adaptive supervision resource allocation could further enhance model autonomy. Ultimately, integrating these ideas could lead to models that autonomously optimize their learning strategies, pushing towards more intelligent and scalable AI systems.
AI Executive Summary
Large language models (LLMs) have demonstrated remarkable capabilities in natural language understanding and reasoning. Post-training strategies like knowledge self-distillation have become essential for further performance gains. Traditional self-distillation methods often treat supervision variables—such as token selection and privileged information (PI)—independently, which can limit their effectiveness. This paper introduces a unified framework that models the coupling between these two variables through a capacity constraint, ensuring supervision aligns with the model's current learning capacity.
The core innovation is the development of USD, an online primal-dual algorithm that dynamically adjusts token weights and PI strength during training. By introducing a single dual variable λ, the method simultaneously governs both supervision aspects, balancing divergence maximization and difficulty limitation. Theoretically, it is shown that λ sets the token selection threshold and PI adjustment direction, maintaining supervision within the student's capacity. Empirical results on Qwen3 models across multiple scales demonstrate that USD consistently outperforms vanilla OPSD and related baselines, with an average accuracy increase of +2.3%. The approach effectively mitigates issues of over- or under-supervision, leading to more robust reasoning abilities.
This work advances the understanding of capacity-aware supervision in self-distillation, providing both theoretical insights and practical algorithms. Its simplicity and effectiveness suggest broad applicability in training large, capable models. Future directions include refining capacity estimation, extending to multi-task learning, and integrating reinforcement learning for adaptive supervision. Overall, this research marks a significant step toward more intelligent, self-regulating AI systems capable of continuous self-improvement in complex reasoning tasks.
Deep Analysis
Background
近年来，深度学习特别是大型语言模型（LLMs）在自然语言理解、推理和生成方面取得了巨大突破。知识蒸馏技术，尤其是自蒸馏（Self-Distillation），通过模型内部privileged信息提升性能，成为后训练的重要手段。早期的蒸馏方法如Hinton提出的知识蒸馏（2015）解决了模型压缩问题，后续如OpenAI的Supervised Fine-Tuning（SFT）和强化学习（RL）方法不断推动模型推理能力提升。OPSD作为一种无需外部教师的自蒸馏策略，利用模型自身生成的rollouts和privileged信息，显著改善了数学推理和编码任务表现。然而，现有方法在令牌选择和privileged信息控制上多为孤立优化，导致 supervision 质量与学生能力不匹配，限制了性能潜力。随着模型规模的不断扩大，如何动态调节 supervision 以适应学生的学习状态，成为亟待解决的问题。
Core Problem
核心问题在于令牌选择（w）和privileged信息（β）控制的耦合性。传统方法在固定或单一调节下，难以实现 supervision 的最优匹配，导致模型在复杂推理任务中表现受限。具体表现为：一方面，过度privileged信息会引入答案泄露或超出学生能力，另一方面，不合理的令牌选择会忽略关键内容或引入噪声。这种不匹配限制了模型的学习效率和推理能力，亟需一种机制实现两者的协同调节，达到能力匹配。
Innovation
本文的创新在于提出容量匹配的联合优化框架，将令牌选择（w）与privileged信息（β）作为两个耦合变量，通过引入单一的拉格朗日乘子λ，实现两者的动态调节。该框架基于容量预算，确保 supervision 既不过度也不不足，最大化模型潜在推理能力。主要创新点包括：1）将两个调节变量统一建模，解决以往孤立优化的局限；2）设计轻量级的在线 primal-dual 算法USD，实时调节λ和β，无需额外模型或复杂调度；3）理论证明λ同时设定令牌阈值和privileged信息调整方向，提供严格的数学基础。
Methodology
- �� 定义令牌权重w和privileged信息强度β作为调节变量。
- �� 构建目标函数，最大化带权 divergence，同时限制学习难度（load）不超过容量ε。
- �� 引入拉格朗日乘子λ，将容量限制转化为惩罚项，形成拉格朗日函数。
- �� 通过在线 primal-dual 方法，实时调整λ和β：
- λ通过梯度上升，调节令牌选择阈值。
- β通过残差调整，确保容量匹配。
- �� 令牌权重通过sigmoid函数实现软阈值，避免非光滑优化。
- �� 最终算法在每个训练批次中，实时更新λ和β，保证 supervision 贴合学生能力。
Experiments
采用Qwen3系列（1.7B、4B、8B）模型，在数学推理任务（AIME24/25、HMMT25）上验证。对比SFT、GRPO、vanilla OPSD、TIP、PAINT等基线，训练300步，调节参数包括：容量ε=0.3，温度τ=0.1，学习率ηλ=0.1，ηβ=0.03。评估指标为Avg@12，结果显示USD在所有模型规模上均优于对比方法，平均提升达+2.3%。 Ablation研究验证了双变量调节的协同作用，单独调节效果有限。
Results
在三种模型规模中，USD平均性能从56.4%提升到58.7%，在复杂推理任务中表现尤为优越。具体数据如Qwen3-8B模型在HMMT25任务中由46.9%提升至49.7%。实验还显示，单独调节令牌或privileged信息效果有限，联合调节显著提升模型推理能力。算法实现简单，训练过程中无需额外模型，只通过动态调节λ和β实现容量匹配，具有良好的实用性。
Applications
该方法适用于大规模语言模型的推理增强，尤其在数学、逻辑推理等任务中。可用于学术研究、工业智能问答、自动推理系统等场景，帮助模型更好理解复杂问题，提升推理准确率。实现条件为模型具备privileged信息和动态调节能力，训练过程中无需额外硬件资源。
Limitations & Outlook
算法对容量预算参数敏感，实际应用中需调优。模型容量估计误差可能影响调节效果，尤其在极端任务或超大模型中表现不佳。此外，调节机制引入训练复杂度，可能影响大规模部署效率。未来需优化容量估计和调节策略，增强鲁棒性。
Plain Language Accessible to non-experts
想象你在准备一场重要的演讲，你希望内容既丰富又不过载。传统方法就像是把所有资料都放在演讲稿里，不管内容是否适合听众理解。而本文提出的方法像是一个智能助手，会根据听众的反应，动态调整内容的深度和重点。它会在你讲得太快或太难时，减缓节奏，简化信息；而在听众准备充分时，提供更深入的细节。这样，演讲既不会太难，也不会太浅，刚好符合听众的接受能力。这个“助手”就是通过调节privileged信息和内容选择，确保每次“讲述”都最有效率，帮助你更好地传达信息。
ELI14 Explained like you're 14
想象你在学校里参加一个问答比赛，你的老师会根据你的水平，给你一些提示和线索。有时候，老师会给你很多线索，帮助你找到答案，但如果线索太多，你可能会依赖得太厉害，反而学不到自己思考的能力。有时候，老师只给你一点点提示，促使你自己努力思考。这个研究就像是让老师在提示和线索之间找到平衡点，确保你既能学到东西，又不会被提示搞得太依赖。它用一种聪明的方法，动态调整提示的多少，让你逐步变得更聪明、更有能力。就像在游戏中逐渐解锁新技能一样，模型也在不断调整 supervision 的强度，帮助它变得更聪明、更会推理。
Glossary
Self-Distillation (自蒸馏)
一种模型通过内部生成的privileged信息自我提升性能的方法，避免外部教师依赖。技术上是模型用自己生成的预测作为 supervision。
论文中利用模型内部的privileged信息，通过自蒸馏提升推理能力。
Privileged Information (特权信息)
在训练时可用，但在推理时不可用的额外上下文信息，用于指导模型学习。
调节privileged信息的强度，避免泄露答案或超出学生能力。
Lagrangian (拉格朗日函数)
一种优化工具，将约束条件融入目标函数，通过引入乘子实现多目标平衡。
本文用拉格朗日乘子λ调节令牌选择与privileged信息控制的关系。
Capacity Budget (容量预算)
模型在训练中可接受的最大学习难度总量，用于调节 supervision 资源分配。
确保 supervision 不超出学生的学习能力，动态调节λ和β。
Open Questions Unanswered questions from this research
- 1 如何在极端任务中准确估计学生的学习容量，确保调节机制的鲁棒性仍需研究。
- 2 在多任务、多模态场景中，容量匹配策略的扩展与优化尚未充分探索。
- 3 未来需结合元学习或自适应机制，提升容量估计的准确性和调节的智能化水平。
Applications
Immediate Applications
大规模推理模型训练
在数学、逻辑推理等任务中，利用USD动态调节 supervision，提升模型推理能力，适用于学术研究和工业应用。
智能问答系统优化
结合privileged信息调节，增强模型理解复杂问题的能力，提升问答准确率，适用于智能客服和教育平台。
Long-term Vision
自主学习与能力提升
未来模型能自主调节 supervision 资源，实现持续学习和能力自我提升，推动人工智能向更智能化方向发展。
Abstract
On-policy self-distillation (OPSD) improves the reasoning abilities of LLMs by internalizing privileged context into model parameters through self-distillation. Two recent research lines promote vanilla OPSD by choosing which tokens to learn from and by controlling how much privileged information the teacher receives, respectively. However, we show that each line optimizes one variable while holding the other fixed, which leads to a suboptimal solution. We argue that the two variables are coupled through the student's learning capacity: the privileged information sets the per-token divergence the teacher prescribes, while token weighting selects which of these the student must absorb. We formalize the two lines of work into a unified optimization framework, which maximizes the aggregate teacher--student divergence, subject to a budget on the aggregate learning difficulty the student can absorb. Under this modelling, we propose Unified On-Policy Self-Distillation (USD), a lightweight online algorithm to solve the Lagrangian. USD reveals that a single dual variable governs both decisions: at one price for learning difficulty, it simultaneously sets the token-selection threshold and the direction of privileged-information adjustment, keeping supervision matched to the student's evolving capacity. Through extensive experiments, USD consistently demonstrates superior performance over OPSD and token- and PI-side baselines across various model scales on various reasoning benchmarks. Code is available at https://github.com/lauvlalala/USD.