SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking
SAS trains a gated sparse attention selector end-to-end with the language modeling loss, beating distillation-based sparsifiers on reasoning, long-context, and agentic tasks.
The paper proposes Simple Attention Sparsification (SAS), which injects a selector's continuous scores into attention softmax logits so the language modeling loss directly optimizes context ranking instead of distilling dense attention distributions. Key design choices include log-form gates inside the softmax, normalized gates calibrated against the current block, and preserved continuous selector scores. A memory-efficient Triton kernel integrates SAS into FlashAttention-style computation for long-sequence training. SAS outperforms trainable sparse attention baselines across budgets, with the largest gains under tight attention budgets.