Attention Head Pruning

The pruning criterion is the heuristic used to score and rank attention heads for removal. Common criteria include:

• Magnitude-based: Removing heads with the smallest L1 or L2 norm of their weight matrices.
• Importance-based: Using metrics like the average attention entropy or the sensitivity of the loss function to a head's removal.
• Gradient-based: Methods like Movement Pruning score heads based on how much their weights change during fine-tuning. The choice of criterion directly impacts the trade-off between compression and retained accuracy.

Attention head pruning is a form of structured pruning, removing entire, coherent computational units. This differs from unstructured pruning of individual weights. The granularity is at the head level, which is advantageous because:

• It produces a smaller, dense model that maintains standard execution patterns.
• It leverages existing, highly optimized matrix multiplication kernels without requiring specialized sparse hardware.
• The reduction in parameters directly translates to lower FLOPs for the attention computation, which scales quadratically with sequence length.

The pruning schedule defines when and how much to prune. Key schedules include:

• One-shot Pruning: Removing a target percentage of heads in a single step after training, often followed by fine-tuning.
• Iterative Pruning: Cyclically pruning a small fraction of heads (e.g., 10%) and then fine-tuning the model to recover accuracy before the next pruning step. This is often more effective at preserving performance.
• Pruning-Aware Training: Incorporating a sparsity-inducing regularization term during training to encourage heads to become prunable.

The core hypothesis enabling this technique is that transformer attention heads exhibit significant functional redundancy. Empirical studies show that many heads learn similar or trivial attention patterns (e.g., focusing on [CLS] tokens or punctuation). Pruning aims to remove these redundant heads while preserving a diverse set of specialized heads (e.g., for syntactic, semantic, or coreference tasks). The goal is to maintain the model's capacity—its ability to represent complex functions—with fewer computational resources.

Before pruning, a sensitivity analysis is often conducted to guide the strategy. This involves:

• Measuring the performance drop from ablating individual heads or layers.
• Identifying which layers are most sensitive to pruning; early and late layers often exhibit different tolerances.
• Using this analysis to implement layer-wise pruning rates, applying more aggressive pruning to less sensitive layers. Evaluation post-pruning must go beyond task accuracy to include metrics like inference latency, throughput, and memory footprint on target hardware.

After pruning, sparse fine-tuning (or rewinding) is critical to recover lost accuracy. The pruned model's remaining weights are fine-tuned on the task data, often with the sparsity pattern fixed. Key considerations:

• Learning Rate: A smaller learning rate is typically used to avoid catastrophic forgetting of the remaining useful features.
• Rewinding: A powerful technique where weights are reset to an earlier checkpoint in the original training (not the final values) before fine-tuning, often leading to better recovery.
• The amount of fine-tuning data and epochs required is proportional to the aggressiveness of the prune.

Structured pruning removes entire, structurally coherent groups of weights—such as filters, channels, or attention heads—resulting in a smaller, dense model that maintains hardware-friendly execution patterns. Unlike unstructured pruning, it produces models that can run efficiently on standard hardware (e.g., GPUs) without requiring specialized sparse computation libraries.

• Key Examples: Removing entire convolutional filters, attention heads, or matrix columns.
• Hardware Efficiency: The resulting dense sub-network leverages optimized BLAS libraries and achieves predictable speedups.
• Trade-off: Often achieves less compression for a given accuracy drop compared to unstructured methods, but delivers more reliable latency reduction.

Model distillation (or knowledge distillation) trains a smaller, more efficient student model to mimic the behavior of a larger, more accurate teacher model. The student learns not just from the original training labels (hard targets) but from the teacher's softened output probabilities (soft targets), which contain richer relational information.

• Objective: Capture the teacher's "dark knowledge" in a compact form.
• Process: The student's loss function combines a distillation loss (vs. teacher logits) and a standard task loss (vs. ground truth).
• Relation to Pruning: An alternative compression paradigm; can be combined with pruning where a pruned model is fine-tuned using distillation from the original dense model.

Quantization reduces the numerical precision of a model's weights and activations (e.g., from 32-bit floating-point to 8-bit integers) to decrease memory footprint, increase memory bandwidth utilization, and accelerate computation on supported hardware.

• Post-Training Quantization (PTQ): Applied after training with minimal calibration data; faster but may impact accuracy.
• Quantization-Aware Training (QAT): Simulates quantization during training, allowing the model to adapt and preserve accuracy.
• Synergy with Pruning: Often used in sequence; pruning reduces the number of parameters, quantization reduces the bit-width of each remaining parameter, leading to multiplicative compression benefits.

A sparse neural network is a model where a significant proportion of its parameters are exactly zero, a state typically induced by unstructured pruning. The sparsity pattern is irregular, requiring specialized software runtimes or hardware (e.g., sparse tensor cores) for efficient computation.

• Contrast with Structured Pruning: Attention head pruning is structured; sparse networks from unstructured pruning have scattered zeros.
• Computational Challenge: Requires sparse matrix multiplication kernels to skip multiplications with zeros.
• High Potential Compression: Can achieve extreme sparsity levels (e.g., 90%+ zeros) while maintaining accuracy, but realizing theoretical speedups on general hardware is complex.

The Lottery Ticket Hypothesis posits that within a dense, randomly-initialized neural network, there exist sparse subnetworks ('winning tickets') that, when trained in isolation from their initial weights, can match the performance of the original network. This finding motivates pruning-at-initialization and provides a theoretical framework for understanding network compressibility.

• Implication: Not all weights are equally important; critical subnetworks exist from the start.
• Iterative Magnitude Pruning (IMP): The algorithm used to discover these winning tickets by iterative pruning and rewinding.
• Relevance: Suggests that attention head pruning may be identifying and preserving critical computational sub-networks within the transformer.

A Mixture of Experts architecture is a conditional computation paradigm where different parts of the model (experts) are activated on a per-input basis via a routing network. During inference, only a subset of the total parameters are used, making it inherently sparse and efficient.

• Sparse Activation: For a given token, the router selects only the top-k experts (e.g., top-2).
• Scale vs. Cost: Enables dramatically increasing model parameter count (e.g., trillion-parameter models) without a proportional increase in inference FLOPs.
• Comparison to Pruning: While pruning removes parameters permanently, MoE keeps all parameters but uses them conditionally. Both aim to reduce active compute per forward pass.