Beyond KV caching: Shared attention for efficient LLMs

Abstract

The rapid scaling of Large Language Models (LLMs) necessitates advancements in computational and memory efficiency during inference. While methods like Multi-Query Attention (MQA), Grouped-Query Attention (GQA), and Cross-Layer Attention (CLA) reduce Key–Value (KV) cache size by sharing K/V pairs, strategies to further reduce computational load include reusing computed attention weights across layers, an idea explored previously (e.g., LazyFormer (Ying et al., 2021)). This paper provides an extensive empirical investigation into the phenomenon of attention weight isotropy—high similarity in attention distributions across layers—within diverse modern LLMs (7B-72B scale). We demonstrate how this isotropy develops during pretraining, offering a fundamental insight into LLM attention dynamics. Leveraging these findings, we systematically evaluate and validate a cross-layer weight sharing technique, termed Shared Attention (SA). SA selectively reuses computed attention weights in layer spans identified as isotropic through our analysis. Our experiments across multiple benchmarks show that strategically applied SA maintains comparable performance to baseline models, particularly in later layers where isotropy is pronounced, while significantly reducing computational FLOPs and key cache requirements associated with attention calculation. This work provides principled guidance for optimizing attention mechanisms based on empirically observed layer dynamics in contemporary LLMs. Code and resources are available at https://github.com/metacarbon/shareAtt.