Top-Theta Attention：用补偿式阈值化稀疏化 Transformer 注意力
核心概要
该工作提出 Top-Theta（Top-j）注意力，一种免训练、在推理阶段用按注意力头校准的静态阈值替代 top-k 搜索来稀疏化注意力元素的方法，并引入 softmax 分母补偿（SDC）与 V 均值补偿（VMC）以维持精度，在 LLaMA2/LLaMA3 的问答、代码生成与长文本摘要任务上实现 3–10 倍 V-cache 行数减少、最多 10 倍注意力元素减少，精度下降不超过 1%。
Figure 1: Two variants of Top-j attention for inference at generative decoding.
· 第 3 页深度剖析
发现并验证了静态阈值可以按注意力头校准，使注意力矩阵每行保留约 k 个显著元素，从而用逐元素比较替代 top-k 搜索。 此前内容相关稀疏化主要依赖 top-k 选择，需要行内排序并带来全行依赖；该工作表明校准后的静态阈值即可近似 top-k，且阈值是模型固有属性、对数据域变化具有韧性。 在 LLaMA2-7b、LLaMA3-8B、LLaMA3-70B 上于 ARC-C、ARC-E、Hellaswag、Human-eval、LongBench 等任务评估，阈值仅需数百个校准样本，跨数据集校准（ARC-C 校准后用于 Human-eval/LongBench）仍保持可比精度。
提出两种数值补偿方法 SDC 与 VMC，在激进稀疏化下几乎完全恢复精度。 SDC 通过估计被丢弃元素的指数和并乘以 R/(R+E) 因子来近似 post-softmax 稀疏化效果；VMC 用 V 矩阵均值行乘以被丢弃注意力质量 β 来补偿缺失的 V 行贡献，并给出期望意义下的形式化证明。 在 Hellaswag 上，未补偿的 Top-j 精度明显下降，而 exp-threshold 与 exact 两种 SDC 变体大幅恢复精度，SDC 与 VMC 结合后几乎弥合与基线的差距；在 ARC-C/ARC-E 上 Top-j 本身已优于基线，补偿收益有限。
Top-j 在预填充阶段减少最多 10 倍注意力元素，在生成解码阶段减少 3–10 倍 V 行读取，精度下降不超过 1%。 相比 Top-k，Top-j 的阈值化是逐元素操作、无行依赖，因而兼容分块（tiling）与分布式推理；同时无需重训练，仅需一次性校准。 问答任务上 Top-k 与 Top-j 均比基线精度提高 0.2%–1% 并减少 2–5 倍活跃元素；Human-eval 上 post-softmax Top-j 在 8B/70B 模型上将 pass@1 保持在基线 1% 内并减少 3 倍 V 行，70B 模型以 1% pass@1 代价换取 5 倍减少；LongBench 上 V 行减少达 10 倍，Rouge-L 保持在基线 1% 内。
阈值具有分布偏移韧性，且按层、按头、按行分别校准均有益。 阈值主要与模型而非数据相关，因此可对每个模型只校准一次；实验显示前两层保持较高 k（如 512）对下游精度关键，post-softmax 稀疏化一致优于 pre-softmax。 在 Human-eval 上，用 ARC-C 校准的 post-softmax Top-j 甚至比用同数据集校准更准确；MedMCQA 上 Top-j 与 Top-k 表现相当；消融实验覆盖层、头、行三个粒度。
启示与展望
该方法面向 decoder-only Transformer 的推理阶段，适用于 LLaMA 系列所覆盖的 GQA 与 MHSA 注意力机制；阈值需针对每个模型一次性离线校准，校准集为数百样本即可，且可存储为模型附带参数（如 LLaMA-3-70B 全部阈值约 11.8 MB）。它使低内存带宽系统部署 LLM 成为可能，并兼容分块与分布式推理；作者明确指出这是有损方法，在安全敏感场景需设置防护。
作者未广泛评估实际墙钟加速，指出有意义的加速需要硬件感知的内核优化，原型 NPU 内核在特定配置下相对 IncrementalFlashAttentionV4 取得 1.17 倍加速；稀疏内核效率高度依赖硬件对稀疏访存的处理。实验限于 LLaMA 系列，模型规模、序列长度与注意力模式如何影响稀疏-精度权衡仍是开放问题。此外，阈值按目标 k 仅保证平均保留 k 个元素，实际数量会随输入波动，作者通过增大校准集缓解而非强制截断。多 k 累积校准（MKC）内存与时间开销较大，其加速仍是待研究方向。