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

- 来源：HuggingFace Daily Papers（社区热门论文）
- 发布时间：2026-09-11 08:00
- AIHOT 分数：40
- AIHOT 链接：https://aihot.news/items/cmu0nreon0383rox2jjtc6qcl
- 原文链接：https://arxiv.org/abs/2609.13141

## AI 摘要

针对可训练注意力稀疏化中 Top-K 硬选择阻断梯度、排序与预测影响不对齐的问题，研究者提出 SAS（Simple Attention Sparsification），一种门控稀疏注意力机制，将选择器的连续分数注入注意力 logits，用语言建模损失端到端优化上下文排序。

## 正文

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.