全注意力的反击：数百步训练即可将全注意力迁移为稀疏注意力 Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps
RTPurbo 用低维索引与动态 top-p，仅数百步训练把全注意力改造为近无损稀疏注意力。
前置知识
稀疏注意力 (Sparse Attention)
标准注意力对所有 token 对做 softmax，计算量随序列长度二次增长。稀疏注意力只让每个 query 关注少量 key（如滑动窗口、top-k 选出的块），把复杂度降到接近线性，是长上下文推理降本的关键手段，但选错 token 会丢信息。
本文的核心命题是全注意力模型本身就内含稀疏结构，能否最小代价地把它显化，决定了整个方法的存在价值，理解稀疏注意力的得失才能看懂它的贡献。
检索头与流式头 (Retrieval vs Streaming Heads)
预训练 LLM 中注意力头分工不同：少数'检索头'会把注意力放到远处语义相关的 token 上负责远距离召回；多数'流式/局部头'只看局部窗口和 attention sink。这种头级特化是 input-agnostic 且稳定的，可通过离线校准测出每个头的'检索得分'。
RTPurbo 的整套设计都建立在'只有少数检索头需要长上下文'这一观察上，头划分是方法的第一块基石，不懂这个无法理解为什么要按头区别对待。
旋转位置编码 RoPE
RoPE 把位置信息通过旋转矩阵注入到 query/key 上，attention 得分只依赖相对偏移 $\Delta=m-n$，由不同频率 $\theta_i$ 的三角项叠加而成。高频项随距离剧烈变化、对远距离检索有害；低频项平滑、保留检索信号，说明检索可压缩到低维子空间。
论文用 RoPE 的频率结构论证'可用 16 维投影近似检索相关度'，这是低维索引器的理论根基，不读 RoPE 部分会觉得 16 维投影是凭空设计。
top-k 与 top-p 选择
稀疏注意力需决定保留多少 token：top-k 固定保留得分最高的 k 个，预算恒定；top-p（也叫 nucleus）保留累积注意力质量达到 p（如 0.9）的最少 token，预算随 query 分布动态变化。检索头所需 token 数因 query 而异，固定 k 会顾此失彼。
动态 top-p 是 RTPurbo 相对静态 top-k 的关键创新，理解两者差异才能看懂论文为何强调'query-dependent'并设计排序无关的 GPU kernel。
自蒸馏 (Self-Distillation)
用原始稠密模型作为 teacher，监督稀疏化后的 student 模型对齐其下一 token 预测（这里只对齐 top-10 logits 的 KL 散度）。相比监督微调，它绕过了特定数据集分布的干扰，能在极小数据量下恢复能力，避免过拟合。
论文声称'仅需约 1M 标签 token'就能近无损恢复，自蒸馏是达成这一点的核心训练手段，不懂它会低估方法的轻量性。
研究动机
长上下文 LLM 推理被全注意力的二次复杂度卡住：1M token 时单层注意力就是吞吐瓶颈。现有替代各有痛点——DeepSeek Sparse Attention、Kimi Delta Attention 需昂贵原生稀疏预训练，难以复用已有全注意力模型；SnapKV、MInference、FlexPrefill、Quest 等启发式 token 驱逐靠'最近 query 估计全局'或'相邻块近似'，在多跳检索任务上崩得很惨。具体证据：RULER 64K 上 MInference 的 multi-K 只剩 40.07、FlexPrefill 的 multi-V 掉到 55.50，而全注意力是 99.66；SnapKV 在 vcsum 从 23.77 暴跌到 11.03。更根本的是固定 top-k 预算与 query 相关需求错配——同一检索头里'大海捞针'查询只需 2 个 token 就覆盖 96.6% 注意力质量，而语义检索'Galápagos'却需要约 8k 个 token。
本文的目标是作者想回答一个更尖锐的问题：把一个已训练好的全注意力模型改造成高稀疏模型，到底需要多大的'手术'？目标是证明无需原生稀疏预训练、只需极轻量的持续训练（几百步、约 1M 标签 token），就能在保留近无损精度的同时，把 prefill 在 1M 上下文加速 9.36×、decode 加速 2.01×。具体要同时解决三个难题：用一个鲁棒指标挑出真正需要长上下文的检索头；用一个轻量选择器高效定位必要 token；用一种自适应稀疏预算适配不同 query 的需求量。最终交付一个'强可解释、训练廉价、精度无损'的长上下文推理方案，让全注意力训练仍是有竞争力且实用的选择。
与已有工作不同的是，本文的独特切入角度是把'LLM 可解释性'和'理论分析'当作设计依据，而不是再发明一种黑盒 token 打分。它先用离线校准量化每个头的检索得分，做出稳定的 retrieval/local 头划分；再从 RoPE 的频率几何推导出'远距离检索受控于低维子空间'，从而敢用 16 维投影逼近检索相关度（实测 16 维就能 >90% 召回）；最后指出 query-dependent 这一被既有 top-k 方法忽视的本质，主张用动态 top-p 取代固定预算。这套组合让它在'训练成本/精度/效率'三角上同时占优，区别于既有方法的单点优化。
核心方法
直觉是：全注意力模型本身已经'内含稀疏'，我们要做的是把它显化而非重训。技术路线分四块。先离线校准：在长文档首尾各插一段相同'针'，测量从后针指向前针的注意力质量得到每头检索得分，据此把头划成约 15% 的检索集 $H_{ret}$ 和 85% 的局部集 $H_{loc}$，这个划分只需一条长序列、跑一次。然后是头级稀疏模式：prefill 阶段检索头做稠密注意力建完整 KV cache，局部头只用 4 个 sink + 8192 窗口的滑动注意力；decode 阶段局部头保持 SWA+Sink，检索头切换到 query 感知的动态稀疏。核心是给检索头配一个 16 维的低维投影索引器，用 pre-RoPE 特征打分，再按 top-p=0.9 动态选 token。最后是轻量两阶段训练 + 自定义 GPU kernel，把打分与选择融合成单次 kernel 启动。
核心创新是把三个独立洞察拧成一个自洽系统：(1) 头级特化作为天然先验——只给少数检索头保留完整 KV cache，其余头本来就是局部的，安全丢弃远端 token；(2) RoPE 诱导的可压缩几何——高频分量对远距离检索无用甚至干扰，所以检索可用 $r \ll d_h$ 的低维子空间近似，作者用 $W_h^Q, W_h^K \in \mathbb{R}^{r \times d_h}$（$r=16$）投影 pre-RoPE 特征来路由，而最终注意力仍在完整 128 维上算，保住精确位置几何；(3) 动态阈值——检索头要保留的 token 数本质上随 query 变化，固定 top-k 必然顾此失彼，因此用查询感知的 top-p 自适应。三者的本质区别在于：路由用低维近似、生成用全维精确，且预算按 query 逐头动态决定，这是已有方法（MInference 靠离线 pattern、FlexPrefill 靠相邻块、Quest 靠块级 min-max）都没有同时做到的。
方法步骤详情
流程分校准、推理、训练三段。校准：在 FineWeb 长文首尾插相同 needle span，按 $R_h = \frac{1}{|N_{post}|}\sum_{t \in N_{post}}\sum_{j \in N_{pre}} A_h(t,j)$ 算每头检索得分，取 top-15% 进 $H_{ret}$，跑一次即可。推理：局部头用 SWA(8192)+Sink(4)；检索头 prefill 稠密、decode 先算低维分 $s_h = (W_h^Q q^{pre}_{m,h})^\top (W_h^K k^{pre}_{n,h})$，据此构造 $S_h(m) = \mathrm{Top}\text{-}P(s_h, 0.9)$ 再在选中 token 上算精确 softmax 注意力。训练两阶段：阶段一冻结骨干只训低维投影（约 840K 参数），最小化 $L_{proj} = \sum_{h \in H_{ret}} \mathrm{KL}(a^{full}_h \| a^{proj}_h)$；阶段二插入投影、切稀疏模式，用 top-10 logits 自蒸馏 $L_{distill} = \mathrm{KL}(\mathrm{softmax}(z^{dense}_{(10)}) \| \mathrm{softmax}(z^{sparse}_{(10)}))$，约 600 步、1.2M 标签 token。
技术新颖性
技术新颖性集中在四点。第一，把可解释性的'检索头/流式头'发现从诊断工具升级为'训练目标的先验'，据此直接剪枝局部头的远端 KV，这是头级稀疏的根本依据。第二，低维投影的'路由用低维、生成用全维'解耦——投影只负责高效挑 token，最终注意力仍用 128 维 $Q,K,V$ 保留精确相对位置，这区别于 DSA/FASA 直接压缩注意力空间的做法，是'近无损'的关键。第三，动态 top-p 取代静态 top-k，论文用 Table 1 给出硬证据：top-16k 比 top-p(0.9) 只多召回 3.8% 注意力质量却要多算约 8k 个 token。第四，自定义 GPU kernel 用 256-bin 直方图 + 原子计数器实现'排序无关的 top-p'，把打分与选择融合成单次 kernel 启动，$O(1)$ 显存开销，绕开了 $O(N_b)$ 的二分搜索显存瓶颈；decode kernel 用无共享内存的单 warp CTA + 2-token 展开 + 半精度 half2 向量加载，最大化 SM 并发与访存重叠。
实验结果
核心发现:'数百步轻量训练即可近无损稀疏化'。精度上 LongBench 16 任务平均 RTPurbo(top-p) 达 54.24，反超全注意力 53.80，稀疏最高；RULER 64K 得 85.49 逼近全注意力 86.23，baseline 集体崩盘（MInference 65.61、Quest 70.60）。超长 multi-hop 上 RTPurbo 在 512K 仍稳住 92.8–98.4，sparsity 推到 97%。AIME24/25 均 86.67 与稠密基线持平，MMLU-PRO 各科目接近全注意力。效率上 prefill 加速从 32K 的 2.83× 涨到 1M 的 9.36×，decode 从约 1.47× 升到 2.01×。消融证明 top-p 优于 top-k（RULER 64K multi-K 从 50.66 跃到 98.60），检索头比例 15% 最佳、dim=16 平衡最好；top-p 按任务动态伸缩：32K 下 niah-S 仅激活 468.8 token、multi-K 激活 2462.1，sparsity 达 78.7%–89.2%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 长上下文 (LongBench 16 任务平均) | 平均准确率 (%) | RTPurbo(top-p) 54.24 | Full Attn 53.80；最强稀疏 baseline RazorAttn 52.98 | 反超全注意力 0.44 点，并以 1.26 点优势领先最强稀疏基线 |
| 长上下文检索 (RULER 64K) | 平均准确率 (%) | RTPurbo(top-p) 85.49 | Full Attn 86.23；MInference 65.61；Quest 70.60 | 相比 MInference 提升 19.88 点，几乎无损逼近全注意力（仅差 0.74） |
| 推理 (AIME24/AIME25) | 准确率 (%) | 86.67 / 86.67 | Full Attn 86.67 / 86.67；Quest 46.67 / 46.67 | 与稠密基线完全持平，相对 Quest 高出 40 点 |
| 超长上下文 multi-hop (512K) | 准确率 (%) | multi-V 92.8 / multi-K 92.8（sparsity >97%） | MInference、FlexPrefill 在极端长度剧烈退化 | baselines 崩溃而 RTPurbo 稳定保持，动态阈值随长度外推 |
| 单层注意力推理加速 | 相对 FlashAttention-2 的加速倍数 | Prefill 2.83×(32K)–9.36×(1M)；Decode 1.47×(32K)–2.01×(1M) | MInference / FlexPrefill（均低于 RTPurbo） | 1M prefill 9.36×，且 decode 内核相对 PyTorch naive 4-stage 快约 1.96–1.99× |
局限与改进
作者承认三点。其一，方法依赖稳定的头特化假设：校准质量在头分工弱或大幅 domain shift 的模型上可能退化。其二，稀疏化不彻底：检索头在 prefill 阶段仍做稠密注意力，所以 prefill 的实际计算量受益主要来自 85% 的局部头；同时实验只覆盖 Qwen3 家族与长上下文/推理负载，跨架构、跨领域的验证不足。其三，超长评测上限是 512K，更极端长度未验证。我额外观察到几点：训练量虽小但仍需 H20 GPU 与两阶段流程，门槛对中小团队不低；top-p 的 GPU kernel 是高度定制的 Triton/CUDA 实现，迁移到非 Hopper 架构或非 Qwen 模型需重写；局部头一律 8192 窗口可能对某些需要中等范围的任务偏保守，限制了潜在收益上限。
独立分析的弱点
第一个弱点是 prefill 稀疏不彻底，检索头仍稠密，这把 prefill 的理论上限卡在 15% 头的稠密开销上——改进方向是给检索头也引入轻量 prefill 稀疏（如分块粗筛后再精确算），作者也把这列为未来工作。第二个弱点是头划分'一次性离线确定、固定比例 15%'，对头分工不明显的小模型或蒸馏后模型可能不适用——改进方向是把比例设为可学习/可在线微调的门控，或引入更鲁棒的检索得分指标（如多 needle、多文档平均）。第三个弱点是低维投影只投 pre-RoPE 的 $Q/K$，对 value 或跨头共享信息未利用——可探索对 KV head 做联合压缩以提升 GQA/MQA 模型的内存收益。第四个弱点是评测偏单一，只测了 Qwen3-30B 系列——建议补 Llama/Mistral 等非 Qwen 架构与对话/Agent 等非推理负载。第五个弱点是自定义 kernel 强绑硬件——可考虑用 Triton 或通用 attention 库封装，降低跨平台部署成本。
未来方向
作者明确提出的方向有三：更自适应的头路由（替代固定 15% 离线划分）、更彻底的 prefill 稀疏化、更广泛的跨模型与跨任务评测。基于成果可延伸的方向包括：把'低维路由+全维生成'的解耦思想推广到其它需要高效近似的模块（如 MoE 的专家选择、检索增强的 RAG 召回打分）；将动态 top-p 的'排序无关直方图 kernel'思路迁移到 KV cache 持续压缩、推测解码的草稿验证等场景；结合本文的'内含稀疏'结论，研究能否在预训练阶段就用极小代价注入稀疏先验，进一步降低原生稀疏训练成本；以及探索 query-dependent 预算在不同 decode 长度下的自适应调度，在长生成（reasoning trace 上万 token）场景里进一步榨取 decode 收益。
复现评估
论文复现门槛中等偏上，配置披露较全但缺关键代码。给出检索头比例 15%、滑动窗口 8192、sink 4、低维 16、top-p 0.9、kernel block 64，两阶段超参（阶段一 lr $10^{-3}$ 余弦、阶段二 lr $3 \times 10^{-6}$，各约 600 步），数据 FineWeb + Dolma3 Longmimo Mix（均开源），评测用 lm-eval，硬件 H20 + Python 3.14 + CUDA 12.8 + PyTorch 2.8，模型 Qwen3-Coder-30B-A3B 与 Qwen3-30B-A3B-Think。难点在于自研排序无关 top-p kernel 与无共享内存单 warp CTA 是高度定制工程，论文给了算法描述（Kernel 1/2）但未确认代码开源，重写工作量大；校准 needle 构造、头划分阈值需自行摸索。整体属于'方法可理解、工程难复刻'的典型前沿工作。
论文图表
示意对比检索头与局部头：大多数头只看局部（A→附近），而检索头会把注意力放到序列远处但语义相关的区域（如相似 pattern A/B/C），说明头级特化天然适合稀疏设计。
这张图直观支撑了'只给检索头留完整 KV cache'这一核心设计动机，是理解头级稀疏先验的最直观证据。