Efficient LLMs with AMP: Attention Heads and MLP Pruning

TL;DR

This work tackles the high computational cost of large language models by introducing AMP, a structured pruning method that simultaneously removes attention heads in MHA and neurons in MLPs. AMP uses activation-based importance scores, derived from projecting inputs onto weights, to rank components and prune uniformly across layers, avoiding the need for specialized hardware. With a lightweight post-training recovery using LoRA, AMP achieves up to 1.25× inference speedups and up to 1.49 percentage-point accuracy gains on select tasks, while maintaining compatibility across LLaMA and Phi families. The method aligns with Green AI by reducing emissions and cost, and is validated through coherence checks and ablation studies, demonstrating that jointly pruning MHA and MLP components is more effective than pruning either alone. These findings suggest AMP as a practical, efficient approach for deploying capable LLMs in resource-constrained environments, with potential for scaling to larger models and datasets.

Abstract

Deep learning drives a new wave in computing systems and triggers the automation of increasingly complex problems. In particular, Large Language Models (LLMs) have significantly advanced cognitive tasks, often matching or even surpassing human-level performance. However, their extensive parameters result in high computational costs and slow inference, posing challenges for deployment in resource-limited settings. Among the strategies to overcome the aforementioned challenges, pruning emerges as a successful mechanism since it reduces model size while maintaining predictive ability. In this paper, we introduce AMP: Attention Heads and MLP Pruning, a novel structured pruning method that efficiently compresses LLMs by removing less critical structures within Multi-Head Attention (MHA) and Multilayer Perceptron (MLP). By projecting the input data onto weights, AMP assesses structural importance and overcomes the limitations of existing techniques, which often fall short in flexibility or efficiency. In particular, AMP surpasses the current state-of-the-art on commonsense reasoning tasks by up to 1.49 percentage points, achieving a 30% pruning ratio with minimal impact on zero-shot task performance. Moreover, AMP also improves inference speeds, making it well-suited for deployment in resource-constrained environments. We confirm the flexibility of AMP on different families of LLMs, including LLaMA and Phi.

Figures (1)

• Figure 1: Pruning impact of AMP on WikiText2 perplexity for LLaMA-2 7B. The red dashed horizontal line indicates the performance of the original model, while the blue and green lines represent the performance of the pruned model with and without fine-tuning, respectively.