SSA: Sparse Sparse Attention by Aligning Full and Sparse Attention Outputs in Feature Space

Sparse attention reduces the quadratic complexity of full self-attention but faces two challenges: (1) an attention gap, where applying sparse attention to full-attention-trained models causes performance degradation due to train-inference distribution mismatch, and (2) a capability gap, where models trained purely with sparse attention lack complete gradient flow, preventing them from matching full-attention performance. We propose SSA (Sparse Sparse Attention), a training framework that integrates both sparse and full attention with bidirectional attention-output alignment. We prove that the approximation error scales linearly with the attention mass dropped under sparse attention, and show that SSA's alignment objective substantially reduces this quantity compared to baselines. Experiments demonstrate that SSA achieves state-of-the-art performance under both inference modes, adapts smoothly to varying sparsity budgets, and demonstrates superior long-context capabilities. Code is available at https://github.com/zhenyi4/ssa.

Today's AI language models work by having every word in a text "look at" every other word to understand context. This is powerful, but costly: doubling the length of a document roughly quadruples the computing required. As AI assistants increasingly read whole codebases, browse the web, and carry out "life-long" tasks, this has become a serious bottleneck. A popular fix, called sparse attention, lets each word look at only a small, carefully chosen subset of other words instead of all of them. This saves a lot of computation, but throwing away information usually makes the model less accurate. The challenge is to save the computation while losing as little accuracy as possible. We studied the two common ways of building such models. One takes a model trained with full attention and simply switches it to sparse attention when used; the other trains with sparse attention from the start. Each has a weakness the other happens to fix, so we combined both during training—using full attention as a quality target while sparse attention keeps training consistent with real use—and gently nudged full attention to behave more like sparse attention, which directly reduces the information lost when words are dropped. The result is a model that runs efficiently on long inputs and comes much closer to the accuracy of full attention than previous methods.