全新注意力机制!Meta重磅新作MTA:多token注意力机制,性能暴涨!

前言 简单的任务,传统的Transformer却错误率极高。Meta FAIR团队重磅推出多token注意力机制(MTA ),精准捕捉复杂信息,带来模型性能飞升!

欢迎关注公众号CV技术指南,专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

本文转载自新智元

仅用于学术分享,若侵权请联系删除

CV方向的准研究生们,未来三年如何度过?

招聘高光谱图像、语义分割、diffusion等方向论文指导老师

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

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

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

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

论文:

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

在实际应用中,很多时候相关的上下文信息没办法通过单一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研究兴趣在于,与探索机器学习模型中的公平性、鲁棒性和问责制相关话题,尤其是在计算机视觉和自然语言处理系统方面。

参考资料:



欢迎关注公众号CV技术指南,专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

计算机视觉入门1v3辅导班

【技术文档】《从零搭建pytorch模型教程》122页PDF下载

QQ交流群:470899183。群内有大佬负责解答大家的日常学习、科研、代码问题。

其它文章

分享一个CV知识库,上千篇文章、专栏,CV所有资料都在这了

明年毕业,还不知道怎么做毕设的请抓紧机会了

LSKA注意力 | 重新思考和设计大卷积核注意力,性能优于ConvNeXt、SWin、RepLKNet以及VAN

CVPR 2023 | TinyMIM:微软亚洲研究院用知识蒸馏改进小型ViT

ICCV2023|涨点神器!目标检测蒸馏学习新方法,浙大、海康威视等提出

ICCV 2023 Oral | 突破性图像融合与分割研究:全时多模态基准与多交互特征学习

听我说,Transformer它就是个支持向量机

HDRUNet | 深圳先进院董超团队提出带降噪与反量化功能的单帧HDR重建算法

南科大提出ORCTrack | 解决DeepSORT等跟踪方法的遮挡问题,即插即用真的很香

1800亿参数,世界顶级开源大模型Falcon官宣!碾压LLaMA 2,性能直逼GPT-4

SAM-Med2D:打破自然图像与医学图像的领域鸿沟,医疗版 SAM 开源了!

GhostSR|针对图像超分的特征冗余,华为诺亚&北大联合提出GhostSR

Meta推出像素级动作追踪模型,简易版在线可玩 | GitHub 1.4K星

CSUNet | 完美缝合Transformer和CNN,性能达到UNet家族的巅峰!

AI最全资料汇总 | 基础入门、技术前沿、工业应用、部署框架、实战教程学习

计算机视觉入门1v3辅导班

计算机视觉交流群

posted @ 2025-05-26 14:53  CV技术指南(公众号)  阅读(31)  评论(0)    收藏  举报