PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Beyond Full Imitation: Task-Preserving Knowledge Distillation

Qianfeng Yuan, Wenbing Tao

arXiv:2609.39338v1cs.LG

TL;DR

Knowledge distillation transfers a teacher’s predictive distribution, but closer imitation can require a student to give back acquired task discrimination. The paper characterizes when full imitation is blocked and how conditional knowledge can still be transferred, then proposes TPKD to preserve label-step margin gains. TPKD improves over standard KD on CIFAR-100 and CLINC150 across three seeds.

  • Problem

    The paper asks what a teacher can still teach when further imitation conflicts with a student’s task progress, a question less explored theoretically than empirical KD objectives and procedures.

    Source: §1 · p. 1

  • Method

    TPKD keeps the label gradient intact and minimally corrects the conditional direction so updates preserve every margin gain of the label step.

    Source: §1 · p. 2, §5.1 · p. 5

  • Results

    88.05% accuracy on CIFAR-100 and 93.81% on CLINC150 exceed standard KD by 0.47 and 0.35 percentage points, respectively, across three seeds.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Full imitation can be blocked while the teacher’s complete conditional distribution over incorrect classes remains learnable without decreasing any correct-class margin.

    Source: §1 · p. 2

    The analysis preserves each correct-versus-incorrect margin, a stricter requirement than merely keeping the answer correct.

    Source: §1 · p. 2, §3 · p. 3

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper asks what a teacher can still teach when further imitation would undo task discrimination the student has gained. It separates full imitation from conditional learning, gives a confidence-compensation cost, proposes TPKD, and reports gains over standard KD on vision and text benchmarks.

Motivation

  • Closer imitation need not produce a better student: label learning can make the student more discriminative than its teacher, so matching the full prediction may give back that discrimination.The teacher can nevertheless encode a different relative ordering among incorrect alternatives.Source: §1 · pp. 1–2

Task preservation

  • The paper preserves each correct-versus-incorrect score gap, rather than only the correct answer, because smaller margins can change whether a prediction requires human confirmation.Different confusions may require different acceptance thresholds.Source: §1 · p. 2

Core result

  • Full imitation is blocked exactly when the student assigns no more probability than the teacher to every incorrect class, although the teacher’s full conditional distribution remains learnable without reducing any correct-class margin.The minimum transfer cost is an increase in correct-class log-odds determined by the largest teacher-to-student conditional-probability ratio.Source: §1 · p. 2

Proposed method

  • TPKD keeps the label gradient intact and minimally corrects the conditional direction to preserve every margin gain of the label step.The correction retains over half of first-order conditional descent with a tight bound; for fixed positive targets and sufficiently small constant output steps, label and conditional errors vanish together.Source: §1 · p. 2

Experimental results

  • 88.05% accuracy on CIFAR-100 and 93.81% on CLINC150 exceed standard KD by 0.47 and 0.35 percentage points, respectively, across three seeds.The experiments also report that removing the CE gradient or replacing conditional learning with protected full imitation reduces accuracy in both domains.Source: §1 · p. 2

2 RELATED WORK

Prior KD work has developed objectives and training procedures, while theoretical analyses have examined several other sources of student–teacher differences. This paper instead studies which teacher signals remain attainable while preserving the student’s acquired class comparisons.

Distillation targets

  • Classical KD matches softened teacher probabilities, while DKD separates target-class and non-target supervision; this paper asks which signals remain attainable under task preservation.The cited methods identify useful teacher signals, whereas the present analysis focuses on preserving existing discrimination.Source: §2 · p. 2

Theoretical perspectives

  • Existing KD theory studies capacity mismatch, teacher quality, statistical supervision, self-distillation regularization, data geometry, and optimization-induced deviations.The paper presents its question as complementary to these theoretical perspectives.Source: §2 · pp. 2–3

Preservation and gradient coordination

  • Related gradient-coordination methods project or balance gradients, filter non-target knowledge, or reduce prediction regressions; TPKD retains the complete label gradient and minimally corrects each conditional increment.The paper first characterizes blocking and transfer cost in the task-preserving output region, then gives tight bounds on retained learning.Source: §2 · p. 3

3 WHAT TO LEARN AND WHAT TO PRESERVE

The section separates label fitting, full imitation, and conditional learning, then defines task preservation through the student’s margins against individual incorrect classes. It explains why confidence and conditional structure can change independently and how margin preservation constrains updates.

Learning targets

  • Label fitting minimizes cross-entropy, full imitation matches the teacher’s confidence and conditional structure, and conditional learning matches only the teacher’s distribution over incorrect classes.The teacher loss decomposes into a binary confidence term and a conditional KL term.Source: §3 · p. 3

Preserving learned comparisons

  • The student’s correct-versus-incorrect margin is Mj(p) = log(py/pj) = zy − zj, and task preservation requires protecting each such comparison.This is stronger than merely keeping the correct answer unchanged.Source: §3 · p. 3
  • Preserving every margin is equivalent to not increasing any weighted loss Lω; unit weights yield cross-entropy, while a single unit weight isolates one competitor.Cross-entropy alone can improve even while one individual comparison worsens.Source: §3 · p. 3

How learning and protection interact

  • At fixed correct-class confidence, increasing a competitor’s conditional probability shrinks its margin, so raising confidence can compensate while conditional learning leaves confidence free.Full imitation also seeks the teacher’s confidence.Source: §3 · p. 3

4 WHEN IMITATION STOPS, WHAT CAN THE TEACHER STILL TEACH?

The section proves that full imitation can become blocked under per-competitor margin preservation even while the teacher’s conditional structure remains learnable. It derives the minimum confidence compensation for exact transfer and illustrates the separation between conditional learning and full KL.

4.1 WHY FULL IMITATION CAN BECOME BLOCKED

  • Full imitation is uniquely blocked within the task-preserving region exactly when the student assigns no more probability than the teacher to every incorrect class.In that state, any permitted change strictly increases full imitation error, even though the conditional distribution may still differ from the teacher’s.Source: §4.1 · pp. 3–4

4.2 LEARNING THE REMAINING STRUCTURE, AND ITS PRICE

  • The teacher’s conditional probabilities among incorrect classes remain fully transferable without decreasing any correct-class margin, even when full imitation is blocked.Matching a different conditional target may require raising correct-class confidence to protect margins.Source: §4.2 · pp. 4–5
  • Exact transfer is task-preserving when the increase in correct-class log-odds is at least D∞(q∥r), the log of the largest teacher-to-student conditional-probability ratio.The minimum confidence is amin = ρR∞/(1 − ρ + ρR∞).Source: §4.2 · p. 4
  • A continuous path reaches amin while every margin remains nondecreasing, and conditional KL falls strictly to zero when the conditional distributions differ.The path moves the student’s conditional distribution toward the teacher while compensating confidence just enough at each point.Source: §4.2 · p. 4
  • In the Figure 1 example, amin = 35/37 ≈ 0.946; the class-3 margin stays fixed while the class-2 margin grows.The student initially favors class 2 over class 3, whereas the teacher favors class 3.Source: §4.2 · p. 4, Figure 1 · p. 4

5 TASK-PRESERVING KNOWLEDGE DISTILLATION

TPKD keeps the label signal intact while correcting only the teacher-conditional direction, preserving label-step margin gains while retaining conditional learning. Under stated conditions, label fitting and conditional matching converge together, and network-parameter updates can achieve conditional and joint progress.

5.1 KEEP LABEL LEARNING INTACT AND CORRECT THE TEACHER SIGNAL

  • TPKD minimally corrects the combined direction so every label-step margin gain and corresponding label-loss ordering are preserved.The corrected direction is the nearest safe direction to h + u/2; the ordering also holds for weighted label losses.Source: §5.1 · p. 5

5.2 PROTECTION RETAINS REAL CONDITIONAL LEARNING

  • Clipping caps the teacher’s strongest requests to raise competing classes and converts the clipped mass into correct-class compensation.The frozen clipping yields a local loss matching the normalized clipped teacher while increasing correct-class log-odds; the most demanding competitors keep the CE margin and others gain more.Source: §5.2 · pp. 5–6
  • The corrected conditional direction retains more than half the original first-order conditional loss reduction at the same step size, with tight bounds.Its conditional component remains nonzero and aligned with the teacher signal whenever conditional error remains.Source: §5.2 · p. 6

5.3 BOTH KINDS OF LEARNING CAN BE COMPLETED TOGETHER

  • For nonuniform q, CE output-gradient flow and fixed steps 0 < η ≤ 1/20 both fit the label but drive r toward uniform, whereas TPKD retains the teacher’s conditional structure.Source: §5.3 · p. 7

5.4 TURNING OUTPUT PROGRESS INTO PARAMETER LEARNING

  • A full-row-rank frozen batch permits a unique minimum-Frobenius-norm head displacement that realizes the output update exactly for every example.It inherits the step’s margin and loss orderings.Source: §5.4 · p. 7
  • Under χc < κc, TPKD yields a nonzero conditional parameter increment beyond compensation alone and strictly lowers batch-average conditional KL.This holds at differentiable parameter states for sufficiently small equal positive parameter steps when some ui is nonzero.Source: §5.4 · p. 7
  • Under χΦ < κΦ, TPKD strictly decreases the batch mean of Φ for sufficiently small equal positive parameter steps.This guarantee requires both joint batch vectors to be nonzero.Source: §5.4 · p. 7

6 EXPERIMENTS

Experiments trace task-preserving conditional learning from ordinary training through exact head and network updates to full training in vision and text. TPKD improves on standard KD across both evaluated domains.

6 EXPERIMENTS

  • Blocking rises from 50.00% to 58.90% across CE states, and every blocked pair retains positive conditional error.Source: §6 · p. 7, Table 1 · p. 8
  • Across 64 batches and 4,096 example states, TPKD learns conditional knowledge beyond CE while preserving every CE-step margin; retention exceeds the approximately 50.51% bound.Source: Table 2 · p. 8, §6 · p. 8
  • All 32 native-optimizer batches pass both geometric tests, and TPKD yields lower conditional KL than compensation-only updates.Source: Table 3 · p. 8, §6 · p. 8

6.1 FULL TRAINING IN VISION AND TEXT

  • 88.05% on CIFAR-100 and 93.81% on CLINC150; TPKD improves over standard KD by 0.47 and 0.35 percentage points across three seeds.Source: Table 4 · p. 9, §6.1 · p. 9

6.2 WHICH COMPONENTS MAKE THE DIFFERENCE?

  • Removing the CE gradient costs 2.72 points in vision and 2.21 in text; safe full KL falls near CE, while unprojected TPKD retains closely matched performance.Source: Table 5 · p. 9, §6.2 · p. 9

7 CONCLUSION

The conclusion presents TPKD as learning conditional structure without surrendering label-learning progress. Appendix A.1 proves that preserving each margin preserves the corresponding weighted task-loss ordering.

7 CONCLUSION

  • TPKD retains the label gradient and adds the nearest safe conditional direction, allowing conditional knowledge to be acquired without surrendering label-learning progress.Source: §7 · p. 9

AI USE STATEMENT

  • The authors report that no generative AI tools were used in conducting the research or preparing the manuscript.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The source bundle includes accuracy summaries, table-generation scripts, and analytical checks; Appendices A–D prove the main-text results and Appendix E specifies experiments.Source: REPRODUCIBILITY STATEMENT · p. 10

A PROOFS SUPPORTING TASK PRESERVATION AND THE SEPARATION

A.1 WEIGHTED TASK LOSSES AND PER-MARGIN PRESERVATION

  • Preserving every correct-versus-incorrect margin preserves every nonnegative weighted task-loss ordering; testing a single competitor establishes the converse.Source: §A.1 · p. 12

A.2 PROOF OF THEOREM 1

The proofs establish when full-KL imitation is blocked under margin preservation, quantify the confidence cost of complete conditional transfer, and trace a path that preserves margins while matching the teacher’s conditional distribution.

A.2 PROOF OF THEOREM 1

  • Full-KL minimization is blocked exactly when the student assigns no more probability than the teacher to every incorrect class; otherwise a margin-preserving change can reduce full KL.Under the blocking condition, the student is the unique task-preserving minimizer when probabilities are strictly positive.Source: §A.2 · p. 12

A.3 PROOF OF THEOREM 2

  • Complete conditional transfer requires increasing correct-class log-odds by the largest relative conditional mismatch, D∞(q∥r), yielding the minimum compensation in Eq. (4).The mismatch is at least one, with equality exactly when the student and teacher conditional distributions match.Source: §A.3 · p. 12

A.4 THE MINIMUM-COMPENSATION PATH

  • Along the minimum-compensation path, every incorrect-class margin is nondecreasing, and the endpoint matches the teacher’s conditional distribution at the minimum required confidence.For q ≠ r, conditional KL decreases strictly before the endpoint.Source: §A.4 · p. 13
  • 0.171739 to zero conditional KL accompanies full teacher KL rising from 0.853727 to 1.099879 along the analytic path.Label CE falls from approximately 0.105361 to 0.055570.Source: §A.4 · p. 13

A.5 THE TWO COMPONENT CONSEQUENCES USED IN THE ABLATION

These proofs distinguish label and conditional learning, characterize imitation under label-only protection and finite confidence budgets, and establish the projection and compensation bounds used by the safe update.

A.5 THE TWO COMPONENT CONSEQUENCES USED IN THE ABLATION

  • At a blocked state with mismatched conditionals, the full-KL teacher increment vanishes under safe correction, while the conditional increment remains available.When conditionals already match, the conditional gradient is zero, but the ordinary label gradient remains nonzero whenever the correct-class probability is below one.Source: §A.5 · p. 14

A.6 THE IMITATION OPTIMUM WHEN ONLY CE IS PROTECTED

  • With only label cross-entropy protected, the KL-optimal prediction matches the teacher’s conditional distribution without lowering correct-class probability, but can reduce an individual competitor margin.In the example, the third-class probability rises above its initial 0.03 despite unchanged CE.Source: §A.6 · p. 14

A.7 THE FINITE-COMPENSATION FRONTIER

  • For a confidence budget B, the minimum remaining conditional error V(B) is nonincreasing and convex, and reaches zero exactly when B ≥ D∞(q∥r).At zero budget, V(0) equals KL(q∥r).Source: §A.7 · p. 14

B.1 THE PROJECTION AND TASK ORDERING

  • The safe update is the nearest direction to the uncorrected combined direction that preserves every margin gain of the label step.The projection can be computed in O(K log K) per example.Source: §B.1 · p. 15

B.2 CLIPPED TEACHER AND THE TIGHT COMPENSATION BOUNDS

  • The corrected direction preserves full task ordering, and compensation is controlled by conditional mismatch: TV(q, r)/(K −1) ≤ ℓ ≤ TV(q, r)/2.Both bounds are tight.Source: §B.1 · p. 15, §B.2 · p. 16

B.3 PROOF OF THEOREM 3

The appendix establishes tight alignment and retention bounds for safe conditional learning, then interprets the projected direction and its cost. It also compares safe conditional progress with unconstrained progress under two local budgets.

B.3 PROOF OF THEOREM 3

  • The Moreau decomposition yields a tight retention bound: a construction attains both bounds simultaneously, with ∥d∥2 / ∥u∥2 = γK.The construction is realizable as the difference of two strictly positive conditional distributions for sufficiently small a.Source: §B.3 · p. 16

B.4 THE ALIGNMENT PART OF THEOREM 3

  • The projected conditional direction retains more than half the original first-order conditional descent, with a uniform alignment bound that is tight.The bound follows from γK ∈ (1/2, 3/4] and is attained by the construction in Appendix B.3.Source: §B.3 · p. 16, §B.4 · p. 17

B.5 THE GEOMETRIC EXAMPLE IN FIGURE 3

  • In the K = 4 illustration, the projected component is orthogonal to the confidence direction and forms an approximately 11.54° angle with u.The two panels show the orthogonal components and their angle within the conditional plane.Source: §B.5 · p. 17

B.6 A LOCAL LOSS REPRESENTATION OF THE CLIPPED TEACHER

  • With reference quantities frozen at z0, the clipped direction is the gradient of a local loss whose terms learn the normalized clipped teacher and increase correct-class log-odds.The two terms are weighted by the retained and clipped probability masses.Source: §B.6 · p. 17

B.7 THE EXACT SIGNAL REMOVED AND THE SIZE OF COMPENSATION

  • The removed conditional signal and extra task-side change are bounded, and the additional confidence effect vanishes as conditional mismatch vanishes.The clipped masses quantify the removed signal; the margin change is bounded using total variation between q and r.Source: §B.7 · p. 18

B.8 OPTIMAL CONDITIONAL PROGRESS UNDER TWO LOCAL BUDGETS

  • Under a local quadratic model, the safe direction is optimal over the task-preserving cone; at equal displacement length its guaranteed-progress ratio is at least √γK, tightly.At equal step size, the corresponding ratio is at least γK.Source: §B.8 · p. 18

C.1 COMPLETE-UPDATE ALIGNMENT IN THEOREM 4

The appendix proves alignment and convergence properties for the complete update, compares its conditional limit with cross-entropy, and extends joint learning to positive conditional weights.

C.1 COMPLETE-UPDATE ALIGNMENT IN THEOREM 4

  • The complete update has a tight alignment guarantee, and a construction shows that no larger uniform angle constant holds for strictly positive probabilities.The equality construction makes the proof bound’s gap approach zero.Source: §C.1 · p. 19

C.2 DESCENT, CONVERGENCE AND THE RESIDUAL RATE

  • For the specified output steps, the joint potential descends and label and conditional errors vanish together; joint and component gradient residuals have vanishing time-averaged squared norms.The limiting state has correct-class probability tending to one and student conditional probabilities tending to q.Source: §C.2 · pp. 19–20

C.3 THE CE CONDITIONAL LIMIT STATED IN SECTION 5.3

  • With finite initial logits, CE gradient flow and fixed steps 0 < η ≤ 1/20 fit the label but drive non-target probabilities to the uniform distribution ν.This differs from TPKD’s conditional limit when q is nonuniform.Source: §C.3 · p. 20

C.4 WHEN A COMPLETE STEP LOWERS ITS OWN CONDITIONAL ERROR

  • Under the stated sufficient conditions, a complete update strictly lowers its own conditional KL while retaining the label-step margin gains.The result applies when the conditional signal dominates the label-gradient contribution sufficiently and the step is sufficiently small.Source: §C.4 · p. 21

C.5 GENERAL POSITIVE CONDITIONAL WEIGHTS

  • For a fixed positive conditional target and sufficiently small constant steps, the general positive-weight update converges to correct-class probability one and conditional probabilities q.The conditional direction also tends to zero.Source: §C.5 · p. 21

D PROOFS SUPPORTING EXACT REALIZATION AND PARAMETER LEARNING

These proofs connect output-space guarantees to parameter updates by establishing exact head realization and conditions for learning through backpropagation. They also give a low-dimensional procedure for checking the required angle–spectrum certificates.

D.1 EXACT HEAD REALIZATION

  • With full-row-rank frozen features, the pseudoinverse realizes the prescribed batch output step through a unique minimum-norm head displacement; repeated updates inherit output convergence.For finite initial logits and 0 < η ≤ 4/5, each row’s correct-class probability tends to one and conditional probabilities tend to its target.Source: §D.1 · p. 22

D.2 THE ANGLE–SPECTRUM ARGUMENT √

  • A positive measured cosine certifies a positive parameter-space inner product when the restricted squared-stretch ratio is below (1 + c)/(1 − c).Each signal pair is tested on its own subspace and spectrum.Source: §D.2 · pp. 22–23

D.3 CONDITIONAL INCREMENT AND FINITE-STEP ADVANTAGE

  • Under the stated restricted-spectrum condition, the TPKD parameter step improves conditional loss over compensation only for sufficiently small positive steps.The finite-step bound supplies a sufficient step-size limit when the loss gradient is Lipschitz on both step segments.Source: §D.3 · p. 23

D.4 JOINT-POTENTIAL DESCENT

  • Under its restricted-spectrum condition, the TPKD parameter step strictly decreases the joint potential for sufficiently small positive steps.The condition bounds the joint signal subspace’s largest-to-smallest squared stretch ratio by κΦ.Source: §D.4 · p. 23

D.5 COMPUTING THE TWO-DIMENSIONAL CERTIFICATES

  • The two-dimensional certificates require vector–Jacobian products and small compressed matrices, not a full Jacobian or kernel matrix.Each signal pair uses its own subspace, spectrum, and measured angle.Source: §D.5 · p. 24

E.1 MODELS, DATA, OPTIMIZATION AND EVALUATION

The appendix specifies the shared data, model, optimization, and evaluation protocols, then details the blocking, update, and training comparisons used to measure TPKD.

E.1 MODELS, DATA, OPTIMIZATION AND EVALUATION

  • The visual protocol uses CIFAR-100, a VOLO-D2 teacher, a PiT-B student, and eight fixed augmented views; CLINC150 uses a fixed 18,000-example training pool and separate 4,500-example test set.Full-training vision and text runs use their specified shared training settings and evaluation procedures.Source: §E.1 · p. 24

E.2 BLOCKING AND EXACT HEAD UPDATES

  • Blocking is checked on 16,000 image–view pairs, while exact head-update candidates use two states and 32 batches per state at output step η = 0.01.The five candidate directions include CE, compensation only, unprojected conditional learning, TPKD, and CE plus full KL.Source: §E.2 · p. 25

E.3 BACKPROPAGATION CONDITIONS AND NATIVE PAIRED UPDATES

  • Native paired updates compare TPKD with compensation only under matched restored training states, while restricted spectra test whether the conditional parameter direction is useful.The geometric tests assess the direction; native pairs measure its additional benefit under the training optimizer.Source: §E.3 · pp. 25–26

E.4 THE 128-ITERATION CONTINUATION

  • The 128-iteration continuation compares CE, compensation only, and TPKD from the same CE60 state, recording conditional KL at six specified iterations.A fixed observation set disjoint from update images measures conditional learning.Source: §E.4 · p. 26

E.5 THE FIVE COMPONENT CONTROLS

  • Five full-training controls isolate projection, the CE gradient, and safe full-KL replacement while keeping initialization, teacher, data order, schedule, and endpoint selection unchanged.Each control recomputes its direction at every iteration.Source: §E.5 · p. 26

E.6 IMPLEMENTATION OF THE PRESCRIBED DIRECTION

  • The implementation forms the conditional mismatch, projects it to d, sets v = h + d/2, and detaches v for the surrogate parameter gradient.Log-softmax values are used for conditional losses and compensation statistics before the native optimizer updates parameters.Source: §E.6 · p. 26
Loading 2609.39338v1…