姚期智团队开源新型注意力，节省90%内存不降性能，一个框架统一MHA/MQA/GQA

梦晨
2025-01-15
17:08:58

来源：量子位

梦晨 发自 凹非寺
量子位 | 公众号 QbitAI

新型注意力机制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

Read More