Skip to content
Residual Stream
Go back

Knowledge Distillation: Learning from a Teacher's Distribution

Updated:
Tan Le
ENVI

Small and Specialized Models found distillation behind several of today’s best small models: Gemma 2’s 2B and 9B, Llama 3.2’s 1B and 3B. This post is how it works. The idea is old and simple: train a small student model to match the output distribution of a large teacher, instead of only the labels in the data. Bucilă, Caruana and Niculescu-Mizil (2006) compressed ensembles this way, and Ba and Caruana (2014) trained shallow networks on deep ones’ outputs; Hinton, Vinyals and Dean (2015) added the temperature and gave it the name distillation. DistilBERT (Sanh et al., 2019) brought it to Transformers, keeping 97% of BERT’s language-understanding score with 40% fewer parameters.

training text “…answer one” Teacher 1.5B · frozen Student 0.5B · trained reads reads ▲ soft label p next token in the text ▲ hard label y: ‘ token’ ▲ student’s guess q loss ½ CE(y, q) + ½ CE(p, q) gradient updates the student only

Distillation in one picture: the student learns from the text’s own next token (the hard label) and from the teacher’s whole distribution over next tokens (the soft label). The student’s bars are illustrative. This post’s experiment trains with each alone and with half of each.

At a glance
  • Soft labels carry more information than hard ones. A teacher’s distribution says which wrong answers were nearly right, at every position: in Hinton et al.’s speech experiment, 3% of the data with soft labels reached 57.0% frame accuracy, against 44.5% with hard labels.
  • Temperature spreads the distribution out so the small probabilities count. Without the τ2\tau^2 factor the gradient falls about 227 times from τ=1\tau = 1 to 16; with it, it stays within 13%.
  • Forward KL covers every mode of the teacher and reverse KL picks one: on a two-bump target, the forward-KL student puts 19% of its mass in the gap where the teacher has 1%.
  • In the experiment, mixing half hard and half soft labels beat hard labels alone and forgot less than half as much: 0.027 bits per byte lost on WikiText, against 0.068.

Table of contents

Open Table of contents

Hard labels and soft labels

Next-token training (Cross-Entropy and the Softmax Gradient) uses a hard label: at each position, the one token the text actually used, with probability 1 for it and 0 for the other 151,664. A teacher model reading the same text produces a soft label: a probability for every token. Here is one, from Qwen2.5-1.5B-Instruct, the teacher used below:

import torch
from transformers import AutoModelForCausalLM
from common import TEACHER, DEVICE, tok

# "Dark knowledge": what a teacher's distribution says beyond its top
# answer. The teacher's next-token probabilities after one sentence of this
# blog, at temperature 1 and softened at temperature 2. A hard label keeps
# only the token the text happened to use.
model = AutoModelForCausalLM.from_pretrained(TEACHER, dtype=torch.float32).to(DEVICE).eval()
text = "A language model generates its answer one"
with torch.no_grad():
    logits = model(tok(text, return_tensors="pt").input_ids.to(DEVICE)).logits[0, -1].float()
print(repr(text))
for tau in (1, 2):
    p = (logits / tau).softmax(-1)
    v, i = p.topk(8)
    print(f"temperature {tau}: " + ", ".join(f"{tok.decode(int(t))!r} {float(x):.3f}" for x, t in zip(v, i)) +
          f"; top 8 hold {float(v.sum()):.0%}")

Full script: dark_knowledge.py

'A language model generates its answer one'
temperature 1: ' word' 0.481, ' character' 0.275, ' sentence' 0.079, ' token' 0.078, ' letter' 0.039, ' step' 0.006, ' symbol' 0.005, ' by' 0.005; top 8 hold 97%
temperature 2: ' word' 0.074, ' character' 0.056, ' sentence' 0.030, ' token' 0.030, ' letter' 0.021, ' step' 0.008, ' symbol' 0.008, ' by' 0.007; top 8 hold 23%

If the text went on “one token at a time”, the hard label is ” token”, probability 1, and says nothing else. The teacher’s distribution says much more: it ranks ” token” fourth, at 8%, behind ” word” (48%), ” character” (28%) and, by a hair, ” sentence” (8%), with ” letter” and ” step” after it. The next word should be a unit of text, and some units are likelier than others. That structure, the same at every position of every training sequence, is what the student can learn from the teacher and can’t learn from the text alone.

At temperature 2 the same eight tokens keep their order but hold only 23% of the probability, against 97% at temperature 1; the other 77% is spread over the rest of the 151,665-token vocabulary, thousands of tokens that each get a sliver. A softened distribution over a vocabulary this large is mostly tail, which matters for the choice of temperature below and is one reason the experiment keeps only the teacher’s top tokens.

‘word’ ‘character’ ‘sentence’ ‘token’ ‘letter’ ‘step’ ‘symbol’ ‘by’ hard label 1 token: 100% 100% 0 100% teacher, τ = 1 top 8 hold 97% 48% 28% 7.9% 7.8% 3.9% 0.6% 0.5% 0.5% 0 100% teacher, τ = 2 top 8 hold 23% 7.4% 5.6% 3.0% 3.0% 2.1% 0.8% 0.8% 0.7% 0 100%

One position, three targets. The hard label puts everything on the token the text used. The teacher ranks a whole family of plausible units of text. At τ = 2 the order is the same but 77% of the mass has moved into the long tail outside these eight.

Data: the teacher's output printed above

That information is what Hinton et al. called dark knowledge: which wrong answers are nearly right. A hard label says one token is correct. A soft label also says which others would have been reasonable, and how reasonable, at every position. Each training example carries far more information, so a student learns more from the same text, or as much from less of it. In the paper’s speech experiment, a model trained on 3% of the data with the teacher’s soft labels reached 57.0% frame accuracy, close to the 58.9% of a model trained on all of it, while the same 3% with hard labels overfit and reached 44.5%.

Temperature, and the τ² factor

A confident teacher puts nearly all its probability on one token, and its soft label is then almost a hard one. Dividing the logits by a temperature τ>1\tau > 1 before the softmax, as in Decoding’s sampling, flattens both distributions and brings the small probabilities up to where they affect the loss. The distillation loss is then the cross-entropy between the softened teacher p=softmax(t/τ)p = \text{softmax}(t / \tau) and the softened student q=softmax(s/τ)q = \text{softmax}(s / \tau), usually mixed with the ordinary hard-label loss.

Its gradient with respect to the student’s logits is (q−p)/τ(q - p) / \tau: the same “prediction minus target” as Cross-Entropy and the Softmax Gradient, divided by τ\tau. And for large τ\tau the difference q−pq - p itself shrinks like 1/τ1 / \tau: expanding ez/τ≈1+z/τe^{z/\tau} \approx 1 + z / \tau for logits that average to zero, qi−pi≈(si−ti)/(Nτ)q_i - p_i \approx (s_i - t_i) / (N \tau) over NN classes. So the gradient shrinks like 1/τ21 / \tau^2, and raising the temperature would quietly turn the soft loss off relative to the hard one. Hinton et al. multiply it by τ2\tau^2 to undo that. Checked numerically on ten classes:

import numpy as np

# Why distillation multiplies its loss by tau^2. With softened teacher
# p = softmax(t / tau) and student q = softmax(s / tau), the gradient of the
# soft cross-entropy with respect to the student's logits is (q - p) / tau.
# For large tau, q - p itself shrinks like 1 / tau, so the gradient shrinks
# like 1 / tau^2 (Hinton et al., 2015): (s - t) / (N tau^2) for zero-mean
# logits over N classes. Multiplying the loss by tau^2 undoes that, so the
# soft term's weight against the hard-label term doesn't depend on tau.
rng = np.random.default_rng(0)
N = 10
t, s = rng.normal(0, 2, N), rng.normal(0, 2, N)
t, s = t - t.mean(), s - s.mean()
softmax = lambda z: np.exp(z - z.max()) / np.exp(z - z.max()).sum()
print(f"{'tau':>4s} {'|gradient|':>11s} {'x tau^2':>9s}   {'large-tau formula':>17s}")
for tau in (1, 2, 4, 8, 16):
    g = (softmax(s / tau) - softmax(t / tau)) / tau
    approx = (s - t) / (N * tau ** 2)
    print(f"{tau:4d} {np.linalg.norm(g):11.5f} {np.linalg.norm(g) * tau ** 2:9.4f}   {np.linalg.norm(approx):17.5f}")

Full script: tau_squared.py

 tau  |gradient|   x tau^2   large-tau formula
   1     0.78426    0.7843             0.89227
   2     0.21055    0.8422             0.22307
   4     0.05381    0.8610             0.05577
   8     0.01366    0.8739             0.01394
  16     0.00345    0.8823             0.00349

Without the factor the gradient falls about 227 times from τ=1\tau = 1 to 16; with it, it stays within 13%, and the large-τ\tau formula is within 4% from τ=4\tau = 4 on. (The experiment below uses τ=1\tau = 1, where the factor is 1.)

0.003 0.01 0.1 1 1 2 4 8 16 temperature τ gradient size × τ²: 0.78 → 0.88 raw: ÷ 227

Without the τ² factor, raising the temperature quietly switches the soft loss off. With it, the gradient keeps its size.

Data: the check above

Forward or reverse KL

Cross-entropy against the teacher is, up to a constant, the forward KL divergence KL(p ∥ q)=∑plog⁡(p/q)\text{KL}(p \,\|\, q) = \sum p \log (p / q): the average, under the teacher, of how badly the student underrates each outcome. It punishes the student hard wherever the teacher has mass and the student doesn’t, and doesn’t care where the student puts mass the teacher doesn’t. The reverse KL KL(q ∥ p)\text{KL}(q \,\|\, p) is the opposite. When the student is too small to match the teacher, the two choose differently. A teacher with two separate modes, fitted by a student that can only be one Gaussian:

import numpy as np

# Forward vs reverse KL. The teacher p is bimodal: two bumps at -2 and +2.
# The student q is a single Gaussian, too simple to match it. Fit q's mean
# and width by minimizing each direction of KL over a grid of candidates.
x = np.linspace(-8, 8, 4001)
dx = x[1] - x[0]
normal = lambda m, s: np.exp(-0.5 * ((x - m) / s) ** 2) / (s * np.sqrt(2 * np.pi))
p = 0.5 * normal(-2, 0.6) + 0.5 * normal(2, 0.6)

def kl(a, b):
    """KL(a || b) = sum a log(a / b), on the grid."""
    mask = a > 1e-300
    return float(np.sum(a[mask] * (np.log(a[mask]) - np.log(b[mask] + 1e-300))) * dx)

best = {}
for m in np.linspace(-3, 3, 121):
    for s in np.linspace(0.3, 4, 75):
        q = normal(m, s)
        for name, value in (("forward KL(p || q)", kl(p, q)), ("reverse KL(q || p)", kl(q, p))):
            if name not in best or value < best[name][0]:
                best[name] = (value, m, s)
for name, (value, m, s) in best.items():
    q = normal(m, s)
    between = q[np.abs(x) < 0.5].sum() * dx                 # mass q puts in the gap between the bumps
    print(f"{name}: best q has mean {m:+.2f}, width {s:.2f}; mass in the gap |x| < 0.5: {between:.0%} "
          f"(teacher: {p[np.abs(x) < 0.5].sum() * dx:.0%})")

Full script: forward_reverse_kl.py

forward KL(p || q): best q has mean +0.00, width 2.10; mass in the gap |x| < 0.5: 19% (teacher: 1%)
reverse KL(q || p): best q has mean -2.00, width 0.60; mass in the gap |x| < 0.5: 1% (teacher: 1%)

Forward KL covers both modes: the student spreads out over both bumps and puts 19% of its probability in the gap between them, where the teacher has 1%. Reverse KL seeks one mode: the student sits exactly on one bump and ignores the other (which one is arbitrary here; the target is symmetric). For a language model the gap is text the teacher would never write, so a forward-KL student that is too small can produce it, and a reverse-KL student plays safe and less varied. MiniLLM (Gu et al., 2024) argued for reverse KL when distilling generative language models for that reason.

-4 -2 0 2 4 gap forward KL(p ‖ q): covers 19% of q in the gap -4 -2 0 2 4 gap reverse KL(q ‖ p): seeks 1% of q in the gap teacher p student q, best fit

The same too-small student, fitted two ways. Forward KL spreads it over both modes, including the gap where the teacher has almost nothing. Reverse KL locks it onto one mode.

Data: the grid search above

Distilling sequences, and on-policy distillation

Matching a distribution at every position needs the teacher’s logits, so teacher and student must share a vocabulary, and the teacher must be run on every training sequence. One variant relaxes that, and another fixes a deeper problem:

  • Sequence-level distillation (Kim and Rush, 2016): let the teacher generate text, and train the student on those outputs with ordinary hard labels. No logits are needed, so it works across vocabularies, and through an API. DeepSeek-R1’s small distilled models were trained this way: about 800,000 samples, some 600,000 reasoning traces from the R1 pipeline and 200,000 other examples, used for plain supervised fine-tuning, with no reinforcement learning, of Qwen2.5 and Llama models from 1.5B to 70B parameters (see Reasoning Models and Test-Time Compute).
  • On-policy distillation (Agarwal et al., 2024, generalized KD): in both setups above the student trains on fixed text, the dataset’s or the teacher’s, but at inference it reads its own text, mistakes included, which it never saw in training. Instead, let the student generate, and have the teacher score each of the student’s tokens with its distribution. The student learns what to do in the states it actually reaches.
logit-level Hinton et al., 2015 dataset text teacher scores every token student matches p needs the teacher’s logits, same vocabulary sequence-level Kim and Rush, 2016 teacher writes text that text is the dataset student trains on it works through an API; DeepSeek-R1’s small models on-policy Agarwal et al., 2024 student writes text teacher scores its tokens student matches p there trains where the student will actually be

Three ways to distill, which differ in who writes the text the student trains on and what the teacher contributes to it.

An experiment: distilling a 1.5B model into 0.5B

Fine-Tuning, and What It Destroys’s fine-tune, with a teacher. The student is Qwen2.5-0.5B (base); the teacher is Qwen2.5-1.5B-Instruct, the same family with the same tokenizer, three times the size. It reads every training batch once, and its top 64 next-token log-probabilities at every position are saved. Storing all 151,665 in 16 bits would take 300 KB per position, 250 GB for the 819,200 training positions; the top 64 take 0.4 KB and hold 92.5% of the teacher’s probability on average. The student then trains on the same 400 batches three ways: on the text’s tokens (hard labels, exactly as in the run in Fine-Tuning, and What It Destroys), on the teacher’s top-64 distribution renormalized, at temperature 1 (soft labels), and on the average of the two losses. The soft target is therefore a truncated teacher: the loss also pushes the student’s probability outside those 64 tokens toward zero, where the real teacher keeps 7.5%. And the teacher is an instruction-tuned model reading raw text, chosen because it was at hand; a base 1.5B model would be the cleaner teacher for raw text. As in Fine-Tuning, and What It Destroys, the student’s embedding table is frozen to fit training in 16 GB, on an Apple M4 laptop. The blog text, blog_train.txt and blog_val.txt, is the same as there: written by the data script of Sequences Before Attention, run on the posts as they stood on 2026-10-04. From common.py, the batches and the bits-per-byte measure:

def batches(seed):
    """The STEPS training batches for a seed, fixed, so teacher and student see the same ones."""
    g = torch.Generator().manual_seed(seed)
    out = []
    for _ in range(STEPS):
        starts = [torch.randint(0, len(blog_ids) - T - 1, (1,), generator=g).item() for _ in range(B)]
        out.append(torch.stack([blog_ids[i:i + T + 1] for i in starts]))
    return out

@torch.no_grad()
def bits_per_byte(model, data, text):
    model.eval()
    total = 0.0
    for s in range(0, len(data) - 1, T):
        x = data[s:s + T + 1].to(DEVICE)
        with torch.autocast(DEVICE, dtype=torch.bfloat16):
            logits = model(x[None, :-1]).logits[0].float()
        total += F.cross_entropy(logits, x[1:], reduction="sum").item()
    model.train()
    return total / math.log(2) / len(text.encode())

Full script: common.py

common.py also downloads the two models and WikiText-2 at pinned revisions, reads the blog text, and defines gap_to_teacher, the distance to the teacher used further down. teacher.py saves the teacher’s top 64 for one seed’s batches:

import sys
import torch
from transformers import AutoModelForCausalLM
from common import TEACHER, DEVICE, TOPK, batches, bits_per_byte, evals

# The teacher, Qwen2.5-1.5B-Instruct, reads every training batch once; its
# top-64 next-token log-probabilities at every position are saved, which is
# what the student is distilled from. Also: the teacher's own bits per byte.
seed = int(sys.argv[1])
model = AutoModelForCausalLM.from_pretrained(TEACHER, dtype=torch.float16).to(DEVICE).eval()
top_lp, top_id, kept = [], [], []
with torch.no_grad():
    for x in batches(seed):
        lp = model(x[:, :-1].to(DEVICE)).logits.float().log_softmax(-1)
        v, i = lp.topk(TOPK, -1)
        top_lp.append(v.half().cpu()); top_id.append(i.int().cpu())
        kept.append(v.exp().sum(-1).mean().item())
torch.save({"lp": torch.stack(top_lp), "id": torch.stack(top_id)}, f"teacher_s{seed}.pt")
if seed == 0:
    scores = {k: bits_per_byte(model, *v) for k, v in evals.items()}
    print(f"teacher: blog {scores['blog held-out']:.3f}, wiki {scores['WikiText-2']:.3f} bits per byte; "
          f"its top {TOPK} tokens hold {sum(kept) / len(kept):.1%} of its probability on average", flush=True)

Full script: teacher.py

distill.py trains the student on one kind of label. Its loop computes both losses on every micro-batch and steps on the chosen one:

for step, x in enumerate(batches(seed) + [None]):
    if step in checkpoints:
        log[step] = {k: bits_per_byte(model, *v) for k, v in evals.items()}
        if trace:
            log[step]["gap"] = gap_to_teacher(model, tv)
    if x is None:
        break
    x = x.to(DEVICE)
    for group in opt.param_groups:
        group["lr"] = lr * min(1.0, (step + 1) / 20)
    opt.zero_grad(set_to_none=True)
    for rows in (slice(0, B // 2), slice(B // 2, B)):               # two micro-batches: half the memory
        with torch.autocast(DEVICE, dtype=torch.bfloat16):
            logits = model(x[rows, :-1]).logits
        lp = logits.float().log_softmax(-1)
        hard = F.nll_loss(lp.reshape(-1, lp.shape[-1]), x[rows, 1:].reshape(-1))
        p = teacher["lp"][step][rows].float().to(DEVICE).softmax(-1)  # top-64, renormalized
        soft = -(p * lp.gather(-1, teacher["id"][step][rows].long().to(DEVICE))).sum(-1).mean()
        loss = {"hard": hard, "soft": soft, "mix": 0.5 * (hard + soft)}[mode] / 2
        loss.backward()
        del logits, lp, loss
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()

Full script: distill.py

three_labels.py runs teacher.py for seeds 0 and 1, then distill.py for each kind of label and seed, and prints the table from their summary lines:

Full script: three_labels.py

teacher: blog 1.038, wiki 0.873 bits per byte; its top 64 tokens hold 92.5% of its probability on average
student before training: blog 1.163, wiki 0.981 bits per byte

student, bits per byte after each step, seed 0 then seed 1
labels          step 50       step 100      step 200      step 400
hard      blog  1.086 1.092   1.072 1.072   1.069 1.065   1.068 1.084
hard      wiki  1.002 1.004   1.012 1.011   1.031 1.026   1.047 1.050
soft      blog  1.166 1.165   1.164 1.163   1.161 1.163   1.156 1.161
soft      wiki  0.991 0.991   0.995 0.993   0.998 0.999   1.003 1.005
half each blog  1.096 1.098   1.081 1.083   1.070 1.070   1.055 1.059
half each wiki  0.992 0.991   0.995 0.993   1.002 1.000   1.008 1.007
hard labels soft labels half each 1.05 1.1 1.15 0 200 400 held-out blog, bits/byte lower is better teacher 1.038 0.98 1 1.02 1.04 1.06 0 200 400 WikiText-2, bits/byte lower: less forgetting 0.3 0.4 0.5 0.6 0.7 0.8 0 200 400 KL to the teacher (nats) seed 0 training step

The experiment at a glance. Mixing learns the blog as well as hard labels, forgets far less, and stays much closer to the teacher (right). Soft labels alone barely move the student at all.

Data: the runs in this post, means of two seeds (KL panel: seed 0)

Half and half beat hard labels on both counts. At step 400 the mixed student is at 1.055 and 1.059 bits per byte on the held-out blog, against 1.068 and 1.084 from hard labels alone, and it kept improving to the end while one hard-label seed had started getting worse. On WikiText, which neither trains on, the mixed student lost 0.027 bits per byte and the hard-label one 0.068 (means of the two seeds): the teacher’s distribution, a model’s view of general English, held the student near what it already knew while the hard labels pulled it toward the blog, the same job replay does in Fine-Tuning, and What It Destroys, without needing any general text.

Soft labels alone taught it almost nothing about the blog. 1.163 became 1.156 and 1.161 after 400 steps, against 1.068 from the text’s own tokens, although the teacher, at 1.038, is better on these posts than any student here.

Was the soft loss doing anything? Seed 0 again, logging how far each student is from the teacher on the held-out posts: the KL divergence over the teacher’s top 64 tokens, both distributions renormalized to them, the part of the soft loss that measures disagreement within those tokens. The teacher’s top 64 on those windows, teacher_val.pt, come from a short script, teacher_val.py, which is not shown here. kl_trace.py runs it, then trains seed 0 again with each kind of label and prints the table:

Full script: kl_trace.py

seed 0, held-out blog: bits per byte / KL to the teacher over its top 64 (nats)
labels             step 0        step 50       step 100       step 200       step 400
soft        1.163 / 0.367  1.166 / 0.386  1.164 / 0.390  1.161 / 0.378  1.156 / 0.371
hard        1.163 / 0.367  1.086 / 0.505  1.072 / 0.578  1.069 / 0.662  1.068 / 0.781
half each   1.163 / 0.367  1.096 / 0.401  1.081 / 0.413  1.070 / 0.418  1.055 / 0.441

Soft labels kept the student where it started, 0.37 nats from the teacher. Hard labels moved it steadily away, to 0.78: learning the blog from its own tokens makes the student less like a teacher that never read it. The mixed student moved only to 0.44, much closer to the teacher, and forgot less. Staying near the teacher isn’t the whole story, though: the soft-only student never moved away from it and still lost 0.02 bits per byte on WikiText. The truncated target is a likely cause: it puts zero on the 7.5% of probability the teacher keeps outside its top 64, so it pushes the student’s own tail toward zero, and the student gets worse wherever the real next token falls there. That is about what the mixed student lost too, so its remaining forgetting is roughly the soft term’s own cost. Keeping the tail as a 65th “everything else” token, so the target stays a proper distribution, is the usual fix.

But soft labels alone never brought the student closer. To check that the loss was being optimized at all, the first 100 steps again, logging the training loss in its two parts: the KL on the 64 tokens, and how much of the student’s probability falls outside them:

    for rows in (slice(0, B // 2), slice(B // 2, B)):
        with torch.autocast(DEVICE, dtype=torch.bfloat16):
            logits = model(x[rows, :-1]).logits
        lq = logits.float().log_softmax(-1).gather(-1, teacher["id"][step][rows].long().to(DEVICE))
        lp = teacher["lp"][step][rows].float().to(DEVICE).log_softmax(-1)
        mass = lq.logsumexp(-1)                                         # log of the student's mass on the 64
        kl = (lp.exp() * (lp - (lq - mass[..., None]))).sum(-1).mean()
        soft = -(lp.exp() * lq).sum(-1).mean()
        (soft / 2).backward()
        acc.append((kl.item(), -mass.mean().item()))
        del logits, lq

Full script: soft_check.py

soft_check_both.py runs it at 10−510^{-5} and at 10−410^{-4}.

Full script: soft_check_both.py

lr 1e-05: held-out KL to teacher 0.367
  steps   1- 10: training KL 0.323, -log mass on top 64 0.150
  steps  11- 20: training KL 0.337, -log mass on top 64 0.091
  steps  21- 30: training KL 0.344, -log mass on top 64 0.091
  steps  31- 40: training KL 0.358, -log mass on top 64 0.105
  steps  41- 50: training KL 0.323, -log mass on top 64 0.084
  steps  51- 60: training KL 0.323, -log mass on top 64 0.087
  steps  61- 70: training KL 0.317, -log mass on top 64 0.090
  steps  71- 80: training KL 0.327, -log mass on top 64 0.095
  steps  81- 90: training KL 0.337, -log mass on top 64 0.090
  steps  91-100: training KL 0.296, -log mass on top 64 0.080
  after 100 steps: held-out KL to teacher 0.390
lr 0.0001: held-out KL to teacher 0.367
  steps   1- 10: training KL 0.367, -log mass on top 64 0.111
  steps  11- 20: training KL 0.439, -log mass on top 64 0.100
  steps  21- 30: training KL 0.579, -log mass on top 64 0.134
  steps  31- 40: training KL 0.600, -log mass on top 64 0.153
  steps  41- 50: training KL 0.563, -log mass on top 64 0.133
  steps  51- 60: training KL 0.616, -log mass on top 64 0.141
  steps  61- 70: training KL 0.587, -log mass on top 64 0.149
  steps  71- 80: training KL 0.606, -log mass on top 64 0.158
  steps  81- 90: training KL 0.608, -log mass on top 64 0.157
  steps  91-100: training KL 0.566, -log mass on top 64 0.135
  after 100 steps: held-out KL to teacher 0.678

At 10−510^{-5} the loss does fall, but only by pulling the student’s probability onto the teacher’s 64 tokens (the second number falls by 40% within 20 steps). The KL over those tokens doesn’t fall in training, and on held-out text it rises slightly, 0.367 to 0.390 by step 100. At 10−410^{-4}, ten times the step, the training loss itself rises, as the fine-tune in Fine-Tuning, and What It Destroys got worse at that rate: the step is too large.

The leading explanation is that there was little here for this student to learn from this teacher. The teacher’s edge over the student is about the same on the blog as on WikiText, 0.125 bits per byte against 0.108: it knows no more about the blog than about English in general, having never read it. So its 0.37 nats of disagreement with the student are general ability, spread thinly over every position, the difference between a 0.5B and a 1.5B model, which 400 batches of 2,048 tokens don’t close; real distillation runs use billions of tokens. The blog-specific vocabulary and phrasing that hard labels teach in a hundred steps aren’t in the teacher’s distribution. This wasn’t tested against the alternatives: learning rates between 10−510^{-5} and 10−410^{-4}, an untruncated target, and a base-model teacher, since some of the 0.37 nats may be the instruction tuning rather than the size. Distillation is used where the teacher knows what the student should learn, a teacher trained on the task or far stronger than the student at it; here its distribution mostly pulled the student back toward general text.

Where this leads

Takeaways

  • A teacher’s next-token distribution says which other tokens were reasonable, at every position; training a student to match it is token-level distillation, and needs the teacher’s logits and a shared vocabulary.
  • Temperature softens both distributions so the small probabilities count; multiplying the soft loss by τ2\tau^2 keeps its gradient from vanishing as τ\tau grows.
  • Forward KL makes a small student cover every mode of the teacher; reverse KL makes it pick one. Sequence-level distillation needs only the teacher’s text, and on-policy distillation trains the student on its own outputs, where it will actually be used.
  • Here, half hard labels and half the teacher’s distribution beat hard labels alone on the blog and forgot less than half as much, the teacher’s distribution acting like replay.
  • Soft labels alone barely changed the student, probably because the teacher’s edge was general ability, the same on the blog as elsewhere, which 400 small steps don’t transfer; truncating the teacher to its top 64 tokens also cost some general text. A teacher helps when it knows what the student needs to learn.

References

Start with Hinton et al., which is short and readable and introduces soft targets, temperature and the τ2\tau^2 factor. Then read Kim and Rush for distilling whole sequences, and Agarwal et al. and Gu et al. for the KL direction and on-policy training.


Written by

Tan Le

AI engineer. I write these notes to learn LLM internals and AI engineering in depth, and to keep what I learn where I can find it again.