Model Compression Quantization Pruning and Knowledge Distillation

# Model Compression: Quantization, Pruning, and Knowledge Distillation

## 1. Introduction & Motivation

Model compression reduces the size, memory footprint, and inference latency of trained neural networks while preserving as much accuracy as possible. As models grow into the billions of parameters, compression is often the difference between a model that ships on a phone or edge device and one that only runs in a data center.

Motivation: enables:
1. Edge Deployment: Run models on phones, IoT devices, and embedded hardware with limited memory/compute
2. Latency Reduction: Smaller/faster models serve requests quicker, critical for real-time applications
3. Cost Reduction: Lower compute per inference reduces cloud serving costs at scale
4. Energy Efficiency: Compressed models consume less power, important for battery-constrained and sustainability-conscious deployment
5. Bandwidth Savings: Smaller models are cheaper to distribute via over-the-air updates

Applications: on-device speech recognition, mobile vision (photo classification, AR filters), compressed LLMs for local inference, and efficient serving of recommendation systems at massive query volume.

## 2. Core Concepts & Theory

### Quantization

Reduce numerical precision of weights/activations from 32-bit floating point to lower-bit representations (16-bit, 8-bit, or even 4-bit/binary), shrinking memory and enabling faster integer arithmetic on supporting hardware.

$$q = ext{round}\left(\frac{x}{s} ight) + z$$

where s is the scale factor and z the zero-point, mapping the continuous range of x to a discrete integer grid.

### Pruning

Remove weights, neurons, or entire structural components (channels, attention heads, layers) that contribute little to model output, based on magnitude, gradient-based sensitivity, or learned importance scores.

Unstructured pruning: Remove individual weights (creates sparse matrices); high compression ratio but requires specialized sparse hardware/kernels for speedup.

Structured pruning: Remove entire channels, filters, or heads; lower compression ratio but yields dense matrices that run fast on standard hardware without special support.

### Knowledge Distillation

Train a small "student" model to mimic the behavior of a large "teacher" model, using the teacher's soft output probabilities (not just hard labels) as a richer training signal.

$$\mathcal{L}_{ ext{KD}} = \alpha \cdot \mathcal{L}_{ ext{CE}}(y, p_{ ext{student}}) + (1-\alpha) \cdot au^2 \cdot \mathcal{L}_{ ext{CE}}(p_{ ext{teacher}}^ au, p_{ ext{student}}^ au)$$

where au is a temperature that softens both distributions, exposing "dark knowledge" — relative similarities between incorrect classes that hard labels discard.

### Low-Rank Factorization

Decompose weight matrices into products of smaller matrices, exploiting the empirical observation that trained weight matrices often have low effective rank:

$$W \approx U V^ op, \quad W \in \mathbb{R}^{m imes n},\ U \in \mathbb{R}^{m imes r},\ V \in \mathbb{R}^{n imes r},\ r \ll \min(m,n)$$

### Neural Architecture Efficiency by Design

Beyond post-hoc compression, some architectures are designed from scratch to be efficient: depthwise-separable convolutions (MobileNet), grouped queries in attention (GQA), and mixture-of-experts routing (activating only a subset of parameters per input).

## 3. Mathematical Formulation

### Uniform Affine Quantization

Given a floating-point tensor x with range [\alpha, \beta] and target bit-width b:

$$s = \frac{\beta - \alpha}{2^b - 1}, \qquad z = ext{round}\left(\frac{-\alpha}{s} ight)$$

$$x_q = ext{clip}\left( ext{round}\left(\frac{x}{s} ight) + z,\ 0,\ 2^b - 1 ight)$$

Dequantization for inference approximation: \hat{x} = s(x_q - z). Quantization error is bounded by |x - \hat{x}| \leq s/2.

### Magnitude-Based Pruning Criterion

The simplest pruning criterion ranks weights by absolute magnitude and removes the smallest fraction:

$$ ext{mask}_i = \mathbb{1}\left[|w_i| > heta ight]$$

where heta is chosen so a target sparsity ratio p (e.g., 90% of weights zeroed) is achieved.

### The Lottery Ticket Hypothesis

Frankle & Carbin's hypothesis: dense, randomly initialized networks contain sparse subnetworks ("winning tickets") that, when trained in isolation from the same initialization, match the accuracy of the full network:

$$f(x; m \odot heta_0) \approx f(x; heta^*)$$

where m is a binary mask found via iterative magnitude pruning, heta_0 the original initialization, and heta^* the fully trained dense weights.

### Distillation Loss Gradient Insight

The soft-target gradient with respect to student logits z_s (for temperature au) is proportional to:

$$\frac{\partial \mathcal{L}_{ ext{soft}}}{\partial z_{s,i}} \propto \frac{1}{ au}\left(p_{s,i}^ au - p_{t,i}^ au ight)$$

As au o \infty, this approaches matching logits directly (up to a constant shift), explaining why distillation transfers relative class-similarity structure, not just the argmax prediction.

## 4. Advanced Theory & Extensions

### Quantization-Aware Training (QAT)

Rather than quantizing a fully trained model post-hoc (post-training quantization, PTQ), QAT simulates quantization noise during training using the straight-through estimator (STE): the forward pass uses quantized values, but gradients flow through as if quantization were the identity function, allowing the model to adapt its weights to tolerate quantization error.

### Mixed-Precision Quantization

Not all layers are equally sensitive to precision loss. Mixed-precision schemes assign higher bit-widths to sensitive layers (e.g., first/last layers, attention scores) and lower bit-widths to robust layers, often guided by Hessian-based sensitivity analysis.

### Structured Sparsity Patterns (N:M Sparsity)

Modern accelerators (e.g., NVIDIA sparse tensor cores) support "2:4 sparsity" — exactly 2 nonzero values in every contiguous block of 4 — striking a middle ground between unstructured pruning (best accuracy, no hardware speedup) and coarse structured pruning (hardware speedup, worse accuracy).

### Self-Distillation & Online Distillation

Self-distillation trains a model to match its own earlier checkpoint or an ensemble of its own predictions across training, improving generalization even without a separate larger teacher. Online distillation trains teacher and student simultaneously, useful when no pretrained teacher is available.

### Distilling Reasoning Chains

For large language models, distillation increasingly transfers not just final-answer probabilities but full chain-of-thought reasoning traces, teaching smaller models to reproduce step-by-step reasoning patterns rather than only output distributions.

## 5. Computational Considerations

### Hardware-Aware Quantization

INT8 speedups depend on hardware support for integer matrix multiplication; not all accelerators accelerate 4-bit or lower uniformly, so compression ratio on paper does not always translate to proportional latency reduction — profiling on target hardware is essential.

### Sparse Kernel Support

Unstructured pruning achieves high theoretical compression, but without a sparse-matrix-multiply kernel tuned for the specific sparsity pattern and hardware, no wall-clock speedup materializes — this is the central practical obstacle to unstructured pruning adoption.

### Calibration Data for Post-Training Quantization

PTQ requires a small calibration dataset (hundreds of samples) to estimate activation ranges for scale/zero-point computation; poor calibration data (unrepresentative of deployment distribution) causes accuracy degradation.

### Distillation Compute Overhead

Distillation requires running the (expensive) teacher forward pass during student training, at minimum for generating soft targets — for very large teachers, pre-computing and caching teacher outputs offline can amortize this cost across multiple student training runs.

## 6. Practical Implementation Strategies

### Compression Pipeline Ordering

A common effective pipeline: (1) knowledge distillation to a smaller architecture, (2) structured pruning of the distilled model, (3) quantization-aware fine-tuning of the pruned model — each stage compounds compression while allowing recovery fine-tuning between stages.

### Iterative Magnitude Pruning with Fine-Tuning

Rather than pruning to target sparsity in one shot, iterative pruning removes a small fraction of weights, fine-tunes to recover accuracy, then repeats — consistently outperforming one-shot pruning at high sparsity ratios.

### Framework Choices

PyTorch Quantization (torch.quantization / torchao): Native PTQ and QAT support with INT8 backends.

Hugging Face Optimum + bitsandbytes: 8-bit and 4-bit quantization for transformer/LLM inference with minimal code changes.

NVIDIA TensorRT: Production inference engine with built-in INT8 calibration and layer fusion for deployment.

Neural Network Distiller / NNI: Research-oriented pruning and quantization toolkits with configurable schedules.

### Compression Trade-off Table

## 7. Benchmark Datasets & Evaluation

### Standard Compression Benchmarks

ImageNet (ResNet/ViT compression): Standard vision compression benchmark; reports top-1 accuracy vs. compression ratio trade-off curves.

GLUE/SuperGLUE (BERT distillation): Standard NLP benchmark for evaluating distilled language models like DistilBERT and TinyBERT against full BERT.

MLPerf Inference: Industry-standard benchmark suite explicitly measuring latency/throughput of compressed models on target hardware.

WikiText/C4 Perplexity (LLM quantization): Standard perplexity-based evaluation for quantized large language models, measuring degradation from FP16 to INT8/INT4.

### Evaluation Metrics

Compression Ratio: Original model size / compressed model size.

Accuracy Retention: Compressed accuracy / original accuracy, often reported as absolute point drop.

Inference Latency: Wall-clock time per inference on target hardware, the metric that ultimately matters for deployment.

FLOPs Reduction: Theoretical compute reduction, useful for hardware-agnostic comparison but imperfectly correlated with real latency.

## 8. Key Challenges & Limitations

### Accuracy Cliff at Extreme Compression

Compression often degrades gracefully until a threshold, beyond which accuracy collapses sharply (e.g., below 2-bit quantization, or above ~95% unstructured sparsity) — practical deployment must stay comfortably before this cliff.

### Outlier Activations in Large Language Models

LLMs develop a small number of activation "outlier" dimensions with much larger magnitude than the rest, which naive uniform quantization handles poorly; specialized methods (e.g., mixed-precision decomposition, per-channel scaling) are required to preserve accuracy.

### Hardware-Software Co-Design Gap

Research routinely proposes novel sparsity/quantization schemes without matching hardware kernel support, meaning many published compression techniques never translate to real deployed speedups.

### Distillation Ceiling

Student model capacity fundamentally limits how much teacher knowledge can transfer; below some architecture size, distillation cannot recover teacher-level accuracy regardless of training signal quality.

## 9. Hyperparameter Tuning

### Distillation Temperature au

Higher au (e.g., 4-20) softens teacher outputs more, exposing more inter-class structure; too high washes out useful signal into near-uniform distributions. Typical effective range: au \in [2, 10].

### Distillation Loss Weight \alpha

Balances hard-label cross-entropy against soft-target distillation loss. Common starting point: \alpha = 0.5, tuned per task; tasks with noisy labels often benefit from lower \alpha (trusting the teacher more).

### Pruning Schedule

Gradual pruning schedules (e.g., cubic sparsity ramp-up over training) outperform pruning all-at-once; typical schedules ramp from 0% to target sparsity over the first 50-80% of training/fine-tuning steps.

### Quantization Bit-Width Selection

Start from INT8 as a safe default (near-lossless for most models); move to 4-bit only for large models (LLMs with billions of parameters) where redundancy is high enough to tolerate the additional error.

## 10. Real-World Applications & Case Studies

### DistilBERT

DistilBERT distills BERT-base into a 6-layer model (40% smaller, 60% faster) retaining 97% of BERT's language understanding performance on GLUE, demonstrating that distillation preserves most task-relevant capability at a fraction of the size.

### Mobile Vision: MobileNet + Quantization

MobileNet architectures combined with INT8 post-training quantization power real-time on-device photo classification and AR effects on billions of smartphones, running efficiently on mobile NPUs.

### LLM Quantization for Local Inference

4-bit quantization methods (GPTQ, AWQ) enable multi-billion-parameter language models to run on consumer GPUs with limited VRAM, making local LLM inference practical outside data centers.

### Speech Recognition on Embedded Devices

Wake-word detection and on-device speech recognition models are aggressively pruned and quantized to run continuously on low-power microcontrollers within tight energy budgets (always-listening scenarios).

## 11. Integration with Other Methods

### Compression + Neural Architecture Search

Jointly searching for architectures that are both accurate and compression-friendly (hardware-aware NAS) yields better Pareto-optimal accuracy/efficiency trade-offs than compressing an architecture chosen without efficiency in mind.

### Compression + Federated Learning

Gradient/model compression techniques (Section 2 of distributed learning literature) share mathematical machinery with model compression — both exploit low-rank structure and quantization to reduce communicated or stored data.

### Compression-Aware Fine-Tuning of Foundation Models

Parameter-efficient fine-tuning methods (LoRA-style low-rank adapters) can be viewed as a compression-inspired alternative to full fine-tuning, updating only a small low-rank delta rather than all weights.

## 12. Future Research Directions

### Sub-4-Bit Quantization Without Retraining

Pushing post-training quantization below 4 bits while retaining acceptable accuracy, without requiring expensive quantization-aware retraining, remains an active challenge especially for LLMs.

### Dynamic/Input-Adaptive Compression

Compressing differently per input (e.g., early-exit networks, adaptive-width inference) rather than applying a single static compressed model to all inputs, trading accuracy for speed only when the input is "easy."

### Automated Hardware-Aware Compression Search

Fully automated pipelines that jointly search pruning ratios, quantization bit-widths, and architecture choices against a target hardware's actual latency measurements rather than proxy metrics like FLOPs.

### Compression-Robust Training from Scratch

Training objectives that make models inherently more compressible from the start (encouraging low-rank, sparse, or quantization-friendly weight structure during pretraining) rather than compressing after the fact.

## 13. Summary & Key Takeaways

Model compression makes large trained networks practical to deploy under memory, latency, and energy constraints. Key insights:

1. Three Core Techniques: Quantization (lower numerical precision), pruning (remove redundant weights/structures), and distillation (transfer knowledge to a smaller model) are complementary and often combined.

2. Structured vs. Unstructured Trade-off: Unstructured pruning achieves higher theoretical compression but requires specialized sparse hardware for real speedup; structured pruning sacrifices some ratio for guaranteed dense-matrix speedup.

3. QAT Beats PTQ at Low Bit-Widths: Quantization-aware training recovers more accuracy than post-training quantization, especially below 8 bits.

4. Soft Targets Carry Dark Knowledge: Distillation's soft-label temperature scaling transfers inter-class similarity structure that hard labels discard.

5. Outliers Complicate LLM Quantization: Large language models require specialized handling of activation outliers for accurate low-bit quantization.

6. Accuracy Cliffs Exist: Compression degrades gracefully until a threshold, beyond which accuracy collapses sharply — stay comfortably before that cliff in production.

7. Hardware Co-Design Matters: Compression techniques only yield real-world speedup when matched with hardware/kernel support; profile on target hardware, not just theoretical FLOPs.

8. Pipelines Compound: Combining distillation, pruning, and quantization sequentially (each stage followed by recovery fine-tuning) achieves greater compression than any single technique alone.

---

## Appendix: Practical Labs

### Lab 1: Uniform Affine Quantization

import torch

def quantize_tensor(x, bits=8):
 """Uniform affine quantization to `bits`-bit integers."""
 qmin, qmax = 0, 2**bits - 1
 x_min, x_max = x.min().item(), x.max().item()

 scale = (x_max - x_min) / (qmax - qmin) if x_max != x_min else 1.0
 zero_point = round(-x_min / scale)
 zero_point = max(qmin, min(qmax, zero_point))

 x_q = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax)
 return x_q, scale, zero_point

def dequantize_tensor(x_q, scale, zero_point):
 return scale * (x_q - zero_point)

def test_quantization():
 torch.manual_seed(0)
 x = torch.randn(1000) * 3.0 # simulate a weight tensor

 for bits in [8, 4, 2]:
 x_q, scale, zp = quantize_tensor(x, bits=bits)
 x_dequant = dequantize_tensor(x_q, scale, zp)
 mse = torch.mean((x - x_dequant) ** 2).item()
 print(f"{bits}-bit: scale={scale:.4f}, MSE={mse:.6f}")

test_quantization()
print("Uniform affine quantization implemented")

### Lab 2: Magnitude-Based Iterative Pruning

import torch
import torch.nn as nn
import copy

def magnitude_prune(model, sparsity):
 """Zero out the smallest-magnitude weights globally, return a mask dict."""
 all_weights = torch.cat([
 p.data.abs().flatten() for name, p in model.named_parameters() if 'weight' in name
 ])
 threshold = torch.quantile(all_weights, sparsity)

 masks = {}
 for name, p in model.named_parameters():
 if 'weight' in name:
 mask = (p.data.abs() > threshold).float()
 p.data *= mask
 masks[name] = mask
 return masks

def iterative_pruning(model, target_sparsity=0.9, n_steps=5, finetune_fn=None):
 """Gradually increase sparsity, optionally fine-tuning between steps."""
 sparsities = [target_sparsity * (i + 1) / n_steps for i in range(n_steps)]
 all_masks = None
 for step, sparsity in enumerate(sparsities):
 all_masks = magnitude_prune(model, sparsity)
 if finetune_fn is not None:
 finetune_fn(model)
 actual_sparsity = compute_sparsity(model)
 print(f"Step {step+1}/{n_steps}: target={sparsity:.2%}, actual={actual_sparsity:.2%}")
 return all_masks

def compute_sparsity(model):
 total, zeros = 0, 0
 for name, p in model.named_parameters():
 if 'weight' in name:
 total += p.numel()
 zeros += (p.data == 0).sum().item()
 return zeros / total if total > 0 else 0.0

def test_pruning():
 torch.manual_seed(0)
 model = nn.Sequential(nn.Linear(64, 32), nn.ReLU(), nn.Linear(32, 10))

 def dummy_finetune(m):
 pass # placeholder for a real training loop

 iterative_pruning(model, target_sparsity=0.9, n_steps=3, finetune_fn=dummy_finetune)

test_pruning()
print("Iterative magnitude pruning implemented")

### Lab 3: Knowledge Distillation Loss

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.5):
 """Combine hard-label cross-entropy with soft-target KL divergence."""
 hard_loss = F.cross_entropy(student_logits, labels)

 student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
 teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
 soft_loss = F.kl_div(student_log_probs, teacher_probs, reduction='batchmean') * (temperature ** 2)

 return alpha * hard_loss + (1 - alpha) * soft_loss

def test_distillation():
 torch.manual_seed(0)
 batch_size, n_classes = 16, 10

 student_logits = torch.randn(batch_size, n_classes, requires_grad=True)
 teacher_logits = torch.randn(batch_size, n_classes) * 2.0 # sharper teacher
 labels = torch.randint(0, n_classes, (batch_size,))

 for temp in [1.0, 4.0, 10.0]:
 loss = distillation_loss(student_logits, teacher_logits, labels, temperature=temp)
 print(f"Temperature={temp}: distillation loss={loss.item():.4f}")

test_distillation()
print("Knowledge distillation loss implemented")

### Lab 4: Low-Rank Weight Factorization

import torch
import torch.nn as nn

def low_rank_factorize(weight, rank):
 """SVD-based low-rank approximation of a weight matrix."""
 U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
 U_r = U[:, :rank]
 S_r = S[:rank]
 Vh_r = Vh[:rank, :]

 # Fold singular values into U for a two-matrix factorization: W ≈ (U_r * S_r) @ Vh_r
 W_approx = (U_r * S_r) @ Vh_r
 compression_ratio = weight.numel() / (U_r.numel() + Vh_r.numel())
 return W_approx, U_r * S_r, Vh_r, compression_ratio

class LowRankLinear(nn.Module):
 """Drop-in replacement for nn.Linear using a low-rank factorized weight."""
 def __init__(self, in_features, out_features, rank):
 super().__init__()
 self.U = nn.Parameter(torch.randn(out_features, rank) * 0.01)
 self.V = nn.Parameter(torch.randn(rank, in_features) * 0.01)
 self.bias = nn.Parameter(torch.zeros(out_features))

 def forward(self, x):
 return x @ (self.U @ self.V).T + self.bias

def test_low_rank():
 torch.manual_seed(0)
 weight = torch.randn(256, 256)

 for rank in [256, 64, 16, 4]:
 W_approx, U, Vh, ratio = low_rank_factorize(weight, rank)
 error = torch.norm(weight - W_approx) / torch.norm(weight)
 print(f"Rank {rank}: relative error={error:.4f}, compression ratio={ratio:.2f}x")

 layer = LowRankLinear(256, 256, rank=16)
 x = torch.randn(4, 256)
 out = layer(x)
 print(f"LowRankLinear output shape: {out.shape}")

test_low_rank()
print("Low-rank factorization implemented")