SAS：通过上下文排序端到端优化的简易注意力稀疏化 SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking
用对数空间软门控让语言建模损失端到端训练稀疏注意力选择器
前置知识
块稀疏注意力与注意力预算
把上下文按块大小 $b$（本文 $b=64$）划分为 $C=n/b$ 个连续块，每个 query 只在被选中的 Top-K 个历史块加上恒保留的当前块上做注意力：$o=\mathrm{softmax}(qK_\mathcal{S}^\top)V_\mathcal{S}$。注意力预算 K 即每个 query 允许关注的块数：单个 query 成本从 $O(n)$ 降到 $O(|\mathcal{S}|)$，长度 n 序列的累计成本从 $O(n^2)$ 降到 $O(n|\mathcal{S}|)$。
本文全部设计围绕「固定预算下如何选出最有用的块」展开，主实验、消融与效率分析都以预算 1024/2048/4096 为轴，不懂块划分与预算的含义就读不懂实验表格。
Top-K 选择的不可微性与直通估计器（STE）
Top-K 输出离散索引集合，选择结果对选择器打分的导数几乎处处为零，语言建模损失的梯度无法经选择路径回传。常见补救是直通估计器：前向用硬 0/1 门 $\hat{g}_{hard}=\mathbb{I}[m\in\mathrm{TopK}(s,K)]$，反向假装它是恒等映射。论文式(7)进一步指出，硬门控下被丢弃块的注意力权重 $\tilde{p}^{hard}_i$ 没有上界、可指数爆炸，导致梯度数量级失控。
理解「硬选择为何阻断梯度、STE 为何不稳定」是体会 SAS 用连续软门控绕开该难题的动机，也是读懂消融 III（连续 vs 离散）的前提。
逐层注意力蒸馏（SeerAttention-R 类方法）
现有可训练稀疏方法在冻结骨干 LLM 后，为每层训练一个轻量选择器（AttnGate），以原稠密模型的逐层注意力分布为教师信号做监督，让选择器按稠密注意力权重给上下文块排序。这是一种代理监督：逐层对齐忽略跨层互补性，且只约束 $qK^\top$ 注意力权重，完全忽略 Value 矩阵对最终预测的贡献，造成「排序错位」。
SAS 的出发点正是批评这种蒸馏监督，主实验最核心的对照就是结构完全相同、只把训练信号换成 LM 损失的 SeerAttention-R。
FlashAttention 与在线 softmax 分块计算
FlashAttention 把 $qK^\top$、softmax 与 $PV$ 融合为 SRAM 内的 tile 级流式计算，配合在线 softmax（滚动维护运行最大值与归一化累加项）避免实体化 $n\times n$ 注意力矩阵，使长序列显存占用近似线性。要做稀疏化或门控注入，通常必须写自定义 Triton 内核在这套流式扫描里插入额外逻辑。
SAS 的训练内核要把逐块 log 门控融进 tile 级 $qK^\top$ 计算，推理要接 SGLang/FlashInfer 稀疏解码后端，不理解分块在线 softmax 就读不懂内核设计与效率分析。
对数空间门控与 softmax 归一化
把选择器打分 $s$ 经 softmax 归一化为门控 $g=\mathrm{softmax}(s)$，再以 $\log g$ 的形式加进注意力 logits：$o=\mathrm{softmax}(qK^\top+\log g)V$。由于 $\log g_m=s_m-\mathrm{LSE}(s)$，这在概率上等价于给注意力权重乘 $g_m$；而共享的 $-\mathrm{LSE}(s)$ 项因当前块门恒为 1 并不抵消，从而真正校准「历史块 vs 当前块」的注意力配额，并对全局平移 $s\to s+c$ 不变。
这是 SAS 的核心公式（式 4–9），论文关于门控位置（内嵌 vs 外嵌）、激活（softmax/sigmoid/原始 logits）与软/硬形式的全部消融都围绕它展开。
GQA（分组查询注意力）
为压缩 KV 缓存，现代 LLM 让一组 query 头共享同一份 key/value 头。稀疏化时必须决定按头还是按组选块：SAS 在解码时让同组 query 共享同一批被选块（per-group selection），推理后端据此只 gather 相应 KV 分页；分析实验中全注意力 oracle 集合也按组内多数投票聚合而成。
论文的推理实现细节、CUDA graph 解码路径以及 oracle 召回分析的集合构造都依赖 GQA 概念。
研究动机
长上下文是 LLM 推理的核心效率瓶颈：自回归生成每产出一个 token 都要对全部前文做稠密注意力，长度为 $n$ 的序列累计成本以 $O(n^2)$ 增长，而已部署的高性能模型几乎都是稠密注意力，业界因此更青睐免重训的后训练稀疏化。可训练稀疏化的主流做法是轻量选择器打分加硬 Top-K，代表工作 SeerAttention-R 用 AttnGate 选择器为每个 query 从若干 64 大小的块中选 Top-32 个。问题在于硬 Top-K 不可微，语言建模损失的梯度被完全阻断，这类方法只能退而求其次，用原稠密模型的逐层注意力分布做蒸馏监督。这种代理监督存在排序错位：它教选择器复制稠密模型把注意力质量放在哪里，而不是在固定注意力预算下哪些块真正影响最终预测。具体缺陷有二：逐层目标忽略跨层互补性（一层漏掉的块可在另一层补上）；注意力匹配只约束 $qK^\top$ 权重，完全忽略 Value 矩阵对输出的贡献。后果在紧预算下被放大：预算 1024 时 SeerAttention-R 在 GPQA-Diamond（Qwen3-4B）只有 39.84，远低于全注意力的 56.19，说明有限预算被浪费在了低价值块上。
本文的目标是本文目标是把块选择器的训练信号从注意力蒸馏换成语言建模损失本身，实现端到端的上下文排序优化：训练时把选择器打分转化为对数空间门控注入注意力 logits，让 LM 损失经标准反向传播更新选择器；推理时再把学到的连续排序离散化为 Top-K 硬选择，不引入额外推理开销。工程上要求长序列训练可行——不实体化注意力矩阵，且训练成本可控。衡量标准是在推理、长文本、智能体三类任务上，跨 1024/2048/4096 三档预算和 Qwen3-4B/8B/14B 三个规模，全面超越 SeerAttention-R 与训练自由基线，并在低预算下显著逼近、局部追平全注意力（如预算 4096 的 AIME24 4B 达 71.72 对全注意力 71.25），同时验证该范式可扩展到骨干与选择器联合训练的持续预训练。
与已有工作不同的是，独特切入是把稀疏选择重构为上下文排序问题：Top-K 的结果完全由打分的排序决定，因此只要让连续打分以可微方式影响注意力计算，排序本身就能被 LM 损失端到端优化，离散性只留在推理阶段。这与两条已有路线都不同：与 STE 硬门控不同，SAS 保留连续软门，避免被丢块注意力权重指数爆炸带来的梯度失控；与把门控乘在注意力输出上的外嵌做法不同，SAS 把 $\log g$ 加进 softmax 内部，让门控直接参与注意力质量的归一化分配。更深一层，LM 损失天然是跨所有层联合的全局目标，而蒸馏是逐层局部目标，这正好解释了后文 SAS 单层注意力质量覆盖更低、但跨层选择并集对 oracle 召回更高的现象。论文的贡献不止一个技巧，而是系统识别出四个缺一不可的设计选择（内嵌 log 门、softmax 归一化、保留连续排序、稀疏训练范围），并用受控消融逐一验证。
核心方法
直觉上，SAS 给每个历史上下文块装了一个连续的音量旋钮：选择器打分决定旋钮开度，注意力权重乘上这个开度；训练时旋钮连续可微，LM 损失能直接告诉选择器哪个块该开大、哪个该关小，推理时再把旋钮二值化为选/不选。技术路线沿用 SeerAttention-R 的块稀疏设定（块大小 64、Top-K 32），冻结 Qwen3 骨干、只训练每层 AttnGate 选择器。训练步（Algorithm 1）：选择器对历史块打分 $s=R_\theta(q,\{K_{B_m}\})$；$g=\mathrm{softmax}(s)$ 且当前块 $g_0=1$；取 Top-K 得索引 $I$，关注集合 $\mathcal{S}=B_0\cup\bigcup_{m\in I}B_m$；前向计算 $o=\mathrm{softmax}(qK_\mathcal{S}^\top+\log g_\mathcal{S})V_\mathcal{S}$；LM 损失的梯度更新 $\theta$。推理步（Algorithm 2）去掉门控，直接 Top-K 后做普通稀疏注意力。为支撑 32K 长序列训练，作者实现了 FlashAttention 风格的 Triton 内核，把门控加法融进 tile 级 $qK^\top$ 的流式扫描，避免实体化注意力矩阵。
核心创新是用 LM 损失替代注意力蒸馏来学块排序，具体靠「门控内嵌 + 归一化 + 保连续」的组合实现。与已有方法的本质区别在于：硬 Top-K 使选择路径梯度为零，前人只能用逐层蒸馏或代理信号；SAS 保持离散选择不变，但给每个被选块附加连续软门 $g$，以 $\log g$ 形式加进 softmax 内的 logits。此时门控梯度为 $\partial g_m\propto\sum_{i\in B_m}\tilde{p}_i\,do^\top(v_i-o)$（式 5），包含重新分配注意力质量的相对信号 $v_i-o$，而外嵌门只有固定的 $\tilde{p}_i do^\top v_i$。softmax 归一化使历史块门控相对恒为 1 的当前块自动校准总注意力配额，避免两类退化解：sigmoid 门饱和到 1、原始 logit 注入塌缩到 0（Figure 3），两者都等效于无门注意力、丧失排序信号。保留连续门控而非 STE 二值化，则保证软权重 $\tilde{p}^{soft}_i\le 1$ 有界，避免硬门控下被丢块权重指数爆炸（式 7）造成梯度数量级失控。四个选择缺一不可，共同让端到端梯度对「哪些块真正支撑预测」变得信息量充分。
方法步骤详情
完整流程分四阶段。阶段一（数据与设定）：在 OpenR1-Math-220k 的 93.7K 样本上训练，序列长 32768，全局 batch 32，AdamW 学习率 1e-3 余弦衰减，块大小 64、Top-K 32，冻结 Qwen3-4B/8B/14B 骨干只训 AttnGate，基于 VeOmni 用 FSDP 分布式。阶段二（训练前向，Algorithm 1）：输入当前 query $q$ 与各块摘要 key，选择器输出打分 $s\in\mathbb{R}^C$；计算归一化门 $g=\mathrm{softmax}(s)$，当前块 $g_0=1$；Top-K$(g,K)$ 得索引 $I$，构造关注集合 $\mathcal{S}=B_0\cup\bigcup_{m\in I}B_m$；Triton 内核计算 $o_{SAS}=\mathrm{softmax}(qK_\mathcal{S}^\top+\log g_\mathcal{S})V_\mathcal{S}$，未选块被掩掉、当前块偏置为 0。阶段三（反向传播）：内核把注意力 logits 梯度按块内求和得 $\partial g_m$，经 softmax 回传到 $s$ 再到 $\theta$；消融表明稀疏范围（只对被选块产生梯度，式 8）起步慢（10 步 GPQA 24.8 对全范围 30.8）但 1 epoch 后相当（54.8 对 54.4）且更省，故主实验采用稀疏范围。阶段四（推理，Algorithm 2）：去掉门控直接 Top-K 选块；实现为 SGLang 原生后端，prefill 稠密、decode 块稀疏，每步按 GQA 组选 Top-K 并 gather 对应 KV 分页交给 FlashInfer 稀疏解码核，支持 CUDA graph 与批量并发，解码复杂度 $O(|\mathcal{S}|)$ 与总上下文无关。
技术新颖性
新颖性有四点。其一，这是后训练稀疏化中首个让选择器完全由 LM 损失端到端训练、彻底去掉教师注意力的范式：在相同骨干、相同 AttnGate 架构、相同训练数据的受控对比下，证明语言建模损失是比逐层注意力蒸馏更有效的监督信号（低预算 GPQA 提升 10.6–15.5 点）。其二，四个设计选择的系统性消融本身是方法学贡献：内嵌优于外嵌（GPQA 54.4 对 41.6）、softmax 归一化优于 sigmoid（54.4 对 17.0）与原始 logits（54.4 对 18.8）、软门优于 STE 硬门（54.4 对 46.0）、稀疏范围以更低成本逼近全范围（54.8），每项都有清晰的机理分析而非经验堆砌。其三，log 域门控带来优雅性质：softmax 门在 log 域是加性偏置 $s_m-\mathrm{LSE}(s)$，共享的 LSE 项因当前块不参与而真正起配额校准作用，且门控对打分全局平移不变；论文还给出式(5)(7)(8)的完整梯度推导。其四，工程贡献：把 per-block 门融进 FlashAttention tile 级计算的 Triton 训练内核，以及可加载不同 gate checkpoint 的 SGLang/FlashInfer 稀疏解码后端，让 SAS 与 SeerAttention-R 跑在同一条推理管线上，对比公平且可直接部署。
实验结果
推理任务（Table 2）：在 Qwen3-4B/8B/14B 上，SAS 几乎在所有预算与任务上超过全部稀疏基线。预算 1024 优势最大：MATH500 上 SAS 为 90.65/91.27/92.93，对 SeerAttention-R 的 84.67/83.57/86.12 提升 6.0–7.7 点；GPQA-Diamond 为 50.41/53.17/61.14，对 39.84/39.43/45.64 提升 10.6–15.5 点；Quest 在预算 2048 的 AIME 上崩至 0 分。预算 2048 的 AIME24（4B）SAS 68.85 对 55.83，+13.0；预算 4096 时 SAS 接近或超过全注意力（AIME24 4B 71.72 对 71.25）。长上下文（Table 3）：仅在数学数据上训练的选择器直接迁移到 LongBench，且输入越长差距越大：14B 预算 2048 的 8K+ 桶 53.9 对 51.5；预算 4096 时 14B 平均 56.2 对全注意力 56.6。智能体（Table 4/5）：BFCL 多轮 4B 预算 2048 为 32.50 对 29.00（+3.5），14B 预算 4096 达 44.00 对全注意力 44.50；VitaBench 预算 4096 时 Delivery 34.2 对 32.5、Instore 28.7 对 26.7。持续预训练（Table 6/7）：OLMo3-7B 约 50B token 联合训练后通用任务平均 43.28，优于 HiLS-Attn-RoPE 41.68，接近稠密基座 43.88；LongBench 平均 30.0 并列最佳且增益集中在 >8K。机制与效率：Figure 5b 显示 SAS 跨层并集对 oracle 的召回达 0.996–1.000（蒸馏约 0.98）；Figure 6 显示生成更短、截断更少（AIME24 截断约 21.7%→12.8%）；Figure 7 显示解码在 64K/256K/512K 加速 2.4×/4.6×/5.6×，batch 8 时 64K 约 13×。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 推理 MATH500（Qwen3-4B，预算 1024） | 准确率 avg（%） | 90.65 | SeerAttention-R 84.67（全注意力 93.93） | +5.98 |
| 推理 GPQA-Diamond（Qwen3-4B，预算 1024） | 准确率 avg@16（%） | 50.41 | SeerAttention-R 39.84（全注意力 56.19） | +10.57 |
| 推理 GPQA-Diamond（Qwen3-14B，预算 1024） | 准确率（%） | 61.14 | SeerAttention-R 45.64（全注意力 65.25） | +15.50 |
| 推理 AIME24（Qwen3-4B，预算 2048） | 准确率（%） | 68.85 | SeerAttention-R 55.83 | +13.02 |
| 长上下文 LongBench 8K+ 桶（Qwen3-14B，预算 2048） | 任务均分 | 53.9 | SeerAttention-R 51.5 | +2.4 |
| 智能体 BFCL 多轮（Qwen3-4B，预算 2048） | 准确率（%） | 32.50 | SeerAttention-R 29.00（全注意力 35.75） | +3.50 |
| 解码延迟（Qwen3-4B，512K 上下文，batch 1，SGLang 单卡） | 相对全注意力的加速比 | 5.6×（64K/256K 分别为 2.4×/4.6×） | 全注意力稠密解码 | 延迟降低 5.6 倍 |
局限与改进
作者承认并可观察到的局限：其一，选择器仅在数学推理数据（OpenR1-Math-220k，93.7K 样本）上训练，虽展现出向 LongBench 与智能体任务的零样本迁移，但训练分布的影响未被系统消融；其二，稀疏范围训练起步慢（10 步时 GPQA 仅 24.8，对全范围 30.8），训练步数受限时可能吃亏；其三，Figure 7c 的延迟分解显示 Top-K 块排序在 512K 上下文时占单步解码延迟的 90%，选择器打分也随上下文线性增长，超长上下文下稀疏收益会被选择开销侵蚀；其四，prefill 阶段仍是稠密注意力，超长输入的首 token 延迟没有改善；其五，持续预训练实验（OLMo3-7B）只是初步证据，MMLU 58.58 低于稠密基座 59.90、CRUX 19.25 明显低于基线 24.62，通用能力仍有折损；其六，主实验集中于 Qwen3 家族与 OLMo3，未覆盖更大规模或 MLA 等其他注意力架构；此外消融只在 Qwen3-4B 单一模型上进行，四个设计选择的结论在更大规模上的稳定性留有疑问。
独立分析的弱点
独立分析的弱点：一是分布依赖，选择器只用数学数据训练，作者声称可迁移但未给出训练分布消融（例如用代码或对话数据训练会怎样），改进方向是混合域训练或域自适应路由；二是选择瓶颈，Figure 7c 表明 512K 时 Top-K 排序占解码延迟 90%，这是对所有候选块全排序的代价，改进方向是层次化/分桶 Top-K、利用块分数的时序局部性做增量更新，或蒸馏一个轻量近似选择器；三是训练-推理不一致，训练用带连续偏置的软门，推理是无线偏置的硬 Top-K，存在分布漂移，改进方向是训练末期把软门退火到硬门或推理时保留低精度门控；四是稀疏范围早期的梯度质量差，Figure 4 显示未选块梯度高度相关且忽略自身分数，改进方向是全范围与稀疏范围的课程式混合，或论文尝试过的加噪全范围近似（full*，1 epoch 52.2 已接近 54.8）；五是跨层互补性证据（Figure 5b）只在 Qwen3-4B 长序列上测得，缺少多模型验证；六是对块大小 64、Top-K 与 GQA 组共享选择等超参的敏感性没有系统讨论。
未来方向
作者指出的方向包括：把端到端选择优化推广到更多架构（如 MLA）与更大规模；优化选择阶段内核（Top-K 排序与选择器打分）以解锁更长上下文；把骨干与选择器联合训练的持续预训练做完整缩放。基于其成果可延伸的方向：一是与 KV 缓存淘汰或检索式方法结合，把稀疏化扩展到 prefill 阶段以消除首 token 延迟瓶颈；二是把 softmax 门控「校准历史 vs 当前块」的思想推广到多模态注意力或扩散 Transformer 的 token 稀疏化；三是利用学到的连续块分数作为难度/显著性信号，做推理时动态预算分配——难题多给预算、易题少给，进一步压缩平均成本；四是在智能体场景用强化学习或在线数据更新选择器，使排序适应任务分布漂移；五是理论层面分析 log 门控下排序的学习动力学，解释为何全局 LM 损失能跨越逐层蒸馏的局部最优，为设计更优的端到端稀疏目标提供依据。
复现评估
复现条件较好。代码已在 GitHub 开源（Tencent-Hunyuan/Simple-Attention-Sparsification），训练数据公开（HuggingFace 上的 OpenR1-Math-220k），骨干为开源的 Qwen3-4B/8B/14B 与 OLMo3-7B，训练超参完整（AdamW、lr 1e-3、余弦衰减、batch 32、序列长 32768、1 epoch；持续预训练为 13000 步、batch 512、约 50B token、lr 2e-4 衰减到 2e-5、序列长 8192）。推理实现基于开源 SGLang 与 FlashInfer，且 SAS 与 SeerAttention-R 共享 AttnGate 架构与选块流程，可加载各自 gate checkpoint 在同一后端直接对比。门槛主要在算力与工程：4B–14B 模型 32K 序列训练需要多卡 FSDP（基于 VeOmni），持续预训练 50B token 需要相当规模的集群；自定义 Triton 门控内核与 SGLang 后端集成需要系统经验。另需注意论文中 Quest 结果取自原图数字化（标注 Quest*）、部分 SeerAttn-R 结果为作者复现（标注 SeerAttn-R#），横向对比存在轻微测量来源差异。总体而言，拥有 8×A100/H100 级资源与 SGLang 使用经验的团队可中等难度地复现主结论。
论文图表
(a) 在各层测 Top-16/32/64 被选块的注意力质量覆盖率，含仅用 $qK^\top$ 与 value 感知 $qK^\top+\log\|V\|_2$ 两种权重：SAS 的单层覆盖普遍低于蒸馏方法。(b) 把各层选择取并集后与全注意力 oracle 比较：SAS 的重叠召回在所有上下文长度与预算下一致更高，K=64 长文本达 0.996–1.000，蒸馏约 0.979–0.983。
这是解释「为什么端到端更好」的核心机制证据：蒸馏逐层拟合注意力质量这一局部目标，SAS 学到跨层互补的选择，单层少覆盖但整体更完整，直接支撑论文中心论点。