Scaling Attention via Feature Sparsity

TL;DR

This paper introduces Sparse Feature Attention (SFA), a novel method that leverages feature sparsity to efficiently scale Transformers to longer contexts while maintaining accuracy.

Contribution

The authors propose SFA and FlashSFA, enabling scalable, accurate attention by exploiting feature sparsity and efficient sparse computations, outperforming existing methods in speed and resource usage.

Findings

SFA matches dense baseline accuracy in GPT-2 and Qwen3 pretraining.

SFA achieves up to 2.5x speedup and nearly 50% reduction in FLOPs and KV-cache.

SFA preserves retrieval accuracy and robustness at long contexts.

Abstract

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 $\Theta(n^2 d)$ to $\Theta(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.