Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

High-Layer Attention Pruning with Rescaling

Songtao Liu, Peng Liu

TL;DR

The paper addresses the high latency of large language models by proposing HARP, a training-free structured pruning method that targets attention heads in the model’s higher layers. It removes the query and key projections in those layers and bypasses the self-attention computation, paired with a non-trainable rescaling parameter alpha in the residual to stabilize representation magnitudes, determined via a greedy, perplexity-based layer-wise search. Across LLaMA3.1-8B, Mistral-7B-v0.3, Qwen2-7B, and Gemma2-9B over 27 tasks, HARP consistently outperforms existing baselines in generation and maintains competitive results on discriminative tasks, with pronounced gains in long-context settings. The method also demonstrates practical efficiency: meaningful speedups for long sequences with modest parameter reduction, and ablations support the value of pruning higher layers and performing alpha search. Together, these results suggest HARP enables faster, more scalable deployment of large LLMs in resource-constrained environments without retraining.

Abstract

Pruning is a highly effective approach for compressing large language models (LLMs), significantly reducing inference latency. However, conventional training-free structured pruning methods often employ a heuristic metric that indiscriminately removes some attention heads across all pruning layers, without considering their positions within the network architecture. In this work, we propose a novel pruning algorithm that strategically prunes attention heads in the model's higher layers. Since the removal of attention heads can alter the magnitude of token representations, we introduce an adaptive rescaling parameter that calibrates the representation scale post-pruning to counteract this effect. We conduct comprehensive experiments on a wide range of LLMs, including LLaMA3.1-8B, Mistral-7B-v0.3, Qwen2-7B, and Gemma2-9B. Our evaluation includes both generation and discriminative tasks across 27 datasets. The results consistently demonstrate that our method outperforms existing structured pruning methods. This improvement is particularly notable in generation tasks, where our approach significantly outperforms existing baselines. Code is available at https://github.com/SongtaoLiu0823/HARP.

High-Layer Attention Pruning with Rescaling

TL;DR

The paper addresses the high latency of large language models by proposing HARP, a training-free structured pruning method that targets attention heads in the model’s higher layers. It removes the query and key projections in those layers and bypasses the self-attention computation, paired with a non-trainable rescaling parameter alpha in the residual to stabilize representation magnitudes, determined via a greedy, perplexity-based layer-wise search. Across LLaMA3.1-8B, Mistral-7B-v0.3, Qwen2-7B, and Gemma2-9B over 27 tasks, HARP consistently outperforms existing baselines in generation and maintains competitive results on discriminative tasks, with pronounced gains in long-context settings. The method also demonstrates practical efficiency: meaningful speedups for long sequences with modest parameter reduction, and ablations support the value of pruning higher layers and performing alpha search. Together, these results suggest HARP enables faster, more scalable deployment of large LLMs in resource-constrained environments without retraining.

Abstract

Pruning is a highly effective approach for compressing large language models (LLMs), significantly reducing inference latency. However, conventional training-free structured pruning methods often employ a heuristic metric that indiscriminately removes some attention heads across all pruning layers, without considering their positions within the network architecture. In this work, we propose a novel pruning algorithm that strategically prunes attention heads in the model's higher layers. Since the removal of attention heads can alter the magnitude of token representations, we introduce an adaptive rescaling parameter that calibrates the representation scale post-pruning to counteract this effect. We conduct comprehensive experiments on a wide range of LLMs, including LLaMA3.1-8B, Mistral-7B-v0.3, Qwen2-7B, and Gemma2-9B. Our evaluation includes both generation and discriminative tasks across 27 datasets. The results consistently demonstrate that our method outperforms existing structured pruning methods. This improvement is particularly notable in generation tasks, where our approach significantly outperforms existing baselines. Code is available at https://github.com/SongtaoLiu0823/HARP.

Paper Structure

This paper contains 37 sections, 5 theorems, 28 equations, 3 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. High-Layer Attention Pruning
  3. Pruning Attention Heads for Long-Context Modeling
  4. Theoretical Computation Complexity Analysis.
  5. High-Layer Attention Heads Are Less Important
  6. Theoretical Analysis.
  7. Empirical Analysis.
  8. Methodology
  9. Rescaling.
  10. Experiments
  11. Experimental Setup
  12. Baselines.
  13. Model Setting.
  14. Baseline Implementation.
  15. Cost analysis and comparison with baselines.
  16. ...and 22 more sections

Key Result

Theorem 1

When $\operatorname{Sim}(\mathbf{H}^{(\ell)}) \to 1$, token representations in layer $\ell$ become nearly parallel. Since queries and keys are linear projections of $\mathbf{H}^{(\ell)}$, their dot products across positions become almost constant, so the row-wise softmax yields nearly uniform attent

Figures (3)

  • Figure 1: Left: Performance comparison of our method against state-of-the-art baselines on LLaMA3.1-8B, demonstrating superior generation capabilities. Right: Ablation study on the rescaling parameter across benchmarks, showing that applying rescaling consistently leads to significant performance improvements over the variant without rescaling.
  • Figure 2: Searched alpha on LLaMA3.1-8B, Mistral-7B-v0.3, Qwen2-7B and Gemma2-9B.
  • Figure 3: Time complexity and runtime scaling with sequence length for the original and pruned models. Solid lines show the absolute runtime (left y-axis), while the green dashed line denotes the relative speedup (%) of the pruned model over the original, plotted on the secondary y-axis (right).

Theorems & Definitions (11)

  • Definition 1
  • Theorem 1
  • Proposition 1
  • Definition 2
  • Definition 3
  • Lemma 1: shi2022revisiting
  • Lemma 2: shi2022revisiting
  • Lemma 3
  • proof
  • proof
  • ...and 1 more