The Distributional View of Knowledge Distillation

2026-08-15

AI Summary

This study addresses the limitation of traditional token-level knowledge distillation in overlooking probability mass allocation by proposing a distribution-centric distillation framework. The method constructs multi-temperature view families with geometry-aware aggregation, leveraging optimal transport techniques—including entropic Wasserstein barycenters and debiased Sinkhorn divergence—to train student models. We establish the equivalence of log-linear pooling and demonstrate that optimal loss depends on the teacher-student performance gap. Furthermore, we uncover a dispersion law and a dual-mechanism landscape, revealing that aggregation benefits are governed by view dispersion while optimal loss adapts dynamically to performance ceilings. These findings effectively overcome the inherent deficiencies of pointwise matching approaches in knowledge transfer.

Abstract

Token-level knowledge distillation (KD) matches two conditional distributions per position, yet the standard objectives compare them pointwise: a Kullback-Leibler gradient is blind to which wrong token receives probability mass. We develop a distributional view in which the teacher is represented not by a single softened output but by a family of multi-temperature views - marginals of the annealing path of its logits - and the student is trained against a geometry-aware aggregate of these views under an embedding-based ground cost. We formalize the resulting design space (mixtures, log-linear pooling, entropic Wasserstein barycenters, and a debiased Sinkhorn-divergence flagship in hub and path forms), prove an exact collapse result showing log-linear pooling of tempered views is equivalent to a single temperature, and give a multi-marginal Schrodinger-bridge reading that yields falsifiable predictions. On instruction-tuned Pythia pairs, experiments yield three empirical laws: (i) dispersion law - the benefit of multi-temperature aggregation grows monotonically with the effective temperature dispersion of the views, not with their number; (ii) dispersed views unlock the aggregation operator - the barycenter separates from the arithmetic mixture exactly when transport-based aggregation starts to beat averaging; and (iii) two-regime picture governed by the ceiling gap $Γ=\mathrm{PPL}_{\mathrm{SFT}}-\mathrm{PPL}_{T}$: when the fine-tuned teacher barely beats a supervised student the gentle transport objective is the best KD loss but no KD beats supervised fine-tuning, whereas at a real ceiling the ranking inverts - and the sign of the fidelity-generalization correlation flips. We argue that "which distillation loss is the best" is not a fixed property of the loss but a function of $Γ$.

Problem

• Knowledge Distillation
• Token-level KD
• Distributional View
• Geometry-aware Aggregation

Innovation

• Distributional Knowledge Distillation
• Multi-temperature Views
• Geometry-aware Aggregation
• Entropic Wasserstein Barycenters
• Ceiling Gap