Technical Summary: Efficient LLMs with AMP: Attention Heads and MLP Pruning

Problem Statement

Large Language Models (LLMs) have achieved human-level performance in cognitive tasks but suffer from extensive parameter counts, leading to high computational costs and slow inference. These factors hinder deployment in resource-constrained environments. While pruning is a promising strategy to reduce model size, existing methods face significant trade-offs:

• Structured Pruning: Removes entire components (e.g., layers, heads) but often requires computationally expensive fine-tuning or leads to performance degradation if critical components are removed.
• Semi-Structured Pruning: Requires specialized hardware (e.g., specific GPU sparsity patterns like 2:4) to achieve speedups, limiting deployment on consumer-grade hardware.
• Unstructured Pruning: Creates irregular sparse matrices that fail to accelerate inference on general-purpose hardware without specialized software optimizations.
• General Limitations: Many existing techniques rely on loss-based metrics, Taylor expansions, or complex heuristics that are computationally expensive, inflexible, or require specialized hardware. There is a need for a hardware-agnostic, efficient method that compresses LLMs with minimal performance loss and tangible inference speedups.

Methodology: AMP (Attention Heads and MLP Pruning)

The authors propose AMP, a structured pruning method that removes less critical attention heads within Multi-Head Attention (MHA) and neurons within Multilayer Perceptrons (MLP). Unlike methods that rely solely on weight magnitudes or loss changes, AMP evaluates structural importance by projecting input data onto weights to assess the contribution of each component to the residual stream.

Core Mechanism

The method operates on the principle that information flows through the residual stream. AMP calculates an importance score for each component based on the magnitude of its activation contribution to the final output of the layer.

1. MHA Component:

   • The MHA output is decomposed into a sum of contributions from individual attention heads ($O = \sum h_n W_n$).
   • The importance of the $n$-th head ($I_n$) is defined as the $\ell_1$-norm of the product of the head's output ($h_n$) and its corresponding output projection matrix ($W_n$):
     $$I_n = \|h_n W_n\|_1$$
   • This approach accounts for the nonlinear interactions between weights and inputs, unlike simple weight-norm pruning.

2. MLP Component:

   • The method targets the SwiGLU mechanism used in models like LLaMA, specifically the Up and Gate projections.
   • Since SwiGLU uses element-wise multiplication, neurons must be pruned in pairs from both projections to maintain architectural integrity.
   • The importance of the $m$-th pair of neurons ($I_m$) is calculated as the average $\ell_1$-norm of the element-wise product of the Gate and Up projection outputs across a sample of input tokens:
     $$I_m = \frac{1}{S} \sum_{s=1}^{S} | \text{SiLU}((x_s W_{\text{Gate}})_m) \cdot (x_s W_{\text{Up}})_m |$$

3. Pruning Process:

   • A small calibration set (e.g., 50 random samples from Alpaca) is used to extract mean activation values.
   • Components (heads and neurons) are ranked by importance across all layers.
   • To achieve a target compression ratio $c$, the bottom $c\%$ of attention heads and $c\%$ of MLP neurons are removed uniformly across all layers.
   • Post-Training: The pruned model undergoes a lightweight fine-tuning process using Low-Rank Adaptation (LoRA) to recover performance.

Key Contributions

• Novel Importance Metric: AMP introduces a data-driven metric that projects input activations onto weights to evaluate component importance, differing from traditional weight-only or loss-based heuristics.
• Hardware-Agnostic Speedup: The method achieves practical inference speedups (up to 1.25$\times$) without requiring specialized hardware or sparse tensor cores, making it suitable for consumer-grade GPUs.
• Efficiency: The pruning decision process takes only minutes on a consumer-grade GPU, and the subsequent fine-tuning requires negligible resources compared to full retraining.
• Coherence Validation: The authors propose and execute a "coherence check" experiment, demonstrating that removing the most important components identified by AMP causes catastrophic performance drops, validating the metric's ability to distinguish critical from non-critical structures.
• Flexibility: The method is applicable across different LLM families (LLaMA, Phi) and allows for flexible compression ratios without being bound by fixed sparsity patterns (e.g., 2:4).

Experimental Results

The authors evaluated AMP on LLaMA 7B, LLaMA-2 7B, Phi-1.5, and Phi-2 models using zero-shot benchmarks (WinoGrande, HellaSwag, ARC-e/c, PIQA) and perplexity (WikiText2).

• Performance:
  • LLaMA 7B: At a 20% pruning ratio, AMP achieved an average accuracy of 64.52%, outperforming state-of-the-art methods like LLM-Pruner (59.71%) and Shortened LLaMA (63.46%).
  • LLaMA-2 7B: At a 30% pruning ratio, AMP achieved 61.02% average accuracy, surpassing SliceGPT (51.50%) and LLM Surgeon (59.03%).
  • Improvement: AMP surpassed existing structured pruning techniques by up to 1.49 percentage points in average task accuracy.
• Inference Speed:
  • AMP achieved a 1.25$\times$ speedup for LLaMA 7B and 1.19$\times$ for LLaMA-2 7B.
  • This outperforms Shortened LLaMA (1.13$\times$) and rivals semi-structured methods like Wanda (1.24$\times$) without requiring specialized hardware.
• Perplexity: Fine-tuned pruned models showed significantly lower perplexity on WikiText2 compared to non-fine-tuned versions, with the gap widening at higher pruning ratios.
• Ablation Studies:
  • Coherence Check: Reversing the pruning order (removing the most important components) resulted in accuracy drops of >25 percentage points, confirming the metric's validity.
  • Component Isolation: Pruning only MHA heads or only MLP neurons resulted in significantly worse performance than pruning both simultaneously, indicating the necessity of a balanced approach.

Significance and Claims

The paper claims that AMP addresses the critical trade-offs in model compression by offering a method that is:

1. Effective: It outperforms current state-of-the-art structured pruning methods in accuracy while maintaining high compression rates.
2. Efficient: It requires minimal computational resources for both the pruning decision (minutes) and performance recovery (lightweight LoRA fine-tuning).
3. Practical: It enables tangible inference speedups on standard hardware, facilitating deployment in resource-constrained settings.
4. Sustainable: Aligned with Green AI principles, the method reduces operational costs by ~17.39% and carbon emissions by ~19.97% due to reduced inference latency and parameter count.

The authors emphasize that their method successfully identifies and removes non-essential structures, allowing for the deployment of highly optimized LLMs without the need for specialized hardware infrastructure. They also highlight the critical role of post-training (fine-tuning) in restoring model performance after pruning.