张量积注意力是你所需要的
计算与语言
2026-01-13 v7 人工智能
机器学习
摘要
将语言模型扩展到处理更长的输入序列通常需要大的 key-value (KV) 缓存，在推理期间导致巨大的内存开销。本文提出了张量积注意力 (TPA)，这是一种新型注意力机制，使用张量分解来紧凑表示查询、键和值，从而在推理时大幅缩小 KV 缓存大小。通过将这些表示因式分解为上下文低秩组件并无缝集成到旋转位置嵌入 (RoPE) 中，TPA 实现了 improved 模型质量和内存效率。基于 TPA，我们引入了张量 ProducT ATTenTion Transformer (T6)，一种用于序列建模的新模型架构。通过对语言建模任务进行大量实证评估，我们展示 T6 在各种指标上（包括 perplexity 和一系列 established 评估基准）优于或匹配包括 Multi-Head Attention (MHA)、Multi-Query Attention (MQA)、Grouped-Query Attention (GQA) 和 Multi-Head Latent Attention (MLA) 等标准 Transformer 基线。值得注意的是，TPA 的内存效率和计算效率在解码阶段使能够在固定资源限制下处理更长的序列，解决了现代语言模型中的关键可扩展性挑战。项目页面: https://github.com/tensorgi/TPA。
关键词
引用
@article{arxiv.2501.06425,
title = {Tensor Product Attention Is All You Need},
author = {Yifan Zhang and Yifeng Liu and Huizhuo Yuan and Zhen Qin and Yang Yuan and Quanquan Gu and Andrew Chi-Chih Yao},
journal= {arXiv preprint arXiv:2501.06425},
year = {2026}
}
备注
Published in NeurIPS 2025 (Spotlight); Project Page: https://github.com/tensorgi/TPA