Research map / Architectures and efficiency / Attention and transformers
Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access
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.
Proposes Hierarchical Sparse Attention (HSA), a chunk-based sparse attention with learnable token-to-chunk relevance and a hardware-aligned Triton kernel, and integrates it into Mamba-2 as RAMba for long-range random access and length generalization.
Problem
RNNs like Mamba cannot randomly access historical context, and existing chunk-based sparse attentions (NSA, MoBA) select chunks inaccurately, both in-distribution and when extrapolating to longer contexts, while remaining efficient.
Approach
Two-stage hierarchical attention: token-level attention within each selected chunk, then chunk-level weighted fusion using stick-breaking chunk weights, so chunk selection receives end-to-end gradients from next-token prediction; combined with CPU offloading of token-level KV cache and a Triton kernel.
Key results
- RAMba achieves perfect passkey retrieval accuracy at 64M context despite pre-training on 4K contexts, while most baselines drop to near zero around 64K.
- On PG19/ArXiv/Code perplexity at 370M, RAMba w/ m.r. reaches 17.01 PG19 and 3.65 ArXiv at 64K eval length, versus Mamba 17.30 and 3.86.
- HSA is reported 3x faster than NSA and 5-25x faster than full attention for contexts of 16K tokens or more in forward pass; RAMba reaches ~90% of Mamba's training throughput.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add HSA to Mamba-2 | Mamba-2 · Long-range language modeling, mostly 16K and 64K contexts | perplexity | Helped |
| add memory reset during training (ablation) | RAMba · Long-range language modeling, 16K and 64K contexts | perplexity: 17.11 -> 17.15; 17.11 -> 17.01 | Helped |
| add memory reset during training (ablation) | RAMba · Long-range language modeling, 4K context | perplexity: 17.63 -> 17.82 | Hurt |
| replace stick-breaking weights with softmax without positional encoding (ablation) | RAMba without memory reset · ArXiv-math, eval_len=16k | perplexity: 3.81 -> 5.01 | Hurt |
| remove the chunk encoder (ablation) | RAMba · RULER, retrieval tasks as context length grows | accuracy: 55.22 -> 23.88 | Hurt |
| enable memory offloading (ablation) | RAMba · Generating 100 tokens, prompt length=16K, batch-size=16 | inference time: 3.05 -> 4.19 | Hurt |
| remove shared chunk selection across HSA layers (ablation) | RAMba with offloading · Generating 100 tokens, prompt length=16K, batch-size=16 | inference time: 4.19 -> 5.87 | Hurt |
| add HSA instead of NSA | Mamba-2 · Attention computation, contexts of 16K tokens or more | forward-pass speed: 3× faster than NSA | Helped |
| add memory reset during training (ablation) | RAMba · RULER, longer-context retrieval | accuracy | Helped |
| add memory reset during training (ablation) | RAMba · RULER, in-domain context length | accuracy: 92.76 -> 91.74; 87.10 -> 80.61 | Hurt |
Surprising findings
- Single-NIAH extrapolates only to 4M context while the simpler passkey template extends to 64M, suggesting template simplicity drives generalization.
- Mamba-based models with attention show higher perplexity at 64K than at 16K without memory reset, suggesting Mamba's memory state hurts attention length generalization.
Limitations
- Experiments focus mainly on Mamba; applicability to other RNNs is theoretical.
- Models larger than 3B parameters were not evaluated due to compute constraints.
- Precise chunk selection for extremely long contexts remains an open challenge; retrieval degrades beyond 1M on harder RULER tasks.
Experimental setup
Datasets: The Pile PG19 ArXiv-math Code RULER synthetic data LongBench V2 XSUM CNN/DailyMail SQuAD HotpotQA QuALITY
Benchmarks: PG19 ArXiv-math Code passkey retrieval RULER LongBench V2 XSUM CNN SQuAD HotpotQA QuALITY
Metrics: perplexity accuracy ROUGE F1 exact match throughput (tokens/s) runtime (ms) memory footprint inference time per 100 tokens
Baselines: Transformer (full attention, YaRN) Mamba-2 Mamba-2 w/ sliding window attention (ALiBi, RoPE) Mamba-2 w/ NSA FlashAttention-2
Builds on: Mamba-2 Native Sparse Attention MoBA stick-breaking attention GCA CEPE
Closest papers by meaning
- RAT: Bridging RNN Efficiency and Attention Accuracy via Chunk-based Sequence Modeling (NeurIPS 2025 · similarity 0.89)
- RAM-Net: Linear-Time Sequence Modeling with Sparsely Addressable State (NeurIPS 2026 · similarity 0.89)
- Efficient Length-Generalizable Attention via Causal Retrieval for Long-Context Language Modeling (ICML 2025 · similarity 0.89)
- Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps (NeurIPS 2026 · similarity 0.88)
- A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention (ICLR 2025 · similarity 0.87)
- Sparse Delta Memory: Scaling the State of Linear RNNs through Sparsity (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.