不止提速，还能兼容预训练模型：聊聊MiTA高效注意力

  Transformer依靠Self-Attention(自注意力)实现强大的上下文理解能力，但它有一个绕不开的痛点：序列越长，计算量就会呈平方级暴涨。当处理高分辨率图像、超长文本时，Vanilla Self-Attention(完整自注意力)的算力开销会迅速变得难以承受。多年来研究人员提出大量高效注意力方案，但各类方法思路分散，缺少统一的理解框架。

  来自北京邮电大学的研究团队在 arXiv 预印本平台提出 Mixture‑of‑Top‑k Attention(简称MiTA)，从 Fast Weight(快速权重)的视角，将五花八门的高效注意力纳入同一套分析体系，归纳出包含 Scaling Strategy(缩放策略)、Expert Count(专家数量)、Expert Type(专家类型)、Expert Construction(专家构造)、Routing Topology(路由拓扑)的五维分类框架。它融合两类主流加速思路，在视觉与长序列任务上兼顾运行速度与模型表现。

  很多人习惯把注意力和Feed-Forward Network(前馈网络)当成完全不同的模块，而快速权重视角揭示了二者的内在联系。普通前馈网络的参数是训练结束就固定的Slow Weight(慢权重);自注意力中的Key-Value Pairs(键值对，KV)，则是根据输入实时生成的快速权重。从数学上看，完整自注意力等价于一个隐藏维度等于序列长度N的双层MLP。序列不断变长，这个虚拟网络也随之变宽，由此带来了平方级的计算负担。

  站在快速权重优化的角度，目前高效注意力大致分为两条路线。第一种是Routing(路由)，代表工作如MoBA，把上下文切分成固定大小的块充当Expert(专家)，Query(查询)只访问少数相关块。它可以保留原始信息，但采用人为划定的刚性分块，语义相关内容很容易被分割到不同块中。第二种是Compression(压缩)，代表如Agent Attention、Linformer，把全部上下文压缩成少量全局表征，计算开销低，但会丢失部分细粒度原始信息。还有一类Top‑k注意力，会让每个查询单独挑选最相关的若干键值对，摆脱固定块限制;但每个查询都维护一套专属集合，重复计算多，硬件执行效率并不理想。过去大部分方案只会选择路由或者压缩其中一条路径，很难同时兼顾全局概括和细粒度检索。

  跳出选择困境：把压缩和路由真正结合起来

  MiTA的核心创新，就是打破二者非此即彼的思路，把压缩、路由两套机制融合，构建出可形变、可复用的Top‑k专家。它引入少量Landmark Queries(地标查询)，数量远小于整体序列长度，论文里简单实现是对图像窗口做 Average Pooling(平均池化)得到。这些地标查询承担双重角色：一方面作为压缩分支，遍历全部键值对，提炼出精简的全局表征，形成Shared Expert(共享专家)，相当于整个上下文的内容摘要;另一方面作为筛选器，每一个地标查询，从全部原始键值对中筛选出匹配度最高的一部分内容，组成Routed Expert(路由专家)。

  虽然地标查询来源于局部窗口，但专家选取的键‑值对不受原始窗口空间约束，主要由地标查询与键的匹配分数决定，也就是所谓Deformable Expert(可形变专家)。实际推理时，查询会先读取共享专家获取全局信息，再通过相似度路由，匹配到提前构建好的其中一组专家(论文当前实现默认每个查询只搭配一个专家)，把两部分键‑值对合并后完成注意力计算。原本每个查询独立生成的多个私有Top‑k集合，变成数量可控、多查询可以共享复用的专家集合，大幅减少重复计算，理论上实现随序列长度近似线性的计算复杂度。论文也提到，这套算法的主要瓶颈并不在理论计算量，而在于Top‑k选取带来的随机内存访问，容易受硬件内存带宽制约，影响实际运行速度;同时论文指出，这一限制可以通过更优化的实现来缓解。

  研究还观察到一个有意思的涌现现象：Token Pruning Effect(令牌剪枝效应)。随着网络层数加深，大量原始Token不会被选入专家集合，注意力会自发过滤无关内容，聚焦任务的关键区域。MiTA还具备不错的Algorithmic Generalization(算法泛化能力)：已经用标准注意力训练好的现成模型，推理阶段可以直接替换为MiTA算子使用，精度损失在一个百分点上下;也可以训练时用小规模专家节省算力，推理阶段放大专家规模进一步提升效果，不必完整重训模型。

  不止更快：效率之外，模型泛化能力更值得重视

  在ImageNet-1K(图像分类)、ADE20K(语义分割)与Long Range Arena(LRA，长序列基准)等多项测试中，MiTA展现出不俗的综合表现。图像分类上，同等算力下它的精度优于线性注意力、Agent Attention等同类高效注意力;真正把算力开销压下来的是语义分割与长序列任务，不做重新预训练、直接替换注意力，语义分割的算力最高可降低42%，长序列基准的总训练时间减少77%。序列越长，相比标准注意力的速度优势就越突出。

  很多高效注意力只追求纸面的理论算力指标，MiTA则体现出算子迁移兼容的价值。现实落地中，大量成熟预训练模型不会重新完整训练，如果新算子只能从零开始训练，实用价值会大打折扣。良好的跨算子泛化能力，拉近了学术创新和工程落地之间的距离。当然这项工作仍有明显局限：目前全部实验集中在视觉任务，还没有在大语言模型Autoregressive Decoding(自回归解码)场景完成验证，这也是团队未来计划探索的方向。

  写在最后

  MiTA的价值，不只是提供了一个性能不错的注意力算子。更重要的是，它跳出堆砌技巧来提速的研究惯性，依靠快速权重视角搭建统一分析框架，梳理清楚过往众多高效注意力方法的固有短板。

  大模型时代，单纯依靠扩大模型规模换取效果的成本越来越高。如何在尽量少损失效果的前提下降低算力开销，是行业长期要面对的命题。MiTA给我们带来启发：不必困在“全局压缩”和“局部稀疏”二选一的局面，融合不同范式的优势，会是长序列Transformer值得继续深挖的方向。目前项目代码已经开源，期待后续研究者把这套思路拓展到大语言模型领域。