Scaling Attention via Feature Sparsity
This addresses the bottleneck of quadratic attention cost in Transformers for researchers and practitioners needing efficient long-context models, though it is incremental as it builds on existing sparsity ideas by focusing on a new axis.
The paper tackles the high computational cost of scaling Transformers to long contexts by introducing Sparse Feature Attention (SFA), which reduces attention cost from O(n^2 d) to O(n^2 k^2/d) and achieves up to 2.5x speed improvement and nearly 50% reduction in FLOPs and KV-cache while matching dense baseline accuracy.
Scaling Transformers to ultra-long contexts is bottlenecked by the $O(n^2 d)$ cost of self-attention. Existing methods reduce this cost along the sequence axis through local windows, kernel approximations, or token-level sparsity, but these approaches consistently degrade accuracy. In this paper, we instead explore an orthogonal axis: feature sparsity. We propose Sparse Feature Attention (SFA), where queries and keys are represented as $k$-sparse codes that preserve high-dimensional expressivity while reducing the cost of attention from $Î(n^2 d)$ to $Î(n^2 k^2/d)$. To make this efficient at scale, we introduce FlashSFA, an IO-aware kernel that extends FlashAttention to operate directly on sparse overlaps without materializing dense score matrices. Across GPT-2 and Qwen3 pretraining, SFA matches dense baselines while improving speed by up to $2.5\times$ and reducing FLOPs and KV-cache by nearly 50\%. On synthetic and downstream benchmarks, SFA preserves retrieval accuracy and robustness at long contexts, outperforming short-embedding baselines that collapse feature diversity. These results establish feature-level sparsity as a complementary and underexplored axis for efficient attention, enabling Transformers to scale to orders-of-magnitude longer contexts with minimal quality loss. Code is available at https://github.com/YannX1e/Sparse-Feature-Attention.