SAS：通过上下文排序的端到端优化实现简单注意力稀疏化

中文摘要 AI 辅助

训练后注意力稀疏化通过为每个查询选择一小部分上下文单元（词元或块）来降低预训练Transformer的二次累积注意力成本。现有的可训练方法通常使用轻量级选择器对上下文单元进行评分，随后进行硬Top-K选择，这会阻断来自语言建模损失的梯度。因此，这些方法通常蒸馏逐层密集注意力分布。虽然这鼓励选择器按原始模型中的密集注意力权重对上下文单元进行排序，但该排序与在固定注意力预算（即每个查询所关注的上下文单元数量）下对预测的影响并不直接对齐，可能将有限的预算浪费在不太有用的单元上。为解决这一错位问题，我们提出了简单注意力稀疏化（SAS），一种门控稀疏注意力机制，通过语言建模损失端到端地优化上下文排序。关键思想是在训练期间将选择器的连续分数注入注意力对数中，使损失能够通过标准反向传播更新选择器。我们确定了该简单设计在实践中良好运行的几个关键选择：将门以对数形式放置在注意力softmax内部，使用归一化softmax门来校准历史上下文与始终保留的当前块，以及保留连续选择器分数，使模型学习相对优先级而不仅仅是硬选择。为支持长序列训练，我们实现了一个内存高效的Triton内核，将SAS集成到FlashAttention风格的计算中。在推理、长上下文理解和智能体任务中，SAS在各种注意力预算下均持续优于可训练的稀疏注意力基线，在严格预算下尤其获得大幅提升，展示了对下游任务更有效的上下文排序。

英文摘要

Post-training attention sparsification reduces the quadratic cumulative attention cost of pretrained Transformers by selecting a small set of context units (tokens or blocks) for each query. Existing trainable methods usually use a lightweight selector to score context units, followed by hard Top-K selection that blocks gradients from the language modeling loss. Consequently, these methods commonly distill layer-wise dense attention distributions. Although this encourages the selector to rank context units by dense attention weights in the original model, the ranking is not directly aligned with their impact on predictions under a fixed attention budget (i.e., the number of attended context units per query), potentially wasting the limited budget on less useful units. To address this misalignment, we propose Simple Attention Sparsification (SAS), a gated sparse attention mechanism that optimizes context ranking end-to-end with the language modeling loss. The key idea is to inject the selector's continuous scores into attention logits during training, allowing the loss to update the selector through standard backpropagation. We identify several choices crucial for this simple design to work well in practice: placing the gate inside the attention softmax in log form, using normalized softmax gates to calibrate historical context against the always-retained current block, and preserving continuous selector scores so the model learns relative priorities rather than only hard selections. To support long-sequence training, we implement a memory-efficient Triton kernel that integrates SAS into FlashAttention-style computation. Across reasoning, long-context understanding, and agentic tasks, SAS consistently outperforms trainable sparse attention baselines across attention budgets, with especially large gains under tight budgets, demonstrating more effective context ranking for downstream tasks.

发表机构

• Tencent HY LLM Frontier（腾讯混元大模型前沿团队）
• Hong Kong University of Science and Technology (Guangzhou)（香港科技大学（广州））
• Hong Kong University of Science and Technology（香港科技大学）

机构由 AI 辅助整理，请以论文原文为准。