Research map / Architectures and efficiency / Attention and transformers
Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Cross-Layer Attention (CLA), which shares KV activations between adjacent layers on top of MQA, halving KV cache size with near-MQA perplexity, validated by training 1B and 3B models from scratch.
Problem
The KV cache memory footprint limits batch size and sequence length during LLM decoding, and existing MQA/GQA already share KV heads only within a layer.
Approach
Share key and value activations across layers (sharing factor, e.g. CLA2 shares each KV projection between two consecutive layers), orthogonal to and combinable with MQA/GQA/MHA, to reduce the number of KV-cache layers.
Key results
- H128-MQA-CLA2 matches the KV cache footprint of H64-MQA (5120 bytes/token) while improving validation perplexity by about 0.21-0.48 points across the design space.
- At 1B scale with tuned learning rates, H128-MQA-CLA2 is within 0.04 validation perplexity of H128-MQA at half the KV cache (12.43 vs 12.39).
- At 3B scale, H128-MQA-CLA2 achieves validation perplexity 9.34 vs 9.52 for H128-MQA with 2x smaller KV cache.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add CLA2 to H128-MQA | 1B MQA · SlimPajama validation | perplexity: 13.54 -> 13.60 | Hurt |
| add CLA2 to H128-MQA while halving KV cache | 1B MQA · validation and Wikitext | perplexity: 12.39 -> 12.43; 19.30 -> 19.29 | No clear effect |
| add CLA2 to H64-MQA at the same KV cache footprint | 1B MQA · validation and Wikitext | perplexity: 12.74 -> 12.43; 20.00 -> 19.29 | Helped |
| add CLA2 to H128-MQA | 3B MQA · validation and Wikitext | perplexity: 9.52 -> 9.34; 13.63 -> 13.25 | Helped |
| add CLA2 to H64-MQA | 3B MQA, second experiment set · Wikitext | perplexity: 12.94 -> 12.99 | Hurt |
| add CLA2 using GQA4 (ablation) | 1B GQA4 · SlimPajama validation | perplexity: 13.36 -> 13.48 | Hurt |
| add CLA2 using GQA2 (ablation) | 1B GQA2 · SlimPajama validation | perplexity: 13.52 -> 13.59 | No clear effect |
| use CLA3 instead of CLA2 (ablation) | 1B MQA · SlimPajama validation | perplexity: 13.60 -> 13.77 | Hurt |
| use CLA4 instead of CLA2 (ablation) | 1B MQA · SlimPajama validation | perplexity: 13.60 -> 13.95 | Hurt |
| use non-uniform CLA2 sharing that keeps end layers separate (ablation) | 1B MQA · SlimPajama validation | perplexity: 13.60 -> 13.62 | Hurt |
| concentrate KV-producing layers at the front (ablation) | 1B MQA-CLA2 · SlimPajama validation | perplexity: 13.60 -> 13.75 | Hurt |
| concentrate KV-producing layers at the back (ablation) | 1B MQA-CLA2 · SlimPajama validation | perplexity: 13.60 -> 14.03 | Hurt |
Surprising findings
- At 3B scale, the d_head=64 MQA baseline outperformed the d_head=128 MQA baseline despite half the KV cache.
- Sharing factors above 2 (CLA3, CLA4) were slightly worse than CLA2 at matched KV cache size, and non-uniform sharing patterns were worse than uniform CLA2.
Limitations
- End-to-end inference efficiency at long context and large batch sizes was not evaluated and is left to future work.
- Experiments are limited to 1B and 3B parameter models trained on 30B to 100B tokens.
- CLA does not reduce core attention memory bandwidth per decoding step.
Experimental setup
Datasets: SlimPajama
Benchmarks: SlimPajama validation perplexity WikiText perplexity Hellaswag PIQA WinoGrande SciQ OpenBookQA BoolQ ARC-E
Metrics: perplexity accuracy KV cache bytes per token
Baselines: MHA GQA MQA
Builds on: multi-query attention grouped-query attention Llama architecture SwiGLU
Closest papers by meaning
- Value-Guided KV Compression for LLMs via Approximated CUR Decomposition (NeurIPS 2025 · similarity 0.90)
- More Value per Key: Asymmetric Sparse Attention for Faster LLM Decoding (NeurIPS 2026 · similarity 0.90)
- Improving Model Representation and Reducing KV Cache via Skip Connections with First Value Heads (NeurIPS 2025 · similarity 0.90)
- ArkVale: Efficient Generative LLM Inference with Recallable Key-Value Eviction (NeurIPS 2024 · similarity 0.89)
- CAKE: Cascading and Adaptive KV Cache Eviction with Layer Preferences (ICLR 2025 · similarity 0.89)
- SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget (ICLR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.