# Knowledge Distillation with Training Wheels

arXiv: https://arxiv.org/abs/2502.17717 · HTML: https://mlsanity.com/paper/2502.17717/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Knowledge distillation for LLMs, one of 12 directions in Knowledge distillation.
- Builds on 12 of its 20 references that are in the map.
- First posted on arXiv 24 Feb 2025.

## Placed in

- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

## Builds on (most cited first, 12 of 12)

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Scaling Instruction-Finetuned Language Models](https://mlsanity.com/paper/2210.11416/) (2022)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization](https://mlsanity.com/paper/1808.08745/) (2018)
- [Accelerating Large Language Model Decoding with Speculative Sampling](https://mlsanity.com/paper/2302.01318/) (2023)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [Sequence Level Training with Recurrent Neural Networks](https://mlsanity.com/paper/1511.06732/) (2015)
- [Reward Constrained Policy Optimization](https://mlsanity.com/paper/1805.11074/) (2018)
- [Bridging the Gap Between Value and Policy Based Reinforcement Learning](https://mlsanity.com/paper/1702.08892/) (2017)
- [DistillSpec: Improving Speculative Decoding via Knowledge Distillation](https://mlsanity.com/paper/2310.08461/) (2023)
- [Online Speculative Decoding](https://mlsanity.com/paper/2310.07177/) (2023)
- [How (not) to Train your Generative Model: Scheduled Sampling, Likelihood, Adversary?](https://mlsanity.com/paper/1511.05101/) (2015)

## Built on by (most cited first, 4 of 4)

- [Speculative Decoding and Beyond: An In-Depth Survey of Techniques](https://mlsanity.com/paper/2502.19732/) (2025)
- [Beyond GRPO and On-Policy Distillation: An Empirical Sparse-to-Dense Reward Principle for Language-Model Post-Training](https://mlsanity.com/paper/2605.12483/) (2026)
- [Revisiting Test-Time Scaling: A Survey and a Diversity-Aware Method for Efficient Reasoning](https://mlsanity.com/paper/2506.04611/) (2025)
- [OptiAgent: End-to-End Optimization Modeling via Multi-Agent Iterative Refinement](https://mlsanity.com/paper/2607.05346/) (2026)

## 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.