Model Pruning Techniques: Reducing Size and Increasing Latency

Previously in this course, we explored Post-Training Quantization (PTQ) to reduce memory footprint by lowering precision. While quantization focuses on bit-width, Pruning attacks model bloat from a different angle: by removing the weights themselves.

Model compression via pruning relies on the observation that deep neural networks are often over-parameterized. Many weights contribute negligible information to the final output. By zeroing these out, we introduce sparsity, which can lead to smaller model files and—with the right hardware support—faster inference.

Magnitude-Based Pruning: First Principles

Magnitude-based pruning operates on a simple heuristic: weights with the smallest absolute values contribute the least to the model's activations. If we set these values to zero, the impact on the overall loss function is theoretically minimized.

We typically define a target sparsity ratio (e.g., 20% of weights removed). The process follows these steps:

1. Ranking: Calculate the absolute value of all weights in a layer.
2. Thresholding: Determine the percentile value corresponding to your desired sparsity.
3. Masking: Create a binary mask where values below the threshold are zeroed out.
4. Fine-tuning: Re-train the model (even briefly) to allow the remaining non-zero weights to compensate for the pruned connections.

Implementing Magnitude Pruning

In PyTorch, the torch.nn.utils.prune module provides a clean interface for this. Instead of manually manipulating tensors, we use structured or unstructured pruning.

PYTHON
import torch
import torch.nn.utils.prune as prune

def apply_magnitude_pruning(model, amount=0.2):
    CE9178">"""
    Applies unstructured L1 magnitude pruning to all linear layers.
    """
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear):
            # Prune 20% of connections in this layer
            prune.l1_unstructured(module, name=CE9178">'weight', amount=amount)
            # Make the pruning permanent by removing the forward pre-hook
            prune.remove(module, CE9178">'weight')
    return model

Evaluating Sparsity Impact

Once you've pruned a model, you've created a sparse representation. However, "sparsity" is not synonymous with "speed."

If you use standard dense matrix multiplication (GEMM) kernels, a zeroed-out weight is still a floating-point operation. The model size on disk might shrink if you compress the weights, but inference latency won't improve unless you use sparse kernels or hardware that supports structured sparsity (like NVIDIA's Ampere architecture, which supports 2:4 structured sparsity).

When evaluating, track two metrics:

• Perplexity/Accuracy Degradation: How much does the model "forget" after pruning?
• Compression Ratio: The ratio of non-zero parameters to total parameters.

Fine-Tuning Pruned Models

Pruning is destructive. You are effectively deleting information. To recover performance, you must perform "recovery training" or fine-tuning. Because the model structure has changed (the mask is now part of the weight matrix), you should use a lower learning rate than your initial training phase to avoid destroying the remaining useful features.

PYTHON
def fine_tune_pruned_model(model, train_loader, optimizer, criterion, epochs=1):
    model.train()
    for epoch in range(epochs):
        for data, target in train_loader:
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            # Ensure pruned weights remain zero
            # This is optional if masks were applied permanently

Common Pitfalls

1. Over-pruning: Pruning beyond 50-70% often leads to catastrophic forgetting that no amount of fine-tuning can recover. Always monitor the "recovery" curve.
2. Ignoring Hardware Constraints: Unstructured pruning is excellent for reducing model storage size but rarely improves GPU latency. If latency is your primary goal, look into structured pruning (removing whole rows or columns).
3. Skipping the Calibration Set: If you don't use a representative dataset for fine-tuning, the model will overfit to the pruned state, leading to poor generalization.

Hands-on Exercise

1. Load your current project model from the Project Milestone: Domain-Specific Fine-Tuning.
2. Implement a function to calculate the global sparsity of the model.
3. Apply 30% unstructured pruning to the k, q, and v projection layers in your attention blocks.
4. Measure the accuracy on your validation set before and after pruning. How much does it drop?
5. Run one epoch of fine-tuning and observe if the accuracy recovers.

Recap

We've moved from managing model complexity via Managing Model Complexity: Pruning and Occam's Razor to actually modifying the weight tensors themselves. Pruning allows us to shed dead weight, making our models leaner for deployment. Remember: always validate the accuracy trade-off, as aggressive pruning is a one-way street unless you maintain the original weights.

Up next: We will explore Knowledge Distillation, where we teach a smaller "student" model to mimic the behavior of our large, pruned "teacher" model to achieve even greater efficiency.