突破性TPA机制将Transformer内存降低90%，姚期智团队开源新成果

突破性TPA机制将Transformer内存降低90%，姚期智团队开源新成果

新型注意力机制TPA（Tensor Product Attention）由姚期智院士团队打造，通过动态张量分解技术，实现了在不牺牲性能的前提下，将内存占用降低90%以上。

创新点：动态张量分解与RoPE兼容

现有的注意力机制虽然在众多任务中表现出色，但仍存在计算和内存开销大的问题。DeepSeek-v2提出的MLA方法虽然压缩了KV缓存，但与RoPE位置编码不兼容，每个注意力头需要额外的位置编码参数。

TPA的核心创新在于：

1. 动态张量分解：在注意力计算过程中对QKV做分解，与LoRA系列低秩分解方法相比，TPA将QKV分别构造为与上下文相关的分解张量，实现动态适应。

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

实验结果：性能与效率的双重提升

研究团队使用FineWeb-Edu 100B数据集训练模型，实验结果显示，TPA在保持较低困惑度的同时，显著降低了内存占用。在ARC、BoolQ、HellaSwag和MMLU等基准测试中，TPA在大多数任务中都优于或匹配所有基线模型。

开源与合作

基于TPA的新模型T6代码已在GitHub开源。论文发布后，业界反响热烈，有创业者表示这将显著降低云计算成本，研究者则期待看到更大规模模型的实验结果。

论文由清华大学、上海期智研究院、UCLA顾全全团队以及心动网络Taptap等机构合作完成，共同第一作者为清华大学博士生张伊凡与姚班校友、现UCLA博士生刘益枫。

论文地址：https://arxiv.org/abs/2501.06425
开源代码：https://github.com/tensorgi/T6

本文原文来自量子位公众号