Research map / Architectures and efficiency / Attention and transformers
AdaSplash: Adaptive Sparse Flash Attention
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 AdaSplash, a hybrid Halley-bisection solver and fused Triton kernels that make alpha-entmax attention efficient enough to rival FlashAttention-2 for long-context training.
Problem
Adaptive sparse attention via alpha-entmax is data-dependent and sparse, but existing implementations are slow, memory-hungry, and cannot exploit sparsity to save runtime or memory at long context lengths.
Approach
A hybrid Halley-bisection root solver that needs about 7x fewer iterations than bisection to compute the alpha-entmax threshold, combined with a FlashAttention-style Triton kernel that skips null blocks via dynamically built block-mask lookup tables in forward and backward passes.
Key results
- Halley-bisection reduces iterations about 7x versus bisection and runs about 15x faster than Torch bisection at n=8192 (2.38 ms vs 36.67 ms).
- Sparse GPT-2 (alpha=1.5) reaches validation loss 3.263 vs 3.283 for softmax GPT-2, with runtime 1.03 s/step vs 0.98 s/step for FlashAttention-2.
- ModernBERT with alpha=1.5 outperforms its dense counterpart on all four BEIR single-vector retrieval tasks (e.g., TREC-COVID 75.2 vs 67.6 nDCG@10).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace Torch bisection with hybrid Halley-bisection | α-entmax computation · Random tensors, sequence length n=8192 | runtime: 36.67 ms -> 2.38 ms | Helped |
| Replace Torch bisection with hybrid Halley-bisection | α-entmax computation · Random tensors, sequence length n=8192 | memory usage: 896.15 MB -> 512 MB | Helped |
| Use α-entmax attention (α=1.5) instead of dense attention | ModernBERT · BEIR SciFact, NFCorpus, FiQA, and TREC-COVID | nDCG@10: 57.7 -> 58.4; 22.4 -> 25.7; 25.7 -> 29.6; 67.6 -> 75.2 | Helped |
| Use α-entmax attention (α=1.5) instead of dense attention | RoBERTa · ECtHR long document classification, sequence lengths 512–4096 | F1 micro: 71.5 -> 71.8; 74.4 -> 75.5; 75.1 -> 76.4; 77.9 -> 78.0 | Helped |
| Use α-entmax attention (α=1.5) instead of dense attention | RoBERTa · ECtHR long document classification, sequence length 8192 | F1 micro: 79.2 -> 78.6 | Hurt |
| Use α-entmax attention (α=1.5) instead of softmax | GPT-2 · GPT-2 language modeling | validation loss: 3.283 -> 3.263 | Helped |
| Use α-entmax attention (α=1.5) instead of softmax | GPT-2 · HellaSwag | accuracy: 30.4 -> 30.6 | Helped |
| Use α-entmax attention (α=1.5) instead of softmax | GPT-2 · GPT-2 training | runtime per step: 0.98 -> 1.03 | Hurt |
| Use α-entmax attention (α=1.5) instead of softmax | GPT-2 · GPT-2 training | peak memory per GPU: 52.5 -> 52.5 | No clear effect |
| Use AdaSplash instead of Torch bisection | ModernBERT-base · Sequence length 8192 | runtime: 4.99 -> 1.53 | Helped |
| Use α-entmax attention (α=2.0) instead of dense attention | ModernBERT · GLUE average | average score: 83.7 -> 83.7 | No clear effect |
Surprising findings
- Sparse ModernBERT (alpha=1.5) outperforms dense ModernBERT on all four BEIR retrieval tasks.
- The sparse attention kernel can surpass FlashAttention-2 runtime at high input sparsity.
Limitations
- Forward pass is always slower than FlashAttention due to extra K reads for computing the threshold tau.
- Benefits depend on high attention sparsity; with dense attention (e.g., GPT-2 at 1k context) the gains over FlashAttention-2 are negligible.
- Experiments are limited to small encoder and GPT-2 models (124M-149M parameters).
Experimental setup
Datasets: FineWeb-Edu FineWeb ECtHR GLUE MS MARCO HellaSwag
Benchmarks: GLUE BEIR (SciFact, NFCorpus, FiQA, TREC-COVID) ECtHR long document classification HellaSwag
Metrics: runtime peak memory nDCG@10 F1 micro accuracy validation loss
Baselines: FlashAttention-2 Torch bisection Torch sorting softmax attention
Builds on: FlashAttention entmax sparsemax Triton
Closest papers by meaning
- Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.89)
- SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking (NeurIPS 2026 · similarity 0.88)
- Sparse Attention as Compact Kernel Regression (NeurIPS 2026 · similarity 0.88)
- SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference (ICML 2025 · similarity 0.88)
- MonarchAttention: Zero-Shot Conversion to Fast, Hardware-Aware Structured Attention (NeurIPS 2025 · similarity 0.87)
- Delta Attention: Fast and Accurate Sparse Attention Inference by Delta Correction (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.