第 1 / 8 章
MiTA：让查询复用候选，
再补上一份全局摘要。
长序列里，每个查询都对所有位置逐一打分很贵。MiTA 先用少量代表查询组织候选 key/value，再让每个查询访问共享的压缩信息 + 选中的原始细节。
这份教程对应 Wen 等的 Mixture-of-Top-k Attention: Efficient Attention via Scalable Fast Weights，固定阅读 v6（2026-09-27）。先从标准 attention 的一次查询算起，再拆开论文的方法。
N 个查询分别比对 N 个 key：N² 对。
用少量全局信息保留覆盖，再读一部分原始 KV。
Q、K、V 分别做什么
一个位置的 query（Q）表达“现在想找什么”；各位置的 key（K）用于匹配；value（V）是最终被加权取回的信息。这些向量通常由输入经过训练过的投影产生。这里直接给小向量，专心看 attention 本身。
aⱼ = exp(scoreⱼ) / Σₗ exp(scoreₗ)
o = Σⱼ aⱼvⱼ
先比对，再把分数变成总和为 1 的权重，最后混合 value。比如两个分数都是 0 时，权重各 1/2；value 分别为 0 和 2 时，输出为 1。这是加权平均，不是一定只拿分数最高的一项。
先看一次完整查询
教学数据有 8 个 token、每个向量 2 维。切换 query，观察所有位置的权重怎样改变。
序列变长一倍，所有查询与 key 的比对数变成四倍。FlashAttention 可以改善精确 attention 的显存访问与中间存储；它并不把所有点对计算改成稀疏访问。
原文定位：§3.1，标准 attention 与 fast weights。
检验一下
一个查询要给 N 个 key 打分，N 个查询一共需要多少对？
第 2 / 8 章
把同一次计算，
看成一个两层网络。
固定当前序列的 K、V，再输入一个 q：先乘 Kᵀ，经过 softmax，再乘 V。于是每组 KV 都像一个隐藏单元，序列长度 N 就是隐藏宽度。
形状比名字更有用
o = V softmax(Kᵀq/√d)
计算图：d → N → d
普通网络的矩阵参数通常在训练后固定下来，这里充当两层矩阵的 K、V 却随当前输入变化，所以被称为 fast weights。产生 Q/K/V 的投影矩阵仍是通常训练的参数，可以叫 slow weights。
一组 KV = 一个隐藏单元
key 决定 q 怎样激活这一项，value 决定这一项向输出贡献什么。序列加一个位置，就多了一组可以检索的信息。
“网络变宽”带来成本
对每个 q 都计算 N 个隐藏激活，N 个 q 合起来便有 N² 项。高效 attention 可以从压缩隐藏单元或少访问部分单元入手。
为什么不能直接叫普通 FFN
softmax 的每一项都依赖所有分数的共同分母。若把激活改写成逐元素 exp，必须同时保留归一化：
o = (V/Z(q)) exp(Kᵀq/√d)
V/Z(q) 仍依赖 q，不能把它当成对所有查询固定的普通 ReLU 前馈网络。论文的 fast-weight 视角是计算结构的重述，方便比较不同压缩和路由方式。
原文定位：§3.1，式1–4。
检验一下
这里 fast weights 的 fast 指什么？
第 3 / 8 章
只保留摘要，
还是只访问一部分原件？
压缩把很多 KV 合并成较少的表示，覆盖范围广却可能淡化局部细节。路由只让 query 读取部分原始 KV，细节直接，但候选怎么分组、怎样选择，决定了效率和灵活性。
| 思路 | 候选怎样来 | 主要取舍 |
|---|---|---|
| 压缩 | 把全库汇成少量表示 | 少量访问覆盖全局，细节可能混合 |
| 固定块路由 | 按位置预分块 | 块内可复用，位置边界限制候选形状 |
| 每 query 独立 top-k | 每个 query 单独搜自己的高分 KV | 细致，但若每次仍全量打分，候选生成也贵 |
| MiTA | 少量 landmark 搜出可复用 KV 子集，并另建共享压缩 | 多个 query 复用搜索结果，另付池化、路由与选择开销 |
“可变形”改变的是候选集合
landmark 对所有 key 打分，选出的 KV 可以分布在远近不同的位置。专家不必是连续矩形，也不是预测坐标偏移的几何采样器。landmark 自己来自局部平均池化，但其候选搜索不被限制在那个池化窗口内。
不同专家的候选可以重叠；一批 query 也可以访问同一专家。它们复用的是一组当前输入的 KV，不是把模型参数划出几套新的 FFN。
论文的五个分类维度
- 缩放策略
用压缩、路由，还是二者结合。
- 专家数量与专家类型
构造多少个模块，模块是线性层、MLP 还是其他结构。这是两个维度，原表把它们放在同一列。
- 专家构造方式
按固定位置、输入内容、landmark 探测还是其他信息构造。
- 路由拓扑
哪个 query 访问哪些专家，是否有所有 query 都访问的共享部分。
原文定位：§3.1、表1与附录B。
检验一下
MiTA 的 routed expert 主要是什么？
第 4 / 8 章
先整理一次，
再被很多查询复用。
记 N 为 token 数，m 为 landmark/路由专家数，k 为每个专家保留的原始 KV 数，s 为每个 query 选择的专家数。论文实验都取 s=1；方法公式允许更一般的 s。
- 1. 从 Q 平均池化出 m 个 landmark
每个代表查询 q̃ᵢ 是若干 Q 的平均。它不是对 K 的简单池化。
- 2. 每个 landmark 对全体 K 打分
用 q̃ᵢ·kⱼ/√d 得 N 个分数，选其中 top-k 的原始 KV，构成第 i 个可复用路由专家。
- 3. 同一套全局分数生成共享压缩
对 N 个分数做 softmax，再加权全部 V，得 ṽᵢ。共享 key 使用 q̃ᵢ；ṽᵢ 不是只对 top-k 做平均。
- 4. 当前 query 选择专家
比较 q·q̃ᵢ，从 m 个 landmark 中选 top-s。这个步骤用分数选索引，没有额外设一个常规 MoE softmax gate 去相加专家输出。
- 5. 拼接共享与原始 KV
所有 query 都访问 m 组 (q̃ᵢ,ṽᵢ)，再加入被选中 s 个专家里的 sk 组原始 KV，共 m+sk 项。
- 6. 对拼接项做一次 attention
全部 m+sk 个 logit 共同 softmax，再加权对应 value。共享和路由两部分共同竞争总权重 1。
V*=[Ṽ, V^(e₁), …, V^(eₛ)]
o=V* softmax((K*)ᵀq/√d)
为什么不能把两边输出直接相加
若分别计算分支内部的归一输出 o₁、o₂，还必须知道各分支未归一分数的总质量 Z₁、Z₂。正确合并为 (Z₁o₁+Z₂o₂)/(Z₁+Z₂)，与联合 softmax 等价；o₁+o₂ 一般不等价。
原文定位：§3.2，式5–12。
检验一下
共享部分和路由部分的输出怎样合并？
第 5 / 8 章
8 个 token，
把每一步都摊开。
选择 query、landmark 数与每个专家的 k。用“下一步”依次查看池化、候选、共享压缩、路由与最终权重。
这里用二维 2×4 网格作平均池化，单头、d=2、s=1，无投影层、位置编码、残差或训练。按论文式5–12统一用 √d 缩放；与公开 CUDA 路径的差别在最后一章说明。
也可点击下方格子切换 query。文字会标出当前 query 与本步候选，颜色仅作辅助。原始网格上排为0–3，下排为4–7；窄屏格子换行不改变池化分组。
展开完整的教学 Q / K / V
检验一下
一个 KV 没被 top-k 选中，就完全不参与 MiTA 了吗？
第 6 / 8 章
少看一些项，
不等于省掉所有准备工作。
单个 query 的最终访问量从 N 项变成 m+sk 项。但产生候选和摘要也要工作：landmark 要对全体 K 打分，query 还要对 landmark 路由。
路由 / 共享 key 分数（可复用）：Nm
路由原始 KV 分数：Nsk
合计独立相似度打分：N(2m+sk)
这个计数把同一 query–landmark 的点积复用于路由与共享 logits，省略向量维度 d。它仍不包含 value 加权求和、top-k、池化、排序、gather、投影、反向传播或通信。
调整预算，再看增长
“线性”需要固定什么
固定 d、m、k、s 时，上面主要矩阵乘法项随 N 线性增长；若 m、k 也随 N 变化，就要重新代入。选择/排序算法与内存访问也可能带来其他开销。
原文图5设置 mk≈N/4，并不是固定 m、k 的线性增长验证。若 m、k 都约随 √N 增长，N(2m+sk) 就约是 N^(3/2)，仍不同于 N²，但不能直接称线性。
原文定位：§3.2 复杂度；§4.3 图5。
检验一下
N(m+sk) 描述最终访问，能直接当完整算法耗时吗？
第 7 / 8 章
先看任务与设置，
再看那一个醒目的倍数。
下面是论文报告的数据，未在本页复现。方法算得快、最终精度怎样、无需训练直接替换是否可行，是三组不同的问题。
图像分类：纯注意力与附加组件分开
表2，ImageNet-1K、DeiT 配方从头训练，224²输入；MiTA 的 m=k=25、s=1。单位为 Top-1 百分比。
| 模型 | dense | 纯 MiTA | MiTA + DWC |
|---|---|---|---|
| DeiT-Tiny | 72.2 | 71.1 | 73.4 |
| DeiT-Small | 79.9 | 79.8 | 80.6 |
纯 MiTA 分别低 1.1、0.1 个百分点；加 DWC 后提高。Small 再加 Gate 达81.2，已经包含额外组件；这个 Gate 也不是选专家的 top-s 路由。
原文定位：§4.1，表2。
分割：计算量下降伴随精度损失
表4，以标准 attention 预训练的骨干替换后进行 ADE20K 下游训练，m=k=49、s=1；不是不训练直接换模块。
| 骨干 | dense FLOPs / mIoU | MiTA FLOPs / mIoU |
|---|---|---|
| Tiny，512² | 12G / 39.1 | 7G / 36.5 |
| Small，512² | 33G / 45.8 | 25G / 44.5 |
| Base，512² | 111G / 49.3 | 95G / 48.1 |
| Large，640² | 622G / 53.3 | 508G / 50.5 |
报告 FLOPs 下降约14%–42%，mIoU 下降1.2–2.8个百分点；不能概括成“精度无损”。
原文定位：§4.2，表4。
LRA 与合成长序列是不同实验
表5的五项 LRA 任务采用共同超参数、RTX4090 24GB，dense 用 fused/FlashAttention。平均准确率58.19→58.91，总训练时长10.6→2.4小时。并非每项任务都提高，也不是所有候选方法中训练最快。
图5另外测三层、维度128的合成长序列推理吞吐，各运行分别调 batch，mk≈N/4，硬件仍为4090 24GB：
| 序列长度 | 相对 dense 的报告吞吐倍数 |
|---|---|
| 16K | 4.7× |
| 64K | 19× |
| 256K | 56× |
| 1024K | 159× |
原文定位：§4.3，表5与图5；附录D。
检验一下
表2里加 DWC/Gate 后的提升能全归功于纯 MiTA 吗？
第 8 / 8 章
带走计算过程，
也带走核对条件的方法。
读到这里，试着不用术语解释一次：少量 landmark 先查全库，建立可复用原始 KV 候选与共享摘要；query 再路由到少量专家，最后对摘要和原始项联合归一。
三个需要主动核对的地方
1. 不训练直接替换，和微调后接近原性能
附录C的文字说保留超过95%性能，但图9中 Standard 训练→MiTA 测试，Tiny 为72.1→58.7，Small 为80.0→71.3，保留比例约81.4%、89.1%。教程按图示数值说明直接替换仍有明显损失。
表7是另一个50 epoch微调实验：Standard/MiTA分别为 Tiny 76.9/75.6、Small 81.2/80.9、Base 84.4/82.8、Large 85.9/85.5。不能把这组数称为不训练替换的结果。
原文定位：附录C，图9与表7。
2. m=k=N，不自动得到数值相同的 dense
附录A有恢复 full attention 的表述；按式10保留共享 KV 的字面计算，严格相等并不普遍成立。
q=0 时 dense 输出=0.5。
共享 values=[0.5, sigmoid(1)]，再拼原始 values=[0,1]。
四项 logits 都为0，MiTA输出≈0.5577646447 ≠ 0.5。
所以滑杆最大端仍按原公式算，不强制回退 dense。严格恢复需要移除/调整共享分支等附加条件；这个反例是算术核对，没有运行论文模型。
原文定位：附录A与正文式10–12。
3. 本页公式演示，与官方代码的共享温度
本页依论文式5–12给拼接项统一除以√d。核对的官方代码把 shared key 预先乘 d⁻¹ᐟ²，mixed attention 中又使用 d⁻¹ᐟ² 缩放，该路径共享 logit 实为 q·q̃/d；原始路由项仍为 q·k/√d。
因此这页不声称逐位复现 CUDA 模块，也不擅自认定该工程差别为 bug。代码参数 pool_size、kv_topk 在基础模块中分别平方得到论文的 m、k，读配置时也要留意。
如果要验证适不适合自己的模型
- 固定比较对象
同一任务、参数量、训练预算、精度和硬件；分清从头训练、下游微调与直接替换。
- 检查完整成本
除主 attention 外测候选生成、排序、gather、显存峰值和端到端时间；不要从打分对数推算实测速度。
- 报告精度与设置
逐任务看损失，记录 N、m、k、s、batch 与数据形态。长序列分类结果不能代替因果解码实验。
来源与本页范围
- Wen 等，Mixture-of-Top-k Attention，arXiv:2602.01219v6
固定版本2026-09-27，论文许可CC BY 4.0。教程中文解释、手工向量、界面与流程图为重新编写，未复制原论文图。
- 作者公开实现，固定提交95460f7
代码MIT许可；本页数值核心独立编写。没有下载权重、训练、评测或复现作者报告的性能。