NeurIPS 2025 Spotlight|还在为KV缓存爆炸苦恼？清华团队重新设计注意力

论文名称：Tensor Product Attention Is All You Need

发布时间：2025年10月23日

👉一键直达论文

👉一键直达Github

👉Lab4AI大模型实验室论文阅读

✨研究背景

大型语言模型在处理长输入序列时面临显著的计算和内存挑战，主要瓶颈在于自回归解码过程中键值（KV）缓存的内存开销随序列长度线性增长。现有方法如稀疏注意力、多查询注意力（MQA）、分组查询注意力（GQA）等虽能部分缓解问题，但存在性能下降、架构修改复杂或与旋转位置编码（RoPE）兼容性差等局限性。

✨研究目的

本文提出一种新型注意力机制——张量积注意力（TPA），旨在通过张量分解对查询、键和值进行上下文低秩表示，显著减少推理时的KV缓存大小，同时提升模型性能，并保持与RoPE的天然兼容性。

✨ 核心贡献

• 提出TPA机制：通过上下文张量分解实现Q、K、V的紧凑表示，相比标准注意力机制将KV缓存减少10倍以上，且性能优于MHA、MQA、GQA和MLA。
• 统一现有注意力机制：揭示MHA、MQA和GQA均为TPA的非上下文特例。
• 引入T6架构：基于TPA的新Transformer模型，在语言建模任务中验证其有效性。
• 无缝集成RoPE：TPA可直接替换现有模型（如LLaMA、Gemma）中的多头注意力层。

✨总结与展望

TPA通过张量分解提供了一种高效且表达力强的注意力机制，显著降低了长序列处理的内存需求，同时提升模型性能。未来工作可进一步探索高阶张量分解、硬件优化及更广泛的应用场景。