Knowledge distillation

Deep dive

How it works

Parts:

• teacher — frozen large model, ensemble or search; supplies logits z_t
• transfer set — labelled or unlabelled inputs, or the student's own samples (on-policy)
• student — smaller model, trained at temperature T, served at T = 1
• loss — α·CE(y, p_s) + (1−α)·T²·KL(p_t^T ‖ p_s^T), p^T = softmax(z/T)
• logit store (optional) — teacher top-k logits precomputed once

Operations:

1. Freeze the teacher; pick inputs from the transfer set.
2. Get z_t live or from the store; soften both with the same T.
3. Gradient on the student logits = α(p_s − y) + (1−α)·T·(p_s^T − p_t^T); the KL gradient is (p_s^T − p_t^T)/T, times T².
4. Update the student; repeat. Serve at T = 1.

 input ─► teacher (frozen) ─ z_t ─► softmax(z_t/T) ─┐
   │        (live, or top-k logit store)            ├─► T²·KL ─┐
   └─► student ───────── z_s ─► softmax(z_s/T) ─────┘          ├─► sum ─► ∇ ─► student
 label y ──────────────────────► CE(y, softmax(z_s)) ───────────┘

Why T². At high T the gradient tends to (z_s − z_t)/K for centred logits: plain logit matching (Caruana's method). The unscaled gradient falls as 1/T², so T² keeps α meaningful when T changes. A 30-unit MNIST student needed T = 2.5-4, larger ones anything above 8 (Hinton et al.).

Sequence level (Kim and Rush). Seq-KD trains the student on the teacher's beam search output as if it were data (the mode stands in for the sequence distribution); black-box teachers work.

Forward vs reverse KL; on-policy. Forward KL is mode-covering: a weak language-model student puts mass on tokens the teacher finds unlikely. Reverse KL is mode-seeking: more precise, less diverse (MiniLLM). GKD trains on the student's own samples scored by the teacher, with a student-data fraction λ and a generalised JSD interpolating both KLs.

Features, self, search. FitNets add hint losses on intermediate layers via a regressor. Born-again networks use a student as large as the teacher and still beat it. AlphaZero trains its policy toward MCTS visit-count probabilities.

Code

Fixed teacher (random tanh MLP), linear softmax students, 40 examples with labels sampled from the teacher. Asserts: gradient vs finite differences, T² scaling, soft beats hard.

const D: usize = 10; // inputs
const K: usize = 5; // classes
const H: usize = 16; // teacher hidden units

struct Rng(u64);
impl Rng {
    fn f64(&mut self) -> f64 { // xorshift64*, uniform in (0,1)
        self.0 ^= self.0 >> 12; self.0 ^= self.0 << 25; self.0 ^= self.0 >> 27;
        ((self.0.wrapping_mul(0x2545F4914F6CDD1D) >> 11) as f64 + 0.5) / (1u64 << 53) as f64
    }
    fn normal(&mut self) -> f64 { (-2.0 * self.f64().ln()).sqrt() * (6.2831853 * self.f64()).cos() }
}
fn softmax_t(z: &[f64], t: f64) -> Vec<f64> {
    let m = z.iter().cloned().fold(f64::MIN, f64::max);
    let e: Vec<f64> = z.iter().map(|v| ((v - m) / t).exp()).collect();
    let s: f64 = e.iter().sum();
    e.iter().map(|v| v / s).collect()
}
fn kl(p: &[f64], q: &[f64]) -> f64 { p.iter().zip(q).map(|(a, b)| a * (a / b).ln()).sum() }
fn argmax(p: &[f64]) -> usize { (0..K).max_by(|&i, &j| p[i].partial_cmp(&p[j]).unwrap()).unwrap() }

/// T^2 * KL(softmax(zt/T) || softmax(zs/T)) and its gradient w.r.t. z_s: T * (ps_T - pt_T).
fn kd_loss_grad(zs: &[f64], zt: &[f64], t: f64) -> (f64, Vec<f64>) {
    let (ps, pt) = (softmax_t(zs, t), softmax_t(zt, t));
    (t * t * kl(&pt, &ps), ps.iter().zip(&pt).map(|(a, b)| t * (a - b)).collect())
}
/// Teacher: fixed 2-layer tanh MLP (u: H x D, v: K x H). Student: linear, w: K x D.
fn teacher(u: &[f64], v: &[f64], x: &[f64]) -> Vec<f64> {
    let h: Vec<f64> = (0..H).map(|j| (0..D).map(|i| u[j * D + i] * x[i]).sum::<f64>().tanh()).collect();
    (0..K).map(|k| (0..H).map(|j| v[k * H + j] * h[j]).sum()).collect()
}
fn student(w: &[f64], x: &[f64]) -> Vec<f64> { (0..K).map(|k| (0..D).map(|i| w[k * D + i] * x[i]).sum()).collect() }

/// Full-batch GD on a*CE(hard) + (1-a)*T^2*KL(teacher_T || student_T); a = 1 is hard labels only.
fn train(xs: &[Vec<f64>], ys: &[usize], zts: &[Vec<f64>], a: f64, t: f64) -> Vec<f64> {
    let mut w = vec![0.0; K * D];
    for _ in 0..400 {
        let mut g = vec![0.0; K * D];
        for ((x, &y), zt) in xs.iter().zip(ys).zip(zts) {
            let zs = student(&w, x);
            let (p1, (_, gk)) = (softmax_t(&zs, 1.0), kd_loss_grad(&zs, zt, t));
            for k in 0..K {
                let dz = a * (p1[k] - (k == y) as u8 as f64) + (1.0 - a) * gk[k];
                for i in 0..D { g[k * D + i] += dz * x[i] / xs.len() as f64; }
            }
        }
        for (wi, gi) in w.iter_mut().zip(&g) { *wi -= 0.5 * gi; }
    }
    w
}

fn main() {
    // Finite-difference check.
    let (zs, zt) = ([0.3, -1.2, 2.0, 0.7, -0.4], [1.5, 0.1, -0.8, 2.2, 0.0]);
    for t in [1.0, 2.0, 4.0, 8.0] {
        let g = kd_loss_grad(&zs, &zt, t).1;
        for k in 0..K {
            let (mut a, mut b) = (zs, zs);
            a[k] += 1e-6; b[k] -= 1e-6;
            let fd = (kd_loss_grad(&a, &zt, t).0 - kd_loss_grad(&b, &zt, t).0) / 2e-6;
            assert!((fd - g[k]).abs() < 1e-6);
        }
    }
    // With T^2 the gradient norm stays near the high-T limit |centred(zs - zt)| / K; without it, it falls ~1/T^2.
    let norm = |v: &[f64]| v.iter().map(|x| x * x).sum::<f64>().sqrt();
    let c = |z: &[f64], k: usize| z[k] - z.iter().sum::<f64>() / K as f64;
    let limit = norm(&(0..K).map(|k| (c(&zs, k) - c(&zt, k)) / K as f64).collect::<Vec<_>>());
    let (n4, n32) = (norm(&kd_loss_grad(&zs, &zt, 4.0).1), norm(&kd_loss_grad(&zs, &zt, 32.0).1));
    println!("scaled grad norm T=4 {n4:.4}, T=32 {n32:.4}, limit {limit:.4}; unscaled {:.5}, {:.6}", n4 / 16.0, n32 / 1024.0);
    assert!((n32 - limit).abs() < 0.03 * limit && (n4 / n32 - 1.0).abs() < 0.1 && n4 / 16.0 > 40.0 * n32 / 1024.0);
    let mut rng = Rng(0x9E3779B97F4A7C15);
    let u: Vec<f64> = (0..H * D).map(|_| rng.normal() * 0.5).collect();
    let v: Vec<f64> = (0..K * H).map(|_| rng.normal() * 1.2).collect();
    let sample = |n: usize, r: &mut Rng| -> Vec<Vec<f64>> { (0..n).map(|_| (0..D).map(|_| r.normal()).collect()).collect() };
    let test = sample(5000, &mut rng);
    let settings = [("hard labels", 1.0, 1.0), ("soft T=1", 0.0, 1.0), ("soft T=4", 0.0, 4.0)];
    let mut res = vec![(0.0, 0.0); settings.len()];
    for _ in 0..20 {
        let xs = sample(40, &mut rng);
        let zts: Vec<Vec<f64>> = xs.iter().map(|x| teacher(&u, &v, x)).collect();
        let ys: Vec<usize> = zts.iter().map(|z| {
            let (p, mut r) = (softmax_t(z, 1.0), rng.f64());
            p.iter().position(|q| { r -= q; r < 0.0 }).unwrap_or(K - 1)
        }).collect();
        for (r, &(_, a, t)) in res.iter_mut().zip(&settings) {
            let w = train(&xs, &ys, &zts, a, t);
            for x in &test {
                let (pt, ps) = (softmax_t(&teacher(&u, &v, x), 1.0), softmax_t(&student(&w, x), 1.0));
                r.0 += kl(&pt, &ps) / 5000.0 / 20.0;
                r.1 += (argmax(&pt) == argmax(&ps)) as u8 as f64 / 5000.0 / 20.0;
            }
        }
    }
    for (s, r) in settings.iter().zip(&res) { println!("{:<12} KL(teacher||student) {:.4}  top-1 agreement {:.3}", s.0, r.0, r.1); }
    assert!(res[2].0 < res[0].0 && res[2].1 > res[0].1); // soft T=4 beats hard labels on both
}

Measured on this Mac (Apple M4 Pro), seed 0x9E3779B97F4A7C15, mean of 20 draws, 5,000 test points, 400 full-batch steps, no regularisation:

Gradient norm: T²-scaled 0.746 (T = 4) and 0.722 (T = 32), limit 0.717; unscaled 0.0466 and 0.0007. With 20 steps (early stopping) hard labels reach KL 0.599, agreement 0.577; soft T = 4 reaches 0.342 and 0.722.

Complexity and constants

• Per example: teacher forward, student forward and backward; softmax O(K), per token O(V) in a language model.
• The teacher needs no gradients or optimiser state: about a third of a training step of equal size (2N vs 6N FLOPs per token, order of).
• Store (arithmetic): V = 32,000 fp16 logits = 64,000 B per token; top-20 as (int32, fp16) = 120 B, 533× smaller; 10⁹ tokens: 64 TB vs 120 GB.

Efficiency

• Cost: training = student step + one teacher forward; serving = student only. DistilBERT: 71 % faster than BERT-base on an iPhone 7 Plus, 207 MB (Sanh et al.).
• Levers:
  • precomputed top-k logits — no teacher pass per epoch (storage above) — fixed, reused data
  • initialise from the teacher — DistilBERT keeps one layer of two; random init costs 3.69 GLUE macro points (Sanh et al.) — equal widths
  • sequence-level KD — student 10× faster than teacher, +4.2/1.7 BLEU over baseline (greedy/beam) (Kim and Rush) — translation, API teachers
• Naive → standard → state of the art:
  • hard labels only · cheapest, baseline quality
  • Hinton loss, live teacher · one extra forward per step
  • student samples scored by the teacher, compressed top-k logits (GKD, DistillKit) · teacher forward only where needed

Implementation notes

• Same T both sides; log_softmax for stability.
• Top-k stores: renormalise over the kept indices or add a tail bucket; fix this before training.

Pitfalls

• Forgetting T²: at T = 4 the soft gradient is 16× smaller than the hard term (arithmetic).
• Copying teacher mistakes: α > 0 lets labels correct them.
• Offline top-k logits only score the teacher's sequences; on-policy training needs a live teacher.

Variants and alternatives

In production

• Speech: a 10-model ensemble distilled into one 85M-parameter acoustic model (frame accuracy 58.9 % → 60.8 %, WER 10.9 % → 10.7 %; Hinton et al.).
• Text: DistilBERT (distilbert/distilbert-base-uncased); Kim and Rush for translation.
• Game search: AlphaZero (policy toward search probabilities).

Implementations

Sources

• Hinton, Vinyals, Dean, Distilling the Knowledge in a Neural Network, 2015. paper
• Romero et al., FitNets: Hints for Thin Deep Nets, ICLR (arXiv 2014). paper
• Kim, Rush, Sequence-Level Knowledge Distillation, EMNLP 2016. paper
• Furlanello et al., Born Again Neural Networks, ICML 2018. paper
• Silver et al., Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm, 2017. paper
• Sanh et al., DistilBERT, 2019. paper
• Gu et al., MiniLLM: On-Policy Distillation of Large Language Models, ICLR 2024. paper · code
• Agarwal et al., On-Policy Distillation of Language Models, ICLR 2024. paper
• TRL GKD trainer docs; DistillKit README