Compressing Models for Production

Why Model Compression?

You've trained a state-of-the-art language model or computer vision network. It achieves 95% accuracy on your test set. Now you need to deploy it to a mobile app, an IoT device, or a server farm handling millions of inference requests per second.

Your model is 500 MB, requires 8 GB of GPU memory, and takes 2 seconds per inference. Your mobile app has 50 MB of storage. Your IoT device has 256 MB of RAM. Your server's inference budget is 50 milliseconds. You can't deploy the full model.

You have three main strategies: quantization (represent weights and activations with fewer bits), pruning (remove less-important weights), and knowledge distillation (train a small model to imitate a large one). Each trades model accuracy for size and speed. You'll use them together: quantize + prune a large model, or distill + quantize a small one, to hit your latency and storage budget while staying above your accuracy floor.

Quantization: Reducing Numeric Precision

Neural networks typically store weights and activations as 32-bit floating-point numbers (float32). Quantization converts them to lower precision: 8-bit integers (int8), 16-bit floats (float16), or even lower. This shrinks model size 4x and accelerates inference on hardware optimized for low-precision math.

Why quantization works: Most neural networks are over-parameterized. Weights cluster around a small set of values; the network doesn't need float32 precision to function. Quantization learns these value clusters and maps them to a discrete integer grid.

There are two main approaches:

Post-training quantization (PTQ): Quantize a trained float32 model without retraining. Fast, simple, works on any model. Most practical for deployment.

Quantization-aware training (QAT): Simulate quantization during training so the model learns to compensate for precision loss. Slower but usually more accurate (better for < 2% accuracy tolerance).

Here's a production example with both approaches:

import torch
import torch.nn as nn
import torch.quantization as q
from torchvision.models import resnet18
import time

# Load a pretrained model
model = resnet18(pretrained=True)
model.eval()

print("=== POST-TRAINING QUANTIZATION (PTQ) ===")

# Step 1: Prepare for quantization
# Fuse operations (e.g., Conv + BatchNorm) to reduce operations
model.qconfig = q.get_default_qconfig('fbgemm')  # Use facebook's efficient quantization
q.prepare_qat(model, inplace=True)

# Step 2: Calibrate on representative data (100-500 samples, not training data!)
# This finds the min/max ranges for each tensor to determine quantization scales
calibration_loader = [torch.randn(16, 3, 224, 224) for _ in range(5)]  # 80 samples
with torch.no_grad():
    for batch in calibration_loader:
        model(batch)

# Step 3: Convert to quantized model
model_quantized_ptq = q.convert(model, inplace=True)

# Measure size reduction
def get_model_size(model):
    total_params = sum(p.numel() for p in model.parameters())
    total_size_mb = total_params * 4 / (1024**2)  # 4 bytes per float32
    return total_params, total_size_mb

params_f32, size_f32 = get_model_size(model)
params_int8, size_int8 = get_model_size(model_quantized_ptq)

print(f"Float32 model: {size_f32:.2f} MB")
print(f"Int8 model: {size_int8:.2f} MB")
print(f"Compression ratio: {size_f32 / size_int8:.1f}x")

# Step 4: Benchmark inference speed
print("\n=== Inference Speed Benchmark ===")
test_input = torch.randn(1, 3, 224, 224)

# Float32 inference
model.eval()
with torch.no_grad():
    start = time.time()
    for _ in range(100):
        _ = model(test_input)
    time_f32 = (time.time() - start) / 100 * 1000  # ms per inference

# Int8 inference (on CPU; GPU quantization is more complex)
model_quantized_ptq.eval()
with torch.no_grad():
    start = time.time()
    for _ in range(100):
        _ = model_quantized_ptq(test_input)
    time_int8 = (time.time() - start) / 100 * 1000  # ms per inference

print(f"Float32 latency: {time_f32:.2f} ms")
print(f"Int8 latency: {time_int8:.2f} ms")
print(f"Speedup: {time_f32 / time_int8:.1f}x")

# Step 5: Validate accuracy on a small test set
# (In practice, you'd use your full test set)
# test_accuracy_f32 = evaluate(model, test_loader)
# test_accuracy_int8 = evaluate(model_quantized_ptq, test_loader)
# accuracy_drop = test_accuracy_f32 - test_accuracy_int8
# print(f"Accuracy drop: {accuracy_drop:.2%}")

print("\n=== QUANTIZATION-AWARE TRAINING (QAT) ===")
# QAT is slower (requires retraining) but can recover accuracy loss from PTQ

model_qat = resnet18(pretrained=True)
model_qat.train()
model_qat.qconfig = q.get_default_qat_qconfig('fbgemm')
q.prepare_qat(model_qat, inplace=True)

# Fine-tune on training data with simulated quantization
# (In practice, you'd use your actual training loop)
# for epoch in range(10):  # Few epochs of fine-tuning
#     for batch, labels in train_loader:
#         optimizer.zero_grad()
#         output = model_qat(batch)
#         loss = criterion(output, labels)
#         loss.backward()
#         optimizer.step()

# Convert to quantized
model_qat.eval()
model_quantized_qat = q.convert(model_qat, inplace=True)

print("QAT offers better accuracy retention but requires retraining.")

Quantization variants:

• Symmetric quantization: Range is [-a, a]; simpler hardware support
• Asymmetric quantization: Range is [a, b]; better accuracy on some layers (weights often one-sided)
• Per-channel quantization: Different scale for each output channel; more accurate than per-tensor
• Mixed precision: Use int8 for most layers, float32 for sensitive layers (e.g., final classification layer)

Trade-offs:

• ✓ 4x size reduction (float32 → int8)
• ✓ 2-5x inference speedup on CPUs with int8 instructions
• ✓ No retraining for PTQ; minimal retraining for QAT
• ✗ Some accuracy loss (usually < 2%, can be > 5% for extreme quantization)
• ✗ GPU inference speedup is hardware-dependent (NVIDIA A100/H100 have int8 cores)
• ✗ Requires careful calibration to find good quantization ranges

When to use:

• PTQ: First-pass compression; try it on any pretrained model. If accuracy drop < 2%, ship it.
• QAT: If PTQ accuracy drop > 2% and you have time to retrain. Usually recovers 1-1.5% of the drop.
• Mixed precision: If specific layers are accuracy-sensitive; trade off compression on non-critical layers.

Pruning: Removing Low-Importance Weights

A trained neural network has many redundant weights—values close to zero that barely affect the output. Pruning removes these weights, sparsifying the network. The remaining sparse network is smaller (fewer non-zero values to store) and can be faster if your hardware supports sparse linear algebra.

Unstructured vs. Structured pruning:

• Unstructured: Remove individual weights; can reach very high sparsity (90%+) but doesn't reduce latency much on generic hardware
• Structured: Remove entire channels, filters, or layers; coarser-grained; actually speeds up inference because you skip computations entirely

Here's a practical comparison:

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
from torchvision.models import resnet18
import time

model = resnet18(pretrained=True)
model.eval()

print("=== UNSTRUCTURED PRUNING ===")

# Clone model for unstructured pruning
model_unstructured = resnet18(pretrained=True)

# Prune 30% of weights in all Conv2d layers (unstructured)
for module in model_unstructured.modules():
    if isinstance(module, nn.Conv2d):
        # L1 magnitude pruning: remove smallest weights first
        prune.l1_unstructured(module, name='weight', amount=0.3)

# Make pruning permanent (compress the sparse tensor)
for module in model_unstructured.modules():
    if hasattr(module, 'weight_orig'):
        prune.remove(module, 'weight')

# Count remaining parameters
def count_nonzero_params(model):
    total = sum(p.numel() for p in model.parameters())
    nonzero = sum((p != 0).sum().item() for p in model.parameters())
    return nonzero, total

nonzero_unstructured, total_params = count_nonzero_params(model_unstructured)
sparsity_unstructured = 1 - (nonzero_unstructured / total_params)

print(f"Remaining parameters: {nonzero_unstructured / total_params:.1%}")
print(f"Sparsity: {sparsity_unstructured:.1%}")

# Benchmark: unstructured sparsity doesn't help much on CPU/GPU without special ops
test_input = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    start = time.time()
    for _ in range(10):
        _ = model_unstructured(test_input)
    time_unstructured = (time.time() - start) / 10 * 1000

print(f"Unstructured pruned latency: {time_unstructured:.2f} ms (little speedup on generic hardware)")

print("\n=== STRUCTURED PRUNING ===")

# For structured pruning, remove entire output channels
model_structured = resnet18(pretrained=True)

# Remove channels where the L1 norm of all weights in that channel is small
for module in model_structured.modules():
    if isinstance(module, nn.Conv2d):
        # Structured: prune by output channel (entire filter)
        prune.ln_structured(module, name='weight', amount=0.3, n=1, dim=0)

# Make pruning permanent
for module in model_structured.modules():
    if hasattr(module, 'weight_orig'):
        prune.remove(module, 'weight')

nonzero_structured, _ = count_nonzero_params(model_structured)
sparsity_structured = 1 - (nonzero_structured / total_params)

print(f"Remaining parameters: {nonzero_structured / total_params:.1%}")
print(f"Sparsity: {sparsity_structured:.1%}")

with torch.no_grad():
    start = time.time()
    for _ in range(10):
        _ = model_structured(test_input)
    time_structured = (time.time() - start) / 10 * 1000

print(f"Structured pruned latency: {time_structured:.2f} ms (meaningful speedup)")

print("\n=== ITERATIVE PRUNING ===")

# Better accuracy: gradually increase sparsity and fine-tune
model_iterative = resnet18(pretrained=True)

for iteration in range(5):
    # Prune 10% of remaining weights each iteration
    for module in model_iterative.modules():
        if isinstance(module, nn.Conv2d):
            prune.l1_unstructured(module, name='weight', amount=0.1)

    # Fine-tune on training data (small epochs)
    # for epoch in range(1):
    #     for batch, labels in train_loader:
    #         output = model_iterative(batch)
    #         loss = criterion(output, labels)
    #         loss.backward()
    #         optimizer.step()

    nonzero, _ = count_nonzero_params(model_iterative)
    print(f"Iteration {iteration+1}: {nonzero / total_params:.1%} params remaining")

# Final: make permanent
for module in model_iterative.modules():
    if hasattr(module, 'weight_orig'):
        prune.remove(module, 'weight')

Pruning strategies:

• Magnitude pruning: Remove weights with smallest |value|. Simple, greedy; often effective.
• Gradients-based: Remove weights with smallest gradient magnitude; depends on loss function.
• Lottery ticket hypothesis: Find a sparse subnetwork initialized at the original weights that reaches similar accuracy when trained from scratch. Works but computationally expensive.
• Iterative pruning: Gradually increase sparsity; fine-tune between iterations. Better accuracy retention than one-shot.
• Layer-wise pruning: Tune sparsity per layer; some layers are more sensitive than others.

Trade-offs:

• ✓ High sparsity possible (90%+)
• ✓ Can pair with quantization for compound compression
• ✗ Unstructured pruning doesn't reduce latency on generic CPUs/GPUs (only model size)
• ✗ Structured pruning has lower maximum sparsity (typically 30-50%)
• ✗ Requires fine-tuning to recover accuracy; iterative pruning is slow
• ✗ Hardware support for sparse operations is limited

When to use:

• Structured pruning: When you need inference speedup and have time to fine-tune
• Unstructured pruning: When model size/storage is the constraint, not latency
• Iterative pruning: When accuracy must stay within 1% of original; time-consuming but reliable
• Combine with quantization: Pruned-and-quantized models beat either technique alone

Knowledge Distillation: Teaching a Small Student Model

Quantization and pruning compress an existing large model. Distillation takes a different approach: train a small "student" model to imitate the outputs of a large "teacher" model. The student learns a simpler, compressed representation of the task by copying the teacher's behavior.

Here's the core idea:

import torch
import torch.nn as nn
import torch.optim as optim

# Large teacher model (pretrained, fixed)
teacher_model = resnet50(pretrained=True)
teacher_model.eval()
for param in teacher_model.parameters():
    param.requires_grad = False  # Don't update teacher

# Small student model (randomly initialized, to be trained)
student_model = resnet18()  # Smaller model, same architecture family

# Distillation loss: KL divergence between teacher and student logits
# Also include a cross-entropy loss on true labels for stability
def distillation_loss(student_logits, teacher_logits, targets, temperature=4.0, alpha=0.7):
    """
    Soft targets from teacher: softmax(teacher_logits / T)
    Hard targets from true labels: cross_entropy(student_logits, targets)
    """
    soft_targets = nn.functional.softmax(teacher_logits / temperature, dim=1)
    soft_loss = nn.functional.kl_div(
        nn.functional.log_softmax(student_logits / temperature, dim=1),
        soft_targets,
        reduction='batchmean'
    ) * (temperature ** 2)

    hard_loss = nn.functional.cross_entropy(student_logits, targets)

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

# Training loop
optimizer = optim.Adam(student_model.parameters(), lr=0.001)
train_loader = ...  # Your training data

for epoch in range(10):
    for images, labels in train_loader:
        # Forward pass through both models
        with torch.no_grad():
            teacher_logits = teacher_model(images)

        student_logits = student_model(images)

        # Compute distillation loss
        loss = distillation_loss(student_logits, teacher_logits, labels)

        # Backprop and update student
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# Student is now trained; smaller, faster, and mimics teacher

The key hyperparameter is temperature: higher values (T=8-20) make the teacher's outputs "softer" (less peaked), giving the student gentler learning targets. Lower temperature (T=2-4) gives sharper targets.

Why it works: The student doesn't need to solve the original task from scratch. Instead, it copies the teacher's intermediate representations and decision boundaries. This is often easier than learning from raw data. The teacher has already distilled the task into a learned representation; the student just memorizes that representation more compactly.

Strengths:

• Flexible: the student can be any smaller architecture (different family, even different modality)
• Accurate: student often retains 95%+ of teacher accuracy with 50% fewer parameters
• Combines well with quantization and pruning

Limitations:

• Requires a good teacher model (if teacher is 90% accurate, student usually won't exceed 90%)
• Training can be slow (need to evaluate teacher on all training batches)
• Requires tuning temperature and alpha (soft/hard loss weight)

When to use: Use distillation when you have a large, accurate model and want to build a smaller one. Distill first, then quantize and prune the student for maximum compression.

Combining Techniques: Compound Compression

For maximum compression, stack multiple techniques. Order matters:

Recommended pipeline:

1. Distill: Train a small student model to match large teacher. (Size: -50%, Accuracy: -2-5%)
2. Prune (structured): Remove low-importance channels. (Size: -30%, Accuracy: -1-2%)
3. Quantize (PTQ or QAT): Convert float32 → int8. (Size: -75%, Accuracy: -0.5-2%)

Cumulative result: Original model → 1/20th size, 20-100x faster, with ~5-10% total accuracy loss.

Example (illustrative estimate) for a 200M-parameter model:

import torch
from torch.nn import Module

# Original teacher: 200M params, 800 MB, 100 ms latency
print("Original teacher model:")
print("  Size: 800 MB, Params: 200M")
print("  Latency: 100 ms")
print("  Accuracy: 95%")

# Step 1: Distill to student
print("\nAfter distillation (50M params):")
print("  Size: 200 MB (4x smaller)")
print("  Latency: 25 ms")
print("  Accuracy: 92% (3% drop)")

# Step 2: Prune student (remove 40% of channels)
print("\nAfter structured pruning:")
print("  Size: 120 MB (1.7x smaller than student, 6.7x smaller than original)")
print("  Latency: 16 ms (1.6x faster than student)")
print("  Accuracy: 90.5% (1.5% additional drop)")

# Step 3: Quantize (float32 → int8)
print("\nAfter quantization:")
print("  Size: 30 MB (4x smaller than pruned, 27x smaller than original)")
print("  Latency: 5 ms (5x faster than pruned, 20x faster than original)")
print("  Accuracy: 89.5% (1% additional drop)")

# Edge hardware benefit
print("\nOn mobile/edge with int8 accelerator:")
print("  Latency: 2-3 ms (30-50x faster than original)")

Why this order?

1. Distill first: Creates a smaller, simpler model that's easier to prune and quantize
2. Prune second: Removes redundancy from the distilled model; easier than pruning the original
3. Quantize last: Fine-grained compression that works well on already-simplified models

Monitoring accuracy during compression:

def compression_pipeline(teacher, train_loader, val_loader):
    # Step 1: Distill
    student = distill_model(teacher, train_loader)
    student_acc = evaluate(student, val_loader)
    print(f"Student accuracy: {student_acc:.1%}")

    # Step 2: Prune
    student_pruned = prune_model(student, amount=0.4)
    pruned_acc = evaluate(student_pruned, val_loader)
    print(f"Pruned accuracy: {pruned_acc:.1%} (drop: {student_acc - pruned_acc:.1%})")

    # If accuracy drop > 3%, fine-tune before next step
    if student_acc - pruned_acc > 0.03:
        student_pruned = finetune(student_pruned, train_loader, epochs=5)
        pruned_acc = evaluate(student_pruned, val_loader)
        print(f"After fine-tuning: {pruned_acc:.1%}")

    # Step 3: Quantize
    student_quantized = quantize_model(student_pruned, calib_loader)
    final_acc = evaluate(student_quantized, val_loader)
    print(f"Final accuracy: {final_acc:.1%} (drop from student: {student_acc - final_acc:.1%})")

    return student_quantized

Hardware-specific optimization:

• CPUs with int8 instructions (x86, ARM Snapdragon): Quantization is essential
• Mobile (iPhone Neural Engine, Qualcomm Hexagon): Distillation + structured pruning + quantization
• Cloud GPUs (NVIDIA, AMD): Quantization helps batch throughput; distillation is less critical
• Edge accelerators (Coral, Jetson): All three techniques compound benefits
• Older GPUs (GTX 1080, etc.): Only distillation helps (not optimized for int8)

Compression Tradeoff Table

| Technique | Model Size Reduction | Latency Reduction | Accuracy Loss | Complexity | When to Use | |-----------|---------------------|-------------------|---------------|-----------:|-----------| | Quantization alone | 4x | 2-5x | < 2% | Low | Quick win on any pretrained model | | Pruning alone | 2-4x | 1-3x (CPU) | 1-3% | Medium | When you control inference hardware | | Distillation alone | 2-4x | 2-4x | 2-5% | High | When you have a large teacher | | Distill + Prune + Quantize | 20-50x | 20-100x | 8-15% | High | Maximum compression for latency-critical systems |

Most production deployments use quantization + pruning or distillation + quantization. Choose based on your accuracy budget and latency requirement.

Hardware-Specific Compression Strategy

| Hardware Target | Best Strategy | Size Reduction | Speedup | Setup Time | |-----------------|---------------|----------------|---------|-----------| | Modern CPU (x86, ARM) | PTQ → Quantization | 4x | 2-3x | < 1 hour | | Mobile (iOS, Android) | Distill + Quantize | 20-40x | 5-10x | 4-8 hours | | Edge accelerator (Coral) | Distill + Quantize + Pruning | 30-50x | 10-20x | 8-16 hours | | Cloud GPU (NVIDIA A100) | Distillation or no compression | 2-4x | 1-2x | 2-4 hours | | Inference appliance (ONNX) | Quantization (INT8) | 4x | 3-5x | 1-2 hours |

Rule of thumb: Compression effort doubles or triples as you target more constrained hardware.

Understanding the Accuracy-Efficiency Frontier

Compression is fundamentally about trading accuracy for size/speed. Different approaches hit different points on this curve. Understanding the frontier helps you make principled choices.

Illustrative estimate (not official benchmarks): On a typical ImageNet classification model:

| Compression Approach | Model Size | Latency (CPU) | Accuracy Loss | |---------------------|-----------|---------------|---------------| | Float32 baseline | 100% | 100% | 0% | | Quantization only (INT8) | 25% | 40% | ~0.5% | | Pruning 50% only | 60% | 70% | ~1% | | Distill (student) + Quantize | 20% | 35% | ~3-5% | | Distill + Prune 50% + Quantize | 12% | 25% | ~5-8% |

These are approximate ranges; actual numbers depend on architecture, dataset, and compression technique. Always benchmark on your specific model.

The key insight: combining techniques compounds the benefit. A 4x compression from quantization × 3x from distillation × 2x from pruning = 24x total, though accuracy loss compounds too.

Choosing Based on Constraints

Different deployment scenarios favor different strategies:

Latency-critical (web API, mobile app):

• Quantization is essential (2-5x latency improvement on modern hardware)
• Pruning helps only if your hardware supports sparse operations
• Distillation + quantization is ideal if you can afford training time

Storage-critical (on-device, embedded):

• Distillation (teaches a small model) is best
• Then quantize the small model for additional 4x compression

Throughput-critical (batch inference, data center):

• Quantization on GPU (huge throughput improvement)
• Batching is more important than compression (see lesson 22)

Energy-critical (IoT, mobile):

• Quantization (less computation = less power)
• Pruning (fewer operations = less power)
• Profile on-device to measure actual power draw

Common Mistake: Not Validating on Real Hardware

You compress a model and check accuracy on your laptop's CPU. Then you deploy to edge devices (phones, IoT) and find a 30% slowdown. Or worse: the quantized model runs slower because your hardware doesn't have INT8 support, so it falls back to float32.

Different hardware benefits from different compression strategies:

• Quantization: CPUs with int8 instructions (modern x86, ARM Snapdragon, Apple Neural Engine) gain significant speedup; older CPUs or unsupported hardware may not. GPU quantization (CUDA) requires NVIDIA's CUTLASS or similar libraries.
• Pruning: Most general-purpose CPUs are not sparse-aware; sparse matrices run slowly unless using special sparse BLAS libraries. Specialized hardware (TPUs, sparse accelerators) benefit more.
• Distillation: Agnostic to hardware; just a smaller model.

Always benchmark on your target hardware before finalizing compression choices.

# Right approach: benchmark on all target devices
import time

for device_name, device in [('cpu', 'cpu'), ('gpu', 'cuda'), ('mobile_emulator', 'cpu')]:
    model = model.to(device)
    model.eval()

    latencies = []
    for _ in range(100):
        start = time.time()
        with torch.no_grad():
            output = model(input_to_device)
        latencies.append((time.time() - start) * 1000)

    latencies = sorted(latencies)
    print(f"{device_name}: p50={latencies[50]:.2f}ms, p95={latencies[95]:.2f}ms, p99={latencies[99]:.2f}ms")

Report percentiles (p50, p95, p99), not just averages. A model that's fast on average but occasionally 10x slower is not suitable for real-time systems.

Compression decisions must be made with real deployment hardware and realistic latency requirements in mind.