浙大校友重磅革新Transformer！多token注意力让LLM开挂，错误率归0

新智元报道

编辑：英智 桃子

【新智元导读】简单的任务，传统的Transformer却错误率极高。Meta FAIR团队重磅推出多token注意力机制（MTA），精准捕捉复杂信息，带来模型性能飞升！

注意力机制就像大模型的「慧眼」，帮模型在大量文本中找到关键信息。

不过，每个注意力权重只取决于单个查询和键向量的相似度。

找信息时，一次只能关注一个「小线索」，好比在庞大的图书馆里找一本书，却只能检索一个关键词。

最近，Meta FAIR团队提出了多token注意力机制，为LLM的性能带来质的飞跃。

论文链接：http://arxiv.org/abs/2504.00927

作者还特别提到，这不是愚人节玩笑，而是一篇真实的论文。

在实际应用中，很多时候相关的上下文信息没办法通过单一token确定。

比如，想查找一个同时提到「Alice」和「rabbit」的句子时，按照传统注意力机制，得把这两个信息压缩到一个查询向量里。

但这样做不仅增加了向量编码的难度，还可能导致信息丢失。

就算用不同注意力头分别查找「Alice」和「rabbit」，也没法有效整合这些注意力权重，这极大地限制了模型处理复杂信息的能力。

MTA是为了解决传统注意力机制的这个问题而设计的。

MTA三大关键创新

MTA让模型能同时依据多个查询和键向量来确定注意力权重，利用更丰富的信息，实现精准的注意力分配。

主要有三个关键部分：键-查询卷积、头混合卷积和带深度缩放的组归一化。

键-查询卷积
键-查询卷积是MTA的一大创新亮点。

通过二维卷积操作，模型可以把多个查询和键token的信息放在一起，计算出更准确的注意力权重：

其中是二维卷积操作，有自己的核权重和大小。

在传统的注意力计算中，注意力权重仅由当前的查询和键向量对决定。

而键-查询卷积打破了这种局限，会综合考虑附近查询和键的信息。

就拿「Where did Alice see the rabbit?」这个问题来说，要是用传统注意力机制，很难同时关注到「Alice」和「rabbit」这两个关键信息。

键-查询卷积通过合理设置卷积核大小，能让模型同时捕捉到这两个信息。

具体计算时，它会在键和查询的长度维度上进行卷积操作。为了不让未来的信息捣乱，计算时只使用过去的查询。

实际操作中，为了简化流程，采用了一种更简便的双重掩码方法。

键-查询卷积有两种方式，一种是在softmax之前进行卷积，另一种是在softmax之后。

在softmax之前卷积，能更好地融合不同查询和键的信息。

在softmax之后卷积，注意力权重之间的相互作用就变成加法了：

在实验中，默认用的是softmax之前卷积的方式。

信息共享：头混合卷积

除了键-查询卷积，MTA还有个很特别的头混合卷积，它的作用是在不同注意力头之间共享信息。

在传统的多头注意力机制里，各个头之间的信息相对独立。

头混合卷积就像一座桥梁，打破了这种孤立，让不同头的注意力权重可以分享信息。

具体计算时，混合注意力权重有两种方式，一种是在softmax之后，像

是卷积核权重。

另一种是在softmax之前，对注意力的对数几率进行混合：

通过信息共享，模型能从多个角度综合信息，理解文本的能力就更强了。

整合关键组件

前面介绍了两种混合注意力权重的方式，MTA把键-查询卷积和头混合卷积结合，形成了强大的信息处理系统。

如果两种混合方法都是在softmax之前，那么它们可以通过单个三维卷积操作实现，如图所示。

模型训练时，随着层数越来越多，会出现一个问题，就是残差流会变得很大，这让梯度传递不太顺畅，影响模型的训练效果。

带深度缩放的组归一化就是解决这个问题的。它会对每个头分别进行归一化操作，而且会根据层的深度进行调整。

这样能让模型训练更稳定，梯度传递得更好。

MTA机制革新Transformer

理论上MTA机制很厉害，实际效果如何呢？

为了验证MTA的有效性，研究人员开展了一系列实验，涵盖了从简单任务到复杂的长上下文任务等。

简单任务：小任务凸显大优势

有一个简单的任务，MTA在这个任务里的表现比传统Transformer好太多了。

模型需要在由多个随机字母组成的块序列中，找到包含特定几个字母的目标块，并输出相关信息。

这个任务看似简单，却能精准暴露传统注意力机制的问题。

传统Transformer不能很好地把多个信息整合起来，要把两个问题字母的信息压缩到一个查询向量中，这对它来说很困难，错误率极高。

MTA在这个任务上的表现堪称惊艳。

它借助键-查询卷积，先分别找到每个问题字母在序列中的位置，然后通过卷积把这些信息整合起来，精准定位目标块。

实验结果令人惊喜，MTA几乎以零错误率完成了任务，这充分展现了MTA在处理多信息检索任务时的实力，与传统注意力机制相比，有着跨越式的进步。

语言建模：全面提升

在语言建模实验中，研究团队训练了880M参数的模型，并与传统Transformer模型、差分Transformer（DIFF Transformer）做比较。

所有模型都在SlimPajama数据集上，用Lingua框架训练。

为提高训练效率，MTA在每第4层应用键-查询卷积，头卷积应用于所有层，并固定了卷积核的维度。

实验结果令人眼前一亮。在验证困惑度方面，用MTA训练的模型在各个验证数据集上都表现更好。

带层缩放的组归一化对MTA和DIFF Transformer的性能提升特别重要。

在一些流行的基准测试中，如BoolQ、PIQA等，MTA模型也超过了基线模型，平均得分更高。

这说明MTA能有效提升模型性能，无论是理解文本含义，还是回答问题，都更出色。

长文本处理的好帮手

研究团队对模型进行了长上下文微调，把上下文的长度从2048增加到4096，还调整了一些训练参数。

用了MTA的模型在困惑度评估里比基线模型好很多。

在LAMBADA任务里，MTA模型预测下一个单词的时候，能更好地利用上下文信息，预测得更加准确。

在写作助手、摘要生成等场景中，MTA能帮助模型更好地理解文章内容和逻辑，生成更符合要求的回复。

精准找到长距离信息

像「大海捞针」（Needle-In-A-Haystack）和BabiLong这样的长距离依赖任务里，MTA的优势就更明显了。

「大海捞针」任务要在很长的文本里找到特定的信息。

MTA模型在这个任务中表现特别好，无论是在2k还是4k的上下文窗口中，准确率都有显著提升。

尤其是在查找隐藏较深的目标信息时，优势更为突出。

BabiLong任务评估模型能不能理解长文本里分散的各种事实，并进行推理。

QA1-5任务中，MTA模型在有很多干扰文本的情况下，也能保持较高的准确率，准确找到关键信息，并进行推理。

消融实验

为了搞清楚MTA各个组件的作用，研究人员做了消融实验。

在键-查询卷积实验里，就算只有2层用了MTA增强，模型就能超过强大的基线模型，6层MTA在性能和复杂程度之间达到了较好的平衡。

在卷积核初始化的实验里，用单位矩阵初始化的MTA模型训练时收敛得更快，性能也更好。

组归一化和指数深度缩放对于提升模型性能很关键，不同大小的卷积核虽然会影响评估结果，但总体的卷积核模式相似。

改变卷积操作和softmax的先后顺序，对模型性能的影响较小。

作者介绍

Meta多token注意力论文中，论文二作是一位华人研究科学家Tianlu Wang。

她曾获得了弗吉尼亚大学计算机科学博士学位，导师是Vicente Ordóñez Román教授。在此之前，她还获得了浙大计算机科学学士学位。

Tianlu Wang研究兴趣在于，与探索机器学习模型中的公平性、鲁棒性和问责制相关话题，尤其是在计算机视觉和自然语言处理系统方面。

参考资料：

http://arxiv.org/abs/2504.00927

https://x.com/jaseweston/status/1907260086017237207