Tensor Product Attention Is All You Need

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

该论文试图解决现代语言模型在处理更长输入序列时面临的内存开销问题。具体来说，随着输入序列长度的增加，传统的自注意力机制需要较大的键值（KV）缓存，这导致了显著的内存消耗。这是一个长期存在的挑战，特别是在资源受限的环境中。

论文提出了一种名为张量积注意力（Tensor Product Attention, TPA）的新颖注意力机制。TPA通过使用张量分解来紧凑地表示查询、键和值，并通过上下文低秩分解（contextual factorization）进一步压缩这些表示。这种方法不仅显著减少了KV缓存的大小，还与旋转位置编码（RoPE）无缝结合，从而提高了模型的质量和内存效率。此外，基于TPA，论文引入了一种新的模型架构——T6（Tensor ProducT ATTenTion Transformer），用于序列建模。

论文通过广泛的实验证明，T6在多个语言建模任务中超越了包括MHA、MQA、GQA和MLA在内的标准Transformer基线模型。实验设计涵盖了多种评估指标，如困惑度以及一系列知名的基准测试。此外，TPA的内存效率使得在固定资源约束下能够处理更长的序列，解决了现代语言模型中的关键可扩展性挑战。代码已开源，可在GitHub上找到：https://github.com/tensorgi/T6。未来值得继续深入研究的方向包括探索TPA在其他任务上的应用以及进一步优化其计算效率。

最近在这个领域中，相关的研究还包括：1.《Longformer: The Long-Document Transformer》——提出了一种可以处理超长文档的Transformer变体；2.《Big Bird: Transformers for Longer Sequences》——通过稀疏注意力机制来处理更长的序列；3.《Performer: Generalized Attention via Performer Kernel Methods》——利用随机特征方法来近似注意力机制，以提高计算效率。