SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking¶

会议: NeurIPS 2026
arXiv: 2609.13141
OpenReview: https://openreview.net/forum?id=Wgm2AWIzrX
代码: https://github.com/Tencent-Hunyuan/Simple-Attention-Sparsification
领域: LLM 效率
关键词: 注意力稀疏化、上下文排序、KV缓存选择、长上下文推理、端到端优化

一句话总结¶

针对后训练注意力稀疏化中离散 Top-K 截断阻断梯度并导致蒸馏排序失准的问题，本文提出 SAS，将轻量选择器的连续分数以归一化对数门控形式注入注意力 Softmax 中，直接由语言建模损失端到端优化上下文排序，在极低注意力预算下显著提升推理与长程任务表现。

研究背景与动机¶

大语言模型在处理超长文本或多轮长链思维推理时，自回归解码过程中每个新 token 都必须对全部前序上下文执行全量稠密注意力计算，使得注意力累计开销随序列长度呈二次方急剧攀升。为了在不改变预训练模型架构的前提下降低推理延迟，后训练注意力稀疏化成为主流范式。这类方案通过设计轻量选择器对上下文块进行评分，并在推理时仅挑出前 $K$ 个关键块参与计算。然而，离散的硬性 Top-K 选择在数学上是分段常数函数，这完全阻断了下游语言建模任务损失反向传播到选择器的梯度路径。

现存可训练稀疏方法为了绕过梯度阻断，普遍退而求其次地采用逐层注意力蒸馏作为代理监督信号，强迫选择器拟合原始稠密模型在各层的注意力权重分布。这种逐层拟合机制引发了严重的上下文排序失准：首先，它只关注稠密模型把注意力权重投放在哪里，却忽略了在固定且受限的注意力预算下，到底哪些上下文块对最终预测输出构成了决定性影响；其次，逐层局部目标忽视了不同网络层之间在上下文使用上的跨层互补性；最后，仅对齐注意力权重矩阵完全忽视了值矩阵 $V$ 对最终表征的实质贡献，极易将宝贵的注意力预算浪费在无益的历史块上。

打破该困境的核心切入点在于将块选择问题重新表述为可微分的上下文排序学习，直接利用最终的自回归语言建模损失来驱动选择器优化。核心 idea：将轻量选择器预测的块级连续分数转化为归一化对数门控并直接融入注意力 Softmax 内核，在保留离散稀疏前向的同时打通语言建模损失的端到端梯度回传，并通过专用 Triton 算子实现兼具跨层互补性与高效性的长文本稀疏注意力训练。

方法详解¶

整体框架¶

SAS 围绕“后训练端到端上下文排序学习”展开构建。在输入阶段，当前 query 向量 $q$ 与历史上下文缓存被划分为若干固定大小的连续块，模型始终无条件保留当前块 $B_0$，而将其余历史块送入轻量级选择器计算连续相关性得分。随后，得分经归一化转换为对数空间门控偏置，直接加至注意力 logits 中，配合离散 Top-K 索引筛选形成最终稀疏计算集合。在训练期间，语言建模交叉熵损失直接通过连续门控反向传播至选择器参数；在推理阶段，则直接按分数离散挑选 Top-K 块走标准稀疏解码流程。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入查询 q 与上下文 KV 块<br/>划分候选历史块与固定保留块 B0"] --> B["内置对数门控机制<br/>计算相关性得分并注入 Softmax"]
    B --> C["归一化 Softmax 门控激活<br/>校准历史块与无偏当前块权重"]
    C --> D["连续排序保留与稀疏训练范围<br/>Top-K 块参与门控注意力与梯度回传"]
    D --> E["融合 Triton 内核执行<br/>流式瓦片计算与端到端反向传播"]

关键设计¶

1. 内置对数门控机制：使语言建模损失端到端反向传播至选择器 针对离散 Top-K 选择阻断梯度、使得主干语言模型损失无法指导选择器的问题，SAS 并没有采用带有启发式假设的辅助蒸馏损失，而是将选择器输出的连续标量转化为乘性注意力调制项。具体而言，若将门控施加于 Softmax 外部（即 $o = \operatorname{softmax}(qK^\top)(g \odot V)$），注意力概率分布被提前冻结，门控只能单纯缩放输出向量，导致梯度形式退化为 $dg_m^{outer} = \sum_{i\in B_m} p_i do^\top v_i$，缺乏对块间相对重要性的比较能力。SAS 选择将正门控向量 $g$ 的对数形式 $\log g$ 嵌入 Softmax 内部参与 logits 计算： $$o = \operatorname{softmax}(qK^\top + \log g) V$$ 由此推导出的门控梯度为： $$dg_m^{inner} = \sum_{i\in B_m} \frac{\tilde{p}_i}{g_m} do^\top (v_i - o)$$ 这一设计使梯度中天然包含了中心化项 $(v_i - o)$，能够直接依据各块所提供值向量与当前混合输出的差异，自适应地在注意力分布中重新拉动权重质量分配，从而让语言建模损失能够精确指导历史块的相对优先级。

2. 归一化 Softmax 门控激活：校准历史上下文与当前块的相对权重 针对门控激活函数在训练过程中容易产生数值饱和或尺度塌陷、从而破坏排序判别力的问题，SAS 对比了 Sigmoid 独立激活、无激活直接注入 logits 以及归一化 Softmax 激活三种形态。实验表明，Sigmoid 激活会使所有历史块门控逐步饱和趋向于 1，而未归一化的原始得分会随着训练塌陷至 0 且方差急剧缩小，两者都会退化为无门控的均质化状态。SAS 强制将恒等门控 $g_0 = 1$（即 $\log g_0 = 0$）赋予必然保留的局部当前块 $B_0$，而对所有候选历史块 $\mathcal{H}=\{B_1, \dots, B_C\}$ 应用竞争性的 Softmax 归一化： $$\log g_m = s_m - \operatorname{LSE}(s), \quad \forall B_m \in \mathcal{H}$$ 项 $-\operatorname{LSE}(s)$ 虽然对所有历史块是共享常数，但由于不作用于当前块 $B_0$，它在 Softmax 中有效地对全局历史上下文与当前局部块的关注比例进行了动态标定。同时，这一归一化使门控具有平移不变性，彻底规避了训练中选择器得分整体漂移导致的退化。

3. 连续排序保留与稀疏训练范围：兼顾优化稳定性与长序列计算效率 针对使用 Straight-Through Estimator（STE）二值化硬门控时引发的梯度爆炸，以及全序列计算导致显存超载的问题，SAS 综合采取了连续权重保留与稀疏域计算策略。在 STE 模式下，被丢弃块在受限归一化分母之外，其虚拟权重上限为 $\exp(z_i - \max_{j\in\mathcal{S}} z_j)$，一旦初始化较差的丢弃块包含较大原始注意力 logit，其反向梯度便会呈指数级剧烈爆炸。SAS 在前向计算中完整保留平滑的连续对数门控，使注意力权重严格上界于 1，从根源上保障了梯度的平稳。在训练范围上，若每步对全部上下文块回传梯度，计算代价极其昂贵；SAS 仅对入选的 Top-K 块 $\mathcal{S} = B_0 \cup \bigcup_{m\in\mathcal{I}} B_m$ 展开注意力计算，未选中的块仅通过归一化分母产生微弱的平滑隐式更新： $$ds_m^{sparse} = -g_m \sum_{\ell \in \mathcal{I}} g_\ell dg_\ell$$ 消融证明，稀疏范围训练在初期收敛略慢，但在 1 个 epoch 内即可完全追平全范围训练的最终精度，大幅降低了长序列训练开销。

损失函数 / 训练策略¶

SAS 在后训练阶段完全冻结主干大语言模型的全部权重，仅更新轻量选择器参数 $\theta$。优化目标完全遵循标准的自回归语言建模交叉熵损失： $$\mathcal{L}_{\mathrm{LM}} = -\sum_{t=1}^T \log P(x_t \mid x_{<t}; \theta)$$ 训练在 OpenR1-Math-220k 数据集上展开，序列长度达到 32,768 tokens，块大小设置为 64，Top-K 设置为 32。使用 AdamW 优化器，最大学习率设为 1e-3，配合余弦退火衰减与 FSDP 分布式策略。为了在超长序列下高效执行稀疏对数门控注意力并避免物化完整的 $N \times N$ 注意力矩阵，研究团队基于 Triton 定制了 FlashAttention 风格的融合内核。该内核将按阈值门控的块跳过机制融合进流式分块扫描流程中，在线执行包含对数门控的 Softmax 递推更新，反向传播时直接在线累加块级梯度，实现高效长程训练。在推理阶段，直接将选择器评分进行离散截断，无缝对齐 SGLang 和 FlashInfer 稀疏解码引擎。

实验关键数据¶

主实验¶

主实验在 Qwen3-4B、8B、14B 模型上针对复杂数学与科学推理基准进行了多预算评测。与密集全量注意力（Full Attn）、静态窗口（Sliding Window）、注意力汇聚（StreamingLLM）、启发式动态缓存选择（Quest）以及同架构但基于蒸馏训练的强基线（SeerAttention-R）展开系统对比。

模型	预算 (Tokens)	方法	MATH500	GPQA-Diamond	AIME24	AIME25
Qwen3-4B	Full	Full Attn	93.93	56.19	71.25	66.41
Qwen3-4B	1024	StreamingLLM	72.72	13.86	-	-
Qwen3-4B	1024	SeerAttn-R	84.67	39.84	-	-
Qwen3-4B	1024	SAS (本文)	90.65	50.41	-	-
Qwen3-4B	2048	Quest	40.74	12.17	0.00	0.00
Qwen3-4B	2048	SeerAttn-R	91.85	49.94	55.83	45.16
Qwen3-4B	2048	SAS (本文)	93.47	54.86	68.85	56.38
Qwen3-4B	4096	SeerAttn-R	94.10	55.40	69.32	58.59
Qwen3-4B	4096	SAS (本文)	93.67	55.11	71.72	59.97
Qwen3-14B	1024	SeerAttn-R	86.12	45.64	-	-
Qwen3-14B	1024	SAS (本文)	92.93	61.14	-	-
Qwen3-14B	2048	SeerAttn-R	93.02	61.68	63.65	48.70
Qwen3-14B	2048	SAS (本文)	93.54	65.09	77.76	64.19
Qwen3-14B	4096	Full Attn	95.22	65.25	78.91	70.21
Qwen3-14B	4096	SAS (本文)	95.23	65.03	78.28	67.29

消融实验¶

在 Qwen3-4B 主干上使用 2048 预算对 GPQA-Diamond 进行系统的门控设计消融（报告 avg@16 准确率 % 与平均生成长度）：

设定编号	门控位置	激活函数	排序连续性	训练范围	100 steps	1000 steps	1 epoch (最终)	说明
(1) 基线	–	–	–	–	–	–	56.1 (len: 8547)	稠密 Full Attention 上限
(2) 默认配置	inner	softmax	连续 (✓)	full	52.2	53.7	54.4 (len: 71091)	全范围端到端门控
(3) 外部门控	outer	softmax	连续 (✓)	full	38.9	43.2	41.6 (len: 13807)	门控仅缩放值向量 V，损失大幅下降
(5) Sigmoid 激活	inner	sigmoid	连续 (✓)	full	20.4	17.6	17.0 (len: 28444)	门控趋向饱和，丧失判别力
(6) 无激活 Logit	inner	none	连续 (✓)	full	19.9	18.9	18.8 (len: 26401)	未归一化导致方差塌陷
(8) 离散 STE	inner	softmax	离散 STE (✗)	full	42.4	49.6	46.0 (len: 72891)	硬门控导致梯度方差激增与欠优
(10) 稀疏训练	inner	softmax	连续 (✓)	sparse	51.0	53.6	54.8 (len: 74101)	节约显存且最终准确率无损

关键发现¶

极紧凑预算下优势显著：在 1024 token 极限注意力预算下，SAS 相比基于蒸馏的 SeerAttention-R 在 Qwen3 各尺度上 MATH500 提升达 6.0–7.7 个点，GPQA-Diamond 更是提升达 10.6–15.5 个点；在 AIME 等深度推理基准上，Quest 等启发式方法彻底崩溃，而 SAS 展现出极高鲁棒性。
跨层互补性打破局部最优：分析发现，虽然 SAS 在各单层所覆盖的注意力质量绝对比例（attention mass coverage）低于蒸馏模型，但其在所有层选出块的跨层并集针对稠密 Oracle 的召回率却显著高于蒸馏基线。这证明端到端语言建模损失能够协调各层做出差异化互补选择，而非盲目复制稠密模型的层级局部注意力。
推理长度缩短与截断率骤降：因选取的上下文更加贴合任务实际需求，SAS 在复杂推理中有效避免了无效推演，平均生成 token 数量更短，且在 32k 最大长度下的生成截断率相比 SeerAttention-R 大幅降低。
端到端加速实测突破：在 SGLang 真实部署中，当上下文扩展至 512K 时，单 GPU batch 1 下 SAS 达到 5.6× 的解码加速比；在 batch 8 负载下，64K 长度即实现 ~13× 的显著吞吐加速。

亮点与洞察¶

抛弃辅助蒸馏的纯粹端到端设计：利用 $\operatorname{softmax}(qK^\top + \log g)$ 的数学形式巧妙连通梯度路径，免除了维护稠密教师模型前向与提取庞大注意力矩阵的高昂开销，使稀疏选择器直接对最终预测质量负责。
内生自洽的对数空间相对调制：通过 Softmax 归一化历史块并保持当前块 $g_0 = 1$ 的无偏锚定，赋予门控天然的尺度平移不变性与竞争性，防止了梯度更新过程中的尺度漂移与模式崩溃。
从局部拟合到全局协同的认知升维：揭示了注意力稀疏化评估中“单层注意力权重覆盖率”与“全局真实预测能力”的脱节，证明大模型稀疏注意力可以通过跨层上下文互补来用更少的总体注意力开销换取更准的推理。

局限与展望¶

超长上下文下 Top-K 排序瓶颈：随着上下文长度向 512K 及以上持续扩展，尽管注意力计算本身的耗时维持常数级别，但对全体候选块进行评分与 Top-K 排序的耗时在解码步中占比急剧攀升至 90%，未来需要两级粗细粒度结合或层次化检索算子来打破排序瓶颈。
选择器结构与预训练主干存在轻微领域漂移：当前后训练实验主要在高质量数学推理数据（OpenR1-Math）上微调选择器，尽管在 LongBench 和 Tool-Use（VitaBench）上展示了极佳泛化性，但在开放领域超长对话中的排序鲁棒性仍有待更大规模多样化数据淬炼。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙提出内部对数门控解决注意力离散选择的梯度阻断问题，重塑了注意力稀疏化的训练范式。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖推理、长文本、智能体工具使用全方位场景，包含 4B 至 14B 模型、消融实验、梯度与层级互补性理论机理分析。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻，从数学推导到工程实现细节详实，图表逻辑清晰自洽。
价值: ⭐⭐⭐⭐⭐ 为长文本大模型高并发部署与推理解码加速提供了无需改动主干权重、精度近乎无损且工程极度简洁的高效基建方案。