Large language models continue to grow in parameter count, but serving them at scale requires more than raw compute. Model pruning removes redundant weights or entire structures from a network to reduce memory bandwidth and inference latency without fully retraining from scratch. For teams building production AI infrastructure, understanding pruning techniques is essential to balancing accuracy against throughput and cost.
What Is Model Pruning?
At its core, model pruning is the process of eliminating parameters that contribute little to a network's output. The technique predates transformers, but LLMs introduce unique challenges due to their depth, attention mechanisms, and sheer scale. Pruning can be applied at various granularities, from individual weights to full attention heads, and can occur during training, after training, or as part of a compression pipeline that includes quantization and distillation.
Unstructured vs. Structured Sparsity
Unstructured sparsity zeros out individual weights based on saliency criteria such as absolute magnitude or gradient information. Because the non-zero elements follow no regular pattern, realizing inference speedups requires sparse matrix kernels and hardware support, such as NVIDIA Ampere's 2:4 structured sparsity constraints or custom CUDA implementations. Without these, dense GEMM libraries often perform the same amount of work as before, meaning memory savings do not automatically translate to lower latency.
Structured sparsity, by contrast, removes coherent blocks: attention heads, feed-forward dimensions, or even whole layers. The resulting model is smaller in every dimension that hardware schedulers care about, so standard dense BLAS routines run faster on smaller matrices. The tradeoff is a steeper accuracy drop per parameter removed, which is why structured pruning is often paired with light recovery fine-tuning or distillation.
Pruning Techniques for Transformers
Several methods have emerged specifically for transformer-based LLMs:
- Magnitude pruning. The simplest baseline: weights with the lowest absolute values are set to zero. It is easy to implement but ignores the interaction between weights and activations.
- Wanda. Prunes weights by multiplying their magnitudes by the corresponding input activation norms, achieving strong one-shot results without retraining.
- SparseGPT. Uses approximate second-order information to adapt remaining weights after pruning, enabling accurate one-shot compression of massive GPT-style models.
- LLM-Pruner. A structured approach that respects transformer dependencies, such as query-key-value groupings and feed-forward expansions, allowing task-agnostic pruning with minimal recovery tuning.
- Sheared LLaMA. Combines structured pruning with continued pre-training on a targeted data distribution, recovering accuracy while permanently reducing model size.
Most production pipelines now treat pruning as one stage in a broader optimization workflow that may also include quantization to INT8 or INT4 and runtime fusion of attention layers.
Code Example: Magnitude Pruning in PyTorch
The following snippet demonstrates unstructured L1 magnitude pruning on a linear layer, followed by a structured variant that removes entire output channels. These patterns generalize to the multi-layer perceptron blocks inside transformer stacks.
import torch
import torch.nn.utils.prune as prune
Unstructured: zero 30% of individual weights with smallest absolute values
linear = torch.nn.Linear(4096, 4096)
Top comments (0)