Title: Lecture 14 — Teaching Smaller Models to Think Like Bigger Ones
Course: AIINFRA 201: Model Adaptation — Fine-Tuning & Quantization > Section: Week 14 — Knowledge Distillation and Quality-Size Tradeoff Analysis

📖 Lecture — Teaching Smaller Models to Think Like Bigger Ones

By now you've spent two weeks learning to shrink models after the fact — fine-tuning them efficiently (Week 12) and compressing their weights through quantization (Week 13). This week introduces a third lever, one that works during training rather than after: knowledge distillation (KD). Instead of taking one model and shrinking its weights, distillation trains an entirely new, smaller model — the student — to reproduce the behavior of a larger, already-trained teacher model. The promise is compelling: get most of the teacher's quality at a fraction of the teacher's inference cost.

Black-box vs. white-box distillation

The first distinction to internalize is what you're allowed to see from the teacher. If the teacher is a closed model you only access through an API (think a commercial frontier model), you can only observe its outputs — the text it generates. Training a student on those outputs is called black-box KD. It's simple and requires no access to the teacher's internals, but it's also the lowest-bandwidth form of supervision: you only get to see the one sequence the teacher happened to sample. If the teacher is open-source and you have its weights, you can do better. White-box KD lets you access the teacher's internal parameters — most importantly, its full probability distribution over the vocabulary at every generation step, not just the single token it sampled. That full distribution is a much richer training signal than a single "correct" token, because it tells the student not just what the teacher said but how confident it was and what else it considered saying.

Aspect | Black-box KD | White-box KD
Teacher access | API outputs only | Full weights and logits
Signal richness | Single sampled sequence | Full next-token probability distribution
Typical use case | Distilling from closed frontier models | Distilling from open-source teachers you host
Compute cost | Lower (no teacher forward pass needed at train time beyond generating data) | Higher (teacher forward passes during training)

Matching the student's own distribution: MiniLLM and reverse KL

A naive way to align student and teacher distributions is forward KL-divergence, KL(teacher ‖ student), which is what standard supervised fine-tuning on teacher-generated text effectively optimizes. The problem: forward KL forces the student to spread probability mass everywhere the teacher does, even over sequences the smaller student is poorly equipped to generate well. This is sometimes called "mode-covering" behavior, and for a lower-capacity student it can produce degenerate, unfocused outputs. MiniLLM proposes flipping the objective to reverse KL-divergence, KL(student ‖ teacher). Reverse KL is "mode-seeking": it rewards the student for concentrating its probability mass on a smaller set of high-quality generations rather than trying to cover the teacher's entire distribution. Because the objective is evaluated with respect to the student's own generation distribution, the training signal is better matched to what the student is actually capable of producing well. This is a subtle but important idea: the best distillation objective isn't necessarily "make the student's distribution identical to the teacher's" in the abstract — it's "make the student generate things the teacher would rate highly," which reverse KL approximates more directly for generative models.

Learning to reason, not just to answer: Distilling Step-by-Step

A second improvement targets what signal the student learns from, not just the divergence metric. Distilling Step-by-Step has the teacher produce chain-of-thought rationales alongside its answers, and trains the student on both the rationale and the final answer, typically as a multi-task objective. Instead of the student memorizing input-output pairs, it learns the reasoning process that produced them. The practical payoff is data efficiency: because the rationale carries more learning signal per example, Distilling Step-by-Step can train smaller students to a target quality bar using substantially less labeled/distillation data than answer-only distillation.

The distribution-mismatch problem and on-policy distillation

There's a structural issue with training a student purely on teacher-generated sequences (whether raw outputs or full distributions over a fixed dataset): at inference time, the student generates its own sequences, token by token, and errors can compound across a generation in ways never seen during training. This is the classic exposure bias / distribution-mismatch problem in autoregressive generation — the model is trained on one distribution (teacher-authored prefixes) but deployed on another (its own prefixes). TRL's GKDTrainer (Generalized Knowledge Distillation) addresses this directly by training on-policy: the student generates its own sequences during training, and the teacher scores those self-generated sequences to provide the distillation loss (typically a generalized Jensen-Shannon divergence, JSD, between student and teacher distributions on the student's own rollouts). A tunable parameter, often called lambda, blends this on-policy loss with a more traditional supervised loss on fixed teacher/dataset sequences. Setting lambda toward the on-policy side helps the student learn to recover gracefully from its own generation trajectory rather than only ever seeing "gold" teacher trajectories.

Why we distill at all: the economics

None of this matters unless it saves money. The core motivation for LLM distillation is inference-cost reduction — producing a materially cheaper model to serve (fewer parameters, lower latency, lower memory) that approximates the quality of a larger teacher. This directly frames the central skill of this week: learning to quantify the quality-vs-size tradeoff. It's not enough to say "the student is smaller"; you need to measure how much quality was given up per unit of size or cost saved — for example, a benchmark-accuracy delta per billion parameters removed, or per dollar of inference cost saved. Framed this way, distillation becomes a design decision with numbers attached, not an act of faith.

Correcting two common misconceptions

Misconception 1: "A distilled student will always outperform a directly-trained baseline of the same size." This is false, and it's important to internalize why. Distillation is a training signal, not magic. If you have abundant compute and abundant high-quality labeled data, supervised training directly on hard labels at that size can match or beat a distilled model — distillation mainly helps when labeled data or compute is the bottleneck, because the teacher's soft signal substitutes for data you don't have. Before you commit to a distillation pipeline, always benchmark against a same-size baseline trained directly on your available data. That comparison is precisely what this week's lab asks you to run. Misconception 2: "Distillation only transfers the correct final answer." In white-box KD especially, the student is trained on the teacher's entire output distribution — including how probability is spread across wrong or lower-quality options, not just the top choice. This is exactly why white-box KD carries richer signal than black-box KD: the "shape" of the teacher's uncertainty is itself informative. But this cuts both ways. If the teacher is biased, miscalibrated, or systematically wrong in some way, the student can silently absorb those same flaws, because it's being trained to match the teacher's whole distribution, not just to be independently correct. Never treat a distilled student as automatically inheriting only the good parts of its teacher — evaluate it critically and independently, the same way you'd evaluate the teacher itself.