Technical Summary: Teacher Supervision over Representation Equivalence Classes

1. Problem Statement

Knowledge distillation is traditionally framed as a choice of what to match within a teacher model: logits (output distributions), hidden features, or sample relations. This framing presupposes that a teacher's representation possesses "absolute coordinates" that a student can directly match.

The paper argues this premise is fundamentally flawed. A pretrained representation is not a fixed matrix but is identifiable only up to an orthogonal-and-isotropic-scaling equivalence class. Because the downstream "reader" (weights) can absorb any rotation or rescaling of the representation without changing the network's output function, the specific coordinates of a feature vector carry no intrinsic meaning. Consequently, attempting to match a teacher's features "basis-for-basis" is an ill-posed objective. The paper posits that capability (e.g., perplexity, top-1 accuracy) is not a property of the representation itself, but an invariant of the representation's equivalence class—specifically, the output function.

2. Methodology and Theoretical Framework

2.1 Representation as an Equivalence Class

The authors define a representation $H$ (layer activations) not as a fixed matrix but as an orbit under the action of the group $G = O(d) \times \mathbb{R}^+$ (orthogonal rotations and isotropic scaling).

• Assumption 2.1 (Absorbable Reparametrization): Any coordinate change $h \mapsto c h Q$ in a representation can be counteracted by a corresponding change in the downstream weights $W_{in} \mapsto c^{-1} Q^\top W_{in}$, leaving the network function $f$ unchanged.
• Proposition 2.1 (Ill-posedness of Absolute Matching): Because the function is invariant under this group action, a loss function that minimizes the distance between absolute coordinates (e.g., $\|H_T - H_S\|_F$) is ill-posed. It penalizes function-preserving reparametrizations and its minimizer depends on an arbitrary choice of representative within the equivalence class.

2.2 The Quotient Space and Capability

The paper distinguishes between two quotient spaces:

1. Representation Quotient ($\pi_{rep}$): Collapses the representation $H$ alone to its equivalence class $[H]$. Invariants here include the Gram matrix, CKA (Centered Kernel Alignment), and principal angles. These measure geometric similarity but do not determine the output function.
2. Joint Quotient ($\pi_{joint}$): Collapses the pair $(H, \theta_{downstream})$ together. The output function $f$ is a well-defined function on this joint quotient.

Proposition 3.2 establishes that capability is a class invariant defined on the joint quotient. Therefore, a distillation objective can only transfer capability if it targets the output-function invariant on this quotient.

2.3 Taxonomy of Distillation Objectives

Based on the quotient framework, the paper categorizes distillation losses into three families:

• (A) Relational Matching: Targets the representation quotient (e.g., CKA, normalized Gram). These are basis-invariant but do not transfer capability on their own because they ignore the reader's coordinate frame.
• (B) Align-then-Match: Estimates an alignment map (e.g., Procrustes rotation or low-capacity linear map $W$) to fix a representative within the class, then matches absolute coordinates. This is necessary for tasks requiring coordinate-specific feature consumption (e.g., grafting) but introduces alignment noise.
• (C) Logit Distillation: Targets the joint quotient directly by matching output distributions. Since the output is already a function of the joint class, this objective is natively defined on the correct quotient and requires no alignment step.

3. Key Contributions

1. Geometric Reframing: The paper reframes teacher supervision from "matching features" to "transferring the representation equivalence class." It identifies capability as a class invariant (the output function) rather than a property of a specific feature matrix.
2. Ill-posedness Proof: It formally proves that absolute feature matching is ill-posed because it separates points within a single functional fiber, penalizing valid reparametrizations.
3. Loss Taxonomy: It provides a unified geometric account for feature matching, relational distillation, alignment, and grafting, distinguishing them by which part of the class structure they fix.
4. Empirical Validation: Through controlled restoration and cross-width distillation experiments on Qwen2.5 and Llama-3.1, the paper isolates the mechanism of capability transfer.

4. Experimental Results

4.1 Restoration Study (Qwen2.5, Llama-3.1)

The authors corrupted a teacher model by re-initializing a fraction of its layers and attempted to restore it using different objectives.

• Result 1: Training with only a basis-invariant relational loss ($L_{CKA}$) successfully restored the representation geometry (CKA $\approx$ 0.99) but failed to restore capability (perplexity ratio $10^2$–$10^4$ worse, top-1 $\approx$ 0).
• Result 2: Training with only a logit loss ($L_{logit}$) restored capability almost perfectly (perplexity ratio $\approx$ 1.01, top-1 $\approx$ 0.98) without aligning the hidden representations (CKA remained low).
• Conclusion: Aligning representations and restoring the function are near-orthogonal tasks. Capability is driven exclusively by the output-function term.

4.2 The Role of Weight Tying

• In models with tied input/output embeddings (Qwen2.5-0.5B), the output head is fixed, so logit matching alone suffices to restore the function.
• In models with untied heads (Qwen2.5-7B, Llama-3.1-8B), restoring capability required explicitly training the output head alongside the logit loss. This confirms that the "reader" must be aligned to the new representation's coordinate frame.

4.3 Cross-Width Distillation

In a genuine cross-width distillation setting (Teacher: 1.5B, Student: 0.5B):

• Feature-only matching (even with a learned projection $W$) caused the student model to collapse (perplexity $> 10^6$).
• Logit-only matching preserved and slightly improved the student's performance.
• Adding a feature term to logit matching provided no benefit and was often mildly harmful.
• A surrogate experiment with an optimally solved alignment map confirmed that feature-only supervision fails structurally, not due to optimization difficulties.

4.4 Grafting and Subspace Overlap

The paper validates that module grafting success is governed by the overlap of boundary subspaces (measured by CKA or Procrustes residual) between the donor and host. High overlap is necessary but not sufficient; grafts near the output layer fail even with high overlap due to sensitivity to small mismatches.

5. Significance and Claims

The paper claims that the intuition that distillation happens "in the basis" is "right in spirit and wrong in letter."

• Core Claim: Teacher supervision should target the equivalence class, not the coordinates. Capability is transferred through the output-function invariant.
• Practical Implication: For transferring capability, logit matching is the only objective natively defined on the correct quotient. Feature matching (even relational or aligned) is insufficient for capability transfer and is only useful for transferring geometric structure or for specific tasks like grafting where coordinate alignment is explicitly required.
• Scope Limitation: The paper emphasizes that distillation is replication, not creation. It transfers the teacher's output function on the subspace covered by the training corpus. It does not claim to generate new capabilities beyond the teacher's existing function or to solve cross-family transfer with mismatched tokenizers.

The paper concludes that "what a teacher knows lives in the class; what transfers is the part of that class the objective and the corpus actually target."