The Efficiency Toolbox: Quantization, Pruning & Distillation

"My teacher had four hundred layers and a god complex. I have eighteen layers and a deadline. Somehow I score within a point of her, and I run on a phone she would not deign to charge."

A Distilled Student Network With No Time to Waste

You have spent Part III making models more accurate. This section is the first where we deliberately make a model worse, in a controlled way, in exchange for speed and size. That trade is the heart of deployment. The model that comes out of Chapter 21 at the top of the validation curve is carrying weights at 32-bit floating-point precision, dense connectivity it does not need, and a parameter count chosen for trainability rather than inference cost. Each of the three techniques here attacks one of those excesses, as the illustration below sketches. We take them in turn, then combine them, and throughout we measure rather than assume, because the only honest claim about a compressed model is one backed by a number on the target metric.

1. Why Over-Parameterization Leaves Room Beginner

A modern vision network has far more parameters than the information content of its task requires. A ResNet-50 from Chapter 20 has roughly 25 million parameters and reaches about 76 percent top-1 accuracy on ImageNet; a MobileNetV2 with under 4 million parameters reaches about 72 percent. The larger model is not six times more accurate; it is a few points more accurate. Most of the extra capacity is there to make optimization easy, not because the final function needs it. This is the slack that compression exploits. After training, the redundancy can be squeezed out: weights can be stored coarsely, many can be removed entirely, and the learned function can be re-expressed in a much smaller network.

It helps to separate the two costs we are trying to reduce. The first is memory: how many bytes the weights occupy, which sets the download size, the on-device storage, and the memory bandwidth burned moving weights to the compute units. The second is compute: how many arithmetic operations a forward pass performs, usually counted in multiply-accumulates or floating-point operations (FLOPs), which together with the hardware's throughput sets latency. Quantization attacks both at once. Pruning attacks both if it removes whole structures, or only memory if it leaves scattered zeros that the hardware still has to multiply. Distillation attacks both by changing the architecture entirely. Figure 28.1.1 places the three techniques on these two axes.

2. Quantization: Computing in Eight Bits Intermediate

Quantization maps a continuous range of float values onto a small set of integers. The standard scheme is affine (asymmetric) quantization to 8-bit integers. Given a tensor of floats with observed range $[\beta_{\min}, \beta_{\max}]$, we pick a positive scale $s$ and an integer zero-point $z$ so that a float value $x$ maps to an integer $q$ and back:

The scale $s = (\beta_{\max} - \beta_{\min}) / (2^b - 1)$ for $b$ bits spreads the float range across the $2^b$ integer levels (256 levels for int8), and the zero-point $z$ is the integer that the real value zero maps to, which lets the scheme represent asymmetric ranges exactly at zero (important because padding and ReLU produce many exact zeros). The reconstructed value $\hat{x}$ differs from $x$ by at most half a step, $s/2$; this rounding error is the entire accuracy cost of quantization, and keeping $s$ small (a tight range) keeps it small.

Why does throwing away 24 of the 32 bits barely hurt? A float32 weight can take any of roughly four billion distinct values, but the network never needed that resolution. Snapping each weight to its nearest of only 256 evenly spaced levels moves it by at most half a step. That perturbation is smaller than the noise the weight already survived during training, so summed over a layer the rounding errors partly cancel and the output barely moves. The model was carrying 32 bits of precision to make optimization smooth, not because the learned function needed it. The code below implements the round trip so you can see the error directly.

import torch

def quantize_affine(x, num_bits=8):
    """Affine int8 quantization of a float tensor; returns ints, scale, zero-point."""
    qmin, qmax = 0, 2 ** num_bits - 1
    beta_min, beta_max = x.min().item(), x.max().item()
    scale = (beta_max - beta_min) / (qmax - qmin)        # float range per integer step
    zero_point = round(qmin - beta_min / scale)          # integer that maps to real 0.0
    zero_point = max(qmin, min(qmax, zero_point))         # clamp into the integer range
    q = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax).to(torch.int32)
    return q, scale, zero_point

def dequantize_affine(q, scale, zero_point):
    """Map integers back to floats: x_hat = scale * (q - zero_point)."""
    return scale * (q.float() - zero_point)

w = torch.randn(1000) * 0.1                  # a typical weight tensor, small magnitude
q, s, z = quantize_affine(w)
w_hat = dequantize_affine(q, s, z)
err = (w - w_hat).abs().max().item()
print(f"scale={s:.6f}  zero_point={z}  max abs error={err:.6f}")
print(f"theoretical bound (scale/2) = {s/2:.6f}")
# scale=0.002545  zero_point=131  max abs error=0.001268
# theoretical bound (scale/2) = 0.001273

The reason this is fast, not just small, is that the actual matrix multiplications can run in integer arithmetic. With both weights and activations quantized, a convolution or linear layer multiplies int8 values and accumulates into int32, and modern CPUs and GPUs run int8 multiply-accumulate at two to four times the throughput of float32, with a quarter of the memory traffic. The scale and zero-point bookkeeping is applied once per tensor (or per channel) to rescale the int32 accumulator back to the next layer's input range.

2.1 Post-Training Quantization Versus Quantization-Aware Training

There are two ways to get a quantized model, and they trade engineering effort against accuracy. Post-training quantization (PTQ) takes an already-trained float model and quantizes it directly, using a small unlabeled calibration set to observe the activation ranges $[\beta_{\min}, \beta_{\max}]$ at each layer. It is fast (minutes, no training) and needs no labels, and for many convolutional networks it costs well under a point of accuracy. Quantization-aware training (QAT) goes further: it inserts simulated quantization (the round-trip above) into the forward pass during a short fine-tune, so the network learns weights that are robust to the rounding. QAT is more work and needs labeled data and a training loop, but it recovers most of the accuracy that PTQ leaves on the table, which matters for sensitive models like detectors and for aggressive sub-8-bit precision.

QAT has one trainability problem to solve first. The key trick that solves it is the straight-through estimator: the non-differentiable rounding is treated as the identity during the backward pass, so gradients flow through it. This trick is necessary because the rounding function is flat between every pair of integer levels, so its true derivative is zero almost everywhere; back-propagating that honest gradient would deliver exactly zero to every weight and the fine-tune could never learn, so the estimator substitutes a derivative of one and lets the forward-pass rounding error guide the update instead. The code below applies PTQ to a torchvision model.

import torch
from torchvision.models import resnet18, ResNet18_Weights

# 1. A trained float model, set to eval mode (quantization fuses BN into conv).
model_fp32 = resnet18(weights=ResNet18_Weights.DEFAULT).eval()

# 2. Configure post-training static quantization for the x86 backend.
from torch.ao.quantization import get_default_qconfig, QConfigMapping
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx

qconfig = get_default_qconfig("x86")              # per-channel weights, per-tensor activations
qconfig_mapping = QConfigMapping().set_global(qconfig)
example = torch.randn(1, 3, 224, 224)

prepared = prepare_fx(model_fp32, qconfig_mapping, example)   # inserts observers

# 3. Calibrate: run a few unlabeled batches so observers learn activation ranges.
with torch.no_grad():
    for _ in range(32):                            # ~32 representative batches is plenty
        prepared(torch.randn(8, 3, 224, 224))

# 4. Convert observers into real int8 quantized ops.
model_int8 = convert_fx(prepared)

def size_mb(m):
    torch.save(m.state_dict(), "tmp.pt")
    import os; mb = os.path.getsize("tmp.pt") / 1e6; os.remove("tmp.pt"); return mb

print(f"fp32 size = {size_mb(model_fp32):.1f} MB")   # fp32 size = 46.8 MB
print(f"int8 size = {size_mb(model_int8):.1f} MB")   # int8 size = 11.9 MB

Notice the calibration step. Post-training quantization needs to see real (or realistic) inputs so its observers can record the activation ranges; quantizing with ranges that do not match deployment is a common and silent accuracy killer. We will return to exactly this point in Section 28.5, because the calibration distribution is the first thing that drifts in production.

from ultralytics import YOLO

model = YOLO("yolo11n.pt")
# int8=True triggers calibration on the given dataset and writes a quantized engine.
model.export(format="engine", int8=True, data="coco128.yaml")
# -> yolo11n.engine, an int8 TensorRT engine ready for the GPU runtime of Section 28.2

3. Pruning: Removing What Does Not Earn Its Place Intermediate

Pruning sets some weights to zero permanently. The simplest and most effective criterion is magnitude pruning: weights with the smallest absolute value contribute least to the output, so remove the bottom $p$ percent by magnitude. The subtlety is in what shape the removed weights take, and it determines whether you save compute or only memory.

Unstructured pruning zeroes individual weights anywhere in a tensor, producing a sparse matrix. It can reach very high sparsity (90 percent and beyond) with modest accuracy loss, but the zeros are scattered, so a dense matrix multiply still touches them; you save storage (a sparse format) but not wall-clock time unless the hardware has dedicated sparse-matrix support. Structured pruning removes whole units, an entire output channel of a convolution, an entire attention head, an entire neuron, so the resulting tensor is simply smaller and dense. Structured pruning saves real compute on ordinary hardware because the network genuinely has fewer channels to convolve, which is why it is the form that matters most for the latency budgets of this chapter. The trade is that it is coarser and tends to cost more accuracy per parameter removed. Figure 28.1.2 contrasts the two on a single weight matrix.

The other half of pruning is recovery. Removing weights damages the network, but a short fine-tune (often called the prune-then-finetune or iterative-magnitude-pruning loop) lets the surviving weights compensate, recovering most of the lost accuracy. The lottery-ticket hypothesis of Frankle and Carbin gave this a striking interpretation: a dense network contains a sparse subnetwork that could have been trained to the same accuracy on its own, and pruning-and-finetuning is one way to find it. The code below applies structured channel pruning to a single convolution and shows the recovery pattern in miniature.

import torch
import torch.nn.utils.prune as prune
from torchvision.models import resnet18, ResNet18_Weights

model = resnet18(weights=ResNet18_Weights.DEFAULT)
conv = model.layer4[1].conv2          # a 512-channel 3x3 convolution

# Structured pruning: remove 30% of OUTPUT channels (dim=0) by L2 norm.
# This is the form that actually shrinks compute, not just storage.
prune.ln_structured(conv, name="weight", amount=0.30, n=2, dim=0)

# Count channels whose entire weight is now zero.
mask = conv.weight_mask                       # prune stores a 0/1 mask
dead = (mask.sum(dim=(1, 2, 3)) == 0).sum().item()
print(f"channels pruned: {dead} of {mask.shape[0]}")   # channels pruned: 153 of 512

# After a short fine-tune you would call prune.remove to bake the mask in:
prune.remove(conv, "weight")                  # makes the pruning permanent
print("weight is now a plain, sparse-by-channel tensor again")

4. Distillation: Learning From a Larger Teacher Intermediate

Distillation trains a small student network to reproduce the outputs of a large, accurate teacher. The insight, due to Hinton and colleagues, is that the teacher's full output distribution carries more information than the hard label alone. When a good classifier sees an image of a husky, it does not output a one-hot vector; it assigns most probability to husky but meaningful residual probability to wolf and malamute and almost none to teapot. That pattern of relative probabilities, sometimes called dark knowledge, tells the student which classes are confusable, a far richer training signal than the single correct answer (see the illustration below). We expose it by softening the teacher's logits with a temperature $T$ in the softmax:

A temperature $T > 1$ flattens the distribution, amplifying the small probabilities that hold the inter-class structure. The student is trained on a weighted sum of two losses: a distillation loss (the Kullback-Leibler divergence between the student's and teacher's softened distributions) and the ordinary cross-entropy against the true labels. A mixing weight (called $\alpha$ in the code below) sets how much the student should trust the teacher's soft guidance versus the hard ground-truth label, with a value near 0.7 leaning on the teacher while still anchoring to the correct answer. The temperature-scaled term is multiplied by $T^2$ to keep its gradient magnitude comparable to the hard-label term. The same softmax and cross-entropy you met in Chapter 18 reappear here, now comparing two networks instead of a network and a label. The code below implements the loss.

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
    """Combine soft-target KL (from the teacher) with hard-label cross-entropy."""
    # Soft targets: both distributions softened by temperature T.
    soft_teacher = F.softmax(teacher_logits / T, dim=1)
    soft_student = F.log_softmax(student_logits / T, dim=1)
    # KL divergence, scaled by T^2 so its gradient stays comparable to the hard loss.
    kd = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (T * T)
    # Standard supervised loss against the ground-truth labels.
    ce = F.cross_entropy(student_logits, labels)
    return alpha * kd + (1.0 - alpha) * ce

# Toy shapes: batch of 16, 1000 classes.
student_logits = torch.randn(16, 1000, requires_grad=True)
teacher_logits = torch.randn(16, 1000)             # teacher in eval mode, no grad
labels = torch.randint(0, 1000, (16,))
loss = distillation_loss(student_logits, teacher_logits, labels)
print(f"distillation loss = {loss.item():.4f}")   # distillation loss = 9.7xxx

Distillation is the most architecturally flexible of the three techniques because the student can be any network you like. It is also how the famous compact models are made: DistilBERT in language, and in vision the small members of the DeiT and EfficientNet families that ship in production. It composes cleanly with the other two: distill first to get a small accurate student, then prune and quantize that student. We will see in Chapter 33 that distillation also drives the step-reduction methods that turn fifty-step diffusion samplers into one-or-two-step generators, the same idea applied to the number of inference iterations rather than the number of parameters.

5. Stacking the Techniques Advanced

In practice you rarely use one technique alone. The canonical deployment recipe applies all three in the order that respects their dependencies: distill into the target architecture first (so pruning and quantization act on the model you will actually ship), then structurally prune with a recovery fine-tune (so quantization calibrates the already-smaller network), then quantize last (because quantization is the final numeric step before the runtime). The three-word recipe is worth memorizing in exactly this order, distill, prune, quantize, because the order is the logic: each step should act on the model the next step will compress, so you choose the architecture before you thin it, and thin it before you round it. Table 28.1.1 shows representative cumulative effects on an ImageNet classifier; the exact numbers depend on the model and data, but the shape of the trade is stable across vision tasks.

Read the table as a path down the accuracy-versus-cost curve. The big size and latency wins come early (distillation changes the architecture); the late stages add multiplicative savings for a fraction of a point each. The final model is about 38 times smaller and roughly 11 times faster than the teacher, at a 5-point accuracy cost, and whether that trade is acceptable is a product question, not a research one. As Table 28.1.1 makes clear, you do not have to take the whole path; you stop at the first row that meets your budget. The histogram-and-statistics view of distributions from Chapter 2 is worth recalling here, because choosing a quantization range is exactly choosing where to clip a distribution's tails.

6. Summary and the Road to a Runtime

We have three composable tools. Quantization lowers numeric precision to int8 (or below), cutting memory by four and unlocking fast integer math, available as cheap post-training quantization or higher-fidelity quantization-aware training. Pruning removes parameters, with structured pruning the variant that actually reduces compute on ordinary hardware. Distillation trains a small student to inherit a large teacher's accuracy through softened soft targets. Stacked in the order distill, prune, quantize, they routinely deliver order-of-magnitude savings for a single-digit accuracy cost. But a compressed PyTorch module is still a Python object running eagerly; to actually realize the speed it promises, it must be exported to a compiled runtime built for the target hardware. That is the subject of Section 28.2, where ONNX, TensorRT, and OpenVINO turn the shrunk graph into a deployable engine. Before moving on, put all three compression tools together in the Hands-On Lab at the end of this section, which builds a small compression studio that shrinks a pretrained classifier and reports the accuracy-size-latency trade for each technique.

pip install "torch>=2.2" "torchvision>=0.17"

import torch, torch.nn as nn, torchvision as tv
from torchvision.transforms import v2

device = "cuda" if torch.cuda.is_available() else "cpu"
tf = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True),
                 v2.Normalize([0.49, 0.48, 0.45], [0.25, 0.24, 0.26])])
train = tv.datasets.CIFAR10("./data", train=True,  download=True, transform=tf)
test  = tv.datasets.CIFAR10("./data", train=False, download=True, transform=tf)
train_dl = torch.utils.data.DataLoader(train, batch_size=128, shuffle=True, num_workers=2)
test_dl  = torch.utils.data.DataLoader(test,  batch_size=256, num_workers=2)

teacher = tv.models.resnet18(weights="IMAGENET1K_V1")
teacher.fc = nn.Linear(teacher.fc.in_features, 10)
# TODO: write train(model, loader, epochs) and evaluate(model, loader) -> accuracy.
# Hint: standard cross-entropy loop; evaluate runs under torch.no_grad() on CPU for fair latency later.
teacher = teacher.to(device)
# TODO: train(teacher, train_dl, epochs=2)

import os, time, copy, statistics

def size_mb(model):
    torch.save(model.state_dict(), "tmp.pt")
    mb = os.path.getsize("tmp.pt") / 1e6
    os.remove("tmp.pt"); return mb

def cpu_latency_ms(model, runs=50):
    model = model.to("cpu").eval()
    x = torch.randn(1, 3, 32, 32)
    with torch.no_grad():
        for _ in range(5): model(x)            # warm up
        # TODO: time `runs` forward passes and return the MEDIAN in milliseconds.
        # Hint: collect time.perf_counter() deltas into a list, use statistics.median.

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
    soft = F.kl_div(F.log_softmax(student_logits / T, dim=1),
                    F.softmax(teacher_logits / T, dim=1),
                    reduction="batchmean") * (T * T)
    hard = F.cross_entropy(student_logits, labels)
    return alpha * soft + (1 - alpha) * hard

student = tv.models.mobilenet_v2(weights=None, num_classes=10).to(device)
teacher.eval()
# TODO: train `student` for ~3 epochs. For each batch, get teacher_logits under
#       torch.no_grad(), then backprop distillation_loss(student(x), teacher_logits, y).

import torch.nn.utils.prune as prune

pruned = copy.deepcopy(student)
for module in pruned.modules():
    if isinstance(module, nn.Conv2d) and module.out_channels > 8:
        # TODO: prune 30% of output channels by L2 norm, then make it permanent.
        # Hint: prune.ln_structured(module, name="weight", amount=0.3, n=2, dim=0)
        #       followed by prune.remove(module, "weight").
        pass
# TODO: fine-tune `pruned` for ~2 epochs to recover accuracy.

from torch.ao.quantization import quantize_dynamic

q_student = copy.deepcopy(student).to("cpu").eval()
# TODO: quantize q_student with quantize_dynamic, targeting {nn.Linear, nn.Conv2d},
#       dtype=torch.qint8. Confirm size_mb(q_student) dropped versus the float student.

rows = []
for name, model in [("Float32 teacher", teacher),
                    ("Distilled student", student),
                    ("+ Pruned (30%)", pruned),
                    ("+ Int8 quantized", q_student)]:
    # TODO: append (name, evaluate(model, test_dl), size_mb(model), cpu_latency_ms(model)).
    pass

print(f"{'Variant':22} {'Acc%':>6} {'MB':>7} {'ms':>7}")
for n, acc, mb, ms in rows:
    print(f"{n:22} {acc*100:6.1f} {mb:7.1f} {ms:7.2f}")

Variant                  Acc%      MB      ms
Float32 teacher          78.4    44.8   18.30
Distilled student        75.1     8.8    6.10
+ Pruned (30%)           73.9     8.7    5.95
+ Int8 quantized         74.8     2.4    4.40

import os, time, copy, statistics
import torch, torch.nn as nn, torch.nn.functional as F
import torchvision as tv
import torch.nn.utils.prune as prune
from torchvision.transforms import v2
from torch.ao.quantization import quantize_dynamic

device = "cuda" if torch.cuda.is_available() else "cpu"

# --- Step 1: data and teacher ---
tf = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True),
                 v2.Normalize([0.49, 0.48, 0.45], [0.25, 0.24, 0.26])])
train = tv.datasets.CIFAR10("./data", train=True,  download=True, transform=tf)
test  = tv.datasets.CIFAR10("./data", train=False, download=True, transform=tf)
train_dl = torch.utils.data.DataLoader(train, batch_size=128, shuffle=True, num_workers=2)
test_dl  = torch.utils.data.DataLoader(test,  batch_size=256, num_workers=2)

def train_model(model, loader, epochs=2, loss_fn=None):
    model.train().to(device)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    for _ in range(epochs):
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            opt.zero_grad()
            out = model(x)
            loss = loss_fn(out, y) if loss_fn else F.cross_entropy(out, y)
            loss.backward(); opt.step()
    return model

def evaluate(model, loader):
    model.eval(); dev = next(model.parameters()).device
    correct = total = 0
    with torch.no_grad():
        for x, y in loader:
            pred = model(x.to(dev)).argmax(1).cpu()
            correct += (pred == y).sum().item(); total += y.numel()
    return correct / total

teacher = tv.models.resnet18(weights="IMAGENET1K_V1")
teacher.fc = nn.Linear(teacher.fc.in_features, 10)
train_model(teacher, train_dl, epochs=2)

# --- Step 2: measurement harness ---
def size_mb(model):
    torch.save(model.state_dict(), "tmp.pt")
    mb = os.path.getsize("tmp.pt") / 1e6
    os.remove("tmp.pt"); return mb

def cpu_latency_ms(model, runs=50):
    model = model.to("cpu").eval()
    x = torch.randn(1, 3, 32, 32)
    times = []
    with torch.no_grad():
        for _ in range(5): model(x)
        for _ in range(runs):
            t0 = time.perf_counter(); model(x)
            times.append((time.perf_counter() - t0) * 1000)
    return statistics.median(times)

# --- Step 3: distillation ---
def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
    soft = F.kl_div(F.log_softmax(student_logits / T, dim=1),
                    F.softmax(teacher_logits / T, dim=1),
                    reduction="batchmean") * (T * T)
    hard = F.cross_entropy(student_logits, labels)
    return alpha * soft + (1 - alpha) * hard

student = tv.models.mobilenet_v2(weights=None, num_classes=10).to(device)
teacher.eval()
opt = torch.optim.Adam(student.parameters(), lr=1e-3)
for _ in range(3):
    for x, y in train_dl:
        x, y = x.to(device), y.to(device)
        with torch.no_grad():
            t_logits = teacher(x)
        opt.zero_grad()
        loss = distillation_loss(student(x), t_logits, y, T=4.0, alpha=0.7)
        loss.backward(); opt.step()

# --- Step 4: structured pruning + recovery ---
pruned = copy.deepcopy(student)
for module in pruned.modules():
    if isinstance(module, nn.Conv2d) and module.out_channels > 8:
        prune.ln_structured(module, name="weight", amount=0.3, n=2, dim=0)
        prune.remove(module, "weight")
train_model(pruned, train_dl, epochs=2)

# --- Step 5: dynamic int8 quantization ---
q_student = copy.deepcopy(student).to("cpu").eval()
q_student = quantize_dynamic(q_student, {nn.Linear, nn.Conv2d}, dtype=torch.qint8)

# --- Step 6: comparison table ---
rows = []
for name, model in [("Float32 teacher", teacher),
                    ("Distilled student", student),
                    ("+ Pruned (30%)", pruned),
                    ("+ Int8 quantized", q_student)]:
    rows.append((name, evaluate(model, test_dl), size_mb(model), cpu_latency_ms(model)))

print(f"{'Variant':22} {'Acc%':>6} {'MB':>7} {'ms':>7}")
for n, acc, mb, ms in rows:
    print(f"{n:22} {acc*100:6.1f} {mb:7.1f} {ms:7.2f}")