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.
- Selector scores injected into attention logits enable end-to-end training
- Avoids distilling dense attention distributions used by prior Top-K methods
- Triton kernel integrates SAS into FlashAttention-style long-sequence training
- Consistent gains on reasoning, long-context, and agentic tasks under tight budgets
From the original article
1 min read at arxiv.orgPost-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.
By Zhiwei Li · arxiv.org
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.