DashAttention: Making Sparse Attention Fully Differentiable with Adaptive Block Selection

Paper: 2605.18753 Authors: Yuxiang Huang, Nuno M. T. Gonçalves, Federico Alvetreti, Lei Li, Xu Han, Edoardo M. Ponti, André F. T. Martins, Marcos V. Treviso Categories: cs.CL, cs.AI, cs.Lg

The Gap

Hierarchical attention methods like NSA and InfLLMv2 try to make long-context processing cheaper by first selecting top-k relevant blocks, then applying full attention only within those blocks. The problem: top-k assumes every query needs exactly k blocks, which is obviously wrong—some queries need more context, some need less. Worse, top-k is non-differentiable, so the coarse selection stage can’t learn from the fine-grained attention stage. You’re training two stages in isolation when they should inform each other.

The paper also identifies a subtle mathematical issue: existing methods are “dispersive”—their attention distributions spread out as context grows, degrading long-range modeling. This happens because they normalize attention scores within selected blocks without accounting for what was discarded.

Problem: Fixed top-k block selection
    |
    v
Assumption: Queries need variable context amounts
    |
    v
Method: Replace top-k with adaptive α-entmax
    |
    +---> Enables gradient flow between stages
    |
    +---> Provides prior for second-stage softmax
    |
    v
Evidence: 75% sparsity, comparable accuracy
    |      Non-dispersive property proven
    |      Better Pareto frontier than baselines
    |
    v
Conclusion: Differentiable hierarchy + adaptivity
            = better long-context modeling

The Increment

One sentence: Before, hierarchical attention used fixed top-k selection with gradient cutoff; after, it uses adaptive sparsity that learns how much context each query actually needs while keeping the entire pipeline differentiable.

Core Mechanism

DashAttention has two stages. Stage 1 divides keys/values into blocks and computes a coarse attention score for each block using α-entmax, a differentiable sparse transformation. Unlike top-k which always picks exactly k blocks, α-entmax can assign zero weight to irrelevant blocks—the number of selected blocks adapts to each query. The α parameter controls sparsity: higher α means sparser selection.

Stage 2 takes the selected blocks and applies standard softmax attention, but with a twist: the coarse scores from Stage 1 become a prior. Specifically, the fine-grained attention logits are added to log(coarse scores) before softmax. This means tokens in blocks that Stage 1 deemed important get a boost, while tokens in barely-selected blocks are downweighted. The entire pipeline is differentiable—gradients from Stage 2 flow back through the log operation into Stage 1’s α-entmax.

Input Query Q, Keys K, Values V
    |
    v
[Stage 1: Block-level Selection]
    |
    +---> Divide K,V into blocks
    |
    +---> Compute coarse scores: Q * block_summary
    |
    +---> Apply α-entmax (adaptive sparse softmax)
    |         |
    |         +---> Output: p_coarse (sparse block weights)
    |
    v
[Stage 2: Token-level Attention]
    |
    +---> For each selected block:
    |         Compute fine logits: Q * K_block
    |         Add prior: logits + log(p_coarse)
    |         Apply softmax
    |
    +---> Weighted sum with V
    |
    v
Output: Attention result
    ^
    |
[Gradients flow back through both stages]

Think of it like a two-stage hiring process. Stage 1 is resume screening: you quickly scan each resume (block) and decide how interested you are. But instead of rigidly picking “top 5 resumes,” you use a flexible threshold—some candidates are clear nos (zero weight), some are maybes (small weight), some are strong yeses (high weight). The number of resumes you advance depends on their quality, not a quota.

Stage 2 is the interview: you carefully evaluate each candidate who passed screening. But you don’t forget Stage 1’s assessment—if someone barely made it through screening, you’re skeptical during the interview (their answers need to be really good to overcome the low prior). If someone had a stellar resume, you’re predisposed to rate them highly (their prior boosts their interview score). Crucially, after interviews conclude, you can reflect back: “I should have screened resumes differently”—the gradient flows backward, improving your screening criteria for next time.

Key Concepts

• α-entmax: Imagine softmax, but with a dial that controls sparsity. Regular softmax spreads probability mass across all options—even terrible ones get a tiny sliver. α-entmax with α > 1 can assign exactly zero to bad options. Higher α means more aggressive pruning. At α=1, you recover standard softmax (no sparsity). At α=2, you get sparsemax (a known sparse transformation). The key property: it’s differentiable everywhere, unlike top-k which has a discontinuous gradient. Concretely, if you have scores [5, 3, 1, 0.1], softmax gives [0.71, 0.19, 0.07, 0.03], but α-entmax with α=1.5 might give [0.82, 0.18, 0, 0]—the bottom two get exactly zero.

• Non-dispersive attention: As context length grows, where does attention probability go? In a dispersive method, it spreads out—if you double the context, each token’s attention weight roughly halves, even if the new tokens are irrelevant. This is bad for long-range dependencies. Non-dispersive means attention stays concentrated on relevant tokens regardless of context length. DashAttention achieves this by using the coarse prior: irrelevant blocks get zero weight in Stage 1, so their tokens can’t dilute Stage 2’s attention distribution. The paper proves this formally: DashAttention’s attention weights don’t depend on the number of irrelevant tokens added to the context.

• Hierarchical prior: Stage 1’s coarse scores aren’t just a filter—they’re a Bayesian prior for Stage 2. In probability terms, Stage 2 computes p(token | query, block is relevant) and Stage 1 provides p(block is relevant | query). The final attention is their product (in log space, a sum). This is why gradients flow: Stage 2’s loss can adjust Stage 1’s beliefs about block relevance. Without this coupling, Stage 1 is blind to whether its selections actually helped Stage 2.

Framework Shift

Before (NSA, InfLLMv2):              After (DashAttention):

Stage 1: Coarse Selection            Stage 1: Adaptive Selection
    |                                     |
    v                                     v
[ top-k operation ]                   [ α-entmax ]
    |                                     |
    +---> Always k blocks                 +---> Variable # blocks
    |                                     |
    +---> No gradient                     +---> Differentiable
    |                                     |
    v                                     v
Stage 2: Fine Attention               Stage 2: Prior-informed Attention
    |                                     |
    v                                     v
[ softmax on selected ]               [ softmax + log(coarse) ]
    |                                     |
    +---> Independent of Stage 1          +---> Coupled with Stage 1
    |                                     |
    v                                     v
Result: Dispersive                    Result: Non-dispersive
        Fixed sparsity                        Adaptive sparsity
        Gradient cutoff                       End-to-end learning

From rigid two-stage pipeline to adaptive coupled hierarchy, the core shift is replacing discrete selection with continuous sparsity that learns.

Expert Assessment

Problem choice: Real gap. Long-context attention is a bottleneck for LLMs, and hierarchical methods are a natural solution. The observation that top-k is both inflexible and non-differentiable is sharp—it’s one of those “obvious in hindsight” insights. The dispersive property is a nice theoretical contribution, though its practical impact is less clear from the experiments.

Method maturity: Elegant. Using α-entmax is clever but not novel (it’s from prior work by Martins et al., one of the co-authors). The real contribution is the hierarchical prior design—adding log(coarse scores) to fine logits is simple but effective. The theoretical analysis (non-dispersiveness proof) is solid. One concern: the method introduces hyperparameters (α, block size) that need tuning. The paper doesn’t deeply explore sensitivity to these choices.

Experimental integrity: Baselines are fair. The comparison with NSA and InfLLMv2 is direct. The Pareto frontier analysis (accuracy vs. sparsity tradeoff) is the right way to evaluate. However, the experiments are mostly on perplexity and a few downstream tasks—would have liked to see more diverse long-context benchmarks (e.g., retrieval, reasoning over long documents). The speedup claims (2.5× over FlashAttention-3) are impressive but only shown at inference; training speedup is unclear. Also, the Triton implementation is mentioned but not open-sourced in the paper, making reproducibility harder.

Writing quality: Clear overall, but the related work section is dense and could be trimmed. The mathematical exposition (especially the non-dispersiveness proof) is rigorous but might lose practitioners. The paper would benefit from more intuitive explanations upfront—the hiring metaphor I used above doesn’t appear in the paper, but something like it should. Figure quality is good, though Figure 3 (Pareto frontier) is doing a lot of heavy lifting and could be broken into separate plots for clarity.

Verdict: Weak accept — solid contribution with good experimental validation, but incremental rather than transformative. The method is practical and the theory is sound, but it’s fundamentally a refinement of existing hierarchical attention rather than a new paradigm.

Takeaways

Adaptive sparsity beats fixed sparsity: If you’re building any kind of sparse selection mechanism (not just attention), consider whether your “top-k” is actually justified. Can you replace it with a differentiable sparse transformation that adapts to input difficulty? α-entmax is a drop-in replacement for softmax in many contexts.

Couple your pipeline stages: If you have a coarse-to-fine architecture where Stage 1 filters and Stage 2 processes, don’t treat them as independent. Make Stage 1’s decisions a prior for Stage 2 (in log space if using softmax). This enables end-to-end learning and often improves both stages.

Non-dispersiveness as a design principle: When designing attention mechanisms for long contexts, check whether your method is dispersive. A simple test: if you append irrelevant tokens to the context, do attention weights on relevant tokens decrease? If yes, you have a problem. The fix often involves explicit normalization or priors that account for what was filtered out.

Hierarchical attention is still underexplored: Despite the hype around linear attention and state-space models, hierarchical attention (coarse-to-fine) remains a practical and effective approach for long contexts. This paper shows there’s still room for improvement in how we design the hierarchy.