Knowledge Distillation with Training Wheels
Not on the map. Why?
Abstract (arXiv)
Knowledge distillation is used, in generative language modeling, to train a smaller student model using the help of a larger teacher model, resulting in improved capabilities for the student model. In this paper, we formulate a more general framework for knowledge distillation where the student learns from the teacher during training, and also learns to ask for the teacher's help at test-time following rules specifying test-time restrictions. Towards this, we first formulate knowledge distillation as an entropy-regularized value optimization problem. Adopting Path Consistency Learning to solve this, leads to a new knowledge distillation algorithm using on-policy and off-policy demonstrations. We extend this using constrained reinforcement learning to a framework that incorporates the use of the teacher model as a test-time reference, within constraints. In this situation, akin to a human learner, the model needs to learn not only the learning material, but also the relative difficulty of different sections to prioritize for seeking teacher help. We examine the efficacy of our method through experiments in translation and summarization tasks, observing trends in accuracy and teacher use, noting that our approach unlocks operating points not available to the popular Speculative Decoding approach.
Key facts. Cited by 4 papers in the map (data through 7 Oct 2026). Builds on 12 of its 20 references that are in the map. First posted on arXiv 24 Feb 2025.
Builds on
12 of 20 references.
References · 12
- 2015Distilling the Knowledge in a Neural Networkcited by 10,076
- 2022Scaling Instruction-Finetuned Language Modelscited by 2,997
- 2022Fast Inference from Transformers via Speculative Decodingcited by 1,615
- 2018Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarizationcited by 1,308
- 2023Accelerating Large Language Model Decoding with Speculative Samplingcited by 975
- 2023On-Policy Distillation of Language Models: Learning from Self-Generated Mistakescited by 840
- 2015Sequence Level Training with Recurrent Neural Networkscited by 568
- 2018Reward Constrained Policy Optimizationcited by 397
- 2017Bridging the Gap Between Value and Policy Based Reinforcement Learningcited by 178
- 2023DistillSpec: Improving Speculative Decoding via Knowledge Distillationcited by 160
- 2023Online Speculative Decodingcited by 109
- 2015How (not) to Train your Generative Model: Scheduled Sampling, Likelihood, Adversary?cited by 76
Built on by
Citing papers per year: 2025: 2; 2026: 2.
4 papers.
- 2025Speculative Decoding and Beyond: An In-Depth Survey of Techniquescited by 22
- 2026Beyond GRPO and On-Policy Distillation: An Empirical Sparse-to-Dense Reward Principle for Language-Model Post-Trainingcited by 9
- 2025Revisiting Test-Time Scaling: A Survey and a Diversity-Aware Method for Efficient Reasoningcited by 8
- 2026OptiAgent: End-to-End Optimization Modeling via Multi-Agent Iterative Refinementcited by 0
- 2026OptiAgent: End-to-End Optimization Modeling via Multi-Agent Iterative Refinement6 Jul 2026
- 2026Beyond GRPO and On-Policy Distillation: An Empirical Sparse-to-Dense Reward Principle for Language-Model Post-Training12 May 2026
- 2025Revisiting Test-Time Scaling: A Survey and a Diversity-Aware Method for Efficient Reasoning5 Jun 2025
- 2025Speculative Decoding and Beyond: An In-Depth Survey of Techniques27 Feb 2025
Similar, not cited
Close in wording, no citation either way
- 2026A Survey of On-Policy Distillation for Large Language Models
- 2025TAID: Temporally Adaptive Interpolated Distillation for Efficient Knowledge Transfer in Language Models
- 2025UNDO: Understanding Distillation as Optimization
- 2026X-KD: General Experiential Knowledge Distillation for Large Language Models
- 2026Knowledge Distillation for Large Language Models
- 2026Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models
- 2026Reinforcement-aware Knowledge Distillation for LLM Reasoning
- 2026KDFP: A first-principles approach to knowledge distillation in large language models