Teacher Supervision over Representation Equivalence Classes

2026-07-03Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors explain that when teaching a smaller model (student) from a larger one (teacher), it's misleading to copy exact internal features because those features can shift in ways that don’t change what the teacher actually does. Instead, they suggest focusing on what the teacher can do (its output behavior) rather than exact feature matching. Their work shows that aligning the teacher's and student's output functions is what really helps the student learn properly, while matching internal features helps match the representation's shape but not the capability. They tested this idea on large language models and found that copying outputs is key to preserving ability, while internal alignment improves the geometric similarity between models.

Knowledge distillationLogitsHidden featuresRepresentation learningOrthogonal transformationCKA (Centered Kernel Alignment)Output functionModel capabilityFeature matchingGeometric alignment
Authors
Sang Il Han
Abstract
Knowledge distillation is usually framed as a choice of what to match in the teacher - its logits, hidden features, or sample relations - which presupposes that the teacher's representation has absolute coordinates to match. It does not: a pretrained representation is identifiable only up to an orthogonal-and-isotropic-scaling equivalence class, so a student should learn the teacher's equivalence class, not its features. The organizing fact is that capability is the teacher's output function, a class invariant that factors through the quotient by the class action, so an objective recovers capability exactly when it is defined there. This makes absolute feature matching ill-posed, and admissible supervision a matter of targeting class invariants (Gram structure, CKA, principal subspaces) or aligning coordinates first, unifying feature matching, relational distillation, alignment, and grafting in one geometric account. We validate our framework on Qwen2.5 and Llama-3.1. A restoration study recovers a corrupted model's representation (CKA ~ 0.99) but not its capability, and an ablation isolates the cause: output-function (logit) matching drives capability, while matching hidden representations aligns geometry without restoring function. Recovery is confined to the corpus-covered region, and a graft study confirms that boundary overlap predicts transplant success but is necessary, not sufficient.