Knowledge Distillation, Implemented on a Tiny Transformer From Scratch

What distillation actually transfers

When you train a language model the usual way, each position has one correct next token and the loss pushes the model toward it. The teacher, however, does not just know the answer; it knows how plausible every other token is. After the word "the," a well-trained model does not place all its mass on a single adjective. It spreads probability across the adjectives it has seen in that position, and the shape of that spread is information. This is what Hinton, Vinyals, and Dean called dark knowledge in their 2015 paper, and distillation is the idea of training the student on it directly.

Concretely, the student minimizes a weighted sum of two losses. The first is the ordinary cross-entropy against the true next token. The second is the divergence between the teacher's softened distribution and the student's, both passed through a temperature. Put as Hinton wrote it, the objective is (1 - alpha) times the hard-label loss plus alpha times the soft-target loss, with the soft term scaled by the temperature squared.

Task

For our experiment, we need a task where the next token is genuinely uncertain but still checkable, because that is exactly where dark knowledge has something to say. We will use small templated language as toy example. Every sentence has the shape determiner, adjective, noun, verb, determiner, adjective, noun, followed by a period, with a short fixed vocabulary for each slot. There are 98,304 possible sentences, so a model trained on a couple of thousand of them has to learn the grammar rather than memorize the set. Whether a generated sentence is grammatical is a matter of a few lines, which gives us a clean metric.

import torch, random
random.seed(0); torch.manual_seed(0)

DET  = ["the", "a"]
ADJ  = ["red", "blue", "tall", "small", "quick", "slow", "bright", "calm"]
NOUN = ["fox", "cat", "dog", "bird", "tree", "rock", "lake", "hill"]
VERB = ["sees", "finds", "likes", "chases", "watches", "holds"]
SLOTS = [DET, ADJ, NOUN, VERB, DET, ADJ, NOUN]

def sentence():
    w = [random.choice(slot) for slot in SLOTS]
    return f"{w[0]} {w[1]} {w[2]} {w[3]} {w[4]} {w[5]} {w[6]}."

def is_grammatical(line):
    line = line.strip()
    if not line.endswith("."):
        return False
    toks = line[:-1].split()
    if len(toks) != len(SLOTS):
        return False
    return all(tok in slot for tok, slot in zip(toks, SLOTS))

# a small slice of the 98,304 possible sentences, so the model must generalize
corpus = "\n".join(sentence() for _ in range(2000))
chars  = sorted(set(corpus)); vocab_size = len(chars)
stoi   = {c: i for i, c in enumerate(chars)}
itos   = {i: c for i, c in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
data   = torch.tensor(encode(corpus), dtype=torch.long)

The reason this task suits distillation is the slot structure. At the first character of an adjective slot, the teacher's distribution is spread across the first letters of all eight adjectives. A hard label names only the one adjective that happened to appear in that training example; the teacher's soft target names the whole set and roughly how likely each is. That is the signal we want the small student to inherit.

The teacher and the student

Both models are the same character-level transformer we built in another cookbook, parameterized by size so we can instantiate a larger teacher and a much smaller student from one definition. Nothing here is specific to distillation yet; it is an ordinary decoder with token and position embeddings, masked self-attention, and a language-model head.

import torch.nn as nn
from torch.nn import functional as F

class Head(nn.Module):
    def __init__(self, n_embd, head_size, block_size):
        super().__init__()
        self.key   = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape
        k, q, v = self.key(x), self.query(x), self.value(x)
        wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5      # scale by head size
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
        wei = F.softmax(wei, dim=-1)
        return wei @ v

class MultiHead(nn.Module):
    def __init__(self, n_head, n_embd, head_size, block_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(n_embd, head_size, block_size) for _ in range(n_head)])
        self.proj  = nn.Linear(n_head * head_size, n_embd)

    def forward(self, x):
        return self.proj(torch.cat([h(x) for h in self.heads], dim=-1))

class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        head_size = n_embd // n_head
        self.sa  = MultiHead(n_head, n_embd, head_size, block_size)
        self.ff  = nn.Sequential(nn.Linear(n_embd, 4 * n_embd), nn.ReLU(),
                                 nn.Linear(4 * n_embd, n_embd))
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x

class TinyTransformer(nn.Module):
    def __init__(self, vocab_size, n_embd, n_head, n_layer, block_size):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.blocks  = nn.Sequential(*[Block(n_embd, n_head, block_size) for _ in range(n_layer)])
        self.ln_f    = nn.LayerNorm(n_embd)
        self.head    = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        x = self.tok_emb(idx) + self.pos_emb(torch.arange(T, device=idx.device))
        x = self.ln_f(self.blocks(x))
        logits = self.head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

We train the teacher with plain next-character cross-entropy. The training helper takes an optional loss function so that, in a moment, we can reuse it unchanged for the student's distillation loss.

device = "cuda" if torch.cuda.is_available() else "cpu"
block_size = 64

def get_batch(bs=64):
    ix = torch.randint(len(data) - block_size, (bs,))
    x = torch.stack([data[i:i + block_size]         for i in ix])
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
    return x.to(device), y.to(device)

def train(model, steps=3000, lr=1e-3, loss_fn=None):
    model.to(device).train()
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    for step in range(steps):
        xb, yb = get_batch()
        loss = model(xb, yb)[1] if loss_fn is None else loss_fn(model, xb, yb)
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
    return model

teacher = TinyTransformer(vocab_size, n_embd=128, n_head=4, n_layer=3, block_size=block_size)
train(teacher, steps=3000)

Temperature, and the squared factor

A well-trained teacher may be too confident, which is a problem for distillation: if it puts almost all its mass on one token, its distribution carries little more than the hard label already did. Temperature fixes this by dividing the logits before the softmax, which flattens the distribution and brings the smaller probabilities up into a range where they matter. As an example, take a teacher whose logits over four options are 2, 1, 0, and -1; raising the temperature spreads its probability mass.

The teacher still ranks the first option highest, just legibly enough that the student can learn the ordering of the others too. The temperature-squared factor handles a side effect of this softening: it also shrinks the gradient of the soft-target loss, by roughly the temperature squared. Taking the same example with a student that starts out flat at 0.25 on every option, the size of that gradient moves like this.

Multiplying the soft loss by the temperature squared puts it back on scale. Without that correction, raising the temperature would quietly turn the soft term down relative to the hard-label loss, and you would be tuning two things at once without realizing it.

def distill_loss(temp=2.0, alpha=0.5):
    teacher.eval()                                  # teacher only provides targets, never trains
    def loss_fn(student, xb, yb):
        logits_s, _ = student(xb)
        with torch.no_grad():
            logits_t, _ = teacher(xb)
        V = logits_s.size(-1)
        # soft term: match the teacher's temperature-softened distribution (forward KL)
        log_p_s = F.log_softmax(logits_s.view(-1, V) / temp, dim=-1)
        p_t     = F.softmax(logits_t.view(-1, V) / temp, dim=-1)
        soft = F.kl_div(log_p_s, p_t, reduction="batchmean") * (temp * temp)
        # hard term: the true next character, at temperature 1
        hard = F.cross_entropy(logits_s.view(-1, V), yb.view(-1))
        return alpha * soft + (1 - alpha) * hard
    return loss_fn

Training the student, with and without the teacher

The comparison is the whole point, so we train two students of identical size: one from the distillation loss above, the other from hard labels alone. Two details matter for getting a result you can trust. First, a model this small is high-variance, so a single run tells you little; we seed several runs and read the average and the spread. Second, within each seed we give the distilled student and the baseline the same initialization and the same batch order, so the only thing that differs between them is the loss.

@torch.no_grad()
def sample_sentences(model, n=500, max_len=64, greedy=False):
    model.eval()
    out = []
    for _ in range(n):
        idx = torch.tensor([[stoi["\n"]]], device=device)        # start of a fresh sentence
        for _ in range(max_len):
            logits, _ = model(idx[:, -block_size:])
            probs = F.softmax(logits[:, -1], dim=-1)
            nxt = probs.argmax(-1, keepdim=True) if greedy else torch.multinomial(probs, 1)
            idx = torch.cat([idx, nxt], dim=1)
            if itos[nxt.item()] == "\n":
                break
        out.append("".join(itos[i] for i in idx[0, 1:].tolist()))
    return out

def grammatical_rate(model, n=500, greedy=False):
    return sum(is_grammatical(s) for s in sample_sentences(model, n, greedy=greedy)) / n

@torch.no_grad()
def mean_entropy(model, n=200):                # average next-char uncertainty, in nats
    model.eval(); total = 0.0
    for _ in range(n):
        xb, _ = get_batch(1)
        p = F.softmax(model(xb)[0][0, -1], dim=-1)
        total += float(-(p * torch.log(p + 1e-9)).sum())
    return total / n

def nparams(m):
    return sum(p.numel() for p in m.parameters())

The grammatical rate uses sampling rather than greedy decoding, because greedy decoding can score well by always emitting the single most likely sentence, whereas sampling tests whether the model learned the grammar's distribution. The helpers also expose a greedy rate and a mean next-character entropy, which are useful for telling whether a distilled student has been pushed too soft to sample cleanly. With those in hand, we train both students across a few seeds and report the mean and spread.

import statistics as stats

def small():
    return TinyTransformer(vocab_size, n_embd=32, n_head=2, n_layer=1, block_size=block_size)

def set_seed(s):
    random.seed(s); torch.manual_seed(s)

def train_student(seed, loss_fn=None):
    set_seed(seed)                             # identical init + batch order across KD and base
    return train(small(), steps=3000, loss_fn=loss_fn)

seeds = [0, 1, 2, 3, 4]
kd, base = [], []
for s in seeds:
    m_kd   = train_student(s, loss_fn=distill_loss(temp=2.0, alpha=0.5))
    m_base = train_student(s)                  # same seed, so only the loss differs
    kd.append(grammatical_rate(m_kd))
    base.append(grammatical_rate(m_base))

def report(name, xs):
    runs = ", ".join(f"{x:.0%}" for x in xs)
    print(f"{name:14} mean={stats.mean(xs):.0%}  std={stats.pstdev(xs):.0%}  runs=[{runs}]")

print(f"teacher        grammatical={grammatical_rate(teacher):.0%}   params={nparams(teacher)}")
report("student (KD)",   kd)
report("student (base)", base)

# the comparison the seeding was for: same start per seed, so only the loss differs
diffs = [k - b for k, b in zip(kd, base)]
per   = ", ".join(f"{d:+.0%}" for d in diffs)
print(f"KD - base      mean={stats.mean(diffs):+.0%}  per seed=[{per}]")

Reading the results

A multi-seed run is useful to assess the difference between methods. In our case the teacher scored 96%; the distilled student averaged 75% with a standard deviation of 20points, over seeds of 92, 83, 80, 86, and 36; the baseline averaged 65% with a deviation of 23 points, over 80, 73, 80, 70, and 20. The first thing to notice is how wide that spread is. The same architecture and the same settings land anywhere from the mid-thirties to the low-nineties depending only on the seed, because a sixteen-thousand-parameter model sits right at the edge of learning this grammar. It is important to read the distribution, not a single run.

The second thing is the point the seeding was for. Because each seed gives the two students the same initialization and the same batches, with only the loss differing, the honest comparison is within a seed rather than across seeds. Done that way the distilled student is ahead in four of the five seeds and level in the fifth, by +12, +10, 0, +16, and +16 points, an average paired gain of about eleven points and never a loss. The across-seed noise is larger than that gain, which is exactly why an unpaired, single-run comparison can show distillation losing when the baseline merely drew a kinder starting point. Even the worst run, where both models collapse to36% and 20%, keeps the distilled student in front.

What governs the size of that gain is worth keeping in view, because none of it is automatic. The soft term enters the loss weighted by alpha times the temperature squared, here 2.0 against 0.5 for the hard term; push that ratio higher, or shrink the student further, and it can become unable to match the softened target while still sampling cleanly, at which point the gain narrows or reverses. Distillation also helps most when the hard labels are an impoverished signal, which is the case here because each example names one word per slot while the teacher names the whole distribution. Enlarge the student, or train on so much data that the baseline saturates the grammar on its own, and the gap shrinks for the opposite reason. If you change alpha, the temperature, the capacity gap, or the amount of data, run the seeds again before concluding anything.

Where this departs from production distillation

The teacher is the first difference. Here we trained one from scratch, so it is not much smarter than the student could be; in practice the teacher is a large pretrained model and the point is to compress capability that the student could not easily learn on its own. Its quality is also the student's ceiling, more or less, so a distilled model inherits the teacher's mistakes along with its strengths.

The tokenizer is the second, and it is a real constraint rather than a detail. Matching soft distributions only makes sense when teacher and student share a vocabulary, because the probabilities have to be over the same set of tokens at the same positions. Our two models share one character vocabulary, so this is free; distilling between models with different tokenizers, say a subword teacher into a different subword student, is an active research problem rather than a setting you can switch on.

The objective is the third. We used the classic forward-KL, token-level loss on a fixed dataset, which is the original Hinton formulation and the clearest one to learn from. For generative language models at scale, two refinements are common. Kim and Rush's sequence-level distillation trains the student on whole sequences the teacher generated, rather than matching per-token distributions on a fixed corpus. MiniLLM and related work swap the forward KL for a reverse KL, which is mode-seeking and pushes the student to concentrate its limited capacity on the teacher's most probable outputs rather than trying to cover everything. Both tend to matter more as the capacity gap and the open-endedness of the task grow.

Finally, the practical machinery is different. A real teacher's distribution is over tens or hundreds of thousands of tokens, so one may store only the top-k logits or recompute them on the fly, and the teacher's forward passes often dominate the training cost. And there is a non-technical consideration that engineers forget at their peril: distilling from a proprietary model may be restricted by its terms of service, so the legality of your transfer set is worth checking before you build on it.

Common questions

Does the student need the same tokenizer as the teacher?

For logit-based distillation like this, yes. The soft target is a probability distribution over the vocabulary at each position, and the student can only match it if its vocabulary and tokenization line up with the teacher's. When they differ, the distributions live over different token sets and there is nothing to match directly, which is why cross-tokenizer distillation needs special alignment methods. In our toy both models share one character vocabulary, so the question never arises.

What temperature should I use?

Temperature controls how much of the teacher's dark knowledge you expose; higher values flatten the distribution and surface the smaller probabilities. Hinton's experiments used values from 1 to 20, and noted that when the student is very small relative to the teacher, lower temperatures tend to work better, because a tiny model may not have the capacity to absorb a very rich distribution. A value of 2 to 4 is a reasonable place to start. Whatever you choose, keep the temperature-squared factor on the soft loss so the gradient scale does not move when you change it.

Can the student end up better than the teacher?

As a rule the teacher is the ceiling, since the student is trained to reproduce its distribution. Distillation buys efficiency, a smaller and faster model at close to the teacher's quality, rather than new capability. A student can sometimes match a same-size model trained from scratch with a clear margin, or generalize a little differently, but expecting it to exceed the teacher it learned from is the wrong mental model.