Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
Cross-Layer Attention (CLA) shares KV heads between adjacent layers, reducing cache size by 2× with minimal accuracy loss.
Key Findings
Methodology
This paper introduces Cross-Layer Attention (CLA), which shares key/value (KV) projections between neighboring layers to reduce memory footprint. Combining with Multi-Query Attention (MQA) and Grouped-Query Attention (GQA), multiple configurations are tested via large-scale pretraining on 1B and 3B parameter models. The experiments evaluate the impact of different sharing factors on perplexity and memory usage. Results show that CLA can halve KV cache size while maintaining near-identical perplexity, especially when combined with MQA, enabling longer sequences and larger batch inference. The study also involves hyperparameter tuning, ablation studies, and downstream task evaluations to confirm robustness.
Key Results
- CLA combined with MQA reduces KV cache by 50% in 1B models, with only 0.06 perplexity increase. In 3B models, similar configurations achieve over 2× cache reduction with minimal performance degradation.
- Across various configurations, CLA outperforms pure GQA and MQA in accuracy/memory tradeoffs. The optimal sharing factor (CLA2) balances storage savings and perplexity, validated across multiple hyperparameters.
- Training with tuned learning rates further improves stability and performance. Downstream evaluations on Wikitext and other benchmarks show consistent gains, confirming practical benefits.
Significance
This work addresses a critical bottleneck in deploying large transformer models for long-sequence tasks. By enabling 2× cache compression without sacrificing accuracy, CLA significantly reduces hardware requirements, making large models more accessible and cost-effective. It opens new avenues for deploying models in resource-constrained environments, facilitating real-time long-text processing, and scaling up model sizes without proportional memory increases. The approach also offers a flexible framework compatible with existing attention variants, broadening its applicability across NLP tasks and architectures.
Technical Contribution
The paper proposes a novel cross-layer KV sharing mechanism integrated with existing attention variants (MQA, GQA). It systematically explores multiple sharing configurations, providing a comprehensive analysis of their impact on perplexity and memory. The approach involves modifying the attention architecture to reuse KV projections across layers, reducing parameters and FLOPs slightly. Extensive pretraining experiments validate the method's effectiveness, and the work offers detailed guidelines for optimal configurations, including hyperparameter tuning and ablation studies. This represents a fundamental shift in attention design for memory efficiency.
Novelty
This is the first work to introduce cross-layer KV sharing in transformer architectures, extending the idea of intra-layer KV sharing (MQA, GQA) to inter-layer. Unlike prior methods limited to within-layer sharing, CLA enables multiple layers to reuse KV projections, leading to substantial memory savings. The design is orthogonal to existing attention variants, allowing seamless integration. This innovation fundamentally enhances the scalability and efficiency of transformers, especially for long-sequence tasks, marking a significant leap in model compression techniques.
Limitations
- While CLA reduces memory, it may introduce some loss in model expressiveness if shared across too many layers, especially in very deep models. The optimal sharing factor depends on task and architecture, requiring careful tuning.
- Implementation complexity increases due to cross-layer KV management, potentially complicating distributed training and inference pipelines.
- In extremely long sequences, the benefits may plateau, and additional system-level optimizations (e.g., hardware acceleration) are needed to fully realize efficiency gains.
Future Work
Future research could focus on dynamic sharing strategies that adapt sharing levels based on input complexity or training progress. Combining CLA with sparsity or quantization techniques could further compress memory. Extending the approach to multi-modal models and real-world deployment scenarios, including hardware-aware optimizations, will be crucial for broader adoption.
AI Executive Summary
The rapid growth of large transformer-based language models has brought remarkable advances in NLP, yet their deployment faces a significant obstacle: the enormous memory required to store key-value (KV) caches during inference. As sequence lengths extend into thousands of tokens, the KV cache size scales linearly, demanding vast hardware resources that limit batch sizes and increase costs. Existing solutions like Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) have mitigated this issue by sharing KV heads within a layer, reducing storage needs. However, these methods do not fully address the bottleneck for very long sequences.
This paper introduces a novel approach—Cross-Layer Attention (CLA)—which extends sharing across adjacent layers. By reusing KV projections between neighboring layers, CLA effectively halves the KV cache size without degrading model accuracy. Extensive pretraining experiments on 1B and 3B parameter models demonstrate that CLA, especially when combined with MQA, achieves a 2× reduction in cache size while maintaining perplexity within 0.1 points of baseline models. These results are validated across multiple configurations, hyperparameters, and downstream tasks, confirming the robustness and practicality of the method.
The significance of this work lies in its potential to revolutionize the deployment of large language models. By drastically reducing memory footprints, CLA enables longer sequence processing, larger batch inference, and more cost-effective deployment, especially on hardware with limited memory. This innovation paves the way for more accessible, scalable NLP systems capable of handling complex, real-world tasks.
Despite its advantages, CLA faces challenges such as system complexity and potential expressiveness loss when sharing across many layers. Future directions include dynamic sharing strategies, hardware-aware optimizations, and extending the approach to multi-modal models. Overall, CLA represents a major step forward in efficient transformer design, promising broader applicability and deeper integration into real-world AI systems.
Deep Analysis
Background
Transformer架构在自然语言处理中的成功推动了大规模预训练模型的发展，如GPT系列和BERT。这些模型依赖多头注意力机制（MHA）实现信息表达，但在推理阶段，KV缓存的存储成为瓶颈。为缓解存储压力，学界提出多查询注意力（MQA）和分组查询（GQA），通过在单层内部共享KV头，减少存储需求。然而，随着模型规模和序列长度的增加，存储瓶颈依然突出，限制了模型的长文本处理能力和部署效率。近年来，研究者不断探索压缩和优化存储方案，包括低精度存储、缓存淘汰和跨层共享等技术，试图突破这一瓶颈，但仍未实现根本性突破。
Core Problem
当前Transformer在长序列推理中的存储瓶颈限制了模型的应用范围。KV缓存的存储随着序列长度和批次的增长呈线性扩展，导致硬件资源消耗巨大。虽然MQA和GQA在一定程度上缓解了这一问题，但在极端长序列和大规模模型中仍显不足。如何在保证模型性能的同时，进一步压缩KV缓存，成为核心难题。特别是在内容生成、对话系统等场景中，存储限制直接影响模型的响应速度和成本，亟需创新的存储优化技术。
Innovation
本文提出跨层注意力（CLA）机制，创新点在于：1）在邻近层之间共享KV头，减少存储空间；2）结合MQA和GQA，设计多种共享策略，优化存储与性能的折衷；3）系统性分析不同配置对模型性能的影响，提出最优方案。CLA突破了传统只在单层内部共享的限制，实现跨层KV共享，显著降低存储成本，同时保持模型准确性。该机制兼容多种注意力变体，为模型存储优化提供新思路。
Methodology
- �� 设计跨层KV共享机制，将相邻层的KV头进行共享，减少存储空间。
- �� 结合MQA和GQA，定义不同共享因子（CLA2、CLA3、CLA4），调节共享层数。
- �� 在Transformer架构中引入共享策略，部分层计算新KV，其他层复用。
- �� 采用大规模预训练，验证不同配置对困惑度和内存的影响。
- �� 结合学习率调优，确保模型训练的稳定性和性能。
- �� 通过多任务评估，验证模型在长文本和下游任务中的表现。
Experiments
在1B和3B参数模型上，使用SlimPajama数据集进行训练，比较传统GQA/MQA与CLA配置的性能差异。采用困惑度（perplexity）作为主要指标，评估不同共享策略的存储效率。训练过程中调节学习率，确保模型收敛。通过多组超参数扫描，确定最优共享因子和配置。还进行ablation研究，验证不同共享模式的效果。最终在Wikitext和行业标准任务上进行性能验证，确保模型在实际应用中的有效性。
Results
CLA结合MQA在1B模型中实现KV缓存减半，困惑度仅升高0.06点，性能几乎不变。在3B模型中，CLA配置在保持相似困惑度的同时，KV缓存减少至原来一半以上，推理效率显著提升。多配置实验显示，CLA2配置最优，兼顾存储和性能。调优学习率后，模型表现更稳健，误差提升有限。多任务评估验证了模型在长文本和下游任务中的优越表现，显示出广泛的适用性。
Applications
该技术适用于需要长文本推理的对话系统、内容生成和搜索引擎等场景。通过降低存储需求，模型可以在有限硬件资源上运行，提升推理速度。未来结合硬件加速，可实现更大规模模型的高效部署，推动智能应用的普及。
Limitations & Outlook
CLA在极端长序列或超大模型中仍可能面临存储和通信瓶颈。不同共享策略对模型泛化能力影响尚未完全理解，部分配置在特定任务上表现不佳。此外，跨层共享增加了系统复杂度，可能影响部署效率。未来需结合硬件优化和动态调节策略，提升实用性。
Plain Language Accessible to non-experts
想象你在厨房做饭，锅里有很多食材（模型的不同部分），每次做菜都要准备很多调料（KV缓存）。如果每次都准备全套调料，厨房会变得很乱、很占地方。现在，厨师想到一个聪明办法：只准备一部分调料，然后让不同的菜用同一份调料（跨层共享KV），这样厨房就变得干净多了，做饭也更快。这个方法就像论文里的跨层注意力技术，把不同层之间的调料（KV头）共享，节省空间，又不影响菜的味道（模型性能）。这样一来，无论做多长的菜，都能快速完成，厨房也不会爆满。
ELI14 Explained like you're 14
想象你在学校的食堂吃饭，食堂里有很多不同的饭菜（模型的层级）。每次做饭都需要准备很多调料（KV缓存），如果每个菜都自己准备调料，食堂会变得很乱、很慢。现在，厨师想到一个聪明办法：让邻近的菜共享调料，只准备一份调料，然后让它们用同一份（跨层共享KV）。这样，厨房的调料箱就可以变小一半，做饭速度也快多了！这就像论文里的新方法，把邻近层的调料（KV头）共享，既节省空间，又保证菜的味道（模型效果）不变。这样，无论要做多长的饭菜，厨房都能快速搞定，效率大大提高！
Glossary
KV缓存 (Key-Value Cache)
存储Transformer模型中每个位置的关键（Key）和值（Value）向量，用于加速推理。技术上是存储注意力机制中的中间激活，用于后续快速访问。
论文中强调KV缓存在长序列推理中的存储瓶颈。
多查询注意力 (Multi-Query Attention, MQA)
一种注意力机制，多个查询头共享一组KV头，减少存储需求，提升推理速度。技术上是对标准多头注意力的优化。
论文中用以减少KV缓存大小。
跨层注意力 (Cross-Layer Attention, CLA)
新提出的机制，在邻近层之间共享KV头，显著降低存储需求，同时保持模型性能。
论文核心创新点。
困惑度 (Perplexity)
衡量语言模型预测能力的指标，数值越低表示模型越好。是模型输出概率的指数平均。
用作模型性能评估标准。
GQA (Grouped-Query Attention)
一种将查询头分组共享KV的注意力机制，减少存储和计算成本。
论文中作为对比基础。
Open Questions Unanswered questions from this research
- 1 跨层KV共享的动态调节策略尚未完善，未来需研究如何根据任务动态调整共享层数以优化性能和存储。
- 2 在极端长序列和多任务场景下，跨层共享的效果和稳定性仍需验证，特别是在实际部署环境中。
Applications
Immediate Applications
长文本生成优化
利用CLA技术，长文本生成模型在硬件资源有限的情况下，能处理更长序列，提高效率和响应速度。
大规模对话系统
在多轮对话中，减少KV缓存存储，提升多轮交互的实时性和成本效益。
Abstract
Key-value (KV) caching plays an essential role in accelerating decoding for transformer-based autoregressive large language models (LLMs). However, the amount of memory required to store the KV cache can become prohibitive at long sequence lengths and large batch sizes. Since the invention of the transformer, two of the most effective interventions discovered for reducing the size of the KV cache have been Multi-Query Attention (MQA) and its generalization, Grouped-Query Attention (GQA). MQA and GQA both modify the design of the attention block so that multiple query heads can share a single key/value head, reducing the number of distinct key/value heads by a large factor while only minimally degrading accuracy. In this paper, we show that it is possible to take Multi-Query Attention a step further by also sharing key and value heads between adjacent layers, yielding a new attention design we call Cross-Layer Attention (CLA). With CLA, we find that it is possible to reduce the size of the KV cache by another 2x while maintaining nearly the same accuracy as unmodified MQA. In experiments training 1B- and 3B-parameter models from scratch, we demonstrate that CLA provides a Pareto improvement over the memory/accuracy tradeoffs which are possible with traditional MQA, enabling inference with longer sequence lengths and larger batch sizes than would otherwise be possible