「注意力也是FFN」：一个重新理解高效注意力的视角，把训练时间从10.6小时砍到2.4小时

每次把序列长度往上推，最先疼的往往是三样东西：显存、训练时间，还有跑完一批之后那笔安静变大的账单。长上下文问题的老解法们——线性注意力、混合专家注意力、Top-k注意力——各讲各的故事，一个比一个像「局部优化」。北京邮电大学团队这篇被 NeurIPS 2026 接收的工作，换了个问法：与其问「注意力怎么加速」，不如先问「注意力到底是什么」。答案有点反直觉，标准注意力也可以写成一个 FFN。

Attention 也就是 FFN

把公式摊开看，对应关系其实很整齐：查询对应输入，变换后的键和值对应两层权重，指数函数就是激活函数，偏置为零。区别在两处。一是权重从哪来，普通 FFN 用训练出来的慢权重，注意力则把当前输入生成的 K、V 当作随上下文变化的快权重。二更关键，是宽度：每个键值对对应一个隐藏单元，序列越长，这个「快权重 MLP」就越宽，访问 N 个隐藏单元、处理 N 个查询，就是 N 平方的开销。这么一讲，长上下文为什么贵就不是玄学了——就是这个网络在变宽。

高效注意力的各种路线，于是都归到同一个问题：怎么让越来越宽的快权重网络保持可扩展。作者把它分成两条主路。压缩，把宽网络弄成紧凑模块，线性注意力、Linformer、TTT 都属于这类，给全局概括但读不到原始键值细节；路由，保留原始内容只访问需要的部分，把键值子集组织成快权重专家，细节在但缺全局概括。再加扩展策略、专家数量、专家类型、构造方式、路由拓扑五个维度，就拼成一套分类框架。这套框架真正的用处，是让「接下来还能怎么设计」有了具体的提问点：专家由什么构造、要几个、哪些查询能复用同一个专家。

从越来越宽的 MLP 到可复用的专家

MiTA 就卡在「复用」这个点上。逐查询 Top-k 的毛病是，N 个查询各挑各的 k 个键值对，等于 N 套私有集合，关注相似内容的查询还得各自重选一遍。MiTA 把流程反过来：先用局部平均池化提取 m 个地标查询，让它们去扫描完整序列、各自挑出最相关的 k 个键值对，组成 m 个可复用、可变形、允许重叠的路由专家；再让真正的查询去选和自己最匹配的专家，最后拼接全局共享专家，用统一的 Softmax 归一化。压缩和路由不再是两个方向，而是共用同一批地标打分，还能互相复用计算。

数字说结果。长序列基准 LRA 上，MiTA 平均准确率 58.91%，完整注意力 58.19%，总训练时间从 10.6 小时砍到 2.4 小时，少了约 77%，跑在单张 24GB 的 RTX 4090 上。推理吞吐量的加速比随序列变长越拉越大，64K、256K、1024K 长度下相对完整注意力分别约 19 倍、56 倍、159 倍。arXiv 长文档分类里它兼顾整体和细节，16K/32K/64K 下准确率 88.0%、92.9%、98.7%，比标准注意力分别高 18.2、3.2、1.7 个百分点。图像这边，DeiT-S 基础配置 79.8% 对完整注意力的 79.9%，只差 0.1 个百分点，FLOPs 从 4.6G 降到 4.4G；加上 DWC、门控这些视觉高效注意力的常见涨点，能到 80.6% 和 81.2%，作者把这些技巧有没有效都摊开摆了出来。

还有两个现象值得单独说。一是没有额外剪枝模块，随着网络加深，选进路由专家的键值位置会逐渐聚到类别相关区域，背景位置越来越少被选中，作者管这叫涌现的词元剪枝效应，没被选中的背景词元还能靠压缩分支参与全局汇总。二是算法泛化：模型参数冻结，只在推理时换注意力机制，MiTA 和完整注意力能相互替换兼容；保持参数不变只调推理时的专家数量和宽度，准确率还能再提一点——这等于把「训练时控制成本、推理时按需调整预算」变成了现实，对落地挺友好。

几句保留

边界也先划清楚。语义分割任务里 mIoU 有回落，ViT-S+CBT 从 45.8% 到 44.5%，FLOPs 减了约 24%；ViT-T 减 42% 也对应 mIoU 掉了 2.6 个百分点——省算力不是白省，准确率得拿一部分去换。目前看到的都是论文和代码里的说法，能核对的都在上面。要动手的话，建议直接跑仓库里那套配置，再拿自己的长文档去试试那 77% 的训练时间和 159 倍推理吞吐量——能兑现多少，得亲手验过才算数。

相关链接
- 论文：https://arxiv.org/pdf/2602.01219
- 代码仓库：https://github.com/QishuaiWen/MiTA