# AdaSplash-2: Faster Differentiable Sparse Attention

arXiv: https://arxiv.org/abs/2604.15180 · HTML: https://mlsanity.com/paper/2604.15180/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 56 references that are in the map.
- First posted on arXiv 16 Apr 2026.

## Builds on (most cited first, 30 of 38)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models](https://mlsanity.com/paper/2512.02556/) (2025)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://mlsanity.com/paper/1602.02068/) (2016)
- [The Impact of Positional Encoding on Length Generalization in Transformers](https://mlsanity.com/paper/2305.19466/) (2023)
- [Repeat After Me: Transformers are Better than State Space Models at Copying](https://mlsanity.com/paper/2402.01032/) (2024)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [Sparse Sequence-to-Sequence Models](https://mlsanity.com/paper/1905.05702/) (2019)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [How to Train Long-Context Language Models (Effectively)](https://mlsanity.com/paper/2410.02660/) (2024)
- [How do Language Models Bind Entities in Context?](https://mlsanity.com/paper/2310.17191/) (2023)
- [NVIDIA Nemotron Nano 2: An Accurate and Efficient Hybrid Mamba-Transformer Reasoning Model](https://mlsanity.com/paper/2508.14444/) (2025)
- [OLMES: A Standard for Language Model Evaluations](https://mlsanity.com/paper/2406.08446/) (2024)
- [Transformers need glasses! Information over-squashing in language tasks](https://mlsanity.com/paper/2406.04267/) (2024)
- [TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training](https://mlsanity.com/paper/2410.06511/) (2024)
- [Softmax is not Enough (for Sharp Size Generalisation)](https://mlsanity.com/paper/2410.01104/) (2024)
- [Scalable-Softmax Is Superior for Attention](https://mlsanity.com/paper/2501.19399/) (2025)

## Built on by (most cited first, 3 of 3)

- [DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention](https://mlsanity.com/paper/2605.18753/) (2026)
- [EntmaxKV: Support-Aware Decoding for Entmax Attention](https://mlsanity.com/paper/2605.21649/) (2026)
- [Vision Transformers Need Better Token Interaction](https://mlsanity.com/paper/2605.23868/) (2026)

## Abstract (arXiv)

Sparse attention has been proposed as a way to alleviate the quadratic cost of transformers, a central bottleneck in long-context training. A promising line of work is α-entmax attention, a differentiable sparse alternative to softmax that enables input-dependent sparsity yet has lagged behind softmax due to the computational overhead necessary to compute the normalizer τ. In this paper, we introduce AdaSplash-2, which addresses this limitation through a novel histogram-based initialization that reduces the number of iterations needed to compute τ to typically 1--2. The key idea is to compute a coarse histogram of attention scores on the fly and store it in on-chip SRAM, yielding a more accurate initialization that enables fast forward and backward computation. Combined with a sparsity-aware GPU implementation that skips zero blocks with low overhead, AdaSplash-2 matches or improves per-step training time relative to FlashAttention-2 when block sparsity is moderate-to-high (e.g., >60%), which often occurs at long-context lengths. On downstream tasks, models trained with our efficient α-entmax attention match softmax baselines at short-context lengths and achieve substantial gains in long-context settings.
