Research map / Architectures and efficiency / Attention and transformers

Overcoming Long Context Limitations of State Space Models via Context Dependent Sparse Attention

NeurIPS 2025 · Architectures and efficiency › Attention and transformers · approach: State space and recurrent sequence models · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Introduces joint recall, a context-dependent extension of associative recall, proves SSMs need sub-quadratic CDSA to solve it, and proposes HAX (LSH attention plus learned key selection attention) hybridized with Mamba, which outperforms Mamba and CISA hybrids on synthetic and real long-context benchmarks.

Problem

State-space models lag Transformers on long-context modeling; the standard associative recall synthetic task does not capture context-dependent key-value associations found in natural language, and SSMs cannot solve multi-query joint recall in sub-quadratic time.

Approach

Extend associative recall to joint recall (values depend on a context), prove SSMs need state size linear in table entries and that SSM + context-dependent sparse attention (LSH) can solve it sub-quadratically, then build HAX, a hybrid of LSH attention and a learned key selection attention that captures globally important "vertical-stripe" keys, added in parallel to Mamba/Mamba2 layers.

Key results

• On Ruler at 2K context, Mamba 790M + HAX averages 44.57 versus 42.43 for Mamba and 41.91 to 42.52 for CISA hybrids.
• On LongBench English, +HAX averages 19.95 versus 19.79 for Mamba, the only hybrid with a clear gain.
• At 4K extrapolation on Ruler, +HAX averages 28.20 versus 27.04 for Mamba, the best among compared models.

Measured effects reported in the paper

Surprising findings

• Plain LSH attention hybrid collapses on Ruler (average 8.36) and LongBench (2.31), far below the Mamba baseline, while the combined HAX recovers performance.
• Dilated attention hybrid nearly fails on most Ruler and LongBench tasks (average 9.78 and 0.51), despite its theoretical CISA status.

Limitations

• Gains on real-world benchmarks are modest (e.g., 19.95 vs 19.79 on LongBench) and many Ruler tasks remain near zero, with no error bars reported.
• Experiments are limited to small Mamba models (130M and 790M) and a short 2K context for training.
• The theoretical results rely on specific constructions and on a uniform distribution assumption for the corollary.

Experimental setup

Datasets: multi-query joint recall synthetic dataset TxT-360 The Pile UltraChat

Benchmarks: Ruler LongBench LM evaluation harness

Metrics: accuracy validation loss average score

Baselines: Mamba Mamba2 dilated attention hybrid sliding window attention hybrid SW+D hybrid A-shaped attention hybrid LSH attention hybrid Samba

Builds on: Mamba Mamba2 locality-sensitive hashing attention (Reformer) associative recall Samba

Closest papers by meaning

• Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps (NeurIPS 2026 · similarity 0.90)
• Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism (ICML 2025 · similarity 0.89)
• FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference (ICLR 2025 · similarity 0.88)
• Efficient Low Rank Attention for Long-Context Inference in Large Language Models (NeurIPS 2025 · similarity 0.88)
• Raven: High-Recall Sequence Modeling via Sparse Memory Routing (NeurIPS 2026 · similarity 0.87)
• Leveraging Attention to Effectively Compress Prompts for Long-Context LLMs (AAAI 2025 · similarity 0.87)