Model Compression and Deployment: Distillation, Pruning, and Serving

A model that reaches its target metric in a notebook still has to meet a serving workload: input sizes, request rates, latency limits, memory, and cost. Compression changes the model; export and serving change how it runs. Each change needs its own quality and performance comparison. This article follows that path and then considers how to detect problems after release.

Knowledge distillation

A teacher supplies training targets for a student, often a smaller model chosen for the deployment budget. For example, probabilities [0.7, 0.2, 0.1] for cat, dog, and truck express a preference among alternatives that a one-hot cat label does not. These scores reflect the teacher’s learned behavior, including its mistakes; they need not measure an intrinsic similarity between classes. Whether they help the student depends on the teacher, data, student capacity, and training procedure.

For classification, let \(z_t,z_s\) be teacher and student logits, \(p_t^T=\operatorname{softmax}(z_t/T)\), and \(p_s^T=\operatorname{softmax}(z_s/T)\). Logits are the scores before softmax. Temperature \(T>0\) controls how spread out the targets are, and \(0\le\alpha\le1\) weights the soft-target term. With the teacher fixed, a common loss is:

\[\mathcal L=\alpha T^2\operatorname{KL}(p_t^T\Vert p_s^T)+(1-\alpha)\operatorname{CE}(z_s,y).\]

KL here sums teacher probabilities times log(teacher probability / student probability); reversing the two distributions changes the objective. The hard-label term CE is cross-entropy against the true class index y. The code averages both terms over examples. It assumes matching class order, finite floating-point logits of shape (batch, classes), and integer labels. Run the Python blocks in order.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels,
                      temperature=4.0, alpha=0.7):
    if not math.isfinite(temperature) or temperature <= 0 or not 0 <= alpha <= 1:
        raise ValueError("Use finite T > 0 and alpha in [0, 1]")
    if (student_logits.ndim != 2 or student_logits.shape != teacher_logits.shape
            or student_logits.size(0) == 0):
        raise ValueError("Matching nonempty (batch, classes) logits required")
    soft = F.kl_div(
        F.log_softmax(student_logits / temperature, dim=-1),
        F.softmax(teacher_logits.detach() / temperature, dim=-1),
        reduction="batchmean") * temperature**2
    hard = F.cross_entropy(student_logits, labels)
    return alpha * soft + (1 - alpha) * hard

student = torch.tensor([[0.0, 0.0, 0.0]], requires_grad=True)
teacher = torch.tensor([[2.0, 1.0, -1.0]], requires_grad=True)
loss = distillation_loss(student, teacher, torch.tensor([0]))
loss.backward()
print("loss", round(loss.item(), 6))
print("student gradient", [round(v, 6) for v in student.grad[0].tolist()])
print("teacher gradient absent", teacher.grad is None)
# loss 0.814555
# student gradient [-0.510466, 0.064662, 0.445804]
# teacher gradient absent True

Before the \(T^2\) multiplier, the soft-loss derivative for one example is \((p_s^T-p_t^T)/T\). At high temperature relative to centered logit differences, the probability difference itself is approximately proportional to \(1/T\), giving the familiar approximate \(1/T^2\) scaling. Multiplying by \(T^2\) compensates for this behavior; it does not make gradients independent of temperature at every \(T\). Tune temperature and loss weight together on validation data.

The negative first gradient component means a gradient-descent step raises the student’s cat logit from its initially equal scores. The teacher receives no gradient. This checks the learning signal on one example; it does not measure student quality. Detaching teacher logits prevents teacher updates through this loss; in a training loop also put the teacher in evaluation mode and compute its outputs under torch.no_grad() to avoid building its graph. The student remains in training mode. Compare the distilled student with the same student architecture trained on labels alone using a stated data and training budget. A successful result for one teacher–student pair does not guarantee a gain for another.

For a generative model, sequence-level distillation uses teacher-generated responses as student training targets. It can work without teacher token probabilities, whereas token-level matching requires compatible output distributions or an explicit mapping. Generated targets can reproduce teacher errors and reduce response diversity. Evaluate held-out tasks and compare against other available training targets. At deployment the teacher is normally removed; student architecture and runtime determine serving cost. Quantization and attention-cache optimizations are covered in LLM Inference Optimization.

Pruning

Unstructured pruning sets selected weights to zero; magnitude is one selection rule. A dense tensor with zeros keeps the same number of stored entries, and an ordinary dense matrix multiply still processes them. A smaller serialized file requires a representation or compressor that exploits the zeros. Runtime gains require suitable sparse kernels and enough useful work to offset indexing and scheduling costs; there is no universal sparsity threshold.

Structured pruning removes entire units such as channels, heads, or layers and rebuilds the affected tensors. This can reduce dense computation, but dimensions, kernel alignment, and other bottlenecks determine latency. A mask alone does not rebuild the network. Removing an intermediate channel also requires removing its input connection in the next layer and updating any associated normalization and branch dependencies.

The small example below removes two hidden units from a Linear–ReLU–Linear network. The retained indices select rows of the first weight matrix and columns of the second. It compares the rebuilt model with the original model whose removed hidden activations are masked to zero. Equality to that masked model does not imply equality to the original unpruned model or preserved task accuracy.

def compact_mlp(first, second, keep):
    if not isinstance(first, nn.Linear) or not isinstance(second, nn.Linear):
        raise TypeError("This example supports two Linear layers")
    if first.out_features != second.in_features:
        raise ValueError("Hidden dimensions must match")
    keep = torch.as_tensor(keep, dtype=torch.long, device=first.weight.device)
    if (keep.ndim != 1 or keep.numel() == 0 or keep.unique().numel() != keep.numel()
            or keep.min() < 0 or keep.max() >= first.out_features):
        raise ValueError("Use distinct, valid hidden-unit indices")
    def linear_like(layer, inputs, outputs):
        return nn.Linear(inputs, outputs, bias=layer.bias is not None,
                         device=layer.weight.device, dtype=layer.weight.dtype)
    left = linear_like(first, first.in_features, keep.numel())
    right = linear_like(second, keep.numel(), second.out_features)
    with torch.no_grad():
        left.weight.copy_(first.weight[keep])
        right.weight.copy_(second.weight[:, keep])
        if first.bias is not None:
            left.bias.copy_(first.bias[keep])
        if second.bias is not None:
            right.bias.copy_(second.bias)
    return nn.Sequential(left, nn.ReLU(), right)

torch.manual_seed(0)
first, second = nn.Linear(3, 4), nn.Linear(4, 2)
x = torch.randn(5, 3)
keep = [0, 2]
small = compact_mlp(first, second, keep).eval()
with torch.no_grad():
    hidden = first(x).relu()
    mask = torch.tensor([1., 0., 1., 0.])
    masked = second(hidden * mask)
    print("matches masked model", torch.allclose(small(x), masked, atol=1e-6))
print("parameters", sum(p.numel() for layer in (first, second) for p in layer.parameters()),
      "->", sum(p.numel() for p in small.parameters()))
# matches masked model True
# parameters 26 -> 14

With biases included, the original layers contain 3×4 + 4 + 4×2 + 2 = 26 parameters; retaining two hidden units leaves 3×2 + 2 + 2×2 + 2 = 14. Selecting which units to remove is a separate problem. Weight norms, activation statistics, and validation-loss changes provide different heuristics. BatchNorm scale magnitudes are not universal importance scores: parameter rescaling and connected layers can change their interpretation. Try a pruning schedule with fine-tuning and compare it with a one-shot cut and a directly trained smaller model when the budget permits. Neither recovery nor a fixed quality cost per removed parameter is guaranteed. Rebuild the optimizer after replacing parameters before continuing training.

Export and numerical parity

A serving artifact includes weights or an executable graph plus its input/output contract, preprocessing, label mapping, and runtime requirements. A Python service can also run PyTorch directly. ONNX export describes supported computations for other runtimes; export itself does not guarantee optimization, support on every device, or unchanged predictions. A runtime such as TensorRT can then compile supported operations for a target configuration.

The following CPU example uses the legacy PyTorch 2.2 exporter, opset 17, ONNX 1.17, and ONNX Runtime 1.20. Install onnx and onnxruntime in the Python environment to run it. It exports a tiny classifier, creates a real runtime session, and checks batches 1 and 4 against PyTorch. Temporary export files are removed when the block finishes. This is a compatibility example, not a recommendation to replace a newer environment with these versions.

import inspect
import tempfile
from pathlib import Path
import numpy as np
import onnx
import onnxruntime as ort

def export_onnx(model, example_input, path):
    model.eval()
    options = {}
    if "dynamo" in inspect.signature(torch.onnx.export).parameters:
        options["dynamo"] = False
    torch.onnx.export(
        model, example_input, str(path),
        input_names=["input"], output_names=["logits"],
        dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
        opset_version=17, **options)

def verify_export(model, session, inputs, rtol=1e-4, atol=1e-6):
    model.eval()
    with torch.no_grad():
        expected = model(inputs).detach().cpu().numpy()
    actual = session.run(["logits"], {"input": inputs.detach().cpu().numpy()})[0]
    if actual.shape != expected.shape:
        raise AssertionError("Output shape mismatch")
    if not np.isfinite(actual).all() or not np.isfinite(expected).all():
        raise AssertionError("Nonfinite output")
    np.testing.assert_allclose(actual, expected, rtol=rtol, atol=atol)
    return float(np.max(np.abs(expected - actual)))

torch.manual_seed(1)
model = nn.Sequential(nn.Linear(3, 4), nn.ReLU(), nn.Linear(4, 2)).eval()
with tempfile.TemporaryDirectory() as directory:
    path = Path(directory) / "model.onnx"
    export_onnx(model, torch.randn(2, 3), path)
    onnx.checker.check_model(onnx.load(str(path)))
    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    for batch in (1, 4):
        error = verify_export(model, session, torch.randn(batch, 3))
        print("batch", batch, "parity", error < 1e-6)
# batch 1 parity True
# batch 4 parity True

Only the batch dimension is dynamic here; the feature width remains three. In newer torch.export-based ONNX export, the dynamo=True path uses dynamic_shapes. Match shape declarations, supported operators, and opset to the chosen exporter and runtime. Declaring an axis dynamic does not make shape-dependent Python logic valid for every size. Unsupported behavior may cause export to fail, and captured control flow can specialize a model to the example path.

Numerical parity on two batches is evidence for those inputs. Test the shapes, branches, dtypes, and extreme values required by the serving contract, including intentional rejection of invalid inputs. Choose tolerances for the output scale and precision, and measure task metrics after conversion: a small logit change near a decision boundary can change a label. Validate on the actual execution provider and target hardware before claiming deployment parity or speed.

Serving

A runtime executes the model; a server handles requests, queues, routing, and model versions. TensorRT is an NVIDIA-oriented inference runtime and compiler, while Triton Inference Server can host models through multiple backends. ONNX Runtime can use different execution providers. Mobile and browser deployment additionally depend on supported operators, download size, memory, and device capabilities. Choose among compatible runtimes by measuring the intended workload rather than assigning one a universal speed ranking.

Dynamic batching collects compatible requests for a short interval and evaluates them together. The extra wait may improve throughput, but it adds to request latency and may bring little benefit when traffic is sparse. Sweep queue delay and maximum batch size under realistic arrival rates. Report throughput alongside latency percentiles, timeouts, and memory; include preprocessing and response handling when measuring the service. LLM continuous batching, discussed in the inference article, also replaces completed sequences between decoding iterations.

Define the request contract before exposing an endpoint: accepted shapes and dtypes, size limits, output meaning, and how invalid requests are rejected. Warm the model before marking an instance ready. Bound queues and use timeouts or admission control so overload does not create an indefinitely growing backlog. These service behaviors need load tests as well as model tests.

Version the model, preprocessing, class or tokenizer mappings, dependencies, and configuration together. First run a candidate on mirrored traffic without letting it affect responses (a shadow deployment), or send it a limited share of real traffic (a canary). Set quality and operational criteria for expansion or rollback. Keep the previous compatible artifact available and test rollback; changing a model version alone may not reverse a feature-schema change.

Training–serving skew

Serving can disagree with training because of resize interpolation, normalization constants, tokenization, channel order, or feature availability. Compare serving preprocessing with the deterministic evaluation path used to assess the trained model. Random training augmentation is intentionally different and should not be required to produce identical tensors.

Share deterministic preprocessing where practical and maintain a golden set of raw inputs with expected tensors. Include missing values, unusual sizes, and representative text or image formats. Compare intermediate tensors and then end-to-end outputs; a single prediction can hide preprocessing differences. Shared code still depends on configuration, library versions, and data access. For time-dependent features, test that historical training features use only information available at the prediction time, with comparable serving freshness. A finite golden set cannot catch every future input or timing bug.

Monitoring

When outcome labels arrive late, distribution and operational checks can raise an alert before labeled evaluation is available. They do not establish that accuracy has fallen. Compare appropriate reference windows and relevant subgroups, accounting for seasonality and sample size.

Monitor missingness, ranges, category frequencies, and prediction distributions. Statistical tests of individual features do not detect every joint-distribution change; with many tests or large samples, even small irrelevant differences can trigger alerts. A changed confidence histogram is a diagnostic signal, not proof of out-of-distribution inputs or a serving bug. Models can remain highly confident on unfamiliar data. Review model-version changes, input processing, and labeled performance to investigate.

Covariate shift changes \(P(X)\) while holding \(P(Y\mid X)\) fixed. It can change average error by moving traffic toward difficult regions; it does not imply that retraining will help. Check support coverage and labeled performance before deciding whether to reweight, gather data, or refit. Concept drift changes \(P(Y\mid X)\). Input statistics alone cannot identify it: updated labels or another valid source of outcome evidence are needed to assess the relationship. Both changes can occur together.

For delayed labels, join outcomes to the prediction-time model version and cohort. Report label coverage and delay as well as task metrics: labels available only for reviewed or completed cases may be a biased sample. Track latency percentiles, error and timeout rates, queue depth, resource use, and cost alongside quality. Alerts should have an owner and a response, such as inspection, rollback, or targeted data collection; an input-shift alert alone is not a reason to retrain automatically.

Retain a permitted, representative sample of requests, predictions, and version metadata for diagnosis, with appropriate access and retention limits. Configuration and checkpoint provenance are covered in Reproducible Training Pipelines.

Before increasing traffic

Check the compressed model’s held-out quality, the exported model’s parity, and the service’s behavior under load separately. Include important subgroups and input sizes, not just an aggregate score. Confirm that the tested artifact and preprocessing are the ones being released, and exercise rollback with their dependencies.

Choose fallback behavior for the application: a previous model, a simpler rule, human review, or an explicit unavailable response. A confidence threshold needs validation for that use; raw confidence alone is not a reliability guarantee. More fallback machinery can itself fail, so test the chosen behavior during timeouts and overload. For prioritizing later improvements, see A Systematic Strategy for Improving Machine Learning Models.

Exercises

1. Distillation temperature. Compute teacher probabilities and entropy at temperatures 1, 3, and 10. Explain what the temperature multiplier compensates for.

logits = torch.tensor([4.0, 2.0, 1.0, 0.5, 0.0])
for T in (1.0, 3.0, 10.0):
    p = (logits / T).softmax(0)
    entropy = -(p * p.log()).sum()
    print("T", T, "probabilities", [round(v, 4) for v in p.tolist()],
          "entropy", round(entropy.item(), 4))
# T 1.0 probabilities [0.8106, 0.1097, 0.0404, 0.0245, 0.0148] entropy 0.6955
# T 3.0 probabilities [0.4071, 0.209, 0.1498, 0.1268, 0.1073] entropy 1.4788
# T 10.0 probabilities [0.2542, 0.2081, 0.1883, 0.1791, 0.1704] entropy 1.5987

2. Zeros versus smaller tensors. Compare dense weight storage and multiplication counts for a 10×10 matrix, the same shape with 90% zeros, and a 5×10 matrix. Count a multiply-add as two FLOPs for one input vector.

W = torch.arange(1., 101.).reshape(10, 10)
masked_W = W.clone()
masked_W.flatten()[:90] = 0
narrow_W = W[:5].clone()
for name, matrix in [("dense", W), ("zeros", masked_W), ("half outputs", narrow_W)]:
    print(name, "entries", matrix.numel(),
          "bytes", matrix.numel() * matrix.element_size(),
          "dense FLOPs", 2 * matrix.numel(),
          "nonzeros", torch.count_nonzero(matrix).item())
# dense entries 100 bytes 400 dense FLOPs 200 nonzeros 100
# zeros entries 100 bytes 400 dense FLOPs 200 nonzeros 10
# half outputs entries 50 bytes 200 dense FLOPs 100 nonzeros 50

3. Training–serving skew. Design tests for preprocessing disagreement, including time-dependent features. What can a golden set establish?

References

• Han, Pool, Tran, and Dally (2015). Learning Both Weights and Connections for Efficient Neural Networks.
• Hinton, Vinyals, and Dean (2015). Distilling the Knowledge in a Neural Network.
• Sculley et al. (2015). Hidden Technical Debt in Machine Learning Systems. NIPS.
• Kim and Rush (2016). Sequence-Level Knowledge Distillation.
• Li, Kadav, Durdanovic, Samet, and Graf (2017). Pruning Filters for Efficient ConvNets.
• Breck, Cai, Nielsen, Salib, and Sculley (2017). The ML Test Score: A Rubric for ML Production Readiness and Technical Debt Reduction. IEEE Big Data.
• Frankle and Carbin (2019). The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks. ICLR.
• Sanh, Debut, Chaumond, and Wolf (2019). DistilBERT, a Distilled Version of BERT: Smaller, Faster, Cheaper and Lighter.
• PyTorch KLDivLoss documentation.
• Triton batching documentation.
• TensorRT sparsity documentation.