[LLM 7/10] Model Distillation: The Teacher's Wrong Answers Are the Most Valuable Part

In chapter 6 we distilled a context into the weights of the same model. In this chapter we distill an entire model into a smaller one. The two chapters are a deliberate pair: context distillation changes what the model knows without having to tell it anymore — model distillation changes the size of the model while trying not to change what it can do. And the heart of this chapter cuts hard against intuition: the most valuable thing a teacher can hand its student is not the right answer, but the way the teacher is wrong — and the dial called temperature is what makes that visible.

1. The Problem​

Suppose you've walked the full six chapters and arrived at a Qwen3-1.7B that handles your organization's Thai workload satisfyingly well. Then one day the infrastructure team asks the one question the model can't answer: "What's the GPU bill per month?"

The 1.7B model eats nearly 3× the VRAM of the 0.6B and answers about 2.5× slower. At 100 requests per second, that difference is not a detail — it's the number of extra cards you buy every month for the life of the system.

The question is what the last row knows that the SFT row doesn't — same training, same data, so where's the difference?

The answer is in the amount of information per token. A hard label is a one-hot vector: it says "the answer is 3," full stop. But the teacher's distribution, softened by temperature, says "the answer is 3 — but 8 is somewhat plausible too, and 'cat' is complete nonsense." This ranking over all the wrong answers is what Hinton called dark knowledge. It encodes the similarity structure of the world (the digit 3 is closer to 8 than to a cat) and it vanishes entirely the moment you keep only the argmax.

At every token position the teacher has 151,936 numbers to offer (Qwen3's vocab size) — a hard label keeps exactly one.

2. What We're Going to Do​

We take Qwen/Qwen3-1.7B as the teacher and Qwen/Qwen3-0.6B-Base as the student, then distill at three levels that pass increasingly fine-grained information from the teacher:

1. SeqKD — have the teacher generate answers, then SFT the student on them (a sample from the distribution)
2. Logit KD — have the student imitate the teacher's entire distribution, token position by token position (the distribution itself)
3. GKD (optional extra) — have the student sample its own answers and let the teacher score the distributions on-policy

And the piece we cannot do without is the control row: SFT the student on the gold answers, same data, same number of steps. Without that row, we have no way to tell whether the gains came from "the teacher's distribution" or from "just training more."

This is not a laboratory curiosity — it's how most of the world's "small but capable" models are actually built. The Qwen3-0.6B we've used all series long was itself trained with strong-to-weak distillation from the larger members of its own family. In this chapter we're doing the same thing, at a scale free Colab can handle.

3. The Equations​

3.1 Hinton's KD loss​

• $z_S, z_T$ = the student's and teacher's logits, at the same token position, on the same input
• $y$ = the gold answer (hard label) — the first term is ordinary cross-entropy, identical to SFT
• $T$ = temperature, divided into the logits on both sides before the softmax
• $\alpha$ = the weight of the soft term (we use 0.9 — listen mostly to the teacher, with the gold answer as a safety line)

Notice the direction of the KL: the teacher comes first. This is forward KL, which forces the student to spread its probability over everywhere the teacher puts weight. Hold on to that — equation 3.4 is about to turn it into "one point on a line."

Now, where does the $T^2$ sitting in front of the KL come from? Nearly every piece of KD code on the internet has this factor, but very few explain why — and if you don't understand it, you will tune T wrongly without ever knowing.

3.2 Deriving where $T^2$ comes from — the two lines that separate "understood" from "copied"​

Line 1 — the gradient of the soft term with respect to one student logit is the standard softmax-CE gradient plus the chain rule through $z/T$, which emits one factor of $1/T$:

Line 2 — as $T$ grows, the softmax flattens toward uniform: $\text{softmax}(z/T)_i \approx \tfrac{1}{K} + \tfrac{z_i - \bar z}{KT}$ (with $K$ = vocab size). So the difference $q^{(T)}_i - p^{(T)}_i$ shrinks by another factor of $1/T$:

The soft term's gradient scales as $1/T^2$, while the hard CE term doesn't depend on $T$ at all. Without multiplying the $T^2$ back in, moving T from 1 to 4 secretly divides the soft term's learning rate by ~16. You'd conclude "high T doesn't work," when in reality you just quietly switched off your own soft loss. Multiplying by $T^2$ makes the gradient scale nearly independent of T — so the $\alpha$ you tuned keeps the same meaning at every T.

A bonus from line 2: in the same limit, the soft loss reduces to matching mean-centered logits in an MSE-like sense — KD is "soft logit regression" that weights the head of the distribution more than the tail.

3.3 SeqKD — the cheap baseline (Kim & Rush, 2016)​

Look familiar? This is plain SFT on answers the teacher generated — nothing more. Instead of sending the whole distribution, the teacher sends "one sample" drawn from its own distribution.

An advantage that often goes unnoticed: SeqKD doesn't care whether the tokenizers match, because it sends text, not logits. Many open-source models advertised as "distilled from GPT-4" are in fact pure SeqKD — collect the teacher's answers through an API, then SFT. That's why it's the baseline we must measure before reaching for anything more expensive.

3.4 GKD and generalized JSD — the line connecting this chapter to chapter 6​

Logit KD per equation 3.1 has a structural weakness: the student learns on sentences someone else wrote (teacher forcing), but at serving time it must generate continuations of its own answers — the accumulating mismatch is called exposure bias. GKD (Agarwal et al., 2023) fixes this by letting the student sample its own answers and having the teacher score those tokens, while also generalizing the distance between distributions to:

with $P$ = teacher, $Q$ = student. The value $\beta$ sweeps from one pole to the other:

• $\beta \to 0$ → forward KL $\mathbb{D}_{\text{KL}}(P\|Q)$ — mass-covering: the student must spread over every mode of the teacher
• $\beta \to 1$ → reverse KL $\mathbb{D}_{\text{KL}}(Q\|P)$ — mode-seeking: the student commits to the modes it can actually handle

To say it as plainly as possible: the reverse KL that chapter 6 chose isn't an oddity from another world — it's the point $\beta = 1$ on this very line, and Hinton's classic KD is the point $\beta = 0$. The two chapters are members of the same family, differing only in "who has to move toward whom" — a student much smaller than its teacher usually benefits from the mode-seeking side, because it never had the capacity to cover every one of the teacher's modes anyway.

4. Seeing the Equations​

Temperature reveals dark knowledge​

The point to read off this figure: temperature adds no information whatsoever — the logits are the exact same set. It only changes how much of the information already there can be seen. At T = 1, the best wrong answer has probability just 0.015 — the gradient flowing through it is essentially zero. At T = 8, the whole ranking becomes a signal the student can genuinely learn from, while "cat" and "!" stay on the floor where they belong.

The missing $T^2$ factor, visible in a single plot​

This is equation 3.2 in visible form: the red line is what happens if you forget $T^2$ — when you sweep for the best T, you are unintentionally sweeping the soft loss's learning rate at the same time. Your whole results table becomes unreadable, because two variables are tangled together. The green line is the only reason we can tune T cleanly.

How differently the teacher and student see the same position​

Have a look at the real thing: the teacher's top-5 versus the student's at a Thai token position from the notebook (the "before" view is the student, the "after" view is the teacher — the gap between the two views is exactly what logit KD tries to close):

5. Setting Up the Environment​

Open Colab and pick Runtime → Change runtime type → T4 GPU (the free tier is enough, but this chapter is unusually tight on VRAM because the teacher and student must sit on the card at the same time).

torch_dtype=torch.float16      # not bfloat16
attn_implementation="sdpa"     # not flash_attention_2
fp16=True                      # in TrainingArguments (not bf16=True)

cap = torch.cuda.get_device_capability(0)
print("compute capability:", cap)                    # T4 = (7, 5)
print("native bf16:", cap[0] >= 8)                   # T4 -> False
print("torch says   :", torch.cuda.is_bf16_supported())  # T4 -> True (counts emulation!)

Loading two models at once — the tightest VRAM budget in the series​

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

teacher = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-1.7B",                # teacher: instruct, ~3.4 GB (fp16)
    torch_dtype=torch.float16,
    attn_implementation="sdpa",
).cuda().eval()                       # always .eval() — the teacher learns nothing

student = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B-Base",           # student: base, ~1.2 GB (fp16)
    torch_dtype=torch.float16,
    attn_implementation="sdpa",
).cuda()

Two sets of weights totaling ~4.6 GB sounds comfortable, but during training you must budget for the activations of both models, plus the LoRA optimizer, plus temporary logits — so on a 16 GB T4 there's room for a batch size of only 2, with gradient accumulation making up the difference. That's the price of having the teacher sitting on the card with you. (Section 6 shows how to "kick the teacher off the card" by precomputing logits offline.)

The assert cell that must come before everything else​

tok_t = AutoTokenizer.from_pretrained("Qwen/Qwen3-1.7B")
tok   = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B-Base")

assert teacher.config.vocab_size == student.config.vocab_size == 151_936
assert tok_t.get_vocab() == tok.get_vocab()
print("vocab matches slot for slot — logit KD is possible")

6. Preparing the Data​

We use 3,000 prompts from airesearch/wangchanx-seed-free-synthetic-instruct-thai-120k, a Thai instruction dataset that ships with complete reference answers:

from datasets import load_dataset

ds = load_dataset("airesearch/wangchanx-seed-free-synthetic-instruct-thai-120k",
                  split="train")
ds = ds.shuffle(seed=42).select(range(3000))

From this single pool we build three ingredients:

Ingredient 1 — the gold answers, used straight from the dataset. They serve both as the control row's (SFT) data and as the sentences logit KD trains on — the two rows deliberately see character-for-character identical text, so the only remaining difference is "does it have the teacher's distribution or not."

Ingredient 2 — the teacher's answers, for SeqKD: have the teacher generate in batches of 16 prompts (max_new_tokens=192, do_sample=False). Takes about 20–25 minutes, done once and saved to disk.

Ingredient 3 — the teacher's top-64 logits, for logit KD: forward the teacher over the sentences of ingredient 1 and keep only the top 64 entries at each position.

Why top-64 — because storing the full vocab is genuinely impossible. Do the arithmetic:

K = 64

@torch.no_grad()                       # the most important line in this cell — see trap 3 in section 9
def teacher_topk(input_ids, attention_mask):
    z = teacher(input_ids=input_ids,
                attention_mask=attention_mask).logits      # [B, L, 151936]
    val, idx = z.topk(K, dim=-1)                           # [B, L, 64]
    return val.half().cpu(), idx.int().cpu()

A small detail that teaches a lot: the indices must be int32, because a vocab of 151,936 overshoots the ceiling of uint16 (65,535) by more than double — so the disk spent on indices is actually larger than the logit values themselves (4 bytes versus 2). The whole set, 3,000 examples × 512 positions × 64 ranks, comes to ≈ 590 MB on disk.

7. The Main Code​

7.1 LoRA on the student + the same fp16 trap from chapter 1​

from peft import LoraConfig, get_peft_model

student = get_peft_model(student, LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05, task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
))

for p in student.parameters():        # the fp16 box from chapter 1, LoRA edition:
    if p.requires_grad:               # cast only the adapter parameters to fp32
        p.data = p.data.float()       # or you'll hit "Attempting to unscale FP16 gradients."

7.2 Writing the KD loss by hand — equation 3.1 line by line​

import torch.nn.functional as F

def kd_loss(z_s, t_val, t_idx, labels, T=2.0, alpha=0.9):
    """z_s: [B, L, V] student logits — t_val/t_idx: [B, L, 64] teacher top-64 (read from disk)"""
    z_s, t_val, t_idx = z_s[:, :-1], t_val[:, :-1], t_idx[:, :-1]
    tgt = labels[:, 1:]                          # position t predicts the token at t+1
    mask = tgt.ne(-100)                          # keep prompt and padding out of the loss

    # ── hard term: cross-entropy against the gold answer, identical to SFT ──
    ce = F.cross_entropy(z_s.flatten(0, 1).float(), tgt.flatten(),
                         ignore_index=-100)

    # ── soft term: KL(teacher ‖ student) on the top-64 axis — divide by T on BOTH sides ──
    p_t   = F.softmax(t_val.float() / T, dim=-1)              # renormalize over 64 dims
    log_q = F.log_softmax(z_s.gather(-1, t_idx.long()).float() / T, dim=-1)
    kl = (p_t * (p_t.clamp_min(1e-9).log() - log_q)).sum(-1)  # KL per position
    kl = (kl * mask).sum() / mask.sum().clamp_min(1)          # average over answer tokens only

    return (1 - alpha) * ce + alpha * (T ** 2) * kl           # ← the T² from equation 3.2

The entire article is compressed into that one final line: (1 - alpha) * ce is the gold answer keeping the student from drifting; alpha * (T ** 2) * kl is the teacher's dark knowledge, complete with the factor we just derived with our own hands.

7.3 A Trainer that feeds teacher logits from disk​

from transformers import TrainingArguments, Trainer

class KDTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        t_val = inputs.pop("teacher_val")     # from disk — not from a teacher on the card
        t_idx = inputs.pop("teacher_idx")
        out = model(input_ids=inputs["input_ids"],
                    attention_mask=inputs["attention_mask"])
        loss = kd_loss(out.logits, t_val, t_idx, inputs["labels"])
        return (loss, out) if return_outputs else loss

args = TrainingArguments(
    output_dir="kd-out",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,        # effective batch = 16
    num_train_epochs=1,
    learning_rate=1e-4,                   # LoRA on the student — tuning only the adapter
    lr_scheduler_type="cosine",
    warmup_ratio=0.05,
    gradient_checkpointing=True,
    fp16=True,                            # the T4 has no bf16
    logging_steps=10,
    remove_unused_columns=False,          # ← do not forget, or the Trainer silently
)                                         #    throws away teacher_val/teacher_idx

remove_unused_columns=False is the line people miss most often in this chapter: by default, Trainer drops any column the model's signature doesn't recognize — which includes our teacher logits. The symptom is KeyError: 'teacher_val' at the first step. Fortunately it fails loudly, not silently.

SeqKD needs nothing new at all — a plain Trainer on ingredient 2 (the teacher's answers), trained like ordinary SFT. And the control row is that same Trainer on the gold answers. Total training time for the two main regimes (SeqKD + logit KD) is around ~17 minutes on a T4 (8–9 minutes per row); the control row costs another ~8 minutes.

7.4 Optional extra (feel free to skip the whole section): on-policy GKD with TRL​

# Run only if you have time left — on-policy is several times slower than offline,
# because the student must generate during training and the teacher must sit on the card the whole time
from trl import GKDConfig, GKDTrainer

cfg = GKDConfig(
    output_dir="gkd-out",
    beta=0.5,                             # the midpoint of the JSD line in equation 3.4
    lmbda=0.5,                            # half of each batch uses answers the student sampled itself
    max_new_tokens=128,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,
    learning_rate=1e-4,
    max_steps=60,                         # just a taste (~25 min), not a real training run
    fp16=True,
)
trainer = GKDTrainer(model=student, teacher_model=teacher,
                     args=cfg, train_dataset=gkd_ds, processing_class=tok)
trainer.train()

8. Results​

The notebook measures four things and writes them to results.json:

1. TH-INSTR from the KobEval-TH benchmark — Thai instruction-following scores, before/after, with a Wilson 95% CI, measured for the teacher, the raw student, and all three student rows, on the same exam.
2. Gap closed — the single number that summarizes the whole chapter:

Why not report raw scores — because "the score went up 4 points" means nothing without knowing how many points of gap there were to close. This metric answers our actual question: what percentage of the student-teacher gap has been closed? 0% is no movement; 100% is catching the teacher exactly. 3. The student's tok/sec, before and after distillation — which should be exactly identical, because neither the architecture nor a single parameter count changed. And that is the point of the whole chapter: quality moves toward the teacher while latency doesn't move at all — if you want one sentence to take back to your team, it's "you get (some of) the teacher's quality at the student's price." 4. A TH-SAFE spot check — the teacher transmits everything through its distribution, including its biases and bad habits. So we measure the distilled student on a Thai safety question set and report honestly what it picked up along the way (details in the limitations box at the end of the chapter).

The actual numbers in the next table are left as ? — they must come from running your own notebook, not from this article.

9. Comparison​

The pattern you should expect: logit KD ≻ SeqKD ≻ SFT control ≻ base, with the bottom four rows identical on tok/sec. If you see something else, read it this way:

• The control row does about as well as logit KD → the KD signal added nothing on this task. Try a higher T (the dark knowledge is still being squeezed flat), or raise α, or the teacher and student are simply too close together.
• SeqKD beats logit KD → genuinely possible when the dataset's gold answers are written worse than the teacher's answers (our logit KD trains on the gold answers) — this is information, not failure. Report it straight.
• No row moves much at all → 3,000 examples may be too few for this particular teacher-student gap. Look at the loss curve before concluding anything, and read the limitations box at the end of the chapter.

Traps to watch for​

1. Forgetting the $T^2$ factor No error appears anywhere. Your T sweep simply becomes fiction, because every time you move T you secretly move the soft loss's learning rate along with it (figure 7.2) — the quietest bug in this chapter.

2. Applying temperature on one side only Writing softmax(z_t / T) but forgetting to divide the student's side by T — the student gets forced to imitate the teacher's flattened distribution using its own sharp logits. The result is that it learns to be genuinely flat, and at serving time (where there is no T) its answers come out bland and oddly dispersed. Equation 3.1 divides by T on both sides, always.

3. Not detaching the teacher's logits Forward the teacher without torch.no_grad() and autograd keeps the teacher's entire activations waiting for a backward pass that never comes — VRAM quietly balloons until an OOM that points at some other line. Our offline path is safe by construction (the logits live on disk; there's no graph to keep). That's the third reason to precompute, on top of time and memory.

4. Letting padding leak into the KL Padding positions have a teacher distribution too — and it's garbage. Fail to mask it out (the mask = tgt.ne(-100) line in section 7.2) and the mean KL gets diluted by meaningless positions. Worse, the dilution ratio varies with the sentence lengths in each batch — the loss will wobble in ways you can never trace.

5. Two models on one card = a vanished batch budget The 3.4 GB teacher sits on what used to be big-batch territory. If you OOM, reduce in this order: per_device_train_batch_size → max_length → stop keeping the teacher on the card (finish the offline precompute first, then del teacher; torch.cuda.empty_cache()) — that last step is already the structure of our notebook.

10. Summary​

• Dark knowledge lives in the teacher's wrong answers — the ranking over wrong choices encodes similarity structure that hard labels can never convey, and temperature is the dial that reveals it
• The $T^2$ factor is not a talisman — the soft loss's gradient scales as $1/T^2$; multiply it back so you can tune T without secretly changing your own learning rate
• SeqKD is SFT on teacher answers — the cheapest baseline, and the only route when tokenizers don't match
• Logit KD requires an identical vocab — always assert first, because KL compares dimension by dimension
• Top-64 is engineering, not theory — one batch of full-vocab logits is a 622 MB tensor; keeping the top 64 leaves 0.79 MB while losing under 1% of coverage
• Forward KL, reverse KL, and GKD's JSD are one line — the β dial sweeps from mass-covering (this chapter) to mode-seeking (chapter 6)
• The SFT control row is what gives the results table meaning — without it, you cannot separate "the teacher helped" from "training more helped"
• Gap closed is the metric that answers the real question — what percentage of the teacher-student gap closed, at exactly the same tok/sec

Next chapter: Guardrails — our student has just inherited its teacher's knowledge and habits in full. Next we build the fence around the model: catch dangerous inputs before they reach the model, catch dangerous outputs before they reach the user, and measure the safety-versus-usability trade-off with real numbers.

References​

1. Hinton et al. (2015). Distilling the Knowledge in a Neural Network — the original KD paper: temperature and the T² factor
2. Kim et al. (2016). Sequence-Level Knowledge Distillation — sequence-level KD -- the cheapest baseline in section 9
3. Agarwal et al. (2023). On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes — GKD: the JSD framework unifying forward and reverse KL
4. Gu et al. (2023). MiniLLM: On-Policy Distillation of Large Language Models — MiniLLM: the case for reverse KL
5. Sanh et al. (2019). DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter — the most widely deployed distillation result
6. Chay-intr et al. (2025). LLaVAC: Fine-tuning LLaVA as a Multimodal Sentiment Classifier — LLaVAC: fine-tuning a multimodal model for a Thai task

The writing, code and notebooks in this series are licensed under CC BY-NC-SA 4.0 — reuse and adapt them freely with attribution, for non-commercial purposes, and share your adaptations under the same terms. Third-party models and datasets referenced here keep their own licences.