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

표준 knowledge distillation은 softened token distribution을 pointwise하게 비교해 token 의미 간 관계와 temperature 전반에서 이용 가능한 정보를 간과한다. 이 논문은 transport-aware geometry로 multi-temperature teacher view를 aggregate하고, 최적 distillation loss가 teacher–student ceiling gap Γ에 따라 달라지며 regime에 따라 순위가 뒤집힌다는 사실을 보인다.

  • 문제

    표준 distillation은 one temperature와 pointwise divergences를 사용해 synonym에 할당된 probability를 서로 무관한 token에 할당된 probability와 동일하게 취급한다.

  • 방법

    이 논문은 teacher를 multi-temperature views로 표현하고, Wasserstein barycenters와 Sinkhorn objectives를 포함하는 transport-geometric aggregate를 대상으로 student를 학습시킨다.

  • 결과

    여섯 번의 runs와 네 가지 settings에서 ceiling gap Γ가 optimal distillation rankings가 reverse되는 regime를 가르며, thin ceiling에서는 어떤 KD도 SFT를 능가하지 못하고 real ceiling에서는 JS가 teacher에 거의 근접한다.

  • 시사점 및 한계

    최적 distillation loss는 not fixed다. teacher가 가르칠 것이 적을 때는 transport가 robust하지만, preferred objective는 ceiling gap Γ에 따라 달라진다.

  • 시사점 및 한계

    주요 regime-II 결과는 one model pair, one dataset, five seeds와 불완전한 large-pair ablations에 기반하며, crossover는 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 서론

이 논문은 knowledge distillation을 하나의 softened distribution과 pointwise하게 매칭하는 대신, transport geometry에서 teacher의 multi-temperature output views를 집계하는 문제로 재구성한다. aggregation의 효과가 언제 나타나는지는 dispersion이 좌우하며, 어떤 KD objective가 효과적인지와 KD가 supervised fine-tuning을 능가할 수 있는지는 ceiling gap Γ가 결정함을 보인다.

  • Distributional view: 제안하는 distributional KD는 teacher를 K개의 tempered views로 나타내고, Wasserstein barycenters 또는 debiased Sinkhorn divergence를 사용해 student를 transport-based aggregate에 맞춰 학습한다.이 연구는 instruction-tuned Pythia pairs [5]를 대상으로 Dolly-15k 에서 이 방법들을 SFT, KL variants, Jensen-Shannon KL, arithmetic-mixture targets와 비교한다.
  • Formal framework: 한 teacher의 tempered views를 log-linear pooling하면 정확히 하나의 temperature로 collapse하므로, multi-view gain은 mixtures 또는 transport barycenters에서 나와야 한다.이 formal framework는 tempered exponential-family views, entropic optimal transport, debiased Sinkhorn divergence [11]도 전개한다.
  • Empirical laws: Multi-temperature aggregation의 효과는 view count가 아니라 effective view dispersion에 따라 커지며, 이는 multi-marginal Schrödinger-bridge 해석 [7]이 예측한 바와 같다.통제된 dispersion sweep과 weights-versus-count probe를 통해 이 monotone law를 검증한다.
  • Empirical laws: Dispersed views는 transport aggregation을 활성화한다. 즉 Sinkhorn-barycenter의 성능이 mixture averaging보다 우월해질수록 barycenter는 arithmetic mixture에서 정확히 분리된다.이 메커니즘은 barycenter–mixture ℓ1 gap과 이에 대응하는 performance gap으로 측정한다.
  • Two-regime picture: 여섯 번의 runs와 네 가지 settings에서 ceiling gap Γ가 두 regime를 가른다. thin ceilings에서는 KD tax를 최소화하는 데 OT가 유리하지만 어떤 KD도 SFT를 능가하지 못하는 반면, real ceilings에서는 ranking이 역전되고 JS가 teacher와 거의 맞먹는다.crossover는 Γ⋆∈(0.10, 1.49)로 한정되며, thin ceilings에서는 fidelity가 generalization과 음의 상관을 보이지만 real ceilings에서는 양의 상관을 보인다. 이는 Stanton et al. 과 일치한다.

2 관련 연구

선행 연구는 softened-softmax 및 policy 기반 knowledge distillation, computational optimal transport, ensemble diversity를 아우르며, 본 논문은 이 흐름들을 한 teacher에서 얻은 tempered views에 결합한다.

  • Knowledge distillation: Knowledge distillation에는 softened-softmax objectives [14], reverse-KL policy-style training, on-policy generalized-JS objectives 가 포함되며, fidelity와 generalization은 서로 분리될 수 있다.
  • Computational optimal transport: Computational optimal transport는 Sinkhorn algorithm [9], debiased Sinkhorn divergence [11], Wasserstein barycenters, fast Bregman-projection computation [4] [10], entropic-bias correction [15]을 제공한다.
  • Ensembles and diversity: 실제로 서로 다른 model을 distilling할 때는 source diversity [16]의 이점을 얻을 수 있지만, 한 teacher의 tempered views에서는 어떤 aggregation operators가 diversity benefits를 만들어낼 수 있는지 입증해야 한다.

3 표기와 예비 지식

이 논문은 annealing path를 따라 temperature-scaled teacher 및 student distribution을 정의한 뒤, 모든 KD method를 masked likelihood와 path에서 유도된 conditional 간 divergence의 결합으로 정식화한다. Geometric method는 이 conditional을 공통 top-k support로 제한하며, transport는 pointwise KL이 token identity에 둔감하다는 문제를 다룬다.

  • Annealing path: annealing path는 β→0에서 uniform distribution으로 시작해 β→∞에서 greedy mode에 이르는 smooth exponential-family curve이며, β가 증가할수록 단조롭게 sharp해진다.score identity는 d/dβ log pβ(v) = z_v − E_pβ[z]이다.
  • KD template: 모든 KD method는 masked ground-truth negative log-likelihood와 teacher 및 student conditional 간 divergence를 결합하며, 비교되는 divergence와 annealing-path functional이 서로 다르다.
  • Divergences: Pointwise KL은 synonym에 probability mass를 부여하는 경우와 임의의 token에 probability mass를 부여하는 경우에 동일한 gradient를 부과하므로, transport-based geometry가 필요하다.관례적인 temperature-squared scaling을 적용한 forward KL gradient는 τ(pS − pT)이다.
  • Support restriction: Geometric method는 position-averaged teacher probability를 기준으로 모든 distribution을 공통 top-k support에 조건화하여, cost matrix를 V×V에서 k×k로 줄인다.여기서는 k=64 대신 k=256을 사용해도 이득이 없었다: Table 2에서 11.40 versus 11.46.

4 BaryKD 계열

BaryKD 계열은 embedding 기반 transport geometry를 사용해 multi-temperature teacher view를 aggregate하며, arithmetic mixture, entropic barycenter, direct Sinkhorn-divergence objective를 포괄한다. 대표 방법인 BaryOT는 미리 계산한 target 없이 student를 view 쪽으로 직접 끌어당긴다.

  • Ground cost: Ground cost는 normalized embedding cosine distance로, symmetric하고 bounded하며 diagonal에서 0이고 triangle inequality를 요구하지 않는다.이 cost는 vocabulary token 사이의 semantic dissimilarity를 제공하며 후속 theory에서 사용하는 가정을 만족한다.
  • 수송 목적함수: 엔트로피 정규화는 수송을 엄밀히 볼록하게 만들고 Sinkhorn 반복을 통해 GPU에서 빠르게 계산되도록 하며, Sinkhorn divergence는 자기 블러를 제거하고 분포가 일치할 때 정확히 사라진다 [11].최적화에서 teacher의 self-term은 상수이므로 계산 그래프에서 제외한다.
  • View와 aggregation target: Teacher view는 기본적으로 endpoint temperature {0.5, 1.5}를 사용하거나 interior temperature {0.7, 1.0, 1.3}를 사용하며, 별도 언급이 없으면 simplex weight는 uniform이다.이 view는 tempered teacher distribution µk = softmax(zT/τk)로 정의된다.
  • Methods: Aggregation option에는 arithmetic-mixture KL, entropic Sinkhorn barycenter KL, BARYGD, direct BaryOT가 포함되며, SFT와 standard KL-family baseline도 함께 사용한다.Entropic barycenter는 iterative Bregman projection [4]으로 계산하는 transport-geometric Fréchet mean이며, 관련 barycenter theory는 Agueh and Carlier, bias correction은 Janati et al. [15]에 따른다.
  • Methods: BaryOT는 Sinkhorn divergence를 통해 student를 multi-temperature teacher view 쪽으로 직접 transport하며, 미리 계산한 target 대신 hub form과 path form을 사용한다.이 계열에는 fixed-point barycenter solver의 15-step unrolled stress test인 BARYGD도 포함된다.
  • Implementation: Geometric method는 τKD^2 scaling 없이 네 개의 uniformly sampled position에서 top-k support를 사용하고, KL loss는 conventional scaling과 함께 full-vocabulary position을 사용한다.그 외에는 모든 method가 training loop, data, α, random seed를 공유한다.

5 이론: 세 가지 평균, 두 가지 목적함수, 하나의 브리지

이 절은 temperature가 조정된 teacher view에 대해 geometric, arithmetic, transport mean을 구분하고, geometric pooling은 하나의 temperature로 붕괴하는 반면 transport는 token geometry를 반영함을 보인다. 이어 surrogate multi-marginal bridge를 통해 hub와 path 목적함수를 정식화하고, view set의 효과가 view의 수가 아니라 dispersion과 ceiling gap Γ에 좌우된다고 예측한다.

  • 평균의 삼분법: Temperature가 조정된 teacher view의 Geometric pooling은 하나의 temperature를 선택하는 것과 정확히 동치이므로, multi-view 이득은 mixture나 barycenter에서 나와야 한다.정규화된 geometric mean은 여전히 temperature가 조정된 view로 남으므로, log-linear pooling이 multi-view 이득의 원천이라는 가능성을 배제한다.
  • 평균의 삼분법: Arithmetic mixture는 일반적으로 tempered-view family를 벗어나고 token geometry를 참조하지 않은 채 tail을 두껍게 만드는 반면, barycenter는 embedding이 유도한 transport geometry 아래에서 평균을 낸다.Mixture는 arg max를 보존하지만, transport는 geometry-aware한 세 번째 평균을 제공한다.
  • Hub와 path 목적함수: Hub mode는 하나의 student view가 모든 teacher view와 만나게 하는 반면, path mode는 shared student logits를 통해 sampled temperature에서 student의 tempered curve를 teacher의 curve에 맞춘다.Path mode는 hub의 일방적인 fixed softening 비대칭을 제거해, geometry와 softening artifact를 구분하는 mechanism test를 가능하게 한다.
  • Multi-marginal bridge 해석: Surrogate bridge 해석은 inverse temperature를 시간으로 취급하며 endpoint view가 가장 강하게 결합될 것이라고 예측하고, view set의 거동은 view 수 K가 아니라 effective dispersion이 좌우한다고 본다.Interior view는 주로 목적함수가 중간 영역을 향하도록 재가중하는 방식으로 기여해야 한다.
  • Ceiling gap과 두 regime: Ceiling gap Γ는 thin-ceiling regime과 real-ceiling regime을 구분하며, Γ⋆를 경계로 crossover가 나타난다. Γ⋆보다 낮으면 어떤 KD도 SFT를 능가하지 못하지만, 그보다 높으면 faithful KL-family KD가 우세하다.Γ < Γ⋆에서는 gentle transport가 KD tax를 최소화하며, 제시된 passage는 대략적인 regime으로 Γ ≲0.1과 Γ ≈1.5를 식별한다.

6 실험 설정

실험은 Dolly-15k에서 fine-tuning한 압축 Pythia teacher–student pair를 사용하며, 포화된 얇은 ceiling과 실제 capacity ceiling을 구분하는 진단을 포함한다. 연구는 teacher-strength 설정을 변화시키면서 fidelity, generalization, calibration, uncertainty measure를 보고하고, training 및 evaluation convention은 고정한다.

  • Model과 data: Dolly-15k에서 Pythia 160M→31M 및 410M→70M pair를 distillation하며, training example 3,000개, held-out example 200개, compression 5.3×–5.8×을 사용한다 [5].Teacher는 matched AdamW loop에서 three epoch, student는 two epoch 동안 training하며, maximum sequence length는 256이고 computation은 fp32다.
  • Metric과 statistics: Evaluation은 held-out PPL로 generalization을, agree@1 및 predKL로 fidelity를, 15-bin ECE로 calibration을 측정한다 [13].Protocol은 seed-paired t-test, Cohen’s d, bootstrap 95% confidence interval도 보고한다.
  • Guard와 probe: Zero-shot capacity probe와 margin-aware guard는 포화된 small-pair ceiling Γ = +0.10 (THIN)과 big-pair real ceiling Γ = +1.49 (OK)를 구분한다.Pretrained small-pair PPL gap은 328.6 대 712.6이지만, fine-tuning은 3,000개 example 이후 이를 압축한다. 이 diagnostic이 후속 결과를 조직한다.
  • 설정: 네 가지 teacher-strength 설정은 under-trained, longer-fine-tuned, saturated-task, capacity-ceiling 조건을 포괄하며, 앞의 두 설정은 student보다 낮은 teacher를 위한 diagnostic으로 유지한다.Saturated-task 조건은 three epoch와 3,000개 example을 사용하며, 두 replication과 flagship endpoint-views run을 포함한다.
  • 보고된 비교: Table 1은 five seed에 걸친 thin- 및 real-ceiling 결과를 비교하고, Table 2는 effective view dispersion, BaryOT 설정, endpoint 선택, debiasing variant를 검증한다.Flagship은 hub mode에서 endpoint view {0.5, 1.5}를 사용한다. Table 2는 geometric probe에 SinkhornBary를 사용하고 BaryOT를 변화시킨다.

7 결과

결과는 두 가지 ceiling-gap regime을 보여준다: 포화에 가까울 때 KD는 supervised fine-tuning을 이길 수 없지만, Γ = +1.49에서는 선택된 KL methods가 SFT를 앞서며 loss ranking을 뒤집는다. Multi-temperature dispersion은 transport aggregation을 활성화하고, path matching은 성능을 높이면서 view sensitivity를 줄인다.

  • Ceiling 문제: Γ = +0.10에서는 어떤 KD method도 SFT를 이기지 못하지만, capacity를 늘리면 Γ = +1.49가 유지되어 실제 distillation ceiling이 드러난다.작은 pair의 pretrained gap은 +384 PPL이지만, 동일한 3k-example fine-tuning으로 그 차이가 거의 사라진다; 410M →70M 단계에서는 fine-tuning 이후의 더 큰 gap이 유지된다.
  • Fidelity-generalization correlation: Fidelity-generalization correlation은 Γ에 따라 뒤집힌다: 0에 가까울 때는 VanillaKD가 out-generalized되지만, Γ = 1.49에서는 VanillaKD와 JSKD가 가장 잘 generalize한다.Γ ≈0에서 VanillaKD는 predKL 0.984, JSKD는 1.018이며 PPL은 각각 11.923과 11.535다; Γ = 1.49에서는 둘 다 predKL 1.43이다.
  • Dispersion law: PPL은 view dispersion이 증가함에 따라 13.08에서 11.47까지 단조롭게 감소하지만, interior views에 균일한 가중치를 주면 endpoint-favoring curve보다 여전히 약 0.8 PPL만큼 나쁘다.예측된 endpoint configuration은 11.44에 0.02 이내로 도달했고, SinkhornBary는 replication을 거치며 12.78에서 11.424로 향상됐다.
  • Flagship pathway: Path mode는 hub mode보다 뛰어나며 view sensitivity를 거의 제거한다: path는 10.88/10.89, hub는 11.28/11.12를 기록해 SFT보다 0.12 PPL 이내에 도달한다.이 결과는 fixed-τKD softening asymmetry보다 path matching을 선호한다.
  • 두 regime과 inversion: Γ = +1.49에서 JSKD는 9.888 ± .027, VanillaKD는 10.588에 도달해 SFT를 이기지만, transport와 averaging methods는 뒤처진다.이는 이 연구의 첫 KD wins이며, ranking은 thin-ceiling regime과 일치한다.
  • Operator effect: 분산된 views는 barycenter operator를 활성화한다: SinkhornBary는 MixtureKD를 −1.30과 −0.99만큼 앞서지만, clustered views에서는 aggregate들이 일치한다.Barycenter-mixture distance는 작은 pair에서 interior views의 0.027에서 endpoints의 0.424로, 큰 pair에서는 0.191로 증가한다.
  • Robustness와 side-channels: KD tax는 thin ceilings에서 더 나은 objectives를 사용할수록 줄어들고, real ceiling에서는 forward-KL과 JS에 대해 음수가 되지만 calibration과 KL-direction orderings는 안정적으로 유지된다.Flagship tax는 0.51에서 0.36, 0.12로 감소한다; ReverseKD는 여전히 +0.07을 지불하며, BaryOT는 모든 run에서 teacher의 ECE보다 우수하다.

8 논의와 한계

논의에서는 교사가 학생을 근소하게 앞설 때 transport의 tolerance가 보호적으로 작용하지만 교사가 실제로 더 강할 때는 잠재적으로 비효율적일 수 있음을 설명하고, 제한된 근거를 인정하며 후속 실험을 제안한다.

  • 해석: Regime I에서 transport tolerance는 teacher noise를 방어해 가장 작은 tax, family 내 최상의 PPL, teacher보다 나은 calibration을 낳는다.transport objective에서는 아슬아슬한 오차의 비용이 작으므로 student가 teacher noise 쪽으로 끌려가지 않는다.
  • 해석: Held-out PPL은 view dispersion이 증가할수록 단조롭게 감소하며, 곡선은 view count가 아니라 effective dispersion으로 설명된다. uniform K=3은 ≈0.8 PPL 더 나쁘다.midpoint-down-weighted K=3 probe는 같은 dispersion curve를 따르는 반면, uniform K=3은 그보다 약 0.8 PPL 높다.
  • 한계: Regime II에 대한 근거는 one model pair, one dataset, five seeds, 그리고 3k examples로 학습한 small models 하나에 한정된다.이 연구는 SFT variance가 ±0.357로 가장 크다고 보고하며, big-pair dispersion과 BaryOT-path ablation은 아직 완료되지 않았다.
  • 향후 연구: 향후 연구에서는 fixed-architecture Γ-sweep, Γ = 1.49에서의 BaryOT-path, hybrid JS-plus-Sinkhorn regularization, 그리고 on-policy reverse-direction variants 를 제안한다.제안된 hybrid는 동일한 PPL에서 작은 γ가 ECE를 개선하는지 검증하며, BaryOT-path는 objective effect와 OT-coupling effect를 분리한다.

9 결론

이 논문은 다중 온도 및 transport-geometric distillation을 정식화하고, teacher-strength 진단을 포함한 통제된 8개 방법 연구를 통해 그 설계 공간을 검증한다. Calibration은 BaryOT가 두 regime 모두에서 averaging보다 개선됨을 추가로 보여주며, PPL 순위가 달라져도 방법 순서는 안정적으로 유지된다.

  • 결론: 이 framework는 multi-temperature distillation에서 endpoint views, path matching, transport-based objectives의 범위별 역할을 정립한다.Endpoint views는 단조로운 dispersion law에 따라 target-side aggregation을 지원하며, path matching은 flagship method에 사용된다.
  • Calibration: Calibration은 BaryOT가 두 regime 모두에서 teacher를 능가하는 반면, averaging family는 이를 해내지 못하는 유일한 접근법임을 보여준다.비교에는 각 regime의 teacher에 대한 8개 student 모두의 15-bin ECE가 사용된다.
  • Calibration: teacher regime이 바뀌면 PPL 순위는 변하지만, calibration 순서는 regime 전반에서 본질적으로 변하지 않는다.모든 model은 next-token top-1 reliability에서 과도한 확신을 보이며, SFT가 diagonal에 가장 가깝다.
Loading 2608.15215v1…