Overcoming Long-Context Limitations of State-Space Models via Context-Dependent Sparse Attention
2025/07/01 by Zhihao Zhan, Jianan Zhao, Zhan, Zhihao +5 · 1 citation
Computer Science · #Big Data and Digital Economy #Computation and Language (cs.CL) #FOS: Computer and information sciences #Generative Adversarial Networks and Image Synthesis #I.2.7 #Machine Learning (cs.LG) #Topic Modeling
paper · pdf · doi:10.48550/arxiv.2507.00449
openalex publication_date 2025/07/01 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28
Abstract
Efficient long-context modeling remains a critical challenge for natural language processing (NLP), as the time complexity of the predominant Transformer architecture scales quadratically with the sequence length. While state-space models (SSMs) offer alternative sub-quadratic solutions, they struggle to capture long-range dependencies effectively. In this work, we focus on analyzing and improving the long-context modeling capabilities of SSMs. We show that the widely used synthetic task, associative recall, which requires a model to recall a value associated with a single key without context, insufficiently represents the complexities of real-world long-context modeling. To address this limitation, we extend the associative recall to a novel synthetic task, joint recall, which requires a model to recall the value associated with a key given in a specified context. Theoretically, we prove that SSMs do not have the expressiveness to solve multi-query joint recall in sub-quadratic time complexity. To resolve this issue, we propose a solution based on integrating SSMs with Context-Dependent Sparse Attention (CDSA), which has the expressiveness to solve multi-query joint recall with sub-quadratic computation. To bridge the gap between theoretical analysis and real-world applications, we propose locality-sensitive Hashing Attention with sparse Key Selection (HAX), which instantiates the theoretical solution and is further tailored to natural language domains. Extensive experiments on both synthetic and real-world long-context benchmarks show that HAX consistently outperforms SSM baselines and SSMs integrated with context-independent sparse attention (CISA).
Citations
- Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
- HashAttention: Semantic Sparsity for Faster Inference
- Expansion Span: Combining Fading Memory and Retrieval in Hybrid State Space Models
- Hymba: A Hybrid-head Architecture for Small Language Models
- MemoryFormer: Minimize Transformer Computation by Removing Fully-Connected Layers
- B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory
- An Empirical Study of Mamba-based Language Models
- Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
- The Expressive Capacity of State Space Models: A Formal Language Perspective
- Mamba-360: Survey of State Space Models as Transformer Alternative for Long Sequence Modelling: Methods, Applications, and Challenges
- The Illusion of State in State-Space Models
- RULER: What's the Real Context Size of Your Long-Context Language Models?
- Repeat After Me: Transformers are Better than State Space Models at Copying
- The What, Why, and How of Context Length Extension Techniques in Large Language Models -- A Detailed Survey
- Zoology: Measuring and Improving Recall in Efficient Language Models
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces
- LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding
- Lost in the Middle: How Language Models Use Long Contexts
- LongNet: Scaling Transformers to 1,000,000,000 Tokens
- Sparse Modular Activation for Efficient Sequence Modeling
- Faster Causal Attention Over Large Sequences Through Sparse Flash Attention
- Enhancing Chat Language Models by Scaling High-quality Instructional Conversations
- Hyena Hierarchy: Towards Larger Convolutional Language Models
- Hungry Hungry Hippos: Towards Language Modeling with State Space Models
- Efficiently Modeling Long Sequences with Structured State Spaces
- Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers
- The Pile: An 800GB Dataset of Diverse Text for Language Modeling
- HiPPO: Recurrent Memory with Optimal Polynomial Projections
- Reformer: The Efficient Transformer
- Analyzing the Structure of Attention in a Transformer Language Model
- Attention Is All You Need
- Using Fast Weights to Attend to the Recent Past
Cited by
Related