全新注意力机制！Meta重磅新作MTA：多token注意力机制，性能暴涨！
前言 简单的任务，传统的Transformer却错误率极高。Meta FAIR团队重磅推出多token注意力机制（MTA ），精准捕捉复杂信息，带来模型性能飞升！
欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。
本文转载自新智元
仅用于学术分享，若侵权请联系删除
招聘高光谱图像、语义分割、diffusion等方向论文指导老师
注意力机制就像大模型的「慧眼」，帮模型在大量文本中找到关键信息。
不过，每个注意力权重只取决于单个查询和键向量的相似度。
找信息时，一次只能关注一个「小线索」，好比在庞大的图书馆里找一本书，却只能检索一个关键词。
最近，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
欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。
【技术文档】《从零搭建pytorch模型教程》122页PDF下载
QQ交流群：470899183。群内有大佬负责解答大家的日常学习、科研、代码问题。
其它文章
LSKA注意力 | 重新思考和设计大卷积核注意力，性能优于ConvNeXt、SWin、RepLKNet以及VAN
CVPR 2023 | TinyMIM：微软亚洲研究院用知识蒸馏改进小型ViT
ICCV2023|涨点神器！目标检测蒸馏学习新方法，浙大、海康威视等提出
ICCV 2023 Oral | 突破性图像融合与分割研究：全时多模态基准与多交互特征学习
HDRUNet | 深圳先进院董超团队提出带降噪与反量化功能的单帧HDR重建算法
南科大提出ORCTrack | 解决DeepSORT等跟踪方法的遮挡问题，即插即用真的很香
1800亿参数，世界顶级开源大模型Falcon官宣！碾压LLaMA 2，性能直逼GPT-4
SAM-Med2D：打破自然图像与医学图像的领域鸿沟，医疗版 SAM 开源了！
GhostSR|针对图像超分的特征冗余，华为诺亚&北大联合提出GhostSR
Meta推出像素级动作追踪模型，简易版在线可玩 | GitHub 1.4K星
CSUNet | 完美缝合Transformer和CNN，性能达到UNet家族的巅峰！
浙公网安备 33010602011771号