跳到正文
原文
36氪·AI·· 21 天前AI 评分32

北邮MiTA从快权重视角重构高效注意力,入选NeurIPS 2026并开源

NeurIPS'26:注意力也是FFN?从快权重出发,重新理解高效注意力

文章要点
  • 1北邮团队提出MiTA,从快权重视角统一理解高效注意力,入选NeurIPS 2026并开源代码
  • 2长序列基准LRA上MiTA平均准确率58.91%,训练时间从10.6小时降至2.4小时,测试用单张RTX 4090
  • 3DeiT-S上加MiTA的Top-1准确率79.8%,与完整注意力差0.1个百分点,FLOPs从4.6G降至4.4G
  • 4论文称64K到1024K序列推理加速约19至159倍,arXiv长文档分类准确率88.0%至98.7%,均为自报数据
决策要点

这项研究把长序列算力开销变成可设计的架构问题,为长上下文成本控制提供了新的技术思路。

  • 影响:长文档与长序列任务的算力开销在论文基准上明显下降;若此类机制进入开源或商用模型,相关场景的推理成本存在压缩空间,但原文未测企业级负载下的表现。
  • 参考:可考虑先算出当前长上下文调用或长文档处理的单次成本作为基线;值得跟踪采用此类注意力机制的开源模型版本,再评估是否替换现有方案。
  • 依据:全部性能数字为论文作者自报,含DeiT-S 79.8%对79.9%、LRA训练时间减少约77%,暂无独立验证。
正文

Image 1: 36氪

账号设置我的关注我的收藏申请的项目退出登录

登录

搜索

搜索

项目推荐

我要入驻

城市合作

NeurIPS'26:注意力也是FFN?从快权重出发,重新理解高效注意力

新智元· 2026年10月12日 03:31

北邮MiTA用快权重视角重构高效注意力,性能优异。

一直以来,在Transformer架构中,注意力(Attention)与前馈网络(Feed-Forward Network,FFN)被认为是两个功能不同且互补的模块:

前者让词元(Token)彼此交互,后者对每个词元进行非线性变换。

但将注意力机制的公式展开,就会看到一条容易被忽略的联系:标准注意力也可以写成一个FFN——只是它的权重由当前输入生成,隐藏层宽度随序列长度增长。

于是,长上下文的开销也有了另一种解释:我们不只是在处理更多词元,还在运行一个不断变宽的网络。那么,线性注意力、混合专家注意力、Top-k注意力这些看似不同的技术路线,能否从「如何组织这个网络」出发重新理解?

围绕这一问题,北京邮电大学团队提出的Mixture-of-Top-k Attention(MiTA),将已有的快权重(Fast Weights)视角发展为一套五维分类框架,并据此设计出可变形、可复用的Top-k专家,既保留全局信息,也按需读取原始细节。该工作已被NeurIPS 2026接收,代码已开源。

Image 2

**论文地址:**https://arxiv.org/pdf/2602.01219

**代码仓库:**https://github.com/QishuaiWen/MiTA

Attention也是FFN,关键在于权重从哪里来

令查询

Image 3

,键和值按列组成

Image 4

。

将常见FFN使用的两层多层感知机(Multi-Layer Perceptron,MLP)与标准注意力放在一起,前者的隐藏层宽度记为D:

Image 5

将Softmax展开,把缩放和依赖查询的归一化吸收到权重中:

Image 6

对应关系随即出现:

Image 7

查询对应输入,变换后的键和值对应两层权重,指数函数对应激活函数,偏置为零,隐藏层宽度则由D变为序列长度N。

区别在于权重的来源:普通FFN使用训练得到的慢权重(Slow Weights),注意力则把当前输入生成的K, V作为随上下文变化的快权重。

更重要的区别在于宽度:**每个键值对对应一个隐藏单元,上下文越长,这个快权重MLP就越宽。**每个查询访问N个隐藏单元,处理N个查询便产生

Image 8

的总体开销。

高效注意力的设计,由此转化为一个统一的问题:如何让不断扩大的快权重网络保持可扩展性?

从一个越来越宽的MLP,看懂压缩与路由

面对这个越来越宽的网络,高效注意力的主要思路可以归纳为两条路线。

**压缩:把宽网络变成紧凑模块(Scaling by Compression)。**线性注意力(Linear Attention)将键值信息压缩到共享线性层,Linformer将其映射为更窄的MLP,测试时训练(Test-Time Training,TTT)则学习紧凑的快权重模块。这条路线提供全局概括,但难以直接读取原始键值对的细节。

**路由:保留原始内容,只访问所需部分(Scaling by Routing)。**将键值子集组织成快权重专家(Fast-Weight Experts),再让查询按需选择。这些专家由当前上下文中的键值对构成:保留了被选中内容的细节,但缺少对完整上下文的概括。

Image 9

图1:上下文增长使快权重MLP变宽。路由选择其中一部分,压缩则构造更紧凑的表示。

这两条路线并不互斥。进一步考虑扩展策略(Scaling Strategy)、专家数量(Expert Count)、专家类型(Expert Type)、专家构造方式(Expert Construction)与路由拓扑(Routing Topology),就得到一套五维分类框架。

Image 10

表1:快权重视角下的高效注意力五维分类框架。

这套框架不仅梳理已有方法,也为新设计提出了具体问题:专家由什么信息构造?需要多少个?哪些查询可以复用同一个专家?

MiTA的切入点:Top-k集合必须每个查询各选一套吗?

以MoBA等分块混合专家注意力(Mixture-of-Experts Attention)为例,常见流程是:先把序列切成连续、不重叠的块,再为每个块生成路由向量,让查询按需选择。

**路由可以随查询变化,专家边界却是预先固定的。**在图像中,同一个目标的相关信息未必落在同一个矩形块内;为了读取少量相关内容,查询可能需要访问整个块。

Top-k注意力走向另一个方向:每个查询单独选择最相关的k个键值对,不再受固定块形状的约束。但相应地,N个查询也对应N套「私有」集合。不同查询即使关注相似内容,仍要分别完成选择和组织。

能否保留Top-k选择的灵活性,又让选择结果被多个查询复用?

MiTA将思路转了过来:不先规定专家的内容,而是先提取少量代表性查询,再由它们从完整上下文中挑选内容,构成数量可调的专家。原始查询只需选择要访问哪个专家,而不必各自构造一套全局Top-k集合。

Image 11

图2:完整注意力共享全部键值对,分块MoE共享固定块,逐查询Top-k构造私有集合,MiTA则构造可复用的可变形专家。

用少量地标查询,同时构造全局概括与细节专家

MiTA通过局部平均池化(Average Pooling)提取m个地标查询(Landmark Queries)。每个地标都探查完整序列,同时完成全局信息汇总与原始细节检索。

Image 12

图3:每个查询既访问压缩后的全局键值对,也访问路由选中的原始键值对。

**先构造可复用的Top-k专家。**第i个地标查询

Image 13

与所有键匹配,选取相关性最高的k个位置:

Image 14

这些键值对组成第i个路由专家(Routed Expert),可以跨越不同空间位置、随输入内容改变形状,不同专家也允许重叠。

随后,查询选择与自己最匹配的地标所对应的专家:

Image 15

不同查询共享候选键值集合,但仍独立计算各自的注意力权重。

**再用同一组地标构造全局共享专家。**地标查询通过交叉注意力(Cross-Attention)读取完整上下文,得到对应的地标值:

Image 16

将地标查询作为键、地标值作为值,就得到全局共享的键值对

Image 17

。

MiTA因而包含m个路由专家与1个共享专家(Shared Expert):前者提供原始细节,后者提供全局概括。

**最后,让概括与细节在同一次注意力中交互。**在每个查询选择一个路由专家,即s=1的配置下:

Image 18

两组键值对拼接后,使用统一的Softmax归一化。Top-k选择与全局汇总共享地标打分,查询到地标的打分也可复用于路由。压缩与路由不仅功能互补,还能复用计算。

以ImageNet-1K的m=k=25为例,每个查询最终访问25个压缩键值对与25个原始键值对,共50个。扫描完整序列和构造Top-k集合的工作,则由少量地标集中完成。

Image 19

算法1:先构造专家,再按路由结果对查询分组计算,通过在线Softmax(Online Softmax)合并结果。

最终的查询—专家注意力开销从

Image 20

降为

Image 21

,并通过FlashAttention与按专家分组计算加速。下面进一步看任务表现和实际运行时间。

实验结果

图像分类:接近完整注意力的基础机制

在ImageNet-1K上,首先采用相同的DeiT训练配置,比较不添加额外组件的注意力机制。

Image 22

表2:DeiT-T与DeiT-S上的分类结果。DWC表示逐通道卷积,Gate表示门控机制。

在DeiT-S上,基础MiTA的Top-1准确率为79.8%,与完整注意力的79.9%相差0.1个百分点;参数量保持不变,浮点运算量(FLOPs)从4.6G降至4.4G。同一设置下,Agent Attention为78.5%。

加入逐通道卷积(Depth-Wise Convolution,DWC)后,MiTA达到80.6%;进一步加入门控(Gating)后达到81.2%。这些视觉高效注意力常用的额外涨点技巧,MiTA诚实地展示了出来。

Image 23

表3:结合ViT-5架构设计的结果,以及与其他高效视觉Transformer的对比。

结合ViT-5的架构设计并加入DWC后,MiTA-ViT-5-S达到81.7%,与ViT-5-S的82.2%相差0.5个百分点,FLOPs则从4.7G降至4.5G。

语义分割:在更长的视觉序列上减少计算

在ADE20K语义分割(Semantic Segmentation)中,将预训练ViT的注意力替换为MiTA,再训练分割模型。设置m=k=49时,每个查询仅访问98个键值对,而完整注意力在两种输入分辨率下分别访问1,024或1,600个键值对。

Image 24

表4:不同规模ViT替换为MiTA后的ADE20K分割结果。

以ViT-S+CBT为例,FLOPs从33G降至25G,减少约24%,平均交并比(Mean Intersection over Union,mIoU)从45.8%变为44.5%。ViT-T配置的FLOPs减少约42%,mIoU下降2.6个百分点。

长序列建模:训练时间从10.6小时降至2.4小时

在长序列基准Long Range Arena(LRA)的五项任务中,MiTA的平均准确率为58.91%,完整注意力为58.19%;总训练时间从10.6小时降至2.4小时,减少约77%。

Image 25

表5:各任务报告准确率与训练吞吐量(steps/s),最后一列为平均准确率与总训练时间(小时)。

测试在单张24GB RTX 4090上进行,完整注意力基线采用FlashAttention。

压缩分支还减少了原始键值对的收集(Gather)开销:在每个查询访问条目数相同的消融实验(Ablation Study)中,相比仅路由版本,加入压缩分支将训练时间从5.5小时缩短至2.4小时,平均准确率从58.81%提升至58.91%。

Image 26

图4:三层、嵌入维度为128的Transformer在RTX 4090上的推理吞吐量(Throughput)。MiTA使用m*k≈N/4,各配置调整批大小以最大化吞吐量。

随着序列增长,推理吞吐量优势进一步扩大:在64K、256K和1024K长度下,相比完整注意力的加速比分别约为19倍、56倍和159倍。

长文档分类:兼顾16K—64K文本的整体与细节

**arXiv长文档分类(arXiv Classification)**要求识别长篇论文的类别,既需要整体理解,也需要捕捉分散在文中的细粒度线索,正好考验全局概括与细节检索的结合。

Image 27

表6:arXiv长文档分类结果(%)。

在16K、32K和64K长度下,MiTA的准确率分别达到88.0%、92.9%和98.7%,均领先于表中其他方法;相比标准注意力,分别提高18.2、3.2和1.7个百分点。

还有两个值得关注的现象

**第一个现象:没有额外剪枝模块,也会逐渐聚焦目标。**随着网络加深,选入路由专家的键值位置逐渐集中到类别相关区域,背景位置越来越少被选中,呈现出涌现的词元剪枝效应(Emergent Token-Pruning Effect)。

Image 28

图5:最左侧为输入图像,其余列展示各层选入专家的键值位置,结果在注意力头之间聚合。后层逐渐聚焦类别相关区域。

未被选中的背景词元仍可通过压缩分支参与全局汇总。

**第二个现象:权重可以复用,推理预算也可以调整。**训练后冻结模型参数,仅在推理时更换注意力机制,MiTA与完整注意力之间展现出较好的相互兼容性。

Image 29

图6:固定模型参数,训练与推理使用不同注意力机制时的分类准确率(%)。左右分别为DeiT-T与DeiT-S。

进一步将完整注意力预训练的ViT替换为MiTA,并在ImageNet-1K上微调,ViT-S和ViT-L分别达到80.9%和85.5%,与完整注意力仅差0.3和0.4个百分点。

此外,保持模型参数不变,只调整推理时的专家数量m与专家宽度k,合适的配置还能进一步提高准确率。这种算法泛化(Algorithmic Generalization)为训练时控制成本、推理时按需调整预算提供了可能。

Image 30

图7:MiTA在不同m,k下的准确率。红框标出训练配置,橙色边界圈出准确率高于原配置的推理设置。

从「Attention是什么」,走向「Attention还能怎样设计」

回到开头的问题:「Attention也是FFN」的价值,不只是两个公式可以对应,而是注意力的效率问题由此获得了一组新的设计维度。

压缩决定怎样保留全局信息,专家构造决定哪些原始内容值得被组织在一起,路由则决定哪些查询可以复用这些内容。这些在FFN中研究过的问题、尝试过的设计,还有积累下来的经验,都可以反哺注意力机制的进一步发展,甚至有一天打破Attention和FFN两者之间的绝对界限。

这项工作提供的不只是一个注意力变体,也是直指架构底层设计信念的展望。

参考资料:

https://arxiv.org/pdf/2602.01219

本文来自微信公众号“新智元”,作者:新智元,36氪经授权发布。

该文观点仅代表作者本人,36氪平台仅提供信息存储空间服务。

商业策划

+1

5

好文章,需要你的鼓励

新智元特邀作者

0收 藏+10评 论

Image 31: 36氪打开微信“扫一扫”,打开网页后点击屏幕右上角分享按钮微 博

沉浸阅读 返回顶部

举报

参与评论

评论千万条,友善第一条

登录后参与讨论

提交评论 0/1000

商业策划

你可能也喜欢这些文章

刚刚,神秘Claude Next偷跑上线,4行提示手搓3D电影级大片

曾经日入过万的网红职业,被月薪800的AI悄悄顶替

AI 给警方编了一条凶案线索,72 天没人发现

清华团队做绳驱柔性太空机器人,连融三轮数亿元,顶级资本押注

98家企业AI应用现状:规模化落地还有多远?

下架8个月,GPT-4o的拯救行动居然还在继续

AI FOMO,正让企业和员工都成为“演员”

Karpathy撕开AI折叠真相,全球20亿人被蒙骗,2000万人陷入智能体眩晕

芯片,如何被卡脖子?美国最新研判

Image 32

新智元

特邀作者

TA没有写简介,但内敛也是一种表达

发表文章 5447 篇

最近内容

刚刚,神秘Claude Next偷跑上线,4行提示手搓3D电影级大片

15分钟前

NeurIPS'26:注意力也是FFN?从快权重出发,重新理解高效注意力

31分钟前

Karpathy撕开AI折叠真相,全球20亿人被蒙骗,2000万人陷入智能体眩晕

2小时前

阅读更多内容,狠戳这里

商业策划

Image 33: 36氪

36氪AI测评

选靠谱AI,看真实评测

查看

36氪AI测评官方交流社区

加入

36氪项目推荐

咨询项目审核和入驻

联系

36氪项目推荐订阅号

关注

企服圈子

Image 34软件选型经验交流社区 加入

Image 35

行业交流群

扫码加我,拉你进对应的行业群,链接一手资源和人脉!已有170000+人深度交流中~

Image 3636氪企服点评订阅号 关注

Image 37

36氪企服-订阅号

「36氪企服」会定期分享行业深度案例、报告以及热点话题,从用户视角出发,帮企业买家选中真正“对”的软件与服务。

Image 3836氪企服点评服务号 关注

Image 39

36氪企服点评-服务号

「36氪企服点评」现已收录近万款产品,致力于帮助每个需求企业服务的人做出正确的决策,携手每个企服行业者为大众提供更高的价值与服务。

企服软件推荐

Image 40: 摄图·云编辑 摄图·云编辑 评Image 42: 金山AIDAtrans 金山AIDAtrans 评Image 44: 三维推 三维推 评

找软件,先查用户评分榜查看

顾问1对1推荐软件免费服务体验

下一篇

十一后,酒店老板开始重算渠道账

旺丁不旺财,单体酒店打算换个活法。

31分钟前

关于36氪

网络谣言信息举报入口

热门推荐

合作伙伴

  • Image 46: 阿里云
  • Image 47: 火山引擎
  • Image 48: 高德
  • Image 49: 个推
  • Image 50: 星球日报
  • Image 51: 鲸准
  • Image 52: 氪空间
  • Image 53: 富途牛牛
  • Image 54: 企服点评
  • Image 55: 人人都是产品经理
  • Image 56: 领氪

36氪APP下载

Image 57

iOS & Android

Image 58: 36氪

本站由阿里云提供计算与安全服务 违法和不良信息、未成年人保护举报电话:010-89650707 举报邮箱:jubao@36kr.com网上有害信息举报

© 2011~2026 北京多氪信息科技有限公司 |京ICP备12031756号-6|京ICP证150143号|京公网安备11010502057322号

意见反馈

Image 59

36氪APP 让一部分人先看到未来

36氪

鲸准

氪空间

Image 60 推送和解读前沿、有料的科技创投资讯

Image 61 一级市场金融信息和系统服务提供商

Image 62 聚焦全球优秀创业者,项目融资率接近97%,领跑行业

来源:36氪·AI · 36kr.com