北邮NeurIPS2026新作：MiTA注意力也是FFN
在传统的Transformer架构设计中，注意力模块与前馈网络长期被视作两类功能互补的独立组件：注意力机制负责建立不同词元之间的全局关联与交互，而前馈网络则针对每个单独的词元执行非线性特征变换。
但如果对标准注意力机制的公式进行展开推导，会发现一个极易被忽略的核心关联：标准注意力本质上可以被等价为一个前馈网络，只不过该网络的权重由当前输入动态生成，且隐藏层的宽度会随着输入序列的总长度同步增加。
这也为长上下文建模的计算开销提供了另一种解读视角：我们面对的不只是数量更多的词元，更是一个随着序列变长而不断扩容的前馈网络。那么，线性注意力、混合专家注意力、Top-k注意力等诸多看似不同的高效注意力技术路线，是否可以从「如何组织这个动态扩容的快权重网络」的角度进行统一理解？
针对这一问题，北京邮电大学的研究团队提出了Mixture-of-Top-k Attention（简称MiTA），他们将已有的「快权重」视角发展为一套包含五个维度的分类框架，并基于此设计出可变形、可复用的Top-k专家模块，既能够保留全局上下文信息，又可以按需读取原始的细节特征。该研究成果已被NeurIPS 2026接收，相关代码已公开开源。
相关研究资料与代码链接：
https://arxiv.org/pdf/2602.01219
https://github.com/QishuaiWen/MiTA
从公式等价的角度来看，标准注意力与前馈网络的对应关系清晰可见：查询向量对应前馈网络的输入，经过变换的键与值对应两层网络权重，指数函数则对应激活函数，且偏置项为零。两者最核心的差异在于权重的来源：普通前馈网络使用训练阶段固定得到的慢权重，而注意力机制则将当前输入生成的键、值作为随上下文动态变化的快权重。
另一个关键差异则在于网络的宽度：每个键值对对应快权重前馈网络的一个隐藏单元，上下文越长，这个网络的隐藏层就越宽。每个查询需要访问全部N个隐藏单元，当处理N个查询时，整体的计算开销会随序列长度呈平方级增长。高效注意力的设计目标，本质上就是解决如何让这个不断扩容的快权重网络保持可扩展性的问题。
针对这一问题，现有的高效注意力技术路线可以大致归纳为两类：
压缩路线：将宽网络压缩为紧凑模块。这类方法通过将键值信息压缩到共享线性层、映射为更窄的网络或者在测试阶段学习紧凑的快权重模块，来实现计算量的降低。这类方法能够提供全局的上下文概括，但难以直接读取原始键值对的细节信息。
路由路线：保留原始内容，仅访问所需子集。这类方法将键值子集组织为快权重专家，再让查询根据自身需求选择对应的专家。这类方法能够保留被选中内容的细节，但通常缺少对完整上下文的全局概括能力。
这两条路线并非互斥，进一步从扩展策略、专家数量、专家类型、专家构造方式与路由拓扑五个维度进行梳理，就可以得到一套完整的高效注意力分类框架。这套框架不仅可以帮助我们理解已有的各类高效注意力方法，也为新的模型设计提供了明确的思考方向：专家由什么信息构造？需要多少个专家？哪些查询可以复用同一个专家？
现有主流的高效注意力方法各有局限：分块混合专家注意力会将序列切分为固定的连续块，路由可以随查询变化，但专家的边界是预先固定的，在处理图像或长文本时，相关信息未必会落在同一个固定块内；而逐查询的Top-k注意力虽然可以让每个查询单独选择最相关的键值对，不受固定块的限制，但每个查询都需要单独构造一套私有集合，即使多个查询关注相似的内容，也需要重复执行选择和组织的过程。
MiTA的核心创新点就在于解决了这一矛盾：如何保留Top-k选择的灵活性，同时让选择的结果可以被多个查询复用。
MiTA的设计思路发生了反转：不再预先规定专家的内容范围，而是先提取少量代表性的地标查询，由这些地标查询从完整上下文中挑选内容，构成数量可调的专家模块。原始的查询只需要选择要访问哪个专家，而不需要各自构造一套全局的Top-k集合。
具体来说，MiTA首先通过局部平均池化提取m个地标查询，每个地标查询都会探查完整的输入序列，同时完成全局信息汇总与原始细节检索。
第一步是构造可复用的Top-k路由专家：第i个地标查询与所有的键进行匹配，选取相关性最高的k个位置，这些键值对就组成了第i个路由专家。这类专家可以跨越不同的空间位置，随输入内容动态调整形状，且不同的专家之间允许存在内容重叠。
接下来，原始的查询会选择与自身最匹配的地标所对应的专家，不同的查询可以共享候选的键值集合，但仍然会独立计算各自的注意力权重。
同时，MiTA还会利用同一组地标查询构造全局共享专家：地标查询通过交叉注意力读取完整的上下文，得到对应的地标值，将地标查询作为键、地标值作为值，就得到了全局共享的键值对。
最终的MiTA模型包含m个路由专家与1个全局共享专家：路由专家负责提供原始的细节特征，共享专家则负责提供全局的上下文概括。
在具体的计算过程中，两组键值对会被拼接在一起，使用统一的Softmax进行归一化。Top-k的选择与全局汇总可以共享地标打分的计算结果，查询到地标的打分也可以复用于路由选择，让压缩与路由的功能不仅互补，还能复用计算资源，进一步降低整体的开销。
以ImageNet-1K的实验设置为例，当m=k=25时，每个查询最终会访问25个压缩后的全局键值对与25个原始的路由键值对，总共50个键值对。而扫描完整序列和构造Top-k集合的工作，都由少量的地标查询集中完成，大幅降低了整体的计算量。
最终的查询-专家注意力开销从传统的O(N²)降低为O(mk + N)，并可以通过FlashAttention与按专家分组计算的方式进一步加速。
图像分类任务：接近完整注意力的性能表现
在ImageNet-1K图像分类任务中，研究团队采用标准的DeiT训练配置，对比了未添加额外组件的基础MiTA模型的性能。在DeiT-S架构上，基础MiTA的Top-1准确率达到79.8%，仅比完整注意力的79.9%低0.1个百分点，且参数量保持不变，浮点运算量从4.6G降至4.4G。同一设置下，同期的Agent Attention模型准确率为78.5%。
如果加入逐通道卷积的优化技巧，MiTA的准确率可以提升至80.6%，进一步加入门控机制后，准确率达到81.2%。结合ViT-5的架构设计并加入逐通道卷积后，MiTA-ViT-5-S的Top-1准确率达到81.7%，仅比ViT-5-S的82.2%低0.5个百分点，浮点运算量从4.7G降至4.5G。
语义分割任务：在更长视觉序列上减少计算
在ADE20K语义分割任务中，研究团队将预训练ViT模型的注意力模块替换为MiTA，再对分割模型进行微调。当设置m=k=49时，每个查询仅需要访问98个键值对，而完整注意力在两种常见的输入分辨率下，分别需要访问1024或1600个键值对。
以ViT-S+CBT配置为例，浮点运算量从33G降至25G，减少了约24%，平均交并比从45.8%变为44.5%。ViT-T配置的浮点运算量减少约42%，平均交并比下降2.6个百分点。
长序列建模：训练时间大幅缩减
在长序列基准Long Range Arena（LRA）的五项任务中，MiTA的平均准确率达到58.91%，而完整注意力的平均准确率为58.19%；总训练时间从10.6小时降至2.4小时，减少了约77%。测试在单张24GB RTX 4090显卡上进行，完整注意力基线采用FlashAttention优化。
压缩分支还进一步减少了原始键值对的收集开销：在每个查询访问条目数相同的消融实验中，相比仅使用路由的版本，加入压缩分支将训练时间从5.5小时缩短至2.4小时，同时平均准确率从58.81%提升至58.91%。
随着序列长度的增长，MiTA的推理吞吐量优势会进一步扩大：在64K、256K和1024K的序列长度下，相比完整注意力的加速比分别约为19倍、56倍和159倍。
长文档分类：兼顾全局理解与细节捕捉
arXiv长文档分类任务需要识别长篇学术论文的类别，既需要对全文有整体的理解，也需要捕捉分散在文中的细粒度线索，非常适合测试MiTA兼顾全局与细节的能力。
在序列长度为16K、32K和64K的设置下，MiTA的准确率分别达到88.0%、92.9%和98.7%，均领先于对比的其他方法；相比标准注意力模型，分别提升了18.2、3.2和1.7个百分点。
两项值得关注的涌现现象
第一个现象是无需额外剪枝模块的词元聚焦效应：随着网络层数的加深，选入路由专家的键值位置会逐渐集中到与类别相关的区域，背景位置被选中的概率越来越低，呈现出自发的词元剪枝效应。未被选中的背景词元仍然可以通过压缩分支参与全局的上下文汇总。
第二个现象是权重复用与推理预算可调的算法泛化性：在训练完成后冻结模型参数，仅在推理阶段更换注意力的配置，MiTA与完整注意力之间展现出较好的相互兼容性。进一步保持模型参数不变，仅调整推理时的专家数量m与专家宽度k，合适的配置还可以进一步提升模型的准确率，这为训练阶段控制计算成本、推理阶段按需调整预算提供了可行的方案。
回到最初的视角问题：「注意力也是FFN」的核心价值，并不只是两个公式可以等价对应，而是为注意力的效率优化提供了一套全新的设计维度。压缩策略决定如何保留全局信息，专家构造决定哪些原始内容值得被组织在一起，路由策略则决定哪些查询可以复用这些内容。
此前在FFN领域中研究过的各类问题、尝试过的设计方案以及积累的经验，都可以反哺注意力机制的进一步发展，甚至有一天可以打破注意力与前馈网络之间的绝对界限。这项研究不仅提供了一个高性能的注意力变体，更直指Transformer架构底层的设计信念，为未来的模型设计提供了全新的思考方向。
