Overcoming Long-Context Limitations of State-Space Models via Hybrid Sparse Attention
Zhihao Zhan, Jianan Zhao, Zhaocheng Zhu, Jian Tang
Long-context modeling is a fundamental challenge in natural language processing (NLP) with numerous practical applications. State-space models (SSMs) have gained attention as a promising alternative to Transformers, offering sub-quadratic efficiency for long-context sequence modeling. Despite this advantage, both empirical evidence and theoretical studies have shown that SSMs are substantially less effective than Transformers in capturing long-range dependencies. While recent approaches such as Transformer-SSM hybrid architectures have yielded better performance, they suffer from poor scalability with quadratic time complexity. In this work, we take an initial step towards identifying design principles for sub-quadratic hybrid architectures. We motivate the study of hybrid sparse attention via a novel synthetic task, joint recall. Through theoretical analysis and empirical evaluation, we demonstrate that hybrid sparse attention significantly improves the long-context modeling capacity of SSMs. Moreover, we highlight the importance of incorporating input-dependent sparse attention patterns. Our findings are further validated on real-world long-context NLP tasks by continual pre-training Mamba augmented with hybrid sparse attention.