TLDR: Researchers introduce HIES (Head Importance-Entropy Score), a novel pruning criterion for Transformer models that combines gradient-based head importance with attention entropy. This unified score provides a more stable and accurate method for model compression, outperforming traditional HIS-only methods by up to 15.2% in model quality and 2.04x in stability, especially at aggressive pruning ratios. HIES ensures that functionally important heads, particularly those with concentrated attention patterns, are preserved, leading to more robust performance across various models and tasks.
Large Language Models (LLMs) and other Transformer-based models have become central to advancements in fields like natural language processing and computer vision. However, their impressive performance often comes with a significant computational cost, making them challenging to deploy, especially on devices with limited resources. To tackle this, researchers have developed various pruning methods, which essentially trim down the model by removing less important parts.
One popular pruning technique focuses on ‘attention heads’ within these models. Attention heads are crucial components that allow Transformers to weigh the importance of different parts of the input data. A common method for deciding which heads to prune uses a ‘Head Importance Score’ (HIS), which measures how much each head contributes to the model’s overall performance, particularly in terms of its impact on the loss function (a measure of how well the model is doing its job).
However, relying solely on HIS has its drawbacks. It primarily captures the gradient-driven contribution of a head, meaning it understands what directly affects the model’s output. What it often misses is the ‘diversity’ of attention patterns. Imagine two attention heads: one that focuses very sharply on a few critical pieces of information, and another that spreads its attention more broadly. Both might have similar HIS values, but their functional roles could be very different. This oversight can lead to instability, especially when models are aggressively pruned, causing a sharp drop in accuracy or unpredictable performance in real-world scenarios where data might vary slightly.
To overcome these limitations, a new pruning criterion called HIES (Head Importance-Entropy Score) has been introduced. HIES unifies the traditional Head Importance Score with ‘Attention Entropy’ (AE). Attention Entropy measures how concentrated or dispersed an attention head’s focus is across the input tokens. A low entropy score means the head is sharply focused on specific tokens, while a high entropy score indicates a more diffuse attention pattern.
The core idea behind HIES is to combine these two complementary signals. HIS tells us how much a head contributes to the task, while AE provides insights into the structural diversity of its attention. By integrating both, HIES can make more informed pruning decisions. For instance, it can help retain low-entropy heads that might be crucial for focusing on decisive tokens, even if their HIS value isn’t exceptionally high. This prevents the ‘collapse of structural diversity’ that often leads to performance degradation in HIS-only methods.
The researchers found that HIES-based pruning significantly improves model quality and stability. Empirically, it yielded up to a 15.2% improvement in model quality and a 2.04 times improvement in stability compared to HIS-only methods. This means models pruned with HIES can be compressed more aggressively without sacrificing accuracy or becoming unstable. The benefits of HIES were observed across various Transformer models, including BERT, LLaMA-27B, ViTLarge, and LLaVA-1.57B, and across diverse tasks like natural language understanding, image classification, and complex multimodal reasoning.
The study also revealed interesting patterns in how HIES prunes heads. While HIS tends to remove heads primarily from the lower layers of a Transformer, HIES promotes a more dispersed selection across lower, middle, and upper layers. This layer-adaptive pruning, driven by the entropy-aware term, helps maintain diversity and robustness throughout the model’s architecture.
Also Read:
- Optimizing LLM Compression: The Surprising Efficacy of Local Reconstruction
- GateSkip: Optimizing LLM Compute with Token-Wise Layer Skipping
In conclusion, HIES offers a more robust and efficient approach to Transformer pruning. By considering both a head’s direct contribution to the loss and the structural diversity of its attention patterns, it enables substantial model compression while preserving both accuracy and crucial inference-time stability, making these powerful models more practical for deployment in resource-constrained environments. You can read the full research paper here.
