AdaSplash-2: Faster Differentiable Sparse Attention

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 $\alpha$-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 $\tau$. 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 $\tau$ 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 $\alpha$-entmax attention match softmax baselines at short-context lengths and achieve substantial gains in long-context settings.

Modern language models decide how much focus each word (or "token") should give to every other word, in a mechanism known as attention. The dominant choice, softmax attention, spreads a little focus everywhere, so models can never fully ignore irrelevant tokens. A sharper alternative, $\alpha$-entmax, can assign exactly zero attention to irrelevant tokens, but it depends on finding a cutoff threshold whose computation is expensive. Previous efforts, such as AdaSplash, brought this method onto modern GPUs, yet still required several passes over the attention scores to pin down that cutoff, holding back its use whenever sparsity is only moderate. We introduce AdaSplash-2 to remove this bottleneck. As scores stream through the GPU's fastest on-chip memory, we sketch their distribution into a small histogram that fits entirely on-chip. This sketch is precise enough to nearly nail the cutoff in one shot, leaving just one or two cheap refinement steps. Combined with a lightweight way to skip blocks of zero attention, AdaSplash-2 keeps up with, and often outpaces, FlashAttention-2, the highly-optimized softmax kernel behind today's large language models, while training models that handle long documents more accurately.