LLM Knowledge Distillation: Teacher Data and Reasoning Traces

Picture a new cook joining a busy dhaba on the AB Road. The head cook has no time for lectures. So the new cook learns in two ways.

Some days he only tastes the finished dal and tries to make it match. Other days he stands beside the head cook and watches every pinch of salt and every pause before the tadka. The second way teaches far more, but only if he can stand right there.

Training a small language model from a bigger one works the same way, and the two learning styles have names.

What is knowledge distillation for LLMs?

Knowledge distillation is a training method in which a small "student" model learns to reproduce the outputs of a larger "teacher" model, so the student picks up the teacher's behaviour at a fraction of the size.

Map the dhaba to the model. Tasting the finished dal is black-box distillation: the student sees only the teacher's final text. Watching every pinch is white-box distillation: the student sees the teacher's probability for every possible next token. The head cook is the teacher, and the new cook is your small student.

In this series you have already built the student. You chose its architecture in the small LLM architecture tutorial, pretrained it and taught it to chat in supervised fine-tuning. Distillation is a way to make that SFT stage stronger.

Black-box distillation needs only text. White-box distillation needs the teacher's full output distribution.

Black-box vs white-box distillation

Most released "distilled" small models are black-box: someone saved a strong model's answers and ran SFT on them.

Why soft labels carry more information

When the teacher writes "Delhi" as the next word, a hard label says "Delhi is 100% right and everything else is 0% right". But the teacher actually believed "New" (as in New Delhi) was quite likely too. That extra belief is what white-box distillation passes on.

Temperature controls how much of that hidden belief you expose. Dividing the logits by a temperature T above 1 flattens the distribution, so the smaller probabilities become visible.

import numpy as np

def softmax(z, T=1.0):
    z = np.asarray(z, dtype=float) / T
    z = z - z.max()
    e = np.exp(z)
    return e / e.sum()

# Teacher logits for the next token after "The capital of India is"
vocab = ["Delhi", "New", "Mumbai", "the", "a"]
teacher_logits = [6.0, 5.2, 2.5, 0.5, 0.2]

for T in (1.0, 2.0, 4.0):
    p = softmax(teacher_logits, T)
    shown = ", ".join(f"{w}={x:.2f}" for w, x in zip(vocab, p))
    print(f"T={T}: {shown}")

hard = np.zeros(len(vocab)); hard[0] = 1.0
print("Hard label:", ", ".join(f"{w}={x:.0f}" for w, x in zip(vocab, hard)))

At T=1 the teacher's second choice already gets 30%. At T=4 even "Mumbai" gets noticeable weight. The hard label throws all of that away. This idea comes from the original distillation paper by Hinton, Vinyals and Dean (2015).

How the distillation loss works

The white-box loss mixes two terms. Cross-entropy keeps the student honest against the real next token. KL divergence pulls the student's softened distribution towards the teacher's.

The KL term is multiplied by T squared. Softening shrinks the gradients by roughly 1/T squared, so the multiplication keeps both terms on a similar scale when you change T. A weight alpha between 0 and 1 balances the two.

import numpy as np

def log_softmax(z, T=1.0):
    z = np.asarray(z, dtype=float) / T
    z = z - z.max(axis=-1, keepdims=True)
    return z - np.log(np.exp(z).sum(axis=-1, keepdims=True))

def distill_loss(student, teacher, target, alpha=0.5, T=2.0):
    # Hard part: cross-entropy against the real next token
    ce = -log_softmax(student)[np.arange(len(target)), target].mean()
    # Soft part: KL(teacher || student) at temperature T, scaled by T^2
    log_pt, log_ps = log_softmax(teacher, T), log_softmax(student, T)
    kl = (np.exp(log_pt) * (log_pt - log_ps)).sum(axis=-1).mean()
    return alpha * ce + (1 - alpha) * T * T * kl, ce, kl

rng = np.random.default_rng(0)
teacher = rng.normal(0, 3, size=(4, 8))          # 4 positions, vocab of 8
target = teacher.argmax(axis=-1)                 # teacher agrees with the data here
far = rng.normal(0, 3, size=(4, 8))              # student that ignores the teacher
near = teacher + rng.normal(0, 0.5, size=(4, 8)) # student close to the teacher

for name, s in (("far", far), ("near", near)):
    total, ce, kl = distill_loss(s, teacher, target)
    print(f"{name:>4} student: total={total:.3f}  ce={ce:.3f}  kl={kl:.3f}")

The student that already resembles the teacher gets a far smaller KL term. In practice, alpha around 0.5 and T between 1 and 2 are reasonable starting points. Tune them on your validation loss rather than trusting any default.

The PyTorch version

Here is a minimal training step. Both models are TinyLMs, which return (logits, loss), and the batches are the (x, y) pairs from the SFT tutorial's SFTDataset, where y is already shifted by one and holds -100 on prompt and padding tokens.

The teacher must share the same 6,400-token tokenizer, for example a wider TinyLM you trained on the same data.

import torch
import torch.nn.functional as F

def distill_step(student, teacher, x, y, optimizer, alpha=0.5, T=1.5):
    # x, y come from SFTDataset: y is already shifted by one, with -100 on prompt and padding
    with torch.no_grad():
        t_logits, _ = teacher(x)                       # teacher is frozen
    s_logits, ce = student(x, y)                       # TinyLM returns logits and masked CE
    s, t = s_logits.float(), t_logits.float()
    mask = y != -100                                   # only answer tokens

    kl_tok = F.kl_div(F.log_softmax(s / T, dim=-1),
                      F.log_softmax(t / T, dim=-1),
                      log_target=True, reduction="none").sum(-1)
    kl = (kl_tok * mask).sum() / mask.sum().clamp(min=1)

    loss = alpha * ce + (1 - alpha) * (T ** 2) * kl
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
    optimizer.step()
    return loss.item(), ce.item(), kl.item()

Start the student from your SFT checkpoint (checkpoints/sft_epoch1.pt) with a small learning rate (around 1e-6 to 1e-5). The teacher roughly doubles memory, so on an 8 to 12 GB GPU you may need a smaller batch with gradient accumulation. See the PyTorch KLDivLoss docs for the argument order, which trips up many people.

How to build reasoning data with think tags

The more popular kind of distillation today is reasoning distillation. You ask a strong teacher to show its working, then you fine-tune the student on those traces. This is black-box distillation, so the training code is your normal SFT loop.

The usual format wraps the working in tags inside the assistant message, so the chat template can show or hide it:

{"messages": [
  {"role": "user", "content": "A shop sells 3 pens for Rs 45. How much do 7 pens cost?"},
  {"role": "assistant", "content": "<think>\nOne pen costs 45 / 3 = 15 rupees. Seven pens cost 7 x 15 = 105.\n</think>\n\nSeven pens cost Rs 105."}
]}

Filtering matters more than volume. A few thousand clean traces beat a large pile of broken ones.

Follow these steps to build the dataset:

1. Collect prompts that suit a small model: arithmetic word problems, short code tasks, everyday questions. Avoid competition maths; the student cannot imitate what it cannot follow.
2. Generate with the teacher. Ask for short working inside think tags, then a clear final answer.
3. Filter mechanically. Drop unclosed tags, empty answers, very long thinking and exact duplicates.
4. Verify what you can. Recompute arithmetic with Python, or run generated code against a test.
5. Mix in direct answers. Keep some samples with an empty think block, so the student learns when not to think.
6. Run SFT with the loss mask on the assistant reply only, exactly as in the SFT tutorial.

Here is step 3 as a runnable filter:

import re

samples = [
    {"q": "What is 12 x 15?", "a": "<think>\n12 x 15 = 12 x 10 + 12 x 5 = 120 + 60 = 180.\n</think>\n\n12 x 15 = 180."},
    {"q": "What is 12 x 15?", "a": "<think>\n12 x 15 = 12 x 10 + 12 x 5 = 120 + 60 = 180.\n</think>\n\n12 x 15 = 180."},
    {"q": "Name a river in Indore.", "a": "<think>\nThe Khan river flows through Indore.\n</think>\n\nThe Khan river."},
    {"q": "Explain RAM.", "a": "<think>\nRAM is memory... RAM is memory... RAM is memory..."},
    {"q": "Say hi.", "a": "Hi! How can I help you today?"},
]

def check(s, max_think_words=120):
    a = s["a"]
    if a.count("<think>") != a.count("</think>"):
        return "unclosed think tag"
    m = re.search(r"<think>(.*?)</think>\s*(.*)", a, re.S)
    if m:
        think, answer = m.group(1).split(), m.group(2).strip()
        if len(think) > max_think_words:
            return "thinking too long"
        if not answer:
            return "no final answer"
    return "ok"

seen, kept = set(), []
for s in samples:
    key = (s["q"], s["a"])
    verdict = "duplicate" if key in seen else check(s)
    seen.add(key)
    print(f"{s['q'][:24]:<26}{verdict}")
    if verdict == "ok":
        kept.append(s)
print(f"kept {len(kept)} of {len(samples)}; with thinking: {sum('<think>' in s['a'] for s in kept)}")

For a model this size, keep thinking short. Your context window during SFT is probably 512 to 1,024 tokens, and a long trace leaves no room for the answer.

Generating teacher answers

Any OpenAI-compatible endpoint works for generation, including a larger open model you run yourself with vLLM. Read the teacher's licence and terms of use first, because some providers forbid using outputs to train other models.

import json
from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed")
SYSTEM = "Think step by step inside <think></think> in under 80 words, then give a short final answer."

with open("prompts.txt") as f, open("distill.jsonl", "w") as out:
    for prompt in (line.strip() for line in f if line.strip()):
        reply = client.chat.completions.create(
            model="teacher", temperature=0.6, max_tokens=400,
            messages=[{"role": "system", "content": SYSTEM},
                      {"role": "user", "content": prompt}],
        ).choices[0].message.content
        out.write(json.dumps({"messages": [
            {"role": "user", "content": prompt},
            {"role": "assistant", "content": reply}]}, ensure_ascii=False) + "\n")

What distillation can and cannot do for a small model

A 64M parameter student trained on reasoning traces learns the format very well: it opens a think block, writes step-like text and closes it. It often gets simple two-step arithmetic right.

It will not reason like the teacher. On harder problems the steps look plausible but the numbers drift. Its knowledge stays roughly what pretraining gave it. Distillation copies habits much better than it copies facts.

Measure the change on a fixed prompt set, comparing the SFT and distilled models, as covered in how to evaluate a small LLM.

After distillation, the next step in this series is preference tuning with the DPO training tutorial. For the theory behind these post-training stages, see SFT vs RLHF vs GRPO. If you prefer to study generative AI in a classroom, Debugshala's in-person GenAI course in Indore offers a free demo class.

Common mistakes in LLM distillation

• Mismatched tokenizers in white-box mode. The KL loss compares token i of the teacher with token i of the student. Different vocabularies make that comparison meaningless.
• Distilling into the base model. Start from your SFT checkpoint; otherwise the student must learn chat format and teacher style at once.
• Learning rate too high. A high rate erases what pretraining taught. Start low and watch validation loss on general text too.
• Unfiltered traces. One broken think tag pattern repeated a thousand times becomes the student's habit.
• Thinking always on. If every sample has long thinking, the student thinks even for "hello".

Key takeaways

• Black-box distillation is SFT on teacher outputs and is what most small "distilled" models use.
• White-box distillation adds a temperature-softened KL loss and needs shared vocabulary and teacher logits.
• Reasoning data uses think tags inside the assistant message; filter, verify and mix with direct answers.
• A tiny student copies format and habits well, but not deep knowledge or reasoning.