Structured Pruning Techniques

Structured pruning is a technique for compressing machine learning models by removing entire, predefined groups of parameters. This approach creates more regular sparsity patterns that standard hardware like GPUs and TPUs can often handle more efficiently, leading to direct reductions in computation and memory bandwidth requirements without needing specialized kernels. Unstructured pruning, in contrast, offers maximum flexibility by removing individual weights, though its irregular sparsity patterns often require specialized hardware or software support (like sparse matrix libraries) to translate into actual inference speedups.

The core idea is to identify and eliminate larger structural components of the network that contribute least to the overall model performance. For Transformer-based LLMs, several structured pruning techniques are particularly relevant:

Neuron Pruning in Feed-Forward Networks (FFNs)

The Feed-Forward Network (FFN) sub-layers within each Transformer block typically consist of two linear transformations with an activation function in between. Neuron pruning targets the intermediate layer of this FFN, removing entire rows from the first weight matrix and corresponding columns from the second weight matrix. This effectively removes specific 'neurons' or dimensions from the intermediate representation.

Consider an FFN defined as:

$$\text{FFN}(x) = \text{Linear}_2(\text{Activation}(\text{Linear}_1(x)))$$

where $x$ is the input, $\text{Linear}_1$ maps from dimension $d_{model}$ to $d_{ffn}$, and $\text{Linear}_2$ maps from $d_{ffn}$ back to $d_{model}$. Pruning a neuron in the intermediate layer means reducing the dimension $d_{ffn}$. If we remove the $k$-th neuron, we remove the $k$-th row of the weight matrix in $\text{Linear}_1$ and the $k$-th column of the weight matrix in $\text{Linear}_2$.

Importance Criteria: How do we decide which neurons to prune? Common criteria include:

• Lp-Norm: Calculate the L1 or L2 norm of the weights associated with each intermediate neuron (the corresponding row in the first matrix and column in the second). Neurons with smaller norms are considered less important.
• Activation Statistics: Analyze the average activation values of neurons over a calibration dataset. Neurons with consistently low activation values might be pruned.
• Gradient Information: Measure the impact of a neuron on the loss function using gradient information or Taylor expansion approximations.

Pruning neurons directly reduces the width ($d_{ffn}$) of the FFN layers, shrinking the weight matrices and decreasing the number of floating-point operations (FLOPs) required during inference.

Attention Head Pruning

Multi-Head Attention (MHA) allows the model to jointly attend to information from different representation subspaces at different positions. Attention head pruning involves removing entire attention heads from the MHA layers.

Recall that in MHA, the input queries ($Q$), keys ($K$), and values ($V$) are linearly projected $h$ times (where $h$ is the number of heads) using different learned projection matrices for each head. The attention mechanism is applied independently for each head, and the results are concatenated and projected back to the model dimension.

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O$$ $$\text{where head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$

Pruning the $j$-th attention head means removing the corresponding projection matrices $W_j^Q, W_j^K, W_j^V$ and the corresponding segment of the output projection matrix $W^O$.

Importance Criteria: Selecting heads for pruning often involves more sophisticated metrics than simple weight norms:

• Head Importance Score: Measure the contribution of each head to the model's output or performance on a specific task. This can be estimated by masking out heads one by one and observing the change in loss or evaluation metrics on a calibration dataset.
• Redundancy Analysis: Analyze the similarity between attention patterns produced by different heads. Highly similar heads might be redundant, allowing some to be pruned.
• Gradient-Based Sensitivity: Assess how sensitive the model's loss is to the parameters of each head.

Removing attention heads directly reduces the computational cost of the attention mechanism, as fewer independent attention computations are performed, and the subsequent concatenation and output projection involves smaller matrices.

Comparison between irregular sparsity from unstructured pruning and regular sparsity resulting from removing entire neurons or attention heads in structured pruning. Pruned components are shown in gray.

Other Structured Approaches

While neuron and head pruning are the most common for Transformers, other granularities exist:

• Layer Pruning: Removing entire Transformer blocks (layers). This is a coarse-grained approach with a significant impact on model capacity and typically leads to substantial accuracy degradation unless the original model was heavily overparameterized.
• Block Pruning: Within weight matrices, removing contiguous blocks of weights. This can sometimes offer a compromise between the fine granularity of unstructured pruning and the hardware friendliness of neuron/head pruning.

Implementation and Fine-Tuning

Structured pruning is often implemented iteratively:

1. Train/Load Model: Start with a pre-trained LLM.
2. Score Structures: Calculate importance scores for the target structures (neurons, heads) across the relevant layers.
3. Prune: Remove the lowest-scoring structures up to the desired sparsity level. This typically involves creating masks or directly modifying the architecture to remove the parameters.
4. Fine-Tune: Retrain the pruned model for a period (often significantly shorter than the original pre-training) to recover accuracy lost during pruning.
5. Repeat (Optional): Steps 2-4 can be repeated to reach higher sparsity levels gradually.

Fine-tuning is a significant step. Removing large chunks of the model can initially cause a noticeable drop in performance. Fine-tuning allows the remaining parameters to adapt and compensate for the removed components.

Hardware Efficiency Gains

The primary advantage of structured pruning lies in its direct mapping to hardware efficiency. Removing an entire attention head or FFN neuron results in smaller dense matrix multiplications and reduced memory requirements. Standard deep learning libraries and hardware accelerators can execute these smaller dense operations efficiently without needing specialized sparse computation support. This often translates to measurable reductions in latency and memory footprint, making it attractive for deploying LLMs in resource-constrained environments.

Trade-offs and Approaches

Structured pruning forces choices at a coarser granularity than unstructured pruning. While it yields more hardware-friendly models, it might remove some parameters that unstructured pruning would have kept, potentially leading to a slightly larger accuracy drop for the same parameter count reduction. The choice between unstructured and structured pruning often depends on the target hardware, the required level of compression, and the acceptable performance degradation. Combining structured pruning with other techniques like quantization can further enhance efficiency, a topic explored in the next section.