Learning Beyond Full Imitation: Task-Preserving Knowledge Distillation

Abstract: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.

Machine Learning (cs.LG)

arXiv:2609.39338 [cs.LG]

(or arXiv:2609.39338v1 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2609.39338 Focus to learn more arXiv-issued DOI via DataCite (pending registration)

Supplementary_Material/README.md

Supplementary_Material/configs/protocols.json

Supplementary_Material/data/accuracy/comparisons.json

Supplementary_Material/data/accuracy/seed_scores.json

Supplementary_Material/data/accuracy/table4.json

Supplementary_Material/data/accuracy/table5.json

Supplementary_Material/data/clinc150/accuracy.json

Supplementary_Material/data/clinc150/seed42/labels.npy

Supplementary_Material/data/clinc150/seed42/logits.npy

Supplementary_Material/data/clinc150/seed42/row_id.npy

Supplementary_Material/data/clinc150/seed43/labels.npy

Supplementary_Material/data/clinc150/seed43/logits.npy

Supplementary_Material/data/clinc150/seed43/row_id.npy

Supplementary_Material/data/clinc150/seed44/labels.npy

Supplementary_Material/data/clinc150/seed44/logits.npy

Supplementary_Material/data/clinc150/seed44/row_id.npy

Supplementary_Material/data/clinc150/settings.json

Supplementary_Material/data/clinc150/teacher_selection.json

Supplementary_Material/data/continuation128/conditional_kl_TPKD_end.npy

Supplementary_Material/data/continuation128/conditional_kl_start.npy

Supplementary_Material/data/continuation128/gain_over_CE_at_end.npy

Supplementary_Material/data/continuation128/gain_over_compensation_at_end.npy

Supplementary_Material/data/continuation128/results.json

Supplementary_Material/data/continuation128/sample_id.npy

Supplementary_Material/data/table1/blocked.npy

Supplementary_Material/data/table1/conditional_kl.npy

Supplementary_Material/data/table1/dinf.npy

Supplementary_Material/data/table1/results.json

Supplementary_Material/data/table1/sample_id.npy

Supplementary_Material/data/table2/batch.npy

Supplementary_Material/data/table2/compensation_minus_CE_conditional_kl.npy

Supplementary_Material/data/table2/conditional_kl_change.npy

Supplementary_Material/data/table2/cosine.npy

Supplementary_Material/data/table2/gain_over_CE.npy

Supplementary_Material/data/table2/joint_descent_check_passed.npy

Supplementary_Material/data/table2/joint_error_change.npy

Supplementary_Material/data/table2/minimum_TPKD_margin_change_vs_CE.npy

Supplementary_Material/data/table2/minimum_unprojected_margin_change_vs_CE.npy

Supplementary_Material/data/table2/numerical_precision.json

Supplementary_Material/data/table2/position.npy

Supplementary_Material/data/table2/results.json

Supplementary_Material/data/table2/retained_fraction.npy

Supplementary_Material/data/table2/sample_id.npy

Supplementary_Material/data/table3/batches.json

Supplementary_Material/data/table3/results.json

Supplementary_Material/requirements.txt

Supplementary_Material/scripts/check_clinc150.py

Supplementary_Material/scripts/check_results.py

Supplementary_Material/scripts/check_theory.py

Supplementary_Material/scripts/generate_tables.py

Supplementary_Material/scripts/results.py

Supplementary_Material/templates/ablation.tex

Supplementary_Material/templates/application.tex

Supplementary_Material/templates/mechanisms.tex

(49 additional files not shown) You must enabled JavaScript to view entire file list.

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.