Model Compression and Knowledge Distillation

Model compression is the general problem of producing a smaller, faster, or more memory-efficient model that approximates the behavior of a larger reference model, subject to an accuracy budget. "Smaller" can mean fewer parameters, fewer floating-point operations (FLOPs), lower memory bandwidth, or lower latency on specific hardware — these are related but distinct objectives, and a compression method optimized for one does not automatically help with another.

Pruning removes elements of a trained network guided by a saliency or importance score. Unstructured pruning zeroes out individual weights, producing a sparse weight matrix; it can achieve very high compression ratios but the resulting sparsity pattern is irregular and often does not speed up computation on standard hardware without specialized sparse-matrix kernels. Structured pruning removes entire rows, channels, filters, or attention heads, keeping the remaining computation dense; it yields smaller and more consistent speedups on ordinary hardware at the cost of typically lower achievable compression for the same accuracy loss. The simplest and most widely used saliency criterion is weight magnitude: weights close to zero are assumed to contribute little to the function the network computes, an assumption justified empirically and, for linear layers, connected to the sensitivity of the loss to weight perturbation via a first- or second-order Taylor expansion.

Knowledge distillation (Hinton, Vinyals, and Dean, 2015) trains a compact student network $S$ to reproduce the output distribution of a larger, already-trained teacher network $T$, in addition to (or instead of) fitting ground-truth labels directly. The key technical device is temperature-scaled softmax. For logits $z = (z_1, \dots, z_K)$ over $K$ classes, the standard softmax is

Distillation instead uses a temperature $T \ge 1$:

As $T \to \infty$, $q_i(T) \to 1/K$ (uniform); at $T = 1$ it recovers ordinary softmax. Raising $T$ softens the distribution, amplifying the relative sizes of small probabilities assigned to non-target classes. Hinton et al. call this extra information — the teacher's implicit ranking of wrong answers — dark knowledge: a teacher that outputs 0.9 for "dog", 0.09 for "wolf", and 0.01 for "car" is saying something informative (dogs resemble wolves more than cars) that a one-hot label "dog" throws away entirely. This is the conceptual core of why distillation can outperform training the small student from scratch on the same labeled data: soft targets carry more bits of information per example than a single hard label.

Magnitude pruning: the mechanism

Magnitude pruning operates on an already-trained (or partially trained) network. The procedure alternates between removing low-magnitude weights and fine-tuning the remaining weights to recover accuracy, because pruning in one shot at high sparsity typically damages accuracy irrecoverably. ``` ALGORITHM: Iterative Magnitude Pruning INPUT: trained network with weights W, target sparsity s_final, number of pruning rounds R, fine-tune epochs per round E OUTPUT: pruned network achieving sparsity s_final mask ← all-ones tensor, same shape as W # 1 = kept, 0 = pruned for r in 1 .. R: s_r ← s_final (r / R) # e.g. linear sparsity schedule for each layer ℓ with weights W_ℓ: threshold_ℓ ← the s_r-quantile of |W_ℓ[mask_ℓ == 1]| mask_ℓ ← mask_ℓ AND (|W_ℓ| > threshold_ℓ) W ← W ⊙ mask # zero out pruned weights for e in 1 .. E: forward pass, compute loss on training data backward pass, compute gradients W ← W − η · gradient # standard SGD/Adam step W ← W ⊙ mask # re-zero pruned weights every step return W, mask ``` Two details matter for correctness. First, the mask must be re-applied after every gradient step*, not only after each pruning round — otherwise gradient noise reintroduces small nonzero values at pruned positions and the sparsity budget silently erodes. Second, the per-layer quantile threshold (rather than a single global threshold) prevents one layer with naturally larger weight magnitudes from monopolizing the pruning budget while another layer is pruned into uselessness; global magnitude pruning is also used in practice but requires normalizing weight scales across layers to be fair.

Structured pruning replaces the per-weight mask with a per-channel (or per-filter/per-head) importance score, most commonly the $L_1$ or $L_2$ norm of all weights feeding into that channel, or a Taylor-expansion estimate of the loss increase from removing it:

where $a_c$ is the channel's activation — a first-order approximation of the loss change from zeroing that channel, derived by treating channel removal as setting $a_c \to 0$ and Taylor-expanding $\mathcal{L}$ around the current activation.

Knowledge distillation: the mechanism

``` ALGORITHM: Knowledge Distillation Training INPUT: trained teacher T (frozen), student S (to be trained), labeled dataset {(x_n, y_n)}, temperature τ, mixing weight α OUTPUT: trained student S for each minibatch {(x, y)}: z_T ← T(x) # teacher logits, no gradient (teacher frozen) z_S ← S(x) # student logits q_T ← softmax(z_T / τ) # softened teacher distribution q_S ← softmax(z_S / τ) # softened student distribution (same τ!) p_S ← softmax(z_S) # ordinary student distribution (τ = 1) L_soft ← CrossEntropy(q_T, q_S) · τ² # distillation loss, τ² correction L_hard ← CrossEntropy(y, p_S) # standard supervised loss L ← α · L_soft + (1 − α) · L_hard backward pass on L w.r.t. student parameters only update student parameters (teacher stays frozen) return S ``` Two non-obvious details deserve explanation. Why multiply by $\tau^2$? The gradient of the softened cross-entropy with respect to a student logit $z_{S,i}$ scales as $1/\tau$ (differentiating $\exp(z_i/\tau)$ introduces a factor of $1/\tau$), so as $\tau$ grows the soft-target gradient magnitude shrinks relative to the hard-label gradient. Multiplying $L_{\text{soft}}$ by $\tau^2$ exactly cancels this scaling in the high-temperature limit — Hinton et al. derive this by expanding the softmax gradient to first order and showing the $\tau^2$ factor keeps the relative contribution of $L_{\text{soft}}$ and $L_{\text{hard}}$ comparable across different choices of $\tau$, so $\alpha$ can be tuned somewhat independently of $\tau$.

Why use the same $\tau$ for teacher and student softmax in $L_{\text{soft}}$? Both distributions must be compared on the same scale; comparing a softened teacher to an unsoftened student would conflate "matching the distribution shape" with "matching sharpness," corrupting the signal that dark knowledge is meant to provide. At inference time, the temperature is set back to $T=1$ (ordinary softmax) since softening was only a training-time device to expose informative gradients — it is not part of the deployed student's forward pass. ```{svg} ```

Pruning — time and space. Computing the magnitude of all weights and selecting a quantile threshold costs $O(P)$ per layer where $P$ is the parameter count, dominated by a partial sort or a selection algorithm ($O(P)$ average case via quickselect, $O(P \log P)$ worst case with a full sort). Across $R$ pruning rounds with $E$ fine-tuning epochs each, total training cost is $O(R \cdot E \cdot (\text{one training epoch}))$ — asymptotically the same order as ordinary training, since pruning itself is cheap relative to gradient computation. The payoff is at inference time: an unstructured sparse network with sparsity $s$ has $(1-s)$ fraction of nonzero weights, so a matrix-vector product can in principle be done in $O((1-s) \cdot P)$ time with a proper sparse format (e.g., CSR), versus $O(P)$ dense — but this speedup is realized only with hardware or libraries that exploit sparsity efficiently; naively storing a sparse tensor as a dense array with zeros gives no runtime benefit at all, only smaller compressed storage on disk. Structured pruning of a channel reduces the actual tensor dimensions, so a network with $c$ of $C$ channels pruned gets a real $O(c/C)$ FLOP reduction on ordinary dense hardware with no special kernel required.

Distillation — time and space. Training cost per step is $O(\text{teacher forward}) + O(\text{student forward+backward})$; since the teacher requires no backward pass (it is frozen), this is cheaper than jointly training two networks of the same sizes, but it is still more expensive per step than training the student alone on hard labels, because of the extra teacher forward pass. If the teacher's outputs are precomputed once and cached (feasible for classification with a fixed dataset), training cost drops to essentially the same as training the student alone, at the cost of $O(N \cdot K)$ storage for $N$ examples and $K$ classes (or full logits for more advanced feature-distillation variants, which cost more). At inference time, only the student is deployed, so cost is purely a function of the student's architecture — the compression benefit is total and permanent.

Why does magnitude pruning preserve accuracy (correctness argument)? Consider a linear layer $y = Wx$ and perturb one weight $w_{ij} \to w_{ij} + \delta$. To first order, the change in the layer's output is $\delta \cdot x_j$, and propagating through a locally smooth loss $\mathcal{L}$, the change in loss is approximately $\delta \cdot \frac{\partial \mathcal{L}}{\partial w_{ij}}$. Near a trained (locally optimal) point, gradients are small, so the linear term is often dominated by the second-order term $\frac{1}{2}\delta^2 H_{ij,ij}$ where $H$ is the Hessian (this is the basis of Optimal Brain Damage/Optimal Brain Surgeon pruning criteria). Magnitude pruning implicitly assumes weights with small $|w_{ij}|$ contribute a small $\delta = -w_{ij}$ change (i.e., setting them to exactly zero), which is a heuristic proxy for low saliency, not a guarantee — it is provably correct only under the added assumption that the Hessian is roughly uniform across weights (e.g., approximately diagonal with similar curvature), which does not always hold. This is why fine-tuning after each pruning step is essential: it lets remaining weights absorb the small function change induced by the removed weights, restoring accuracy that a pure "prune once" approach would lose.

Why does distillation work (the information argument)? A one-hot label carries $\log_2 K$ bits of information at most, entirely determined by the class index. A softened teacher output distribution over $K$ classes carries continuous-valued information about relative similarities between classes, which — assuming the teacher's confidences are well-calibrated relative to true class similarity — acts as a regularizer that discourages the student from being overconfident and guides its decision boundaries toward the teacher's, which are typically smoother and more generalizable because the teacher was trained with far more capacity and (often) more data or ensembling. Formally, minimizing $L_{\text{soft}}$ is equivalent to minimizing the KL divergence $D_{KL}(q_T \| q_S)$ up to a constant (since $\text{CE}(q_T, q_S) = H(q_T) + D_{KL}(q_T \| q_S)$ and $H(q_T)$ does not depend on the student), so distillation exactly performs distribution matching rather than merely label matching — a strictly richer training signal whenever $q_T$ is not one-hot.

Tradeoffs, stated honestly. Distillation does not guarantee the student surpasses training from scratch on hard labels; if the student's capacity is too small to represent the teacher's function even approximately, forcing it to match soft targets can hurt compared to hard-label training, because the mismatch introduces an additional approximation error term. Pruning at very high sparsity (>90%) on some architectures causes a sharp accuracy cliff rather than a graceful degradation, and the sparsity level at which this cliff occurs is architecture- and task-dependent, discoverable only empirically. Unstructured pruning's storage savings do not translate into speed savings without specialized sparse kernels, a common practical trap.

Consider a 3-class classification problem (dog, wolf, car) and a single test image of a wolf. The teacher network, a large ResNet, produces logits $z_T = (2.0, 1.5, -3.0)$ for (dog, wolf, car). The student is a small CNN with logits $z_S = (1.0, 0.2, -0.5)$ before training.

Step 1 — ordinary softmax ($T=1$) on the teacher, for comparison:

Already fairly informative, but let's see what temperature reveals.

Step 2 — soften with $T = 4$:

Compare the ratio of "wolf" to "car" mass at $T=1$: $0.373 / 0.011 \approx 34$. At $T=4$: $0.358/0.238 \approx 1.5$. Softening compresses the ratio between confident and less-confident classes, which is precisely what exposes the relative structure (dog and wolf are much closer to each other than either is to car) instead of letting the dominant class swamp the signal.

Step 3 — student's softened distribution at the same $T=4$:

Step 4 — soft loss. Cross-entropy $\text{CE}(q_T, q_S) = -\sum_i q_{T,i}\log q_{S,i}$:

So $L_{\text{soft}} = 1.093$, scaled by $T^2 = 16$ for the gradient-magnitude correction: contribution to total loss is $16 \times 1.093 = 17.5$ before mixing weight $\alpha$ is applied (in practice $\alpha \cdot L_{\text{soft}}$ with $\alpha$ typically 0.1–0.9, tuned per task).

Step 5 — hard loss. True label is "wolf" (one-hot $y = (0,1,0)$). Ordinary softmax of student logits ($T=1$): $p_S = \text{softmax}(1.0, 0.2, -0.5) = (0.554, 0.249, 0.135)$ (rounding), so $L_{\text{hard}} = -\log(0.249) = 1.390$.

Step 6 — total loss with $\alpha = 0.7$:

This is backpropagated only through the student's parameters — the teacher's logits $z_T$ are treated as fixed constants (no gradient flows into the teacher). Notice qualitatively what happened: even though the student's hard-label loss already correctly ranks wolf above car, the soft-target loss additionally pushes the student toward the teacher's belief that dog is nearly as plausible as wolf and that car is a distant third but not negligible — exactly the "dark knowledge" that a one-hot label could never express.

Model compression sits at the intersection of decades of work on generalization theory, neural network regularization, and — for distillation specifically — a small but highly influential 2015 paper.

The Softmax Function is the mathematical hinge connecting these ideas: it is what gets temperature-scaled in distillation, and it is the same function whose Boltzmann-distribution origins Hinton had studied since his work on Boltzmann machines decades earlier. Pruning's Hessian-based saliency criteria connect directly to second-order optimization ideas explored throughout the history of Neural Network Training and regularization.

On-device assistants (voice recognition, keyboard prediction, camera scene detection) run compressed models because cloud round-trips are too slow or violate privacy constraints; distillation and pruning are core to the model pipelines behind mobile speech and vision features. Large language model deployment increasingly relies on distilling huge foundation models into smaller, faster variants for latency-sensitive applications like real-time chat or search ranking, where a student with a fraction of the parameters must approximate the teacher's behavior within a strict latency budget. Recommendation systems at web scale distill large ranking models into lighter models that can score millions of candidates within milliseconds. Edge and embedded ML (wearables, IoT sensors, autonomous drones) depend on structured pruning to fit models within kilobytes of memory. Compression is also the natural bridge to two further stages of the deployment pipeline: Quantization of Neural Networks, which reduces numeric precision after pruning/distillation shrink the model's structure, and Neural Architecture Search, which can be framed as searching for architectures that are compression-friendly or efficient by design, rather than compressing after the fact.

Ask the tutor: How would the distillation loss and pruning strategy change if the student and teacher had different output vocabularies, as in cross-lingual or cross-architecture distillation?