用张量分解压缩注意力内存，让大模型跑更长文本
Tensor Product Attention Is All You Need
- 用张量分解压缩查询/键/值表示，大幅减少推理时缓存大小
- 在语言建模任务中性能超越或持平主流注意力机制，支持更长序列
- 适合需要高效处理长文本的部署场景，如实时对话系统
将语言模型扩展至处理更长输入序列通常需要庞大的键值（KV）缓存，导致推理阶段内存开销巨大。本文提出张量积注意力（TPA），一种新型注意力机制，通过张量分解紧凑表示查询、键和值，显著降低推理时的KV缓存规模。通过将这些表示分解为上下文相关的低秩分量，并与旋转位置编码（RoPE）无缝集成，TPA在提升模型质量的同时实现内存高效。基于TPA，我们提出了张量积注意力变压器（T6）这一新型序列建模范式。在多种语言建模任务上的广泛实证评估表明，T6在困惑度及多个标准评估基准上优于或等同于标准Transformer基线，包括多头注意力（MHA）、多查询注意力（MQA）、分组查询注意力（GQA）和多头潜在注意力（MLA）。值得注意的是，TPA在解码阶段兼具内存与计算效率，使在固定资源约束下处理更长序列成为可能，解决了现代语言模型中的关键可扩展性挑战。
原文摘要 · Abstract (English)
Scaling language models to handle longer input sequences typically necessitates large key-value (KV) caches, resulting in substantial memory overhead during inference. In this paper, we propose Tensor Product Attention (TPA), a novel attention mechanism that uses tensor decompositions to represent queries, keys, and values compactly, substantially shrinking the KV cache size at inference time. By factorizing these representations into contextual low-rank components and seamlessly integrating with Rotary Position Embedding (RoPE), TPA achieves improved model quality alongside memory efficiency. Based on TPA, we introduce the Tensor ProducT ATTenTion Transformer (T6), a new model architecture for sequence modeling. Through extensive empirical evaluation on language modeling tasks, we demonstrate that T6 surpasses or matches the performance of standard Transformer baselines including Multi-Head Attention (MHA), Multi-Query Attention (MQA), Grouped-Query Attention (GQA), and Multi-Head Latent Attention (MLA) across various metrics, including perplexity and a range of established evaluation benchmarks. Notably, TPA's memory efficiency and computational efficiency at decoding stage enables processing longer sequences under fixed resource constraints, addressing a critical scalability challenge in modern language models. Project Page: https://github.com/tensorgi/TPA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。