SSA: Sparse Sparse Attention by Aligning Full and Sparse Attention Outputs in Feature Space
Zhenyi Shen, Junru Lu, Lin Gui, Jiazheng Li, Yulan He, di yin, Xing Sun
Abstract
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 4440239d-758a-4a84-8f52-2eb465bd979aBuilds on22
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah et al.NeurIPS 2020 · 64,255 citations
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- FlashAttention-2: Faster Attention with Better Parallelism and Work PartitioningTri DaoICLR 2024 · 2,600 citations
- SWE-bench: Can Language Models Resolve Real-world Github Issues?Carlos E. Jimenez, John Yang, Alexander Wettig, Shunyu Yao et al.ICLR 2024 · 2,082 citations
- Efficient Streaming Language Models with Attention SinksGuangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han et al.ICLR 2024 · 1,714 citations
Related papers
- Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse AttentionJingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo et al.ACL 2025 · 334 citations
- SEA: Sparse Linear Attention with Estimated Attention MaskHeejun Lee, Jina Kim, Jeffrey Willette, Sung Ju HwangICLR 2024 · 12 citations
- Elastic Attention: Test-time Adaptive Sparsity Ratios for Efficient TransformersZecheng Tang, Quantong Qiu, Yi Yang, Zhiyi Hong et al.ICML 2026 · 4 citations
- Delta Attention: Fast and Accurate Sparse Attention Inference by Delta CorrectionJeffrey Willette, Heejun Lee, Sung Ju HwangNeurIPS 2025 · 9 citations
- DuSA: Fast and Accurate Dual-Stage Sparse Attention Mechanism Accelerating Both Training and InferenceChong Wu, Jiawang Cao, Renjie Xu, Zhuoheng Ran et al.NeurIPS 2025 · 5 citations
