# Learning Beyond Full Imitation: Task-Preserving Knowledge Distillation

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in On-policy and self-distillation for LLMs, one of 12 directions in Knowledge distillation.
- Builds on 19 of its 24 references that are in the map.
- First posted on arXiv 30 Sep 2026.

## Placed in

- Knowledge distillation / [On-policy and self-distillation for LLMs](https://mlsanity.com/distill/on-policy-and-self-distillation-for-llms/)

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Improved Knowledge Distillation via Teacher Assistant](https://mlsanity.com/paper/1902.03393/) (2019)
- [Decoupled Knowledge Distillation](https://mlsanity.com/paper/2203.08679/) (2022)
- [An Evaluation Dataset for Intent Classification and Out-of-Scope Prediction](https://mlsanity.com/paper/1909.02027/) (2019)
- [On the Efficacy of Knowledge Distillation](https://mlsanity.com/paper/1910.01348/) (2019)
- [Rethinking Spatial Dimensions of Vision Transformers](https://mlsanity.com/paper/2103.16302/) (2021)
- [Does Knowledge Distillation Really Work?](https://mlsanity.com/paper/2106.05945/) (2021)
- [Self-Distillation Amplifies Regularization in Hilbert Space](https://mlsanity.com/paper/2002.05715/) (2020)
- [VOLO: Vision Outlooker for Visual Recognition](https://mlsanity.com/paper/2106.13112/) (2021)
- [Towards Understanding Knowledge Distillation](https://mlsanity.com/paper/2105.13093/) (2021)
- [Supervision Complexity and its Role in Knowledge Distillation](https://mlsanity.com/paper/2301.12245/) (2023)
- [Fantastic Gains and Where to Find Them: On the Existence and Prospect of General Knowledge Transfer between Any Pretrained Model](https://mlsanity.com/paper/2310.17653/) (2023)
- [On student-teacher deviations in distillation: does it pay to disobey?](https://mlsanity.com/paper/2301.12923/) (2023)
- [Selective Classification Under Distribution Shifts](https://mlsanity.com/paper/2405.05160/) (2024)
- [Dual-Head Knowledge Distillation: Enhancing Logits Utilization with an Auxiliary Head](https://mlsanity.com/paper/2411.08937/) (2024)
- [DeepKD: A Deeply Decoupled and Denoised Knowledge Distillation Trainer](https://mlsanity.com/paper/2505.15133/) (2025)
- [Mitigating Negative Flips via Margin Preserving Training](https://mlsanity.com/paper/2511.08322/) (2025)

## Abstract (arXiv)

Knowledge distillation transfers knowledge by encouraging a student to match a teacher's predicted class probabilities. These probabilities express not only confidence in the correct class, but also relations among incorrect alternatives. Yet closer imitation does not necessarily yield a better student. A student may already distinguish the correct class more sharply than its teacher, so further imitation can require giving back discrimination it has acquired. Our main result is an exact separation between full imitation and conditional learning. When the correct class's score advantage over each alternative must be preserved, full teacher-to-student KL minimization is blocked exactly when the student assigns no more probability than the teacher to every incorrect class. Crucially, the teacher's relative probabilities among incorrect classes remain fully learnable. We characterize the exact price of this transfer: a minimum increase in correct-class log-odds that compensates for the largest conditional-probability mismatch. Label fitting and conditional matching can therefore be completed even as full teacher KL diverges. This separation motivates task-preserving knowledge distillation (TPKD), which keeps the label gradient intact and minimally corrects the conditional gradient so that its output update preserves the label step's gains against every incorrect alternative. The corrected conditional direction retains more than half of the original first-order conditional descent at the same step size, with a tight bound. For a fixed positive conditional target and sufficiently small constant output steps, label and conditional errors vanish together. Experiments trace this learning from exact head updates to ordinary network training. TPKD reaches 88.05% accuracy on CIFAR-100 and 93.81% on CLINC150, improving over standard distillation by 0.47 and 0.35 percentage points across three seeds.