04. Magnitude Pruning

Chapter 4 of 18 · 15 min

KEY INSIGHT

Magnitude pruning eliminates low-magnitude weights based on the hypothesis that smaller weights contribute less to model predictions. Magnitude pruning ranks weights by absolute value and removes the smallest fraction. The intuition: weights near zero have minimal influence on the dot products and transformations that constitute neural network computation. Removing them should little affect the model's behavior. The original magnitude pruning research demonstrated surprising effectiveness. Networks pruned to 50% sparsity after training and fine-tuned for several epochs matched dense network performance. This discovery contradicted prior assumptions that random initialization matters critically for performance. The typical magnitude pruning schedule follows a iterative pattern: prune for a training period, then restore pruned connections and re-train. The standard schedule from research: train for one epoch, prune 20% of remaining weights, re-train for a fraction of an epoch, and repeat. This gradual pruning allows the network to adapt to structural changes while maintaining performance. ```python def iterative_magnitude_pruning(model, train_loader, optimizer, prune_fraction=0.2, iterations=5, rewind_epochs=0.1): """ Standard iterative magnitude pruning procedure. Args: model: Trained PyTorch model train_loader: Training data optimizer: Optimizer for re-training prune_fraction: Fraction of weights to prune per iteration iterations: Number of prune-retrain cycles rewind_epochs: Epochs to re-train after each prune """ for iteration in range(iterations): # Calculate threshold for pruning all_weights = torch.cat([p.data.flatten() for p in model.parameters()]) threshold = torch.quantile( all_weights.abs(), prune_fraction ) # Apply unstructured pruning below threshold for name, module in model.named_modules(): if hasattr(module, 'weight'): prune.custom_from_mask( module, name='weight', mask=module.weight.abs() >= threshold ) # Re-train after pruning model.train() for batch in train_loader: optimizer.zero_grad() loss = compute_loss(model, batch) loss.backward() optimizer.step() ``` A critical failure mode involves re-winding to initial weights rather than trained weights. After pruning, networks sometimes perform better when weights are reset to their values from an earlier training stage. This rewinding suggests that magnitude pruning creates configurations dependent on training trajectory, and earlier configurations may generalize better. Hardware limitations become apparent at extreme sparsity levels. Below 95% sparsity, memory bandwidth rather than computation becomes the bottleneck. Magnitude pruning achieves this sparsity, but the practical speedup requires hardware optimized for sparse computation—hardware that remains uncommon in production environments.

EXERCISE

Compare magnitude pruning at 50%, 70%, and 90% sparsity levels on a classification model. Track accuracy across sparsity levels and identify the sparsity threshold where accuracy degrades significantly.