前言#

在 Transformer 中，Attention 负责从上下文中读取信息，FFN 则用训练好的网络处理每个位置的表示。它们看起来是两种不同的操作，却可以放到同一个视角下理解：Attention 也是一个网络，只是它的部分权重由当前上下文临时构造。

这个临时网络有一个特别之处：上下文越长，网络就越宽。既然如此，能否借鉴 MoE 的思路，让每次计算只使用其中一部分？

Mixture-of-Top-k Attention，简称 MiTA，给出了一种做法：让少量代表先搜索上下文，把选出的候选共享给大量查询，再用全局摘要补充候选之外的信息。理解这篇论文，关键在于看清楚：哪些工作可以共享，哪些信息值得保留原样，哪些信息适合压缩。

接下来以一张“鸟停在树上”的图片为例，先解释注意力在读取什么，再沿着一次信息读取的过程理解 MiTA。

注意力的检索过程#

视觉 Transformer 会把图片切成许多小块，每块对应一个 token。假设模型正在更新鸟头所在位置的表示，它可能需要读取鸟身、翅膀等其他位置的信息，才能更好地识别这个物体。

注意力为每个位置生成三种向量，分别承担不同职责：

Query，查询向量：这个位置想读取什么信息。
Key，匹配向量：让其他位置判断，这里是否有它们需要的信息。
Value，内容向量：匹配后实际读取的内容。

这三个向量都由当前位置的表示经过可学习的线性投影得到，通常记作 Q、K、V。它们是数值特征，下面借助“鸟头”“翅膀”等语义来解释它们的作用。

鸟头位置的 query 会与各个位置的 key 比较，得到一组匹配分数。softmax 把这些分数转换成总和为 1 的权重，再按权重汇总各位置的 value。相关性越高的位置，通常对输出贡献越大。

下面把“比较 key”和“读取 value”分开画出来。它们分别回答去哪里找，以及找到后读到什么。

一个位置用 query 匹配其他位置的 key，再加权读取 value

对每个 query，完整注意力都会访问全部位置。假设图片切成 100 块，每块都要与 100 个 key 比较，整张图片就是 10000 次匹配。分辨率提高后，小块越来越多，成对比较的数量会增长得很快。

问题由此出现：很多 query 最后重点读取的，可能只是其中少数位置，但完整注意力仍为每个 query 检查了全部候选。

快权重视角#

回头看这次检索，其实可以把它画成一个两层网络。第一层用 key 判断输入 query 会激活哪些位置，第二层用相应的 value 汇总出结果。每个 KV 对，都对应这个网络中的一个中间节点。

这与 FFN 的计算骨架相似：输入先经过第一层变换，产生中间响应，再经过第二层变换得到输出。区别在于，FFN 使用训练得到的权重，而这里的 key 和 value 来自当前上下文。换一张图片，这个网络的权重也会改变。

这就是快权重（fast weights）的含义。“快”描述的是权重随输入变化；生成 Q、K、V 的投影参数仍然通过训练学习。下面对比两种网络的构成。

FFN 的连线由训练参数决定，Attention 的连线由当前上下文中的 K、V 决定，每个 KV 对对应一个中间节点

一个图片块提供一个 KV 对，也就提供一个中间节点。图片块越多，这个网络越宽。于是，高效注意力可以转化成一个网络设计问题：怎样减少每次需要访问的中间节点，同时保留足够的信息？

一种办法是从原始节点中选择一部分，另一种办法是把许多节点汇总成少量新节点。论文把它们分别称为路由和压缩，MiTA 会把两者结合起来。

MiTA 的共享选择#

搜索成本#

先考虑最直接的选择办法：让每个 query 只读取匹配分数最高的几个 KV。“最高的几个”就是 Top-k，其中 k 表示保留的数量。

但只保留少数结果，不代表搜索便宜。如果仍先让每个 query 与全部 key 比较，再取 Top-k，前面的 10000 次匹配一次也没省。

MiTA 想共享的，正是这部分搜索工作。比如，鸟头与翅膀位置的 query 可能都需要鸟身的信息，它们需要的候选存在重叠，没有必要各自从头扫描整张图片。可以先由少量代表准备几组候选，再让不同 query 选择合适的一组。

地标查询#

MiTA 把这些代表称为地标查询（landmark queries）。论文默认把图像上的 query 分成若干局部区域，每个区域内的 query 向量取平均，得到一个地标。

平均后的向量仍然可以像普通 query 一样，与 key 计算匹配分数。它代表这一片区域的查询特征，但只需要由它扫描一次全局，就能为后续的多个 query 准备候选。

假设一张图片有 100 个位置，生成 4 个地标。每个地标与全部 100 个 key 比较，再保留最匹配的 8 个位置，就得到 4 组候选 KV。每组候选，就是一个路由专家。

这里的“专家”是一组从当前输入中选出的 KV。每组候选的数量相同，包含哪些位置则由匹配结果决定，会随图片内容改变。

地标来自局部区域，搜索范围却覆盖整张图片。因此，一个地标选出的候选可以包含相隔很远的鸟头、鸟身和翅膀，不必是连续的图片块。

查询路由#

候选准备好后，每个原始 query 先与这 4 个地标比较，找到最匹配的一个，再读取那个地标选出的 8 个 KV。这一步就是路由。

例如，鸟头和翅膀位置的 query 可能选择同一个专家，于是复用同一组候选。但鸟头可能更关注鸟身，翅膀可能更关注另一侧翅膀，所以它们还要分别与候选 key 比较，计算各自的权重和输出。

MiTA 共享的是候选名单，每个 query 如何使用这份名单，仍然由它自己决定。

这也是近似发生的地方：地标选出的 Top-k，不一定恰好是每个 query 单独搜索会得到的 Top-k。MiTA 用少量可复用的候选集合换取更低的搜索成本，同时保留 query 在集合内部精细区分的能力。

全局信息的补充#

原始细节与全局摘要#

只读取一组候选，还有一个问题：如果某条重要信息没有被选中，query 就无法直接读取它。

比如，某个专家选中的位置主要在鸟身上，而树枝与背景信息可能有助于判断鸟的姿态或所处环境。把所有候选都加回来，就又回到了完整注意力。

MiTA 的办法是再准备一份紧凑的全局信息。地标已经与全部 key 比较过，除了取出匹配最高的几个，还可以把完整分数经过 softmax，用来汇总全部 value。每个地标因此得到一个摘要向量，4 个地标就产生 4 个摘要。

摘要是有侧重的加权汇总：不同地标关注的内容不同，形成的摘要也不同。它让未进入 Top-k 的位置仍有机会影响输出，但许多位置的信息会混合在少量向量里，无法像原始 KV 那样逐条区分。

下图沿用 100 个位置、4 个地标、每组 8 个候选的例子，展开两条路径的内部计算。从中间的分数表往下看：左边按分数选出原始 KV，右边把完整分数转成权重，汇总全部 value。

同一份地标匹配分数用于 Top-k 选择和全局加权汇总，再由当前 query 读取一组原始 KV 与全部摘要

路由路径提供少量原始细节，压缩路径提供全局摘要。MiTA 让每个 query 同时使用这两种信息，这组所有 query 都能访问的摘要，在论文中称为共享专家。

两条路径的协作#

对刚才的鸟头 query，它最终读取的是：选中专家里的 8 个原始 KV，以及 4 个共享摘要，共 12 个条目。

读取摘要也需要 key 来匹配。MiTA 直接用生成摘要的地标充当 key，摘要本身充当 value。这样，query 与地标的比较结果既能用来选择路由专家，也能决定它对不同摘要的关注程度。

最后，12 个条目一起参与 softmax。原始细节与全局摘要在同一组权重里竞争，再加权形成输出。它们的贡献取决于当前 query 与各个 key 的匹配程度。

为什么同一位置的信息可能既出现在原始候选里，又进入摘要？

因为两条路径保留的是不同粒度的信息。原始路径允许 query 单独读取某个位置，摘要路径则把它放进更大的上下文中。这两种表示可以同时参与最终计算。

一次信息读取#

现在把前面的步骤连起来。仍以 100 个图片块、4 个地标、每组 8 个候选为例：

生成代表：对 query 做局部平均池化，得到 4 个地标。
准备信息：每个地标扫描全部 key，一方面挑出 8 个原始 KV，另一方面汇总全部 value，形成一个摘要。
选择候选：每个 query 比较 4 个地标，选择其中一组原始 KV。
读取信息：query 在那 8 个原始条目和全部 4 个摘要之间分配注意力，得到自己的输出。

下图中，两个 query 选择同一组候选，却形成不同的读取权重。地标搜索发生在上游，搜索结果则在下游反复使用。

地标共享搜索得到候选与摘要，两个 query 复用同一组候选，各自分配注意力并得到输出

与完整注意力相比，变化主要发生在“给 query 准备什么信息”这一步。最后的读取仍然是根据 query 匹配 key，再对 value 加权汇总。MiTA 用共享搜索和分层的信息表示，缩小了每次读取的范围。

效率与效果#

节省的工作#

用这个例子比较主要的 query-key 匹配次数，就能看到开销去了哪里：

工作	完整注意力	MiTA
地标扫描全部位置	无	4 个地标各比较 100 次，共 400 次
原始 query 匹配地标	无	100 个 query 各比较 4 次，共 400 次
原始 query 匹配原始 key	100 个 query 各比较 100 次，共 10000 次	100 个 query 各比较 8 次，共 800 次

MiTA 在这个示意配置中把主要匹配次数从 10000 次变为 1600 次。地标扫描得到的分数同时用于选候选和做摘要，query 匹配地标的分数同时用于路由和读取摘要，两处都复用了计算。

这些数字用来说明节省发生在哪里。实际运算还包括 value 的加权汇总、Top-k 选择和数据收集，所以不能把匹配次数的比例直接当作运行速度的比例。

更一般地，设位置数为 $N$ 、地标数为 $m$ 、每组候选数为 $k$ 。完整注意力要做 $N^2$ 次匹配；每个 query 选择一个路由专家时，MiTA 的主要匹配次数是 $2Nm+Nk$ 。固定地标与候选数量后，这部分开销随位置数线性增长。

官方实现还会把选择同一专家的 query 放在一起计算，以便复用那组 KV。反过来，候选可能分散在不同内存位置，收集这些数据也会产生开销。因此，实际加速既依赖算法减少了多少工作，也依赖硬件能否高效读取这些候选。

实验证据#

这篇论文的主要实验在视觉任务上完成。先看一个图像分类消融：作者在同一个小型视觉 Transformer DeiT-T 中，分别使用两条路径，比较 ImageNet-1K 分类准确率。数值表示模型把图片类别判断正确的比例。

注意力方式	分类准确率
完整注意力	72.2%
仅压缩	70.3%
仅路由	70.8%
MiTA，压缩与路由结合	71.1%

组合两条路径优于单独使用任一条，支持了“原始细节与全局摘要互补”的设计。但这组结果仍低于完整注意力，说明共享候选与压缩确实存在信息取舍。在稍大的 Small 模型上，MiTA 与完整注意力分别为 79.8% 和 79.9%，差距更小。

长序列方面，作者在 RTX 4090 上测试了一组长序列理解任务 LRA：

注意力方式	平均准确率	总训练时间
完整注意力	58.19%	10.6 小时
MiTA	58.91%	2.4 小时

总训练时间减少约 77%，平均准确率略有提高。共享搜索与压缩在这组任务中带来了实际加速，同时保留了接近的任务效果。

这些结果支持的是视觉与 LRA 任务中的效果。论文尚未评估 LLM 的自回归解码；用于语言模型生成时，还需要让地标、摘要和候选选择都遵守不能读取未来 token 的因果约束。

实现与验证#

把例子缩小到 8 个位置、2 个地标，每组保留 2 个原始 KV，看看共享候选之后，query 实际怎样读取信息。运行后，地标 1 选中位置 1、3，地标 2 选中位置 2、4。

query 1 和 query 2 都选择地标 1，因此读取的条目完全相同：原始 KV 1、KV 3，以及两个全局摘要。它们分配的注意力权重如下：

读取条目	query 1 的权重	query 2 的权重
原始 KV 1	66.64%	55.70%
原始 KV 3	16.20%	27.46%
摘要 1	16.20%	13.54%
摘要 2	0.96%	3.29%

两个 query 都主要读取位置 1，但位置 3 在 query 2 中的权重提高到了 27.46%。复用候选名单之后，每个 query 仍然可以有自己的读取侧重。两个摘要也参与了这次读取，四个条目在同一次 softmax 中共同分配注意力。

展开查看实现与数值对拍

实验运行在 RTX 3070 Laptop 上，使用 PyTorch 2.9.1、CUDA 12.6。把候选直接拼接后调用标准注意力作为对照，FP32 前向与输入梯度的最大绝对误差均小于 $7\times10^{-7}$ 。测试还覆盖多头、不同候选数量、不等长池化和空专家。

输入 q、k、v 的形状均为 [B,H,N,d]，依次表示 batch、注意力头、位置数和头维度。landmarks 是池化得到的地标，形状为 [B,H,m,d]，topk 是每个地标选择的候选数。

先看实现中的四组操作：probe 让地标扫描全部 key；indices 保存候选位置，shared_v 保存摘要；route 为每个 query 选择专家；最后用同一个 softmax 为摘要与原始候选分配权重。

1
import torch
2

3

4
def mita(q, k, v, landmarks, topk):
5
    # q,k,v: [B,H,N,d]; landmarks: [B,H,m,d]
6
    b, h, n, d = q.shape
7
    m = landmarks.size(-2)
8
    assert k.shape == q.shape and v.shape == q.shape
9
    assert landmarks.shape[:2] == (b, h)
10
    assert landmarks.size(-1) == d and 1 <= topk <= n
11
    scale = d ** -0.5
12

13
    probe = landmarks @ k.transpose(-2, -1) * scale
14
    indices = probe.topk(topk, dim=-1).indices   # [B,H,m,k]
15
    shared_v = probe.softmax(dim=-1) @ v        # [B,H,m,d]
16

17
    route_scores = q @ landmarks.transpose(-2, -1) * scale
18
    route = route_scores.argmax(dim=-1)        # [B,H,N]
19
    chosen = torch.gather(
20
        indices, 2, route[..., None].expand(b, h, n, topk)
21
    )                                         # [B,H,N,k]
22
    # Each query gets the original KV pairs selected by its landmark.
23
    gather_idx = chosen[..., None].expand(b, h, n, topk, d)
24
    routed_k = torch.gather(
25
        k.unsqueeze(2).expand(b, h, n, n, d), 3, gather_idx
26
    )
27
    routed_v = torch.gather(
28
        v.unsqueeze(2).expand(b, h, n, n, d), 3, gather_idx
29
    )
30
    routed_scores = (q.unsqueeze(-2) * routed_k).sum(-1) * scale
31

32
    # One softmax over the shared m entries AND the routed k entries.
33
    scores = torch.cat((route_scores, routed_scores), dim=-1)
34
    weights = scores.softmax(dim=-1)
35
    out = weights[..., :m] @ shared_v
36
    out = out + (weights[..., m:, None] * routed_v).sum(-2)
37
    return out, (indices, route, shared_v, scores)

这是一份方便阅读的实现，会为每个 query 收集候选。官方实现按专家组织 query，并分别计算分支后保留归一化信息再合并，以减少数据复制。两种方式遵循相同的联合 softmax 规则；分别归一化后直接相加会改变结果。

完整实验文件见 mita-reference.py，运行 python mita-reference.py --device cuda 即可复现。

结语#

MiTA 把高效注意力拆成了两个相互配合的决定：用少量地标共享候选搜索，用原始细节与全局摘要共同服务每次查询。

地标替大量 query 扫描上下文，query 则保留在候选内部重新分配注意力的能力。路由缩小了直接读取的范围，压缩补上了范围之外的全局信息。从快权重视角看，这就是让一个很宽的临时网络，只为每次查询启用少量原始节点和一组紧凑摘要。

How MiTA Works

前言#