Fast attention mechanisms: a tale of parallelism
Jingwen Liu, Hantao Yu, Clayton Sanford, Alexandr Andoni, Daniel Hsu
cs.LG
2025-09-11
哥伦比亚提出亚二次近邻注意力ANNA，仍能与MPC互相模拟；k-hop只需O(log k)层，低秩注意力则要线性深度，合成任务误差可到0。
标准注意力对长度 N 是二次的，长上下文很快扛不住。低秩近似、稀疏注意力、基于哈希的近邻注意力都在减这个开销，但少有人证明：砍掉二次扫描之后，Transformer 还保不保得住原来的表达能力。
哥伦比亚大学和 Google Research 把这条线接到 Massively Parallel Computation（MPC）上。MPC 是 MapReduce 那一类模型：很多机器、每台本地内存亚线性、一轮本地计算再一轮通信。此前工作已经证明，标准 Transformer 能模拟 MPC 算法，反向模拟却要大约 N² 台机器。这个缺口几乎必然存在：注意力评测在顺序机上被认为需要接近二次时间，一轮 MPC 若本地计算很快，机器数就得补上这个二次。
问题变成：有没有一种亚二次注意力，能更紧地刻画高效 MPC？
Approximate Nearest Neighbor Attention（ANNA）把 query 当成近邻查询、key 当成数据库点。每个 query 只给距离足够近的 key 分配权重：r-近邻必须拿到足够大的权重，距离超过 c·r 的 key 权重必须为零（以高概率）。实现走 locality-sensitive hashing：ℓ 张哈希表，每张把 key-value 累加进桶，query 只读碰撞到的桶并做平均。当近似因子 c 大于 √3、哈希质量 ρ<1/3 时，运行时间为 O(m N^{1+3ρ} log N)，空间可压到 Õ(mN)。ρ 可小到约 1/c²，于是整体接近线性。
极限 c→∞、r=0 时，ANNA 退化成 Exact-Match Attention：只聚合 key 与 query 完全相同的 token。证明 MPC 模拟时，这一特例已经够用。每个 token 扮演一台机器，逐元素的 Q/K/V 做本地计算，精确匹配负责投递消息。反向：用 MPC 模拟 ANNA 的 LSH 实现，机器数是 N^{1-δ+O(1/c²)}，在 c 较大时接近线性，不再需要 N²。
同一套 MPC 桥还给出：常数层 ANNA 能模拟常数层低秩注意力 Transformer。Reformer 那类「先哈希再在固定大小块里做标准注意力」的能力，主要来自分块前的排序；拿掉排序后，常数层甚至算不了全局平均。
理论主结论是双向模拟，深度只差常数倍，宽度只多一个任意小的 N^δ。具体任务上：
| 任务 | ANNA 深度 | 对照 |
| Match2 | 1 层、1 头、嵌入维 1 | 标准 Transformer 同样常数层可解 |
| induction heads | 常数层、宽度亚线性 | 由常数轮 MPC 转写 |
| k-hop | O(log k) 层、宽度亚线性 | 低秩注意力或 RNN 需要 Ω(k) 层，或宽度 Ω(N/k⁶) |
Match2 问每个位置 xi 是否存在 xj 使 xi+xj ≡ 0 (mod M)。k-hop 是把 induction head 的「找最近一次相同 token 并取其下一个」复合 k 次。低秩注意力在 k-hop 上已被证明需要线性深度或接近线性宽度，ANNA 走的是对数深度。
实验把不可微的 LSH 实现换成可训的单位范数 softmax 代理，再蒸馏回 ANNA。Match2 用 N=32、M=37，1 层、8 张哈希表、每表 1 个哈希函数可到零误差。induction heads 用 N=100、字母表大小 4、k=1，2 层在哈希表足够多时误差可到约 0.1。
亚二次注意力常常被当成工程近似。这篇给出一条可证的表达力边界：ANNA 仍能承接 MPC 能做的事，并且比标准注意力更贴 MPC 的机器数。对做高效注意力的人，k-hop 这类组合任务是筛选器：低秩路线在这里先天吃亏，近邻/哈希路线至少在理论上站得住。常数层 ANNA 还能模拟低秩层，低秩能表示的，ANNA 也能表示。
这是理论论文。它没有给出可直接替换 FlashAttention 的训练好的大模型。
ANNA 的定义是一组权重约束，不是单一可微算子。训练走 softmax 代理再蒸馏，作者自己说需要更原则性的直接优化。实验只在 N=32 和 N=100 的合成任务上，没有语言建模或长上下文基准。LSH 实现是随机的，模拟假设各层随机种子已知。Q/K/V 被允许是任意逐元素函数（受位宽限制），这是为了对齐 MPC 的任意本地计算，距离真实 MLP 仍有一步。条件性的 MPC 下界会传给 ANNA 的深度，图连通性、k-hop 的对数深度下界仍然成立。