Distillation

How the student is scored

The student also produces logits, and they're softened with the same temperature. Training then compares two lists of chances: the teacher's and the student's. The score is the cross-entropy you met in Phase 2, with the teacher's list in place of a single right answer.

Write $q_i$ for the teacher's softened chance for option $i$. Write $p_i$ for the student's softened chance for the same option. The loss for one position is:

Here $\log$ is the natural log, as in the cross-entropy lesson. The factor $T^2$ is the temperature times itself.

In plain words: for every option, take the log of the student's chance and weight it by the teacher's chance. Add those up and flip the sign. The loss is small when the student puts high chances where the teacher does. Then multiply by T times T. In training, this is averaged over every position in the batch.

With a hard label, the teacher's list is 1 for the true token and 0 for the rest. Then the sum keeps only one term. At $T = 1$, that's ordinary cross-entropy again.

import math

teacher = [0.5, 0.4, 0.1]    # silnylon, ripstop, canvas: the softened target
student = [0.8, 0.1, 0.1]    # the student's softened guess
T = 2

hard = -math.log(student[0])                                  # the true word was silnylon
soft = -sum(q * math.log(p) for q, p in zip(teacher, student))
print(round(hard, 3), round(soft, 3), round(soft * T * T, 3))

copy = -sum(q * math.log(q) for q in teacher)                 # a student that matches the teacher
print(round(copy, 3))

Distilling a smaller GPT

Now the real experiment. The teacher is the course's tiny GPT. The student has about a quarter of the teacher's parameters. It has one layer of width 16 instead of two layers of width 32. (Width is the length of the list of numbers that stands for each token inside the model.)

The student sees only 800 tokens of catalog text, about 25 product descriptions. Two students train on that same text:

• One trains on the hard labels: the true next token.
• One trains on the teacher's soft targets, at temperature 2.

The teacher's soft targets are computed once, up front. Offline distillation usually works this way: run the teacher over the data first, store its outputs, then train the student.

Both students are then tested on descriptions they never saw. The test uses ordinary cross-entropy against the true next tokens, and lower is better. The cell takes a while to run. The Run button counts the seconds.

import requests
import torch
import torch.nn.functional as F
from tinygpt import GPT, load_pretrained

teacher, tok = load_pretrained("tiny-gpt")
data = torch.tensor(tok.encode(requests.get("https://api.gradient.test/v1/corpus").json()["text"]))
train, held_out = data[:800], data[int(0.9 * len(data)):]       # the student sees 800 tokens
V, T = len(tok), 2.0

windows = torch.stack([train[i:i + 33] for i in range(len(train) - 33)])   # every 32-token window + next token
with torch.no_grad():
    soft_targets = torch.softmax(teacher(windows[:, :-1]) / T, dim=-1)    # the teacher's view, computed once

@torch.no_grad()
def held_out_loss(model):
    torch.manual_seed(123)
    losses = []
    for _ in range(8):
        starts = torch.randint(0, len(held_out) - 33, (32,)).tolist()
        x = torch.stack([held_out[i:i + 32] for i in starts])
        y = torch.stack([held_out[i + 1:i + 33] for i in starts])
        losses.append(F.cross_entropy(model(x).reshape(-1, V), y.reshape(-1)).item())
    return sum(losses) / len(losses)

def train_student(soft, steps=300):
    torch.manual_seed(0)
    student = GPT(V, context=64, d=16, n_layers=1, n_heads=2)
    opt = torch.optim.AdamW(student.parameters(), lr=3e-3, weight_decay=0.01)
    torch.manual_seed(1)
    for step in range(steps):
        b = torch.randint(0, len(windows), (16,))
        logits = student(windows[b, :-1])
        if soft:                                              # match the teacher's softened distribution
            loss = -(soft_targets[b] * torch.log_softmax(logits / T, dim=-1)).sum(dim=-1).mean() * T * T
        else:                                                 # match the one true next token
            loss = F.cross_entropy(logits.reshape(-1, V), windows[b, 1:].reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
    return student

count = lambda m: sum(p.numel() for p in m.parameters())
hard, soft = train_student(soft=False), train_student(soft=True)
print(f"teacher: {count(teacher):,} parameters; each student: {count(hard):,}")
print(f"student trained on hard labels:    held-out loss {held_out_loss(hard):.3f}")
print(f"student trained on teacher (T=2):  held-out loss {held_out_loss(soft):.3f}")

The line starting loss = -(soft_targets[b] ... is the formula from above, written in PyTorch.

From the same 800 tokens, the distilled student does noticeably better on new text. Each example carries the teacher's view of what else could have come next, and the teacher learned that view from the whole catalog.

With plenty of data, the gap narrows. In testing, with nearly ten times as much text, the two students were nearly tied. That matches practice: distillation helps most when the student can't see everything the teacher learned from. (torch here is the course's in-browser PyTorch subset.)

Distillation for language models in practice

There are two main kinds, and they need different access to the teacher.

• Logit distillation is what the cell did. It needs the teacher's full lists of chances, and both models must use the same tokenizer. So it's used when one organization owns both models. Google, for example, trained its smaller Gemma 2 models with distillation from a larger model. DistilBERT (2019) kept about 97% of BERT's performance while being 40% smaller and 60% faster.
• Sequence-level distillation trains the student with SFT on text the teacher generated. You ask the big model thousands of questions and fine-tune the small one on its answers. This needs no access to the chances, only the teacher's text. Stanford's Alpaca (2023) did this with 52,000 answers from an OpenAI model. DeepSeek fine-tuned small Qwen and Llama models on about 800,000 samples from R1, and got strong small reasoning models.

A common business pattern works like this:

1. Prototype with a frontier model and a careful prompt.
2. Collect its outputs on real inputs, and have people review them.
3. Fine-tune a small model on the approved outputs for the high-volume job.
4. Keep evaluating the small model against the big one.

Know the terms. Many commercial providers' terms of service ban using their models' outputs to build competing models. They generally allow narrower internal uses. Open-model licenses vary too. Check the teacher's terms before you distill. Also check the license of any dataset that was itself generated by a model.

Know the limits. In general, a student can't beat its teacher on the teacher's own tasks. It copies the teacher's mistakes along with its skills. And a small model distilled on a narrow task can lose general abilities it had before.

For marketers: say a big model writes good product descriptions, but it costs too much at catalog scale. Generate a few thousand, have editors approve them, and fine-tune a small model on the approved set. Check first that your provider's terms allow it.