一直以来,注意力机制和前馈网络被当成两个不同的模块。
前者让词元彼此交流,后者对每个词元做非线性变换。
把公式展开,这道界限就没那么绝对了。
标准注意力其实也能写成一层前馈网络。
区别只在权重从哪来。
前馈网络的权重是训练时学出来的,落定就固定。
注意力里的权重由当前输入现算,跟着上下文走。
一个稳定,一个随内容漂移。
这也解释了一件反直觉的事。
AI对话 里动辄几万字的上下文,压力不只来自词元数量。
我们还在运行一个不断变宽的网络。
每个键值对都占一个隐藏单元,上下文越长,这层网络越宽。
真正吃不消开销的根子在这里。

围绕这个问题,北京邮电大学的一个团队提出了一套办法。
他们把已有的快权重视角往下推,整理成一套五维分类框架。
在这套框架里,他们把按序挑选的专家做成可变形也能复用的形态。
既留住全局概括,也能按需翻出原始细节。
这项工作被一个顶级学术会议接收,代码已经开源。
项目代号叫 MiTA。
论文地址和代码仓库都公开在文里。
注意力其实也是一层前馈网络
先把两者的写法并排放在一起。
普通前馈网络用的是两层结构,隐藏层宽度记作 D。
注意力稍作变形,也能呈现出同样的骨架。
把归一化那一步吸收进权重,对应关系就清楚了。
查询对应输入,变换后的键和值对应两层权重。
指数函数对应激活函数,偏置取零。
宽度那一项换了个身份。
前馈网络的隐藏层是固定的 D,注意力里则是序列长度。
差别因此收拢到两点。



第一点是权重的来源。
前馈网络用训练得到的慢权重,落定就不动。
注意力把当前输入算出来的键和值当快权重,跟着上下文走。
第二点是宽度。
每个键值对占一个隐藏单元,上下文越长,这层网络越宽。
每个查询要访问全部隐藏单元,查询一多,开销就上去了。
高效注意力的设计问题因此被重新表述。
怎么让这个不断变宽的网络保持可扩展。
压得更窄,或者只读需要的部分
面对越来越宽的网络,现成办法大致分成两条路。
一条是压缩。
把宽网络折成紧凑的模块。
线性注意力把键值信息压进一个共享的线性层。
另一种做法把它映射成更窄的两层网络。
还有一类方法专门学一个紧凑的快权重模块。
这条路能给全局的概括。
代价是难以直接翻出某个原始键值对的细节。
另一条是路由。
原始内容全留着,只让查询去读它真正需要的那部分。
把键值对组织成一个个快权重专家,查询按需挑。
专家由当前上下文里的键值对构成。
被选中的内容保留了细节,却缺了对整段上下文的概括。
两条路都不完美。

其实它们并不互斥。
再往下拆,还有扩展策略和专家数量几个维度。
专家的类型也算一维。
专家的构造方式算一维。
路由的拓扑还是一维。
五个维度摆在一起,就得到一张分类图。
它不只是在给已有方法归档。
它把新设计该回答的问题摆到了台面上。
专家用什么信息造出来。
要造几个。
AI编程 场景里那些超长的代码上下文,该怎么分给它们。

一个具体例子是分块混合专家注意力。
它常见的流程是先切块。
序列被切成连续又不重叠的段,每段生成一个路由向量。
查询再按需挑块。
路由可以随查询变,块的边界却早早定死了。
在图像里,同一个目标的信息未必落在同一块内。
想读一点点相关内容,可能得把整块翻一遍。
另一条路是按序选前 k 个最相关的键值对。
它摆脱了固定块形状的约束。
代价是每个查询都得自己挑一套。
哪怕两个查询关注的东西很像,也得各选一遍。
能不能既留下这份灵活,又让结果被多个查询共用。
先用少量地标把专家造出来
这套方法把思路倒了过来。
不先规定专家里装什么,而是先挑出少量有代表性的查询。
让它们从完整上下文里选内容,组成数量可调的专家。
普通查询只需挑去哪个专家,不必各自造一套全局集合。

它用局部平均池化挑出若干个地标查询。
每个地标都扫一遍完整序列。
一边汇总全局信息,一边检索原始细节。
先造可复用的按序专家。
第 i 个地标查询和所有键做匹配,挑出相关性最高的 k 个位置。
这些键值对组成第 i 个路由专家。
专家可以跨不同的空间位置,形状随输入内容改变。
不同专家之间也允许重叠。



随后,查询挑一个和自己最匹配的地标,用后者对应的专家。
不同查询共用同一批候选键值,各自的注意力权重还独立算。
接着用同一组地标造一个全局共享专家。
地标查询通过交叉注意力读完整上下文,拿到对应的地标值。
把地标查询当键、地标值当值,就得到一份全局共享的键值对。
整套方法因此含若干路由专家,外加一个共享专家。
前者给原始细节,后者给全局概括。

最后让概括和细节在同一次注意力里碰头。
每个查询挑一个路由专家,两组键值对拼起来,走统一的归一化。
按序挑选和全局汇总共用同一套地标打分。
查询到地标的打分也能复用到路由上。
压缩和路由不只功能互补,还能共用计算。
以图像分类为例,设地标数和宽度都是 25。
每个查询最终读 25 个压缩键值对,再读 25 个原始键值对。
扫完整序列和造集合的活,交给少量地标集中干完。
在线归一化和按专家分组计算,把整条流水线进一步压快。

实验里跑出来的差距
图像分类是第一站。
在同一套训练配置下,先比不加额外组件的注意力机制。

在一档小模型上,基础版准确率和完整注意力只差 0.1 个百分点。
参数量没变,浮点运算量从 4.6G 降到 4.4G。
同一设置下,另一个方法的成绩落后一截。
再叠上逐通道卷积,成绩升到 80.6%。
继续加门控,能到 81.2%。
这些视觉高效注意力常用的涨点技巧,团队照样诚实摆了出来。
换一套更新的主干设计,加入逐通道卷积后可到 81.7%。
和对照版本只差 0.5 个百分点,运算量从 4.7G 降到 4.5G。

第二站是语义分割。
把预训练主干的注意力换掉,再训分割模型。
设地标数和宽度都是 49,每个查询只读 98 个键值对。
完整注意力在两种输入分辨率下,分别要读 1024 或 1600 个。

一个中等配置里,运算量从 33G 降到 25G,少了约四分之一。
分割质量从 45.8% 落到 44.5%。
换成更小的配置,运算量减少约 42%,质量掉 2.6 个百分点。
第三站长序列。
在五项任务里,这套方法的平均准确率是 58.91%。
完整注意力是 58.19%。
总训练时间从 10.6 小时降到 2.4 小时,少了约 77%。

测试在一张 24GB 消费级显卡上完成。
压缩那条支路还省下了收集原始键值对的开销。
在访问条目数相同的消融实验里,加入压缩支路后训练从 5.5 小时缩到 2.4 小时。
平均准确率同时从 58.81% 抬到 58.91%。
序列越长,推理吞吐的优势越明显。
在 64K、256K 和 1024K 长度下。
加速比分别约为 19 倍、56 倍和 159 倍。


最后一站是长文档分类。
任务要求识别长篇论文的类别。
既得看懂整体,也得抓住散落文中的细粒度线索。
这恰好考的正是概括与细节的结合。
在 16K、32K 和 64K 长度下。
准确率分别是 88.0%、92.9% 和 98.7%。
三个长度段都领先表里其他方法。
相比标准注意力,分别高出 18.2、3.2 和 1.7 个百分点。
还有两个值得留意的地方。
第一个是它不用额外的剪枝模块,自己就会逐渐聚焦。
网络越深,选进路由专家的键值位置越往类别相关区域集中。
背景位置被选中的次数越来越少。
一种自发的词元剪枝效应浮现出来。
没被选中的背景词元,仍能通过压缩支路进入全局汇总。

第二个是权重能复用,推理预算也能调。
训练完把参数冻住,只在推理时换注意力机制,两边的兼容性不错。
把完整注意力预训练的主干换成这套方法,再在图像分类上微调。
两种规模分别拿到 80.9% 和 85.5%,和完整注意力只差 0.3 和 0.4 个百分点。

参数不动,只调推理时的专家数量和宽度,合适的配置还能再涨一点。
这份泛化能力让训练时控成本、推理时按需调预算有了落点。

回到开头那个问题。
把注意力看成一眼前馈网络,价值不只在两套公式能对上。
它给注意力的效率问题一组新的设计维度。
压缩决定怎么留住全局信息。
专家构造决定哪些原始内容值得组织在一起。
路由决定哪些查询可以复用这些内容。
这些在前馈网络里研究过的问题和攒下的经验,都能反哺注意力机制。
说不定哪天,两者之间那道绝对的界限也会被打破。
说到底,这类底层方法最终都会落到我们天天在用的 AI工具 里。
它不显眼,却决定了对话能不能撑住长上下文。
