PaperlayerView exact revision on arXiv

Source-linked AI summary

Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access

Xiang Hu, Jiaqi Leng, Jun Zhao, Kewei Tu, Wei Wu

arXiv:2504.16795v2cs.CLcs.AI

TL;DR

Transformers face quadratic self-attention costs on long sequences, while chunk-based sparse attention can select chunks inaccurately. HSA learns token-to-chunk relevance through hierarchical attention, and RAMba combines HSA with Mamba. RAMba outperforms baselines across long-context modeling and downstream tasks and achieves perfect passkey retrieval accuracy at 64M context.

  • Problem

    Transformers incur quadratic self-attention costs on long sequences, while chunk-based sparse attention can select chunks inaccurately, including at longer out-of-domain contexts.

    Source: §1 · pp. 1–2

  • Method

    HSA learns token-to-chunk relevance through two-stage chunk-level attention and weighted fusion; RAMba integrates HSA into selected Mamba-2 layers.

    Source: §1 · p. 2

  • Results

    RAMba consistently outperforms baselines in long-context modeling and downstream tasks, and is the first Mamba-based model with perfect passkey retrieval accuracy at 64M context.

    Source: §1 · p. 2

  • Takeaways & Limitations

    RAMba offers a foundation for language models with permanent memory, balancing performance, efficiency, long-range access, and length generalization.

    Source: §5 · p. 10

    Experiments mainly focus on Mamba, and models larger than 3B parameters are not evaluated.

    Source: Limitations · p. 17

Abstract

from arXiv · show

A key advantage of Recurrent Neural Networks (RNNs) over Transformers is their linear computational and space complexity enables faster training and inference for long sequences. However, RNNs are fundamentally unable to randomly access historical context, and simply integrating attention mechanisms may undermine their efficiency advantages. To overcome this limitation, we propose Hierarchical Sparse Attention (HSA), a novel attention mechanism that enhances RNNs with long-range random access flexibility while preserving their merits in efficiency and length generalization. HSA divides inputs into chunks, selects the top-$k$ chunks and hierarchically aggregates information. The core innovation lies in learning token-to-chunk relevance based on fine-grained token-level information inside each chunk. This approach enhances the precision of chunk selection across both in-domain and out-of-domain context lengths. To make HSA efficient, we further introduce a hardware-aligned kernel design. By combining HSA with Mamba, we introduce RAMba, which achieves perfect accuracy in passkey retrieval across 64 million contexts despite pre-training on only 4K-length contexts, and significant improvements on various downstream tasks, with nearly constant memory footprint. These results show RAMba's huge potential in long-context modeling.

1 Introduction

RNNs offer efficient sequential processing but lack random access, while adding attention can compromise their efficiency and length generalization. HSA and RAMba address this trade-off with chunk-aware retrieval and a hardware-aligned design.

1 Introduction

  • RNNs process sequences efficiently but lack historical random access, while attention can bring quadratic compute, memory overhead, and poor length extrapolation.These constraints are especially consequential for passkey retrieval, where performance degrades as sequence length increases.Source: §1 · p. 1
  • HSA selects chunks using token-to-chunk relevance learned from whole-chunk feedback, then attends within selected chunks and fuses their information with chunk weights.This chunk-aware selection remains accurate in passkey retrieval at context lengths over 10,000 times the pretraining length; a hardware-aligned algorithm supports efficient parallel computation.Source: §1 · p. 2
  • RAMba integrates HSA into Mamba and offloads token-level KV cache to CPU memory, keeping compact chunk representations on GPU and loading selected chunks as needed.Sharing an intermediate-layer KV cache across subsequent HSA layers reduces memory swaps to one chunk selection and swap per step.Source: §1 · p. 2
  • RAMba achieves perfect passkey-retrieval accuracy at 64M context; HSA is 3× faster than NSA and 5–25× faster than full attention at 16K tokens or more, and RAMba outperforms baselines downstream.The speed comparisons refer to the forward pass.Source: §1 · p. 2

2 Related Works

Prior work reduces attention cost or extends context access through local windows, clustering, cache eviction, and chunk selection, each with trade-offs. This paper positions HSA against these approaches through its learnable chunk retrieval and token-level attention within chunks.

Sparse Attentions

  • Local-window attention [9] [14] limits each token to nearby context, while clustering [64] and cache eviction trade quality or arbitrary access for efficiency.Combiner compresses each chunk to one key-value pair, sacrificing token-level information.Source: §2 · pp. 2–3
  • NSA and MoBA select chunks using summed token-level attention scores, but can struggle to identify important chunks accurately.HSA instead learns token-to-chunk relevance while preserving token-level information within chunks.Source: §2 · p. 3

RNN-base Language Models

  • RNN-based models retain constant per-token inference cost as sequences grow, while hybrid approaches combine recurrent long-range memory with attention for short-range patterns.HSA uses RNNs to retain short-range information and sparse attention to access arbitrary long-range context; GCA supports 16M retrieval contexts but retrieves only once every S tokens, usually 64.Source: §2 · p. 3

3 Methodology

RAMba pairs an RNN working-memory representation with HSA as an extendable memory for retrieving long-range context. Its method combines chunk-based architecture, hierarchical attention, hardware-aligned kernels, and training and inference optimizations.

3 Methodology

  • RAMba uses Mamba to compress context into a finite working-memory representation and HSA to retrieve from an extendable KV cache as long-term memory.The paper describes HSA as weighting tokens within chunks and then weighting across chunks.Source: §3 · p. 3

3.1 Model Architecture

  • The architecture encodes chunk memories independently, then alternates HSA and Mamba layers in the upper decoder; a selection layer chooses relevant chunks for long-range attention.The lower decoder's outputs are divided into chunks and passed through a bidirectional Transformer encoder to form chunk memories.Source: §3.1 · pp. 3–4

3.1.1 Hierarchical Sparse Attention

  • HSA selects top-K past chunks per query group, assigns chunk weights, attends separately to tokens within each selected chunk, and fuses the resulting representations.Its learnable chunk retrieval module is the primary distinction identified relative to MoBA and MiniCPM4.Source: §3.1.1 · pp. 4–5, §3.1 · p. 4

3.1.2 Kernel Design

  • Hardware-aligned Triton kernels load each step's query and selected-chunk KV pairs for attention; the backward pass accumulates gradients for Q and weights before K and V.The design addresses the substantial memory footprint of naive per-token chunk selection.Source: §3.1.2 · p. 5

3.2 Training & Inference

  • Training uses random segment resets to improve length generalization; inference offloads chunk KV to CPU and retrieves K chunks per step, while chunk-selection representations still grow with sequence length.The paper says those representations are sufficiently small in practice that offloading them is unnecessary.Source: §3.2 · p. 6

4 Experiments

The experiments pre-train 370M models on 4K contexts and evaluate language modeling, downstream tasks, and efficiency. RAMba achieves perfect passkey retrieval accuracy at 64M context, while the efficiency results report near-constant inference memory usage.

4.1 Setups

  • The 370M models are pre-trained from scratch on 4K contexts using the same 60-billion-token Pile subset, with HSA chunks of 64 tokens and eight selected chunks.HSA uses a single-layer bidirectional Transformer chunk encoder and is inserted every eight Mamba layers.Source: §4.1 · p. 6

4.2 Long-range Langauge Modeling

  • RAMba performs better than baselines across most language-modeling datasets at in-domain 4K and out-of-domain 16K and 64K lengths.With memory reset, Mamba variants using NSA or HSA extrapolate more strongly but lose in-domain performance.Source: §4.2 · p. 7, Table 1 · p. 7

4.3 Downstream Tasks

  • At 64M context, RAMba achieves perfect passkey retrieval accuracy, while most baselines fall to nearly zero around 64K.Removing the chunk encoder reduces RAMba’s accuracy as context length grows.Source: §4.3 · p. 7, Figure 3 · p. 8
  • On RULER retrieval tasks, HSA performs comparably to full attention in-domain and better than NSA; Mamba extrapolates to 256× pre-training length on FWE.Memory reset improves extrapolation but reduces in-domain performance; precise chunk selection at extremely long contexts remains open.Source: §4.3 · p. 8, Table 3 · p. 8
  • RAMba outperforms in most fine-tuned downstream tasks and maintains a significant lead among 2.7B models.Mamba with sliding-window attention has an advantage on SQuaD, whose texts are shorter than the window.Source: §4.3 · pp. 8–9, Table 4 · p. 8, Table 5 · p. 9

4.4 Efficiency Analysis

  • RAMba reaches 90% of Mamba’s training throughput, and HSA is faster than NSA through shared chunk selection, with near-constant inference memory.The runtime comparison uses three attention layers; HSA shares one chunk selection across them, and the additional encoder can offset speed gains at some scales.Source: §4.4 · pp. 9–10, Figure 4 · p. 9, Table 6 · p. 9

5 Conclusion

The paper presents HSA and builds RAMba by combining an RNN backbone, length-generalizable sparse attention, and a forgetting mechanism. It describes this architecture as balancing performance, efficiency, long-range access, and length generalization, and as a foundation for language models with permanent memory.

5 Conclusion

  • RAMba combines an RNN backbone, length-generalizable sparse attention, and a simple forgetting mechanism.Source: §5 · p. 10
  • The architecture is presented as balancing performance, efficiency, long-range access, and length generalization.Source: §5 · p. 10
  • The authors offer RAMba as a foundation for language models with permanent memory.Source: §5 · p. 10

Limitations

Although the method is theoretically applicable to all RNNs, experiments mainly focus on Mamba and exclude models larger than 3B parameters because of computational-resource constraints.

Limitations

  • The method is theoretically applicable to all RNNs, but experiments mainly focus on Mamba.Source: Limitations · p. 17
  • The experiments do not establish performance across RNN architectures beyond the Mamba-focused setting.Source: Limitations · p. 17
  • Models larger than 3B parameters are not discussed because of computational-resource constraints.Source: Limitations · p. 17

A How HSA achieves accurate chunk retrieval

The appendix explains why approximating chunk relevance can select the wrong chunks, and how HSA instead learns token-to-chunk relevance through hierarchical attention. This lets chunk contributions inform the weights used in the forward pass and their gradients.

A How HSA achieves accurate chunk retrieval

B Hyper-parameters

Table 8 is titled as a listing of hyperparameters for 370M models.

B Hyper-parameters

C Training hyper-parameters

The 370M models used AdamW, with a cosine-decay learning-rate schedule and 60B training tokens. The listed settings also specify batch size, gradient clipping, and linear biases.

C Training hyper-parameters

D Templates for passkey retrieval and Single-NIAH tasks

The appendix specifies the passkey and Single-NIAH prompt templates and gives architectural settings for RAMba 2.7B. HSA is inserted at regular intervals in the Mamba backbone.

D Templates for passkey retrieval and Single-NIAH tasks

E RAMba 2.7B

  • RAMba 2.7B follows Mamba-2 2.7B’s embedding dimension of 2560 and uses 64 layers.Source: E RAMba 2.7B · p. 19
  • Its two-layer Transformer-based encoder has hidden size 2560 and intermediate dimension 3392.Source: E RAMba 2.7B · p. 19
  • HSA starts at Mamba layer 32, with one HSA layer inserted every eight layers.Source: E RAMba 2.7B · p. 19
  • HSA uses group size 5 and query-head size 16 per group, adding approximately 110M parameters.Source: E RAMba 2.7B · p. 19

E.1 Post-training

The appendix describes post-training procedures for RAMba and a separate synthetic-data evaluation setup for RULER. It also states that the abstract and introduction’s claims are supported by Section 4 experiments.

E.1 Post-training

  • The Mamba-2 2.7B base model was post-tuned with BPTT on 32K-token sequences for 3K steps, totaling 1.5B tokens.This stage took five hours on 32 PPUs.Source: §E.1 · p. 19
  • After freezing most Mamba backbone parameters, the HSA modules were trained in a warmup stage and a post-training stage.Source: §E.1 · p. 19
  • Warmup trained only HSA-related parameters to copy the first half of each sentence from an appended duplicate, using 32K contexts for 16K steps.The stage totaled 8B tokens and took around 24 hours on 32 PPUs.Source: §E.1 · p. 20
  • Post-training used LoRA to fine-tune 5% of Mamba-2 parameters while keeping all HSA parameters trainable, for 32K steps on 32K contexts.This stage totaled 16B tokens and took around 48 hours on 32 PPUs.Source: §E.1 · p. 20
  • The post-trained model described in this section is the model used for LongBench evaluation in Table 5.Source: §E.1 · p. 20

E.2 RULER finetuning

  • For RULER, RAMba was trained from scratch on 60B tokens, fine-tuned on 1B synthetic tokens, and compared with Mamba-2 2.7B fine-tuned on the same data.The stated motivation was to test stable convergence and long-range retrieval on synthetic tasks such as NIAH and FWE.Source: §E.2 · p. 20

NeurIPS Paper Checklist

2. Limitations

The checklist reports that code was not openly available at submission, although its release was planned after de-anonymization. It also notes that single checkpoints prevented error-bar evaluation.

2. Limitations

  • The checklist marks open access to data and code as unavailable at submission, with code to be released after de-anonymization.It says data sources are explained in the paper.Source: §2 · p. 22
  • The authors mark experimental details as specified and say the reported results are reproducible from disclosed information.They point to Section 4.1 and Appendices A, B, and D for experiment settings.Source: §2 · pp. 21–22
  • The authors report that each pre-trained model has a single checkpoint, making error-bar evaluation infeasible.They state that multiple datasets or task sets and varied metrics provide support for assessing statistical significance.Source: §2 · p. 23

10. Broader impacts

  • The broader-impacts checklist response is marked not applicable.Source: §10 · p. 24

11. Safeguards

  • The safeguards checklist response for high-risk releases is marked not applicable.Source: §11 · pp. 24–25

12. Licenses for existing assets; 13. New assets

  • The checklist says existing assets are credited in Section 4, while new code assets will be documented after de-anonymization.Source: §12 · p. 25, §13 · p. 25

14. Crowdsourcing and research with human subjects; 16. Declaration of LLM usage

  • The authors state that the work involved neither crowdsourcing nor human subjects, and that its core method did not involve LLM usage.Source: §14 · pp. 25–26, §16 · p. 26

Related papers

Loading 2504.16795v2…