Research map / Architectures and efficiency / Attention and transformers

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

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse attention for long-context LLMs · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes SAS, which trains a block-sparse attention selector end-to-end with the language modeling loss by injecting softmax-normalized log-gates into attention, avoiding SeerAttention-R's layer-wise attention distillation.

Problem

Post-training attention sparsification needs a selector that picks the most useful context blocks per query under a budget, but hard Top-K selection blocks gradients from the language modeling loss, so existing trainable methods rely on layer-wise attention distillation.

Approach

Keep hard Top-K selection but add softmax-normalized selector scores as log-space gates inside the attention softmax for selected blocks, so LM-loss gradients train the selector end-to-end; a fused Triton kernel makes it practical.

Key results

• On MATH500 at 1024 budget, SAS improves over SeerAttention-R by 6.0-7.7 points and on GPQA-Diamond by 10.6-15.5 points across Qwen3-4B/8B/14B.
• At 2048 budget on Qwen3-4B, SAS reaches 55.83-ish vs SeerAttention-R 55.83 on AIME24 with +13.0 gain reported in text (68.85 vs 55.83).
• Decode speedup over full attention on Qwen3-4B in SGLang: 5.6x lower latency at batch 1 at 512K and ~13x at batch 8 at 64K.

Surprising findings

• Sigmoid-activated gates and unnormalized raw-logit injection collapse toward ungated attention and perform far worse than softmax gates.
• Sparse-scope training (updating only selected blocks) starts slower but reaches comparable final accuracy to full-scope training at lower cost.
• SAS covers less per-layer attention mass than distillation yet achieves higher cross-layer overlap recall with the full attention oracle.

Limitations

• Performance degrades substantially at long context lengths (RULER 128K), leaving a large gap to full attention, attributed to pooling-based block summaries.
• Top-K selection and selector scoring dominate decode latency at very long contexts.
• Continued-pretraining results are described as initial evidence only.

Experimental setup

Datasets: OpenR1-Math-220k OLMo3 pretraining corpus dolma3_longmino_mix-100B-1125

Benchmarks: MATH500 GPQA-Diamond AIME24 AIME25 LongBench BFCL (Multi-Turn) VitaBench RULER MMLU GPQA HellaSwag ARC-c

Metrics: accuracy avg@16 avg@4 pass@4 generation length truncation rate overlap recall

Baselines: SeerAttention-R Sliding Window Attention StreamingLLM Quest HiLS-Attn-RoPE OLMo3-512SWA

Builds on: SeerAttention-R FlashAttention StreamingLLM HiLS-Attention

Closest papers by meaning

• Selective Attention: Enhancing Transformer through Principled Context Control (NeurIPS 2024 · similarity 0.90)
• SEA: Sparse Linear Attention with Estimated Attention Mask (ICLR 2024 · similarity 0.89)
• Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps (NeurIPS 2026 · similarity 0.89)
• A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention (ICLR 2025 · similarity 0.89)
• Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study (ICLR 2025 · similarity 0.89)
• Punctuation-aware Hybrid Trainable Sparse Attention for Large Language Models (NeurIPS 2026 · similarity 0.89)