姚期智院士团队提出了新型注意力机制TPA，通过动态张量分解，显著降低了内存占用，节省高达90%甚至更多。TPA不存储完整的静态KV，而是保留分解的版本，同时保持了性能。该论文证明了流行的MHA、MQA、GQA都是TPA的特殊情况，实现了现代注意力设计的统一。基于TPA训练的新模型T6已在GitHub开源。TPA与LoRA等方法不同，它将QKV构造为上下文相关的分解张量，实现动态适应，并能与RoPE位置编码无缝集成。实验表明，TPA在多个基准测试中性能优异。

💡TPA通过动态张量分解，在不牺牲性能的前提下，将内存占用降低了90%以上，解决了现有注意力机制计算和内存开销大的问题。

🧮 TPA将QKV分别构造为与上下文相关的分解张量，实现动态适应，与LoRA等低秩分解方法相比，更具灵活性。

🔗 TPA可以与RoPE位置编码无缝集成，以较低的成本旋转分解KV，无需进行复杂的调整，便于实际应用。

🏆 实验结果表明，TPA在多个基准测试中表现优异，包括ARC、BoolQ、HellaSwag和MMLU等，性能优于或匹配所有基线。

新型注意力机制TPA，姚期智院士团队打造。

TPA对每个token做动态的张量分解，不存储完整的静态KV，而是保留分解的版本，内存占用节省90%（或者更多），而不会牺牲性能。

论文中还证明了流行的MHA、MQA、GQA都是TPA的特殊情况，用一个框架统一了现代注意力设计。

用此方法训练的新模型T6，代码已在GitHub开源。

论文发布后，有创业者表示，终于不用付那么多钱给云厂商了。

也有研究者认为，论文中的实验看起来很有希望，不过实验中的模型规模有点小，希望看到更多结果。

动态张量分解，无缝集成RoPE

尽管现有的注意力机制在众多任务中取得了不错的效果，但它还是有计算和内存开销大的缺陷。

DeepSeek-v2中提出的MLA压缩了KV缓存，但与RoPE位置编码不兼容，每个注意力头需要额外的位置编码参数。

为了克服这些方法的局限性，团队提出张量积注意力（TPA，Tensor Product Attention）。

新方法在注意力计算过程中对QKV做分解。

与LoRA系列低秩分解方法相比，TPA将QKV分别构造为与上下文相关的分解张量，实现动态适应。

通过只缓存分解后的秩，设置合适的参数可使内存占用降低90%或以上。

对于流行的RoPE位置编码，TPA可以与之无缝集成，实现以较低的成本旋转分解KV，无需进行复杂的调整。

在实验中，使用FineWeb-Edu 100B数据集训练模型，TPA与其他注意力设计相比始终保持较低的困惑度。

在ARC、BoolQ、HellaSwag和MMLU等基准测试中测试了零样本和少样本性能。TPA和TPA-KVonly在大多数任务中都优于或匹配所有基线。

论文由清华&上海期智研究员团队、UCLA顾全全团队合作，共同一作为清华博士生张伊凡与姚班校友、现UCLA博士生刘益枫。

此外还有来自心动网络Taptap的Qin Zhen。

论文地址：

https://arxiv.org/abs/2501.06425

开源代码：

https://github.com/tensorgi/T6

参考链接：

[1]https://x.com/yifan_zhang_/status/1879049477681741896

本文来自微信公众号“量子位”，作者：梦晨 ，36氪经授权发布。