新智元报道  

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

前者让词元（Token）彼此交互，后者对每个词元进行非线性变换。

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

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

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

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

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

令查询，键和值按列组成。

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

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

对应关系随即出现：

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

先构造可复用的Top-k专家。第i个地标查询与所有键匹配，选取相关性最高的k个位置：

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

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

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

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

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

MiTA因而包含m个路由专家与1个共享专家（Shared Expert）：前者提供原始细节，后者提供全局概括。

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

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

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

算法1：先构造专家，再按路由结果对查询分组计算，通过在线Softmax（Online Softmax）合并结果。

最终的查询—专家注意力开销从降为，并通过FlashAttention与按专家分组计算加速。下面进一步看任务表现和实际运行时间。

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

表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诚实地展示了出来。

表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个键值对。

表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个百分点。

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

表5：各任务报告准确率与训练吞吐量（steps/s），最后一列为平均准确率与总训练时间（小时）。

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

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

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

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

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

表6：arXiv长文档分类结果（%）。

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

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

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

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

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

图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）为训练时控制成本、推理时按需调整预算提供了可能。

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

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

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

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

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

参考资料：

https://arxiv.org/pdf/2602.01219

编辑：LRST