网易首页 > 网易号 > 正文 申请入驻

注意力也是FFN?长文本开销从10.6小时砍到2.4小时

0
分享至

把注意力的公式展开,会看到一条容易被忽略的联系:标准注意力也可以写成一个前馈网络(FFN),只是它的权重由当前输入生成,隐藏层宽度随序列长度增长。于是,长上下文的开销有了另一种解释——我们不只是在处理更多词元,还在运行一个不断变宽的网络。

北京邮电大学团队提出的 Mixture-of-Top-k Attention(MiTA),把已有的快权重(Fast Weights)视角发展成一套五维分类框架,并据此设计出可变形、可复用的 Top-k 专家。该工作已被 NeurIPS 2026 接收,代码已开源。


注意力为什么能写成FFN

在 Transformer 架构里,注意力负责让词元彼此交互,FFN 负责对每个词元做非线性变换,两者长期被视为功能不同且互补的模块。但把 Softmax 展开,把缩放和依赖查询的归一化吸收到权重中,对应关系就出现了:查询对应输入,变换后的键和值对应两层权重,指数函数对应激活函数,偏置为零,隐藏层宽度则由 D 变为序列长度 N。

区别在于权重的来源。普通 FFN 使用训练得到的慢权重(Slow Weights),注意力则把当前输入生成的 K、V 作为随上下文变化的快权重。更重要的区别在于宽度:每个键值对对应一个隐藏单元,上下文越长,这个快权重 MLP 就越宽。每个查询访问 N 个隐藏单元,处理 N 个查询便产生 N² 的总体开销。

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

压缩与路由,两条路线并不互斥

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

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

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

两条路线并不互斥。进一步考虑扩展策略、专家数量、专家类型、专家构造方式与路由拓扑,就得到一套五维分类框架。这套框架不仅梳理已有方法,也为新设计提出了具体问题:专家由什么信息构造?需要多少个?哪些查询可以复用同一个专家?

Top-k集合必须每个查询各选一套吗

以 MoBA 等分块混合专家注意力为例,常见流程是先把序列切成连续、不重叠的块,再为每个块生成路由向量,让查询按需选择。路由可以随查询变化,专家边界却是预先固定的。在图像中,同一个目标的相关信息未必落在同一个矩形块内;为了读取少量相关内容,查询可能需要访问整个块。

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

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

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

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

先构造可复用的 Top-k 专家。第 i 个地标查询与所有键匹配,选取相关性最高的 k 个位置,这些键值对组成第 i 个路由专家,可以跨越不同空间位置、随输入内容改变形状,不同专家也允许重叠。随后,查询选择与自己最匹配的地标所对应的专家。不同查询共享候选键值集合,但仍独立计算各自的注意力权重。

再用同一组地标构造全局共享专家。地标查询通过交叉注意力读取完整上下文,得到对应的地标值;将地标查询作为键、地标值作为值,就得到全局共享的键值对。该方案因而包含 m 个路由专家与 1 个共享专家:前者提供原始细节,后者提供全局概括。

最后,让概括与细节在同一次注意力中交互。在每个查询选择一个路由专家的配置下,两组键值对拼接后使用统一的 Softmax 归一化。Top-k 选择与全局汇总共享地标打分,查询到地标的打分也可复用于路由。压缩与路由不仅功能互补,还能复用计算。

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

算法上,先构造专家,再按路由结果对查询分组计算,通过在线 Softmax 合并结果。最终的查询—专家注意力开销从 N² 降为 Nk,并通过 FlashAttention 与按专家分组计算加速。下面进一步看任务表现和实际运行时间。

实验结果

图像分类方面,在 ImageNet-1K 上首先采用相同的 DeiT 训练配置,比较不添加额外组件的注意力机制。在 DeiT-S 上,基础版本的 Top-1 准确率为 79.8%,与完整注意力的 79.9% 相差 0.1 个百分点;参数量保持不变,浮点运算量(FLOPs)从 4.6G 降至 4.4G。同一设置下,Agent Attention 为 78.5%。

加入逐通道卷积(DWC)后,该方案达到 80.6%;进一步加入门控后达到 81.2%。这些视觉高效注意力常用的额外涨点技巧,作者也如实展示了出来。结合 ViT-5 的架构设计并加入 DWC 后,MiTA-ViT-5-S 达到 81.7%,与 ViT-5-S 的 82.2% 相差 0.5 个百分点,FLOPs 则从 4.7G 降至 4.5G。

语义分割方面,在 ADE20K 上,将预训练 ViT 的注意力替换为该方案,再训练分割模型。设置 m=k=49 时,每个查询仅访问 98 个键值对,而完整注意力在两种输入分辨率下分别访问 1,024 或 1,600 个键值对。以 ViT-S+CBT 为例,FLOPs 从 33G 降至 25G,减少约 24%,平均交并比(mIoU)从 45.8% 变为 44.5%。ViT-T 配置的 FLOPs 减少约 42%,mIoU 下降 2.6 个百分点。

长序列建模方面,在 Long Range Arena(LRA)的五项任务中,该方案的平均准确率为 58.91%,完整注意力为 58.19%;总训练时间从 10.6 小时降至 2.4 小时,减少约 77%。测试在单张 24GB RTX 4090 上进行,完整注意力基线采用 FlashAttention。

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

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

长文档分类方面,arXiv 长文档分类要求识别长篇论文的类别,既需要整体理解,也需要捕捉分散在文中的细粒度线索,正好考验全局概括与细节检索的结合。在 16K、32K 和 64K 长度下,该方案的准确率分别达到 88.0%、92.9% 和 98.7%,均领先于对比方法;相比标准注意力,分别提高 18.2、3.2 和 1.7 个百分点。

两个值得关注的现象

第一个现象:没有额外剪枝模块,也会逐渐聚焦目标。随着网络加深,选入路由专家的键值位置逐渐集中到类别相关区域,背景位置越来越少被选中,呈现出涌现的词元剪枝效应。未被选中的背景词元仍可通过压缩分支参与全局汇总。

第二个现象:权重可以复用,推理预算也可以调整。训练后冻结模型参数,仅在推理时更换注意力机制,该方案与完整注意力之间展现出较好的相互兼容性。进一步将完整注意力预训练的 ViT 替换为该方案,并在 ImageNet-1K 上微调,ViT-S 和 ViT-L 分别达到 80.9% 和 85.5%,与完整注意力仅差 0.3 和 0.4 个百分点。

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

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

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

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

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

特别声明:以上内容(如有图片或视频亦包括在内)为自媒体平台“网易号”用户上传并发布,本平台仅提供信息存储服务。

Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.

相关推荐
热点推荐
鲁尼重申拒绝追溯冠军:我们当时是第二,为啥第二要拿奖牌?

鲁尼重申拒绝追溯冠军:我们当时是第二,为啥第二要拿奖牌?

懂球帝
2026-10-12 15:35:06
消失的彭加木:一场精心策划的逃亡?隐藏在罗布泊的离奇秘密!

消失的彭加木:一场精心策划的逃亡?隐藏在罗布泊的离奇秘密!

莫地方
2026-10-12 23:08:40
局势彻底恶化了!莫迪贸然开战,印军跨境开火,美国出动F-35A支援,中印边境恐怕要生变?舆论瞬间炸锅!

局势彻底恶化了!莫迪贸然开战,印军跨境开火,美国出动F-35A支援,中印边境恐怕要生变?舆论瞬间炸锅!

谈史论今1
2026-10-10 18:54:23
广东3消息!杜锋重返宏远,胡明轩官宣担任队长,三外援首秀时间确定

广东3消息!杜锋重返宏远,胡明轩官宣担任队长,三外援首秀时间确定

多特体育说
2026-10-12 20:37:40
CBA最新消息!观看直播或收费,广东签弗格另有深意,布朗回到北京

CBA最新消息!观看直播或收费,广东签弗格另有深意,布朗回到北京

中国篮坛快讯
2026-10-12 20:00:58
午后突发利空!大盘尾盘加速下探,又出啥事了?

午后突发利空!大盘尾盘加速下探,又出啥事了?

星图金融研究院
2026-10-12 15:50:25
上甘岭坑道最禁忌的一夜:17岁女卫生员用嘴吸出战友尿液,秦基伟回忆录曝光真相

上甘岭坑道最禁忌的一夜:17岁女卫生员用嘴吸出战友尿液,秦基伟回忆录曝光真相

纪史行者
2026-10-08 05:10:03
周军任神农架林区党委书记,王桂峰任神农架林区党委副书记

周军任神农架林区党委书记,王桂峰任神农架林区党委副书记

界面新闻
2026-10-12 10:52:23
宜信CEO,直接抓了

宜信CEO,直接抓了

投资者内参
2026-10-12 13:58:44
对于今天周一A股,我只说3句话:第一,A股大概率迎更大调整行情?

对于今天周一A股,我只说3句话:第一,A股大概率迎更大调整行情?

趋势清风侠
2026-10-12 08:53:48
浙江男子花30万买新能源车自驾去西藏,半路车突发故障,安全辅助功能直接罢工;4S店:暂时没货,配件到货后为车主更换维修

浙江男子花30万买新能源车自驾去西藏,半路车突发故障,安全辅助功能直接罢工;4S店:暂时没货,配件到货后为车主更换维修

大风新闻
2026-10-11 08:45:09
侯震与二叔侯耀华正式“开撕”! 侯耀华连线曝家族内幕,侯震制作搞笑视频回应

侯震与二叔侯耀华正式“开撕”! 侯耀华连线曝家族内幕,侯震制作搞笑视频回应

裕丰娱间说
2026-10-12 18:15:48
一枚导弹干掉20名军官!俄第41集团军指挥所被端

一枚导弹干掉20名军官!俄第41集团军指挥所被端

子桑说
2026-10-12 10:33:19
局势彻底变天了!中日对峙一天天拖下去,越来越多国家终于意识到:中日若兵戎相见,不会像俄乌冲突那样给第三方牟利的机会,而是一场血战

局势彻底变天了!中日对峙一天天拖下去,越来越多国家终于意识到:中日若兵戎相见,不会像俄乌冲突那样给第三方牟利的机会,而是一场血战

林杰论事
2026-10-03 02:28:50
官宣!弗格加盟宏远,陈老板签NBA首轮秀,4外援剑指CBA总冠军!

官宣!弗格加盟宏远,陈老板签NBA首轮秀,4外援剑指CBA总冠军!

体坛卡卡说
2026-10-12 20:59:42
特朗普选举造势翻车!两度遭到当众受辱,现场爆发激烈对峙

特朗普选举造势翻车!两度遭到当众受辱,现场爆发激烈对峙

军情观察家
2026-10-12 13:23:48
赖清德,自取灭亡!双十讲话狂言台湾地位,解放军战机已进台岛空域

赖清德,自取灭亡!双十讲话狂言台湾地位,解放军战机已进台岛空域

无情有思ss
2026-10-12 22:33:00
郑钦文退出武网引热议:小米辣也退了,赛程太紧签位烂,退的漂亮

郑钦文退出武网引热议:小米辣也退了,赛程太紧签位烂,退的漂亮

球场没跑道
2026-10-12 12:00:28
狂丢17次球权!利物浦亿元水货惹众怒:浪射不传球 2队友气炸挥拳

狂丢17次球权!利物浦亿元水货惹众怒:浪射不传球 2队友气炸挥拳

我爱英超
2026-10-12 07:32:07
李娜没做到的,郑钦文做到了!夺冠是第一步,三大好消息才是重点

李娜没做到的,郑钦文做到了!夺冠是第一步,三大好消息才是重点

天马幸福的人生
2026-10-12 22:59:34
2026-10-12 23:47:00
硅屿手记
硅屿手记
有态度网友ytd
228文章数 82关注度
往期回顾 全部

科技要闻

尊界:10月底确认制动踏板支架升级方案

头条要闻

胡塞武装24小时内四袭沙特 专家:沙特存在战略误判

头条要闻

胡塞武装24小时内四袭沙特 专家:沙特存在战略误判

体育要闻

一周吃6顿披萨的胖子,也能单防姆巴佩?

娱乐要闻

张艺谋身体状况曝光 陈婷不再顾及体面

财经要闻

尊界,是江淮的救命车还是失血引擎?

汽车要闻

小鹏G9L巴黎车展全球上市 法国市场售价74990欧元起

态度原创

健康
数码
旅游
公开课
军事航空

孩子运动莫名晕倒?当心这类疾病

数码要闻

AMD发布智能体AI计算机架Helion,纯CPU版本可扩展至32,768核心

旅游要闻

游客付费入住特定酒店即可进入“中国天眼”核心区参观?中国科学院发布情况通报

公开课

李玫瑾:为什么性格比能力更重要?

军事要闻

视频:沙特方面公布利雅得机场遭袭画面

无障碍浏览 进入关怀版