TravisML Distilling Generative Models
GitHub ↗

TravisML

Distilling Generative Models

Knowledge Distillation for Language Models, from the Objective to a Reproducible Study

Travis Lelle

Copyright TravisML.ai 2026. All rights reserved.

First edition, August 2026. Companion to the thirteen-lab course Distillation of Generative Models.

Every measured number in this book came out of an execution. Every citation was verified against its source. Errors that remain are mine.

Preface

I built the course this book accompanies because I wanted to be comfortable with distillation, and reading about it was not producing comfort.

The literature is good. Hinton’s paper is short and clear, the GKD paper is careful, and the surveys are thorough. But a paper tells you what worked in someone else’s setting, and it compresses out everything that made the result hard to obtain. It does not tell you that the temperature-squared factor is missing from a third of the implementations you will read. It does not tell you that your loss curve will descend beautifully while training on a mask that is off by one position. It does not tell you that the parameter your library calls beta is documented backwards in more places than it is documented correctly, including, for a while, in my own code.

So I built thirteen laboratory notebooks that assert every claim they make, and worked through them. This book is what I wish I had been reading alongside them.

What this book is#

It is a textbook on knowledge distillation for generative language models, written for someone who can read PyTorch and has trained a transformer, and who knows nothing beyond that. Everything else is built here: the information theory, the floating point, the serving arithmetic, the reinforcement learning vocabulary that the on-policy methods borrow. I define every term of art at its first use and I do not assume folklore.

It is organized around building blocks rather than around methods. The first part is foundations: what a token distribution is, how the numbers that hold it behave, how to measure the distance between two of them, and how to estimate that distance when you cannot see both. The second part is the objective itself. The third part makes it real, which means tokenizers, a first run, the economics, and the pipeline I would actually build first. The fourth part is the method space. The fifth is judgment: serving, evaluation, security, and how to turn a question into a study another person could rerun.

Each chapter climbs. It starts at a level that presumes nothing and ends at a level where the open questions are visible. That means some chapters are long. The instruction I gave myself was not to be stingy with explanation, and I have kept to it.

What this book is not#

It is not a survey. Gou and colleagues wrote a good one for the general case and Xu and colleagues wrote a good one for language models, and both are cited where they belong. This book covers fewer methods in far more depth, and it is opinionated about which ones are worth your time on a single workstation.

It is not a reproduction of the labs. The labs execute; the book explains and derives. Where the same fact appears in both, the lab asserts it mechanically and the book says why it is true. I have tried hard not to duplicate a single paragraph between the two.

And it is not neutral about hardware. Every number here that depends on a machine refers to one machine: a workstation with 128 GB of unified memory at roughly 273 GB/s. Distillation advice that does not name its hardware is close to useless, because the ordering of methods by cost inverts between hardware classes. Appendix B gives you the four numbers you need to redo every calculation in the book for whatever you have.

On the numbers, and on being wrong in public#

Every measured number in this book came out of an execution. Where a number is a prediction rather than a measurement, the text says so, and the figures draw predictions differently from measurements. Where a course notebook made a claim and the data refused it, I have kept the refusal and the diagnosis rather than quietly rewriting the claim. There are four of those in here. They are among the most useful pages in the book, because the shape of being wrong is more transferable than the shape of being right.

The same standard applies to citations. Every reference in the bibliography was verified against its source: the identifier resolves, the title matches, and where a paper has since been published in a venue, the venue is given. Three of the works I lean on are unrefereed preprints, and the text says so at the point where it uses them rather than only in a footnote.

How to read it#

Straight through, if you want the arc. The chapters are ordered by what they ask of you, and each one uses the previous ones.

By part, if you have a specific problem. Parts III and IV are close to self-contained if you have Part I and Chapter 5 under your belt.

Alongside the labs, if you want to come out of it able to do the work rather than discuss it. Appendix E maps every chapter to the notebooks and exercises that make it concrete, in both directions. The order I would suggest is chapter, then lab, then back to the chapter’s exercises, which are written to be answerable only after you have run the thing. One convention travels with the notebooks and is easy to lose: every one of them sets HF_HUB_DISABLE_PROGRESS_BARS=1 in its header, because widget progress bars crash some notebook stacks and the plain log lines carry the same information. If you lift a cell out of a lab and into your own environment, take that line with it.

The exercises are thinking problems, not coding problems. The labs own the coding. Several of them ask you to make a prediction before reading a result, and those are worth doing honestly, because the gap between what you expected and what happened is the only reliable measure of what you have actually understood.

Travis Lelle August 2026

Part I · Foundations

1

What Distillation Is, and What It Is Not

You have a model that works and a constraint that says you cannot ship it.

The constraint is usually one of four things. The model costs too much per token to serve at the volume you need. It does not fit on the hardware where the work has to happen, which might be a phone, a car, a laptop, or a rack you already paid for. It is too slow, because a user waiting on a response has a shorter attention span than a benchmark does. Or it belongs to someone else, and you are renting access to it through an API on terms that could change.

In every one of those cases you want a smaller model that behaves like the larger one on the work you actually care about. Distillation is the family of techniques for getting one.

The name comes from Hinton, Vinyals, and Dean’s 2015 paper, which framed it as extracting the useful structure out of a large model the way you extract a spirit from a mash.1 The idea is older than the name. Buciluǎ, Caruana, and Niculescu-Mizil compressed an ensemble of classifiers into a single neural network in 2006 by having the small model fit the ensemble’s predictions on a large pool of unlabeled data.2 What Hinton’s paper added was the observation that the large model’s probabilities, not just its decisions, carry information the small model can use, and a specific mechanism for exposing that information. That observation is still the center of the subject twenty years later, and Chapter 5 spends most of its length on it.

What has changed since 2015 is scale and stakes. When the teacher is a 70-billion-parameter language model and the student is a 1.7-billion-parameter one, the practical questions stop being about a loss function and start being about which of the two models is generating text at any given moment, how many bytes per second your memory bus can move, and whether the corpus you trained on has quietly contaminated the benchmark you are about to report. Most of this book is about those questions. But you cannot answer them without first being precise about what distillation is, so that is what this chapter does.

1.1 The setup, stated once#

There are two models. The teacher is fixed: you do not train it, and often you cannot, because you do not have its weights. The student is trained. Somewhere there is a set of inputs, and for each input the teacher produces something the student is asked to reproduce.

Definition

Teacher

The model whose behavior is being transferred. Fixed during distillation. Usually larger than the student, though not always, and not necessarily accessible beyond its outputs.

Definition

Student

The model being trained. Its parameters are the only ones that change. Its architecture, size, and tokenizer are all design choices, and Chapters 13 and 14 are about what those choices cost.

The whole design space of distillation falls out of three questions about that picture. Get all three in view before any of them gets an answer.

What does the teacher produce that the student has to match? The candidates run from a single sampled token at the cheap end to the full probability distribution at every position at the expensive end, with truncated distributions, whole generated sequences, and internal activations from the middle of the network in between. Where you land on that range fixes two things at once: how much information flows per training example, and how much it costs to obtain.

Where do the inputs come from? Either from a fixed corpus you already have, or from text the teacher generated, or from text the student generated and the teacher scored after the fact. The last of those is on-policy training and the other two are off-policy, a distinction borrowed from reinforcement learning that matters more than most people expect. Chapter 12 is about it.

How is the mismatch measured? Cross-entropy against a hard label is one answer. A divergence between two full distributions is another, and there is a whole family of those with genuinely different behavior. Chapters 3 and 6 handle this.

Everything else in the subject is a combination of answers to those three questions, plus engineering to make the combination affordable.

1.2 What “knowledge” means here#

The word “knowledge” in “knowledge distillation” is doing real work, and the reason distillation works at all is not obvious, so pin the word down first.

Suppose you have a corpus of text with known continuations. You could train the small model directly on that corpus with ordinary next-token cross-entropy, no teacher involved. That is just language model training. Why would routing the same corpus through a teacher first do anything?

The answer Hinton gave is that the teacher’s probabilities on the wrong continuations carry information that the correct continuation alone does not. Consider a position in a sentence where the true next token is cat. A hard label says: cat is correct, everything else is equally incorrect. The teacher says something much richer. It might put 0.62 on cat, 0.19 on dog, 0.08 on kitten, 0.03 on pet, and spread the remaining 0.08 across forty thousand other tokens including 4e-7 on refrigerator. The ratios among those wrong answers encode a similarity structure the teacher learned: that dog is a plausible substitute here and refrigerator is not, that kitten is close to cat in a way pet is not quite.

Definition

Dark knowledge

The information carried by a teacher’s probabilities on the incorrect outputs. A hard label assigns zero to all of them and therefore says nothing about how they relate to each other; the teacher’s relative probabilities among wrong answers encode a learned similarity structure. The term is Hinton’s.

That is the intuition, and it is a good one. It is also not the whole story, and one of the things this book will keep doing is telling you where the standard intuition stops being reliable. The places it stops are complicated enough to need the rest of the book, so here they are in outline, now, rather than in a section at the end.

The dark knowledge lives in very small numbers. A probability of 4e-7 contributes almost nothing to a cross-entropy loss dominated by a term of size 0.62. Hinton’s fix is temperature, which flattens the distribution before the loss is computed and therefore amplifies the small values relative to the large ones. Temperature is not a cosmetic knob; without it the mechanism does not fire. Chapter 5 derives exactly what it does, including the factor of that has to accompany it and that gets left out of implementations often enough to be worth its own warning.

Distillation frequently fails to make the student match the teacher, and helps anyway. Stanton and colleagues measured this directly: students trained with distillation often have worse agreement with their teacher on held-out data than you would expect from their improved generalization, and improving the optimization does not close the gap.3 Whatever distillation is buying, “the student becomes a copy of the teacher” is an incomplete description of it. Some of the benefit looks more like a regularization effect from training against a smooth target than like information transfer.

A better teacher is not always a better teacher. Past some gap in capacity, a larger teacher produces a worse student than a smaller one would have. Cho and Hariharan showed this on image classifiers and traced it to the student being unable to fit the larger teacher’s function at all, so the distillation loss stops providing useful gradient.4 Chapter 5 covers the competing explanations and what the course’s own capacity-gap probe measures.

I am putting these up front rather than in a “limitations” section at the end because the sequencing matters pedagogically. If you learn the clean story first and the complications later, you will spend a period confidently applying a model of the world that is wrong in ways you cannot see. Better to know from the first chapter that the clean story is a first approximation.

1.3 Where this came from#

The subject has a lineage, and it is worth walking it once, because every step in it was solving a different problem and the shape of current practice is the residue of all of them. There is also a pattern running underneath the usual retelling. Each era assumed a particular level of access to the teacher, mostly without saying so, and the eras where the field moved fastest are the ones where that unstated assumption happened to match what practitioners actually had.

2006: an ensemble you cannot ship. The winning entry in a machine learning competition in the mid-2000s was reliably an ensemble: dozens or hundreds of separately trained models whose predictions are averaged. Ensembles win because averaging cancels errors that individual models make independently of each other. They are also close to unshippable, since serving one means running every member for every request and storing one means storing all of them.

Buciluǎ, Caruana, and Niculescu-Mizil asked whether the accuracy was intrinsic to the size. Their answer was no. The function an ensemble computes is usually not complicated enough to require hundreds of models to express; a single network of modest size has the capacity to represent it. What the single network lacks is a way to find it, because the original labeled training set is small and noisy and ordinary supervised training on it lands somewhere worse. So they used the ensemble as a labeling machine: run it over a large pool of inputs, record what it says, train the single network on that. Where no pool of unlabeled inputs was available, they manufactured one by perturbing the data they had, which is the ancestor of every synthetic corpus in this book.

Hold onto two things from that paragraph. The bottleneck they identified was the training signal rather than the model class, and that diagnosis is still the one the subject rests on. And they needed to run the ensemble as often as they liked, on inputs of their choosing. That is the access assumption the field inherited without stating it, and it took fifteen years for it to become the binding constraint.

2015: probabilities instead of decisions. What the 2006 work transferred was the ensemble’s output in the ordinary sense: the predicted label, or the predicted score. By 2015 the models had changed. Deep networks with softmax heads over hundreds or thousands of classes were standard, and Hinton, Vinyals, and Dean noticed that the output of such a model is not a decision but a distribution, and that nearly all of the distribution lives in the entries nobody looks at.

The reframing did two things. It changed the object being transferred from a function’s decisions to a distribution’s shape, which is the content of §1.2. And it supplied a mechanism for making that shape visible to a loss function, because at a confident classifier’s wrong-answer probabilities are so small that they contribute almost nothing to a gradient. The mechanism matters as much as the observation. Without temperature the information sits in the target and stays invisible to the optimizer, and a reader who takes away “use soft labels” and nothing else will build something that does not work.

Notice the precondition. In a two-class problem there is one wrong answer, and one wrong answer has no internal structure, because there is nothing for it to be relative to. It is the thousand-way softmax of the image classification era, and later the fifty-thousand-way softmax of a tokenizer, that makes the wrong answers worth transmitting at all.

2014 to 2019: everything else the teacher knows. Once the target was a distribution, the obvious next question was whether the output was the only thing worth copying. A trained network is a stack of representations, and the final distribution is a lossy summary of the last one. FitNets asked the student to match an intermediate activation of the teacher through a learned projection, under the name “hints,” and the reason it worked was that a thin deep student got a training signal at its middle rather than only at its end. Attention transfer changed the target again, from what the network computes to where it looks, on the argument that attention maps are indexed by position rather than by channel and therefore compare across networks of different widths.20 Relational methods changed it a third time, to the geometry among examples rather than the examples themselves.

Every one of those methods requires the teacher’s weights and the ability to read out its insides. The vision era took that for granted, because vision checkpoints were files you downloaded, and it is worth registering how much of the classical literature is white-box work assuming a world where the teacher belongs to you.

2019: the recipe arrives in language. Transformer encoders got good and immediately got expensive to serve at web volume, which recreated the 2006 problem with different models. DistilBERT is the compact answer: initialize the student from every other layer of the teacher, train on the teacher’s output distribution plus a term that aligns the two models’ embeddings, and keep most of the quality at roughly half the depth.21 TinyBERT is the maximal answer: match embeddings, hidden states through learned projections, attention matrices, and predictions, in two stages. The pair is instructive because the simple recipe does not lose to the elaborate one by much, and Chapter 14 returns to that comparison with the machinery to make it fairly.

2016, on a separate track: sequences. Everything above concerns a model that emits one distribution per input. A translation model emits a distribution over sequences, and the set of sequences is exponentially large, so the object you would like to match cannot be written down. Kim and Rush’s move was to approximate that intractable object with a single sample: take the teacher’s most likely output sequence, treat it as if it were training data, and fit the student to it with ordinary cross-entropy. The approximation is crude, the results were good, and the method has a property nobody was optimizing for at the time. It needs nothing from the teacher but text.

2023 onward: the generative era. Three things changed at once, and their interaction is why current practice looks different from the classical literature rather than like a scaled-up version of it. Teachers got large enough that generating a corpus with one is the dominant cost of a project rather than a preprocessing step, which inverts the usual ordering of what is expensive. Teachers moved behind APIs, so access stopped being an assumption and became a constraint that varies by vendor and by quarter. And the output moved from a label to a long generation judged by how it reads, which means the quantity you optimize and the quantity you are graded on stopped being the same quantity.

The methods track those changes. MiniLLM argued that a student too small to cover everything the teacher can say should not be trained on an objective that asks it to try, and replaced the forward divergence with the reverse one, which is Chapter 6’s subject.22 Generalized knowledge distillation moved the training inputs onto the student’s own generations, which is §1.4.4 and Chapter 12. DeepSeek’s R1 report went the other way entirely and distilled a series of small models with supervised fine-tuning on teacher-generated reasoning traces and nothing else, a method Kim and Rush would recognize on sight, and the resulting models were strong enough that the simple method is now the reference point rather than the fallback. Meanwhile the cross-tokenizer problem arrived, because a practitioner who wants a specific small model and a specific large one frequently cannot have them share a vocabulary; Universal Logit Distillation is the current answer and Chapter 14 covers what it gives up to get there.23

2026. On-policy methods are the active area, active enough to have accumulated a survey of their own, although that survey describes itself as ongoing work rather than a settled account.24 The honest one-sentence summary of twenty years is that the object of study moved from “compress a classifier” to “transfer a behavior,” and the binding constraint moved from compute to access.

2026-08-01T09:02:42.732338 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 2005 2007 2009 2011 2013 2015 2017 2019 2021 2023 2025 2027 publication year ensemble compression vision era encoder era generative era White box weights and activations 5 of 12 Grey box output distributions 5 of 12 Black box generated text only 2 of 12 before 2020: 4 of 7 works here from 2020: 1 of 5 9 years, two entries an API teacher leaves you in this lane Buciluă et al., model compression FitNets Hinton et al. Kim and Rush, sequence-level KD Attention transfer DistilBERT TinyBERT MiniLLM GKD Minitron Universal Logit Distillation DeepSeek-R1 distilled series
Figure 1.1 Each era of the field produced the methods its access to the teacher allowed, and the assumed access has drifted from full weights toward generated text as teachers moved behind APIs.

Read the figure against your own situation before you read another method paper. If a method sits in a lane you are not in, its results do not transfer to you, however good they are.

1.4 A taxonomy that survives contact with practice#

Textbook taxonomies of distillation usually sort methods by what kind of signal the teacher provides, following the survey literature.5 That axis is real and I will use it. But it is not the axis that determines what you can actually build, so I want to lead with a different one and come back to the classical split afterward.

Before the axes, one scenario to hang them on, because a taxonomy is easy to agree with and hard to use. I am going to carry a single project through all four axes and say what each one answers for it and what each answer eliminates. The project is deliberately ordinary.

The product. A team ships an assistant that currently runs in the cloud against an 8-billion-parameter instruction-tuned model. They want the same assistant on the handset: rewriting and shortening messages the user is drafting, pulling dates and tasks out of pasted text, answering questions about content already on the device. The requirement that starts the project is that it has to work with the network off.

The constraint. The handset can give a model roughly 4 GB and needs the first token back in well under a second. At half a byte per weight, which is what a 4-bit format costs, four gigabytes is an 8-billion-parameter model on paper and considerably less once the key-value cache, the runtime, and the operating system take their share. The target lands at 1.7 billion parameters: 3.4 GB in bf16 during development on the reference machine, near 1 GB once quantized for the device. Nothing about that number came from the distillation literature. It came from a memory budget, which is where student sizes usually come from.

The access. The 8B teacher is an open-weights model, downloaded, sitting on disk, with its weights readable and its activations available to anything that asks.

The corpus. 400,000 prompts logged from the cloud version over six months, de-identified. That is the real asset here: they are drawn from the distribution the student will actually be asked about, which no public instruction corpus is, and §1.8’s second limit is about the teams that do not have such a thing and use a general corpus instead.

1.4.1 Sort by what access you have#

This is the first question to ask about any distillation problem, because it eliminates most of the method space immediately.

Table 1.1 What access to the teacher buys you.

Access level What you can read Methods available Typical situation
White box Weights, activations, gradients, full logits Everything, including representation matching and pruning the teacher into the student You own or downloaded the teacher
Grey box Full or top- output probabilities, no internals Token-level distillation on any divergence, logit caching A served open model, or an API that returns log-probabilities
Black box Generated text only Sequence-level KD, trace fine-tuning A commercial API with no logprob endpoint

The boundaries move. APIs that once returned log-probabilities have stopped; models that were closed have been released. But at any moment, for any given teacher, you sit in exactly one of those rows, and the row determines the chapter of this book that applies to you.

The white-box row deserves one note now because it contains a method people forget exists: if you have the teacher’s weights, one of your options is to build the student out of the teacher. Delete some layers, keep the rest, and every parameter you kept is already trained. Chapter 13 covers this, and the economics are startling enough to change how you plan a project. Minitron reported matching a from-scratch model of the same size at a small fraction of its training compute by pruning and then distilling briefly.6

The black-box row deserves a note too, because it is where a great deal of current practice actually lives, and because the method it forces you into is the simplest one in the book. If all you can get is text, you generate a corpus with the teacher and fine-tune the student on it with ordinary cross-entropy. No divergences, no temperature, no logits. DeepSeek’s R1 report describes exactly this for its distilled model series: supervised fine-tuning on teacher-generated reasoning traces, with no reinforcement learning stage applied to the students.7 It works well enough to be the reference point that fancier methods have to beat, and Chapter 11 is about when they do.

The on-device assistant sits in the white-box row, and the instructive thing about that answer is how little it constrains. Every method in this book is available, so the access axis eliminates nothing and the design has to be settled on other grounds. The top row is the comfortable row to be in and the hard row to plan in.

It is also less stable than it looks. The teacher is someone else’s model under a license, and a license can be superseded. If a re-licensing or an internal legal review moved this project to the bottom row next quarter, everything built on logits would stop existing overnight while the trace fine-tuning path kept working unchanged. That asymmetry is why I now build the black-box baseline first on any project whose teacher belongs to someone else, even when I have logits. It survives every change of access, and it is the number the sophisticated methods have to beat anyway.

1.4.2 Sort by what the student is asked to match#

This is the classical axis, and the three categories the surveys use are worth having.5

Response-based distillation matches the teacher’s outputs. For a language model that means the next-token distribution, or a truncation of it, or a sampled token, or a whole generated sequence. This is where the field spends most of its time, and Parts II through IV of this book are almost entirely about it.

Feature-based distillation matches internal representations. Pick a layer in the teacher and a layer in the student, and add a loss term that pushes their hidden states together, usually through a learned linear projection because the dimensions differ. FitNets introduced this in 2014 under the name “hints,” and TinyBERT built a full recipe on it for transformer language models.89 Chapter 14 covers it, including the check that tells you whether the projection you need can even exist before you spend a training budget finding out.

Relation-based distillation matches the relationships between examples rather than the examples themselves: if the teacher considers inputs A and B similar and A and C distant, the student should too, regardless of where in representation space it puts them. Park and colleagues formalized this, and Tian and colleagues gave a contrastive version.1011 This branch is less developed for generative language models than for vision, and I will be honest about that in Chapter 14 rather than implying a maturity the literature does not have.

For the on-device assistant, response-based is the default and feature-based is a live option rather than an obvious one. The teacher and student differ in both width and depth, so a feature-matching term needs a learned projection between hidden states of different sizes, and Chapter 14’s ridge check tells you whether one can exist before you spend a training budget finding out that it cannot. Relation-based is eliminated, and the reason is worth stating precisely: not that the situation forbids it, but that the methods are immature for generative models and the project has a delivery date. Ruling out a branch on maturity grounds is legitimate; pretending the elimination was technical is not.

1.4.3 Sort by granularity#

For a model that produces sequences, there is a third axis with no equivalent in classification.

Token-level distillation treats each position independently. At position , the teacher has a distribution over the vocabulary, the student has one, and the loss compares them. Sum over positions. This is the direct translation of Hinton’s objective into the sequence setting and it is what most people mean when they say “distillation” about a language model.

Sequence-level distillation treats the whole output as the unit. Kim and Rush’s version, which is the original, approximates the intractable sum over all possible sequences by using the teacher’s most likely output as a single sample and training the student to reproduce it with ordinary cross-entropy.12 The approximation is crude and it works remarkably well, which is one of several places in this subject where theory and practice have an uncomfortable relationship.

The distinction matters because the two methods fail differently. A token-level student can match the teacher at every position and still produce incoherent text, because matching a conditional distribution at every position under teacher-provided context says nothing about what happens when the student has to condition on its own output. Chapter 12 is about exactly that gap.

The on-device assistant starts token-level. The corpus is prompts, the teacher scores each one in a single parallel forward pass, and the replies the product needs are short and structured: a rewritten message, an extracted date, a two-sentence answer. Length is what makes the token-level failure mode bite, and a hundred-token reply accumulates far less drift than a two-thousand-token chain of reasoning. That is an argument for starting token-level rather than a proof that it will suffice, because the acceptance test is what the generations read like while the objective is measured on conditional distributions under context the teacher supplied.

1.4.4 Sort by where the training inputs come from#

The last axis is the one that has moved the most in the last three years, and it is the reason Chapter 12 exists.

Off-policy distillation trains the student on text it did not produce: a fixed corpus, or the teacher’s generations. The name is borrowed from reinforcement learning, where a policy is the thing that chooses actions and “off-policy” means learning from data generated by some policy other than the one you are training.

Definition

Off-policy distillation

Training the student on inputs drawn from a distribution other than the student’s own outputs. The corpus is fixed before training starts, which makes the pipeline cheap, restartable, and easy to reason about.

On-policy distillation trains the student on text the student generated, with the teacher scoring those generations after the fact. The student produces a rollout, the teacher says what it would have done at each position of that rollout, and the loss compares them.

Definition

On-policy distillation

Training the student on the student’s own generated outputs, scored by the teacher. The training distribution moves as the student learns, which addresses exposure bias at the cost of a generation step inside the training loop.

The argument for on-policy training is exposure bias: a model trained only on text it did not produce never sees its own mistakes, and therefore never learns to recover from them.13 The argument was made for sequence models long before distillation adopted it, and scheduled sampling was an early attempt at a fix.14 Agarwal and colleagues built the modern distillation version, generalized knowledge distillation, which interpolates between the two regimes with a single parameter.15

The on-device assistant goes off-policy first. Scoring 400,000 prompts once with the teacher is a prefill workload, the cheapest row of Table 1.3, and what comes out is a durable asset: the cached distributions survive every later change to the student’s learning rate, divergence, initialization, and schedule, so the teacher is paid for once. Chapter 10 is that pipeline. On-policy becomes the second phase if the first leaves a gap that shows up in generations rather than in the loss. That ordering is about more than cost, since the off-policy phase is also the cheapest way to learn whether a 1.7B student can hold this capability at all, which is §1.9’s third failure class and the one you want to meet in week one rather than week six.

There is a cost inversion sitting under that ordering, and it is the reason Chapter 9 exists in the shape it does. On-policy distillation is the cheap path on this machine and a teacher-generated corpus is the expensive one, because what decides the price is which model decodes, and in on-policy training the decoding is done by the student while the teacher only ever scores text that already exists. I had it backwards for a while, for reasons §9.5 dissects; until you have read that arithmetic, take the ordering in Table 1.3 as a claim and not a result.

1.5 Why it works, as honestly as the evidence allows#

There are four explanations in circulation for why a student trained against a teacher beats the same student trained against hard labels. They are not mutually exclusive and the evidence supports different amounts of each.

The information argument. The teacher’s full distribution contains more bits per training example than a one-hot label does. A vocabulary of 49,152 tokens with a one-hot label carries at most bits; the full distribution over the same vocabulary carries substantially more, because it specifies a value for every entry. More bits per example means fewer examples needed for the same amount of learning. This is the argument Hinton made and it is the easiest to state.

The regularization argument. Soft targets are smoother than hard ones, and training against a smooth target reduces overfitting the same way label smoothing does. If this were the whole story, label smoothing should be a substitute for distillation, and it is not: Müller, Kornblith, and Hinton showed that a teacher trained with label smoothing distills worse than one trained without, even though the smoothed teacher is itself more accurate.16 The reason is that label smoothing collapses the very structure among wrong-answer probabilities that distillation transmits. That result is one of the cleanest pieces of evidence that the information argument is doing real work, and Chapter 5 reproduces the reasoning.

The optimization argument. The teacher’s outputs define a smoother loss landscape than hard labels do, so the student’s optimization problem is easier even when the optimum is the same. Beyer and colleagues pushed this line hard, showing that treating distillation as function matching, meaning making the teacher and student see identical inputs under identical augmentation and training for a very long time, produces much better students than the usual recipe.17 Their headline finding is that patience matters more than most architectural choices, which is not what the field expected.

The curriculum argument. In the sequence setting specifically, the teacher’s generated text is easier to learn from than natural text, because it is more consistent, less noisy, and drawn from a distribution a model can actually represent. This is folklore rather than settled science, but it is the most common explanation offered for why sequence-level KD works as well as it does despite its crude approximation, and Kim and Rush’s original analysis gestures at it.12

Watch out

Notice what none of the four arguments claims: that the student ends up computing the same function as the teacher. Stanton and colleagues measured teacher-student agreement directly and found it worse than the generalization improvement would suggest, in settings where the student had enough capacity to match the teacher and the optimizer was given every advantage.3 If your mental model of distillation is “the student learns to be a small copy of the teacher,” you will be repeatedly surprised. A better model is “the teacher’s outputs are a training signal with useful properties, one of which is that they came from something that solves the task.”

1.6 What makes a good teacher#

The teacher is usually chosen in about four seconds. Somebody asks which model is the strongest one the project has access to, that model becomes the teacher, and every later decision happens downstream of a choice nobody wrote down. It is the wrong procedure, and the reasons are specific enough to list, because each one is a case where the obvious criterion and the correct one point in different directions.

Teacher accuracy is not the selection criterion. The distillation loss never reads the teacher’s accuracy. It reads the teacher’s distribution, and nearly all of that distribution sits in the entries that are not the answer. Two teachers can agree on the top-1 prediction at every position of your corpus, score identically on every benchmark you would put in a report, and differ by an order of magnitude in how much structure their tails hold. The first has something to teach; the second is an expensive way to produce hard labels. What separates them is a one-number summary of how unequal the wrong-answer probabilities are, derived in Chapter 5, costing one forward pass over a few hundred prompts.

A label-smoothed teacher is a worse teacher and a better model. §1.5 gave the result: the smoothed teacher is more accurate and better calibrated and it distills worse, because the smoothing objective rewards making every wrong answer equally probable and that equality is the absence of the thing being transferred. The rule generalizes past label smoothing to any procedure that pushes the wrong answers toward each other, and what makes it unsettling is that the damage is invisible to everything anyone reports, because the argmax does not move.

A bigger teacher can be a worse teacher. §1.2 introduced the capacity gap and Chapter 5 covers the competing explanations for it. The planning consequence is that the parameter which matters is the ratio between teacher and student rather than the teacher’s absolute size, so a 70B teacher for a 1.7B student can produce a worse student than an 8B teacher would have. Under the explanation Cho and Hariharan give, the student cannot represent the larger teacher’s function at all, so the distillation term spends its gradient pointing at an unreachable target.

An early-stopped teacher can be a better teacher. This is their remedy, and it breaks people’s intuition hardest: take the teacher’s checkpoint from partway through training, the one you would refuse to ship, and it distills better than the finished model does. If you have intermediate checkpoints of a model you trained, you have candidate teachers, and losing the internal bake-off is not evidence about how they will do here.

The teacher does not have to be larger at all. Born-again networks distill a model into a fresh copy of its own architecture, same size, same everything, and the students come out ahead of the teachers.25 Whatever is happening there, “the teacher must be big” is a convention inherited from compression being the original motivation rather than a requirement of the method.

One framing ties those together. Soft labels are a trade: lower variance in the training signal against a bias toward whatever the teacher is wrong about.26 You want the teacher whose trade lands well for your student on your corpus, and there is no leaderboard for that.

So, a procedure that fits inside a normal project.

  1. Enumerate the candidates you can actually run: every size in the family you have access to, the base and instruction-tuned variants separately, and any intermediate checkpoints of your own.
  2. Score a few hundred prompts from your own corpus with each candidate, one prefill pass. Record three numbers per candidate: top-1 agreement with whatever reference you have, mean entropy of the distribution, and the spread of log-probabilities across the non-top-1 tokens.
  3. Discard the candidates whose tails are flat, however they rank on accuracy. They have nothing to transfer and you would be paying full price to receive it.
  4. Among the survivors, take the two that differ most on those numbers and run the shortest distillation you can afford with each, at matched budget, one key different.
  5. Decide on the student’s numbers, not the teacher’s.

Field note

On the on-device assistant my first instinct was the strongest 8B checkpoint in the family, the instruction-tuned one, and the instinct went unexamined longer than it should have. The family part is right: the corpus is logged product prompts, and a base checkpoint answers a product prompt with a plausible continuation rather than an answer. What I had not thought about is that the tuning which makes a model follow instructions also sharpens it. The tuned checkpoint is more confident at nearly every position than the base model it came from, and a more confident teacher at a fixed temperature has less to say about the wrong answers. That does not make the base model the better teacher. It means the right temperature is a function of which checkpoint I picked, so comparing the two at a shared temperature would have measured the temperature rather than the teacher.

Every step in that procedure ends up measuring the student. The teacher’s own scores are a proxy for the thing the project cares about, and this section is a list of four named ways that proxy has been shown to fail.

1.7 The student’s side of the design#

Four decisions define the student: its architecture family, whether it shares the teacher’s tokenizer, how its weights are initialized, and how big it is. Three of those four are settled before a single training step runs, and none of the three can be revised afterward without throwing away everything trained so far. That makes them the decisions with the most riding on them, and it puts them in an awkward relationship with a literature that spends most of its length on the divergence, the temperature, and the mixing coefficient. Those three are real and Chapters 5 and 6 treat them seriously. They are also revisitable in an afternoon.

Size comes from the deployment constraint, not from the method. Work backwards from where the model has to run, as the on-device assistant did in §1.4: available memory, minus the runtime, minus the key-value cache at the context length you need, divided by the bytes per parameter of the format you will ship in. What comes out is a parameter count, and the distillation design begins after it. Size also fixes the teacher-to-student ratio, which is the capacity-gap parameter, so §1.6 and this section are two halves of one decision rather than two.

Architecture family decides whether teacher and student are relatives. Taking the teacher’s own small sibling buys you a shared tokenizer, shared chat-template conventions, similar width-to-depth proportions, and a live option on feature matching because the layers correspond in a way you can at least argue about. Taking another vendor’s model of the same size buys you whatever that model is better at, which for on-device work is often a genuinely different thing: kernels hand-tuned for the target chip, or a smaller vocabulary. The vocabulary point is not cosmetic at this scale, since the embedding and output matrices are a large fraction of a 1.7B model’s parameters. Note also that “a small model” is itself two different objects: one obtained by training at the target size on a carefully constructed corpus, and one obtained by pruning a larger model down to the target size and continuing to train it.2728 Both hand you a checkpoint. They hand you different checkpoints.

Tokenizer sharing is not an implementation detail. If the two models tokenize differently, then teacher position 7 and student position 7 are not the same span of text, and a position-wise loss is comparing distributions over different things while reporting a number that looks fine. Chapter 7 proves that no general position-wise alignment exists across tokenizer families, and Chapter 14 covers the methods that recover a usable signal anyway and what each one gives up. Everything in Parts II and III of this book assumes a shared tokenizer, so this single decision determines whether two of the book’s five parts apply to you. It is also a serving decision, because the tokenizer sets how many tokens a given reply costs, and tokens are what decoding time is measured in.

Initialization is the one you can revisit. Random, a published checkpoint, or the teacher with layers removed. It changes the starting weights and nothing about the interfaces, so it is the one of the four you can change on Tuesday without invalidating Monday. Chapter 13 measures what each is worth in training steps, and the answer is large enough that treating initialization as a default rather than a decision is the most expensive habit in this list.

For the on-device assistant, three of the four fell out of decisions already made. The size came from the handset. The family came from the teacher, because the 8B model has a 1.7B sibling and sharing the tokenizer keeps the whole book available. That leaves initialization as the only one still open, which is the arm worth running: the published sibling checkpoint against the teacher pruned to the same size, at matched budget. Three decisions made before training and one experiment worth doing is the typical distribution, and it is close to the inverse of where a first-time reader of the literature would expect the consequence to sit.

1.8 What distillation cannot do#

A book that only says what a technique achieves is not useful for planning, so here are the three limits, stated plainly.

It does not create capability the student’s architecture cannot hold. If a task requires more computation per token than the student can perform, no amount of teacher signal supplies it. The capacity gap is the mild version of this; the hard version is that some behaviors seem to require scale and do not survive compression at any ratio. The field does not have a crisp characterization of which behaviors those are, which is an honest gap rather than a failure of this book to look.

It does not transfer what the teacher does not express on your inputs. Distillation is supervised by teacher outputs on the corpus you chose. Teacher behavior that your corpus never elicits does not transfer. This has a practical consequence people discover late: a student distilled on a general instruction corpus can be excellent on general instructions and terrible on your domain, not because the teacher was bad at your domain but because you never asked it.

It does not launder provenance. Whatever is in the teacher can arrive in the student, including things you did not want. Chapter 17 covers the measured version of this: ordinary backdoors mostly do not survive distillation, but triggers built from tokens that are common in distillation corpora do, and the difference is measurable on your own pipeline before you ship.18 The same chapter covers the other direction, where distillation is the attack rather than the pipeline, and someone is extracting a model you serve.

1.9 The four ways a distillation project fails#

Those are limits of the technique. Failures of projects are a different list and a more useful one, because a limit is something you plan around and a failure is something you diagnose at 11pm with a training run three days in. Across the work I have done and the work I have watched, distillation projects go wrong in one of four ways. The value of having exactly four is that they are diagnosed differently and in a definite order, and each has a chapter that handles it.

The pipeline is silently wrong. Nothing crashes, the loss descends, the curve looks like a curve. Meanwhile the labels are shifted by one position relative to the logits, or the prompt tokens are in the loss when they should be masked out, or the teacher and the student tokenize the same string into different numbers of tokens and the loss is comparing position 7 to position 8, or the cache was built from a corpus revision that is not the one being trained on, or the temperature is applied to one side of the divergence and not the other. The signature is a run that trains to a plausible loss and produces a student no better than the hard-label baseline, or one that looks better on the loss and worse on anything you generate from it. You will not find these by reading the code, because the code looks correct, which is why it shipped. You find them by assertion: recompute the framework’s loss externally and compare to float precision, fingerprint the corpus and check it at load time, compute the set of configuration keys by which two arms differ and assert it has exactly one element. Chapter 7 is where these bugs live, Chapter 8 is the discipline that catches them before a run, and Chapter 10 covers the cache-specific ones.

The method does not suit the access you have. You designed around output distributions and the vendor removed the log-probability endpoint. You designed around weights and the legal review came back saying generated text only. You designed around a shared tokenizer and the model the device team ships uses a different one. It presents as a project that stalls with a plan nobody can execute, at the point where the plan meets a constraint that was decided somewhere else in the organization. The defense is to settle the access row first, as §1.4.1 argues, and to build the black-box baseline early because it is the one method available in every row. Chapter 11 is that baseline, and Chapter 12 is what query access buys on top of it.

The student cannot hold the capability. The capacity gap is the mild form; the hard form is a behavior that does not survive compression at any ratio you can afford. The signature is distinctive once you have seen it. The student closes most of the gap quickly and then stops, and the residual does not move for more steps, more data, a different divergence, or a better teacher. It also concentrates rather than spreading, so the remaining errors pile up on a recognizable subset of the evaluation, usually the items needing the most computation per token: long chains of arithmetic, multi-step retrieval across a long context, anything where an intermediate result has to be held and reused. The defense is a small-budget probe before the budget is committed, and the honest response when it fires is to change the size, the initialization, or what you are promising. Chapter 5 covers the mechanism and Chapter 13 covers the two responses that are about the student.

The evaluation cannot see the difference. The subtlest one, the most common in published comparisons, and the reason Part V exists. The loss went down and top-1 agreement with the teacher went up, and neither is evidence, because the loss is computed on the distribution the student was trained toward and agreement counts argmax matches, which is close to the statistic the objective optimized. Both numbers can improve while the model gets worse in ways nobody measured: calibration degrades, so the model’s confidence stops meaning anything;29 output diversity collapses; the eval set overlaps the distillation corpus and the score is partly a memorization score. The defense is to build the evaluation before the run and decide in advance what would count as a win, which sounds like bureaucracy and is the difference between a result and an anecdote. Chapter 16 is the evaluation, Chapter 18 is the practice of registering the decision beforehand, and Chapter 17 is the version where the thing you cannot see is a security property rather than a quality one.

Table 1.2 The four failure classes, in the order to check them.

Failure What you observe Cheapest test Where it is handled
Pipeline silently wrong Loss descends, student is no better than baseline, or loss and generations disagree External recomputation of the loss to float precision; one-key config diff Chapters 7, 8, 10
Evaluation cannot see it Every number improves and nobody trusts the model Measure one quantity the objective does not optimize: calibration, diversity, or a generation-based probe Chapters 16, 18
Method does not suit the access The plan cannot be executed as written Name your row in Table 1.1 and check every planned method against it Chapters 11, 12
Student cannot hold it Gap closes fast, then stops, and the residual concentrates Small-budget probe at the intended ratio before committing Chapters 5, 13

The table’s ordering is diagnostic and differs from the order I introduced them in. When a run disappoints, check the pipeline first, because it is cheap and it is the most likely explanation. Check the evaluation second, because it is also cheap and because being wrong about it wastes everything downstream. Check the method against your access third, and capacity last, because ruling capacity in or out costs a training run and it is the only one of the four whose remedy is to change what you promised.

1.10 The economics, in one page#

Everything in Part III and Part V of this book rests on a small number of facts about hardware, so here is the shape of the argument in advance. Chapter 9 derives it.

Running a transformer has two very different modes. Prefill processes text that already exists, which means every position can be computed at once, which means the arithmetic units stay busy. Decode generates one token at a time, and each step has to read the entire weight matrix out of memory to produce a single token, which means the memory bus is the bottleneck and the arithmetic units idle.

Definition

Prefill

Running a model forward over a sequence that already exists, scoring all positions in parallel. Compute bound, and fast.

Definition

Decode

Generating tokens one at a time, each conditioned on the ones before it. Memory-bandwidth bound, and on large models, slow.

The ratio between them on real hardware is large. Measurements on a 20-billion-parameter model in a 4-bit format on this class of machine put prefill near 2,053 tokens per second and decode near 49.7, a ratio of roughly 40 to 1. That single number reorganizes the whole method space:

Table 1.3 Distillation workloads sorted by who generates.

Workload Mode Who decodes Cost
Teacher scores a fixed corpus to build a logit cache Prefill Nobody Cheap, and you pay once
Teacher scores student rollouts Prefill Nobody Cheap
Student generates rollouts for on-policy training Decode The student, which is small Moderate
Teacher generates a corpus for sequence-level KD Decode The teacher, which is large Expensive

Read the table as a planning tool. Two of those rows are nearly free, one is affordable, and one costs two orders of magnitude more than the others. If your plan involves the expensive row, you should know it before you start, price it explicitly, and then decide whether to buy the corpus once and keep it forever, use someone else’s, or rent different hardware for that stage alone.

The memory side is simpler and just as binding. A parameter in bf16 costs 2 bytes to hold for inference. Full fine-tuning costs about 16 bytes per parameter once you count bf16 weights, bf16 gradients, and the fp32 optimizer moments and master weights that a standard Adam setup keeps. Low-rank adaptation collapses the optimizer share to nearly nothing, which is why it shows up whenever a configuration would otherwise not fit.19 Appendix B has the tables.

1.11 The reference machine#

Every number in this book that depends on hardware refers to one specific machine unless it says otherwise: a single workstation with 128 GB of unified CPU-GPU memory at roughly 273 GB/s of bandwidth, on arm64, CUDA compute capability sm_121.

I am being specific on purpose, for the reason the preface gives: the same method is cheap on one machine and impossible on another, and the ordering of methods by cost inverts between hardware classes. The machine above has generous capacity and modest bandwidth, which is a combination that makes large models fit and makes them slow to generate from. On a machine with the opposite profile you would make different choices, and Chapter 9 gives you the arithmetic to redo the comparison for whatever you have rather than inheriting mine.

2026-08-01T07:26:24.555677 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 0 32 64 96 128 160 memory, GB 32B teacher + 4B student full fine-tune 32B teacher + 8B student LoRA 8B teacher + 4B student full fine-tune 8B teacher + 1.7B student full fine-tune 128 GB ceiling 128 GB 82 GB 80 GB 43 GB does not fit once KV cache and activations are counted LoRA: 1.6 GB of optimizer state, for a student twice the size teacher bf16 weights +grads optimizer state
Figure 1.2 The reference machine's memory budget under four teacher-student configurations, showing where full fine-tuning stops fitting and where low-rank adaptation buys the room back.

1.12 How this book is organized#

The table of contents lists the parts and I am not going to restate it here. Two things about the arrangement are not visible from a contents page, and both change how you should read.

The first is that Part III is the part I would put in a practitioner’s hands if I could only give them one. It covers tokenizers and alignment, where the silent bugs live; the first complete run and the discipline around it; the hardware economics; and then the cached-logit pipeline, which is the cheapest correct thing you can build on this machine and the one I would reach for first on a new project. Parts I and II exist to make Part III readable rather than as ends in themselves, and Parts IV and V are what you reach for once the cheap thing is working and you want to know whether it is good.

The second is that the book has a companion: thirteen laboratory notebooks that execute everything described here. The relationship is deliberate. The book explains and derives; the labs assert. Almost every claim in Part I and Part II is checked mechanically in a notebook against autograd or a closed form, so if a claim ever stops being true the notebook fails loudly. Where a chapter states a measured number, that number came out of one of those executions, and the chapter says which.

1.13 Where this lands in the labs#

Lab 00 is the entry point and it presumes nothing from this chapter beyond the vocabulary. The taxonomy here is not itself asserted anywhere, because taxonomies are not the kind of thing you can assert; what the labs check is every specific claim the taxonomy organizes, and they are scattered across the whole course rather than gathered into a Lab 01. The access rows of Table 1.1 turn into concrete pipelines in Labs 04, 06 and 07, one per row, and running any two of them back to back teaches the ordering in Table 1.3 faster than the arithmetic in Chapter 9 does. If you want the fastest possible confirmation that the subject is real before committing to any of that, Lab 01’s temperature-limit cell takes under a second and shows the factor doing its job on a fixed logit vector you can read in one line.

1.14 Exercises#

  1. For each of the four constraints in this chapter’s opening paragraph, name a distillation method from Table 1.1 that addresses it and one that does not. Where a method addresses the constraint only partially, say what remains.

  2. A colleague proposes distilling a 70-billion-parameter teacher into a 1-billion-parameter student by having the teacher generate 10 million tokens of training corpus. Using only Table 1.3 and the bandwidth figure from §1.11, estimate the wall-clock cost of the generation step to within a factor of two, and state every assumption you made. Chapter 9 will give you the tools to do it properly; do it roughly now and compare later.

  3. The regularization argument and the information argument make different predictions about what happens when you distill from a teacher trained with heavy label smoothing. State both predictions before reading §1.5’s answer, then say which prediction the measured result supports.

  4. You have API access to a teacher that returns the top 20 log-probabilities per position. Which row of Table 1.1 are you in, and which methods from §1.4.2 and §1.4.3 become available or unavailable? Name one piece of information you are missing and one way to bound its size.

  5. Stanton and colleagues found that students often generalize better while agreeing with their teacher less than expected. Propose two experiments that would distinguish “the student learned a different function that happens to be better” from “the student learned the teacher’s function badly in a way that happens to help.” Say what each experiment would measure.

  6. Give a task where you would expect distillation to fail outright, using the three limits in §1.8. Be specific enough that someone could test your prediction.

  7. You are handed four candidate teachers for the on-device assistant in §1.4: the 8B instruction-tuned model, the same model’s base checkpoint, a 32B instruction-tuned model from the same family, and an 8B checkpoint saved 40 percent of the way through instruction tuning. Rank them using only the five facts in §1.6, saying for each which fact applies and in which direction. Then name the one measurement that could overturn your ranking and say what it costs.

  8. A run of the §1.4 scenario finishes. The distillation loss is a third of where it started, top-1 agreement with the teacher on held-out prompts is 71 percent against 44 percent for the hard-label baseline, and the product team reports that the replies read worse than the baseline’s. Assign this to one or more of the four classes in §1.9, say which you would test first and why, and give the cheapest test for each class you named.



  1. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean, “Distilling the Knowledge in a Neural Network,” arXiv:1503.02531 (2015). Presented at the NIPS 2014 Deep Learning Workshop. https://arxiv.org/abs/1503.02531 ↩

  2. Cristian Buciluǎ, Rich Caruana, and Alexandru Niculescu-Mizil, “Model Compression,” Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (2006), 535-541. The originating work for the idea, predating the term “distillation.” ↩

  3. Samuel Stanton, Pavel Izmailov, Polina Kirichenko, Alexander A. Alemi, and Andrew Gordon Wilson, “Does Knowledge Distillation Really Work?” arXiv:2106.05945 (2021), NeurIPS 2021. https://arxiv.org/abs/2106.05945 ↩↩

  4. Jang Hyun Cho and Bharath Hariharan, “On the Efficacy of Knowledge Distillation,” arXiv:1910.01348 (2019), ICCV 2019. https://arxiv.org/abs/1910.01348 ↩

  5. Jianping Gou, Baosheng Yu, Stephen J. Maybank, and Dacheng Tao, “Knowledge Distillation: A Survey,” International Journal of Computer Vision 129, no. 6 (2021): 1789-1819. https://arxiv.org/abs/2006.05525. For the language-model-specific landscape see Xiaohan Xu et al., “A Survey on Knowledge Distillation of Large Language Models,” arXiv:2402.13116 (2024). https://arxiv.org/abs/2402.13116 ↩↩

  6. Saurav Muralidharan et al., “Compact Language Models via Pruning and Knowledge Distillation,” arXiv:2407.14679 (2024), NeurIPS 2024. https://arxiv.org/abs/2407.14679 ↩

  7. DeepSeek-AI, “DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning,” Nature 645 (2025): 633-638, https://doi.org/10.1038/s41586-025-09422-z; preprint arXiv:2501.12948. The distillation section describes supervised fine-tuning on teacher traces with no reinforcement learning stage for the student models. ↩

  8. Adriana Romero et al., “FitNets: Hints for Thin Deep Nets,” arXiv:1412.6550 (2014), ICLR 2015. https://arxiv.org/abs/1412.6550 ↩

  9. Xiaoqi Jiao et al., “TinyBERT: Distilling BERT for Natural Language Understanding,” arXiv:1909.10351 (2019), Findings of EMNLP 2020. https://arxiv.org/abs/1909.10351 ↩

  10. Wonpyo Park, Dongju Kim, Yan Lu, and Minsu Cho, “Relational Knowledge Distillation,” arXiv:1904.05068 (2019), CVPR 2019. https://arxiv.org/abs/1904.05068 ↩

  11. Yonglong Tian, Dilip Krishnan, and Phillip Isola, “Contrastive Representation Distillation,” arXiv:1910.10699 (2019), ICLR 2020. https://arxiv.org/abs/1910.10699 ↩

  12. Yoon Kim and Alexander M. Rush, “Sequence-Level Knowledge Distillation,” arXiv:1606.07947 (2016), EMNLP 2016. https://arxiv.org/abs/1606.07947 ↩↩

  13. Marc’Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba, “Sequence Level Training with Recurrent Neural Networks,” arXiv:1511.06732 (2015), ICLR 2016. https://arxiv.org/abs/1511.06732 ↩

  14. Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer, “Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks,” arXiv:1506.03099 (2015), NeurIPS 2015. https://arxiv.org/abs/1506.03099 ↩

  15. Rishabh Agarwal et al., “On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes,” arXiv:2306.13649 (2023), ICLR 2024. https://arxiv.org/abs/2306.13649 ↩

  16. Rafael Müller, Simon Kornblith, and Geoffrey Hinton, “When Does Label Smoothing Help?” arXiv:1906.02629 (2019), NeurIPS 2019. https://arxiv.org/abs/1906.02629 ↩

  17. Lucas Beyer et al., “Knowledge distillation: A good teacher is patient and consistent,” arXiv:2106.05237 (2021), CVPR 2022. https://arxiv.org/abs/2106.05237 ↩

  18. Giovanni De Muri, Mark Vero, Robin Staab, and Martin Vechev, “Pay Attention to the Triggers: Constructing Backdoors That Survive Distillation,” arXiv:2510.18541 (2025), ICLR 2026. https://arxiv.org/abs/2510.18541 ↩

  19. Edward J. Hu et al., “LoRA: Low-Rank Adaptation of Large Language Models,” arXiv:2106.09685 (2021), ICLR 2022. https://arxiv.org/abs/2106.09685 ↩

  20. Sergey Zagoruyko and Nikos Komodakis, “Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfer,” arXiv:1612.03928 (2016), ICLR 2017. https://arxiv.org/abs/1612.03928 ↩

  21. Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf, “DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter,” arXiv:1910.01108 (2019), 5th Workshop on Energy Efficient Machine Learning and Cognitive Computing, NeurIPS 2019. https://arxiv.org/abs/1910.01108 ↩

  22. Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang, “MiniLLM: Knowledge Distillation of Large Language Models,” arXiv:2306.08543 (2023), ICLR 2024. https://arxiv.org/abs/2306.08543v2 The arXiv landing page currently shows a later title; the ICLR 2024 version of record is the one cited here. ↩

  23. Nicolas Boizard, Kevin El Haddad, Céline Hudelot, and Pierre Colombo, “Towards Cross-Tokenizer Distillation: the Universal Logit Distillation Loss for LLMs,” arXiv:2402.12030 (2024), Transactions on Machine Learning Research, January 2025. https://arxiv.org/abs/2402.12030 ↩

  24. Mingyang Song and Mao Zheng, “A Survey of On-Policy Distillation for Large Language Models,” arXiv:2604.00626 (2026). An unrefereed preprint whose own comment field describes it as ongoing work; read it as a map of activity rather than a settled account. https://arxiv.org/abs/2604.00626 ↩

  25. Tommaso Furlanello, Zachary C. Lipton, Michael Tschannen, Laurent Itti, and Anima Anandkumar, “Born Again Neural Networks,” arXiv:1805.04770 (2018), ICML 2018. https://arxiv.org/abs/1805.04770 ↩

  26. Helong Zhou et al., “Rethinking Soft Labels for Knowledge Distillation: A Bias-Variance Tradeoff Perspective,” arXiv:2102.00650 (2021), ICLR 2021. https://arxiv.org/abs/2102.00650 ↩

  27. Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen, “Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning,” arXiv:2310.06694 (2023), ICLR 2024. https://arxiv.org/abs/2310.06694 ↩

  28. Loubna Ben Allal et al., “SmolLM2: When Smol Goes Big, Data-Centric Training of a Small Language Model,” arXiv:2502.02737 (2025). https://arxiv.org/abs/2502.02737 ↩

  29. Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger, “On Calibration of Modern Neural Networks,” arXiv:1706.04599 (2017), ICML 2017. https://arxiv.org/abs/1706.04599 ↩

Part I · Foundations

2

Distributions Over Tokens, and the Numbers That Hold Them

Every method in this book bottoms out in the same operation. A model emits a vector of raw scores, one per vocabulary entry. That vector becomes a probability distribution over the next token. Some other distribution, obtained the same way from another model, is compared against it, and the comparison produces a number that gets differentiated. That pattern is what the surveys call response-based distillation.1415 Chapters 3 and 6 are about which comparison to use. This chapter is about what comes first: the distribution itself, and the floating-point arithmetic that has to hold it.

The second half of that sentence is the part people skip, and the part that costs them a week. A next-token distribution has between 50,000 and 256,000 entries spanning forty orders of magnitude: the token the model is confident about sits near 0.9, and one it has ruled out sits at or lower. Distillation is the business of caring about the small entries, because that is where the teacher’s learned structure lives. So it is the business of doing arithmetic near the floor of whatever number format you are in, and the format has a floor, and things below the floor become exactly zero, and the logarithm of exactly zero is negative infinity.

Here is what that looks like from the outside. You launch a run and the first loss prints as nan. Or worse, a thousand losses print as ordinary numbers, then one prints as inf, and everything after it is nan. The natural reading of that dashboard is “the optimizer diverged,” and the natural response is to lower the learning rate, add gradient clipping, and try again. That response is wrong, it costs hours, and the cause is one function call that took a logarithm of something which had underflowed to zero. I have watched people, myself included, chase the wrong hypothesis here.

So this chapter builds the object, then the number system under it, and tries to leave you able to predict which arithmetic survives before you run it.

2.1 What a next-token distribution is#

A language model, given a sequence of tokens, produces at each position a vector of real numbers with one entry per vocabulary item. Call the vocabulary size and the vector $z \in \mathbb{R}^V$. These numbers are not probabilities. They are unnormalized scores, and they are called logits.

Definition

Logit

A raw, unnormalized score emitted by a model’s output layer, one per vocabulary entry. Logits live on the whole real line, bounded neither above nor below, and become probabilities only after a softmax.

The map from logits to probabilities is the softmax:

Definition

Softmax

The function that turns a vector of real-valued logits into a probability distribution: . Every output is positive because , and the outputs sum to one because the denominator is the sum of the numerators.

One property of the softmax determines most of what follows, and it is worth deriving rather than accepting.

The softmax is shift invariant. Add the same constant to every logit and the distribution does not change:

$$\mathrm{softmax}(z + c)_i = \frac{e^{z_i + c}}{\sum_j e^{z_j + c}} = \frac{e^{c}\,e^{z_i}}{e^{c}\sum_j e^{z_j}} = \frac{e^{z_i}}{\sum_j e^{z_j}} = \mathrm{softmax}(z)_i.$$

The factor appears once in the numerator and once in every term of the denominator, so it factors out and cancels. Nothing about this is deep, and it has three consequences that are.

First, only logit differences carry information. If someone tells you a model assigned a logit of 14.2 to a token, you have learned nothing, because 14.2 relative to what is the only question. Comparing two models’ raw logits entry by entry is meaningless; comparing their log-probabilities is fine.

Second, the softmax is not injective. A whole line of logit vectors, ${z + c : c \in \mathbb{R}}$, maps to the same distribution, so “match the teacher’s distribution” and “match the teacher’s logits” are different objectives and the second is strictly stronger. Chapter 5 shows they converge at high temperature, which is a fact about the limit rather than the general case.

Third, and this is the one that saves the computation, you are free to subtract anything you like from all the logits before exponentiating. Section 2.2 uses that freedom.

A note on notation, fixed here for the rest of the book: is always the teacher’s distribution and is always the student’s. The course’s code follows the same convention, while the divergence functions in kd_core take arguments in the order (student, teacher), the opposite of the order they appear in the mathematics. Chapter 3 puts a warning on that.

The vocabulary sizes involved are large enough to change the arithmetic. The course prices three: Llama 3 at 128,256 entries, Qwen2.5 and Qwen3 at 151,936, Gemma 2 and 3 at 256,000.1 The 360M-parameter student the labs train has 49,152.2 A dense distribution over 151,936 entries in a 16-bit format is about 302 KB per token position, the number that makes Chapter 10 exist. That figure rounds the vocabulary to 151k, which is the form the labs carry it in; multiply the exact 151,936 out and you get 303.9 kB, and the difference has never changed a decision made with it.

2.2 Overflow, and why the textbook formula does not run#

Write out the softmax exactly as it appears above, in code, and it will fail on real inputs.

Every floating-point format has a largest finite value. A calculation whose true result exceeds it does not wrap around or saturate; it produces inf, and inf divided by inf produces nan, and nan propagates through every later operation until the whole tensor is nan.

Definition

Overflow

A calculation whose true result is larger in magnitude than the largest value the number format can store. The result becomes inf, which then contaminates everything computed from it.

The single-precision format, fp32, tops out at about . So the question “which logits overflow when you exponentiate them” has an exact answer:

That is the whole derivation. No rule of thumb is involved: 88.7 is the natural logarithm of fp32’s ceiling, so exp of any logit above roughly 88 overflows and every logit below it is fine.

Half precision, fp16, has a much lower ceiling of 65,504, and the same derivation gives

so exp of any logit above about 11 overflows fp16. Eleven. A trained language model routinely emits logits well past that.

Now put realistic numbers in. Logits from a trained checkpoint sit around 10 to 40 in magnitude, comfortable in fp32 and already dangerous in fp16. Then divide by a temperature, an operation Chapter 5 performs constantly and Section 2.4 defines. A temperature of 0.25 multiplies every logit by 4, so logits near 30 become logits near 120, well past 88.7, and the naive softmax overflows in fp32. Lab 00 runs this case on the vector at and asserts that the naive implementation produces nan while the stable one matches PyTorch’s F.softmax to .

The fix is the third consequence of shift invariance. Let and compute

By shift invariance with , this is the same distribution, exactly rather than approximately, and the largest exponent is now , so no term can overflow. Terms can still underflow, which is Section 2.3’s problem, but they cannot blow up. The transformation costs one pass to find the maximum, and it is what every library does inside F.softmax.

The listing below is the shape of the idea. Look at the single line of difference between the two functions and at what it does to the exponent.

import torch

def naive_softmax(z):
    e = z.exp()                                  # e^124 has nowhere to go in fp32
    return e / e.sum(-1, keepdim=True)

def stable_softmax(z):
    m = z.max(-1, keepdim=True).values           # shift invariance says this is free
    e = (z - m).exp()                            # largest exponent is now exactly 0
    return e / e.sum(-1, keepdim=True)

z = torch.tensor([31.0, 29.5, 27.0, 12.0]) / 0.25   # temperature 0.25 -> logits near 120
print(naive_softmax(z))    # tensor([nan, nan, nan, nan])
print(stable_softmax(z))   # tensor([9.9753e-01, 2.4726e-03, 1.1226e-07, 9.8298e-34])

The two functions implement the same formula and differ in whether the intermediate quantities fit in the format. That distinction, between the mathematics being right and the arithmetic being survivable, is the subject of this chapter.

You will not write naive_softmax on purpose. You will meet it when you implement anything below the level of the framework, and in a nastier form when someone builds a logit cache and stores instead of to save an exponential at training time. The stored values carry no record of what was subtracted, the protection is gone, and nothing warns you. Chapter 10 specifies what a cache record should contain, and log-probabilities are the answer, for this reason among others.

2.3 logsumexp, and the rule about never taking the log of a softmax#

The denominator of the softmax has a name and a life of its own. The log-partition function, called logsumexp in code, is

Definition

Log-partition function

Written and called logsumexp in every numerical library. It is the logarithm of the softmax’s normalizing denominator, and it is the primitive that log-probabilities, cross-entropy, and every divergence in this book are built from.

Take the logarithm of the softmax formula and you get, with no approximation,

That identity is the definition of log_softmax. A log-probability is a logit minus a single scalar shared by the whole vector.

LSE has the same overflow problem as the softmax, and the same fix. With ,

$$\mathrm{LSE}(z) = \log \sum_j e^{z_j} = \log \left( e^{m} \sum_j e^{z_j - m} \right) = m + \log \sum_j e^{z_j - m}.$$

The first step pulls out of the sum, valid because it is constant with respect to ; the second uses . The right-hand side never exponentiates anything above zero, and the remaining term is bounded between 0 and . Same quantity, computed in an order the number format can survive.

Now the rule, which is the single most useful sentence in Lab 00:

Never take the logarithm of a probability you obtained from a softmax. Get log-probabilities directly from log_softmax.

The reason is underflow, which is overflow’s mirror image and much less widely understood.

Definition

Underflow

A calculation whose true result is smaller in magnitude than the smallest positive value the number format can store. The result becomes exactly 0.0. Unlike overflow, underflow produces a perfectly ordinary-looking number, and you find out about it later, when something takes its logarithm or divides by it.

Here is the argument in full, because it is the load-bearing one.

The fp32 format’s smallest positive value is about . (That figure is the smallest subnormal, which Section 2.7 defines; the smallest value with full precision is much larger, at .) Anything below the floor is stored as exactly 0.0, with no warning and no exception.

Take logits . This is not an exotic input. It is what a trained model looks like at a position where it is confident: one plausible continuation and one it has ruled out, separated by a large margin. The true probability of the second token is

That is below by a factor of about a thousand. Run the composed path, softmax and then log: the softmax stores as 0.0, the logarithm of 0.0 is , and you have a negative infinity sitting in your log-probability tensor.

Run the direct path, log_softmax, and the identity is evaluated instead. Here , so the answer is , a subtraction of two ordinary-sized numbers. The tiny probability is never materialized in any intermediate at any point, and the result is a normal float that behaves correctly downstream.

Same mathematics. Different order of operations. Only one order fits inside the number format.

What happens next is worth memorizing, because it is how you will encounter this. The flows into a divergence. Forward KL contains a term ; if the student’s log-probability is where the teacher’s probability is positive, that term is and the loss is inf. Subtract inf from inf downstream and you get nan. One backward pass later every parameter is nan, and every subsequent loss is nan regardless of the data.

On a dashboard that is a curve which behaves and then goes vertical, and the label your brain attaches to it is “the run became unstable.” Everything you know about unstable training tells you to reduce the learning rate, add warmup, clip the gradient norm. None of it helps. The tell is that the failure reproduces at the same step on the same data regardless of the learning rate, which is not how genuine optimizer divergence behaves.

The listing below is the smallest complete demonstration. Watch the second entry of each output.

import torch
import torch.nn.functional as F

z = torch.tensor([0.0, -110.0])          # a token the model has ruled out

composed = F.softmax(z, -1).log()        # materialises exp(-110) ~ 1.7e-48, stores 0.0
direct   = F.log_softmax(z, -1)          # evaluates z_i - logsumexp(z), all ordinary sizes

print(composed)   # tensor([0., -inf])
print(direct)     # tensor([   0., -110.])

assert torch.isinf(composed[1])
assert torch.isfinite(direct[1])

The composed path is wrong by an infinite amount on an input a real model produces at a real position, which is the justification for a rule that otherwise sounds like style preference.

Both of those have a consequence past this one function. If you need for small , use log1p(-x) rather than composing, for the same class of reason. And if you are handed probabilities rather than log-probabilities, by an API or a cache or a colleague, you have lost information you cannot recover, because every value that underflowed is now indistinguishable from one that was genuinely zero. Serving stacks return log-probabilities for this reason, and Chapter 15 covers the shape of that response.16

2.4 Temperature#

Temperature is one scalar divided into every logit before the softmax.

Definition

Temperature

A positive scalar that rescales logits before the softmax: . Values below 1 sharpen the distribution toward its largest entry; values above 1 flatten it toward uniform; leaves it unchanged.

The name is borrowed from statistical physics, where the same expression describes the occupancy of energy states at temperature , and the analogy carries: high temperature means the system spends time in states it would otherwise never visit.

What the knob does is pinned down by its two limits, both derivable in a few lines.

As , the distribution converges to a one-hot vector on the largest logit. Write the softened probability with the max subtracted, which shift invariance permits. Let , attained at a unique index . Then

Every exponent is zero or negative. For the numerator is . For every other , , so as and the numerator goes to zero while the denominator goes to 1. So and everything else goes to 0. Ties for the maximum spread the limit uniformly over the tied set, by the same argument.

As , the distribution converges to uniform. Now every exponent $(z_i - m)/T \to 0Vp_i \to 1/V$, and the distribution has forgotten the logits.

Between those limits it moves smoothly, and the way it moves is why temperature is in this book at all. Take the seven-token vector Lab 01 uses, with logits standing for the, a, an, cat, dog, xylophone, qq.

Table 2.1 One logit vector, six temperatures. Probabilities rounded; entropy in nats.

entropy the a an cat dog xylophone qq
0.25 0.0006 0.99995 0.000045 0.000006 0 0 0 0
0.5 0.057 0.9909 0.0067 0.0025 0.000001 0 0 0
1.0 0.448 0.8823 0.0724 0.0439 0.0008 0.0005 0.00004 0.00002
2.0 1.053 0.6353 0.1820 0.1417 0.0192 0.0149 0.0043 0.0026
4.0 1.578 0.4031 0.2158 0.1904 0.0700 0.0618 0.0331 0.0258
10.0 1.875 0.2387 0.1859 0.1768 0.1185 0.1128 0.0878 0.0795

Read the top row. At the distribution is one-hot to five decimal places, so a student trained to match it learns which token was correct and nothing else, which is what a hard label would have taught it. Everything distillation is supposed to transfer has been annihilated by the temperature.

Read the bottom row. At , xylophone is at 0.0878 against the at 0.2387, a ratio of about one to three. The teacher does not believe xylophone is a third as likely as the; it assigned them logits of and . What the student is being asked to reproduce, with a substantial share of the loss, is the teacher’s noise floor, the region where it has ruled everything out and the remaining differences are close to arbitrary.

The middle rows are where the mechanism operates. Dark knowledge, defined in Chapter 1, is carried by the ratios among the teacher’s rejected tokens. In the row, everything below the top token amounts to 0.118 of the mass, and the informative part of it, the difference between cat at 0.0008 and xylophone at 0.00004, contributes a cross-entropy term four orders of magnitude below the top token’s. The gradient will not notice. Raise the temperature and those ratios survive while their absolute sizes grow: at the same two tokens sit at 0.0192 and 0.0043, large enough to move a loss.

That is the mechanism: temperature does not create the information in the teacher’s tail, it rescales the tail into the range where a gradient-based optimizer can respond to it. Hinton, Vinyals, and Dean introduced this move, and it is the one hyperparameter their method could not do without.3 A factor of has to accompany the softened loss term to keep its gradient scale comparable across settings of , and Chapter 5 derives it.

2026-08-01T07:26:27.421249 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 1 0 − 6 1 0 − 4 1 0 − 2 1 0 0 probability one-hot: a hard label would have taught the same thing T = 0.25 H = 0.0006 nats T = 0.5 H = 0.0573 nats T = 1 H = 0.4479 nats the a an cat dog xylophone qq 1 0 − 6 1 0 − 4 1 0 − 2 1 0 0 probability T = 2 H = 1.0527 nats the a an cat dog xylophone qq T = 4 H = 1.5783 nats the a an cat dog xylophone qq xylophone within a factor of 2.7 of `the` T = 10 H = 1.8748 nats
Figure 2.1 One fixed logit vector under six temperatures: as T rises the top token's mass drains into the tail, and the ratios among rejected tokens grow into a range a gradient can act on.

One clarification, because the same word covers three jobs. The temperature on an inference API applies this formula to change what gets sampled. The temperature in a distillation loss changes what the target looks like, with no sampling involved. And temperature scaling as post-hoc calibration fits one scalar on held-out data so that a model’s confidences match its accuracy, which makes it measured rather than chosen; Chapter 16 uses that version.4

The useful range of depends on how peaked your teacher is on your data. Peaked domains, where the continuation is nearly forced (code, structured output, arithmetic), tolerate low temperatures because there is little tail to expose. Diffuse domains, where many continuations are reasonable, need more. That is a property of your corpus and your teacher, measurable in an afternoon, and not a constant to copy out of someone else’s configuration file.

2.5 Entropy, in nats#

Temperature moves a distribution between “certain” and “uniform.” Entropy says where on that axis a distribution currently sits.

Definition

Nat

The unit of information you get when your logarithms are natural. One nat is $1/\ln 2 \approx 1.4427$ bits. Everything in this book is in nats, because every loss in this book uses natural logarithms and because log returns the natural logarithm in every framework you will use, and mixing bases silently changes every number by a factor of 0.693.

Definition

Entropy

, measured in nats. It is the average number of nats of surprise you get per sample from , and operationally it measures how spread out is. Zero for a distribution certain of its answer, for the uniform distribution over outcomes, and nothing outside that range.

Both extremes are worth confirming. For a one-hot distribution one term is and every other is , taken to be zero by the limit . So : a distribution that always gives the same answer carries no surprise. For the uniform distribution every , so

and a standard argument by concavity shows no distribution does better. Over a 151,936-token vocabulary that ceiling is nats.

The teacher’s entropy is one of the two terms in the identity connecting entropy, cross-entropy, and KL, and Chapter 3 lives inside that identity. What I want here is the shape of as a function of temperature. The limits are established: as and as . In between, rises monotonically, and the rate is worth knowing at both ends.

At high temperature, expand for small , where is the mean logit. Then , and carrying the expansion of to second order gives

where is the variance of the logits. Entropy approaches its ceiling like , at a rate set by how spread out the logits were. On Table 2.1’s seven-token vector that variance is 14.48, so the prediction at is nats against an exact 1.8748. It lands that close because is already deep in the asymptotic regime for logits of this size.

At low temperature, let be the gap between the largest logit and logit . Each rejected token contributes about to the entropy, so

which goes to zero exponentially fast in . Entropy does not decline gently as you cool a distribution; it falls off a cliff. On Table 2.1, dropping from 1.0 to 0.5 cuts entropy by a factor of eight, and dropping to 0.25 cuts it by another factor of ninety.

2026-08-01T07:26:29.727698 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 0.1 0.25 0.5 1 2 4 10 temperature T 0.0 0.5 1.0 1.5 2.0 entropy of softmax(z / T), nats uniform ceiling, log 7 = 1.9459 one-hot floor 0.0006 0.0573 0.4479 1.0527 1.5783 1.8748 exact entropy low-T expansion ∑ i D T i e D T − / i ( 1 + / ) high-T expansion l o g   V   , − / 2 σ T 2 2 σ 2   =   1 4 . 4 8 where T is actually chosen
Figure 2.2 Entropy in nats of a fixed logit vector as a function of temperature, with the two asymptotes marked: exponential decay to zero as T falls and 1/T-squared approach to log V as T rises.

Entropy earns its place here for a reason beyond Chapter 3’s identity: it is the cheapest useful thing you can monitor during a run that is not the loss. The loss tells you the objective is going down, which it does whether or not anything good is happening. Mean entropy tells you whether the student is becoming more decisive, and how fast.

A healthy student’s entropy declines somewhat and then plateaus, because training on a teacher’s targets concentrates its mass. A student in trouble keeps going, and the failure has a name. Entropy collapse is the runaway version, where the distributions narrow until the output is effectively deterministic and the text loses all variety. It is well documented in reinforcement-learning post-training of reasoning models, characterized there as an exchange between entropy and short-term performance, and the same mechanism appears in on-policy distillation, where the entropy worth watching is measured on the student’s own rollouts.5612 Chapter 12 covers the monitors that catch it and the distinction that makes one useful: a healthy decline of ten to thirty percent is not collapse, and a monitor that fires on that is worthless. The course’s own pairs a floor of 0.15 nats with a rule about entropy lost inside a trailing window.

Entropy also has a diagnostic use that costs nothing. If your student’s mean entropy is larger than your teacher’s on the same inputs, the student is hedging, which is the signature of a mode-covering objective (Chapter 6). If it is much smaller, the student has committed, which is the signature of a mode-seeking one. That one number, read before any sample text, usually tells you which situation you are in.

2.6 The tail#

A next-token distribution has a head and a tail, and much of distillation depends on which one you are paying attention to.

The head is easy to describe. At most positions in ordinary text a trained model is confident: the top token takes most of the mass, a handful of alternatives take most of what is left, and tens of thousands of remaining entries share a remainder that rounds to nothing. On Table 2.1 at , the top token holds 0.8823 and the other six share 0.1177.

The tail is the rest of the vocabulary, and it has two properties that pull in opposite directions.

The tail holds almost no probability mass. The course measures this on real data: Lab 02 runs the teacher over the corpus Lab 03 trains on and computes how much of each position’s mass the top entries capture. At , mean retained mass exceeds 0.99. Ninety-nine percent of the teacher’s belief lives in 64 out of 49,152 entries, 0.13 percent of the vocabulary. Chapter 10 turns that measurement into a storage decision.

The tail holds most of the entries, and all of the information about what the model considers absurd. Dark knowledge is a statement about ratios among rejected tokens, and rejected tokens are by definition in the tail. A student that matches the head perfectly and ignores the tail has copied the teacher’s decision and learned nothing about its similarity structure.

Both statements hold at once, which is why the tail is the interesting part of the object and the hardest to work with. A quantity can be numerically negligible and informationally significant at the same time, and temperature is the device that converts the second property into something the first does not suppress.

The tail’s third property is the one this chapter cares about most. The tail is where the underflow lives: a probability of is a fine fp32 number and a hard zero in fp16, and is a hard zero in both. Every claim about the tail is therefore also a claim about your number format, and a pipeline that computes the head correctly can be silently wrong about the tail on the same batch.

One more thing the tail does will surprise you if you have not seen it, and Chapter 3 makes it precise. A tiny amount of mass in the wrong place can produce an unbounded KL divergence, because the divergence weights each token by the logarithm of a probability ratio and a logarithm has no ceiling. Lab 00’s second solution exercise builds distribution pairs whose total variation distance goes to zero while their KL goes to infinity, using nothing but a shrinking mass on a token where the other distribution is astronomically small. So “the tail holds one percent of the mass” does not license “the tail contributes one percent of the loss.” It usually does, it can fail arbitrarily badly, and the case where it fails is the one where the student puts near zero probability on something the teacher likes, which is what a freshly initialized student does everywhere.

Carry three things out of this section and into Chapter 3. The tail exists, its mass is measurable on your own corpus with one prefill pass, and you should measure it. When the only thing you have is the token that was sampled, the tail becomes something you estimate rather than compute, and Chapter 4 is about how badly that can go.13 And the split between the top-1 token and everything else is the split between what a hard label tells you and what distillation is for. Work on generation quality has circled the same object from the sampling side for years: nucleus sampling exists because the shape of the tail decides whether sampled text degenerates.7

2.7 Floating point, in enough detail to predict failures#

Everything above assumed the numbers behave. This section says what they actually do.

2.7.1 The anatomy of a float#

A floating-point number is stored as three bit fields: a sign bit , an exponent field , and a mantissa field with bits. In the normal range, the value is

The factor in parentheses is the significand. It lies in , and the leading 1 is not stored because it is always there, which buys one bit of precision for free. That implicit bit is the source of the documentation conflict this section resolves later.

That layout has two consequences, and between them they explain every property of every format.

The exponent field buys range. With exponent bits the exponent covers roughly powers of two, so each additional exponent bit squares the ratio between the largest and smallest representable magnitude. Range is cheap.

The mantissa field buys precision, and it buys the same relative precision everywhere. With stored mantissa bits the significand is quantized into steps across , so within any binade (any interval ) the representable values are spaced apart. Spacing scales with magnitude: near 1.0 the gap is , near 1000 it is a thousand times larger, near a million times smaller. A float has constant relative precision and wildly varying absolute precision, and most floating-point surprises reduce to expecting the opposite.

That constant has a name.

Definition

Machine epsilon

The gap between 1.0 and the next representable number above it. For a format with stored mantissa bits, machine epsilon is exactly . Read it as a digit budget: a format with an epsilon of carries about seven reliable decimal digits, and one with an epsilon of carries between two and three.

Machine epsilon is measurable without knowing anything about the bit layout, which makes it the honest way to settle an argument about a format’s precision. If it comes out to , then , whatever the documentation says.

2026-08-01T07:26:28.743315 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 1 0 − 4 5 1 0 − 3 0 1 0 − 1 5 1 0 0 1 0 1 5 1 0 3 0 magnitude of the number 1 0 − 4 5 1 0 − 3 0 1 0 − 1 5 1 0 0 1 0 1 5 1 0 3 0 gap to the next representable value where a rejected token's probability lives fp32 fp16 bf16 1.4e-45 9.2e-41 6.0e-8 flat: the subnormal floor. Below it, everything is zero. fp16 overflows to inf at 65504 fp32 and bf16 overflow at 3.4e38 at x = 1 the gap is the machine epsilon: 7.8e-3 bf16, 9.8e-4 fp16, 1.2e-7 fp32
Figure 2.3 Spacing between adjacent representable values against magnitude for fp32, fp16, and bf16: constant relative precision shows up as parallel lines on log-log axes, and the vertical cliffs are where each format overflows and where each one underflows to zero.

2.7.2 The three formats you will actually use#

Table 2.2 The formats, with everything that follows from the bit split.

format total bits sign exponent bits stored mantissa bits significand precision max finite machine epsilon smallest normal smallest subnormal
fp32 32 1 8 23 24 bits 3.4e38 1.2e-7 1.2e-38 1.4e-45
fp16 16 1 5 10 11 bits 65504 9.8e-4 6.1e-5 6.0e-8
bf16 16 1 8 7 8 bits 3.4e38 7.8e-3 1.2e-38 9.2e-41

Read the table as one decision made twice with opposite answers. Both 16-bit formats have 16 bits to spend. fp16 puts 10 into mantissa and 5 into exponent. bf16 puts 7 into mantissa and 8 into exponent, which is exactly fp32’s exponent field.

Definition

bf16

Brain floating point, a 16-bit format with fp32’s 8 exponent bits and only 7 stored mantissa bits. It has fp32’s dynamic range and about a third of its precision: nothing that fits in fp32 overflows or underflows in bf16, and every value carries only two to three reliable decimal digits. Converting fp32 to bf16 is a truncation of the low 16 bits, which is why the conversion is so cheap.

The consequences follow from the table. fp16 overflows at 65,504, which Section 2.2 showed is exp of a logit above 11, and it flushes any probability below about to exact zero. Both thresholds sit inside a language model’s operating range, which is why fp16 needs loss scaling to train at all and why it produces exact zeros where a distribution should have small positive numbers. bf16 cannot overflow on anything fp32 could hold and cannot underflow before fp32 does. What it does instead is round: every bf16 result is good to about three decimal digits and noise after that.

The failure signatures differ accordingly. fp16 fails loudly, with inf, nan, and a dead run. bf16 fails quietly: the curve is smooth, the numbers are plausible, and the third digit is garbage. On a dashboard that looks like a plateau, or like a run that gives slightly different numbers each time you repeat it with the same seed. You conclude your intervention did not work, when your instrument could not resolve the effect you were measuring.

2.7.3 Subnormals and flush-to-zero#

Section 2.7.1’s formula has a lower limit: once the exponent field hits its minimum you cannot go smaller by decrementing it. IEEE 754 handles that with a second regime.

Definition

Subnormal

A floating-point value below the format’s smallest normal magnitude, represented with an implicit leading 0 instead of an implicit leading 1 and a fixed smallest exponent. Subnormals extend the range downward toward zero at the cost of losing precision progressively: the smallest subnormal has one significant bit. Also called denormal.

The gap between the two regimes is larger than people expect. fp32’s smallest normal value is and its smallest subnormal is , seven orders of magnitude further down, obtained by giving up the 23 mantissa bits one at a time. Below there is nothing but zero.

Operationally, subnormals matter in two places. First, they are why Section 2.3’s underflow threshold is rather than , seven orders of magnitude in how far a probability can fall before becoming a hard zero. Second, they are not always enabled. Many accelerators and some compiler settings implement flush-to-zero, replacing any subnormal result with zero because subnormal arithmetic is expensive in hardware. The effective floor then rises to the smallest normal, and probabilities between and that survived on one device become exact zeros on another. That is one mechanism by which a loss that was finite on your laptop becomes inf on your accelerator with no change to the code. The defense is the one this chapter keeps giving: stay in log space.

2.7.4 The 7-versus-8 question, resolved#

The course’s own materials disagree with themselves here, and rather than quietly fixing one side I want to state the resolution, because this is the kind of disagreement that looks like a contradiction and is not.

Watch out

Lab 00’s format table says bf16 has 7 mantissa bits. The docstring of kd_core.topk_forward_kl says bf16 has 8 mantissa bits. Both statements are in the course, both are about the same format, and neither is a typo.

The resolution: bf16 stores 7 mantissa bits and has 8 bits of significand precision, because the leading 1 of a normal float is implicit and does not occupy a bit. The same relationship holds for the other formats: fp32 stores 23 and has 24 bits of precision, fp16 stores 10 and has 11.

When the distinction matters, use the one that answers your question. Storage layout and machine epsilon depend on the stored count: bf16’s epsilon is , not . Rounding error and spacing arguments depend on the significand count: the spacing of bf16 values immediately below 1.0 is , because those values live in the binade and their spacing is . Section 2.8’s bug turns on that second number.

The habit worth building is to say which count you mean. “bf16 has 8 mantissa bits” will make a reader compute the wrong machine epsilon by a factor of two, and a factor of two in a precision budget is the difference between a clamp that works and a clamp that does nothing.

2.7.5 The dtype policy: model in bf16, loss in fp32#

The rule the course follows, and the one Hugging Face and TRL implement internally, is one line:

Run the model in bf16. Compute the loss in fp32.

The first half is forced by memory. A parameter in bf16 costs 2 bytes instead of fp32’s 4, so a 32-billion-parameter teacher occupies 64 GB rather than 128 GB, which is the difference between fitting on the reference machine alongside a student and not fitting at all. Chapter 8 does that arithmetic in full.

The second half is forced by precision, and the argument is the digit budget. bf16’s machine epsilon of means every stored value carries a relative error of up to about 0.4 percent, and a loss accumulates over tens of thousands of terms. Lab 00 measures the effect: a KL computed through fp32 log_softmax has relative error under against an fp64 reference, while the same KL through bf16 log_softmax is more than ten times worse, and the lab asserts both bounds.

Now ask what you were going to do with that loss. An ablation compares arms whose difference late in training is a few percent of a loss value, and a run’s improvement over its own previous checkpoint is smaller still. If the third digit is noise you cannot resolve either. The policy is not about correctness of the forward pass, which bf16 handles fine. It is about whether your loss is a usable instrument.

The policy costs one fp32 copy of the logit tensor at the loss site, which at batch 8, sequence 512, and vocabulary 151,936 is 2.5 GB. That cost is why people skip the upcast. When the copy does not fit, the move is not to compute the loss in bf16; it is to reduce the batch or chunk the vocabulary dimension, both of which preserve the answer.

One more consequence, because it returns in Chapter 10. If you cache teacher outputs in bf16, which you should for storage reasons, you are storing three good digits. Three digits is plenty for log-probabilities, which are ordinary-sized negative numbers, and nothing at all for probabilities near 1. The format is identical; what changed is where on the number line you put the quantity.

2.7.6 A prediction that failed#

Lab 00’s first exercise asks the reader to break bf16 on purpose: sweep the logit scale and find where bf16’s loss arithmetic first disagrees with fp32’s in the second significant digit. The premise is plausible, since bf16 keeps two to three significant digits and larger logits produce larger absolute rounding errors.

The solution notebook ran the sweep at scales 1, 2, 4, 8, 12, 20, and 40 against an fp64 reference, and the premise did not survive. bf16’s relative error stayed between about and throughout, peaking at scale 12 and getting better at scale 40. The third digit went bad at scales 2 and 12, with no monotone trend between. The second digit was never corrupted at any scale tested, and the notebook asserts it: maximum bf16 relative error below .

The reason is a cancellation the premise did not account for. Scaling the logits up grows the absolute rounding error, because float spacing is proportional to magnitude. It also grows the KL itself, from about 1 nat at scale 1 to about 115 nats at scale 40. Relative error is the ratio of those two growing quantities, and they grow at comparable rates, so it wanders inside the digit budget instead of climbing out. Averaging over the test problem’s 64 rows claws back a little more, which is why the measured error sits below machine epsilon rather than at it.

Field note

The corrected statement is not “bf16 gets relatively worse as logits grow.” It is “bf16 holds the loss to about three digits regardless of scale, and the third digit is already noise at realistic post-temperature logit scales.”

I am keeping the failed prediction rather than quietly replacing it, because the failure is more instructive than the result. The reasoning behind it was not sloppy. It was a correct statement about absolute error plus an unstated assumption that the quantity being measured holds still while the error grows, and that assumption is false here and in many precision arguments. When you reason about relative error, write down the numerator and the denominator and ask what each does under the change you are making.

The corrected version is still ample reason for the dtype policy. A loss whose third digit is noise cannot resolve the small late-training differences a distillation study is built to measure, at any scale rather than only at extreme ones.

2.8 A clamp that did nothing#

This is a real bug from the course’s own history, and the cleanest illustration in the book of why the previous section’s arithmetic is worth carrying around.

Chapter 10 covers top-k logit caching properly. The piece needed here: when you keep only the teacher’s top log-probabilities, you have to decide what to do about the missing mass, and one standard answer is to keep the total discarded mass as a single extra bucket. The loss then gains a term comparing the teacher’s tail mass against the student’s, where the student’s tail mass is one minus its mass on the retained tokens:

That needs care in floating point, because the sum can be very close to 1 and can exceed it by a rounding error. The obvious defense is to clamp the sum below 1 before subtracting, with a bound like , and to use log1p rather than composing a subtraction with a logarithm. That is what the implementation did.

Field note

The bug: in bf16, that clamp is a no-op, and the guarded expression produces -inf anyway.

Work through the numbers. On teacher-forced text the student’s mass on the teacher’s top- tokens routinely runs above 0.998. bf16 values immediately below 1.0 live in the binade , where the spacing is . The largest bf16 value below 1.0 is therefore 0.99609375, and the midpoint between it and 1.0 is 0.998046875. Round-to-nearest sends anything above that midpoint to exactly 1.0.

So a student mass of 0.999 is stored as exactly 1.0, and the clamp bound is also stored as exactly 1.0, being above the same midpoint. The minimum of two values that are both exactly 1.0 is exactly 1.0: the clamp compares against a bound that is not representable and therefore cannot bind. Then log1p(-1.0) is , the tail term is , and the loss is nan from the first step onward.

What made this hard to find is that the code looks correct. There is a clamp, there is a log1p, both are the textbook defenses, and a reviewer sees a guarded computation. The guard fails not because the logic is wrong but because the bound was chosen in the wrong precision: is a sensible margin in fp32, where the spacing near 1.0 is , and it is four orders of magnitude below the resolution of the format the tensor was in.

The fix is Section 2.7.5’s dtype policy applied at the top of the function rather than patched at the point of failure: upcast to fp32 first, whatever the student’s dtype, which costs one copy of a [batch, sequence, k] tensor. Raising the clamp margin to something bf16 can represent also stops the nan, and is worse, because it leaves a tail term carrying the resolution of a three-digit format.

The general lesson, which I now apply reflexively: an epsilon is a number relative to a precision. Any constant of the form is a statement about the format the computation runs in, and moving that computation to another dtype silently changes what the constant means. When you meet a magic small constant in numerical code, the first question is which dtype it was chosen for.

The listing below is the whole bug in six lines. Look at the comparison in the last two prints.

import torch

CLAMP = 1.0 - 1e-6                       # intended as "strictly below 1"
student_mass = torch.tensor([0.999])     # student's probability on the teacher's top-k

for dtype in (torch.float32, torch.bfloat16):
    s = student_mass.to(dtype)
    bound = torch.tensor([CLAMP], dtype=dtype)
    guarded = torch.minimum(s, bound)
    print(dtype, float(guarded), float(torch.log1p(-guarded.float())))

# torch.float32   0.99900001   -6.90777
# torch.bfloat16  1.0               -inf

The same expression, with the same guard, on the same input, is finite in one dtype and infinite in the other. Nothing in the source distinguishes the two cases, which is why this class of bug survives review.

2.9 What all of this costs#

A short accounting, in the spirit the rest of the book keeps.

Max-subtraction costs one reduction over the vocabulary dimension and is already inside your framework. Log space costs nothing either, since log_softmax and softmax do the same work; its price is conceptual, because the identities you know for probabilities (they sum to one, they multiply) become different ones (they logsumexp to zero, they add). The fp32 loss policy costs one copy of the logit tensor, 2.5 GB at [8, 512, 151936], which is one more reason Chapter 10’s top-k pipeline is attractive: it shrinks the tensor that has to be upcast.

What the policy gives up is the last few digits. fp32 relative error here lands near , fine for a training loss and not fine for a claim about a difference of between two arms. If you make that kind of claim, compute the measurement in fp64 and say so, which is what the labs do when they need a ground-truth reference.

2.10 What to be suspicious of#

The presentation above makes three things sound more settled than they are.

The set of 16-bit formats is not stable, and neither format was designed for this. fp16 came from graphics and bf16 from accelerator design, and both were compromises that happened to be available when large-model training needed them. Formats below 16 bits are already in serious use, including 8-bit floating-point variants and 4-bit schemes that are not IEEE floats at all: QLoRA’s NF4 assumes the stored values are approximately normally distributed rather than spread over a uniform exponent grid, and post-training quantization methods choose for themselves which values deserve precision.8910 Chapter 15 covers what those cost in quality. Section 2.7’s analysis applies to them in structure, and none of its numbers do.

“Compute the loss in fp32” rests on a threshold nobody has characterized properly. I have given you the digit-budget argument and one measurement. The argument is sound, and nobody has systematically mapped which distillation objectives at which scales degrade in bf16 and by how much. The rule is cheap enough that following it without that map is the right call, and you should know you are following an argument rather than evidence.

The relationship between measured entropy and anything you care about is loose. Entropy is cheap to compute and it is the standard early-warning signal for collapse. A run can also have healthy entropy and produce bad text, or declining entropy and produce good text. The work on when diversity actually collapses during post-training is active and does not agree with itself yet.11

2.11 Where this lands in the labs#

Lab 00 is where every claim in this chapter gets asserted rather than argued, on no GPU and with no training: the naive softmax on temperature-scaled logits, asserted nan; both paths to a log-probability, asserted -inf on one and on the other; a KL computed in fp64, fp32, bf16, and fp16 with the error ordering asserted; fp16 turning a probability of into an exact zero. What the lab does that this chapter cannot is let you change one number and watch the failure move. Set the temperature to 0.3 and find where the overflow starts, or change the to and watch the composed path start working. Lab 00’s first solution exercise is the bf16 sweep whose premise failed in Section 2.7.6, and the fifteen minutes you spend expecting the second digit to break are the point of it. Lab 01 §1 is Table 2.1, and it runs in under a second.

2.12 Exercises#

  1. A checkpoint’s largest logit at some position is 34.2, and you are about to compute a softmax at temperature 0.2 in fp16 without max-subtraction. Write the arithmetic out, say whether it overflows, and give the smallest temperature at which it would not. Repeat for fp32. Then say what happens to both answers under max-subtraction, and prove it from shift invariance.

  2. Using , predict the entropy of Table 2.1’s logit vector at before looking at Figure 2.3. State the condition under which the two-term expansion should be trusted, in terms of and , and say whether satisfies it.

  3. A colleague reports a run whose loss is stable at 2.417 for a hundred steps, reads inf at step 101, and is nan afterward, and that cutting the learning rate by 10x reproduces the failure at step 101 exactly. Give the most likely cause, name the line of code you would read first, and say what single measurement would confirm or refute your hypothesis before you change anything.

  4. Section 2.8’s clamp bound of fails in bf16. Someone proposes changing the bound to , which bf16 can represent. Say what that does to the tail term when the student’s true tail mass is , whether the resulting loss is biased upward or downward, and which of the two fixes you would ship.

  5. You inherit a logit cache whose README says it stores probabilities in fp16, top 32 entries per position. List what you can no longer determine from it that a cache of fp32 log-probabilities would have given you, and say for each item whether the loss is recoverable.

  6. Argue both sides of this: “since the top 64 tokens hold more than 99 percent of the teacher’s probability mass, discarding the rest costs at most one percent of the training signal.” Use Section 2.6 for the case against, and be specific about the quantity that makes the claim fail. Then say what you would measure on your own corpus to settle it there.



  1. Vocabulary sizes as priced in the course’s cache-cost table (Lab 01 §6). For the Qwen family’s tokenizer and model configuration see Qwen Team, “Qwen2.5 Technical Report,” arXiv:2412.15115 (2024). https://arxiv.org/abs/2412.15115 ↩

  2. The course’s student is from the SmolLM2 family: Loubna Ben Allal et al., “SmolLM2: When Smol Goes Big, Data-Centric Training of a Small Language Model,” arXiv:2502.02737 (2025). https://arxiv.org/abs/2502.02737 ↩

  3. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean, “Distilling the Knowledge in a Neural Network,” arXiv:1503.02531 (2015), §2. The softened-softmax formulation and the argument that the relative probabilities of incorrect answers carry the transferable structure are both there. https://arxiv.org/abs/1503.02531 ↩

  4. Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger, “On Calibration of Modern Neural Networks,” arXiv:1706.04599 (2017), ICML 2017. Temperature scaling as a single-parameter post-hoc calibration method, fitted on a validation set. https://arxiv.org/abs/1706.04599 ↩

  5. Ganqu Cui et al., “The Entropy Mechanism of Reinforcement Learning for Reasoning Language Models,” arXiv:2505.22617 (2025). The standard reference for entropy collapse in reinforcement learning with verifiable rewards, including the entropy-performance exchange law. https://arxiv.org/abs/2505.22617 ↩

  6. Renren Jin et al., “Revisiting Entropy in Reinforcement Learning for Large Reasoning Models,” arXiv:2511.05993 (2025), ACL 2026 Findings. https://arxiv.org/abs/2511.05993 ↩

  7. Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi, “The Curious Case of Neural Text Degeneration,” arXiv:1904.09751 (2019), ICLR 2020. Nucleus sampling truncates the distribution at a cumulative probability threshold, which is a statement about tail shape. https://arxiv.org/abs/1904.09751 ↩

  8. Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer, “QLoRA: Efficient Finetuning of Quantized LLMs,” arXiv:2305.14314 (2023), NeurIPS 2023. The 4-bit NormalFloat data type is built around an assumed distribution of the values being stored rather than around a uniform exponent grid. https://arxiv.org/abs/2305.14314 ↩

  9. Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh, “GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers,” arXiv:2210.17323 (2022), ICLR 2023. https://arxiv.org/abs/2210.17323 ↩

  10. Ji Lin et al., “AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration,” arXiv:2306.00978 (2023), MLSys 2024. https://arxiv.org/abs/2306.00978 ↩

  11. Constantinos Karouzos, Xingwei Tan, and Nikolaos Aletras, “Where does output diversity collapse in post-training?” arXiv:2604.16027 (2026); Longfei Yun et al., “The Price of Format: Diversity Collapse in LLMs,” arXiv:2505.18949 (2025). Both are preprints without a peer-reviewed venue at the time of writing. https://arxiv.org/abs/2604.16027 · https://arxiv.org/abs/2505.18949 ↩

  12. Rishabh Agarwal et al., “On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes,” arXiv:2306.13649 (2023), ICLR 2024. The on-policy setting where the student’s own rollouts are scored, and where entropy on those rollouts is the quantity worth watching. https://arxiv.org/abs/2306.13649 ↩

  13. John Schulman, “Approximating KL Divergence,” blog post, joschu.net, 7 March 2020, accessed 1 August 2026. http://joschu.net/blog/kl-approx.html ↩

  14. Jianping Gou, Baosheng Yu, Stephen J. Maybank, and Dacheng Tao, “Knowledge Distillation: A Survey,” International Journal of Computer Vision 129, no. 6 (2021): 1789-1819. https://arxiv.org/abs/2006.05525 ↩

  15. Xiaohan Xu et al., “A Survey on Knowledge Distillation of Large Language Models,” arXiv:2402.13116 (2024). https://arxiv.org/abs/2402.13116 ↩

  16. Woosuk Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention,” arXiv:2309.06180 (2023), SOSP 2023. The serving path that returns log-probabilities rather than probabilities, and the shape of that response, are covered in Chapter 15. https://arxiv.org/abs/2309.06180 ↩

Part I · Foundations

3

Measuring the Distance Between Two Distributions

At one position in one sequence, the teacher has a probability distribution over its vocabulary and the student has another one. Two vectors of 151,936 numbers each, both summing to one. Your job is to collapse the difference between them into a single scalar, because a scalar is what an optimizer can descend.

There are a lot of ways to do that, and the choice is not a formality. Two students trained on the same corpus, from the same initialization, against the same teacher, for the same number of steps, with the same learning rate, will produce visibly different text depending on which scalar you chose. One will hedge. One will commit. One will occasionally produce something the teacher would never have said. Chapter 6 runs that comparison as a controlled experiment and reports what happens. This chapter builds the objects the comparison is made of.

I want to be careful about the order here. Almost everything written about divergence choice in distillation starts from behavior: forward KL is mode covering, reverse KL is mode seeking, use reverse KL if you want crisp output. Those statements are true and I will get to them. But they are consequences of arithmetic, and if you learn the consequences without the arithmetic you will not be able to predict what a new objective does when someone publishes one next month. So this chapter stays with the arithmetic. Every behavioral claim I make here is derived from the shape of a function, and every one of them is checkable by hand on a four-token vocabulary.

Throughout, and throughout the rest of the book, is the teacher and is the student. That convention comes from the labs and it is worth fixing now, because half the confusion in this subject comes from people writing and not saying which letter is which model.

3.1 One identity, and everything that follows from it#

Chapter 2 defined entropy. Here are the other two quantities in the same family, all three measured in the same unit.

Every divergence below is in nats, as Chapter 2 defined them, because that is what log returns in every framework you will use and because it is the unit the entropy in the identity is already in.

For two distributions and over the same finite set of outcomes indexed by :

The first is entropy, the uncertainty in by itself. The second is cross-entropy, which is the quantity your training loop already minimizes: the average number of nats it costs to encode samples from using a code built for . The third is the Kullback-Leibler divergence.

Definition

KL divergence

, measured in nats. The extra cost, per outcome, of describing using a code optimized for rather than a code optimized for . It is zero exactly when , positive otherwise, asymmetric in its two arguments, and unbounded above.

The three are related by one identity, and it is the single most useful line in this chapter:

The derivation is two steps of algebra, and it is worth doing rather than quoting, because the two steps tell you what the identity means.

Rearranged, that is the identity. The first step splits the log of a ratio into a difference of logs. The second step recognizes each piece. Nothing else happens.

Read it as an accounting statement. Cross-entropy is the total bill. Entropy is the part of the bill you were always going to pay, because is genuinely uncertain and no model can remove uncertainty that lives in the target. KL is the remainder, the part that is your model’s fault. Drive KL to zero and you have paid exactly the irreducible price and not one nat more.

Now the consequence that makes the identity load-bearing for distillation. During training the teacher is fixed. Its parameters do not move, so is a constant with respect to everything the optimizer touches. Which means:

Minimizing cross-entropy against a fixed reference is minimizing forward KL against that reference. They differ by a constant, they have identical gradients, and they reach their minimum at the same parameters. Reporting one and optimizing the other is fine. Reporting cross-entropy and claiming it is a divergence is not fine, because cross-entropy has a floor at that varies from position to position, so a cross-entropy of 2.1 nats means something different on a high-entropy position than on a low-entropy one. If you want a number you can compare across positions, subtract the entropy and report KL.

Three corollaries fall out of that identity immediately, and I want all three on the table before the next section.

Hard-label training is distillation with a one-hot teacher. A one-hot distribution puts probability 1 on a single outcome. Its entropy is , because there is nothing uncertain about it. So the identity collapses to , and the ordinary next-token cross-entropy loss you have been using since your first language model is forward KL against a degenerate teacher. Distillation with a soft teacher is not a different objective. It is the same objective with a better reference distribution, and every piece of machinery in this chapter applies to both.1 What makes a reference distribution better is itself a variable rather than a given: a teacher whose training procedure flattened the structure among its wrong-answer probabilities transmits less through this channel than one that did not, which is why label smoothing in the teacher measurably degrades the student it produces.2 Chapter 5 reproduces that argument in full.

Perplexity is a KL statement wearing different clothes. Per-token perplexity is , read as “the model is as uncertain as if it were choosing uniformly among this many options.” Substituting the identity gives . When someone reports that a distilled student has 1.4x the perplexity of its teacher on a corpus, they have reported averaged in a particular way, and the KL in question is about nats. The perplexity ratio is the interpretable quantity; the absolute perplexity mixes in , which is a property of the data and not of your model.

KL is never negative. This is Gibbs’ inequality, and it follows from concavity of the logarithm:

where the inequality is Jensen’s applied to the concave function . Equality holds exactly when is constant across all with , which for two distributions means . So a loss of zero is achievable in principle, and any negative KL appearing in a training log is a bug, with no exceptions worth discussing. Usually it is a sampled estimator (Chapter 4) or a sign error.

3.2 The two conventions that make the sum well defined#

The formula has two edge cases, and both of them show up in real training runs rather than in textbook footnotes.

Where , the term is 0. The convention is , which is not arbitrary: the limit , so defining the term this way makes continuous at the origin. Operationally it means forward KL charges the student nothing at all for its behavior on tokens the teacher has ruled out. The student can put whatever it likes there, and the forward-KL loss will not notice. Hold onto that; it is half of §3.4.

Where and , the term is . There is no convention that rescues this one. The ratio diverges, the logarithm of it diverges, and the whole sum is infinite. The extreme case has a name.

Definition

Disjoint support

Two distributions have disjoint support when every outcome that one of them assigns positive probability to, the other assigns exactly zero. KL is infinite in both directions on disjoint support.

I used to think of the infinity as a mathematical corner case, a thing that happens on the whiteboard and never in float32. It happens constantly. Two mechanisms produce it. The first is cold start: a freshly initialized student, or a student that has been pruned down from its teacher and not yet recovered, will assign genuinely negligible probability to tokens the teacher is confident about. The second is the number format. Chapter 2 established that fp16 stores nothing smaller than about 6e-8, so in fp16 a probability of 2e-16 is exactly zero rather than a small number. A distribution that was perfectly healthy in fp32 arrives at the loss site in fp16 with holes in it, and the KL that was 14 nats becomes inf, and the gradient becomes NaN, and the run dies. That is the most common way a distillation run fails on the first day.

Lab 00 makes the infinity concrete by clamping the zeros rather than pretending they are not there. On the four-token pair and with zeros replaced by , the KL comes out at nats. Halve and the answer grows by . There is no value it converges to. The clamp chooses how large you would like your infinity to be, and its size is set by a constant nobody in the room chose deliberately.

3.3 KL is not a distance, and both failures matter#

A distance in the mathematical sense is a metric, and a metric has to satisfy four conditions: non-negativity, zero exactly on identical arguments, symmetry, and the triangle inequality.

Definition

Metric

A function that is non-negative, zero exactly when , symmetric (), and satisfies the triangle inequality ().

KL satisfies the first two and fails the last two. Both failures have operational consequences and both are demonstrable on distributions you can hold in your head.

3.3.1 Asymmetry, with numbers#

Take a four-token vocabulary. Let be the softmax of the logits , which is a teacher that is confident about token 0, and let be the softmax of , which is a student that has no opinion at all. Writing them out:

Compute both directions by hand. Forward first:

Now reverse:

The same pair of distributions, 1.38 nats one way and 4.61 nats the other, a factor of 3.35. And the asymmetry has a direction you can read off the arithmetic. The forward direction is dominated by a single moderate term, because concentrates almost all its weight on the one token where the ratio is only . The reverse direction is dominated by three terms of , because spreads weight onto three tokens where has almost nothing, and is a large number. Each direction sums over the weight of its first argument. §3.4 works through the same mechanism one term at a time.

3.3.2 The triangle inequality, with a triple you can verify in your head#

Asymmetry is the failure everyone knows about. The triangle failure gets less attention and it matters more for how you plan a project, because it invalidates a chain of reasoning that multi-stage pipelines invite you to make.

Here is the cleanest counterexample I know. Take a two-token vocabulary and three distributions:

Then:

The direct route costs 0.8318 and the route through costs . Going around is cheaper than going straight, by a factor of exactly two.

The factor of exactly two is not a coincidence and the derivation is short enough to do here. Let , , and for any . Then

and

Adding the last two, the terms cancel and expands to , leaving

So for this entire one-parameter family, the detour through the uniform distribution costs exactly half of the direct path, for every . The triangle inequality is not violated marginally on adversarial inputs; it is violated by a factor of two on an entire family of the most ordinary distributions imaginable. Lab 00 finds violations by random search, drawing 2000 triples from a four-token simplex and asserting that at least one violates, which it always does. The closed-form family above is why the search never comes up empty.

The consequence is a piece of reasoning you cannot use. Suppose you distill a 70B teacher into a 14B intermediate model, then distill the intermediate into a 1.7B student. Both measured KLs come out small, and you would like to conclude that is small. You cannot. Nothing supports that step, and the family above shows the error can be a clean factor of two in the friendliest case. If you want the conclusion you have to measure it directly against the teacher, which costs one forward pass of the teacher over your evaluation set and is worth paying for. Teacher-assistant setups that insert a mid-sized model between an enormous teacher and a small student are a standard recipe, and the prune-then-distill pipeline of Chapter 13 has the same shape.

Watch out

The chained-distillation argument is invalid for KL and it is invalid for every divergence in this chapter except total variation and the square root of Jensen-Shannon, which are the only two that are metrics. If your project plan contains a sentence of the form “A is close to B and B is close to C, therefore A is close to C,” and the closeness is measured by a KL, that sentence is doing no work.

3.4 Forward and reverse KL, stated as facts about the integrand#

Since KL is asymmetric, there are two of it, and both are used.

Definition

Forward KL

with the teacher first: the expectation, under the teacher, of the log ratio of teacher to student. Also called the mode-covering or zero-avoiding direction. In this book and in the course code, “forward” always means teacher first.

Definition

Reverse KL

with the student first: the expectation, under the student, of the log ratio of student to teacher. Also called the mode-seeking or zero-forcing direction.

Now the mechanical statement, which I want on the page before any word about training behavior.

Forward KL sums over . Look at what each term does at the boundaries.

Where the teacher has mass and the student does not, and , the term is . The student is charged without limit.

Where the student has mass and the teacher does not, and , the term is by the convention in §3.2. The student is charged nothing.

Reverse KL sums , and the same two checks come out the other way around.

Where the teacher has mass and the student does not, , the term is , because . Charged nothing.

Where the student has mass and the teacher does not, with fixed, the term is . Charged without limit.

That is the entire asymmetry, stated as two sentences with no metaphor in them:

Forward KL penalizes the student for putting near-zero mass where the teacher has mass. Reverse KL penalizes the student for putting mass where the teacher has none.

Everything anyone has ever written about mode covering and mode seeking is a restatement of those two sentences in the presence of a student too small to satisfy both constraints. If the student can represent the teacher exactly, both directions are minimized at the same place, namely , and the choice does not matter. The choice matters exactly when the student cannot, which is the situation you are always in when distilling.

2026-08-01T07:26:30.655902 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 0.00 0.04 0.08 teacher p(x) valley: p(29) = 1.6e-06 mode 1 mode 2 4 15 28 44 56 student centre mu 4 6 8 10 12 forward KL(p||q), nats forward minimum: mu = 28.0, KL = 3.6657 forward KL reverse KL 1 2 3 4 reverse KL(q||p), nats KL = 0.9590 KL = 1.1588
Figure 3.1 Forward and reverse KL, evaluated over the same one-parameter family of student distributions against the same bimodal teacher, disagree about which student is best: the forward objective is minimized in the valley between the teacher's two modes, where the teacher has almost no mass, while the reverse objective has a local minimum on each mode and no stationary point in between.

Figure 3.1 makes the two sentences visible without any training happening. The teacher is a fixed bimodal distribution over 60 outcomes, with peaks near index 15 and index 44 and a deep valley between them. The student family is a single bump of fixed width whose center is the only free parameter, so it is constitutionally incapable of having two peaks. Sweeping across the whole range and evaluating both divergences at every gives two curves over the same set of candidates. The forward curve has one minimum, at , which is in the valley: the teacher assigns a probability of about there, so this is the student that the forward objective prefers even though it puts its own peak on a region the teacher considers nearly impossible. The reverse curve has two local minima, one sitting on each of the teacher’s modes, and a large hump between them.

Read that as a statement about the objectives, not about optimizers. Nobody has run a training loop. Two functions have been evaluated on the same 200 candidate students and they rank those candidates differently, with different argmins. The training consequences, what happens to a real student’s entropy, its output diversity, its length, and its confidence, are Chapter 6’s subject and they are measured there rather than asserted. The one-sentence preview: forward KL tends to produce a student that hedges, and reverse KL tends to produce one that commits, and which of those you want is a product decision rather than a mathematical one.3 The runaway version of commitment has a name, entropy collapse, and a literature of its own that grew up around reinforcement-learning fine-tuning before distillation adopted the vocabulary.4

3.5 One formula that contains all of them#

Every divergence in this book is an instance of a single expression. Pick a function and define

Definition

f-divergence

, where the generator is convex on and satisfies . Introduced independently by Csiszár and by Ali and Silvey in the 1960s, which is why the family is sometimes called the Ali-Silvey-Csiszár divergences.

Definition

Generator

The function that determines an f-divergence. It must be convex and satisfy . Two generators that differ by an affine term define exactly the same divergence.

The two conditions are doing specific work, and neither is decoration.5

Why . The ratio equals 1 exactly where the two distributions agree. The condition says agreement costs nothing, which is what makes .

Why convexity. Convexity buys non-negativity, through Jensen’s inequality applied in one line:

So convexity plus gives you exactly the two properties every divergence needs: it is zero on identical inputs and it is never negative. Nothing more is required, and nothing more is guaranteed. Symmetry is not guaranteed. The triangle inequality is not guaranteed. Boundedness is not guaranteed. Which of those you get is decided by the shape of , and §3.6 shows how to read them off.

There is one piece of gauge freedom worth knowing. Replacing by for any constant changes nothing, because

The affine term integrates to zero against any pair of distributions. This looks like a technicality and it is the trick behind Schulman’s k3 estimator, which Chapter 4 derives: you add an affine term that is free in expectation in order to change the variance of a sampled estimate without changing what is being estimated.6

Here is the family, written out. The direct form is the way you would implement it; the generator is what you would look for in a paper.

Table 3.1 Six divergences and their generators.

Divergence Direct form Generator
Forward KL,
Reverse KL,
Total variation
Chi-squared
Squared Hellinger
Jensen-Shannon ,

Check a couple of rows by hand to see that the correspondence is real rather than notational. Forward KL: , the cancels. Reverse KL: , a sign flip inside the log. Chi-squared: , expand the square and multiply through. The other three are the same kind of manipulation.

Notice that reverse KL is generated by while forward KL is generated by . The family is closed under swapping the arguments: if generates , then the conjugate generator generates . Apply it to and you get , which is exactly the reverse-KL row. So forward and reverse KL are not two unrelated objectives that happen to be similar; they are one generator and its conjugate.

The reason to hold the family in your head rather than a list of formulas is that it turns reading a paper into a mechanical operation. When a new distillation objective is announced, the first move is to find its generator, because two properties then follow without running anything: which direction of disagreement it charges most heavily, and whether it is bounded. The 2026 on-policy distillation literature organizes methods this way, and there is an explicit line of work on minimizing general f-divergences for sequence-level knowledge distillation.78

Verification is the other reason. The generator form and the direct form are different code paths computing the same number, so running both and asserting agreement catches an entire class of mistakes: a sign error in a hand-written direct formula, a mixed-up argument order, a normalization factor of dropped from Jensen-Shannon, an implementation that silently computes reverse KL when its argument name says forward. That last one is not hypothetical. The course’s own kd_core.kl_divergence takes the student first and the teacher second, while its direction argument names the KL order, so kl_divergence(student, teacher, direction="forward") computes . Reading argument order as KL order gets you the wrong objective with no error and no warning.

The listing below is the check in its smallest honest form. Watch that every generator is a one-liner and that the direct formulas share no code with them.

import torch

def f_divergence(p, q, f):
    """Generic f-divergence: sum_i q_i f(p_i / q_i), over the last axis."""
    return (q * f(p / q)).sum(-1)

def kl(a, b):
    return (a * (a.log() - b.log())).sum(-1)

torch.manual_seed(3)
p = torch.softmax(torch.randn(6, 48), -1)
q = torch.softmax(torch.randn(6, 48), -1)
m = 0.5 * (p + q)

cases = [
    ("forward KL",   lambda t: t * t.log(),                    kl(p, q)),
    ("reverse KL",   lambda t: -t.log(),                       kl(q, p)),
    ("total var",    lambda t: 0.5 * (t - 1).abs(),            0.5 * (p - q).abs().sum(-1)),
    ("chi-squared",  lambda t: (t - 1) ** 2,                   ((p - q) ** 2 / q).sum(-1)),
    ("sq Hellinger", lambda t: 0.5 * (t.sqrt() - 1) ** 2,      0.5 * (p.sqrt() - q.sqrt()).pow(2).sum(-1)),
    ("Jensen-Shannon", lambda t: 0.5 * (t * t.log() - (1 + t) * ((1 + t) / 2).log()),
                                                              0.5 * kl(p, m) + 0.5 * kl(q, m)),
]
for name, f, direct in cases:
    assert torch.allclose(f_divergence(p, q, f), direct, atol=1e-6), name
    print(f"{name:>15}: {float(f_divergence(p, q, f).mean()):.6f}")

Six divergences, one summation, and every one of them agrees with an independently written closed form to within floating-point noise. Lab 00 runs this check as an assertion so that it fails loudly if either path ever drifts, and Solutions 00 extends the same check with a chi-squared row as an exercise.

3.6 Reading behavior off the generator#

Two numbers determine most of what an f-divergence does, and both are limits of at the edges of its domain.

The first is , the value of the generator as the ratio goes to zero. That is the regime where the teacher has no mass and the student does: with fixed, so , and the term contributes .

The second is the recession constant . That is the regime where the student has no mass and the teacher does: with fixed, so , and the term tends to . Writing $q_i f(p_i/q_i) = p_i \cdot \frac{f(t)}{t}t = p_i/q_i$ makes the limit visible.

Table 3.2 What each generator does at the boundaries, and what follows.

Divergence Value on disjoint support Bounded?
Forward KL no
Reverse KL no
Total variation yes, by 1
Chi-squared no
Squared Hellinger yes, by 1
Jensen-Shannon yes, by

The first two rows are §3.4 restated in the new notation, and they now read as a single line each. Forward KL has and : free where the teacher is silent, infinite where the student is. Reverse KL has exactly the opposite pair.

The last column follows from the first two by an argument worth doing once. Suppose and have disjoint support. Split the sum into the outcomes where lives and the outcomes where lives. On the first group , so each term contributes , and summing gives since sums to one. On the second group , so each term contributes , and summing gives . Total:

For total variation that is . For Jensen-Shannon it is . Those are exactly the familiar bounds, and now they are derived rather than memorized, and the derivation tells you they are attained rather than approached in the limit. It also gives the general criterion.

Definition

Bounded divergence

An f-divergence whose value cannot exceed a finite ceiling for any pair of distributions. The condition is that both and are finite, in which case the maximum is , attained on disjoint support. Jensen-Shannon, total variation, and squared Hellinger are bounded; forward KL, reverse KL, and chi-squared are not.

2026-08-01T07:26:31.408499 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 0.0 0.5 1.0 1.5 2.0 2.5 3.0 t = p(x) / q(x) 0 1 2 3 4 generator f(t) f(1) = 0 for all six forward KL reverse KL total variation chi-squared squared Hellinger Jensen-Shannon student overweights (t < 1) student underweights (t > 1)
Figure 3.2 The six generators on the same axes: all pass through zero at t=1 and all are convex, but they charge disagreement in visibly different ways, and the two boundary limits f(0+) and f(t)/t as t grows determine whether the divergence is bounded.

Figure 3.2 puts the six generators on one pair of axes over . Three things are visible at a glance and none of them are visible from Table 3.1. Every curve touches zero at and curves upward on both sides, which is the two-condition definition drawn. The curves on the left of , the region where the student overweights relative to the teacher, are ordered completely differently from the curves on the right, which is the asymmetry made geometric. And climbs to a vertical asymptote as while approaches zero smoothly, which is the single picture that makes the forward-reverse distinction stick.

3.7 Total variation distance#

Definition

Total variation distance

, equivalently over all events . The largest amount by which the two distributions can disagree about the probability of anything. Symmetric, bounded by 1, and a genuine metric.

The two forms in that definition are worth connecting, because the second one is what makes TVD interpretable and the first one is what you compute.

Take any event , meaning any subset of the vocabulary, and ask how much the two distributions disagree about it: . To make that as large as possible, put into exactly the outcomes where is positive, since including a negative term would cancel some of your total. Call that set . Then

Now note that the positive and negative parts of have equal total size, because both distributions sum to one, so forces . Each half is therefore exactly half of . So the maximum disagreement over all events equals , which is the computational form. The factor of that looks arbitrary in the formula is what converts an L1 norm into a probability.

That interpretation is the reason TVD is the divergence to quote when you want to make a claim someone can act on. A TVD of 0.03 between teacher and student at a given position means: for any question you might ask about the next token, any set of tokens you care to define, the two models’ answers differ by at most three percentage points. No vocabulary size appears in that sentence. It holds for a 32,000-token vocabulary and a 256,000-token one alike.

TVD is a metric, satisfying all four conditions including the triangle inequality, which makes it one of exactly two quantities in this chapter you may legitimately chain. It is also almost never used as a training loss, and the reason is its generator. The derivative of is for and for , so the loss carries information about the sign of each disagreement and nothing about its size. A student that is off by 0.4 on some token receives the same gradient signal as one that is off by 0.001. Compare with forward KL, whose gradient with respect to the student’s logits is exactly the residual (Chapter 5 derives this), which is large where the disagreement is large and vanishes as the disagreement closes. TVD’s L1 structure does show up in distillation, though not as the primary objective: the Universal Logit Distillation loss for cross-tokenizer distillation sorts both probability vectors and takes an L1 distance between them, precisely because the sorted L1 distance survives a change of vocabulary in a way that a position-wise KL does not.9 Chapter 14 covers that construction.

3.8 Chi-squared, and why it is the tail-sensitive one#

Definition

Chi-squared divergence

, generated by . Equal to the variance, under , of the importance ratio . Unbounded above and dominated by whichever single outcome has the largest ratio.

The identity in the second sentence is the one to remember, and it takes three lines. Write and treat as a random variable under . First, , always, for any pair of distributions. Second, expand the generator form:

using in the last step. So the chi-squared divergence is literally the variance of the importance ratio. That fact reappears in Chapter 4 as the thing that governs whether a sampled KL estimator is usable, and it is the reason chi-squared is worth understanding even though nobody trains against it.

Chi-squared sits above KL, and the relationship is a two-line consequence of Jensen’s inequality. Start from the definition of forward KL as an expectation under and push the logarithm outside:

Then note that , since the cross term is . So

the last step because . Chi-squared upper bounds KL, and the logarithm in the middle expression is the entire difference between them. KL charges one power of a log-ratio; chi-squared charges the ratio itself, squared. On a token where the student is off by ten orders of magnitude, that difference is ten orders of magnitude.

Solutions 00 measures the gap on a construction designed to isolate it. Teacher logits are , which puts 0.998995 of the teacher’s mass on token 0. The student’s logit on token 0 is walked from down to while everything else stays at zero, so the log-ratio on that one token grows by 25 nats across the sweep and nothing else in the problem changes.

Table 3.3 The unboundedness race, from Solutions 00 Exercise 3. One token’s log-ratio grows by 25 nats; watch what each divergence does about it.

Student logit on token 0 Forward KL JSD
6.09 0.6812
11.08 0.6891
21.07 0.6892
31.06 0.6892

Forward KL grows by a factor near 5 across the sweep, which is linear in the logit gap because KL charges and grows linearly with the logit difference by construction. Chi-squared grows by more than ten orders of magnitude, which is exponential in the logit gap because . Jensen-Shannon does not move, because it has already saturated at its ceiling of and has nowhere to go. Three divergences, one moving variable, three qualitatively different responses.

The verdict Solutions 00 records is not to train against chi-squared, and the reason is in the moment form. The gradient of is dominated by whichever token currently has the largest ratio, squared. A single rare token where the student badly lags the teacher contributes a term exponentially larger than everything else in the sum combined, and one optimizer step later a different token plays that role. The compounding problem is that estimating from samples requires averaging , whose variance is governed by the fourth moment of , and Chapter 4 shows that even estimators involving to the first power are already in trouble whenever the models disagree. A divergence whose value is the variance of a quantity that was already too noisy to estimate is a poor candidate for a training signal.

3.9 Jensen-Shannon, and the metric that lives inside it#

Definition

Jensen-Shannon divergence

where is the equal mixture. Symmetric in its arguments, bounded above by (one bit), zero exactly when , and finite even on disjoint support.

The construction fixes both of KL’s structural problems at once by routing through the mixture. Symmetry is immediate from the definition, since swapping and leaves unchanged and swaps the two terms of a sum. Finiteness comes from the mixture too: wherever , the mixture has , so the ratio is at most 2 and its logarithm is at most . No term can blow up, because the mixture cannot have a hole anywhere either distribution has mass. That argument also gives the bound directly:

which agrees with the general calculation from §3.6. In bits the bound is exactly 1, which is where the “JSD is at most one bit” phrasing comes from. Equality holds exactly on disjoint support, which is the useful shape: total disagreement registers as a specific finite number rather than as an infinity, and every partial disagreement registers as something less than that number.

JSD is not a metric. It is symmetric and zero on identical inputs, but it fails the triangle inequality. Its square root does not.

is a metric on the space of probability distributions. Endres and Schindelin proved this directly, and Österreicher and Vajda established it as the case of a broader family of metric divergences, both in 2003.1011 The result generalizes: powers are metrics for .12 That is the reason JSD shows up whenever someone needs a number that behaves like a distance, such as clustering checkpoints from a training run or plotting how far a model has drifted from its initialization. Those operations assume a triangle inequality somewhere, and is one of the very few things in this chapter that supplies one.

Lab 00 checks the claim the same way it checks the KL violation, by drawing 2000 random triples from a four-token simplex and testing the triangle inequality on each. KL produces violations. The square root of JSD produces zero violations at a tolerance of . That is not a proof and the lab does not present it as one; it is a guard, so that if someone edits the JSD implementation and breaks it, the notebook fails on the next run.

3.9.1 The generalized version, and the parameter that interpolates#

JSD as defined uses an equal mixture. Replacing with a free parameter gives a family that contains both KL directions as limits, and this is the object Chapter 6 spends real time with.

Definition

Generalized Jensen-Shannon divergence

with and . At it is ordinary JSD. As it degenerates to zero, but ; as , .

The limits are the part worth checking rather than believing, so here is why they hold. As the mixture tends to . The second term goes to , and it does so quadratically in because KL behaves quadratically near its minimum. The first term is . So the whole thing vanishes linearly in with slope , which is what dividing by recovers. The end is the mirror image.

Table 3.4 on the fixed pair from Lab 01, teacher logits and student logits , for which and nats.

0.001 0.001838 1.8379 0.0018
0.01 0.018088 1.8088 0.0183
0.25 0.310710 1.2428 0.4143
0.5 0.412470 0.8249 0.8249
0.75 0.334102 0.4455 1.3364
0.99 0.022527 0.0228 2.2527
0.999 0.002318 0.0023 2.3180

Read down the third column and it climbs toward 1.8412, the forward KL, as falls. Read down the fourth column and it climbs toward 2.3256, the reverse KL, as rises. Read the second column and notice that itself is largest in the middle and vanishes at both ends, which is the trap: the raw divergence value does not tell you which direction you are heading toward, and two very different settings can produce the same loss magnitude.

I am defining here and stopping. The convention question, meaning which end of the range corresponds to which KL direction in which library, is genuinely inconsistent across the literature and the tooling, and getting it backwards trains the opposite objective with no error, no warning, and a perfectly normal-looking loss curve. That is Chapter 6’s problem and it gets a full treatment there, including the numerical procedure for determining which convention a library you have inherited actually implements. What matters here is that the object exists, that it is a one-parameter path through the divergence family connecting the two KL directions, and that the two limits above are the way you identify which end is which.

3.10 Bounded and unbounded, and what the bound buys you#

Table 3.2 sorted the six divergences into bounded and unbounded. Here is why the sorting matters during a training run.

Take a batch of 4,096 supervised positions. At 4,095 of them, teacher and student agree reasonably well and the per-position forward KL is around 0.05 nats. At one position, the student has assigned near-zero probability to a token the teacher is confident about. Lab 01 constructs exactly this position: an eight-token vocabulary, a teacher logit of 12 on token 3 so the teacher’s probability there is 0.99996, and a student logit of on the same token so the student’s probability there is .

At that position the forward KL is 31.94 nats, because the log-ratio on token 3 is and the teacher’s mass on that token is essentially all of it. The batch mean becomes instead of 0.05, a 16% jump from a single position out of four thousand.

Now change the student’s logit from to . Nothing else about the run changes; this is a student that has become slightly more wrong at one position out of four thousand. The forward KL at that position becomes 301.9 nats and the batch mean becomes 0.1237, which is two and a half times the healthy value. Change it to and the position contributes 3001.8 nats, dragging the batch mean to 0.78: more than fifteen times healthy, from one position in four thousand. Push the student logit to about and that single position carries the batch mean past 1.0 on its own. There is no value the loss converges to, because there is no ceiling on the term, and the position’s share of the mean grows linearly in the logit while everything else stays where it was.

Under JSD the same position contributes 0.6929 nats and cannot contribute more, no matter how far the student’s logit falls, because 0.6931 is the ceiling. A batch mean of against a healthy 0.02 is a 1% jump. Total disagreement at one position out of four thousand moves the batch loss by one percent, and that is the safety valve: no single position can dominate the batch.

Field note

I had the mechanism wrong for a while, and the error is worth writing down because the correct version is more interesting than the version I believed.

My assumption was that an unbounded loss means an unbounded gradient, so the loss spike and the training instability were the same event. That is false for forward KL. Chapter 5 derives the gradient of with respect to the student’s logits and it is exactly , the residual between the two distributions. Both are probabilities, so the residual is bounded in no matter how catastrophic the disagreement. At the position above, the gradient on token 3 is , which is the largest it can ever be, and it would be whether the student’s logit were or .

So why does the unboundedness matter? Three reasons, and they are all real.

First, the loss value is what you monitor, what your early-stopping rule reads, and what your gradient-clipping and loss-scaling thresholds were calibrated against. A metric that can move by a factor of twenty because of one position out of four thousand is a bad thermometer, and Chapter 8 is about reading these curves.

Second, once the student’s probability underflows to exactly zero, which fp16 accomplishes at and fp32 at around , the loss stops being a large number and becomes inf, and inf times anything in the backward pass is NaN. The gradient bound does not save you from a NaN.

Third, and this is the part that changed my mind, the bound on the gradient is specific to forward KL. For reverse KL the logit gradient is , which contains an unbounded log-ratio and is genuinely unbounded when the teacher has a hole where the student has mass. For chi-squared it is worse: the gradient contains , which diverges as the student’s probability falls. So “unbounded loss, bounded gradient” is a property of one row of Table 3.2, not of the family, and I had generalized from the one case I had checked.

2026-08-01T07:26:32.360934 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 1 0 − 1 1 0 − 3 1 0 − 5 1 0 − 7 1 0 − 9 1 0 − 1 1 overlap parameter eps (decreasing to the right) 0 5 10 15 20 25 forward KL, nats forward KL has no ceiling: at eps = 1e-8 it is 17.03 nats, while JSD sits at 0.6931 and TVD at 1.0000 0.00 0.25 0.50 0.75 1.00 JSD (nats) and TVD ln 2 = 0.6931 1 JSD TVD
Figure 3.3 As two distributions are pushed toward disjoint support, forward KL grows without limit while JSD and TVD saturate at their ceilings of ln 2 and 1, which is the entire content of the bounded-divergence safety valve.

The safety valve is not free, and the cost is the mirror image of the benefit. A bounded divergence cannot tell you how bad your worst position is. Once JSD has saturated at , a student that is wrong by 14 nats and a student that is wrong by 1400 nats produce the same loss and, more to the point, nearly the same gradient. If the thing you need to fix is a small number of catastrophically wrong positions, a bounded objective has stopped providing signal about exactly them. This is the same trade in both directions: the property that keeps one bad token from destroying your batch is the property that keeps you from noticing it.

3.11 Pinsker’s inequality, and the direction it does not run#

Everything above says the divergences differ. Pinsker’s inequality is the one theorem that ties two of them together, and it is worth stating carefully because it is quoted carelessly.

Read it as a promise. Drive KL down and every event probability the student assigns converges to the teacher’s at a known rate, for every event simultaneously. A KL of 0.02 nats gives , so no matter what set of tokens you ask about, the two models’ probabilities for that set differ by at most 10 percentage points. No vocabulary size enters. That is a genuinely useful thing to be able to say about a trained student, and it is the main reason to care about a measured KL as a number rather than as a loss.

The citation on this one requires care.13 Pinsker’s 1964 book gives the inequality and the name, but with a weaker constant than the one everybody uses. The constant, meaning the form above, is due independently to Csiszár in 1967 and Kullback in 1967, with Kemperman arriving at it in 1969, which is why the result is properly called the Csiszár-Kullback-Pinsker inequality. Citing Pinsker alone for the constant is technically wrong, and it is wrong in almost every machine-learning paper that cites it. Sharper refinements exist.14

The bound is also loose almost everywhere, which is worth knowing before you rely on it. On the symmetric two-token family and , the TVD is exactly and the KL is , which for small expands to , giving . So the bound is attained in the limit along that family, and Pinsker cannot be improved as a universal statement. But sample 3,000 random pairs of 16-outcome distributions and the ratio rarely exceeds 0.9 and sits near 0.6 in the middle of the distribution. Pinsker gives you a real guarantee and it typically overstates the TVD by a substantial factor.

2026-08-01T07:26:33.996169 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 1 0 − 6 1 0 − 5 1 0 − 4 1 0 − 3 1 0 − 2 1 0 − 1 1 0 0 1 0 1 1 0 2 1 0 3 KL(p||q), nats 0.00 0.25 0.50 0.75 1.00 total variation distance Pinsker's bound TVD = sqrt(KL / 2) symmetric two-token family: attains the bound only as d -> 0 Pinsker in reverse (Table 3.5): TVD 0.1 -> 0.001 while KL 9.7 -> 1000. No bound runs in this direction. 3,000 random pairs on 16 outcomes, every one below the bound
Figure 3.4 Pinsker's inequality drawn over random distribution pairs: every point falls below the bound, the bound is attained only along a narrow symmetric family as the two distributions converge, and the horizontal spread at any fixed TVD shows there is no bound in the other direction.

3.11.1 Pinsker in reverse: there is no such inequality#

The important thing about Pinsker is the direction it does not run. There is no bound of the form for any function . A tiny TVD certifies nothing whatsoever about KL, and the gap can be made arbitrarily large.

Solutions 00 constructs the counterexample explicitly, and the construction is worth carrying around because it is the compressed form of a real failure mode. Work on a two-token vocabulary. Let

The teacher puts a small mass on token 2. The student puts an astronomically smaller mass there. Now compute both quantities.

TVD sums absolute differences, so , because is negligible. As , the TVD goes to zero at exactly the rate of .

KL sums mass times log-ratio, and the log-ratio on token 2 is . So the KL is dominated by

The mass vanishes and the divergence diverges, on the same sequence, because one of them is linear in the misplaced mass and the other is linear in the misplaced mass times an unbounded log-ratio.

Table 3.5 Pinsker in reverse, from Solutions 00 Exercise 2. TVD falls to while KL climbs to on the same sequence of distribution pairs.

TVD KL (nats) KL / TVD
0.1 0.1000 9.67 96.8 2.20
0.03 0.0300 33.20 1,107 4.07
0.01 0.0100 99.94 9,994 7.07
0.003 0.0030 333.31 12.91
0.001 0.0010 999.99 22.36

The last column is Pinsker’s bound, and it holds at every row, and it is useless at every row: it promises TVD is at most 22.36, and TVD is 0.001. That is the point. Pinsker constrains TVD from above given a small KL, and here the KL is not small, so the bound has nothing to say.

There is a numerical detail in the construction that is a lesson in itself. The quantity underflows every float format almost immediately: at it is , which is zero in fp64. So the KL cannot be computed by forming and taking a ratio. It has to be computed from directly, in log space, which is exactly the discipline Chapter 2 established for a different reason. The solution notebook cross-checks the log-space formula against a plain tensor computation at , where is still representable in fp64, and asserts agreement to a relative difference below , so the shortcut is verified rather than trusted.

The listing below is the construction in the form that survives the number format. Watch that never appears anywhere except inside log1p, where its underflow to zero is harmless.

import math

def pinsker_pair(eps):
    """p = [1-eps, eps], q = [1-a, a] with log a = -1/eps^2. Returns (TVD, KL)."""
    log_a = -1.0 / eps ** 2
    a = math.exp(log_a)                     # underflows to 0.0 for small eps; that is fine
    tvd = eps - a                           # 0.5 * (|p1-q1| + |p2-q2|)
    kl = (1 - eps) * (math.log(1 - eps) - math.log1p(-a)) \
       + eps * (math.log(eps) - log_a)      # log_a used directly; a is never divided by
    return tvd, kl

for eps in (0.1, 0.03, 0.01, 0.003, 0.001):
    tvd, kl = pinsker_pair(eps)
    print(f"eps={eps:<6} TVD={tvd:9.6f}  KL={kl:9.2f}  ratio={kl/tvd:12.1f}")
    assert tvd <= math.sqrt(kl / 2), "Pinsker still holds, and is still useless here"

The output is Table 3.5, and the assertion at the bottom is the part that makes the point: Pinsker is never violated along the sequence. Both things are true at once. The inequality holds, and it tells you nothing you wanted to know.

The operational conclusion is a sentence you should be able to produce on demand. A measured TVD of 0.001 between student and teacher certifies nothing about their KL, because the student may have assigned an astronomically small probability to a rare token the teacher still cares about, and that single token carries unbounded KL while contributing 0.001 to the TVD. This becomes concrete in Chapter 10, where a top- logit cache discards the teacher’s tail by construction: the tail is exactly where tiny masses with enormous log-ratios live, so any cache or comparison that truncates the tail is blind to unbounded KL hiding inside it.

3.12 What all of this is for#

The choice of divergence is a design decision with measurable consequences, and it is one of the few in this subject where the mechanism is fully understood in advance. You are not guessing. Given a generator you can say, before writing any code, whether the objective will charge the student more for missing the teacher’s mass or for inventing its own, whether one bad position can dominate a batch, whether the resulting number can be chained through an intermediate model, and whether a sampled estimate of it will have usable variance.

What you cannot say in advance is which of those properties you want, because that depends on what the student is for. A student that will be sampled from at temperature 1 to produce diverse completions, or used as the starting point for a reinforcement-learning stage that needs multiple plausible continuations left intact, wants a different objective from a student that will be run greedily on a narrow structured task. The literature reflects the split rather than resolving it: MiniLLM argues for reverse KL on the grounds that a small student should not be forced to cover regions of the teacher it cannot represent, GKD makes the interpolation parameter itself a design knob alongside the on-policy/off-policy mix, and DistiLLM introduces a skewed variant specifically to control how much of each direction you get.151617 These are not competing claims about a fact. They are different answers to a question about products. The same reverse-KL structure appears one field over, in preference optimization, where a policy is regularized toward a fixed reference distribution, which is worth knowing because the failure modes transfer.18 The survey literature organizes the whole method space if you want a map before the details.19

A caution before Chapter 6 does the work. It is tempting to read the divergence choice as a knob that controls how faithfully the student reproduces the teacher, with the only question being which kind of infidelity you prefer. The evidence does not support the clean version of that reading: students trained by distillation frequently fail to match their teachers on held-out data by a margin their improved generalization does not explain, and improving the optimization does not close the gap.20 Whatever the divergence is doing, “it determines how closely the student copies the teacher” is an incomplete account of it, and Chapter 6’s measurements are on generated text for exactly that reason.

Chapter 6 runs the ablation that makes the answer measurable for a specific student on a specific task, with fixed seeds, fixed data, fixed steps, and measurement on generated text rather than on loss values. Chapter 4 handles the problem that comes first in practice, which is that you almost never have both full distributions in hand and must estimate a divergence from samples, at which point the variance properties matter more than the boundary behavior.

3.13 Where this lands in the labs#

Lab 00 §3 through §6 is this chapter executed with assertions attached. The lab does three things the chapter cannot: it verifies the six-generator table by computing every divergence twice through independent code paths and asserting agreement, so the table is a checked claim rather than a transcription; it finds the triangle-inequality violations by random search over 2,000 triples and confirms that survives the same search with zero violations; and it sweeps 3,000 random pairs to confirm the JSD, TVD, and Pinsker bounds hold empirically over a range of peakedness. Solutions 00 Exercise 2 builds the Pinsker-in-reverse table, including the log-space computation the construction forces, and Exercise 3 adds chi-squared to the generator check and races it against KL and JSD on a single moving logit. If you have time for one cell, run the disjoint-support demo at the end of §6: it takes under a second and shows KL at 26.9 nats next to JSD at exactly on the same pair of distributions.

3.14 Exercises#

  1. Prove that when is fixed, and state precisely which step of your proof fails when is not fixed. Then consider self-distillation, where the teacher is an earlier checkpoint of the same model: is fixed during a training step? Is it fixed across a training run? Say what quantity you would have to monitor to detect the difference from a loss curve alone.

  2. §3.3.2 showed that for , , , the detour through costs exactly half the direct path, for every . Derive this. Then determine whether the factor of two is special to the uniform midpoint: replace with for arbitrary and find the set of for which the triangle inequality is violated. State what happens as and explain why that is the answer you should have expected.

  3. Using only the argument from §3.6, compute the maximum value of the squared Hellinger divergence as defined in Table 3.1, and verify your answer against the direct formula evaluated on a disjoint pair. Then do the same for the generator : is it a legal generator, and if so, is the divergence it produces bounded? Say which of the six divergences in Table 3.1 it most resembles in behavior and why.

  4. Someone hands you an objective with generator . Before computing anything numerically, state: whether it is a legal generator; whether it is bounded and if so by what; which direction of disagreement it charges more heavily; and whether it is symmetric under swapping and (use the conjugate from §3.5). Write your four predictions down, then verify them.

  5. A colleague reports that their distilled student has a total variation distance of 0.002 from the teacher, averaged over 50,000 held-out positions, and concludes the two models are interchangeable for production purposes. Using §3.11.1, describe a concrete situation in which that number is accurate and the student is nonetheless badly wrong, name the single additional measurement that would settle it, and say why that measurement is the expensive one to obtain.

  6. A forward-KL distillation run trains normally for 800 steps at a loss around 0.4 nats, then spikes to 8.1 at one step, comes back to 0.5 for a few steps, then goes to NaN and stays there. The model is held in bf16 and the loss is computed in bf16. Using §3.2, §3.10, and Chapter 2’s number-format material, give the most likely sequence of events, and name two changes to the configuration that would each have prevented it. Then say which of the two you would prefer and what it costs.

  7. You are asked to review a plan that distills a 70B teacher into a 14B intermediate, then the intermediate into a 1.7B student, and reports and as evidence that the student is close to the teacher. Write the objection in three sentences. Then propose the smallest change to the measurement protocol that would make a defensible claim, and state what it costs in teacher forward passes. Finally, say what claim you could legitimately make using instead, and whether it is a claim anyone would care about.



  1. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean, “Distilling the Knowledge in a Neural Network,” arXiv:1503.02531 (2015), §2. https://arxiv.org/abs/1503.02531. The soft-target objective is a cross-entropy against a temperature-softened teacher distribution, which by the identity in §3.1 is a forward KL against that distribution. ↩

  2. Rafael Müller, Simon Kornblith, and Geoffrey Hinton, “When Does Label Smoothing Help?” arXiv:1906.02629 (2019), NeurIPS 2019. https://arxiv.org/abs/1906.02629 ↩

  3. The mode-covering and mode-seeking vocabulary predates its use in distillation and is standard in variational inference. Its distillation-specific consequences are measured in Chapter 6 and in Lab 05; the summary claim carried forward from Lab 01 is that the choice is a product decision rather than a mathematical one. ↩

  4. Ganqu Cui et al., “The Entropy Mechanism of Reinforcement Learning for Reasoning Language Models,” arXiv:2505.22617 (2025). https://arxiv.org/abs/2505.22617. The standard reference for entropy collapse in reinforcement learning with verifiable rewards; the mechanism it describes is the one Chapter 12 watches for in on-policy distillation. ↩

  5. Imre Csiszár, “Information-type measures of difference of probability distributions and indirect observations,” Studia Scientiarum Mathematicarum Hungarica 2 (1967): 299-318. The family was co-discovered independently by S. M. Ali and S. D. Silvey, “A general class of coefficients of divergence of one distribution from another,” Journal of the Royal Statistical Society Series B 28, no. 1 (1966): 131-142, https://doi.org/10.1111/j.2517-6161.1966.tb00626.x. Csiszár’s earlier German-language precursor appeared in Publ. Math. Inst. Hungar. Acad. Sci., Ser. A, 8 (1963): 85-108. Neither Csiszár paper has a DOI; cite by volume and page. ↩

  6. John Schulman, “Approximating KL Divergence,” joschu.net, 7 March 2020, http://joschu.net/blog/kl-approx.html (accessed 2026). A personal blog post rather than a refereed paper, and the standard source for the k1/k2/k3 estimators. Chapter 4 derives all three. ↩

  7. Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou, “f-Divergence Minimization for Sequence-Level Knowledge Distillation,” arXiv:2307.15190 (2023), ACL 2023. https://arxiv.org/abs/2307.15190 ↩

  8. Mingyang Song and Mao Zheng, “A Survey of On-Policy Distillation for Large Language Models,” arXiv:2604.00626 (2026). https://arxiv.org/abs/2604.00626. An ongoing preprint rather than a published survey, and it should be read as a snapshot of a moving field; its organizing axis is the f-divergence family. ↩

  9. Nicolas Boizard, Kevin El Haddad, Céline Hudelot, and Pierre Colombo, “Towards Cross-Tokenizer Distillation: the Universal Logit Distillation Loss for LLMs,” arXiv:2402.12030 (2024), Transactions on Machine Learning Research, January 2025. https://arxiv.org/abs/2402.12030 ↩

  10. Dominik M. Endres and Johannes E. Schindelin, “A new metric for probability distributions,” IEEE Transactions on Information Theory 49, no. 7 (2003): 1858-1860. https://doi.org/10.1109/TIT.2003.813506. Proves specifically that the square root of the Jensen-Shannon divergence satisfies the triangle inequality. ↩

  11. Ferdinand Österreicher and Igor Vajda, “A new class of metric divergences on probability spaces and its applicability in statistics,” Annals of the Institute of Statistical Mathematics 55, no. 3 (2003): 639-653. https://doi.org/10.1007/BF02517812. Establishes the broader family of metric divergences , of which is the case. The two 2003 results are conventionally cited together. ↩

  12. Frank Nielsen et al., “Metrization of powers of the Jensen-Shannon divergence,” arXiv:2302.10070. The modern generalization of the Endres-Schindelin result. ↩

  13. M. S. Pinsker, Information and Information Stability of Random Variables and Processes (San Francisco: Holden-Day, 1964), translated and edited by A. Feinstein from the 1960 Russian original. Pinsker’s own constant is weaker than the form used here; the optimal constant is due independently to Csiszár (1967, cited above) and to S. Kullback, “A lower bound for discrimination information in terms of variation (Corresp.),” IEEE Transactions on Information Theory 13, no. 1 (1967): 126-127, https://doi.org/10.1109/TIT.1967.1053968, with a 1970 correction, https://doi.org/10.1109/TIT.1970.1054514, and to J. H. B. Kemperman (1969). The result is therefore properly the Csiszár-Kullback-Pinsker inequality. ↩

  14. A. A. Fedotov, P. Harremoës, and F. Topsøe, “Refinements of Pinsker’s inequality,” IEEE Transactions on Information Theory 49, no. 6 (2003): 1491-1498. https://doi.org/10.1109/TIT.2003.811927 ↩

  15. Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang, “MiniLLM: Knowledge Distillation of Large Language Models,” arXiv:2306.08543 (2023), ICLR 2024. https://arxiv.org/abs/2306.08543v2. Note that the arXiv landing page now shows a later title, “MiniLLM: On-Policy Distillation of Large Language Models”; the ICLR 2024 version of record uses the title given here. ↩

  16. Rishabh Agarwal et al., “On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes,” arXiv:2306.13649 (2023), ICLR 2024. https://arxiv.org/abs/2306.13649. The method is generalized knowledge distillation; the name does not appear in the title. ↩

  17. Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se-Young Yun, “DistiLLM: Towards Streamlined Distillation for Large Language Models,” arXiv:2402.03898 (2024), ICML 2024. https://arxiv.org/abs/2402.03898. The follow-up is Jongwoo Ko et al., “DistiLLM-2: A Contrastive Approach Boosts the Distillation of LLMs,” arXiv:2503.07067 (2025), ICML 2025. https://arxiv.org/abs/2503.07067 ↩

  18. Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn, “Direct Preference Optimization: Your Language Model is Secretly a Reward Model,” arXiv:2305.18290 (2023), NeurIPS 2023. https://arxiv.org/abs/2305.18290 ↩

  19. Jianping Gou, Baosheng Yu, Stephen J. Maybank, and Dacheng Tao, “Knowledge Distillation: A Survey,” International Journal of Computer Vision 129, no. 6 (2021): 1789-1819, https://doi.org/10.1007/s11263-021-01453-z. For the language-model-specific version see Xiaohan Xu et al., “A Survey on Knowledge Distillation of Large Language Models,” arXiv:2402.13116 (2024). https://arxiv.org/abs/2402.13116 ↩

  20. Samuel Stanton, Pavel Izmailov, Polina Kirichenko, Alexander A. Alemi, and Andrew Gordon Wilson, “Does Knowledge Distillation Really Work?” arXiv:2106.05945 (2021), NeurIPS 2021. https://arxiv.org/abs/2106.05945 ↩

Part I · Foundations

4

Estimating Divergences from Samples

Chapter 3 gave you a formula for the distance between two next-token distributions and treated both distributions as things you can read off a page. Every divergence in that chapter is a sum over the vocabulary, and every sum assumed you had all numbers from the teacher and all numbers from the student, at every position, whenever you wanted them.

You usually do not.

Take the smallest concrete case this book uses. SmolLM2’s tokenizer has a vocabulary of 49,152 entries.1 A modest training batch of 8 sequences at 1,024 tokens each is 8,192 positions, so one full-vocabulary logit tensor for that batch holds 8,192 × 49,152 ≈ 402 million numbers. In bf16 that is 805 megabytes. You need one for the teacher and one for the student, and the loss computation reads both, writes at least one intermediate the same size for the log-softmax, and then reduces. On the reference machine’s 273 GB/s of memory bandwidth, moving a single 805 MB tensor once costs about 2.9 milliseconds, and a naive implementation moves several of them per step. None of that is fatal. It is a real cost that scales linearly in , and it is the reason Chapter 10 exists.

The harder problem is not cost. It is access. In on-policy distillation the student generates a token, and what you get back from the teacher is often the log-probability of that token and nothing else. In an API setting you may get the top 20 log-probabilities, or the log-probability of the sampled token alone, or nothing but text.2 The full-vocabulary sum you would need to evaluate exactly is not available at any price, because the numbers were never sent. That situation is the direct consequence of training on the student’s own generations, which the sequence-modeling literature has wanted to do since long before distillation adopted the idea.16

So the question this chapter answers is: given samples and log-ratios at those samples, how do you estimate a divergence, how wrong is your estimate, and how would you know? The answers are more interesting than they sound. One of them is a measured result from the course’s own labs that contradicts the advice you will find in most places this topic gets discussed, including the advice I wrote into the lab before the data changed my mind.

4.1 Three situations, and only one of them is comfortable#

Before any estimator, sort the problem. There are exactly three access patterns for divergence computation in distillation, and which one you are in determines everything downstream.

Both distributions in full. You have the teacher’s complete logit vector and the student’s complete logit vector at every position. The divergence is an exact sum. This is the situation when you own both models, run them in the same process, and are willing to pay the memory traffic described above. Everything in Chapter 3 applies without modification, and the number you compute is the number, not an estimate of it.

Definition

Dense logits

The teacher’s complete score vector over the whole vocabulary, at every position. Having dense logits is what makes an exact divergence computation possible. Most of the machinery in this chapter exists because you often do not have them.

One in full, one truncated. You have the student’s full distribution because you are training it, and a top- slice of the teacher’s because that is what got cached, or that is what the API returned. The missing tail has some mass, and every divergence you compute is off by whatever the tail would have contributed. This is not a sampling problem and not a Monte Carlo problem; it is a deterministic truncation with a bias you can sign and sometimes bound. Chapter 10 handles it in full, including the two conventions for what to do with the missing mass and the direction each one errs in.

Only sampled log-ratios. You have a token that was sampled from the student, the student’s log-probability of it, and the teacher’s log-probability of it. One number per position, per model. This is the on-policy case and the API case, and it is where this chapter lives.

Table 4.1 What you can compute from each access pattern.

Access pattern What you hold per position Divergence status Where it comes up
Dense both sides teacher logits, student logits Exact, cost in traffic Both models in one process
Dense student, top- teacher student, teacher plus tail mass Truncation bias, deterministic Cached-logit pipeline (Ch. 10), logprob APIs
Sampled log-ratios , at the sampled Monte Carlo estimate, random On-policy rollouts, black-box scoring

Read the third row carefully, because the change it introduces is qualitative rather than quantitative. In the first two rows, running the same computation twice on the same inputs gives the same answer. In the third row it does not. Your “KL” becomes a random variable with a mean, a variance, and a sampling distribution, and every property you learned about KL in Chapter 3 (nonnegativity, in particular) is a property of the mean rather than of any number you will actually see.

4.2 Monte Carlo, from the ground#

Suppose you want the average value of some function under a distribution , written . If is a categorical distribution over 49,152 tokens, that expectation is a sum of 49,152 terms. If you cannot afford the sum, or cannot see all of its terms, you can draw independent samples from and average over them:

Nothing about this is specific to divergences. Sequence-level distillation is the same construction with : the intractable sum over all output sequences gets replaced by the teacher’s single most likely one.13

Definition

Monte Carlo estimator

An estimate of an expectation formed by averaging a function over random draws from the distribution the expectation is taken under. It approaches the true value only as draws accumulate, and at any finite it is a random number with a distribution of its own.

That average is useful because of two facts about it. The first is that exactly, for every , including . Each term in the average has the right expectation by construction, and expectation is linear, so the average does too. That property has a name.

Definition

Unbiased estimator

An estimator whose expected value equals the quantity it estimates, at every sample size. Unbiased does not mean accurate: a single draw from an unbiased estimator can be arbitrarily far from the truth, and in the cases this chapter cares about, routinely is.

The second fact is that the variance of the average is the variance of one term divided by , so the standard deviation of is where is the standard deviation of under . That quantity is the standard error, and it is the honest one-number summary of how much you should trust the estimate. The is unforgiving: ten times the accuracy costs a hundred times the samples.

Now specialize. Reverse KL, in the convention this book uses where is the teacher and is the student, is

That is already an expectation under . And is the student, which is the model that generated the token. So if you sample tokens from the student and average the log-ratio over them, you get an unbiased estimate of the reverse KL with no further machinery. This is the whole reason reverse-flavored objectives dominate on-policy pipelines, and it is worth saying plainly now because it will be re-derived from a different direction in Chapter 12: the reverse direction is the one whose defining expectation is taken under the distribution you are already sampling from.

Forward KL is an expectation under , the teacher, and you are not sampling from the teacher. To estimate it from student samples you need to reweight, and §4.5 shows what that costs.

It is convenient to work with the ratio in the other orientation, matching Schulman’s note and the lab.

Definition

Importance ratio

For a sample drawn from , the quantity : how much more probable the teacher considers this token than the student does. Its expectation under is exactly 1, because . That identity is the hinge the rest of this chapter turns on.

With , the log-ratio is , so

Note also that holds no matter how different and are. It says nothing about how concentrated is. Chapter 3’s exercise on the chi-squared divergence establishes the companion fact, which will matter in §4.4: . The variance of the importance ratio is a divergence in its own right, and it is one of the badly behaved ones.

4.3 Three estimators of the same number#

John Schulman’s blog post on approximating KL divergence is the standard reference for what follows, and the names k1, k2, k3 come from it.3 Each is a function of the importance ratio at a single sampled token, each is meant as an estimator of from draws , and each makes a different trade along the bias-variance axis that runs through most design decisions in distillation.15

Everything worth knowing about them follows from writing and expanding. I will do the three in order.

4.3.1 k1: unbiased, noisy, and sometimes negative#

is the definition read literally. Its expectation under is , which is the identity from §4.2, so is unbiased by construction and there is nothing further to derive.

The problem is what a single sample looks like. is negative exactly when , meaning exactly when the teacher assigns the sampled token more probability than the student did. That happens constantly. Even for two distributions that are quite close, roughly half the tokens will have on one side of 1 and half on the other. In the course’s own measurement on Lab 00 §9’s close regime, 43 percent of the 200,000 individual samples came out negative, with the smallest at , against a true KL of nats.

That is not a bug in . KL is nonnegative as an average, by Gibbs’ inequality, and Gibbs’ inequality is a statement about the full sum. A single term of that sum carries no such guarantee. The negativity is the sampling noise being larger than the signal, which is the ordinary condition when two distributions are close: the true value is small precisely because the positive and negative contributions nearly cancel, and a small number obtained by cancellation of large ones is exactly the situation where sampling hurts most.

Quantitatively, . That is the variance of a log quantity, which grows slowly. If the student’s logit on some token sits 25 nats below the teacher’s, on that token is 25, large but finite, and it enters linearly. Hold on to that; it is the reason survives the far regime in §4.4 when does not.

4.3.2 k2: low variance, and biased#

is nonnegative on every sample by construction, which is appealing, and it has visibly less spread than because squaring a small number makes it smaller. It is also biased, and the bias has a clean closed form worth deriving because it tells you exactly when to distrust it.

Start from and write :

Expand the exponential inside the expectation:

so

Since and , rearrange:

So ’s bias is exactly the negated sum of the third and higher moments of the log-ratio, with factorial denominators. When and are close, is small, the true KL is second order in , and the bias is third order, so the relative bias is first order in the typical size of and shrinks as the distributions converge. When and are far apart, the higher moments stop being small and the whole expansion stops being informative.

The measurements bear this out with unusual clarity. On Lab 00 §9’s close regime, ’s mean is against a true KL of , an overestimate of about 5 percent. On the far regime, its mean is against a true KL of : not 5 percent high, but 164 percent high. An estimator whose bias is “small when the thing you are measuring is small” is a different animal from one whose bias is small.

4.3.3 k3: unbiased and nonnegative at the same time#

is with the term added. Two properties follow immediately.

It is unbiased. $\mathbb{E}_q[k_3] = \mathbb{E}_q[r] - 1 - \mathbb{E}_q[\log r] = 1 - 1 + \mathrm{KL}(q|p) = \mathrm{KL}(q|p)$. The added term contributes exactly zero to the expectation, because . The correction is free.

It is nonnegative on every single sample. The inequality holds for all , with equality only at ; it is the statement that the logarithm lies below its tangent line at 1, which follows from concavity. Therefore pointwise, and for every draw.

Getting both of those at once is the structural property that makes the usual default. An unbiased estimator that can go negative will occasionally print a negative KL, and anyone reading that dashboard will either conclude the code is broken or, worse, stop trusting the panel. An estimator that never goes negative but is biased will report a number that is systematically wrong in a direction that changes as training proceeds. refuses both failure modes.

The reason it can is worth naming, because the idea generalizes far past this estimator.

Definition

Control variate

A quantity added to an estimator that has known expectation (usually zero) and is correlated with the estimator’s noise. Adding it changes nothing in expectation and can cancel a large part of the variance. In , the control variate is , whose expectation under is zero because , and which is correlated with because both are functions of the same ratio.

There is also a nice reading of as an f-divergence generator rather than as a trick. Chapter 3 introduced for convex with , and gave $f(t) = -\log ta(t-1)$ to a generator changes no divergence value, because .4 The choice produces , which is the unique member of that affine family that touches zero at and is therefore nonnegative everywhere. is not a clever hack bolted onto reverse KL; it is reverse KL’s generator written in the normalization where every pointwise term is a nonnegative contribution.

Table 4.2 The three estimators, before the regime question.

Formula Bias Sign of a single sample Structural appeal
Zero Either It is the definition, read literally
Always Small spread when
Zero Always Both properties at once

The listing below is the shape of all three in code, on synthetic distributions. Watch the third column of the output, the standard deviation, since that is the quantity Table 4.2 deliberately leaves blank.

import torch, torch.nn.functional as F

def estimator_stats(p, q, n, seed):
    """Sample n tokens from q and return the three KL(q||p) estimators."""
    g = torch.Generator().manual_seed(seed)
    x = torch.multinomial(q, n, replacement=True, generator=g)
    log_r = (p.log() - q.log())[x]        # one number per sampled token
    r = log_r.exp()
    return {
        "k1": -log_r,
        "k2": 0.5 * log_r ** 2,
        "k3": (r - 1) - log_r,
    }

def report(p, q, n=200_000, seed=7):
    truth = float((q * (q.log() - p.log())).sum())   # exact, full vocabulary
    for name, k in estimator_stats(p, q, n, seed).items():
        se = float(k.std()) / n ** 0.5
        print(f"{name}  mean {k.mean():+.5f}  std {k.std():8.5f}"
              f"  bias {k.mean() - truth:+.5f}  se {se:.5f}")

The listing proves nothing on its own; it is the instrument. What it measures is the subject of the next section, and the measurement is the reason this chapter exists.

4.4 The regime dependence, which is the result that matters#

Here is the claim you will find nearly everywhere this topic is discussed, including in the first draft of the course’s own lab: is unbiased like and low-variance like , so use and stop thinking about it.

That claim is false in a way that matters enormously for distillation, and the course measured it.

Lab 00 §9 builds two regimes over a 500-token vocabulary and draws 200,000 samples in each. Regime A is late training: a teacher and a student for standard normal , so the two are the same distribution with a small perturbation. Regime B is cold start: and drawn independently from the same random-distribution generator, with no relationship between them at all. Both regimes use the same seed and the same sampling call.

Table 4.3 Lab 00 §9’s measured estimator statistics, , , seed 7.

Regime True
A (close) 0.05328 mean 0.05233 0.05617 0.05323
std 0.33106 0.07118 0.06122
B (far) 4.06400 mean 4.06940 10.72473 3.96930
std 2.21122 7.62477 23.85423

Read the two bolded cells. In the close regime, ’s standard deviation is 0.061 against ’s 0.331, a factor of 5.4 in ’s favor, which is the result everyone quotes. In the far regime, ’s standard deviation is 23.85 against ’s 2.21, a factor of 10.8 in the opposite direction. The estimator that is five times better in one regime is eleven times worse in the other, on the same code with the same seed.

Both remain unbiased. Both means sit on the truth in both regimes, within the lab’s assertion tolerance of five standard errors. The bias story does not change. Only the variance does, and it inverts.

Field note

I wrote the “k3 always wins” cell first, with a single regime and an assertion that std(k3) < std(k1). It passed, because the regime I happened to construct was the close one, since that is the regime you naturally reach for when you want a demonstration to look clean.

Then I added a second regime to show the effect was general, and the assertion failed. My first instinct was that I had a bug: sign error somewhere, or the sampler reusing state across calls. There was no bug. I had constructed two independent random distributions over 500 tokens, and on those, reaches into the thousands on tokens that the student happens to hit and the teacher happens to like, and carries every one of those spikes into at full size.

The lab now demonstrates both regimes and asserts the inversion explicitly: std(k3) < std(k1) in the close regime, std(k3) > std(k1) in the far one. That inverted assertion is the most useful line in Lab 00, because it is the one that would fail loudly if someone “simplified” the cell back to the version I originally wanted to write.

4.4.1 Why the correction stops correcting#

The mechanism is visible in three lines of algebra. Write again, and note

Expand around :

The linear term is gone. That is the entire content of the control variate. fluctuates at first order in the log-ratio; adding , whose own expansion is , cancels the first-order term exactly, leaving to fluctuate at second order. When is typically 0.2, a second-order fluctuation is roughly five times smaller than a first-order one, and you get the close-regime result.

The same expansion gives two more results. First, near , so and are the same estimator to second order, which explains why their close-regime standard deviations in Table 4.3 (0.0712 and 0.0612) sit so near each other. Second, the expansion is a statement about small , and it has no force otherwise.

For large the two estimators diverge in behavior completely. grows linearly in the log-ratio. contains , which grows exponentially in the log-ratio. A token where the teacher assigns probability 0.1 and the student assigns 0.0001 has : contributes to the average, while contributes . One such draw in a batch of a few hundred moves the average by several units and the average by a hundredth of one.

The formal version is the connection to chi-squared. , from Chapter 3, and ’s variance is dominated by the variance of its term once has any appreciable spread. Chapter 3’s exercise measured the chi-squared divergence growing by more than ten orders of magnitude over a logit gap where forward KL grew by a factor of five, because chi-squared charges where KL charges . inherits that growth rate. does not.

There is one more way to see it that generalizes, and I will state it here because §4.8 needs it. The family of unbiased estimators

is unbiased for every real , since the added term always has expectation zero. is and is . Standard control-variate theory says the variance-minimizing choice is

and computing that quantity on Lab 00 §9’s own two regimes gives in the close regime and in the far one. In the close regime the optimal correction is almost exactly the one hardcodes. In the far regime the optimal correction is to apply almost no correction at all, which is to say, to use . is not a general-purpose improvement over ; it is a fixed choice of control-variate coefficient that happens to be optimal in one regime and roughly 75 times too large in the other.

2026-08-01T07:26:37.246756 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 0.001 0.01 0.1 1 true KL(q||p), nats 1 0 − 3 1 0 − 2 1 0 − 1 1 0 0 1 0 1 1 0 2 per-sample standard deviation k1 k2 k3 k3 wins k1 wins crossover at true KL = 0.90 nats regime A, true KL 0.0533: std(k1) 0.331 against std(k3) 0.061 regime B, true KL 4.064: std(k1) 2.211 against std(k3) 23.85
Figure 4.1 Across a sweep from near-identical to unrelated distributions, k3's standard deviation crosses above k1's at a true KL near one nat, so the estimator with the lower variance is a function of how far apart the models already are.

4.4.2 What this means for a training run#

Now connect it forward, because the regimes are not abstract.

“The student is close to the teacher” is the description of a run in steady state, some thousands of steps in, when the loss has come down and the two distributions overlap heavily. “The student is far from the teacher” is the description of the first few hundred steps of a run that started from a base model, or from random initialization, or from a pruned checkpoint whose surgery has not yet been recovered. Chapter 12 calls that the cold start and spends real length on it, because it is the regime where on-policy distillation is at its most fragile.

Put those together and the conclusion is uncomfortable. The estimator you would trust during steady-state training is the one that is least trustworthy at the beginning of a run, and the beginning of a run is when you most need a reliable monitor, because that is when things actually break. Sampled-ratio estimators are dependable exactly when the two models already roughly agree, and a freshly initialized student is the case where they do not.

That is not a reason to abandon them. It is a reason to know which regime you are in before you read the number, and to prefer (or a full-vocabulary computation on a held-out slice) for the early phase of a run.

2026-08-01T07:26:40.765721 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 0 50 100 150 200 k1 = -log r mean 0.0539 sd 0.0412 9.4% of replicates estimate a negative KL 0 200 400 600 800 replicates per bin k2 = 0.5 (log r)^2 mean 0.0563 sd 0.0089 −0.10 −0.05 0.00 0.05 0.10 0.15 0.20 estimate of KL(q||p) from 64 sampled tokens, nats 0 250 500 750 1000 k3 = (r - 1) - log r mean 0.0533 sd 0.0076 true KL = 0.05328
Figure 4.2 The sampling distribution of each estimator at 64 samples in the close regime: k1 spreads widely and puts substantial mass below zero, k2 and k3 are tight and strictly nonnegative, and all three are centered on the true value.

4.5 The forward direction, where importance weights are unavoidable#

Everything so far estimates , the reverse direction, and got away with it because the samples came from . Forward KL is the direction distillation most often wants, since it is the mass-covering one and the one that punishes a student for abandoning something the teacher cares about. Estimating it from student samples is a different problem, and Solutions 00 Exercise 4 works it through.

The derivation starts by writing forward KL as an expectation under and then changing the measure:

The step in the middle is importance weighting: multiply and divide by , so the sum becomes an expectation under with the extra factor inside. That factor is the price of sampling from the wrong distribution, and it is charged multiplicatively rather than inside a logarithm.

So the naive forward estimator is the average of over draws from . It is unbiased, and it is bad. It can also be negative on a single sample: for all , with a minimum of at . Hold on to that number; §4.6 needs it.

Apply the same control variate as before. still, so is still free to subtract:

This is unbiased by the same one-line argument, and it is nonnegative on every sample because for all . That inequality is the same fact as before in a different costume: it says lies above its tangent at , which is convexity, and Chapter 3’s generator table gave as the forward-KL generator. Subtracting the affine term changes no divergence value and shifts the generator to touch zero at .

The lab runs the identical construction to §9, same , same , same seed, and adds reverse on the same draws as a yardstick.

Table 4.4 Solutions 00 Ex4’s measured forward-direction estimators, on the same draws as Table 4.3.

Regime True Estimator mean std
A (close) 0.05147 naive 0.05235 0.33631
0.05145 0.05893
reverse (yardstick) 0.05233 0.33106
B (far) 4.74743 naive 4.06521 174.06
4.16531 150.08
reverse (yardstick) 4.06940 2.21122

Regime A behaves. The control variate cuts the standard deviation by a factor of 5.7, both means sit on the truth, and every sample is nonnegative. This mirrors the reverse-direction result exactly, for the same reason: near , is quadratic in and therefore tiny.

Regime B is a wreck, and the details are instructive. The naive estimator’s standard deviation is 174, which is 37 times the true value of 4.75 it is trying to estimate. The control variate reduces that to 150, a 14 percent improvement, against the 5.7-fold improvement it delivered in regime A. Subtracting cannot rescue anything here, because the noise is the heavy tail of , and appears in both terms. And on the identical draws, reverse has a standard deviation of 2.21: the forward estimator is 79 times noisier than the reverse one on exactly the same sampled tokens.

Look at the far-regime means as well. Both forward estimators come in low: 4.07 and 4.17 against a truth of 4.75. They are unbiased, and the run still undershoots, because the expectation is carried by astronomically rare draws with enormous that 200,000 samples did not happen to include. This is the sentence from the solutions worth memorizing: an unbiased estimator you cannot afford enough samples for behaves, in practice, like a biased one. Unbiasedness is a statement about an average over infinitely many hypothetical repetitions of your experiment, and you are running the experiment once.

The compressed reason forward is harder than reverse: the reverse estimators need only inside a logarithm, with the linear as an optional correction, while the forward estimator needs multiplicatively as the importance weight itself. So the forward estimator’s variance is governed by , driven by the same second moment of (the chi-squared divergence) that explodes whenever the models disagree.

That has a consequence for how pipelines are built which people usually attribute to philosophy. On-policy distillation pipelines default to reverse-flavored objectives, or to interpolations weighted toward reverse, not because reverse KL is conceptually superior but because it is the direction whose sampled estimator does not require importance weights.5 Chapter 6 argues the mode-seeking versus mode-covering case on its own merits, and Chapter 12 covers the design space that GKD’s two parameters open up.6 This chapter is only pointing out that the estimator mathematics has a vote, and it votes before the philosophy gets a turn.

def forward_kl_estimators(p, q, n, seed):
    """Estimate KL(p||q) from samples of q. Note that r enters multiplicatively."""
    g = torch.Generator().manual_seed(seed)
    x = torch.multinomial(q, n, replacement=True, generator=g)
    log_r = (p.log() - q.log())[x]
    r = log_r.exp()
    return {
        "naive": r * log_r,                # unbiased, can be negative, heavy tailed
        "k3fwd": r * log_r - (r - 1),       # unbiased, nonnegative, same heavy tail
        "reverse_k1": -log_r,               # different quantity, shown for scale
    }

The listing is deliberately close to the reverse-direction one so the single structural difference is visible: r appears outside the logarithm, and that is the whole cause of Table 4.4’s right-hand column.

4.6 The sampled-token trap, in the form you will meet it#

Lab 07 puts the failure in its practical clothes, and the numbers are worth walking through in detail because the mechanism is not obvious from the general variance argument.

Definition

Sampled-token estimator

A divergence estimate computed from the teacher’s and student’s log-probabilities at only those tokens the student actually sampled, one token per position, with no access to the rest of the vocabulary. It is the only estimate available when the teacher scores rollouts rather than returning full distributions.

The setup is ten tokens. The teacher has one strong mode: logits are zero everywhere except token 0, which sits at 3.0, giving and for the other nine. The student has abandoned that mode completely: it is uniform except that before renormalization, which leaves on each of the nine others. This is a caricature of a real failure, a student that has stopped putting any mass at all on something the teacher considers most likely.

The exact forward KL is nats, and it is almost entirely one term: . That is what forward KL is for. It is the direction that charges you enormously for putting near-zero probability where the teacher puts mass.

Now estimate it from 5,000 tokens sampled from the student. The student’s probability of drawing token 0 is , so across 5,000 draws the expected number of times token 0 appears is . In the lab’s run it appeared zero times, and the lab asserts this explicitly. Every one of the 5,000 sampled tokens is one of the nine others, where the ratio is

identical for all nine. Since every sample has the same ratio, the importance-weighted estimate is not an average of varying terms at all. It is exactly

The reported estimate is against a true value of . It lands on the wrong side of zero, which makes it a different kind of answer and not a degraded version of 13.6.

Two observations sharpen this. First, is very close to , the global minimum of . The estimate did not land low by a little; it landed near the most negative value that a single term can take. Second, the reverse estimate on the very same 5,000 draws lands within 10 percent of its closed form, which is the property Lab 07 asserts. The samples are fine. The direction is the problem.

4.6.1 How a nonnegative quantity gets a negative estimate#

There are two distinct mechanisms in play in this chapter, and it is worth separating them because the fix differs.

The first is ’s ordinary negativity from §4.3.1: a single log-ratio has an arbitrary sign, and averaging a few of them can land below zero even though the average of all of them cannot. This is symmetric noise around a small true value. More samples fix it, and fixes it structurally.

The second is what Lab 07 shows, and it is not symmetric noise. The estimator is unbiased over the full sampling distribution, and the full sampling distribution includes draws of token 0. Those draws carry an importance weight of and would contribute to the average. The expectation works out to 13.6 precisely because that gigantic contribution is multiplied by the chance of seeing it. But you will not see it. To see it once in expectation you would need on the order of samples, and you drew 5,000. So what you actually observe is the conditional distribution given that the mode was missed, whose mean is the small negative number computed above.

That is the shape of the whole trap: an unbiased estimator whose expectation is carried entirely by events your sample size cannot reach. It will report a plausible-looking small number, indefinitely, with no variance to warn you, because the variance you would need to see also lives in the draws you are not getting. The 2026 survey of on-policy distillation catalogues this class of failure among the practical problems the on-policy literature keeps rediscovering, and the framing is worth adopting: these are properties of estimating an expectation under one distribution using samples from another that has abandoned part of its support, and no amount of debugging removes them.7

Watch out

A negative number on a panel labeled “KL” is not always a bug, and treating it as one will send you looking in the wrong place. What the number means depends entirely on which estimator produced it, so work through the three cases in order. If the panel computes an exact full-vocabulary sum, a negative value is a bug and Gibbs’ inequality says so. If it averages a sampled , negative values are expected whenever the true KL is small relative to the sampling noise, and the fix is more samples or , not a code review. If it is an importance-weighted forward estimate, a small negative value signals that the student is not sampling the region where the divergence lives, which is the most serious of the three possibilities and the one that looks the most benign.

4.6.2 What to do about it on a real dashboard#

What helps, in increasing order of what it costs you:

Label the panel with the estimator, not with the quantity. A panel called “KL” that is actually a sampled average over the current batch’s rollout tokens should say so, in the title. This sounds like a documentation nicety and is not: the label is what determines whether a reader interprets a downward move as the student improving or as the sampler drifting.

Never monitor forward KL from student samples alone. If you want the mass-covering direction on a dashboard, compute it densely on a small fixed held-out batch, at whatever interval you can afford, rather than estimating it from rollouts. A hundred positions of exact forward KL every 50 steps is worth more than every position of a sampled estimate, because the sampled estimate is systematically blind to exactly the failure forward KL is meant to detect.

Watch a quantity that does not have this pathology, alongside. Rollout entropy is computed from the student’s own distribution, requires no teacher and no ratio, and collapses visibly when the student abandons modes.8 Chapter 12 builds a monitor on it and calibrates the monitor against healthy runs so it does not fire on ordinary entropy decline. The pairing matters: a sampled KL that looks stable while entropy is falling off a cliff is the exact signature of the Lab 07 situation.

It is worth adding that a well-built on-policy trainer avoids the estimator problem entirely for the loss itself. GKD scores the student’s rollouts with the teacher and computes the per-position divergence over the full vocabulary, so the objective is an exact sum at every rollout position.6 What stays student-driven is which positions exist at all, meaning the state distribution rather than the per-position divergence. The sampled-token trap bites when you are monitoring, when the teacher is remote and returns only the sampled token’s log-probability, or when you have written the estimator yourself from a survey formula.

4.7 Practical guidance#

4.7.1 How many samples#

The standard error of a Monte Carlo estimate is , so pick a tolerance and invert. Using Lab 00 §9’s measured per-sample standard deviations from Table 4.3, and asking for a standard error of 0.05 nats:

Table 4.5 Samples required for a standard error of 0.05 nats, from Table 4.3’s measured spreads.

Regime Estimator Per-sample std Required
A (close) 0.331 44
A (close) 0.061 2
B (far) 2.211 1,956
B (far) 23.854 227,610

Read the table twice. In the close regime, a batch of a few hundred rollout tokens is plenty for either estimator, which is why sampled monitoring works at all in steady state. In the far regime, needs about two thousand tokens and needs about a quarter of a million, which no ordinary batch supplies. Note also that a standard error of 0.05 nats is a much weaker demand in regime B, where the true value is 4.06, than in regime A, where it is 0.053; asking for a fixed relative precision makes the far-regime numbers worse still.

A general rule that survives both regimes: compute the standard error from the samples you have, every time, and refuse to interpret any movement smaller than a couple of standard errors. The computation costs one line.

4.7.2 Seeds and paired comparison#

Every sampled divergence in the labs uses an explicit torch.Generator with a stated seed rather than the global random state, and the reason is not superstition about reproducibility. It is that you often want to compare two things measured on the same draws.

Table 4.4’s yardstick row is the example. The claim “the forward estimator is 79 times noisier than the reverse one” is only meaningful because both were computed from the identical 200,000 sampled tokens. Had they been computed from independent draws, part of the difference would be sampling noise in the comparison itself, and the factor of 79 would need its own error bar. Reusing the same draws across estimators is the technique statisticians call common random numbers, and it is nearly free: draw once, compute every estimator on the same x.

The same discipline applies across arms of a training experiment. If arm A and arm B are being compared on a sampled divergence, and both are evaluated on a fixed held-out prompt set with a fixed evaluation seed, then a difference between them is more likely to be real than if each arm’s evaluation drew its own tokens. Chapter 18 makes this part of the pre-registration.

4.7.3 Reporting#

Report a sampled divergence as a value and a standard error, in that order, with the sample count: “reverse KL 0.053 ± 0.001 nats (k3, N = 200,000)”. Every part of that string does work. The estimator name tells a reader whether to expect nonnegativity. The tells them whether the standard error is believable. The standard error tells them what movements to ignore.

And one discipline that is more important than all of the above, stated as a rule.

Never compare an estimated divergence against an exactly computed one across arms of an experiment. If arm A reports dense forward KL on a held-out set and arm B reports a sampled estimate on rollouts, the difference between the two numbers contains the difference between the arms, plus the estimator’s bias, plus the difference between the two evaluation distributions, and you cannot separate them after the fact. This sounds like an obvious hygiene rule, and it is violated constantly, because the two arms are often built at different times by different people and the panel is named “KL” in both cases. Table 4.3 makes the magnitude concrete: overstates the far-regime KL by 164 percent. If arm A used and arm B used a dense sum, arm A would appear to be more than twice as far from the teacher, entirely as an artifact.

The same rule has a weaker cousin worth stating: do not compare sampled divergences computed at different sample counts without noting it, since the bias of does not shrink with but the apparent stability of any estimator does, and a smoother line reads as a better run.

2026-08-01T07:26:44.358626 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 1 0 1 1 0 2 1 0 3 1 0 4 1 0 5 samples N 0.02 0.04 0.06 0.08 0.10 running estimate of KL(q||p), nats true KL 0.05328 regime A: close, true KL 0.05328 k1 k2 k3 1 0 1 1 0 2 1 0 3 1 0 4 1 0 5 samples N 2 4 6 8 10 12 k3 still 3.969 at N = 200,000, approaching from below, band +/-0.053 true KL 4.06400 k1 settles by a few thousand samples, band +/-0.0049 k2 converges to 10.73, which is not the KL regime B: far, true KL 4.06400
Figure 4.3 Running estimates against sample count in both regimes: in the close regime all three converge quickly and k3's band is tightest, while in the far regime k3's band stays wide and its running mean approaches the truth from below long after k1 has settled.

4.8 Variance reduction as a general tool#

is one instance of a technique that has a name and a theory, and it is worth having the general form because distillation training runs are full of high-variance estimates that are not divergences.

The general control-variate construction: you want , you have some whose expectation you know exactly, and you form

This is unbiased for every , since the subtracted term has expectation zero. Its variance is

a quadratic in minimized at , where the variance becomes with the correlation between and . The reduction is governed entirely by how correlated your control variate is with the thing you are estimating. A perfectly correlated control variate eliminates the variance; an uncorrelated one does nothing; and a badly chosen can make things much worse, which is the case lands in during cold start.

Reading through that formula: , , , and fixed at in the sign convention above (equivalently, is added rather than subtracted). §4.4.1 gave the measured on the lab’s two regimes, 1.011 and 0.0133, and the far-regime number says the correlation between and has collapsed relative to ’s own variance. That is what a heavy tail does: ’s variance is dominated by a handful of enormous values, and barely moves on those same draws, so the covariance grows far slower than the variance.

You could estimate from the batch and use it, which would give an estimator at least as good as in both regimes. Almost nobody does, for a reason worth knowing: estimating from the same samples you apply it to introduces a small bias, and the resulting estimator is no longer exactly unbiased. In practice the honest simple options are “use and know your regime” or “use and pay in variance.”

Control variates show up in the distillation literature in a place you might not expect. MiniLLM formulates reverse-KL distillation as a policy-gradient problem, since minimizing reverse KL over sequences is an optimization over the student’s own sampling distribution, and policy-gradient estimators are notoriously high variance. The paper’s contribution is as much the set of variance-reduction strategies that make the gradient estimator usable as the objective itself.9 The general lesson transfers: when your objective is defined as an expectation under the model you are training, the estimator’s variance becomes a first-class design problem rather than an implementation detail, and methods that ignore it fail for reasons that look like optimization instability. The same structure appears wherever a model is optimized against its own samples under a fixed reference distribution, which is why preference optimization runs into a closely related set of estimator questions.14

Three other variance reducers are worth naming briefly, since they appear elsewhere in the book.

More samples, spent where they matter. The law is harsh but it is a law. If you have a fixed budget of teacher scoring calls, spending them on more positions from fewer sequences reduces divergence-estimate variance more than spending them on more sequences, because the estimate is a per-position average.

Dense computation on a subset. For monitoring, the exact answer on 100 positions usually beats a sampled answer on 10,000, because the exact answer has zero variance and no estimator bias. This is the same argument that makes Chapter 10’s cached-logit pipeline attractive: paying once for dense teacher outputs buys you exactness forever after.

Changing the objective so the estimator is easier. Skew KL, as used in DistiLLM, evaluates the divergence against a mixture of teacher and student rather than against either alone, which bounds the ratio and therefore bounds the variance of any ratio-based estimate.10 Bounded divergences in general, from Chapter 3, have this property: a divergence no single token can blow up is also a divergence no single sample can blow up. Wen and colleagues’ f-divergence work makes the same move from the generator side.11 When a divergence is chosen partly for estimator behavior rather than purely for what it does to the student’s distribution, that is a legitimate reason, and it should be stated as one.

4.9 What to be suspicious of#

None of what follows is settled advice. It is the honest state of the matter.

“Use k3” as unqualified advice. It is correct in the regime most people who give it are working in, which is fine-tuning-scale RL where the policy is initialized close to the reference model and stays there because a KL penalty keeps it there. Distillation from a base or pruned student is a different regime and the advice does not transfer. This chapter’s measurement is the evidence, and it is the course’s own, not a published result; I do not know of a paper that isolates the crossover cleanly, which is a gap rather than a claim that nobody has noticed.

Reported KL numbers in papers, when the estimator is not stated. A great deal of published distillation work reports “KL” on training curves without saying whether it is a dense sum, a sampled , a sampled , or an importance-weighted forward estimate. Table 4.3 shows those can differ by factors of two to three on the same distributions. Comparing across papers on this quantity is close to meaningless unless both stated their method, and the surveys that tabulate such numbers across dozens of methods inherit the problem wholesale.17

Small measured divergences as evidence of similarity. Chapter 3’s Pinsker exercise already made this point from the deterministic side: a total variation distance of 0.001 certifies nothing about KL, because the KL can be a thousand nats hiding on a token with negligible mass.12 The sampled version is worse, because the tokens carrying the divergence are exactly the ones your sampler is least likely to visit. Both failures have the same shape and both live in the tail.

The assumption that a stable monitor means a stable run. The Lab 07 case produces a sampled forward KL that is stable, small, and completely uninformative, while the student is doing the worst thing it can do. Stability of an estimate is a property of the estimator’s variance, not of the system being measured, and the two are only linked when the estimator is actually seeing the thing that would change.

4.10 Where this lands in the labs#

Lab 00 §9 is the cell that produces Table 4.3, and it is thirty lines that run in a couple of seconds; the thing it does that this chapter cannot is fail. The assertion std(k3) > std(k1) in the far regime is written the “wrong” way around on purpose, so that anyone who tidies the cell into the version the folklore predicts breaks the build. Solutions 00 Exercise 4 extends the same construction to the forward direction and produces Table 4.4, including the yardstick row that makes the 79-fold gap between forward and reverse estimation on identical draws a measured quantity rather than an argument. Solutions 00 Exercise 3, on the chi-squared divergence, supplies the missing piece of the variance story by measuring growing by ten orders of magnitude over a logit range where KL grows by a factor of five. Lab 07’s Exercise 2 puts the trap in production clothing with the against demonstration, and it is worth running that one before Chapter 12 rather than after.

4.11 Exercises#

  1. §4.3.2 derives where . Table 4.3 shows overestimating in both regimes. What does that tell you about the sign of for these distributions, and is there a reason to expect that sign generally? Construct a two-token example where underestimates, or argue that none exists.

  2. By analogy with for the reverse direction, propose a low-variance biased estimator of from samples of . State its bias to leading order using the same expansion technique as §4.3.2. Then say whether it inherits the heavy-tail problem of §4.5, and why.

  3. Before looking at Figure 4.1, predict at roughly what true KL the standard deviations of and cross, using only the expansion in §4.4.1 and the fact that ’s variance is governed by . You will need a rough relationship between KL and chi-squared for nearby distributions; derive one for the case where is small and roughly symmetric. Then compare with the figure and account for the gap.

  4. A colleague’s on-policy run shows a panel labeled “KL” that has sat between and for six hundred steps, while the sample text has become repetitive and the average generation length has fallen by two thirds. The panel is computed as (w * (p_log[x] - q_log[x])).mean() where w = (p_log[x] - q_log[x]).exp() and x are the student’s sampled tokens. Say exactly what quantity that panel is estimating, what its true value is likely doing, and what single additional panel you would add to make the situation legible. Then say what you would change about the loss, if anything.

  5. You are training a 1.7B student against an 8B teacher on the reference machine, batch of 8 sequences at 1,024 tokens, vocabulary 49,152. Estimate the additional memory traffic per step of computing a dense divergence over the full vocabulary versus a sampled estimate at one token per position, in bytes and in milliseconds at 273 GB/s. Then say at what monitoring interval the dense computation becomes free in practice, and what that implies about §4.6.2’s second recommendation.

  6. Derive from the quadratic in §4.8. Then explain, in terms of what “unbiased” means, why plugging in a estimated from the same batch you apply it to breaks unbiasedness. Propose a scheme that recovers unbiasedness at some cost, and say what the cost is.

  7. You are 200 steps into an on-policy run from a pruned student and you want to know whether you are in the close regime or the far one, using only sampled quantities. Design a test. It should be cheap, should not require dense teacher logits, and should give an answer you would defend. Say what would make it give the wrong answer.



  1. Loubna Ben Allal et al., “SmolLM2: When Smol Goes Big: Data-Centric Training of a Small Language Model,” arXiv:2502.02737 (2025). https://arxiv.org/abs/2502.02737 The 49,152-entry vocabulary is the tokenizer used by the student models throughout this book’s labs. ↩

  2. The shape of what a serving stack returns is not incidental to this chapter. vLLM’s completion response returns log-probabilities for the sampled token and optionally a small number of alternatives per position, which is the grey-box row of Chapter 1’s access table. Woosuk Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention,” SOSP 2023, pp. 611-626, arXiv:2309.06180. https://arxiv.org/abs/2309.06180 Chapter 15 covers the serving path in detail. ↩

  3. John Schulman, “Approximating KL Divergence,” blog post, joschu.net, 7 March 2020. http://joschu.net/blog/kl-approx.html Accessed 1 August 2026. This is a personal blog post with no venue or DOI, and it is the standard source for the k1, k2, and k3 names and for the observation that is unbiased and pointwise nonnegative. ↩

  4. The invariance of under adding an affine term to the generator is a standard property of the f-divergence family introduced by Csiszár. Imre Csiszár, “Information-type measures of difference of probability distributions and indirect observations,” Studia Scientiarum Mathematicarum Hungarica 2 (1967): 299-318. Independently developed by S. M. Ali and S. D. Silvey, “A general class of coefficients of divergence of one distribution from another,” Journal of the Royal Statistical Society, Series B 28, no. 1 (1966): 131-142. ↩

  5. The clearest large-scale statement of the reverse-direction preference in language-model distillation is Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang, “MiniLLM: Knowledge Distillation of Large Language Models,” ICLR 2024, arXiv:2306.08543. https://arxiv.org/abs/2306.08543v2 Note that the arXiv landing page currently shows a later retitled version; the ICLR 2024 title is the one used here. ↩

  6. Rishabh Agarwal et al., “On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes,” arXiv:2306.13649 (2023), ICLR 2024. https://arxiv.org/abs/2306.13649 The method is generalized knowledge distillation, GKD; the name does not appear in the title. ↩↩

  7. Mingyang Song and Mao Zheng, “A Survey of On-Policy Distillation for Large Language Models,” arXiv:2604.00626 (2026). https://arxiv.org/abs/2604.00626 This is an unrefereed preprint whose arXiv comment reads “Ongoing Work”; I cite it for its catalogue of reported failure modes rather than for any empirical claim of its own. ↩

  8. Entropy as the monitored quantity during training on a model’s own samples is best documented in the reinforcement-learning-with-verifiable-rewards literature, where the collapse dynamics were characterized directly. Ganqu Cui et al., “The Entropy Mechanism of Reinforcement Learning for Reasoning Language Models,” arXiv:2505.22617 (2025). https://arxiv.org/abs/2505.22617 ↩

  9. Gu et al., “MiniLLM,” §3. The paper derives a policy-gradient form of the sequence-level reverse-KL objective and introduces single-step decomposition, teacher-mixed sampling, and length normalization specifically to control the gradient estimator’s variance. https://arxiv.org/abs/2306.08543v2 ↩

  10. Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se-Young Yun, “DistiLLM: Towards Streamlined Distillation for Large Language Models,” arXiv:2402.03898 (2024), ICML 2024. https://arxiv.org/abs/2402.03898 The skew KL objective evaluates the divergence against an interpolation of the two distributions, which bounds the ratio that appears in any sampled estimate. ↩

  11. Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou, “f-Divergence Minimization for Sequence-Level Knowledge Distillation,” arXiv:2307.15190 (2023), ACL 2023. https://arxiv.org/abs/2307.15190 ↩

  12. The constant in the inequality Chapter 3 uses is due independently to Csiszár and to Kullback rather than to Pinsker’s original 1964 statement, which had a weaker constant. S. Kullback, “A lower bound for discrimination information in terms of variation (Corresp.),” IEEE Transactions on Information Theory 13, no. 1 (1967): 126-127. DOI: 10.1109/TIT.1967.1053968; M. S. Pinsker, Information and Information Stability of Random Variables and Processes (San Francisco: Holden-Day, 1964). ↩

  13. Sequence-level knowledge distillation is itself a Monte Carlo approximation with a sample size of one: the intractable sum over all output sequences is replaced by the teacher’s single most likely output. Yoon Kim and Alexander M. Rush, “Sequence-Level Knowledge Distillation,” arXiv:1606.07947 (2016), EMNLP 2016. https://arxiv.org/abs/1606.07947 Chapter 11 examines how much that approximation gives up and why it works anyway. ↩

  14. The general structure of learning from one’s own samples under a fixed reference distribution, and the estimator questions it raises, is shared with preference optimization. Rafael Rafailov et al., “Direct Preference Optimization: Your Language Model is Secretly a Reward Model,” arXiv:2305.18290 (2023), NeurIPS 2023. https://arxiv.org/abs/2305.18290 ↩

  15. The bias-variance framing applied to soft targets in distillation generally, rather than to divergence estimation specifically, is developed in Helong Zhou et al., “Rethinking Soft Labels for Knowledge Distillation: A Bias-Variance Tradeoff Perspective,” arXiv:2102.00650 (2021), ICLR 2021. https://arxiv.org/abs/2102.00650 ↩

  16. Exposure bias, the reason on-policy sampling is attractive in the first place and therefore the reason this chapter’s estimators are needed at all, was formalized for sequence models long before distillation adopted the idea. Marc’Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba, “Sequence Level Training with Recurrent Neural Networks,” arXiv:1511.06732 (2015), ICLR 2016. https://arxiv.org/abs/1511.06732 ↩

  17. For the broader map of where sampled estimation sits within language-model distillation methods, see Xiaohan Xu et al., “A Survey on Knowledge Distillation of Large Language Models,” arXiv:2402.13116 (2024). https://arxiv.org/abs/2402.13116 The pre-language-model literature is surveyed in Jianping Gou, Baosheng Yu, Stephen J. Maybank, and Dacheng Tao, “Knowledge Distillation: A Survey,” International Journal of Computer Vision 129, no. 6 (2021): 1789-1819. https://arxiv.org/abs/2006.05525 ↩

Part II · The Objective

5

The Classical Objective

Here is a run that fails in a way you cannot see from the outside.

You have a 1.7-billion-parameter teacher, a 360-million-parameter student, and a loss that mixes the teacher’s softened predictions with the ordinary hard-label cross-entropy. You set the temperature to 2 and the mixing coefficient to 0.5, train for fifteen hundred steps, and the student comes out a few points ahead of the hard-label baseline. Encouraged, you sweep the temperature: 1, 2, 4, 8. The results form a clean curve with a peak at 2, and you write down that intermediate temperatures win, which is what the 2015 paper said, and you move on.

Nothing about that sequence of events is wrong except the conclusion. If your implementation of the soft loss omits a factor of , and many do, then every step of that temperature sweep also changed the effective learning rate of the soft term by a factor of about sixty-four across the range you swept. You did not measure a temperature effect. You measured a temperature effect convolved with a learning-rate effect, and you cannot tell from the numbers which one produced the peak. The loss ran without complaint. The curve looked like the curve in the paper. The only way to have known was to derive the gradient before running anything.

That is the chapter. The classical distillation objective is four lines of algebra, everyone believes they already understand it, and the places it goes wrong live in the derivation rather than in the code. So this chapter derives everything: the loss, its gradient, the factor, the high-temperature limit, and the reason the mixing coefficient has the shape it has. Then it covers what a decade of trying to use this objective has taught the field about when it stops working, which is a less tidy body of knowledge.

5.1 Two classes, then three#

Start with the smallest case that is not degenerate.

Suppose the model has to choose between two outcomes. It produces two logits, and , and the softmax turns them into probabilities:

where is the logistic function. The whole distribution is one number, and that number depends only on the difference of the logits. Adding a constant to both logits changes nothing, which is the shift invariance from Chapter 2 showing up in its simplest form.

Now bring in a teacher. The teacher has its own two logits, and , and its own probability . Say the correct answer is outcome 1 and the teacher assigns it 0.9. Ordinary supervised training would hand the student the label “outcome 1,” which is a distribution with 1.0 on the first entry and 0.0 on the second. Distillation hands the student the teacher’s 0.9 and 0.1 instead.

Definition

Soft target

The teacher’s full probability distribution over the output space at a given position, used as the training target in place of, or alongside, the one-hot hard label. The name distinguishes it from a hard label, which puts all its mass on one outcome and says nothing about the others.

The student gained exactly one number: the teacher’s confidence. That is worth something, because a target of 0.9 stops the student from driving its logit gap to infinity the way a target of 1.0 does, and that alone is a regularizer. It is not the thing Hinton’s paper is about. With two outcomes there is one wrong answer, and one wrong answer has no internal structure, because there is nothing for it to relate to.

Three outcomes is the smallest case where the interesting thing exists. Suppose the answer is dog, and the alternatives are cat and car. Let the teacher’s logits be in that order. Softmax them:

The teacher puts 1,097 times more probability on cat than on car. That ratio is a claim about the world: given that the answer is not dog, the teacher considers cat roughly a thousand times more plausible than car, presumably because it has learned that cats and dogs appear in the same slots and cars do not. The hard label assigns zero to both and makes no claim at all. That ratio, and the ten thousand ratios like it in a real vocabulary, is what Chapter 1 introduced as dark knowledge, and it is the entire payload distillation is trying to move.

Two facts about it matter before the algebra starts. The ratio is a difference of logits, since , so the similarity structure lives in the teacher’s logit gaps rather than in its probabilities; the probabilities are a monotone re-encoding of the gaps that squashes the small ones toward zero. And the squashing is exactly the problem, because the KL divergence weights the discrepancy at outcome by , the teacher’s probability there. With , the car term contributes about two parts in a million to the loss. The information is present and carries no weight. A student trained on this target at face value spends all of its gradient on getting dog right and never hears the claim about cat and car.

Hinton’s fix is to raise the temperature.

5.2 Temperature, and what it actually changes#

Definition

Temperature-softened distribution

The distribution obtained by dividing every logit by a constant before the softmax: . Larger moves mass from the top of the distribution into the tail; smaller concentrates mass on the argmax. recovers the model’s own distribution.

Softening the three-class teacher at gives

and the cat-to-car ratio drops from 1,097 to 5.755. Note what happened to both quantities. The ratio fell, because and dividing the logits by 4 divides every log-ratio by 4. The tail’s share of the mass rose, from two millionths to three percent.

This is the whole mechanism, and it is worth stating in a form that makes the trade explicit. Temperature neither creates structure nor destroys it. Matching a softened student to a softened teacher at temperature imposes the constraint for every pair, which is the same set of constraints at every . What temperature changes is the weight the loss puts on each constraint, because forward KL weights the discrepancy at outcome by , and is what temperature moves.

Read the two limits off that statement.

As , converges to a one-hot distribution on the teacher’s argmax. All the weight lands on one outcome and the objective becomes cross-entropy against the teacher’s most likely token, which is a hard label that happens to have been produced by the teacher rather than by an annotator. Distillation at very low temperature is not distillation; it is training on the teacher’s predictions as labels, which is the sequence-level idea of Chapter 11 arrived at from the wrong direction.14

As , converges to uniform. Every outcome gets weight , including the forty thousand tokens the teacher has effectively ruled out and whose logits are close to arbitrary. Section 5.4 shows that this limit has a clean closed form, and that the closed form is plain regression on logits.

Neither endpoint is where you want to be, which is why an intermediate temperature wins. The argument for the interior is an observation rather than folklore: you want enough weight on the tail to hear the teacher’s real preferences, and not so much that you also hear its noise floor.

2026-08-01T07:26:46.343346 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 0.25 0.5 1 2 4 10 t e m p e r a t u r e         ( l o g   s c a l e ) T 1 0 − 9 1 0 − 8 1 0 − 7 1 0 − 6 1 0 − 5 1 0 − 4 1 0 − 3 1 0 − 2 1 0 − 1 1 0 0 softmax probability p(the) = 0.8823 p(qq) = 1.5e-05 p(the)/p(xylophone) = 2.72 at T = 10 the a an cat dog xylophone qq
Figure 5.1 Raising the temperature lifts the teacher's tail into the range where it can influence a loss: on a fixed seven-token logit vector, the lowest-scoring token moves from a probability of 1.5e-5 at T=1 to within a factor of three of the top token at T=10.

The seven-token example the labs use makes the movement visible. With logits standing in for the, a, an, cat, dog, xylophone, qq, the top token holds 0.8823 of the mass at and the last holds . At , xylophone is within a factor of three of the, which is a different failure: the student is being asked to model the teacher’s opinion about a token the teacher has no opinion about.

5.2.1 Quantifying dark knowledge#

“Dark knowledge” is easy to gesture at and worth measuring, because if you cannot measure it you cannot tell whether your temperature is doing anything. Two numbers do the job.

The first is the mass outside the top-1 entry, . That is an upper bound on how much of the loss can be about anything other than the argmax.

The second is sharper. Decompose the teacher’s entropy by conditioning on whether the outcome is the argmax:

where is the binary entropy of the split between “argmax” and “not argmax,” and is the teacher’s distribution over the non-argmax outcomes, renormalized to sum to one. The second term is exactly the information about which wrong answer, given that it is a wrong answer. That is dark knowledge in nats, and it is precisely the quantity a hard label destroys and a top-1 agreement metric cannot see.

Table 5.1 Dark knowledge in the seven-token teacher of Figure 5.1, computed from the logit vector, in nats.

mass outside top-1 dark knowledge share of
0.25 0.000052 0.0006 0.0000 3%
0.5 0.0091 0.0573 0.0053 9%
1.0 0.1177 0.4479 0.0856 19%
2.0 0.3647 1.0527 0.3966 38%
4.0 0.5969 1.5783 0.9041 57%
10.0 0.7613 1.8748 1.3252 71%

At this teacher spends 81 percent of its uncertainty on the single question “is it the argmax or not,” which a hard label already answers. At that is down to 62 percent. The temperature knob is buying the loss access to the other half of the teacher’s uncertainty, and the table tells you how much you bought.

Run this on your own teacher and your own corpus before choosing . A teacher on code or structured output will be far more peaked than this toy and will need more heat to expose the same fraction of its tail; a teacher on open-ended prose will already be diffuse at and may need almost none. The useful temperature is a property of the teacher and the domain, not a constant. Pick the setting where a meaningful fraction of the teacher’s entropy has moved outside the top-1 entry and the tail has not yet approached uniform; for most instruction-following work that lands between 2 and 4. At with a confident teacher you are training against something close to a hard label, and the run looks broken because it is barely distilling.

That measurement is also the check to run when a temperature sweep surprises you. If wins your sweep outright, measure the teacher’s entropy before you believe the number: a teacher that is unusually diffuse to begin with leaves less room for flattening to hurt, so a high optimum on that teacher is a property of the teacher and not a refutation of Hinton’s argument. The same sweep on a peaked teacher will put its optimum several settings lower.

Here is the measurement, written to run on a batch of real teacher logits.

import torch

def dark_knowledge_budget(logits, mask, temperatures=(1.0, 2.0, 4.0, 8.0)):
    """Mass and nats living outside the top-1 entry, per temperature.

    logits: [B, T, V] teacher logits.  mask: [B, T] bool, positions to count.
    Returns one row per temperature: (T, mean tail mass, mean total entropy,
    mean dark-knowledge nats).
    """
    rows = []
    flat = logits[mask].float()                    # [N, V]
    for temp in temperatures:
        logp = torch.log_softmax(flat / temp, dim=-1)
        p = logp.exp()
        H = -(p * logp).sum(-1)                    # total entropy, nats
        top = p.max(-1).values
        tail = (1.0 - top).clamp_min(1e-12)
        # entropy of the renormalized non-argmax distribution, via the
        # conditional decomposition H(p) = H_b(top) + tail * H(tilde p)
        Hb = -(top * top.clamp_min(1e-12).log() + tail * tail.log())
        dark = H - Hb
        rows.append((temp, float(tail.mean()), float(H.mean()), float(dark.mean())))
    return rows

The last column is the one to watch. If it is near zero at every temperature you tried, your teacher has nothing to teach beyond its argmax on this corpus and you should expect distillation to look like training on teacher labels.

5.3 The loss#

With soft targets defined, the objective is short to state. Let be the student’s logits at a position, the teacher’s, the one-hot hard label, and write

The soft term is the forward KL from teacher to student at temperature , scaled by :

The hard term is ordinary cross-entropy at :

And the combined objective is a convex mixture of the two:

Definition

Mixing coefficient (alpha)

The weight on the soft-target term of the combined distillation loss, with on the hard-label cross-entropy. is ordinary supervised training with no teacher; is pure distillation with no ground-truth signal; the interior is a mixture.

Before the derivation, four notes on what the definition commits you to.

The direction of the KL is forward, meaning the teacher is the reference distribution and the student is the argument. This is Hinton’s choice and it is the classical objective. It is also a design decision with visible consequences for what the student’s text looks like, and Chapter 6 is about that decision, along with the reverse-KL alternative,24 the wider f-divergence family,20 and the skewed variants that address the unbounded penalty forward KL charges when the student puts near-zero probability where the teacher has mass.23 For this chapter, forward KL is the objective.

Some implementations write the soft term as a cross-entropy rather than a KL. Since the teacher is fixed, the two differ by the teacher’s entropy , which does not depend on the student’s parameters, so their gradients are identical. The loss values differ, and that matters when you compare a number in your log to a number in a paper, but nothing about training changes.

The hard term is computed at , always. Softening the hard label is meaningless because a one-hot distribution is invariant to temperature, but softening the student in the hard term is a real and different objective, and it is not what the classical loss does. If your implementation divides the student’s logits by before the cross-entropy against , it is computing something else.

The objective as written says nothing about where the inputs come from. It assumes a corpus exists and that teacher and student are both scored on it, which is the response-based, off-policy corner of the taxonomy from Chapter 1.15 DistilBERT is an early large-scale instance of exactly this loss on a transformer, with the soft term mixed against a masked-language-modeling loss instead of next-token cross-entropy.21 Replacing the fixed corpus with text the student generates gives the on-policy family of Chapter 12,19 and the language-model-specific survey literature organizes the whole space around those two axes.18

In the labs: Lab 03

The lab re-verifies both endpoints of the mixture at the door of the training loop, on synthetic logits, before any model loads: hinton_kd_loss at matches F.cross_entropy to within , and at matches the standalone forward-KL function with the scaling on. The point is that the hard arm really is a no-teacher baseline and the soft arm really has no ground-truth term, so the comparison between them measures the thing it claims to.

5.4 The gradient, derived#

Everything interesting about this objective is in its derivative with respect to a student logit, so take it.

Fix a position. The soft term without the factor is

where the constant is and is independent of . Write the student’s softened log-probability using the log-sum-exp form from Chapter 2:

Differentiate with respect to :

where is 1 when and 0 otherwise. Substitute:

The teacher’s probabilities sum to one, so the second sum is 1 and

That expression repays being read three ways.

At it is the residual. The gradient is , the difference between what the student currently believes and what the teacher believes, entry by entry. Nothing else appears; the update pushes the student’s probability toward the teacher’s by exactly their current gap. If the teacher is a one-hot label, this is , the ordinary softmax cross-entropy gradient that every deep learning course derives in its first month. The distillation gradient is that same expression with a soft distribution in place of the one-hot.

The whole temperature effect is a factor of out front, with the functional form unchanged. This is the most useful fact in the chapter for debugging, because any custom distillation loss can be checked in five lines: compute the two softened distributions, form , and compare to what autograd gives you. If they disagree, your loss is not the loss you think it is.

Masking is part of the formula. In a real training loop you average over supervised positions rather than summing over all of them, so with a boolean mask over positions and supervised positions, the gradient is

where broadcasts the position mask across the vocabulary axis. That is not cosmetic. It means long sequences in a batch contribute more total gradient than short ones, and the same arithmetic reappears on the evaluation side in Chapter 16. Sorting an eval set by completion length and scoring each stratum separately, which is the method §16.3 sets out, produces a long stratum that scores several points higher than the short one, and the mechanism is this weighting rather than anything about the model: a per-token agreement average mechanically favours long strata, because deep positions inside a long completion are pinned down by local context while the first few tokens after a prompt are the hardest ones in the set. Chapter 16 also gives the failure signature that tells you the effect has gone the wrong way, which is a KD advantage concentrated on short completions and usually means prompt tokens have leaked into the mask.

In the labs: Lab 01

The identity is checked against autograd rather than trusted. On a batch with a deliberately short second row so the mask is exercised, at , the largest absolute difference between PyTorch’s gradient and the analytic is under . That assertion is the reason this chapter can state the identity flatly.

5.5 Where the comes from#

The in front of the residual is the visible half of the temperature’s effect on gradient magnitude. The other half hides inside the residual itself.

Softening pushes both distributions toward uniform. Pushing two distributions toward the same place pushes them toward each other, so shrinks as grows. How fast? Expand the softmax for large . With small,

Sum over the vocabulary. If the logits are zero-mean, meaning , the sum of the first-order terms vanishes and , where is the vocabulary size. So

and the residual is

which is , as promised. Multiply by the explicit from the derivative:

The soft term’s gradient falls off as . The hard term’s gradient does not move at all, because the hard term is computed at and knows nothing about the temperature you chose for the soft term. So changing silently changes the ratio between the two terms’ contributions to the update, even though has not moved. Multiplying the soft term by cancels the scaling exactly, and that is where the factor comes from.

Definition

T² correction

The factor of applied to the soft-target term of the distillation loss. The softened KL’s gradient with respect to student logits scales as ; the hard-label term’s does not. Multiplying the soft term by restores comparable gradient magnitudes so that the mixing coefficient means the same thing at every temperature.

Hinton’s paper states the correction in one sentence and gives the same reason: because the soft gradients scale as , they should be multiplied by so that the relative contributions of the hard and soft targets stay roughly unchanged when the temperature is varied during experimentation.1 It is one line in a workshop paper, and it is the line most often dropped in reimplementation.

2026-08-01T07:26:48.053975 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 1 2 3 4 5 6 7 8 t e m p e r a t u r e   T 1 0 − 2 1 0 − 1 1 0 0 ‖ ∇ ‖ z s o f t 2 a s y m p t o t e       =   0 . 1 8 3 7 ‖ ̄ − ‖ ̄ / z v V corrected T q p T 2 ‖ ( − ) / ‖ T T uncorrected ‖ ( − ) / ‖ q p T T T falls 86x from T=1 to T=8 +16% from T=1 to T=2, then drifts down: the first-order expansion is loose at low T
Figure 5.2 On a fixed pair of eight-entry logit vectors, the uncorrected soft-loss gradient norm falls by a factor of 86 between T=1 and T=8, while the T-squared-corrected norm stays inside a factor of 1.6 across the same range.

Figure 5.2 shows the correction working on a fixed toy pair. Two details in it are worth naming because they are places the clean story is approximate.

The first is that the uncorrected gradient falls by 86 times between and , not by the 64 times that predicts. The law is the high-temperature asymptote, and at a logit vector with a spread of six is nowhere near the regime where is small. Below the asymptote the decay can be steeper or shallower depending on how spread out the logits are. Lab 01 measures roughly a 64-fold drop on random unit-scale logits over the same range, which is closer to the asymptote because those logits are less spread out. Both are the same phenomenon at different points on the curve.

The second is that the corrected line is not flat. It rises 16 percent from to , then drifts down toward , its asymptote, where the bars denote mean-centered logits. “Roughly constant” is the honest description, and roughly constant is all the correction needs to be, since its job is to keep interpretable rather than to hold a physical constant.

Field note

I had this half-wrong for longer than I want to admit. I could see the explicit in the gradient, so I believed the correction should be a factor of , and I spent an afternoon convinced that the paper had a typo. The missing piece is that the residual is not a fixed quantity that the scales; it is itself a function of , and it shrinks at the same rate. Two factors of , one visible in the chain rule and one buried in how softening moves the two distributions toward each other. The general lesson is the one from Lab 01: when a scaling argument involves a quantity that itself depends on the parameter you are scaling, derive it, do not count the factors you can see.

5.6 The high-temperature limit is logit matching#

Cash in the second half of the previous section’s expansion now. It found that at large the gradient of the uncorrected soft loss approaches . Now ask what objective has that gradient. The derivative of

with respect to is , the same thing. So in the high-temperature limit, minimizing the softened KL is minimizing the squared error between the student’s logits and the teacher’s. All the structure the softmax imposes, the competition among outcomes for a fixed budget of probability mass, has dropped out, leaving plain regression.

Apply the correction to both sides and the temperature disappears entirely:

which is a -independent objective. That is the cleanest possible statement of what the correction does: it makes the loss converge to something finite and fixed as grows, instead of converging to zero.

In the labs: Lab 01

The limit is asserted numerically, not argued. At on mean-centered synthetic logits, the gradient of the softened KL and the gradient of agree to within everywhere. Solutions 03 makes the same point with a cosine similarity between the two gradient directions, which lands above 0.99 at and at 0.82 at , confirming that the alignment is a high-temperature effect rather than something that was always true.

5.6.1 The zero-mean assumption#

The expansion above used and . Hinton’s paper makes this assumption explicitly, zero-meaning the logits separately for each transfer case, and it is worth being precise about what it costs.2

Redo the expansion without it. If is the mean logit, then , so and the residual becomes

The limit is still logit matching, but on mean-centered logits. That is a consequence of softmax’s shift invariance rather than a failure of the theory: the mean logit has no effect on any probability, so no probability-space objective can constrain it.

The practical trap is in the other direction. Some pipelines skip the softmax and regress the student’s raw logits onto the teacher’s raw logits, on the grounds that the high-temperature limit says this is what distillation converges to anyway. Without centering, that objective penalizes the student for having a different mean logit than the teacher, a quantity that does not affect the student’s predictions at all. The student spends capacity satisfying a constraint with no behavioral content and, depending on your weight decay, fights the optimizer over it. If you are going to run the high-temperature limit on purpose, center both sides.

The other place the assumption fails is about rather than the mean. The first-order expansion is accurate once is small for every , and real language-model logits have a spread of tens, so “small” means a temperature well into the double digits. The lab verifies the limit at and finds the gradient alignment only 0.82 at . Nobody trains at . The logit-matching interpretation is therefore a statement about a limit that practice does not visit, and you should not reason about a run at as though it were doing logit matching.

What the limit is good for is explaining why very high temperatures hurt. Logit MSE weights every vocabulary entry equally. The teacher’s logit on its thirty-thousandth-choice token gets exactly as much of the student’s capacity as its logit on the answer, and that logit is close to noise: the teacher was never trained to place it carefully, because nothing in its own loss ever depended on it. Hinton makes the same observation from the other side, arguing that intermediate temperatures work well precisely because they partially ignore the very negative logits, which carry little signal and considerable noise.3 Solutions 03 reports the consequence on the synthetic teacher it builds: by , more than 90 percent of the teacher’s mass has moved onto wrong answers and the entropy is approaching nats, the uniform ceiling for that toy vocabulary. There is not much teacher left in a target like that.

5.7 What happens when the factor goes missing#

Drop the from the soft term and keep everything else. The combined loss becomes

and its gradient is

The first term now carries a factor that shrinks as relative to the second. Set and go from to : the soft term’s contribution to every update drops by about sixteen while the hard term’s holds still. In effect, has moved from 0.5 to something near 0.06 without you touching it.

The consequence for experiments is the thing to internalize. Sweeping temperature at fixed and fixed learning rate with the correction missing sweeps two variables at once: the shape of the target and the effective step size on the soft term. Any peak you find is uninterpretable, because it might be the temperature optimum of §5.2 or it might be one temperature happening to land the soft term’s effective learning rate in a good place. You cannot separate them after the fact, because you did not vary them independently.

Watch out

A missing produces no error, no warning, and a loss curve of entirely normal shape. It is in the same family as an inverted divergence direction or an off-by-one mask: the run trains, the number goes down, and the thing being trained is not the thing you specified. Check the factor before you trust any temperature comparison, including one in a paper that does not say which convention it used.

There are three ways to notice, and none of them requires a training run.

Log the two terms’ gradient norms separately. This is the direct measurement and it takes about ten lines. Compute the soft loss alone, take its gradient norm, zero the gradients, compute the hard loss alone, take its gradient norm, and log the ratio. If the correction is present, the ratio should be roughly stable across temperatures. Lab 03’s pre-flight asserts exactly this in the form for the soft term, a band chosen to pass with the correction and to fail without it, since the uncorrected ratio would be near 4. Solutions 03 runs the wider version: between and the raw ratio should land in the band , since , while the corrected ratio should stay in .

The measurement is short. Look at how the returned soft norm behaves as moves.

import torch

def term_gradient_norms(student_logits, teacher_logits, labels, mask, T, alpha):
    """Gradient norm of each loss term, taken separately, at one temperature.

    Run this at a few temperatures before a sweep. With the T^2 correction in
    place the soft norm should be roughly flat in T; without it, it falls ~T^2.
    """
    z = student_logits.detach().clone().requires_grad_(True)
    soft = (T ** 2) * kl_softened(z, teacher_logits, mask, T)     # your loss
    soft.backward()
    soft_norm = z.grad.norm().item()

    z.grad = None
    hard = cross_entropy_masked(z, labels, mask)                  # your loss
    hard.backward()
    hard_norm = z.grad.norm().item()

    return {"T": T, "soft": alpha * soft_norm, "hard": (1 - alpha) * hard_norm,
            "ratio": (alpha * soft_norm) / max((1 - alpha) * hard_norm, 1e-12)}

The ratio in the last field is the number that is supposed to control. If it moves when you change , does not mean what the config file says it means.

Watch for a sweep where wins outright. With the correction missing, low temperature gets the largest effective learning rate on the soft term, so it has a structural advantage that has nothing to do with the target. Solutions 03 predicts the opposite pattern for a correct sweep: the arm should land closest to the hard-label baseline of all of them, because a peaked teacher at is nearly a hard label. Getting the reverse ordering is a signal to check the loss before believing the result.

Report it either way. A temperature sweep must either keep the factor or say that it did not, and there is no third option that produces an interpretable number. If you keep it, run the gradient-norm ratio check above before you launch, so that you know the factor is in the code and not only in the documentation. If you sweep temperature without the correction, say so in the writeup, and say what the effective soft-term learning rate was at each temperature. That converts an uninterpretable result into an interpretable one about a joint sweep. The failure is not omitting the factor, which is sometimes a defensible choice if you retune the learning rate per temperature. The failure is omitting it silently.

5.8 The mixing coefficient#

The last free parameter of the classical objective is , and its behavior is less settled than the temperature’s.

Start with what each endpoint gives up. At the teacher is disconnected from the gradient and you are running ordinary supervised fine-tuning. At the ground truth is disconnected and the student is trained only to reproduce the teacher, including the teacher’s mistakes; on any position where the teacher is confidently wrong, a pure-soft student is trained toward the wrong answer with no counterweight.

There is a clean way to see the interior, and it is exact at . Take the gradient of the combined loss with the correction in place:

Set and the two softened distributions collapse to the same :

At unit temperature, the whole combined objective is a single cross-entropy against the blended target : a distribution that is the teacher’s, with extra mass moved onto the correct token in proportion to . That makes the trade concrete. Moving down increases the mass the target puts on the ground truth and shrinks the tail structure by the same factor, linearly, entry by entry. Above the two terms no longer collapse into one target, but the intuition survives: is trading a guarantee about the correct token against the amount of the teacher’s tail that reaches the gradient.

2026-08-01T07:26:49.270363 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 0.0 0.2 0.4 0.6 0.8 1.0 α       ( w e i g h t   o n   t h e   t e a c h e r ) 0.4 0.5 0.6 0.7 0.8 0.9 1.0 target mass on the correct token ground-truth mass 1.0000 -> 0.4726 dark knowledge 0 -> 0.6406 nats both exactly linear in alpha c o m p u t e d   e x a c t l y ,   T = 1 0.0 0.5 1.0 α       ( w e i g h t   o n   t h e   t e a c h e r ) held-out quality hard mixed soft Lab 03's three arms peak somewhere in here; the ordering inside is not settled expected shape, not a measurement 0.0 0.2 0.4 0.6 dark-knowledge content of the target (nats)
Figure 5.3 The alpha trade, exactly computed on the left and expected on the right: the blended target's ground-truth mass falls linearly in alpha while its dark-knowledge content rises linearly, and held-out quality is expected to peak somewhere in the interior, though the interior peak is not guaranteed.

5.8.1 The U-shape, stated as an expectation#

The standard claim is that held-out quality as a function of is an inverted U: worse at both endpoints, better in the middle, because pure hard labels throw away the teacher’s information and pure soft targets throw away the ground truth. That claim is repeated more confidently than the evidence supports, so it is worth splitting into halves.

The left half is on solid ground. If soft targets do not beat hard labels at all, with an in-family teacher, in-distribution data, and a healthy teacher in eval mode, the problem is the setup rather than the theory. Lab 03 expects the distilled arms to beat the hard-label arm by 2 to 8 points of top-1 agreement with the teacher at equal steps, and by a visible margin on expected calibration error, the binned gap between a model’s stated confidence and its realized accuracy.13

The right half is not. Whether beats depends on how much the ground truth adds beyond what the teacher already knows, and for a strong teacher on in-distribution data the answer is often “very little.” Hinton’s own guidance points this way: the paper reports that the best results generally came from putting a considerably lower weight on the hard-label objective than on the soft one, which is a claim that the optimum sits well toward the soft end.4 Lab 03 expects mixed at and soft at to land genuinely close, and says either ordering is a legitimate result.

Where the interior does tend to earn its keep is calibration rather than agreement. The hard term keeps pulling probability onto tokens that are actually correct, which keeps the student’s stated confidence tied to its realized accuracy, and a student trained purely against a teacher can drift into confidence its accuracy does not support. Lab 03’s expectation is that ’s advantage over shows up on ECE more than on agreement. That is a case where two metrics disagree about the ordering of two arms, and if you only logged one of them you would report a clean result and be half wrong.

Zhou and colleagues frame the soft-target weight as a bias-variance trade rather than a single global constant, arguing that the right weight differs across training examples.5 That is one explanation for why the U is so often flat: if the optimal varies per example, any fixed is a compromise, and small changes to a compromise do not move much.

So start at 0.5, which is the labs’ value, and move from there for reasons you can name. Move toward 1 when the teacher is much stronger than the labels, which is the usual situation when the labels are a scraped corpus and the teacher is a well-trained instruction model. Move toward 0 when the teacher is mediocre, out of domain, or confidently wrong on the cases you care about. Expect the curve to be flat between 0.5 and 1 and to fall off steeply below about 0.2, and expect to learn more from the ECE column than from the agreement column while you do it.

In the labs: Lab 03

The alpha comparison is three arms rather than a sweep: hard at 0.0, mixed at 0.5, soft at 1.0, with the same student, the same 1.7B teacher, the same seed, and throughout. Three points can establish that the interior beats the left endpoint. They cannot resolve the shape of the curve between 0.5 and 1.0, and the lab does not claim they can. If you want the shape you need a denser sweep and more than one seed, which is Chapter 18’s material.

Sweeping either knob has the same precondition, and it is worth stating once for both. Everything else has to be held fixed and asserted rather than assumed: seed, data, data order, batch size, gradient accumulation, learning rate, warmup, schedule shape, total steps, sequence length, and dtype. Lab 03’s discipline is to compute the set of configuration keys by which two arms differ and assert that it has exactly one element, so a sweep over must produce {"T"} and nothing else. An ablation that varies two things measures neither of them.

5.9 A better teacher can be a worse teacher#

The classical objective has an implicit assumption in it: that a more accurate teacher produces better soft targets. That assumption is false in at least two distinct ways, and both of them are instructive.

5.9.1 Label smoothing#

Label smoothing replaces the one-hot training target with a mixture: instead of 1.0 on the correct class and 0 elsewhere, it trains toward on the correct class and spread uniformly over everything else, so with every wrong token gets exactly of the target mass. It is a standard regularizer, it usually improves accuracy slightly, and it reliably improves calibration.

Now read that target with distillation in mind. A model optimized against it is being explicitly rewarded for making every wrong answer equally probable, and the entire content of a soft target is that the wrong answers are unequally probable. Label smoothing trains away the exact structure distillation exists to transfer, on purpose, as its objective.

Müller, Kornblith, and Hinton measured this and traced the mechanism.6 Their finding is that label smoothing tightens the clusters of penultimate-layer representations: examples of the same class are pulled into tighter, more equidistant groups, which erases information about how a given example resembles the templates of the other classes. That information is what the soft targets encode. So the smoothed teacher scores better on its own metrics and distills worse, and the effect is not small.

Solutions 03 makes the mechanism visible without training anything, on a synthetic 1000-entry teacher with dog at logit 9, cat at 4, and car at . It reads two dials: the ratio , a plausible wrong answer against an implausible one, which the original teacher puts near 1,100; and the standard deviation of the log-probabilities across all wrong answers, a one-number summary of how much wrong-answer structure exists at all, which starts around 2 nats and would be zero if every wrong answer were interchangeable.

At the label-smoothing optimum the ratio collapses to exactly 1.0 and the spread to exactly 0. Every wrong answer has become the same wrong answer. Halfway toward the optimum, standing in for a brief smoothed fine-tune, the ratio is already cut by a factor of about thirty and the spread by half. And the argmax never moves. stays high, the teacher’s accuracy is intact, and any evaluation that looks at the teacher’s top-1 predictions sees nothing wrong at all.

Watch out

Teacher accuracy is the wrong selection criterion for a teacher. Two teachers can have identical top-1 accuracy, identical loss, and identical benchmark scores, and one of them can have almost no distillable structure in its tail. If you are choosing between candidate teachers, measure the thing distillation actually reads: the wrong-answer spread from §5.2.1, or the dark-knowledge nats, on your own corpus. It takes one forward pass over a few hundred prompts.

The full result has a sharper form: a teacher can get better and become a worse teacher at the same time, because teacher quality for the purposes of distillation is measured in the wrong answers. This is also, as Chapter 1 noted, the cleanest available evidence that the information argument for why distillation works is doing real work. If the benefit were entirely a regularization effect from training against a smooth target, then label smoothing, which makes the target smoother, could only help.

5.9.2 The capacity gap#

The second way a better teacher can be a worse teacher is about size rather than training.

Definition

Capacity gap

The finding that a teacher much larger than the student can produce a worse student than a smaller, less accurate teacher would have. The student quality as a function of teacher size rises, peaks, and then declines, so the best teacher for a given student is often not the best available model.

Cho and Hariharan established the result on image classifiers and, more usefully, tried to find out why.7 They showed that increasing the teacher’s size past a point degrades the student, and they tested the obvious explanation directly by checking whether the student could match the teacher on the training data at all. It could not, even when given every advantage. Their proposed remedy is the part that makes the result memorable: stop the teacher’s training early. An early-stopped teacher is less accurate by every conventional measure, and it distills better.

2026-08-01T07:26:50.131737 image/svg+xml Matplotlib v3.10.9, https://matplotlib.org/ 100M 1B 10B 100B teacher parameter count (log scale) student held-out quality peak location depends on the student and the task gap-small 360M 2.7x student gap-large 1.7B 12.6x student student held fixed at 135M three explanations for the decline: representation, optimization, degenerate targets expected shape
Figure 5.4 The expected shape of the capacity-gap result, with the course's two probe points marked: student quality as a function of teacher size rises to a peak and then declines, so the largest available teacher is not automatically the right one.

There are three competing explanations in circulation, and the evidence does not cleanly separate them.

Representation. The student’s architecture cannot express the function the teacher computes, so the target is unreachable and the loss stops providing useful gradient. This is the explanation people reach for first and the hardest to test, because “cannot express” is a claim about a hypothesis class nobody can characterize for a transformer.

Optimization. The student could express something close to the teacher’s function but cannot find it. Cho and Hariharan’s inability to make the student fit the teacher even on training data points here rather than at representation, and §5.10 widens the seam considerably.

Degenerate targets. A very large, very well-trained teacher is extremely confident. Its distributions are close to one-hot, the tail carries almost no mass, and the soft target degenerates toward a hard label. On this account the capacity gap is not about capacity at all; it is the dark-knowledge measurement from §5.2.1 going to zero as teachers get better, and it predicts that raising the temperature should partially rescue the large-teacher case. I have not seen that prediction tested cleanly.

The three explanations are not exclusive and probably all contribute. What matters for planning is that they make different recommendations: the first says pick a smaller teacher, the second says train longer, and the third says raise the temperature. A fourth response, older than the diagnosis, is to stop relying on output matching alone and supervise the student’s intermediate layers as well, which is what FitNets introduced and what Chapter 14 develops.22

5.9.3 Patience and consistency#

Beyer and colleagues argued the second explanation hard enough to change what the recipe looks like.8 Their framing is that distillation should be understood as function matching, and that if you take that framing literally the standard recipe is doing it wrong in two ways.

Definition

Function matching

The interpretation of distillation as approximating the teacher’s input-output function rather than as a form of regularized supervised training. Taken literally it requires that teacher and student see identical inputs, including identical data augmentation, so that every training signal is an evaluation of the same function at the same point, and that training run long enough for the approximation to converge.

The first fix is consistency. If the teacher and the student see different views of an example, because augmentation was applied independently to each, then the student is being trained to match the teacher’s answer to a question the teacher was never asked. Fixing this means passing exactly the same augmented input to both.

The second fix is patience. Function approximation is a hard optimization problem, and Beyer and colleagues found that the gains kept arriving far past the point where a supervised recipe would have stopped, with training schedules long enough that the usual intuitions about overfitting stop applying. Their headline is that this patience is worth more than most of the architectural choices the field argues about, including the size gap between teacher and student.

The consistency half translates directly to language models. Teacher and student must see the same tokens under the same chat template with the same masking, which sounds obvious and is exactly what Chapter 7 shows is hard to get right. The patience half translates too, and it is the basis for the course’s own test of the capacity gap, which runs on the SmolLM2 family because that family publishes checkpoints at 1.7B, 360M, and 135M parameters along with the instruction corpus they were trained on, so teacher and student are in-family and the data is in-distribution for both.16

In the labs: Lab 03

The capacity-gap probe is two arms with one key different. gap-small distills the 360M model into the 135M model, a ratio of 2.7. gap-large distills the 1.7B model into the same 135M model, a ratio of 12.6. Same student, same , same , same seed, same data, same 1,500 steps: the configs differ in the teacher key and nothing else, and the lab asserts that before it runs. The expectation is that gap-large beats gap-small by less than the teacher-size ratio suggests, may tie it, and may lose. A small loss for the larger teacher is the capacity gap appearing in your own logs, and it is a real result. A large win for the larger teacher would be the surprising outcome, and would say your student had headroom the literature did not expect.

Solutions 03 turns the patience claim into an experiment on that probe, and the design is worth walking through because the obvious version of it is wrong. The question is whether doubling the training budget on gap-large closes the gap. Comparing gap-large at 3,000 steps to gap-large at 1,500 confounds patience with compute: the longer arm got twice the gradient updates, so its advantage tells you nothing about the capacity gap. Three comparisons do answer something. The patient arm at 1,500 steps against the original at 1,500 must match, since the configs agree up to that point and the seed is fixed, so any drift is a reproducibility bug rather than a finding, and it is free to check. The patient arm at 3,000 against the original at 1,500 asks whether more time helps at all. And the patient arm at 3,000 against gap-small at 1,500 is the capacity-gap question: does patience buy back what the oversized teacher cost. Beyer’s line predicts partial closure, roughly 1 to 3 points of top-1 agreement with eval KL still declining slowly, putting the patient large-teacher arm between the two originals. Full closure would say the gap in your setup was mostly an optimization-time artifact, which is reportable in either direction.

Two traps in that design deserve naming. Doubling max_steps stretches the cosine learning-rate schedule, so the patient run spends its extra steps at moderate learning rates rather than in the near-zero tail; resuming a finished run and bolting on 1,500 more steps with a fresh schedule breaks the comparison. And if agreement improves while ECE worsens through the extension, the student is memorizing the teacher’s confidence faster than its ranking, which is the known late-phase failure of long distillation runs. Stop at the ECE inflection rather than at the step budget.

Teacher size is the fourth knob in this chapter and the last one to get a diagnostic, so this is the place to put all four side by side. Every row is a failure I have watched happen, and the right column is what separates it from the failures it resembles.

Table 5.2 What each knob costs when you get it wrong, and the diagnostic that catches it.

Knob Set too low Set too high Diagnostic that catches it
Soft target is nearly a hard label; distil