Compression & distillation

A model that is accurate and too expensive to serve is not finished. This section covers the three ways of making one smaller — training a small model to imitate a large one, removing weights that contribute little, and storing the weights at lower precision — and what each actually buys on real hardware, which is often less than the numbers suggest. The discipline is in the last step: every one of these trades quality for cost, so you need a measurement that tells you what you gave up before it ships.

Which lever, and in what order#

The three levers do different jobs and they are not interchangeable. Quantization changes how each number is stored and nothing else; the model has the same shape and the same arithmetic, done at lower precision. Pruning deletes weights, either individually or in whole units. Distillation throws the model away and trains a different, smaller one to behave like it. Read the table as a recognition rule, then the sections below measure each claim.

symptom                                      reach for
model will not fit in the memory you have    quantization, then a smaller model
generation is slow, one token at a time      quantization (it is bandwidth bound)
batch throughput is too low                  structured pruning or a smaller model
cost per request on a hosted API is too high distillation onto a small open model
you have a big model and one narrow task     distillation
the model is already small and still slow    none of these; look at the serving stack

Order of effort matters as much as order of effect. Quantization needs no training data and no training run, so it is hours of work. Pruning needs a fine-tuning run to recover what it broke, so it is days. Distillation needs data, a training run, an eval and a decision about licensing, so it is weeks. Start at the top.

Everything below runs against one trained model, so the numbers can be compared with each other. It is a two-hidden-layer network of width 512 on a synthetic tabular problem: 284,674 parameters, test accuracy 0.9467.

import torch, torch.nn as nn
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

X, y = make_classification(n_samples=12000, n_features=40, n_informative=15,
                           n_redundant=10, class_sep=0.8, flip_y=0.03, random_state=0)
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=4000, random_state=0)
scaler = StandardScaler().fit(Xtr)
as_t = lambda a: torch.tensor(scaler.transform(a), dtype=torch.float32)
Xtr, Xte = as_t(Xtr), as_t(Xte)
ytr, yte = torch.tensor(ytr), torch.tensor(yte)

def build(width):
    return nn.Sequential(nn.Linear(40, width), nn.ReLU(),
                         nn.Linear(width, width), nn.ReLU(), nn.Linear(width, 2))

def fit(m, epochs=25, lr=3e-3, seed=0):
    torch.manual_seed(seed)
    opt = torch.optim.AdamW(m.parameters(), lr=lr, weight_decay=1e-2)
    g = torch.Generator().manual_seed(seed)
    for _ in range(epochs):
        m.train()
        perm = torch.randperm(len(Xtr), generator=g)
        for i in range(0, len(Xtr), 128):
            idx = perm[i:i + 128]
            opt.zero_grad()
            nn.functional.cross_entropy(m(Xtr[idx]), ytr[idx]).backward()
            opt.step()
    m.eval()
    return m

def accuracy(m):
    with torch.no_grad():
        return (m(Xte).argmax(1) == yte).float().mean().item()

torch.manual_seed(0)
teacher = fit(build(512))
print(sum(p.numel() for p in teacher.parameters()), round(accuracy(teacher), 4))
# 284674 0.9467

Quantization: what it does to the numbers#

Store weights at lower precision. Usually the best ratio of saving to effort, since it needs no retraining. 8 bit is close to free; 4 bit costs something real and task-dependent.

Concretely, a weight in fp32 is a 32-bit floating point number that can represent a huge range at fine resolution. An 8-bit integer has 256 values. Quantization picks one step size for a block of weights, divides every weight by it, rounds to the nearest integer, and stores the integers plus the step size. Getting a weight back means multiplying the integer by the step size, which is not the number you started with.

def quantize(W, bits):
    """Symmetric per-tensor quantization: integers, step size, and the values you get back."""
    qmax = 2 ** (bits - 1) - 1
    scale = W.abs().max() / qmax
    q = torch.clamp(torch.round(W / scale), -qmax - 1, qmax)
    return q, scale, q * scale

W = teacher[0].weight.data                       # the 512 x 40 first layer
q, scale, back = quantize(W, 8)
print("largest weight", round(W.abs().max().item(), 4), "step size", round(scale.item(), 6))
print("fp32", [round(v, 6) for v in W[0, :5].tolist()])
print("int8", [int(v) for v in q[0, :5].tolist()])
print("back", [round(v, 6) for v in back[0, :5].tolist()])
# largest weight 0.5442 step size 0.004285
# fp32 [-0.038751, 0.016033, -0.121082,  0.002565, -0.158385]
# int8 [-9,         4,       -28,         1,       -37]
# back [-0.038569,  0.017142, -0.119992,  0.004285, -0.158561]

That last observation is the whole story of quantization error. The step size is set by the largest weight in the block, and small weights are the ones it hurts. Measured across the matrix, and then across the whole model.

import copy

for bits in (8, 4):
    _, _, back = quantize(W, bits)
    err = (back - W).abs()
    rel = (back - W).pow(2).mean().sqrt() / W.pow(2).mean().sqrt()
    print(bits, round(err.max().item(), 5), round(err.mean().item(), 5), round(rel.item(), 4))
# 8  0.00214  0.00107  0.0088      error is 0.88% of the typical weight size
# 4  0.03887  0.01949  0.1605      error is 16% of it

for bits in (8, 4):
    m = copy.deepcopy(teacher)
    for layer in m:
        if isinstance(layer, nn.Linear):
            layer.weight.data = quantize(layer.weight.data, bits)[2]
    print(bits, round(accuracy(m), 4))
# 8  0.9467      identical to fp32
# 4  0.9457      one tenth of a point

So on this model int4 is nearly free, which is not what happens to large language models, and the difference is worth understanding rather than memorising. The damage from low-bit quantization is caused by outliers: one weight far larger than the rest drags the step size up and every ordinary weight then rounds to almost nothing. Large transformers have exactly that structure in some layers and activations. Here is the mechanism, forced by hand.

def rel_rms(a, b):
    return ((a - b).pow(2).mean().sqrt() / b.pow(2).mean().sqrt()).item()

Wo = W.clone()
Wo[0, 0] = W.abs().max() * 20                    # one weight twenty times bigger than any other
ordinary = torch.ones_like(W, dtype=torch.bool)
ordinary[0, 0] = False

print(round(rel_rms(quantize(W, 4)[2][ordinary], W[ordinary]), 4))
print(round(rel_rms(quantize(Wo, 4)[2][ordinary], Wo[ordinary]), 4))
s = Wo.abs().amax(dim=1, keepdim=True) / 7       # one step size per output unit
per_channel = torch.clamp(torch.round(Wo / s), -8, 7) * s
print(round(rel_rms(per_channel[ordinary], Wo[ordinary]), 4))
# 0.1605   int4, no outlier, one step size for the tensor
# 1.0000   int4, one outlier, one step size for the tensor: every ordinary weight rounds to zero
# 0.1033   int4, one outlier, one step size per output unit: the damage is contained

This is why real quantization schemes never use one step size for a whole model. Per-channel scaling gives each output unit its own step size; group-wise schemes give every block of 64 or 128 weights its own. That metadata costs a little: at group size 64 with a 16-bit scale per group, you are paying an extra 0.25 bits per weight, so int4 is really 4.25 bits. It is worth it.

What you actually buy is memory and bandwidth, and the arithmetic is a multiplication. A 7-billion-parameter model needs 4 bytes per weight in fp32, so 28 GB; 14 GB in fp16; 7 GB in int8; 3.5 GB in int4. That decides which accelerator it fits on, and it is also most of the latency story for generation, because producing one token reads every weight once and does very little arithmetic with each. Generation is bandwidth bound, so roughly halving the bytes roughly halves the time per token. Prompt processing and large-batch serving are compute bound instead, and there the gain depends on whether your kernels have low-precision arithmetic at all, or dequantize back to fp16 and do the maths as before. Measure rather than assume; a quantized model that dequantizes in the kernel saves memory and not time.

Pruning, and why structured wins in practice#

Remove weights that contribute little. Unstructured pruning zeroes individual weights, which reduces size and rarely speeds anything up because the hardware still processes the zeros. Structured pruning removes whole units or layers, which is a real speedup and a larger quality cost.

allw = torch.cat([l.weight.data.abs().flatten() for l in teacher if isinstance(l, nn.Linear)])

def unstructured(frac):
    """Zero the smallest frac of weights. Shapes are unchanged."""
    m = copy.deepcopy(teacher)
    thr = torch.quantile(allw, frac)
    for layer in m:
        if isinstance(layer, nn.Linear):
            layer.weight.data = layer.weight.data * (layer.weight.data.abs() > thr)
    return m

def structured(keep):
    """Keep the most useful hidden units and rebuild a genuinely smaller network."""
    l0, l2, l4 = teacher[0], teacher[2], teacher[4]
    imp1 = l0.weight.data.norm(dim=1) * l2.weight.data.norm(dim=0)
    k1 = torch.topk(imp1, keep).indices.sort().values
    imp2 = l2.weight.data[:, k1].norm(dim=1) * l4.weight.data.norm(dim=0)
    k2 = torch.topk(imp2, keep).indices.sort().values
    m = build(keep)
    m[0].weight.data, m[0].bias.data = l0.weight.data[k1], l0.bias.data[k1]
    m[2].weight.data, m[2].bias.data = l2.weight.data[k2][:, k1], l2.bias.data[k2]
    m[4].weight.data, m[4].bias.data = l4.weight.data[:, k2], l4.bias.data
    m.eval()
    return m

for frac in (0.5, 0.8, 0.9, 0.95, 0.99):
    print(frac, round(accuracy(unstructured(frac)), 4))
for keep in (256, 128, 64, 32):
    print(keep, sum(p.numel() for p in structured(keep).parameters()),
          round(accuracy(structured(keep)), 4))

unstructured, no fine-tuning        structured, no fine-tuning
50% zeroed     0.9455               512 -> 256   76,802 params   0.9240
80% zeroed     0.9240               512 -> 128   22,018 params   0.8662
90% zeroed     0.8508               512 ->  64    6,914 params   0.7552
95% zeroed     0.7235               512 ->  32    2,434 params   0.5500
99% zeroed     0.5065

baseline: 284,674 params, 0.9467

Now the part that decides which one you use. Time all four on the same batch, single threaded.

dense, width 512          0.884 ms
95% of weights zeroed     0.923 ms     slower than dense
structured to width 128   0.146 ms     6.1x faster
structured to width 64    0.083 ms    10.7x faster

The 95%-sparse model is a fraction of a millisecond slower than the dense one it came from. It has the same tensor shapes, so the same multiply-accumulate operations happen; the zeros are multiplied like any other number, and the masking added a little work. This is the single most useful fact about pruning, and it is why the literature on unstructured sparsity and the practice of serving models diverge. Unstructured sparsity pays off only where something downstream exploits it: a sparse storage format that shrinks the file, or hardware with a sparsity mode, such as the 2:4 pattern some accelerators support, which requires exactly two of every four adjacent weights to be zero rather than any 50% you like.

Structured pruning is slower to recover from and is the one that pays. Fine-tuning after the cut is not optional: it is where the accuracy comes back.

                    after pruning   after 8 epochs   width trained
                                    of fine-tuning   from scratch
512 -> 128              0.8662          0.9450           0.9482
512 ->  64              0.7552          0.9370           0.9337

That last column is the comparison people skip, and on this problem it kills the case for pruning at width 128. Pruning is worth the trouble when the large model encodes something the small one cannot learn on its own: not enough labelled data to train the small model directly, a very expensive pretraining run you are not going to repeat, or a compression ratio steep enough that training from scratch does not converge. Otherwise the honest answer is to train the smaller model.

Distillation#

Train a small model to imitate a large one's outputs. The small model learns from the large model's full probability distribution rather than only hard labels, which carries more information and is why it can approach the teacher's quality at a fraction of the size.

Distillation earns its keep in one specific situation, and the measurement is in the next section of this topic. The short version: on a six-class problem with 8,000 labels available, a small network trained on those labels scored 0.8220 and the same network distilled from a large teacher scored 0.8294, which is three quarters of a point for a lot of machinery. With only 500 labels and 6,000 unlabelled inputs, the same small network scored 0.4854 on the labels alone, 0.7949 on the teacher's hard guesses for the unlabelled pool, and 0.8205 on the teacher's full distributions, which is 0.15 points off having all 8,000 real labels. Distillation converts a teacher plus unlabelled data into labelled data. Where you already have the labels, it buys very little.

Measure what you lost#

                  size    latency   your eval
teacher            26 GB    840 ms     0.91
distilled           1.4 GB   95 ms     0.88
distilled + int8    0.7 GB   61 ms     0.87

Whether losing 0.04 is acceptable is a product decision,
and it should be made with the numbers in front of you.

Published compression benchmarks are aggregates and hide uneven degradation: some capabilities survive and others fall off sharply. Always measure on your own task.

Break the eval down rather than reporting one number. The pattern that shows up repeatedly is that easy, frequent cases survive compression and rare or multi-step cases do not, so an aggregate that is 96% easy cases moves by half a point while the capability your product depends on has halved. Slice by input length, by task type, by language, and by whatever your traffic is heavy in. Compare the compressed and original models on the same inputs and look at where they disagree, not only at the totals.

When compression is not the answer#

• The cost is in the prompt, not the model. If two thousand tokens of retrieved context go in on every request, halving the model does not halve the bill. Shorten the input first, and cache what repeats.
• Traffic is low. A compression project is weeks of engineering; below a few million requests a month a larger instance is cheaper than your time, and the comparison is worth doing explicitly in money.
• The bottleneck is elsewhere. If p50 latency is 900 ms and 700 of it is a database call, a 6x faster model buys 170 ms.
• Quality has no headroom. If you are already at the edge of acceptable, there is nothing to trade. Fix quality first; compression only ever spends it.
• A smaller off-the-shelf model already does the job. Try it before building anything. It costs an afternoon and it frequently ends the project.

Distillation is training a small model to behave like a large one, either by matching its full output distribution or, for language models, by generating outputs from the large model on your real inputs, filtering them, and fine-tuning the small one on what survives. This section does both, then insists on the comparison that decides whether it was worth doing: against a smaller model with a better prompt, and against a mid-sized model, either of which may reach the same quality with no training at all. Skipping the filter is the other common way these projects fail, because the student copies the teacher's mistakes as faithfully as its successes.

Soft targets#

A teacher model's full output distribution carries more information than the correct label alone. For an image of a cat, the teacher might put a small probability on 'dog' and almost none on 'truck'. Training a student to match that distribution teaches it how classes relate. Raising the softmax temperature spreads the distribution out so this information is visible.

import torch.nn.functional as F

def distill_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.5):
    soft = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),
        F.softmax(teacher_logits / T, dim=-1),
        reduction="batchmean",
    ) * T * T                                   # keeps gradient size comparable across T
    hard = F.cross_entropy(student_logits, labels)
    return alpha * soft + (1 - alpha) * hard

The T squared factor is the detail that gets dropped and then quietly breaks the recipe. Dividing logits by T shrinks the gradient of the soft term by roughly a factor of T squared, so without it the soft loss fades out as you raise the temperature and alpha stops meaning what you think it means. Put it in.

How much extra information is really in those soft targets is measurable, and the answer decides whether the technique is worth setting up. Train a width-512 teacher on a six-class problem, then look at what its output distributions contain at different temperatures.

import torch, torch.nn as nn, numpy as np
from sklearn.datasets import make_classification
from sklearn.preprocessing import StandardScaler

X, y = make_classification(n_samples=20000, n_features=40, n_informative=20, n_classes=6,
                           n_clusters_per_class=2, class_sep=1.0, flip_y=0.02, random_state=1)
scaler = StandardScaler().fit(X[:8000])
as_t = lambda a: torch.tensor(scaler.transform(a), dtype=torch.float32)
Xtr, ytr = as_t(X[:8000]), torch.tensor(y[:8000])       # the teacher's training data
Xun = as_t(X[8000:14000])                               # unlabelled inputs nobody trained on
Xte, yte = as_t(X[14000:18000]), torch.tensor(y[14000:18000])

def build(width):
    return nn.Sequential(nn.Linear(40, width), nn.ReLU(),
                         nn.Linear(width, width), nn.ReLU(), nn.Linear(width, 6))

def new(width, seed=0):
    """Fixed initialisation, so two runs differ only in what we changed."""
    torch.manual_seed(seed)
    return build(width)

def fit(m, Xa, ya, epochs=25, lr=3e-3, seed=0, teacher=None, T=4.0, alpha=0.7):
    torch.manual_seed(seed)
    opt = torch.optim.AdamW(m.parameters(), lr=lr, weight_decay=1e-2)
    g = torch.Generator().manual_seed(seed)
    if teacher is not None:
        with torch.no_grad():
            tlogits = teacher(Xa)
    for _ in range(epochs):
        m.train()
        perm = torch.randperm(len(Xa), generator=g)
        for i in range(0, len(Xa), 128):
            idx = perm[i:i + 128]
            opt.zero_grad()
            out = m(Xa[idx])
            loss = nn.functional.cross_entropy(out, ya[idx])
            if teacher is not None:
                loss = alpha * distill_loss_only(out, tlogits[idx], T) + (1 - alpha) * loss
            loss.backward()
            opt.step()
    m.eval()
    return m

def distill_loss_only(student_logits, teacher_logits, T):
    return F.kl_div(F.log_softmax(student_logits / T, dim=-1),
                    F.softmax(teacher_logits / T, dim=-1),
                    reduction="batchmean") * T * T

def accuracy(m):
    with torch.no_grad():
        return (m(Xte).argmax(1) == yte).float().mean().item()

teacher = fit(new(512), Xtr, ytr, epochs=30)
print(round(accuracy(teacher), 4))                      # 0.8727

with torch.no_grad():
    tlogits = teacher(Xte)
for T in (1.0, 2.0, 4.0):
    p = torch.softmax(tlogits / T, dim=-1)
    H = -(p * p.clamp_min(1e-12).log()).sum(-1).mean()
    print(T, round(H.item(), 4), round(p.max(-1).values.mean().item(), 4))
# 1.0  0.0993  0.9601      almost all the mass is on one class
# 2.0  0.2228  0.9165
# 4.0  0.5304  0.8150      now there is a distribution to learn from
# for reference, the most entropy six classes can carry is ln(6) = 1.7918

This is also the fastest way to know whether distillation is worth trying at all. A binary classifier's teacher has almost nothing to say beyond its top choice, and distilling one usually buys nothing measurable. The more classes, and the more genuinely confusable they are, the more there is in the distribution.

Does it help when you already have the labels#

Mostly not, and it is better to find that out in ten lines than after a fortnight. Train a small student on the same 8,000 labelled examples the teacher saw, with and without the soft targets, three seeds each.

scratch = [accuracy(fit(new(48, s), Xtr, ytr, seed=s)) for s in range(3)]
distil = [accuracy(fit(new(48, s), Xtr, ytr, seed=s, teacher=teacher)) for s in range(3)]
print([round(v, 4) for v in scratch], round(float(np.mean(scratch)), 4))
print([round(v, 4) for v in distil], round(float(np.mean(distil)), 4))
# [0.819,  0.8255, 0.8215]  0.8220     student trained on the labels
# [0.8255, 0.8298, 0.833 ]  0.8294     same student, plus the teacher

The case it is actually for#

Distillation is a way of turning a teacher plus unlabelled inputs into labelled data. Give the student 500 labels and 6,000 inputs nobody has labelled, and let the teacher fill them in.

Xs, ys = Xtr[:500], ytr[:500]
with torch.no_grad():
    yun = teacher(Xun).argmax(1)                # the teacher's guesses, right 86.9% of the time
Xboth, yboth = torch.cat([Xs, Xun]), torch.cat([ys, yun])

few = [accuracy(fit(new(48, s), Xs, ys, epochs=60, seed=s)) for s in range(3)]
hard = [accuracy(fit(new(48, s), Xboth, yboth, seed=s)) for s in range(3)]
soft = [accuracy(fit(new(48, s), Xboth, yboth, seed=s, teacher=teacher)) for s in range(3)]
for r in (few, hard, soft):
    print(round(float(np.mean(r)), 4))

500 real labels only                         0.4854
500 + 6,000 of the teacher's hard guesses    0.7949
500 + 6,000 of the teacher's distributions   0.8205
the same student with all 8,000 real labels  0.8220

Two things in that table are the reason this technique is used at all. The teacher's labels were 86.9% correct and the student trained on them reached 79.5%, which is far better than the 48.5% it could manage from the labels it had. And the soft targets were worth another 2.6 points over the same guesses as hard labels, because a distribution tells the student which mistakes are near misses. The student ends up almost as good as if somebody had labelled 8,000 examples by hand, which is exactly the trade an LLM distillation project is making: the teacher's time is cheap, and human labelling is not.

Distilling language models with generated data#

For LLMs, the common form is simpler: have a large model produce high-quality outputs for many inputs from your task, filter them, and fine-tune a small model on the result. The small model learns to do your task the way the large one does, at a fraction of the serving cost.

• Use real inputs from your task, with the variety of production traffic.
• Filter the teacher's outputs with your eval checks. Distilling mistakes teaches mistakes.
• Hold out a test set the teacher never saw, and compare student and teacher on it.

Choosing the student#

The student's size is the decision that determines whether the project is worth doing, and it has to be swept rather than guessed. Running the same comparison at six widths on the problem above shows something the single-width experiment could not.

width   parameters   from labels   distilled   difference
   12          726      0.6563       0.6498      -0.0065
   24        1,734      0.7688       0.7617      -0.0071
   48        4,614      0.8220       0.8294      +0.0074
   96       13,830      0.8307       0.8465      +0.0158
  192       46,086      0.8512       0.8475      -0.0037
  384      165,894      0.8642       0.8575      -0.0067

teacher, width 512, 286,726 parameters:            0.8727

There is a band where it pays and it is narrower than the technique's reputation suggests. A student too small cannot represent what the teacher is telling it, so the soft targets are a distraction from the labels. A student large enough to learn the task well does better from the real labels than from an imperfect teacher, because the teacher is only 87% right and the labels are 100% right. In between, where the student has the capacity to use extra information but not enough to find it alone, distillation is worth one and a half points. Sweep the widths, and note that the sign of the effect changes, so a single run at an arbitrary width can tell you that distillation helps, or that it hurts, depending on which width you picked.

For language models the same decision is made between released models rather than by choosing a width, and the comparison to run is fixed: the large model prompted, the small model prompted, the small model distilled, and a mid-sized model prompted. The fourth is the one people skip and it frequently wins, because it needs no training run and no maintenance. A distilled model is a training pipeline you now own: it has to be rebuilt when the task changes, re-evaluated when the teacher is upgraded, and hosted by somebody.

Also decide what the student is allowed to be worse at. Distillation transfers the teacher's behaviour on the distribution you sampled, so the student inherits the common cases and not the general ability that handled the rare ones. Route what the student has not seen to the teacher rather than pretending the gap does not exist: a confidence threshold on the student, or a simple check on input length or type, recovers most of the loss for a small fraction of the traffic and therefore a small fraction of the cost.

Check the terms#

Some model providers' terms restrict using their outputs to train other models, and open model licences vary. Check before building a product on distilled weights.

The comparison that decides it#

On your task eval:              quality   cost per 1k requests
large model, prompted             0.91          high
small model, prompted             0.78          low
small model, distilled            0.88          low      <- worth it only if this beats the row above clearly