Complementary Attention Head Pruning for Efficient Transformers

2026-06-17

AI Summary

Transformer models are challenging to deploy in resource-constrained environments due to their large parameter counts, and existing pruning methods often suffer from instability, structural degradation, and reliance on manual hyperparameter tuning. This work proposes a novel approach that formulates attention head pruning as a global graph structure optimization problem. By leveraging graph clustering and information-theoretic distance metrics, the method automatically selects a topologically diverse and functionally complementary subset of attention heads—without requiring a predefined pruning ratio or gradient-based signals—thereby avoiding “proximity bias” and adaptively determining the number of heads to retain per layer. Integrated with marginal performance decay analysis and a post-processing pruning framework, the approach significantly outperforms current techniques on SST-5 and MNLI benchmarks, maintaining high performance even under aggressive compression and effectively preserving critical functional heads in intermediate layers.

Abstract

The remarkable success of Transformer-based models in natural language processing stems from architectural scaling, which leads to a large number of parameters and hinders deployment in resource-constrained environments. While structured pruning offers a pathway to compression, existing state-of-the-art methods often rely on gradient-based importance ranking or stochastic gating, which suffer from instability, structural degeneration, and the need for extensive manual hyperparameter tuning. In this paper, we introduce CAHP (Complementary Attention Head Pruning), a novel post-hoc framework that redefines head selection as a global graph-theoretical problem. Rather than evaluating heads in isolation, CAHP utilizes graph-based clustering combined with information-theoretic distance measures to identify and preserve a topologically diverse subset of complementary attention heads. Without requiring a predefined sparsity level or pruning ratio, the framework automatically determines the number of selected attention heads across layers by identifying a diminishing marginal performance curve, where pruning additional heads leads to a sharp degradation in performance, as determined by the chosen polynomial degree. Extensive evaluations on the SST-5 and MNLI benchmarks, across different Transformer model scales, demonstrate that CAHP consistently outperforms competitive baselines, particularly in high-compression regimes. Furthermore, our structural analysis shows that CAHP avoids the "proximity bias" of gradient-based pruning methods, which tend to preserve heads mainly in layers close to the output, and instead retains a functionally critical set of attention heads in the model's intermediate layers.

Problem

• Transformer pruning
• attention head
• model compression
• structured pruning
• resource-constrained deployment

Innovation

• Complementary Attention Head Pruning
• Graph-theoretical Pruning
• Information-theoretic Distance
• Post-hoc Structured Pruning
• Transformer Compression