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.

翻译：扩展语言模型以处理更长输入序列通常需要大型键值（KV）缓存，导致推理过程中产生显著的内存开销。本文提出张量积注意力（TPA），一种新颖的注意力机制，利用张量分解紧凑地表示查询、键和值，从而在推理时大幅压缩KV缓存大小。通过将这些表示分解为上下文低秩分量，并与旋转位置编码（RoPE）无缝集成，TPA在提升模型质量的同时实现了内存效率。基于TPA，我们引入了张量积注意力Transformer（T6），一种用于序列建模的新模型架构。通过在语言建模任务上进行广泛的实证评估，我们证明T6在包括困惑度和一系列既定评估基准在内的各项指标上，均超越或匹配了标准Transformer基线模型，包括多头注意力（MHA）、多查询注意力（MQA）、分组查询注意力（GQA）和多头潜在注意力（MLA）。值得注意的是，TPA在解码阶段的内存效率和计算效率，使得在固定资源约束下能够处理更长的序列，解决了现代语言模型中的一个关键可扩展性挑战。项目页面：https://github.com/tensorgi/TPA。