PaperlayerView exact revision on arXiv

Source-linked AI summary

The Distributional View of Knowledge Distillation

Gordei Verbii, Juho Lee

arXiv:2608.15215v1stat.MLcs.LG

TL;DR

Standard knowledge distillation compares softened token distributions pointwise, overlooking relationships between token meanings and the information available across temperatures. This paper aggregates multi-temperature teacher views with transport-aware geometry and finds that the best distillation loss depends on the teacher–student ceiling gap Γ, with rankings reversing across regimes.

  • Problem

    Standard distillation uses one temperature and pointwise divergences that treat probability assigned to synonyms like probability assigned to unrelated tokens.

  • Method

    The paper represents teachers with multi-temperature views and trains students against their transport-geometric aggregate, including Wasserstein barycenters and Sinkhorn objectives.

  • Results

    Across six runs and four settings, the ceiling gap Γ separates regimes where optimal distillation rankings reverse, with no KD beating SFT under a thin ceiling and JS nearly matching the teacher at a real ceiling.

  • Takeaways & Limitations

    The best distillation loss is not fixed: transport is robust when the teacher has little to teach, while the preferred objective depends on the ceiling gap Γ.

  • Takeaways & Limitations

    The main regime-II result rests on one model pair, one dataset, five seeds, and incomplete large-pair ablations, while the crossover is bracketed by only two points.

Abstract

from arXiv · show

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 $Γ$.

1 Introduction

This paper recasts knowledge distillation as aggregating a teacher’s multi-temperature output views in transport geometry rather than matching one softened distribution pointwise. It shows that dispersion governs when aggregation helps, while the ceiling gap Γ determines which KD objective is effective and whether KD can beat supervised fine-tuning.

  • Distributional view: The proposed distributional KD represents the teacher with K tempered views and trains the student against a transport-based aggregate, using Wasserstein barycenters or debiased Sinkhorn divergence.The study compares these methods with SFT, KL variants, Jensen-Shannon KL, and arithmetic-mixture targets on instruction-tuned Pythia pairs [5] over Dolly-15k.
  • Formal framework: Log-linear pooling of one teacher’s tempered views collapses exactly to a single temperature, so multi-view gains must arise from mixtures or transport barycenters.The formal framework also develops tempered exponential-family views, entropic optimal transport, and the debiased Sinkhorn divergence [11].
  • Empirical laws: Multi-temperature aggregation helps according to effective view dispersion rather than view count, as predicted by the multi-marginal Schrödinger-bridge reading [7].A controlled dispersion sweep and weights-versus-count probe test this monotone law.
  • Empirical laws: Dispersed views unlock transport aggregation: the barycenter separates from the arithmetic mixture precisely as the Sinkhorn-barycenter performance advantage over mixture averaging grows.The mechanism is measured through the barycenter–mixture ℓ1 gap and the corresponding performance gap.
  • Two-regime picture: Across six runs and four settings, the ceiling gap Γ separates regimes: thin ceilings favor OT for minimizing KD tax but no KD beats SFT, whereas real ceilings reverse rankings and make JS nearly match the teacher.The crossover is bracketed by Γ⋆∈(0.10, 1.49); fidelity anti-correlates with generalization in thin ceilings but correlates positively at real ceilings, consistent with Stanton et al..

2 Related work

Prior work spans softened-softmax and policy-based knowledge distillation, computational optimal transport, and ensemble diversity; this paper combines these strands for tempered views from one teacher.

  • Knowledge distillation: Knowledge distillation includes softened-softmax objectives [14], reverse-KL policy-style training, and on-policy generalized-JS objectives, while fidelity and generalization can dissociate.
  • Computational optimal transport: Computational optimal transport provides the Sinkhorn algorithm [9], debiased Sinkhorn divergence [11], Wasserstein barycenters, fast Bregman-projection computation [4] [10], and entropic-bias correction [15].
  • Ensembles and diversity: Although distilling genuinely different models can benefit from source diversity [16], tempered views of one teacher require proving which aggregation operators can create diversity benefits.

3 Notation and preliminaries

The paper defines temperature-scaled teacher and student distributions along an annealing path, then frames every KD method as a masked likelihood plus a divergence between path-derived conditionals. Geometric methods restrict these conditionals to a shared top-k support, where transport addresses pointwise KL’s token-identity blindness.

  • Annealing path: The annealing path is a smooth exponential-family curve from the uniform distribution at β→0 to the greedy mode at β→∞, sharpening monotonically as β increases.Its score identity is d/dβ log pβ(v) = z_v − E_pβ[z].
  • KD template: Every KD method combines masked ground-truth negative log-likelihood with a divergence between teacher and student conditionals, differing in the divergence and annealing-path functionals compared.
  • Divergences: Pointwise KL charges the same gradient for assigning mass to a synonym or an arbitrary token, motivating transport-based geometry.The forward KL gradient under conventional temperature-squared scaling is τ(pS − pT).
  • Support restriction: Geometric methods condition all distributions on the shared top-k support of position-averaged teacher probabilities, reducing the cost matrix from V×V to k×k.Using k=256 instead of k=64 showed no benefit here: 11.40 versus 11.46 in Table 2.

4 The BaryKD family

The BaryKD family aggregates multi-temperature teacher views using embedding-based transport geometry, spanning arithmetic mixtures, entropic barycenters, and direct Sinkhorn-divergence objectives. Its flagship BaryOT avoids precomputed targets by pulling the student directly toward the views.

  • Ground cost: The ground cost is normalized embedding cosine distance, symmetric, bounded, and zero on the diagonal, without requiring the triangle inequality.This cost supplies semantic dissimilarity between vocabulary tokens and satisfies the assumptions used by the subsequent theory.
  • Transport objective: Entropic regularization makes transport strictly convex and GPU-fast via Sinkhorn iterations, while Sinkhorn divergence removes self-blur and vanishes exactly when distributions match [11].In optimization, the teacher self-term is constant and omitted from the computational graph.
  • Views and aggregation targets: The teacher views use endpoint temperatures {0.5, 1.5} by default or interior temperatures {0.7, 1.0, 1.3}, with simplex weights that are uniform unless stated.These views are defined as tempered teacher distributions µk = softmax(zT/τk).
  • Methods: Aggregation options comprise arithmetic-mixture KL, entropic Sinkhorn barycenter KL, BARYGD, and direct BaryOT, alongside SFT and standard KL-family baselines.The entropic barycenter is a transport-geometric Fréchet mean computed by iterative Bregman projections [4], with relevant barycenter theory from Agueh and Carlier and bias correction from Janati et al. [15].
  • Methods: BaryOT directly transports the student toward multi-temperature teacher views through Sinkhorn divergence, using hub and path forms rather than a precomputed target.The family also includes BARYGD as a 15-step unrolled stress test of the fixed-point barycenter solver.
  • Implementation: Geometric methods use top-k supports at four uniformly sampled positions without τKD^2 scaling, whereas KL losses use full-vocabulary positions with conventional scaling.All methods otherwise share the training loop, data, α, and random seeds.

5 Theory: three means, two objectives, one bridge

The section distinguishes geometric, arithmetic, and transport means for tempered teacher views, showing that geometric pooling collapses to one temperature while transport respects token geometry. It then frames hub and path objectives through a surrogate multi-marginal bridge and predicts view-set effects depend on dispersion and the ceiling gap Γ.

  • The trichotomy of means: Geometric pooling of tempered teacher views is exactly equivalent to choosing a single temperature, so multi-view gains must come from mixtures or barycenters.The normalized geometric mean remains a tempered view, ruling out log-linear pooling as the source of a multi-view benefit.
  • The trichotomy of means: Arithmetic mixtures generally leave the tempered-view family and thicken tails without consulting token geometry, whereas barycenters average under the embedding-induced transport geometry.Mixtures preserve the arg max, while transport provides the geometry-aware third mean.
  • Hub vs. path objectives: Hub mode makes one student view meet all teacher views, whereas path mode matches the student’s tempered curve to the teacher’s at sampled temperatures through shared student logits.Path mode removes hub’s one-sided fixed softening asymmetry, enabling a mechanism test that separates geometry from softening artifacts.
  • The multi-marginal bridge reading: The surrogate bridge reading treats inverse temperature as time and predicts endpoint views bind most strongly, while effective dispersion—not view count K—controls view-set behavior.Interior views should contribute mainly by reweighting the objective toward the middle.
  • The ceiling gap and the two regimes: The ceiling gap Γ separates thin-ceiling and real-ceiling regimes, with a crossover where no KD beats SFT below Γ⋆ but faithful KL-family KD wins above it.For Γ < Γ⋆, gentle transport minimizes the KD tax; the supplied passage identifies approximate regimes Γ ≲0.1 and Γ ≈1.5.

6 Experimental setup

Experiments use compressed Pythia teacher–student pairs fine-tuned on Dolly-15k, with diagnostics separating a saturated thin ceiling from a real capacity ceiling. The study fixes training and evaluation conventions while varying teacher-strength settings and reporting fidelity, generalization, calibration, and uncertainty measures.

  • Models and data: The study distills Pythia 160M→31M and 410M→70M pairs on Dolly-15k, using 3,000 training examples, 200 held-out examples, and 5.3×–5.8× compression [5].Teachers train for three epochs and students for two under matched AdamW loops, with maximum sequence length 256 and fp32 computation.
  • Metrics and statistics: Evaluation measures generalization with held-out PPL, fidelity with agree@1 and predKL, and calibration with 15-bin ECE [13].The protocol also reports seed-paired t-tests, Cohen’s d, and bootstrap 95% confidence intervals.
  • Guard and probe: A zero-shot capacity probe and margin-aware guard distinguish a saturated small-pair ceiling, Γ = +0.10 (THIN), from the big-pair real ceiling, Γ = +1.49 (OK).The pretrained small-pair PPL gap is 328.6 versus 712.6, but fine-tuning compresses it after 3,000 examples; this diagnostic organizes downstream results.
  • Settings: Four teacher-strength settings span under-trained, longer-fine-tuned, saturated-task, and capacity-ceiling conditions, with the first two retained as diagnostics for teachers below their students.The saturated-task condition uses three epochs and 3,000 examples, including two replications and a flagship endpoint-views run.
  • Reported comparisons: Table 1 compares thin- and real-ceiling results over five seeds, while Table 2 probes effective view dispersion, BaryOT settings, endpoint choices, and debiasing variants.The flagship uses endpoint views {0.5, 1.5} in hub mode; Table 2 uses SinkhornBary for the geometric probe and varies BaryOT.

7 Results

Results reveal two ceiling-gap regimes: KD cannot beat supervised fine-tuning near saturation, but at Γ = +1.49 selected KL methods surpass SFT and reverse the loss ranking. Multi-temperature dispersion activates transport aggregation, while path matching improves performance and reduces view sensitivity.

  • Ceiling problem: At Γ = +0.10, no KD method beats SFT, whereas increasing capacity preserves Γ = +1.49 and exposes a real distillation ceiling.The small pair’s pretrained gap is +384 PPL, but identical 3k-example fine-tuning nearly eliminates it; the 410M →70M step preserves the larger post-fine-tuning gap.
  • Fidelity-generalization correlation: Fidelity-generalization correlation flips with Γ: VanillaKD is out-generalized near zero, but VanillaKD and JSKD are best generalizers at Γ = 1.49.At Γ ≈0, VanillaKD has predKL 0.984 versus JSKD 1.018 and PPL 11.923 versus 11.535; at Γ = 1.49, both have predKL 1.43.
  • Dispersion law: PPL decreases monotonically with view dispersion, from 13.08 to 11.47, while uniformly weighting interior views remains about 0.8 PPL worse than the endpoint-favoring curve.The predicted endpoint configuration reached 11.44 within 0.02, and SinkhornBary improved from 12.78 to 11.424 across replications.
  • Flagship pathway: Path mode outperforms hub mode and nearly removes view sensitivity: path gives 10.88/10.89 versus hub 11.28/11.12, reaching within 0.12 PPL of SFT.The result favors path matching over fixed-τKD softening asymmetry.
  • Two regimes and inversion: At Γ = +1.49, JSKD reaches 9.888 ± .027 and VanillaKD 10.588, beating SFT while transport and averaging methods fall behind.These are the study’s first KD wins; the ranking mirrors the thin-ceiling regime.
  • Operator effect: Dispersed views unlock the barycenter operator: SinkhornBary beats MixtureKD by −1.30 and −0.99, while clustered views make the aggregates coincide.The barycenter-mixture distance grows from 0.027 for interior views to 0.424 for endpoints in the small pair and 0.191 in the big pair.
  • Robustness and side-channels: The KD tax shrinks with better objectives at thin ceilings and becomes negative for forward-KL and JS at the real ceiling, while calibration and KL-direction orderings remain stable.The flagship tax falls from 0.51 to 0.36 to 0.12; ReverseKD still pays +0.07, and BaryOT beats the teacher’s ECE in every run.

8 Discussion and limitations

The discussion frames transport’s tolerance as protective when the teacher barely surpasses the student but potentially wasteful when the teacher is genuinely stronger, while acknowledging limited evidence and identifying targeted next experiments.

  • Interpretation: Transport tolerance protects against teacher noise in Regime I, yielding the smallest tax, best-in-family PPL, and better-than-teacher calibration.Near-misses are cheap under the transport objective, so the student is not dragged toward teacher noise.
  • Interpretation: Held-out PPL falls monotonically as view dispersion increases, with effective dispersion—not view count—explaining the curve; uniform K=3 is ≈0.8 PPL worse.The midpoint-down-weighted K=3 probe follows the same dispersion curve, while uniform K=3 sits approximately 0.8 PPL above it.
  • Limitations: The evidence for Regime II is limited to one model pair, one dataset, five seeds, and small models trained on 3k examples.The study reports SFT variance of ±0.357 as its largest, while the big-pair dispersion and BaryOT-path ablations remain unfinished.
  • Future work: Future work proposes a fixed-architecture Γ-sweep, BaryOT-path at Γ = 1.49, hybrid JS-plus-Sinkhorn regularization, and on-policy reverse-direction variants.The proposed hybrid tests whether a small γ improves ECE at equal PPL, while BaryOT-path separates objective effects from OT-coupling effects.

9 Conclusion

The paper formalizes multi-temperature, transport-geometric distillation and validates its design space through a controlled eight-method study with teacher-strength diagnostics. Calibration further shows that BaryOT improves on averaging in both regimes, while the method ordering remains stable despite PPL-rank changes.

  • Conclusion: The framework establishes scope-specific roles for endpoint views, path matching, and transport-based objectives within multi-temperature distillation.Endpoint views support target-side aggregation through a monotone dispersion law, while path matching is used for the flagship method.
  • Calibration: Calibration shows BaryOT beats the teacher in both regimes, whereas the averaging family is the only approach that fails to do so.The comparison uses 15-bin ECE for all eight students against each regime’s teacher.
  • Calibration: Although PPL rankings change when the teacher regime flips, the calibration ordering is essentially unchanged across regimes.All models are over-confident in next-token top-1 reliability, with SFT closest to the diagonal.
Loading 2608.15215v1…