This section is about making a model smaller, so that every token it generates costs less. There are two ways to do it. Distillation trains a small model to imitate a large one. Pruning deletes parts of a trained model. Nearly every small model that labs ship today was made with one or both, and they are how a frontier model's abilities reach a phone, a laptop, or a cheap tier of an API.

Why does the number of weights matter so much? A model writes text one token at a time. Each step of this decoding produces one new token for every conversation in the batch, and to take it the GPU must read every weight from its memory once (§20.1). With only a few conversations in the batch, the reading takes far longer than the arithmetic. So a step lasts about as long as it takes to read the weights.

Take Llama 3.1 8B on one H100 GPU. Its 8.03 billion weights, stored as 16-bit numbers of 2 bytes each, fill 16 GB. The H100 reads 3.35 TB per second from its memory. So one decode step takes at least 16 GB divided by 3.35 TB/s, about 4.8 ms, and a single conversation gets at most about 208 tokens per second. A model with 4 billion weights reads half as many bytes and runs about twice as fast.

Quantization, in §20.2 and §20.3, made each weight take fewer bytes. This section makes the number of weights smaller, which cuts both the bytes read and the arithmetic. The same reasoning is why labs train small models on far more data than the compute-optimal amount (§9.4). A model that will answer billions of requests should be small, and the question becomes how to make a small model as good as possible.

Here is the plan. We derive the loss that trains a small model on a large model's predictions, and see why it is multiplied by the square of a temperature. We look at the three ways to distill a language model and which divergence each one minimizes. Then we turn to pruning. We find which weights can go by asking how much removing each one changes its layer's outputs. Then we measure what removing them buys on a real GPU, which is often far less than the count of weights suggests. Then comes the recipe that combines the two. Prune a trained model to the shape you want, then repair it by distillation. Last, we see how distillation makes better drafts for speculative decoding.

Two ways to make a model smaller

Distillation trains a new, smaller model, the student, to imitate a larger trained model, the teacher. The student can have any architecture. What the teacher adds is a richer training signal than raw data gives. At every position the data gives only the one token that happened to come next. The teacher gives a whole probability distribution over the next token.

Pruning starts from the trained model and removes parts of it. Unstructured pruning removes individual weights, so every matrix keeps its shape and fills up with zeros. Structured pruning removes whole rows, columns, attention heads, or layers, and leaves smaller dense matrices. A middle ground, 2-of-4 sparsity, keeps exactly two weights in every group of four. NVIDIA GPUs since the Ampere generation (the A100) can skip the missing two in hardware.

The two are not rivals. The strongest recipes prune first and then distill, and a distilled model is usually quantized for serving as well. For each method an interviewer will ask which cost it cuts, the bytes read per step or the arithmetic, and this section answers that for each one.

Soft targets carry more than labels

A label says "this image is a 2". A trained teacher network looking at the same image says something richer, for example "a 2 with probability 0.9, a 7 with probability 0.09, a 3 with probability 0.01". The label carries one fact. The teacher's answer also says that this particular 2 looks a little like a 7 and hardly at all like a 3. That is knowledge about how the digits relate, and a student trained only on labels never sees it.

There is a catch. A very confident teacher puts nearly all its probability on the right class. It might give the 3 a probability of 0.0001 instead of 0.01. Probabilities that small barely change the training loss, and the student learns almost nothing from them. The fix is to soften the teacher's distribution before the student copies it.

Here is the setup. For one input, the teacher produces logits $\mathbf{u} \in \R^V$, one real score per class for $V$ classes. The student produces its own logits $\mathbf{z} \in \R^V$. For a language model, $V$ is the vocabulary size, and there is one such pair of vectors at every position. A softmax turns logits into probabilities by exponentiating each one and dividing by the total. A temperature $\tau > 0$ divides the logits first, as in sampling (§6.4),

Throughout this section $p$ is the teacher and $q$ the student. At $\tau = 1$ these are the models' ordinary predictions. A larger $\tau$ shrinks every gap between two logits by the factor $\tau$, so the distribution flattens. The ratio of two probabilities, $p_{\tau,i}/p_{\tau,j} = e^{(u_i - u_j)/\tau}$, moves toward 1, but the order of the classes never changes.

The cell checks the example and adds two more temperatures.

The relative probabilities of the wrong answers show how the teacher generalizes. Hinton, Vinyals & Dean (2015) made this point, and Hinton's talks called this information dark knowledge. Raising the temperature is how you make it visible. The student is trained at the same temperature, so it is asked to match the softened distribution. Afterwards it is used at $\tau = 1$.

The distillation loss

The student is pulled toward two targets. One is the true label, at the ordinary temperature. The other is the teacher's softened distribution, at temperature $\tau$. A weight $\lambda_{\text{KD}}$ between 0 and 1 balances the two.

Two measures of mismatch from §0.5 make up the loss. The cross-entropy $H(\mathbf{p}, \mathbf{q}) = -\sum_i p_i\log q_i$ is the student's average surprise when outcomes are drawn from $\mathbf{p}$. Against a one-hot label for the true class $c$ it is $-\log q_c$, the usual classification loss. The KL divergence $D_{\mathrm{KL}}(\mathbf{p} \Vert \mathbf{q}) = \sum_i p_i\log(p_i/q_i)$ measures how far $\mathbf{q}$ is from $\mathbf{p}$. It is never negative, and it is zero only when the two distributions are equal.

In plain words, the student pays for missing the label and for differing from the softened teacher, and the second payment is multiplied by $\tau^2$. The next part explains that factor.

Papers write the soft term in two forms that train the same student. The KL equals the soft cross-entropy minus the teacher's entropy $H(\mathbf{p}_\tau) = -\sum_i p_{\tau,i}\log p_{\tau,i}$, its own uncertainty, so $D_{\mathrm{KL}}(\mathbf{p}_\tau\Vert\mathbf{q}_\tau) = H(\mathbf{p}_\tau, \mathbf{q}_\tau) - H(\mathbf{p}_\tau)$. The teacher's entropy does not depend on the student. So the two forms have the same gradients, and papers write either one. The KL has the advantage of being zero when the student matches the teacher.

The order of the two arguments matters too. The teacher sits in the first slot, which makes this the forward KL. It is large wherever the teacher puts probability and the student puts little. So it pushes the student to cover everything the teacher considers plausible, a behavior called mode covering (§0.5). The reverse KL, with the student first, behaves differently, and on-policy distillation later in this section relies on that.

Why the soft term is multiplied by τ²

Dividing the logits by $\tau$ weakens the soft term's gradient twice over. The chain rule through $\mathbf{z}/\tau$ contributes one factor of $1/\tau$. And softened distributions are closer to uniform, so the student's and the teacher's differ by about $1/\tau$ as much. Together the gradient shrinks like $1/\tau^2$, and multiplying by $\tau^2$ undoes it. Here is the exact statement.

Proof of part 1. Write the soft cross-entropy with the log-softmax, $H(\mathbf{p}_\tau, \mathbf{q}_\tau) = -\sum_k p_{\tau,k}\big(z_k/\tau - \log\sum_j e^{z_j/\tau}\big)$. The derivative of $-\sum_k p_{\tau,k} z_k/\tau$ with respect to $z_i$ is $-p_{\tau,i}/\tau$. Since $\sum_k p_{\tau,k} = 1$, the derivative of the log-sum-exp term is $e^{z_i/\tau}/(\tau\sum_j e^{z_j/\tau}) = q_{\tau,i}/\tau$. So $\partial H/\partial z_i = (q_{\tau,i} - p_{\tau,i})/\tau$. The KL has the same gradient, because it differs from $H$ by a constant. ∎

This is the familiar "predicted minus target" gradient of a softmax with cross-entropy (§0.2), with an extra $1/\tau$ from the chain rule. It is equation 2 of Hinton et al.

In plain words, at high temperature the raw gradient of the soft term has size about $\lVert \mathbf{C}_V(\mathbf{z} - \mathbf{u})\rVert/(V\tau^2)$, so it shrinks like $1/\tau^2$. The hard-label gradient, $(1 - \lambda_{\text{KD}})(\mathbf{q}_1 - \mathbf{e}_c)$, does not depend on $\tau$ at all. Without the factor, raising the temperature from 2 to 10 would cut the soft term's influence 25-fold. Distillation would quietly turn back into training on labels. With the factor, the two terms keep their relative weight as $\tau$ changes, so $\lambda_{\text{KD}}$ means the same thing at every temperature. That is the reason Hinton et al. give for it, and it is the answer an interviewer wants, with both factors of $1/\tau$ named.

Part 2 also shows that matching logits is a special case of distillation. Caruana and his collaborators first trained small models to imitate large ones (Buciluǎ et al. 2006). Later they did it by copying the large model's logits with a squared error (Ba & Caruana 2014). Distillation at high temperature does exactly that, on centered logits. The centering matters. Softmax ignores a constant added to every logit, so a model's logits are only defined up to such a constant. Hinton's equation 4, $(z_i - u_i)/(V\tau^2)$, assumes both vectors already have mean zero.

At lower temperatures the loss pays less attention to logits far below the mean. Hinton et al. note that this can help. A teacher's very negative logits are barely constrained by its own training, so they may be noise.

The cell reproduces those entries. Then it takes a random teacher and student with $V = 10$ and checks the theorem as $\tau$ grows. The raw gradient should fall like $1/\tau^2$. The scaled gradient should settle to a constant. And the error of the logit-matching approximation should fall like $1/\tau$.

From $\tau = 1$ to $\tau = 64$ the raw gradient falls from 0.784 to 0.00022, more than three thousandfold. The scaled one moves only from 0.784 to 0.890. The relative error of logit matching halves each time $\tau$ doubles, from 0.098 at $\tau = 8$ to 0.050 at 16.

import numpy as np

def kd_loss(z, u, tau):
    raise NotImplementedError

def kd_grad(z, u, tau):
    raise NotImplementedError

import numpy as np

def high_temp_limit(z, u, tau):
    raise NotImplementedError

The widget shows both effects of the temperature at once, the wrong classes rising into view and the soft gradient fading without the factor. Drag the teacher's logits, slide the temperature, and compare the soft gradient with and without $\tau^2$ against the hard-label gradient.

Soft targets help most when data is scarce

The argument so far says that soft targets carry more information per example than labels. A small experiment tests that claim. The teacher is a fixed network with one hidden layer of 64 ReLU units that maps 8 input features to 6 classes. The student is weaker by design. It is a linear softmax model on 32 fixed features, each a ReLU of a random projection of the input.

The student is trained on $n$ inputs in one of three ways. One uses the teacher's labels, its most likely class. The other two use its soft distributions, at $\tau = 1$ or at $\tau = 4$. Training uses L-BFGS, a standard quasi-Newton optimizer, with a small L2 penalty. We measure agreement, how often the student's prediction matches the teacher's on 2,000 fresh inputs, averaged over two draws of the training set.

That is the pattern to expect. Soft targets are worth the most when examples are scarce compared with what the student must learn, and the advantage fades as data grows. The temperature matters little here, because the teacher's logits are moderate. The labels here are the teacher's own top classes, so the two targets never disagree. The label-trained student is the case $\lambda_{\text{KD}} = 0$, and the distilled students use the soft term alone, $\lambda_{\text{KD}} = 1$.

Distilling a language model, token by token or sequence by sequence

A language model is a classifier applied once per position, so the loss above applies to every token. But a sequence raises a question that a single classification does not. At position $t$ the student predicts the next token from a prefix, the tokens $\mathbf{y}_{<t} = (y_1, \dots, y_{t-1})$ that came before. Where do those prefixes come from? They can come from a training corpus, from the teacher, or from the student itself. That choice decides which divergence you minimize. §14.2 named the three resulting methods and compared what each needs from the teacher. Here we derive which divergence each one minimizes and what it costs, because those are the facts that decide which to use when the goal is a small, fast model.

Write $p(\mathbf{y}\mid x) = \prod_t p(y_t\mid x, \mathbf{y}_{<t})$ for the teacher's probability of a response $\mathbf{y} = (y_1, \dots, y_T)$ to a prompt $x$. Write $q_\theta(\mathbf{y}\mid x)$ for the student's, with parameters $\theta$. From here on $\mathbf{y}$ is a sequence of tokens, not a label vector. The chain rule for KL divergence (§0.5) splits a KL between whole sequences into a sum of KLs between next-token distributions,

In plain words, the KL between whole sequences is the per-token KL, summed over positions and averaged over prefixes. The prefixes are drawn from whichever model sits in the first slot. The same identity holds with $p$ and $q_\theta$ swapped, and then the prefixes come from the student.

Word-level KD, also called token-level, applies the distillation loss above at every position of a training corpus, usually at $\tau = 1$. Gemma 2's small models were pretrained this way. Sequence-level KD replaces the corpus with text the teacher wrote and trains on it with the ordinary next-token loss, which is fine-tuning on teacher outputs. On-policy KD lets the student write and has the teacher grade what it wrote. Kim and Rush named the first two. The next three parts take them in turn.

Storing the teacher's top k probabilities per token

Word-level KD needs the teacher's full distribution at every position. One way is to compute it online, by running the teacher alongside the student. A forward pass costs about $2N_T$ FLOPs per token for a teacher of $N_T$ parameters. A training step of a student with $N_S$ parameters costs about $6N_S$, two for the forward pass and four for the backward pass (§9.2). So a teacher three times the student's size doubles the cost of training. The other way is to compute the distributions once and store them, and then the problem is storage.

So labs keep a few hundred entries per token. Gemma 3, for example, samples 256 vocabulary ids per token in proportion to the teacher's probabilities and renormalizes the teacher over them. Suppose you stored $k$ ids per token, the set $K$, with their true log-probabilities $\log p_i$ for $i \in K$ under the teacher's full softmax. What loss should the student minimize? There are two natural choices.

Renormalize the teacher. Set $\tilde p_i = p_i/\sum_{j\in K}p_j$ on $K$ and zero elsewhere. Then minimize $H(\tilde{\mathbf{p}}, \mathbf{q}) = -\sum_{i\in K}\tilde p_i\log q_i$, with the student's full softmax over all $V$ ids.

Keep a tail bucket. Merge all the ids you did not store into one extra outcome, the tail. The teacher's tail mass $p_{\text{tail}} = 1 - \sum_{i\in K}p_i$ is known exactly from the stored log-probabilities. The student's is $q_{\text{tail}} = \sum_{j\notin K}q_j$. Minimize

In plain words, renormalizing teaches the student to be overconfident, putting all its probability on the teacher's top $k$. The tail bucket keeps the total right. The measured cost of renormalizing is real but small. Busbridge et al. found that keeping the top 128 entries and renormalizing raised the student's validation loss by 0.11 nats against the full distribution. Mixing the truncated loss with the ordinary next-token loss ($\lambda_{\text{KD}} = 0.7$) brought it back to within 0.01 nats.

A tail bucket has one numerical trap. If you compute $q_{\text{tail}}$ as $1 - \sum_{i\in K}q_i$, it cancels to exactly zero once the student is confident, and its log is $-\infty$. Compute it instead with a log-sum-exp over the logits you did not store.

import numpy as np

def topk_kd_loss(z, top_idx, top_logp, mode):
    raise NotImplementedError

Storage is not the only practical constraint. Logit distillation compares two distributions token id by token id, so the teacher and the student must share a tokenizer (§2.6). With different vocabularies the per-token losses above are undefined. You then fall back to sequence-level KD on the teacher's text, or to methods that map one vocabulary onto the other.

Sequence-level distillation is fine-tuning on the teacher's text

The sequence-level objective is the cross-entropy between the teacher's and the student's distributions over whole responses,

The sum runs over every possible response, $V^T$ terms for responses of $T$ tokens, far too many to compute. But the expectation on the right says what to do. Sample responses from the teacher, then train the student on them by maximum likelihood. Kim and Rush (2016) trained instead on one response per source sentence, the teacher's beam-search output, which approximates its single most likely response. On their English-German translation data, the teacher's greedy output alone held 1.3% of its probability on average. That is a lot for a distribution over sentences.

Two facts follow, and both come up in interviews. The first is that if you have the teacher's full distributions, you should use them.

In plain words, training on the teacher's sampled tokens and training on its full distributions along the same text aim at the same target. At each position the full distribution gets there with no sampling noise. The cell shows it at one position, with a teacher distribution $(0.6, 0.3, 0.1)$ and a student $(0.5, 0.3, 0.2)$.

The one-token loss has the right mean, 0.937 against 0.938, but a variance of 0.10. The full-distribution loss has no variance at all at this prefix.

The second fact is that sequence-level KD is what most of the field now calls "distillation". Responses written by a stronger model are now the bulk of nearly every open instruction-tuning mix (§14.2). It needs only the teacher's text. So it works through an API, across tokenizers, and with a teacher you cannot run yourself.

On-policy distillation

Both objectives above train the student on prefixes it did not write. At generation time it conditions on its own tokens. After its first mistake it is in a state the training data never showed it. This mismatch between training prefixes and generation prefixes is exposure bias, and §3.5 showed how the errors it causes compound with the length of the output. The fix there was to put the learner's own states into training, and that is what on-policy distillation does. Sample a response from the student. Ask the teacher for its log-probabilities on the same prefixes. Then move the student toward the teacher at the places the student actually goes.

Which divergence fits student samples? The reverse KL,

is an average over the student's outputs, so it can be estimated from exactly the samples we have. Gu et al. (2024) chose it for a second reason. It is mode seeking (§0.5). It punishes the student for putting probability where the teacher puts little. It does not punish the student for ignoring some of what the teacher can say. So a student too small to cover everything the teacher does is not pushed to spread probability over continuations it cannot produce well. A two-token example shows the difference starkly.

The cell finds both optima by searching a grid of $(\theta_1, \theta_2)$.

Now the mechanics. Write $s_t = (x, \mathbf{y}_{<t})$ for the state at position $t$, the prompt and the tokens so far. By the chain rule, the reverse KL over sequences is a sum of per-token terms,

In plain words, minimizing the reverse KL is maximizing a reward $r_t$ earned at every token. The reward is high when the teacher likes the student's token more than the student does. §15.7 trained a language model to maximize a task reward minus a penalty, a coefficient times the KL from a reference model. On-policy distillation is that objective with no task reward, a coefficient of 1, and the teacher in place of the reference model. So the policy-gradient methods of §15.3 apply unchanged.

Two simplifications are standard in practice. The gradient is not propagated through the sampling of earlier tokens, so each state $s_t$ is treated as given (generalized knowledge distillation, or GKD, by Agarwal et al. 2024, states this explicitly). And each token is credited only with its own reward, not with the rewards of the tokens after it. Kevin Lu and colleagues at Thinking Machines (2025) call this a discount factor of zero, and report that larger discounts did not help. What remains is one problem per token, and it has a clean answer.

In plain words, the student needs only the teacher's log-probability of the one token the student sampled. REINFORCE, the score-function estimator of §15.3, estimates the gradient of an expected reward as a sample's reward times the gradient of that sample's log-probability. So a REINFORCE step with reward $-a_y$ moves the student, on average, exactly along the negative gradient of the per-token reverse KL. That is cheap, one number per token from the teacher. It works with a teacher served by any engine that returns log-probabilities. With the teacher's full distribution in hand you can instead compute the per-token KL exactly and differentiate it, with no sampling noise at that position. That is what GKD does.

import numpy as np

def on_policy_kd_step(student_logits, teacher_logits, sampled_tokens):
    raise NotImplementedError

On-policy distillation has three costs. The teacher must run during training, a forward pass over every student sample. The student must generate, which needs an inference engine inside the training loop. That is the same machinery as group relative policy optimization (GRPO), the reinforcement learning method of §18.4, which also samples from the model being trained. And the reverse KL's mode seeking costs diversity, which matters for a model that will be sampled many times. An interviewer who asks "why not always distill on-policy?" wants those three costs.

When a stronger teacher stops helping

A stronger teacher is not always a better one. Hinton et al. already saw that a very small student wants different settings, the intermediate temperatures that pay less attention to the teacher's least informative logits. A related effect has a name, the capacity gap. Busbridge et al. (2025) measured it in a controlled sweep of language-model teachers and students. As the teacher's own loss falls, the student's loss falls too, until the teacher becomes too capable for the student. After that point a better teacher makes a worse student. Gemma 3's report adds a twist. For short training runs a smaller teacher gave the better student, and for long runs the larger teacher won.

Busbridge et al. fit a distillation scaling law, a formula for the student's loss as a function of its size, its training tokens, and the teacher's loss. They drew two conclusions from it.

This is why distillation is everywhere in practice. Labs already have the large model. They produce whole families of small models from it, and they train each small one on far more tokens than its size would justify. That is over-training, trading extra training compute for a cheaper model to serve (§9.4). Gemma 2 trained its 2B and 9B models by word-level KD from a larger teacher, on 2 and 8 trillion tokens, more than 50 times the compute-optimal amount. In its ablation, a 2B model trained on 500 billion tokens averaged 60.3 on three benchmarks from scratch and 67.7 when distilled from a 7B teacher. Earlier, DistilBERT (Sanh et al. 2019) distilled BERT during pretraining into a model 40% smaller and 60% faster that kept 97% of its language-understanding scores.

Pruning, which weights can go

The second way to a smaller model removes parts of the trained one. Which parts can go? §20.3 answered a similar question for rounding, and the same idea works here. A weight matters only through the outputs it helps compute, so a changed layer is good if it gives the same outputs on real inputs. To measure that, run the model on a few hundred sample sequences, called calibration data. Record each layer's inputs. Then measure how much a change moves that layer's outputs.

Take one linear layer $\mathbf{Y} = \mathbf{X}\mathbf{W}$ in the book's layout. The weights are $\mathbf{W} \in \R^{d_{\text{in}}\times d_{\text{out}}}$, and the calibration inputs are $\mathbf{X} \in \R^{n\times d_{\text{in}}}$, one row per token. A pruned layer replaces $\mathbf{W}$ by $\hat{\mathbf{W}}$ and should keep the outputs, so it minimizes

where $\mathbf{w}_j$ is column $j$ of $\mathbf{W}$ and $\lVert\cdot\rVert_F$ is the Frobenius norm, the square root of the sum of squared entries. In plain words, the error splits into one term per output column, and every column shares the same matrix $\mathbf{H}$. It is the Hessian of the error, its matrix of second derivatives. Its diagonal entry $H_{kk} = 2\lVert\mathbf{X}_{:,k}\rVert^2$ measures how large input channel $k$ is on the calibration data. Its off-diagonal entries measure how the input channels move together.

The cost of removing one weight

The cell checks both numbers by least squares.

The general version of this calculation prices both ways of removing a weight $w_q = W_{qj}$ from column $j$. The second way is the update §20.3 used to put a weight on a quantization grid, with the grid value replaced by 0. It removes the weight and then shifts all the others in its column to make up for the loss.

In plain words, removing a weight costs the square of the weight times the energy of its input. If its input is correlated with others, the remaining weights can absorb part of the loss, and the optimal brain surgeon says exactly how much. In the two-input example, $\mathbf{H} = 2n\begin{pmatrix}1 & 0.8\\ 0.8 & 1\end{pmatrix}$ gives $[\mathbf{H}^{-1}]_{11} = 1.389/n$. The formula then gives $0.25/(2 \times 1.389) = 0.09$ per token, the number the cell measured. In general $1/[\mathbf{H}^{-1}]_{qq} = 2\lVert(\mathbf{I} - \boldsymbol\Pi)\mathbf{X}_{:,q}\rVert^2$, where $\boldsymbol\Pi$ projects onto the span of the other input columns (§20.3 proves it). So the cost is $w_q^2$ times the energy of the part of input $q$ that the other inputs cannot reproduce. An interviewer will ask when compensation helps, and part 3 is the answer. It helps only when inputs are correlated.

Scoring a weight by its size times its input's size

The classic method, magnitude pruning, removes the weights with the smallest $|W_{kj}|$. With retraining, Han et al. (2015) used it to cut AlexNet's parameters ninefold. It works for networks whose inputs all have similar scales. Large language models do not. §20.2 showed their outlier features, a few input channels whose activations are often tens of times larger than the rest. A tiny weight on such a channel can carry more of the output than a large weight elsewhere, and magnitude pruning removes it.

The fix is to score each weight by its size times the size of its input,

where $\lVert\mathbf{X}_{:,k}\rVert_2$ is the norm of input channel $k$ over the calibration tokens. Then remove the lowest-scoring fraction within each output column. This method is called Wanda, for pruning by weights and activations (Sun et al. 2023). It needs one forward pass to collect the input norms and no weight updates.

Proof. The inverse of a diagonal matrix is diagonal with reciprocal entries. Substituting into part 2 of the saliency theorem gives $W_{qj}^2/(2/H_{qq}) = \tfrac12 H_{qq}W_{qj}^2$. Since $H_{qq} = 2\lVert\mathbf{X}_{:,q}\rVert^2$, this is $\lVert\mathbf{X}_{:,q}\rVert^2W_{qj}^2$, the square of $S_{qj}$. Squaring preserves the order of nonnegative numbers, so the rankings agree. ∎

So Wanda scores a weight by the cost of removing it without compensation, part 1, which ignores how inputs correlate. Magnitude pruning goes one step further and assumes every input channel has the same norm. One more choice matters in practice. Wanda compares weights within each output column, so every output neuron loses the same fraction of its inputs. Its authors found this consistently better than comparing across the whole layer for LLMs, though not for the image classifiers they tried.

The cell builds a layer with correlated inputs and four outlier channels 20 times larger than the rest. It prunes the layer four ways. Then it refits the surviving weights of each column by least squares. That refit is the exact version of the compensation the optimal brain surgeon performs one weight at a time.

import numpy as np

def wanda_mask(W, X, sparsity):
    raise NotImplementedError

def mask_2of4(scores):
    raise NotImplementedError

Pruning with compensation at scale

Wanda ignores the off-diagonal entries of the Hessian, and with them the compensation that the refit in the cell measured. Exact optimal brain surgeon steps for every weight of a 175-billion-parameter model are out of reach. After each removal the update needs the inverse Hessian of only the weights that remain, and each column removes different weights. So every column would need its own sequence of matrix inverses.

SparseGPT (Frantar & Alistarh 2023) makes it tractable with the trick of GPTQ, the quantizer of §20.3. GPTQ quantizes a layer one input row at a time, in a fixed order shared by every column. After each row it pushes that row's error onto the rows not yet processed, with the optimal brain surgeon update. SparseGPT walks the input rows the same way. In each block of 128 rows it chooses the mask by the compensated cost of part 2, $w_q^2/[\mathbf{H}^{-1}]_{qq}$ up to the factor $\tfrac12$, with $\mathbf{H}^{-1}$ the inverse Hessian of the rows not yet processed. Then it zeroes the masked weights and applies the update to the rows not yet processed. It reuses GPTQ's precomputed factorization of the inverse Hessian, which makes each update cheap. For 2-of-4 sparsity the mask block is four rows. The same sweep can also quantize the surviving weights.

What sparsity buys on a GPU

A model with half its weights set to zero has half the parameters on paper. On a GPU it is usually not faster at all. A dense matrix multiply reads every stored number and multiplies every pair, zeros included. The tensor cores of §11.1, the units that do a GPU's matrix math, have no way to skip an arbitrary zero. A sparse format would have to store the position of every surviving weight, and those indices can cost more than the zeros save. Mishra et al. (2021) put the index overhead of a general compressed format at up to 200% for 8-bit weights.

Unstructured sparsity does pay on CPUs, whose kernels can skip work. SparseGPT measured end-to-end speedups of 1.57, 1.82, and 2.16 times at 40, 50, and 60% sparsity for OPT-2.7B on a CPU inference engine. On GPUs the hardware supports one pattern.

Why is the end-to-end gain so much less than 2? To answer it we need the step-time model of §20.1, which we restate. A decode step for $B$ sequences does two jobs. It does arithmetic, about two FLOPs per weight per sequence plus the attention over each sequence's cached keys and values. And it reads memory, every weight once plus each sequence's cache. Each job runs at the GPU's peak, and the step lasts as long as the slower one,

Here $N_{\text{act}}$ is the number of parameters each token is multiplied by, $L$ the number of layers, $d_{\text{attn}}$ the total width of the attention heads, and $\bar T$ the mean context length. $M_W$ is the bytes of weights and $m_{\text{KV}}$ the bytes of cached keys and values per token. The replica's $n_{\text{gpu}}$ GPUs each do $\pi$ FLOPs per second and read $\beta$ bytes per second. Here $\beta$ is a memory bandwidth, not the weight of the GKD box earlier in this section. The last two terms are communication between GPUs and fixed overhead, which we set to zero. Prefill, the pass that reads the whole prompt at once, multiplies each weight by every prompt token, so it is bound by arithmetic.

A 2-of-4 matrix halves the FLOPs of the linear layers, if the kernel runs at the sparse rate, and cuts their weight bytes to 0.5625. Decode at small batch is bound by bytes. Prefill is bound by FLOPs. And neither attention nor the KV cache gets any cheaper.

The cell works this out for Llama 3.1 8B on an H100 (989 TFLOP/s dense BF16, 3.35 TB/s), with 2-of-4 sparsity on the linear layers of its 32 blocks.

The lesson is general, and it is the answer an interviewer wants. A sparsity pattern speeds up only the term of the step-time model it shrinks, and only in the regime where that term is the larger one.

The widget puts the whole story on one small layer. Choose a pruning method and a sparsity, and watch the output error and the two step times together.

Removing whole neurons, heads, and layers

Removing whole structures leaves smaller dense matrices. They are faster in every regime and need no special kernels. The price is accuracy. A whole neuron or layer carries far more of the function than a scattered half of its weights. So structured pruning is almost always followed by retraining, and the best retraining is distillation from the unpruned model.

Width pruning removes neurons, attention heads, or embedding channels. Recall the feed-forward block of §5.3, $f(\mathbf{X}) = \phi(\mathbf{X}\mathbf{W}_1)\mathbf{W}_2$, with $\mathbf{W}_1 \in \R^{d\times d_{\text{ff}}}$, $\mathbf{W}_2 \in \R^{d_{\text{ff}}\times d}$, and an activation function $\phi$. Neuron $i$ is column $i$ of $\mathbf{W}_1$ together with row $i$ of $\mathbf{W}_2$. Its contribution to the output is the rank-one term $\phi(\mathbf{X}\mathbf{W}_{1,:i})\,\mathbf{W}_{2,i:}$. So removing it changes the output by exactly that term, and a neuron that is rarely active costs little. A gated feed-forward block also removes the matching column of its gate projection.

One recipe scores every structure from its activations on a small calibration set, with no gradients. It is Minitron (Muralidharan et al. 2024), whose full recipe closes this part of the section. It ranks neurons by their activations, averaged over the sequence and combined over the batch with an L2 norm. It ranks heads by the norm of their outputs, and embedding channels by the output of the normalization layer. The authors cite evidence that weight magnitude is a poor guide to importance in LLMs.

Depth pruning removes whole blocks, and two cheap scores are standard. One removes each block in turn and measures the perplexity of what remains. The other is the block influence of ShortGPT (Men et al. 2024). Recall that each block reads the residual stream, the running vector of every token, and adds its output to it (§5.1). Block influence is $1 - \mathbb{E}_t[\cos(\mathbf{x}_t^{(\ell)}, \mathbf{x}_t^{(\ell+1)})]$, the average cosine distance between the stream before block $\ell$ and after it. It is small for blocks that barely change the stream, and one forward pass computes it for every block.

import numpy as np

def prune_neurons(W_up, W_down, X, k):
    keep = np.arange(k)                                 # placeholder: keeps the first k neurons
    return W_up[:, keep], W_down[keep], keep, 1.0

So when does pruning and distilling beat training a small model from scratch? The distillation scaling law answers in the same terms. The pruned student inherits the teacher's weights, a head start no initialization gives. The teacher already exists, so its cost is already paid. When both hold, and the small model does not need more tokens than the threshold for its size, prune and distill. When you must train the large model only to make one small one, or the small model will be trained on many trillions of tokens anyway, training it directly is competitive. Labs do both. Gemma and Qwen distill their small models from their large ones, and NVIDIA's Minitron line prunes and distills.

Distillation makes better drafts

One use of distillation goes the other way. It makes the large model faster without replacing it. Speculative decoding, the subject of §20.5, lets a small draft model propose several tokens, which the large target model checks in one pass. Its speedup is set by how often the target accepts the draft's tokens. That depends on how closely the draft's next-token distribution matches the target's. Training the draft by distillation from the target raises that agreement directly. DistillSpec (Zhou et al. 2024) found that the two choices of this section matter there too, on-policy data generated by the draft and a divergence tuned to the decoding strategy. It reported 10 to 45% speedups over speculative decoding with an off-the-shelf draft.

The real way in PyTorch

torch.nn.functional.kl_div computes the KL divergence from log-probabilities, with the target in the second slot, so the boxed loss is one line. reduction="batchmean" divides by the batch size, which matches the per-example mean of the definition. log_target=True takes the teacher's log-probabilities directly, which avoids a log(0). The on-policy step uses detach() as the stop-gradient. In a real pipeline the teacher's log-probabilities come from an inference engine such as vLLM or SGLang, which return them per token. The student's samples come from the same kind of engine running the current weights.

For 2-of-4 sparsity, PyTorch's torch.sparse.to_sparse_semi_structured converts a weight that already has the pattern into the compressed format. It sends matrix multiplies to the sparse kernels of NVIDIA's cuSPARSELt or CUTLASS libraries. It needs an Ampere or newer GPU, so it cannot run here.

Choosing a method

The table puts the methods side by side in the terms of the step-time model. "Bytes" and "FLOPs" are those of the layers the method touches.

The methods compose. Kim and Rush pruned their distilled student, SparseGPT quantized its sparse model, and Minitron distilled its pruned one. A typical modern small model is a distilled student that is then quantized for serving.

Read more

• Hinton, Vinyals & Dean (2015). Distilling the Knowledge in a Neural Network. Soft targets, the temperature, the $\tau^2$ factor, and the proof that logit matching is the high-temperature limit, in eight readable pages.
• Kim & Rush (2016). Sequence-Level Knowledge Distillation. Names word-level and sequence-level distillation, and shows that training on the teacher's beam output removes the need for beam search.
• Agarwal et al. (2024). On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes. GKD, training on student samples with forward, reverse, or generalized Jensen-Shannon divergences. Read it with Lu et al. (2025), On-Policy Distillation, for the per-token reverse-KL recipe and its compute accounting against reinforcement learning.
• Busbridge et al. (2025). Distillation Scaling Laws. When distillation beats ordinary training, the capacity gap, and careful ablations of temperature, top-k truncation, and KL direction for language models.
• Hassibi & Stork (1993). Second order derivatives for network pruning: Optimal Brain Surgeon, with LeCun, Denker & Solla (1990), Optimal Brain Damage. The two saliencies this section derives, in their original setting.
• Sun et al. (2023). A Simple and Effective Pruning Approach for Large Language Models. Wanda, its connection to the optimal brain surgeon, and measured 2-of-4 speedups. Read it with Frantar & Alistarh (2023), SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot.
• Mishra et al. (2021). Accelerating Sparse Deep Neural Networks. The 2-of-4 format, its metadata, the sparse tensor cores, and a retraining recipe that recovers dense accuracy.
• Muralidharan et al. (2024). Compact Language Models via Pruning and Knowledge Distillation. The Minitron recipe and its best practices, followed by Sreenivas et al. (2024), LLM Pruning and Distillation in Practice: The Minitron Approach.

Watch next

▶

▶

▶

▶