Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access
Authors: Xiang Hu
Jiaqi Leng
Jun Zhao
Kewei Tu
Wei Wu
Venue: NeurIPS 2025 poster ()
Keywords: long-context modeling
language modeling
attention mechanism
sparse attention
RNN
Primary area: deep_learning
TL;DR: A sparse attention mechanism balances efficiency, long-range random access flexibility and length generalization ability

Abstract: 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 \textbf{H}ierarchical \textbf{S}parse \textbf{A}ttention (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, selecting 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.

Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access
Xiang Hu', Jiaqi Leng", , Jun Zhao', Kewei Tu Wei Wu
Ant Group, FFudan University, "ShanghaiTech University
{aaron hx, congyue... wanteewwwwwa.tt.nuuoooom. com, tukw@shanghaitech.. com
com httpaa/rooooaae/ooooooooooooooo-
context-
Abstract
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-chunl. relevance based on fine- ane-grained token-level information inside each chunk. This approach enhances the precision of chunk selection across both in in-domain and out-of-domain context lengths. To make HSA efficient, we further introduce hardware-aligne kernel design. By combining HSA with Mamba, we introduce RAMba, which achieves perfect accuracy 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 's huge potential in long-context modeling.
Introduction
The success of Large Language Models (LLMs) 59 has been largely driven by the Transformer architecture [60]. However, the quadratic computational and memory costs of self-attention make inefficient for processing long sequences. Moreover, Transformers often struggle with inputs that exceed their pre-training length. These limitations have renewed interest in alternative architec- tures such as Recurrent Neural Networks (RNNs) [10 24 28 33) that enable efficient, finear-tima processing sequential data while retaining degree of extrapolation capability.
However, RNN-hased models suffer from critical limitation: the information bottleneck 58 caused by compressing variable-length contexts into dime--iieeni representations. Unlike attention mechanisms, they lack random access to contextual information, which becomes especially problem- atic in tasks like passkey retrieval, where performance degrades as sequence length increases (62]. While augmenting RNNs with attention mechanisms can help mitigate this limitation, it introduces drawbacks such as poor length extrapolation, quadratic computational complexity, and substantial memory overhead during inference, ultimately undermining the original efficiency advantages of RNNs. Consequently, there remains no satisfactory RNN- solution that can simultancously achieve length generalization, random-access flexibility, and efficiency.
To address the trilemma, we propose a novel Hierarchical Sparse Attention (HSA) mechanism. Existing sparse attentions like NSA and MoBA 32 70 typically divide sequences into chunks,
"Corresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

allowing each token to attend to concatenation of k selected chunks, thus can potentially achieve both efficiency and random access. While this design offers a promising step toward resolving the trilemma, a closer examination reveals a critical weakness: these methods often suffer from inaccurate chunk selection, both within the training distribution and when generalizing to longer, out-of-domain contexts. As illustrated in Figure TA), the issue may stem from learning token-to-chunk relevance in a chunk-unawan way, relying on token-to-token gradients instead chunk-level feedback.
Based on this insight, HSA introduces a two-stage hierarchical mechanism for the selected chunks enable end-to-end learning of token-to-chune vance. As shown in Figure B)), in the first stage, it applies attention separately to the tokens within each chunk to capture chunk-level information. the second stage, it fuses the chunk-level informa- tion by applying weighted summation using token- to-chunk weights. During the backpropagation pass, token-to-chunk weights are adjusted based on the contribution of the entire chunk to the next token prediction, achieving Shunk-aware learning. Experimental results show that this chunk selee- tion mechanism remains accurate even when con- text lengths exceed pre-training lengths by over 10,000 times in passkey retrieval. Since each to- ken corresponds to different &o chunks, naive im- plementation would require substantial memory. Thus, we further propose a hardware-aligned algorithm achieve efficient parallel computation.
(A)
rele-
netionnal QU k., M
concattration
Q
K
top k
5,7
Q
softmax
top-k
Way Wy
Tex chunk
5-th chunk
7-th ChARTA
Inguul
21 22 25 X
Figure 1: K,, V; are the i-th chunk's 's key and value, with K; the mean pooling of K, In (A), the chunk se- lection scores Q K. are learned from token-to-token interactions (chunk-unaware). In HSA (B), Q K; are guided by the feedback from the entire chunk (chunk-aware), with O,, the chunk-level information obtained from the i-th chunk by the f-th token.
With HSA, we propose Random Access Mamba (RAMba), an extension of Mamba-2 [6] that incorporates HSA into specific layers. To maintain a constant memory footprint during inference, RAMba offloads the token-level key-value (KV) cache to CPU memory while retaining compact chunk-level representations on the GPU for efficient chunk selection. At each time step. only the KV cache of the selected chunks is loaded onto the GPU, ensuring efficient memory usage. To further minimize GPU-CPU memory swaps, RAMba leverages the hidden states derived from an intermediate layer as shared KV cache for all subsequent HSA layers, requiring just one chunk selection and swap per step. This architecture efficiency, length generalization, and long-range random access flexibility.
simultaneously balance training and inference
In our experiments, we compare RAMba with baselines such as Transformers, Mamba-2, and their variants with sliding window attention [14] and NSA (70), evaluating performance across long-range language modeling, downstream tasks, and efficiency. RAMba consistently outperforms the baselines in long-context modeling and downstream tasks while exhibiting exceptional length generalization. Notably, it is the first Mamba-based model to achieve perfect accuracy passkey retrieval task. In terms of efficiency, HSA is 3x faster than NSA and 5-25x faster than full attention for contexts of 16K tokens more during the forward pass. Additionally, when memory offloading is enabled, RAMba maintains nearly constant memory usage. These results demonstrate RAMba's superior capability in long-text modeling. In summary, our contributions are threefold:
64M context in the
1. We propose HSA, a novel hierarchical attention mechanism paired with a hardware-efficient algorithm that simultaneously enables efficiency, length generalization, and flexible long-range random access.
2. Based on HSA, we introduce RAMba, which integrates the advantages of the attention mechanism into Mamba while maintaining a nearly constant memory footprint during inference.
3. We conducted comprehensive experiments on the length generalization of Mamba with various attention mechanisms. The results show that HSA excels in both performance and efficiency.
2 Related Works
Sparse Attentions. Sparse attention aims to reduce computational complexity by focusing limited number of tokens. For example, sliding window attention and its variants [9] 14 41 7 restrict computations to fixed-size local window for each token. Such methods often sacrifice

---

the ability to capture long-ram onngraange information. Clustering based approaches [30 52 61 64 employ mecalty-sensitivv hashing or K-Means for token clustering and perform attention with clusters. These methods often trade off efficiency for quality due to the limited accuracy of clustering. Cache eviction maintain constant memory costs by retaining only the most important tokens in the KV cache, but this limits the model's ability to access arbitrary contexts randomly. Combiner [50] utilizes a hierarchical attention mechanism by compressing tokens within chunks into single key-value pairs via max-pooling. However, this approach sacrifices the ability to capture token-level information inside the chunk, resulting in reduced accuracy. Recently, NSA [70] and MoBA [32 achieve sparsity by dividing a sequence into chunks and dynamically selecting relevant chunks for each token based on summed token-level attention scores. However, these approaches struggle with accurately identifying important chunks. Our work introduces two-stage hierarchical attention mechanism with sparse chunk selection, achieving end-to-end learning of token-to-chunk relevance, while maintaining the ability to capture token-level information, resulting in attention flexibility, accuracy, and efficiency.
approaches [21 31 74
RNN-base Language Models. RNN-based language models have recently gained increased atten- tion because their per-token inference cost remains constant as the sequence length grows. Linear Attention [28] shows that replacing softmax attention with kernel-basee approximations allows Trans- formers be reformulated as RNNs, achieving similar performance while benefiting from recurrent properties. Recent advancements in RNNs, such as RWKV series
442 43), Mamba series [1 23)
x1 STST & and others 66 73 continue to push the boundaries of this approach. Our work builds on these RNN models and is in principle applicable across different RNN architectures. Meanwhile, other studies [88 37 51 explore hybrid architectures using RNNs for long-range memory and attention for short-range patterns. In contrast, our approach leverages RNNs to adaptively retain essential short-range information for next-token prediction, while sparse attention selectively accesses raagerange context at arbitrary positions.
GLA 666, HGRN series 44 47
Context Length Extrapolation. The full attention mechanism struggles generalize contexts longer than those observed during pretraining [45]. even with techniques like relative positional encoding
(54) or attention weight scaling (63] While methods such as Landmark Attention (35)
and DCA 3 have been proposed to address this limitation, they still encounter a sharp rise in
perplexity when extrapolating beyond a certain length, typically 32x. Recently, GCA (27) achieves perfect accuracy on 16M context lengths despite being pre-trained on only 4K lengths by integrating end-to-end retrieval within attention. However, GCA limits retrieval once every S (usually 64)
tokens, reducing its flexibility. The two-stage attention mechanism in HSA is inspired by GCA, but achieving per-token retrieval through kernel optimization.
3
Methodology
From the psychological perspective, human memory primarily comprises two main systems: working memory (6. which temporarily stores and manipulates information with limited capacity, and long- term memory (5). which stores information indefinitely with virtually unlimited capacity. Inspired by this, RAMba applies Mamba simulate working memory by compressing variable-length contexts into a finite representation for manipulation. Mcanwhile, it uses HSA to model long-term memory as an extendable KV cache, enabling efficient retrieval and attention computation. In HSA, the key innovation lies in its two-stage hierarchical attention mechanism: first, weighting over tokens within a chunk, and then weighting over chunks. By offloading KV cache to CPU memory or disk, it can theoretically maintain unlimited memory. In the following sections, we elaborate the model architecture, HSA, kernel design, and optimizations for training and inference.
3.1 Model Architecture
As shown in Figure 2 a), RAMba contains I. Mamba layers, equally divided into upper and lower decoders, akin to previous works (2] 53 Given the input sequence {x1,..,, 2.), their corresponding embeddings are fed into the lower decoder and the outputs are divided into [ chunks according to chunk size S. These chunked representations are then passed through a Transformer-basee bi- directional encoder independently to form chunked memories, used for chunk selection and sparse attention in subsequent HSA layers. The upper decoder takes the outputs of the lower decoder as inputs, and alternates between one HSA layer and G Mamba layers for L times. Between the
3

---

Oulput to HBM
Upper
E
XG
AL To SRAM
2G
Chunk Encoder
HSA
Grid Loop
Attn
ViA
Krtt
Chunk Selection
decoder
da
Mamba
select
22
Inner Loop
(b)
Figure (a) Model architecture for RAMba. (b) Kernel design for HSA.
lower and upper decoders, chunk selection layer selects the most relevant k chunks based on the dot-product similarity between token chunk representations at each decoding step. HSA layers capture long-range dependencies by attending to the selected chunks.
Notations. We focus on architectures like GQA 22,. where each group of h query heads shares common KV head, and selects chunks independently. The hidden size d is divided into groups, satisfying g x dg d, with dg dimension per head. We denote the m-th layer output as H Ix with the f-step representation H. E Rdd and the encoded token representations of the i-th chunk E, E Rxdd To simplify the notation, use as generic symbol to represent one of the groups, as each group behaves in the same way. For example, H E iRd, denotes one of the grouped representations.
dn. Here, dg is the dimension of each group, and d₁ the
Chunk Selection. At each time step token uses the representation derived from the lower decoder to select top-k past chunks. Each query group independently selects chunks, which are shared across / heads within the group. Formally, we have:
Q w (H, K
Be, SQ
K ddg
i. rank(( K).
W
-00,
where norm RMSNorm Wale Wale Raxd conduct linear transformations, rank(-) denotes the ranking position in descending order, E, E gd is the mean-pooled representations of E. dt K.de E Rd are representations used for chunk selection, with Qu K E IR as their grouped representations. For each group, št, is the relevance score of to the i-th chunk, and Σ, is the selected K chunk indices for Zt.
Q
3.1.1 Hierarchical Sparse Attention
Chunk Weights. To avoid the impact of position encoding on length generalization 29 we discard it and instead model the ordinal relationships of distances akin to the stick breckbreaking attention 55). In the stick-breaking process, participants sequentially take portion from stick, with later participants dividing the remaining portion left by earlier ones, thus introducing a "most recent" bias. Let the total weights serve as the stick and Wt, denote the weight assigned to the k-th chunk selected by the f-th token. For each group, we have:
z' sort(ī.), w.sh B.. [I (1---1)) Brs)) 0
[[ - 1[1,-0(5,22)) tck
where 3r, is the proportion taken from the remaining attention, calculated via the sigmoid function a. The sort(.) function rearranges t, in descending order to prioritize chunks closer to I; in attention allocation. z denotes the the k-th value T.-
Hierarchical Attention. After obtaining the weights for each chunk, 2; performs attention with tokens in each retrieved chunk separately. We denote the representation of information collected from the k-th selected chunk as Or E 1R These representations are then fused using the chunk weights

---

Formally, for the 1-th HSA layer, we have:
Q = We norm(h)
K. WAE.
V. WvE, = +
, atin Q R V x
Σ
di O
& &K
attention group
the comnatenation of all groups
chunt:
group
where W E Rddd WK, Wy E (Rxx())d)) apply linear transform to obtain query, key, and value representations, with h query head shares the same KV head. K,, V, E RSx((x)a) are the key and value corresponding to the i-th chunk, while K,, V. E RSSxd are their grouped representations, shared across all HSA layers. Q is the query representation of x4 at the 1-th layer.
The key Difference with MoBA/NSA. The primary distinction of HSA lies in the learnable chunk retrieval module. To illustrate the difference between existing sparse attentions, such as in MoBA [32] and MiniCPM4 [56], we provide detailed analysis with specific examples in the Appendix
3.1.2 Kernel Design
In HSA, each token corresponds to distinct set of K chunks, which can lead to substantial memory footprint naive implementation. Inspired by NSA, we address this issue by implementing hurdware-aligned HSA kemnlls based Triton [57]. As illustrated in F each GPU thread loads the query representations for a single step along with the selected chunks KV pairs for attention computation. Algorithm demonstrates a parallel forward pass for group of query heads, with multiple groups processed in parallel following the same approach. Particularly, we use softmax- off-by-one (34] to allow tokens in the current chunk to ignore any retrieved tokens. During the back- prokkpropagttion process, we adopt a two-phased backward pass inspired by Dao [[].. The first stage accumulates gradients for and w. followed by K and V in the second stage. The pseudo-code for this process shown in Algorithm 2 and 3 which M,,, E (0, 1 $1 denotes whether the f-th
token selects the i-th chunk, and RG, E 2 2 represents the index of the i-th chunk in I/-
Algorithm 1 FORWARD thread t
I: O +
Initialize O2 E RAxd, 1/ load Qq to Static RAM (SRAM), Q E T Rlxhxdd Q' E Rhxdd
2: Q' + load Q:
3: for < k < K do
load It,,, w +- load w 5: K' load K; V' +- load V;
4:
#IZZIK E Bixk HK. E Ixxxdd K' V B I1 Inter-chunk token-level attention, no online softmax required. I/ Chunk-level attention via weighted sum.
6: O' (Q K"" yv'
O2 O2 + wO'
end for
9: O + write to O,
// Write to High Bandwidth Memory (HBM) from Static RAM (SRAM).
Algorithm 2 BBACWWARD Q, w thread
Algorithm BACKWARD-K, thread
VQ' + 0, Q'Q'= < load Q,, VO' +- load VO.
K' V' <- load K, V; VK', VV
I/K V E ERSxda /VK' VV VV'' 9xda
for < k < K do
- load Ir,h. +- load wi.k
for
l do
K', V + load K; V; // K', V
E 8xda I/P
if M, is true then
P < softmax, (Q (Q'K"))
O/ dhxd hxs
k + load Rr.i. w + load Wt,i
O' PV'
Q' < load Q
1/ Rlxdd
40') /D E I²,, VO' hxdd
D' rowsum(O'
pointwise multiply
VO + load VO,
#V0 R ghxdn
D' + load Dr.k
// D' E
D' + write to D
I/D E RlxXxh 1/ Vw' ER // /ww RIxK
softmax / (Q'K"T)
//P E TRAXS
Vw' &- rows
VV + VV' +MP VO'
Vw' write VWEs
VP t VO'V'T
/VP E IgAxs //S hgx
VP < VO'V"T
// /VP ghxs
VS wP (VP- D')
VS wP (VP- -D'
// VS hxs
VK' + VK' + V
VQ. < VQ. + VSK'
// //Q. E R hxdd
end if
end for
VK', VV write VK, VV.
VQ' write VQ,
end
5

---

3.2 Training & Inference
Training. Many efforts have been made to address Mamba's length generalization issues. A simple yet effective solution is truncated backpropagation through time (BPTT) [3 66], in which the first state of a sequence is initialized as the final state of its preceding sequence. We follow this approach in training RAMba and other baselines. However, even though RAMba demonstrates certain extrapolation capabilities, we still observe an increase in perplexity on longer contexts. We hypothesize that RNNs' memory state provides certain shortcuts [22] for long-range attention, degrading performance contexts largely exceeding the pre-trained length. Thus attempt to introduce an appropriate forgetting mechanism into the memory state to disrupt the shortcuts. A straightforward way is memory reset, where sequences are divided into equal segments, and the initial state of each segment is reset to zero. To align with BPTT, we set the initial state as the last hidden state of random segment in the previous step. In other words, for RNNs, the previous segment is randomly replaced, while the attention mechanism retains access to the original context. Our experiments show that this method effectively improves length generalization for RAMba.
Inference. A key challenge in introducing attention mechanisms to RNNs is managing the mem- ory footprint, as the KV cache scales linearly with sequence length. Prior works [27 35] have demonstrated the feasibility of offloading the KV cache CPU memory and selectively loading chunks during inference. For prompt of length L, all chunk representations are offloaded to the
CPU after prefilling, while only Kale E R xx is kept in GPU memory for chunk selection. During decoding, RAMba retrieves and loads chunks at each step for each group. Since the KV cache is shared across all HSA layers, the number of parameters exchanged between the CPU and GPU totals g X de x K x S. Our efficiency analysis in $ 4.4 demonstrates the overhead for memory exchange is fully acceptable. Since the memory footprint of kafe still increases with the sequence length, to theoretically achieve constant memory, a straightforward way is to offload it to a FAISS [19 database. However, in practice, the memory footprint of K is very limited, making such offloading unnecessary. Detailed analysis is elaborated in the experiments section.
4 Experiments
4.1
Setups
To ensure fair comparison, we pre-train all 370M models from scratch with 4K context length to observe their performance and extrapolation capabilities across various tasks. For 2.7B models, the training details are presented in Appendix
Baselines. We adopt the Mamba-2 architecture as the backbone of the RNN model and YaRN [44] as the Transformer baseline. The parameter size of all models trained from scratch is 370M, with detailed parameters provided in Appendix [ We experiment with Mamba variants with different exeerimeet
attention mechanisms, including sliding window attention, native sparse attention (NSA), and HSA. For sliding window attention, the window size is set to 512, incorporating two position encoding schemes: ALiBi 45 and RoPE, the latter following the settings in Samba [5T]. We set the chunk size of HSA to 64 following NSA. To ensure that the field of view for sparse attention matches the sliding window size (64 * 8 = 512)) we set the number of selected chunks to For NSA, we use its efficient open-source implementation | To isolate the effects of the sparse attention components, we disable the sliding window attention in NSA. The HSA incorporates bi-directional encoder for chunk memory encoding, accounting for 5.4% of the total parameters, whose impact on fairness is minimal. HSA layers are inserted into the upper decoder every G = 8 Mamba layers, with other attention mechanisms like SWA and NSA following the same pattern. These settings remain consistent across all subsequent 370M models. Since the compressed attention in NSA functions similarly to Combiner [50], we do not conduct separate comparison against Combiner. Some other related works 36 are not included in the experiments due to the lack of open-source implementations.
linglellyered Transformer-
Pre-training. All models are pre-trained on the same 60-billion-toker subset of the Pile dataset 20 Detailed training hyper-parameters are provided in Appendix C

---

Pg
arxiv
code
Pg
arxiv
code
Pg
arxiv
code
Models(370M)
eval eval_een=ae Jen=4k
eval LaalJee-lll 16k
eval Jen_len=644
Transformer
18.61
4.23
3.28
539.15
199.42
62.17
>10
>104
2865.51
Mamba
17.92
4.24
3.28
17.38
3.91
3.09
17.30
3.86
3.05
w/ SWAAL.Bi
17.82
4.21
3.26
20.48
5.01
3.53
23.86
6.46
3.96
w/ SWAroo
17.82
4.21
3.26
17.50
4.03
3.19
17.80
4.25
3.35
w/ NSA
17.87
4.20
3.25
17.31
3.87
3.06
17.31
3.87
3.05
w/ NSAvfo
17.74
4.18
3.24
17.56
4.29
3.26
17.62
4.35
3.28
RAMbau/ m.r.
17.82
4.15
3.23
17.15
3.73
3.04
17.01
3.65
3.07
RAMbuuio m.r.
17.63
4.13
3.21
17.11
3.81
3.08
17.11
3.87
3.21
RAMbu m.r, s.
18.07
4.52
3.34
17.61
5.01
3.17
17.61
6.05
3.16
Table 1: Perplexity for long-range language modeling. We highlight the best results in bold and underline the second best. All models are pre-trained on 4K contexts.
Ablation Studies. We denote memory reset usage as w/ m.r and its absence as w/o m.r. Addition- ally, w/o s.b. refers using softmax without positional encoding instead of stick-breaking weights. For an 8K token context, the memory is reset every 4K tokens, which aligns with the context length of other baselines. However, the chunk selection scope for sparse attention spans 8K tokens, which might be unfair to other baselines. To ensure fair comparisons, we apply the same settings to both HSA and NSA.
4.2 Long-range Langauge Modeling
Datasets. We evaluate long-range language modeling on PG19 488 ArXiv-math
4 and Code 66
Results. As shown in Table ! RAMba performs better than the baselines across most datasets, both in-domain (4K) and oft-of-domain (16K & 64K) lengths. addition, we have two findings. First, Mamba's perplexity decreases the context length increases from 4K to 64K, whereas Mamba with attentions shows higher perplexity at 64K compared to 16K when not trained with memory reset. This suggests that Mamba's memory state may influence the generalization ability of attention mechanisms. Secondly, when trained with memory reset, Mamba with NSA or HSA all exhibit stronger length extrapolation capabilities but show decline in in-domain performance. This result validates the effectiveness of memory reset, which essentially constrains the model to rely entirely on the content of the text for retrieval, thereby preventing the model from learning shortcuts. Since the model becomes adapted to retrieving information from longer contexts, it is reasonable that its performance declines when handling shorter contexts.
4.3 Downstream Tasks
Tasks. We evaluate various models' long-context modeling abilities classic tasks like passkey retrieval 35 and the LongBench V2 dataset (7). To increase task difficulty, we replace numbers in passkey retrieval with random token sequences. Since passkey retrieval is relatively simple, we further fine-tuned the models using synthetic data following RULER (25). We use a context length of 4K for fine-tuning with total training step size equivalent to 5% of the pre-training stage. Evaluations were conducted across different lengths on four RULER tasks: Single NIAH (S-N), Multi-queries NIAH (MQ-N), Variable Tracking (VT), and Frequent Words Extraction (FWE). To align with passkey retrieval, keys in Single NIAH were also replaced with random token sequences. We adopt Cloze format for LongBench evaluation. following Waleffe al. [62. address the instruction cntruction-followinn following challenges of small models. Since LongBench V2 is zero-shot benchmark and thus small models may exhibit randomness, we additionally evaluate on fine-tunable datasets, including summarization tasks like XSUM [39 and CNN [88] and QA tasks like SQual) [49], HotpotQA [68) and QuALITY [4). NSA's implementation currently does not support generation, so its results on generative tasks are not reported.
Results. First, Figure 3 shows that RAMba achieves perfect accuracy in the passkey retrieval task with 64M context, even without memory reset, while most baselines drop to nearly zero around 64K length. However, removing the chunk encoder causes RAMba's accuracy to decrease as context length grows, highlighting the encoder's importance HSA and length generalization. One possible reason is that the hidden states primarily optimized for predicting the next token, thus
7

---

Model Perfomance by Context Size
(77els(770M)
Overall
Fasy
Hard
Short
Medium
Long
Transformeroot
24.8
28,5
Accuracy (%6
Mamba
22.4
16.4
26.1
19.6
22.7
26.3
SWAALBI
23.3
21.0
26.8
22.1
RoMbe
W/ SWAnge
24.2
25.7
23.2
20.3
26.3
26.3
NS AS
23.5
21.5
25.8
22.1
w/ NSAutome
22.1
26.3
19.6
17.1
26.3
22.1
RAMbe,
25.7
23.9
20.9
27.8
29,5
RAMb
23.5
26.3
21.7
20.9
23.7
27.4
Conteat Size (K]
Table 2: LongBench V2 results.
Figure 3: Passkey retrieval results.
(370M)
S-N
MQ-N ctx
VT
FWE
S-N
MQ-N
FWB
S-N
MQ-N
VT
FWE
S-N
FWB
ctx-len=44K
ctx-len-64K
ctx-len-256K
ctx-len-
Transformenutaan
95.08
88.59
97.12
60.02
0.00
0.00
0.00
0.00
0.00
0.00
0.90
0.00
2
96.66
1.86
61.78
64.75
37.31
0.00
0.00
6,25
0.00
40.00
w/ SWA
91.84
6.96
37.20
72.08
28.36
0.00
17.91
0.00
6.25
0.00
0.00
0.00
w/ NSA (wlo
92.58
39.09
57.70
66. 06.98
13:43
0.00
0.00
2.99
6.25
0.00
0.00
6.25
w/ NSA (w/ m.r.)
84.79
4,08
31,63
49.54
4,48
0.00
5.97
26.87
0.90
0.00
0.00
12,50
RAMbua
91.74
80.61
95.18
49.81
85.07
55.22
55.22
53.73
62.50
62.50
37,50
37.50
24.00
20.00
RAM
92.76
87.10
66
76.81
55.22
29.85
68.66
8.96
18.75
12.50
31.25
0.00
RAM
92.67
81.08
57.79
75.97
41.79
23.88
14,92
26.87
6.25
12.50
0.00
6.25
Table 3: Results for selected sub-tasks in RULER. All models are pre-trained on 4K contexts.
requiring adapter to extract information representing the current context. Second, LongBench's s evaluation results in Table 2 are generally consistent with those of language modeling. However, since LongBench is zero-shot benchmark, most baselines performed below the random guess rate of 25%, with the results across different sub-tasks show high variability. Thus, we believe the following SFT tasks can provide deeper insights into architectural capabilities for small models.
3 presents the evaluation results of various models fine-tuned on RULER synthetic data across different context lengths. We have four interesting findings. First, HSA demonstrates more precise chunk selection. HSA performs comparably to full-attention on retrieval-related tasks (S-N, MQ-N, VT) over in-domain context length, while NSA lags behind both, supporting our argument that estimating chunk importance using to-ee-to-token attention scores is inaccurate. Meanwhile, HSA achieves strong performance across retrieval tasks despite selecting chunks only once, validating the effectiveness of the hierarchical attention mechanism in learning token-to-chunk. relevance. Second, Mamba excels in sequential statistical tasks. While Mamba-based models underperform in retrieval tasks, they outperform in Frequent Word Extraction (FWH), successfully extrapolating up to 256x the pre-training length. This advantage likely stems from Mamba's ' continuous memory flow mechanism, which is also partially inherited by RAMba. Third, memory reset helps length generalization. Models with memory reset consistently outperform those without in length extrapolation. Though performance drops significantly beyond IM context length on more challenging retrieval tasks, it still achieves 256x extrapolation. Fourth, templates can make a significant difference. Single-NIAH extrapolates only to 4M, while the similar task of passkey retrieval extends to 64M, with the sole difference lying in the template of passkey retrieval being simpler. We elaborate the templates of these two tasks in Appendix [D] It suggests that a simpler pattern may facilitate learning more generalizable patterns. Still, precise chunk selection for extremely long contexts remains an open challenge for future work.
Models (370M)
XSUMR.18..R.L.
CNNR.1RA2RLL avg-len=883
SQuaDEwWFI wyg-len=174
HolpotQAAEMF) avg-len=1428
QuALITY Aue avg-len=7745
AVG.
4vg-len=498
Mambu-2 w/ SWALLBi
30.40/1 19//24553
37.97/1 07.97711.54/35.0
41.33/52.03
18.70/26.20
33.32
31.03
30.78/12. 10012.11//22 11/24.79
73 38.96/17
50.54/61.22 .23
19.96/27.5 .58
29.48
32.57
w/ SWArpp
30.88/12.33/24.99
18.78/17.67/36.5
51.46/61911
19.89/28.00
32.89
32.61
w/ NSA
31.40
Transformermii..oo
30. 10.10//1.63/24.25
38.25/17. 14
45.50/56.13
31.90/29.40
33.46
32.82
RAMba
30.81/12.25/24.80
39.11/18.04/36.85
48.24/59.17
22.30/30.53
34.13
33.64
'Table 4: Downstream task evaluations.
Table shows the results of the model after fine-tuning on summarization and generation tasks.
RAMba still outperforms in the vast majority of tasks, but two observations are worth noting. First, Mamba with SWA demonstrates advantages over other models on the SQual) task, One potential reason is that the SQuaD dataset contains shorter texts, even less than the sliding window, making
8

---

SWA able to randomly access all contexts. This phenomenon shows that integrating random access can improve RNN performance even for short contexts. Second, despite lagging behind Mamba in perplexity (PPL., the Transformer demonstrates strong downstream performance, ranking just behind RAMba. This further validates the importance of random access capability for downstream tasks.
Models
Aparams.
Overall
Easy
LongRench
Shot
Medium
Long
cts-loimSK
RRI sivp NYTEPA
Mrmmbs-2
27.3
89.247K 69/77.99
0.000
RAMbo
2.70 110M
27.3
25.1
28.5
30.4
27.3
22.1
21.00/6636666808 47722977
81.877.12/0000000.088
Table 5: Results of 2.7B models trained on 4K contexts.
Table Shoows the results of the 2.7B models. These results show that even after scaling up, RAMba still maintains a significant lead. Details about training can found Appendix
4.4 Efficiency Analysis
Experimental Setup. To evaluate the scalability and efficiency of HSA. measure the runtime and memory footprint of FlashAttention-22, NSA, and HSA operators across different sequence lengths, the throughput of various models different scales, and the per-token time consumption during inference. Runtime measurements were conducted with three attention layers only, excluding additional components such MLPs or Mamba. In HSA, a single chunk selection across all three layers, while SWA remains disabled in NSA. The memory footprint is reported as ratio of the memory occupied by Mamba 370M. We use w/ offloading and w/o offloading to denote whether memory offloading is enabled. To further evaluate the benefits of shared chunk selection, we introduce an ablation group (w/o sharing), where chunk selection and CPU-GPU memory exchanges happen in each HSA layer. When measuring training throughput, we enable FSDP 75 and gradient checkpointing [12), running models on Physies Processing Units (PPUs), each with approximately half the computational power of an A100 GPU.
shared
Forward Pass Performance Comparison
Backward Pass Performance Comparison
Comparison Nemory Footprint
HSA (fwd)
HSA (bwd)
Window Att offioad
(twd) flesivAttertion222 (fwd)
NSA (bwd)
Merory bost 0000000 0A00r 10me0
FlashAttention-2 (bwd)
Time (mis)
16K
32K
64K
16K
32K
64K
120K
Context Length
Context Length
Length
Figure 4: Comparison attention computation time: attention layers per group. (The lower the better)
Parameter Size
Prompt-L.ength
Models
Models
370M
780M
1.4B
3B
4K,
16K.
64K_
Mamba
16.44
9.71
6.53
3.46
Transformernut atm
2.26
8.90
32.12
SWA
15.64
9.24
6.27
3.31
Mamba-2 -
2.82
2.84
w/ full attn
10.00
5.73
4.31
2.07
RAMba
3.14
3.05
2.76
w/ NSA
14.35
8.70
5.46
3.18
RAMba,,
hlbba,//fllooling ofBooling
4.19
4.02
RAMba
14.80
8.76
5.70
3.18
RAMbu, ofBouding, wlo shuring
5.87
5.95
Table 6: Training throughput (10" tokens/s) with con-
Table 7: Inference time cost (seconds, prefilling time
text length=32K.
excluded) for generating tokens (batch-size=16)
Results. Figure compares the time consumption and memory footprint of three operators across different context lengths. Both NSA and HSA, as sparse attention mechanisms, significantly out- perform Flash-Attention in terms of speed. HSA is faster than NSA because it performs chunk selection- the only operation with quadratic complexity-one layers. In contrast, NSA involves computations with quadratic complexity like compressed token attention and chunk selection every layer, which increases time consumption. For memory foot- print, enabling memory offloading drastically reduces GPU memory usage and slows its growth with
once, and shares across HSA

---

increasing context length. Despite introducing CPU-GPU memory exchange, the impact on inference speed is limited, reported in Table 7 Table 6 shows that RAMba achieves 90% of Mamba's training throughput. Although HSA is faster than NSA, the additional encoder sometimes offsets this advantage at certain scales. Nonetheless, the results clearly highlight HSA's high efficiency, excellent scalability during training, and near-constant memory usage during inference.
Conclusion
In this work, we present Hierarchical Sparse Attention (HSA) and build on it to propose RAMba, which integrates an RNN backbone, seggt--eneraalzable sparse attention, and simple forgetting mechanism. This architecture strikes strong balance between performance, efficiency, long-range access, and length generalization, offering a foundation for language models with permanent memory.
10

---

References
[1] Josh Achiam, Steven Adler. Sandhini Agarwal, Lama Ahmad, Igee Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmirt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
[2] Joshua Ainslie, James Lee-Thorp, Michiel Jong, Yury Zemlyanskiy. Federico Lebron, and Sumit Sanghai. GQA: Training generalized multi-query transformer models from multi-head checkpoints. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023. URL https //openreview http://oruunnnddddnmmie
[3] Chenxin An, Fei Huang, Jun Zhang, Shansan Gong, Xipeng Qiu, Chang Zhou, and Lingpeng Kong. Training-free long-context scaling of large language models. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024, OpenReview.nct, 2024. URL
[4] Zhangir Azerbayev, Edward Ayers, and Bartosz Piotrowski. Proof A pre-training dataset of mathe- matical text, 2023, URL httL::stt/h/hggggggggee htt.://cc..cocatttttt///soaaaaaa center/proof pile
[5] Alan Baddeley. The episodic buffer: a new component of working memory? Trends in Cognitive Sciences, 4(11):417-423, 2000. ISSN 1364-6613. doi: https://doioorr/100001///1111111 16666600016611166666666666660001 3(0)) 538-2, URL.
[6] Alan D. Baddeley and Graham Hitch. Working memory. volume 8 of Psychology of Learning and Motivation, pages 47- 89, Academic Press, 1974, doi: https:://oi..or//0.. 10000000010001111111 142 1. URL http:tt/tt// wwtp://wwwww..eee......... com
[7] Yushi Bai, Shangqing Tu, Jiajie Zhang, Hao Peng, Xiaozii Wang, Xin Lv, Shulin Cao, Jiazheng Xu, Lei Hou, Yuxiao Dong. Jie Tang, and Juanzi Li. Longhench v2: Towards deeper understanding and reasoning realistic long-context multitasks, 2025. httt:/tttt////rrtttt org/ ab2/242.. 15204
[8] Maximilian Beck, Korbinian Pöppel. Markus Spanring, Andreas Auer. Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory, 2024. /ttpssss:///rxx org/ abs/ 2405. 2401104117
[9] Iz Beltagy, Matthew F. Peters, and Arman Cohan. Longformer: The long-document transformer. CoRR, abs/2004.051500 2020. URL hRRl/httta///aatt httpaa/////axv///eeeeaeee wttp:/oor/aabaaaaaaaaa 05150
[10] Yoshua Bengio, Patrice Y. Simard, and Paolo Frasconi. Learning long-temn dependencies with gradient descent is difficult. IEEE Trans. Neural Networks, 5(2):157-166, 1994. doi: 10.1109/72.27918.. URL
https://doi. org/10. 10. 1109/72. 279181
[11] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, 1lya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle. Mare' Aurelio Ranzato, Raia Hadsell, Maria- Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Process- ing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtual, 2020. URL httph: htap::/:///occeiings.. weeurips.cc/papeee 2atppper/2220/aaaa hash/ (ww.0008...c099777181118........
Abstract.html
[12] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. CoRR, abs/1604.061 74, 2016. http:///t///rrivvooo org/ abs// www.aaae/.......11
[13] Yingfa Chen, Xinrong Zhang, Shengding Hu, Xu Han, Zhiyuan Liu, and Maosong Sun. Stuffed mamba: State collapse and state capacity of ran-hased long-context modeling. CoRR, abs/2410.071445 2024, doi: 10. 100-85500AARXI.221.0077.
CRIl/htts:::/// //doi. org/10. 38550/arXiv 2410.07145
[14] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse 10509
transformers. CoRR, abs/1904. 10509, 2019. httpt:p//aaxiv.org...to
[15] Tri Dao. Flashattention-2. 22 Faster attention with better parallelism and work partitioning. In The Tweifth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. httpp:////ppppor.....eee///////p

---

[16] Tri Dao and Albert Gu. Transformers are ssms: generalized models and efficient algorithms through structured state space duality. In Proceedings the 41 International Conference Machine Learning. ICML'24, org, 2024.
[17] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, In Forty-first International Conference on Machine Learning. ICML 2024. Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. DRR. htw.pppnneereeee.eettttuummaaaa net/ffrummid- stn8FCR1td
[18] Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murarte Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, and Caglar Gulcehre. Griffin: Mixing gated linear recurrences with local attention for efficient language models, 2024. URL https
httgg/abs/24422 19427
[19] Mathiis Douze, Alexandr Guzhva, Chengqi Deng. Jeff Johnson, Gergely Szilvasy, Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. The faiss library. CoRR, abs/2401 .02281, 2024. doi: 10.48550/ARXIV.2400 008281. http::///////...... org/ 100 48550/ 38550/arXiv. arxiv 2401. .02881
[20] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe. Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Comnor Leahy. The pile: An 800gb dataset of diverse text for language modeling, 2020, httpa:tp::///
org/ abs/ 2ttp//111...m 00027
[21] Suyu Ge, Yuman Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive KV cache compression for llms. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL
htcps://openreviett wtt.:///ww....eeooooooooooo///ap
[22] Robert Geirhos, Jörn-Herrik Jacobsen, Claudio Michaelis, Richard S. Zemel, Wieland Brendel, Matthias Bethge, and Felix A. Wichmann. Shortcut learning in deep neural networks, Nat. Mach. In- rell., 2(11):665-667 2020. doi: 10. 1003///422660000000000 842256-020
URL httpt::t///// http::::://oo..1/11...o 1038/
[23] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. CoRR, abs/2312.00752, 2023. doi: 10448550/ARXXI22211.00022 URL htt.h:://p/// org/ 10. 48550//55aari
2312.00752.
[24] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 18)):1335 November 1997. ISSN 0899-7667. doi: 10.1162/neco. http//eco..19...88. 199..9.8.77355. URL https::/doi..org org/ 1162/neco. 1997.
[25] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. RULER: What's 's the real context size of your long-context language models? In First Conference on Language Modeling, 2024. URL htt.aaaa//ppeeeeveeee..tee http://wwwwerr.eeeeeoeoodeddddee
[26] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL
https://openreviies wttp:eewwwee//foueeeeeem de=n2eWeeeFYf9
[27] Xiang Hu, Zhihao Teng, Jun Zhao, Wei Wu, and Kewei Tu. Efficient lengthhgeneraaizzale attention via causal retrieval for long-context language modeling, 2025. htt.:/tttt//axxii.rrrrge 2410
01651
[28] Angelos Katharopoulos, Apoorv Vyas. Nikolaos Pappas, and François Fleuret. Transformers are rans: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICMI 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 51 5156-5165. PMLR, 2020. URL http://prrcceeedingg mlr www.ss/esssviee vvess//v19// katharopoulos220 os20a. html html
[29] Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023, htt :/ttt::://paprre nips cc/ caper Abstreaatttccetttttttcaaacaaatce www.eeeeeeooooeeen...c.o html
12

---

[30] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference Learning Representations, ICLR 2020, Addix Ababa, Ethiopia, April 26-30. 2020. OpenReview.nel, 2020. URL /tLLiattt////ppeeeevieer net wwt.:/orunnddddkgllatttttt...
[31] Chao Lou, Zixia Jia, Zilong Zheng, and Kewei Tu. Sparser is faster and less is more: Efficient sparse attention for long-range transformers. CoRR, abs/2406. 16747, 2024. doi: 10.48550/ARXI..4066600 167477
URL https://doll org/ 10. 0el.448550///0
2406 16747
[32] Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du Tao Jiang. Chao Hong, Shaowei Liu, Weiran He, Enmiing Yuan, Yuzhi Wang, Zhiqi Huang, Huan Yuan, Suting Xu, Xinran Xu, Guokun Lai, Yanru Chen, Huabin Zheng, Junjie Yan, Jianlin Su, Yuxin Wu, Neo Y. Zhang, Zhilin Yang. Xinyu Zhou, Mingxing Zhang, and Jiezhong Qiu. Moba: Mixture block attention for long-contess Ilmss CoRR, abs/2502.13189, 2025. doi: 10.48550/ARXTV25500 13189. URL BtL.sttps: //dopps://oi...o org/10. 1..5/../ arxiv. 2502. 13189
[33] Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In Inter- national Conference on Learning Reprexentations, 2018. URL http://pppnnreeeeeeeeeeereereree
HyUNmuIC-
[34] Evan Miller. Attention is off by one, 2023. by- one html
URL. hhttpa://wwi fttp::://wwweevmmmllee.. .or.g.rrrr.r//o..ooo
attention- off ww-.oof--oooooooooooo
[35] Amirkeivan Mohtashami and Martin Jaggi. Random-access infinite context length for transform- ers. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, Decem- ber 10 16, 2023, 2023. URL http::/ppaper.a...////// wttpppprrrrrppp.pppppppppppppepe http/pppprffieeeeeeeeeeeeeeeeeee wwwpr//app///ae h.t.fllee//pppe//22/haaaaa
htmll
[36] Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave context behind: Efficient infinite context transformers with infini-attention. CoRR, abs/2404.07144. 2024. doi: 10.48550/ARXIV.2404,, 07143. (ttps:s::://dod . org/ 10. 48550/ /r.550//rriv.ee 2404 07143
[37] Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave context hehind: Efficient infinite context transformers with infni-attention, 2024. URL, https: aawiiviiivvv.i org/ /bba///aa///// 2////4444.0001
[38] Ramesh Nallapati, Bowen Zhou, Cícero Nogueira dos Santos, Çaglar Gülçehre, and Bing Xiang. Ab- stractive text summarization using sequence-to tereceeeto-sqquenee rnns and beyond. In Yoav Goldberg and Stefan Riezler, editors, Proceedings of the 20th SIGNLL Conference on Computational Natural Lan- guage Learning. CoNLL 2016, Berlin, Germany, August 11-12, 2016, pages 280 290. ACL, 2016. doi: 10. 10.18653/V1/K 16- 1028. URL httpt/ttt////oooi. org/ 10.
[39] Shashi Narayan, Shay B. Cohen, and Mirella Lapata. Don't give the details, just the summary! topic-aware convolutional neural networks for extreme summarization. In Ellen Riloff. David Chiang. Julia Hockenmaier, and ichi Tsujäi, editors, Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 November 4, 2018, pages 1797- 1807. Association for Computational Linguisties, 2018. doi: 10.18653/V 1206. URL https //dei.org/10. 10. 11653//1//11- 1206
[40] Richard Yuanzhe Pang. Alicia Parrish, Nitish Joshi, Nikita Nangia, Jason Phang, Angelica Chen, Vishakh Padmakumar, Johnny Ma. Jana Thompson, He He, and Samuel R. Bowman. Quality: Question answering de Marneffe, and Iván Vladimir Meza
with long input texts, yes! In Marine Carpuat, Marie-Catherie Ruíz, editors, Proceedings of the 2022 Conference of the North American Chopter of the Association for Computational Linguistics: Human Language Technologies, NAACL 2022, Seattle WA, United States, July 10-15, 2022, pages $336-5358. Association for Computational Linguisties, 2022. doi: 2a.11865//11 1/2022.. NAACL-MAIN.391 URL (att:::ttt///ttt org/ 10. /1/1////////2222 maacl-main.nn.- main.391
[41] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässann Stockholm. Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 4052-4061. PMLR, 2018. ww..pprrss..................m www.ppre//e/ee p00/paaaaar/// h8t.
URL htt//ttt//ppcceedinng
mlr.
[42]
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung. Leon Derezynski. Xingjian Du, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Jiaju Lin, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang. Johan S. Wind, Stanislaw Wozniak, Zhenyuan Zhang, Qinghua Zhou, Jian Zhu, and Rui- Jie Zhu. RWKV: reinventing mns
13

---

for the transformer era. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Findings of the Association for Computational Linguistics: EMNLP 2023, Singapore, December 6-10, 2023, pages 14048-14077. EMNLP936. URL
Association for Computational Linguistiis. 2023, doi: 10.18663/1/2232///NDDN wttpo///.....////////0///0 findings emnlp. 936
[43] Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Xingjian Du, Teddy Ferdinan, Haowen Hou, Przemyslaw Kazienko, Kranthi Kiran GV, Jan Kocon, Bartlomiej Koptyra, Satyapriya Krishna, Ronald McClelland Jr., Niklas Muennighoff, Fares Oheid, Atsushi Saito, Guangyu Song, Haoqin Tu, Stanislaw Wozniak, Ruichong Zhang, Bingchen Zhao, Qihang Zhao, Peng Zhou, Jian Zhu, and Rui-Jie Zhu. Eagle and finch: RWKV with matrix-valued states and dynamic recurrence. CoRR. abs/2404. 008524444408822, 2024. doi: 1044550ARXXV2244005555 URL https://doi. ort/ 48850/4450//5.00 arw.v...vvi..eei. 05892
[44] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. YaRN: Efficient context window extension of large language models. In The Twelfth International Conference Learning Representations, 2024. URL httpttt/oppeeeeeiiiiiieee
[45] Ofir Press, Noah A. Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25- 2022. OpenReview.net, 2022. URL http:://///oooeeeeetttoeeeeeee/e R8sQPpGCvO
[46] Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling. In Thirry-seventh Conference on Neural Information Processing Systems, 2023. URL hRh.https
[47] Zhen Qin, Songlin Yang. Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. HGRN2: Gated linear RNNs with state expansion. In First Conference on Language Modeling, 2024. URL https://openrrevier. net//forumnddryy$5o.GGG.
[48] Jack W. Rac, Anna Potapenko, Siddhant M. Jayakumar, Chloc Hillier, and Timothy P. Lillicrap. Com- pressive transformers for long-range sequence modelling, In International Conference on Learning Representations, 2020. URL https:://openreevie net htte://ormmidd-SyKK.aa
[49] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100, 000+ questions for machine comprehension of text. In Jian Su, Xavier Carreras, and Kevin Duh, editors, Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 144, 2016, pages 2383 2392. The Association for Computational Linguistics, 2016. doi: 10. 1..1863//1/D16--- 1264. URI. URI htt.hhtttp://... orgg 18653/11/11v66 1264
[50] Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovec, Dale Schuurmans, and Bo Dai. Combiner: Full attention transformer with sparse computation cost. In Mare Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NewrlPS 2021, December 6-14, 2021, virtual, pages 22470-22482, 2021. httttts://poceeeeeeee neur ips.
2tt///2222///aaaa hwwphhhhshhhh4444e00060646011010 wtt.rrrrttttttt-tttaaaratet Abstract Abtt html
[51] Liliang Ren, Yang Liu, Yadong Lu, yelong shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling. In The Thirteenth International Conference on Learning Representations, 2025. URL htttps://openrevieee net. wtte///orrriddddllyoiliii
[52] Aurko Roy. Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Trans. Assoc. Comput. Linguistics, 9:53-68, 2021. doi: 10.1162/
TACLI Al 00353. URL http:/tttpp//.oooooo org/10. tatp://aacl____005_5_
[53] Ohad Rubin and Jonathan Berant. petrrivall transformer: Long- language modeling with self-retrieval. Transactions of the Association for Computational Linguistics, 12:1 12:1197-1213, 2024. doi:
10. 10.11162200 10.1162/tac___00000 L_00693, URL httph:httt/aatta http:///ooaaaaooogoo.ooooo org/ arg/2024. 2024. t..ll1.66) 1.66/
[54] Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. CoRR, abs/2104.09864, 2021. URL https: org/ wwwpbb/4. 09864
[55] Shawn Tan, Songlin Yang. Aaron Courville, Rameswar Panda, and Yikang Shen. Scaling dick-breaking attention: An efficient implementation and in-depth study. In The Thirteenth International Conference on Learning Representations, 2025. URL htttps://openreviek: http://www....eeeeemmm/r//eeaaaa
14

---

[56] MiniCPM Team, Chaojun Xiao, Yuxuan Li, Xu Han, Yuzhuo Bai. Jie Cai, Haotian Chen, Wentong Chen, Xin Cong. Ganqu Cui, Ning Ding, Shengda Fan, Yewei Fang, Zixuan Fu, Wenyu Guan, Yitong Guan, Junshao Guo. Yufeng Han, Bingxiang He, Yuxiang Huang, Baoxi Ji. Cunliang Kong, Qiuzuo Li, Siyuan Li, Wenhao Li, Xin Li. Yanghao Li, Yishan Li, Zhen Li, Dan Liu, Biyuan Lin, Yankai Lin, Xiang Long, Quanyu Lu, Yaxi Lu, Poiyan Luo, Hongya Lyu, Litu Ou, Yinxu Pan, Lushi Pu, Zekai Qu, Qundong Shi, Zijun Song, Jiayuan Su, Zhou Su, Ao Sun, Xianghui Sun, Peijun Tang, Fangzheng Wang. Feng Wang. Shuo Wang. Yudong Wang. Zheng Wang. Yesai Wu, Zhenyu Xiao, Jie Xie, Zihao Xie, Xiaoyue Xu, Yukun Yan, Jiarui Yuan, Jinqian Zhang. Kaihuo Zhang. Lei Zhang, Linyue Zhang, Xueren Zhang, Yudi Zhang, Hengyu Zhao, Weilin Zhao, Weilun Zhao, Yuanqian Zhao, Zhi Zheng, Chuyue Zhou, Ge Zhou, Jie Zhou, Wei Zhou, Yanghao Zhou, Zihan Zhou, Zixuan Zhou, Zhiyuan Liu, Guoyang Zeng. Chao Jia, Dahai Li, and Maosong Sun. Minicpm4: Ultra-efficient Ilms end devices, 2025.
[57] Philippe Tillet, Hsiang- -uunn Kung, and David D. Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Tim Mattson, Abdullah Muzahid, and Armando Solar-Lezama, editors, Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Pro- gramming Languages, MAPL@ PLDI 2019, Phoemix, AZ, USA, June 22, 2019, pages 10-19. ACM, 2019. 3329973
doi: 10.1 1111443333508.33299993
[58] Naftali Tishby, Fernando C. Percira, and William Bialek. The information bottleneck method, 2000. URL https://arxis org/ wwa.abaa/paaaee hhtpp/ssssscs/00007
[59] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307, 092888 2023.
[60] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg. Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA. pages 5998-6008 2017. htttthss://proceediinnee Abstract.html
neurips wttp://pppppppeeeeeeeeeeeeeeeep 2atppe/r/21//aah////223 http://hhhhhhfeeeeee
[61] Apoorv Vyas, Angelos Katharopoulos, and François Flcuret. Fast transformers with clustered attention. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Pro- cessing Systems 2020, NewrlPS 2020, December 6- 12, 2020, virtual, 2020. URL sttctps://proceedinng. Abataatttt
[62] Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shocybi, and Bryan Catanzaro. An empirical study of mamba-based language models, 2024. htRpsstts:// //axxxv///axxv..oom www. httgg/bb///00 078777
[63] Jie Wang, Tao Ji, Yuanbin Wu, Hang Yan, Tao Gui, Qi Zhang, Xuanjing Huang, and Xiaoling Wang. Length generalization of causal transformers without position encoding. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar, editors, Findings of the Association for Computational Linguistics: ACL 2024. pages 14024-14040, 14040, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10. htt.s///w/12222..nnnnnn--- 833.33 UtIpstts:://// http:///llaanntoooo...o.o org 2024. f findings- ac1.834/
[64] Shuohang Wang, Luowei Zhou, Zhe Gan, Yen-Chun Chen, Yuwei Fang, Siqi Sun. Yu Cheng, and Jingjing Liu. Cluster-former: Clustering- based sparse transformer for long-range dependency encoding. CoRR, Abs/2009.06097 2020. URL httpp://t// orw.org//aa..o.oo 2009. 06097
[65] Thomas Wolf, Loubna Ben Allal, Leandro von Werra, Li Jia, and Armel Zebaze. A dataset of python files from github, 2023. URL http:////gggchbbboooooggggggggge wttp://.cccchhctceee
wwwpp/////.oo.o///////.////aae/e wwwpaiiinnodddrrrrr.oomo md
[66] Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware -ffiiciiee training. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.nes, 2024. URL https:///peerrviee http:orrummidd1aa5xxxEvE
[67] Songlin Yang, Bailin Wang. Yu Zhang, Yikang Shen, and Yoon Kim. Parallclizing lincar transformers with the delta rule over sequence length, 2025. /ttppst::/// //axiiiv////xv.../// org/abs/ abs/ 2406. 06484
[68] Zhilin Yang, Peng Qi Saizheng Zhang, Yoshua Bengio, William W. Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. In Ellen Riloff, David Chiang. Julia Hockenmaier, and Jun'ichi Tsujii, editors, Proceedings of the 2018
15

---

Conference Empirical Methods in Natural Language Processing. Brussels, Belgium, October 31 November 4, 2018, pages 2369-2380. Association for Computational Linguistics, 2018. doi: 10.18653/V1/ D18-1259. URL httpttttp/s//oo/ wttp./......////////t....... 1259
[69] Howard Yen, Tianyu Cao, and Danqi Chen. Long-context language modeling with parallel context encoding. In Lun-Wei Ku. Andre Martins, and Vivek Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers). pages 2818. Bangkok, Thailand, August 2024, Association for Computational Linguistics. doi: 10. 118866331122224 acl-444ccllloggg long.142. URL WtRpsttps: httpe///aaaaaaaahaaaaaaaa
org/ /204//2224.22 acl- mol-long... 142/
[70] Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang. Zhenda Xie, Y. X. Wei, Lean Wang. Zhiping Xiao, Yuqing Wang. Chong Ruan, Ming Zhang, Wenfeng Liang, and Wangding Zeng. Native sparse attention: Hardware-aligned and natively trainahle sparse attention. CoRR, ahs/2502. 1089, 2025, doi: 10.48550/ARXIV25022 1089. URL C httpp::t/////... hwaiddii.rgggg/// org/10. 48550/arXiv. 250.. 11089
[71] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslic, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang. Yang. and Amr Ahmed. Big bird: Trans- formers for longer sequences. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, Decem- ber 6-12, 2020, virtual, 2020. URL attpss:/procceeding htg.:edddnnnn.ouuss....///ooeeee c8612211422228477255553388
20////2222aaa//a/ hash/
a--.ab9-6bbbbbb Abstract html
[72] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Curran Associates Inc., Red Hook,
NY, USA, 2019.
[73] Zhang, Songlin Yang. Rui-Jie Zhu, Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang. Freda Shi, Bailin Wang. Wei Bi, Peng Zhou, and Guohong Fu. Gated slot attention for efficient linear-time sequence modeling, In The Tiirth-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://openreviex net/forumiiteeeeeeeeeee jy4Phhiiimg
[74] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song. Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen. H2O: hitterr oracle for efficient generative inference large language models. Alice Oh, Tristan Naumann, Amir Gloherson, Kate Sacnko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Anmual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL ltttp://paperss nips lalllppppplleeeeee pappr//023/nab//eeeaeta//// http:////////////////////tteeeer htw.:e/eaaeeeeeeaaaaaaaaaareeeee
html
[75] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang. Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer. Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Pytorch fsdp: Experiences on scaling fully sharded data parallel, 2023. URL JtLtssss:///art //tr:i///..x/////aa orreiiiiorrg/ee aa////bb/2334 11277
16

---

Limitations
Although the method proposed in this work is theoretically applicable to all RNNs, the experimental section of this work mainly focuses on Mamba.
Considering computational resources, this work does not discuss the performance of models larger than 3B parameters.
A How HSA achieves accurate chunk retrieval
The core principle of previous sparse attentions' chunk selection lies in approximating the token-to- chunk relevance using unnormalized attention scores, In self-attention, the relevance of token j to token is defined as:
logits:..
Pi.j
Z.
logits..,,
Z.
jei
where logits,,, 9: k, is the dot product between the query q of token and the key k, of token j.
The relevance of chunk to token is ideally the sum of the relevance of all tokens within that chunk:
Tí.c Σ
logits...
PLJ
Z:
JEC
JEC
where C denotes tokens in chunk However, calculating Tie requires computing Z; (the full softmax normalization across all tokens), which would necessitate full computation and thus undermine the computational efficiency.
To ensure efficiency, they instead approximate chunk relevance using the mean-pooled representation of keys:
1 Σ - S
The = q. K.KK = - S
k, jEC
logits jEC
jEC
where K. represents the mean-pooling of key representations within the e-th chunk. While this approximation bypasses the need for normalized softmax scores across all tokens, introduces a discrepancy between r/,c (the approximation) and ri,c (the ideal one).
Let's consider the following example shown in Figure assuming every 2 tokens form a chunk.
chunk ids
3
token positions
6
logits;jj
20.0
20.0
1.0
0.0
2.0
1.0
1.0
1.0
Pij
0.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
TU
1.0
0.0
0.0
0.0
T'u
0.0
0.5
1.5
1.0
Figure 5: How unnormalized scores mislcad the chunk selection.
only the top-2 chunks can be selected, they would choose chunks 2 and according to r' thus missing the chunk with the highest sum of attention weights.
Assuming they select the top-2 chunks for sparse attention, the subsequent attention process is shown Figure Throughout this process, r/.. only participates chunk selection not the forward computation, nor does it receive gradients. The inaccurate chunk selection issue stems from using unnormalized attention logits to estimate chunk importance, with r/.c not learnable, making the inaccuracy of chunk selection unavoidable.
17

---

chunk ids
2
token positions
5
logits
-2.0
10
1.0
softmax
p'L,
0.49
0.17
0.17
0.17
Figure 6: Applying attention over the concatenation of selected chunks.
chunk ids
2
3
token positions
1
4
7
logitstj
-20.0
20.0
1.0
0.0
2.0
1.0
1.0
1.0
Pt.j
0.0
0.0
0.0
0.0
0.0
Si,c
$1,0
S1,1
$1,2
$1,3
Figure 7: The chunk selection stage HSA
For the same example, HSA works as shown Figure 7 where 81.e represents learnable token-to- chunk relevance scores. HSA selects chunks 2 and 3 according 8t,es and applies the hierarchical attention over selected chunks as shown in Figure g In this process, Wi,es derived from & par-
chunk ids
2
3
token positions
5
6
7
logitst,
intra-chunk weighting
softmax
softmax
token-level weights
Sic
$12
S13
inter-chunk weighting
softmax
chunk-level weights
W1,2
Wi3
final weights
0L73ML2
0.27w
0.5w
0.5m
Figure 8: Applying hierarchical attention selected chunks.
ticipates the final attention weight computation and the whole forward pass. This allows it to receive gradients and allocate higher weights to chunks with more important tokens. Even if random initialization initially misses the most relevant chunk continuous training enables the model to gradually learn to select the most relevant chunks.
18

---

B Hyper-parameters
Architecture
Transformer
Mamba
Mamba+NSA
Mamba+SWA
RAMba
Total Params (M)
372
368
375
375
385
Hidden size,
1024
1024
1024
1024
1024
Mamha Layers
48
48
48
48
Attention Layers
16
3xNSA
3xSWA
3 HSA IxChunk Selection, Excoderr
Other Layers
16xMLP
MLP hidden size
5504
1344
Query heads
32
16
16
KV heads
16
1
Vocab size
50280
50280
50280
50280
50280
Table Hyper-parameters 370M models.
C Training hyper-parameters
All 370M models used the AdamW optimizer with
linear learning rate warmup with warmup ratio 0.02, cosine decay to 4e 5.
peak learning rate 3.
total tokens 60B, batch size 1M tokens.
gradient clip value 1.0
no dropout
no linear bias terms
weight decay
Adam hyperparameter = (.9, (the GPT3 value)
All models are pre-trained on 16 PPUs, with each taking approximately 60 hours.
D Templates for passkey retrieval and Single-NIAH tasks
The passkey retrieval template is structured as follows: "(essays) The pass key is <PASS KEY>. (essays) What is the passkey? The passkey is". The single NIAH template follows this format: "(essays)... One the special magic numbers for long-context is: <PASS KEY>. (essays)... What is the special magic number for long-context mentioned in the provided text? Answer:"
E RAMba 2.7B
We follow the hyperparameters of Mamba-2 2.7B, where the embedding dimension is 2560, and the total number of layers is 64. We use a two-layer Transformer-basee encoder with hidden size of 2560 and an intermediate dimension of 3392. HSA is inserted into Mamba starting from the 32nd layer, with one HSA layer inserted every 8 layers. In HSA, the group size is 5, with query head size h of 16 for each group. The total additional parameters amount approximately I 110M.
E.1 Post-training
Due to the length generalization issues inherent in the Mamba-2 2.7B, we post-train it to stabilize its perplexity on longer contexts. Specifically, we utilized BPTT for post-tuning the base model. We trained the model on sequences of 32K tokens with batch size of 16 for 3K steps. totaling 1.5B tokens. This stage takes hours on 32 PPUs.
Then we follow CEPE 69 by freezing most parameters of the Mamba backbone and tuning the parameters of the additional HSA modules for post- post-training. This process primarily involves the following two stages:
19

---

Warmup. We employ warmup initialization method by simply training the model to copy the first half of the sentence. Specifically, we append identical copy of each sentence to itself and train the model to locate the distant context and replicate it. At this stage, all parameters of Mamba-2 are frozen, with only the HSA-related parameters remaining tunable. We train the model with a 32K context length, batch size of 16, for 16K steps, with a peak learning rate of 2 x 10-5, totaling 8B tokens. This stage takes around 24 hours on 32 PPUs.
Post-Training We use LoRA [26] to fine-tune 5% of the parameters in the Mamba-2 module, while keeping all parameters in the HSA module fully trainable. The model is trained with context length of 32K, a batch size of 16, for 32K steps, using peak learning rate of 2 x 10-5, totaling 16B tokens. This stage takes around 48 hours on 32 PPUs.
This model is used for the LongBench evaluation in Table
F.2 RULER finetuning
Since the RULER tasks like NIAH and FWE do not have high requirements for intrinsic knowledge of L.LMs, we opt to train RAMba from scratch and fine-tune it on RULER's synthetic dataset. This approach aims evaluate whether RAMba trained from scratch can stably converge and demonstrate long-range retrieval capabilities. We conduct pre-training on 60B tokens, which amounts to one-tenth of the Mamba-2 2.7B model, followed by fine-tuning on 1B synthetic data, which takes around 200 hours on 32 PPUs. We also fine-tuned the Mamba-2 2.7B model on the same synthetic dataset for comparison.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made in abstract and introduction are supported by experimental results in Section 4.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitation section is presented in page 16.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
20

---

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: This paper provides detailed experiment settings in Section 4.1, Appendix A,B, and D, which should be sufficient for reproducing the main experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often
21

---

one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification: The sources of all data are explained in the paper, and the code will be open-sourced on GitHub after de-anonymization.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips. //nips.cc/pppp
ssiniooiccyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: This paper provides detailed experiment settings in Section 4.1, Appendix A,B, and D, which should be sufficient for reproducing the main experimental results.
22

---

Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Since each pre-trained model only has single checkpoint, it is not feasible to evaluate error bars. However, all experiments, we include multiple datasets or multiple sets of tasks and evaluate using various metrics. Therefore, the statistical significance of the experiments can be verified and supported.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
In Appendix B,D and Section 4.4.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
23

---

Question:
Does the research conducted in the paper conform, in every respect, with the
NeurIPS Code of Ethics
Answer: [Yes]
Justification: Our work adheres all the ethical guidelines outlined by NeurIPS.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation. generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

---

We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All dataset details and original authorship are cited in Section 4.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. cow datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The code will be published and well documented after de- de-anonymization.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
25

---

According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects, hence do not require IRB approval.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: Our core method doesn't 't involve the usage of LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.MM policy for what should or should not be described.
26


Reviews:
[{"review_id": "qZd7FvTxO7", "reviewer": "Reviewer_fVSb", "summary": "This paper introduces Hierarchical Sparse Attention (HSA), a novel attention mechanism designed to provide RNNs with efficient, long-range random access capabilities, addressing the 'trilemma' of balancing efficiency, random-access flexibility, and length generalization. \n\nHSA operates by dividing the input into chunks, selecting the most relevant chunks, and then applying a two-stage attention process: first within tokens of a chunk, and then across the selected chunks. The authors integrate HSA with a Mamba-2 backbone to create RAMba. To manage memory, RAMba employs techniques like hardware-aligned kernels and offloading the KV cache to CPU memory. \n\nThe paper's experiments demonstrate that RAMba achieves near state-of-the-art performance on long-context tasks, such as perfect accuracy on a 64-million-token passkey retrieval task , while maintaining high throughput and a near-constant memory footprint during inference.", "questions": "1. Can the authors also reports the results for pure NSA w/o SWA as used in MiniCPM4? I think it is an important baseline besides transformer.\n2. Could you elaborate more on BPTT? How did you process the last hidden states after param optimized, If I understand correctly, the saved states does not align with the ones produced by the updated model weights, will this hurt the performance?\n3. I did not figure out the difference between the alg in  Fig. 2 (b) and Algorithm 1 with NSA.\n4. How about the performance of MK-NIAH on RULER?\n5. I think some more results beyond RULER and long-bench could be present to verify the claims of random long-context access abilities. Wondering if it's possible to report the results of Phonebook lookup used in \"Repeat After Me\", which is training-free\n\nMiniCPM4: https://arxiv.org/abs/2506.07900\n\nRepeat After Me: https://arxiv.org/pdf/2402.01032", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "RAMba is a novel architecture that integrates HSA with a Mamba backbone. The goal is to equip Mamba with efficient, long-range random access capabilities, a direction of significant interest for sequence modeling. While the approach is promising and the results are strong, I have a few concerns regarding the architectural design and the justification for its components.\n\n* On the Source of Performance Gains in HSA: A central claim of HSA is its \"chunk-aware\" learning mechanism, which is presented as its core innovation over prior work. However, the mechanics of how this is achieved beyond standard backpropagation through the hierarchical structure could be elaborated upon. This point is crucial because the ablation results in Table 1 for RAMba w/o s.b. show a significant performance degradation when stick-breaking attention is removed. This raises the question of whether the primary performance gains stem from the specific properties of the stick-breaking mechanism itself, rather than the proposed two-stage hierarchical attention. Further clarification or a more targeted ablation could help disentangle these effects and solidify the paper's core claim.\n* Justification for the Transformer-based Encoder: The choice of a Transformer-based bi-directional encoder for chunk memory is another key design decision. What was the specific motivation for this choice over other alternatives, such as using another Mamba layer? Given that RAMba is a Mamba-centric architecture, integrating a Transformer component introduces additional complexity. An ablation study comparing different encoder designs would significantly strengthen the justification for this specific architectural choice and its contribution to the model's impressive length generalization capabilities.\n\nGenerally, I think the the overall design of RAMba is quite sophisticated, with several novel components (HSA, chunk encoder, memory reset) interleaved with the Mamba backbone. \nWhile the results on very long contexts are impressive, a discussion on the model's practical efficiency (latency and throughput) on shorter sequences—which are common in many real-world scenarios—would be beneficial. It is unclear how the overhead from the added components impacts performance at smaller scales.\n\nTo make the contribution more convincing and the design principles clearer, the authors could add a more detailed architectural justification. For instance, explaining the rationale behind alternating one HSA layer with G Mamba layers in the upper decoder could help readers understand the intended synergy between the components. A more minimalist design, supported by careful ablations that justify each component's inclusion, would make the final architecture even more compelling.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your insightful feedback.\n\nQ3,W1. **I did not figure out the difference between HSA and NSA.  How \"chunk-aware\" learning is achieved beyond standard backpropagation through the hierarchical structure could be elaborated upon**\n\nThis is the most important contribution of this work, so please allow us to explain it first.  \n\nWhy NSA suffer from inaccurate chunk selection, and why can HSA learn from back-propagation?\n\nThe core principle of NSA's learnable chunk selection lies in approximating the token-to-chunk relevance **using unnormalized attention scores**. In self-attention, the relevance of token j to token i is defined as:\n\np[i,j] = $\\frac{e^{logits_{i,j}}}{Z_i}, \\quad Z_i = \\sum_{j < i}{e^{logits_{i,j}}}$,\n\nwhere logits[i,j] = $q_i^T k_j$ is the dot product between the query $q_i$ of token i and the key $k_j$ of token j.\n\nThe relevance of chunk c to token i is ideally the sum of the relevance of all tokens within that chunk:\n\nr[i,c] = $\\sum_{j \\in \\mathcal{C}}{p_{i,j}}= \\frac{1}{Z_i} \\sum_{j \\in \\mathcal{C}} e^{logits_{i,j}}$, \n\nwhere $\\mathcal{C}$ denotes tokens in chunk c.\nHowever, calculating r[i,c] requires computing $Z_i$ (the full softmax normalization across all tokens), which would necessitate full $ Q^TK$ computation and thus undermine the computational efficiency.\n\nTo ensure efficiency, NSA instead approximates chunk relevance using the mean-pooled representation of keys:\n\nr'[i,c] = $q_i^T K_c= q_i^T \\frac{1}{S} \\sum_{j \\in \\mathcal{C}}{k_j} = \\frac{1}{S} \\sum_{j \\in \\mathcal{C}} q_i^T k_j = \\frac{1}{S} \\sum_{j \\in \\mathcal{C}}{logits_{i,j}}$,\n\nwhere $K_c$ represents the mean-pooling of key representations within chunk c. While this approximation bypasses the need for normalized softmax scores across all tokens, it introduces a discrepancy between r'[i,c] (the approximation) and r[i,c] (the ideal one).\n\nLet's consider the following example, assuming every 2 tokens form a chunk, separated by |:\n```\nchunk ids       |       0      |      1    |      2    |     3     | ...\ntoken positions |   0   |  1   |  2  |  3  |  4  |  5  |  6  |  7  | ... , i (current token)\nlogits[i,j]     | -20.0 | 20.0 | 1.0 | 0.0 | 2.0 | 1.0 | 1.0 | 1.0 |\np[i,j]          |   0.0 |  1.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |\nr[i,c]          |      1.0     |    0.0    |    0.0    |    0.0    |\nr'[i,c]         |      0.0     |    0.5    |    1.5    |    1.0    |\n```\n\nIf only the top-2 chunks can be selected, NSA would choose chunks 2 and 3 according to r' [i,c], thus missing the chunk 0 with the highest sum of attention weights. \n\nAssuming NSA selects the top-2 chunks for sparse attention, the following process would be:\n``` \nchunk ids         |      2      |      3      | \ntoken positions   |   4  |   5  |   6  |   7  |\nlogits[i,j]       |  2.0 |  1.0 |  1.0 |  1.0 |\n                  |<-----     softmax   ----->|\np'[i,j]           | 0.49 | 0.17 | 0.17 | 0.17 |\n``` \n\nThroughout this process, r'[i,c] only participates in chunk selection but not in the forward computation, nor does it receive gradients. \n\n**The NSA's inaccurate chunk selection issue stems from using unnormalized attention logits to estimate chunk importance, with r'[i,c] not learnable, making the inaccuracy of chunk selection unavoidable.**\n\nFor the same example, HSA works as follows:\n```\nchunk ids          |       0      |      1     |      2    |     3     | ...\ntoken positions    |   0   |   1  |  2   |  3  |  4  |  5  |  6  |  7  | ...\nlogits[i,j]        | -20.0 | 20.0 |  1.0 | 0.0 | 1.0 | 1.0 | 1.0 | 1.0 | ...\np[i,j]             |  0.0  |  1.0 |  0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | ...\ns[i,c]             |    s [i,0]   |  s [i,1]   |  s [i,2]  |  s [i,3]  | ...\n```\nWhere s [i,c] represents learnable token-to-chunk relevance scores as described in Line 155 of the paper. For example, HSA may also select chunks 2 and 3 when parameters are randomly initialized, we have:\n```\nchunk ids            |           2           |          3          |\ntoken positions      |     4     |     5     |     6    |     7    |\nattn_logits [i,j]    |   2.0     |    1.0    |    1.0   |    1.0   |\nintra-chunk weighting|<--     softmax      ->|<--    softmax     ->|\ntoken-level weights  |   0.73    |    0.27   |    0.5   |    0.5   |\ns[i,c]               |        s[i,2]         |        s[i,3]       |\ninter-chunk weighting|<---            stick-breaking           --->|\nchunk-level weights  |        w[i,2]         |        w[i,3]       |   // s.t. w [i,2} + w [i,3] <= 1\nfinal weights        |0.73×w[i,2]|0.27×w[i,2]|0.5×w[i,3]|0.5×w[i,3]|\n```\n\nIn this process, w[i,c], derived from s[i,c], participates in the final attention weight computation and the whole forward pass. **This allows it to receive gradients and allocate higher weights to chunks with more important tokens.** Even if random initialization initially misses the most relevant chunk 0, continuous training enables the model to gradually learn to select the most relevant chunks.\n\nQ1. **Can the authors also report the results for pure NSA w/o SWA as used in MiniCPM4?**\n\nSince MiniCPM4 had not been released at the time of submitting the paper, we did not include a comparison. However, this suggestion was very helpful. We tested it, and the results are as follows:\n\nPerplexity:\n|Models| 4K | pg19 | arxiv | code | 16K | pg19 | arxiv | code | 64K | pg19 | arxiv | code | \n|-|-|-|-|-|-|-|-|-|-|-|-|-|\n|Transformer full-attn| | 18.61 | 4.23 | 3.28 | | 539.15 | 199.42 | 62.17 | | >104 | >104 | 2865.51 |\n|RAMba w/o m.r. | |17.63 | 4.13 |3.21 | | 17.11| 3.81 |3.08 | | 17.11 |3.87 |3.21|\n|Transformer w/ NSA| | 18.43 | 4.18 | 3.24 | | 20.74 | 4.55 | 3.46 | | 23.61 | 5.66 | 4.02|\n\nRULER:\n| Models | 4K | S-N | MQ-N | VT | FWE | 64K | S-N | MQ-N | VT | FWE |\n|-|-|-|-|-|-|-|-|-|-|-|\n|Transformer full-attn| | 95.08 | 88.59 | 97.12 | 60.0 | | 0.00 | 0.00 | 0.00 | 0.00 |\n|Ramba w/o m.r | | 92.76 | 87.10 | 96.66 | 76.81 | | 55.22 | 29.85 | 68.66 | 8.96 |\n| Transformer w/ NSA | | 46.85 | 5.66 | 19.48 | 51.95 | | 0.00 | 0.00 | 0.00 | 7.46 |\n\nThis result further validates our argument that the NSA has issues with chunk selection, even for in-domain length.\n\nW1. **Whether the primary performance gains stem from the stick-breaking mechanism itself, rather than HSA.**\n\nRAMba$_{w/o s.b}$ means using softmax without position encoding instead of the stick-breaking mechanism (please refer to Line 242). We believe that the performance degradation is due to the inability to distinguish the relative distance between chunks, rather than additional improvement brought by the stick-breaking mechanism.  \n\nRegarding the extrapolation ability,  the experiments in the original stick-breaking attention paper[1] (Figure 6)  show that stick-breaking attention itself cannot resolve the NIAH extrapolation issue, which demonstrates that the extrapolation capability is not brought by the stick-breaking mechanism. Moreover, the stick-breaking mechanism is not applied to tokens, but only to chunks, as we mentioned above.\n\nIn summary, stick-breaking is one component of the overall design, and it undoubtedly contributes to our results to some extent. However, the core contributions of this paper, the extrapolation capability and the high performance of sparse attention, are still primarily achieved by HSA.\n\nW2. **Justification for the Transformer-based Encoder.**\n\nSince the intra-chunk attention of HSA is conducted with each chunk separately,  an intuitive idea is to use a bi-directional encoder to enhance chunk-level representations. Please note that it doesn't break causality, as each token is not allowed to access the chunk it belongs to.\n\nUsing bi-directional Mamba layers may be another choice. Here are the results for RULER:\n\n|Models| 4K | S-N | MQ-N | VT | FWE | 64K | S-N | MQ-N | VT | FWE |\n|-|-|-|-|-|-|-|-|-|-|-|\n|RAMba tfm enc | | 91.74 | 80.61 | 95.18 | 49.81 | | 85.07 | 55.22 | 55.22 | 53.73 |\n|RAMba mamba enc | | 92.25 | 78.45 | 94.28 | 45.35 | | 83.58 | 56.72 | 53.73 | 47.76 |\n\nThe results are quite close, which validates that the choice of encoder has minimal impact.\n\nQ2. **Could you elaborate more on BPTT?**\n\nExisting parallelizable RNNs often suffer from Oversized States[3]. A common solution is to use BPTT (Backpropagation Through Time) for pretraining. \n\nSimply put, the last state from the previous step is used, and after gradient truncation, it becomes the initial state for the current step.\n\nThis aligns with the recent work[4] by the author of Mamba, where they term this mechanism as 'state passing.'\n\nAs demonstrated in the experiments in [4] (Fig. 5), state passing has a negligible impact on results but significantly enhances extrapolation capability.  \n\nQ5. **Phonebook lookup**\n\nHere are the results:\n\n| Models | 4k | #entries=10 | #entries=50 | 16K | #entries=10 | #entries=20 | 64K | #entries=10 | #entries=20 |\n|-|-|-|-|-|-|-|-|-|-|\n| Transformers full-attn| | 87.76 | 77.46 | | 0.0 | 0.0 | | 0.0 | 0.0 |\n| Mamba | | 0.0 | 0.0 | | 0.0 | 0.0 | | 0.0 | 0.0 |\n| Mamba w/ SWA | | 0.01 | 0.0 | | 0.0 | 0.0 | | 0.0 | 0.0 |\n| Mamba w/ NSA | | 14.65 | 1.11 | | 5.94 | 4.08 | | 0.0 | 0.0 |\n| RAMba | | 87.01 | 71.52 | | 75.80 | 63.56 | | 52.24 | 35.82 | \n\nQ4. **MK-NIAH on RULER**\n\nThe MQ-NIAH test in Table 3 covers the multi-key NIAH test. It contains 6 key-value pairs in the context. Moreover, the phonebook lookup task is a harder version of MK-NIAH.\n\n[1] Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study\n\n[2] Retrieval-Pretrained Transformer: Long-range Language Modeling with Self-retrieval\n\n[3] Stuffed Mamba: Oversized States Lead to the Inability to Forget\n\n[4] Understanding and Improving Length Generalization in Recurrent Models\n\nYour suggestion to compare Transformers with the NSA is extremely helpful. We have provided the relevant data to address your concerns and hope our responses adequately resolve the issues. We would greatly appreciate it if you could reconsider and raise the score."}, {"review_id": "Jk2NLimHNO", "reviewer": "Reviewer_hWyn", "summary": "This paper introduces a novel Hierarchical Sparse Attention (HSA) mechanism to give RNN-based models like Mamba random-access capability while maintaining linear complexity. By integrating HSA with a hardware-aligned kernel and memory offloading, the resulting RAMba model shows exceptional length generalization and efficiency. Notably, it achieves perfect accuracy on a 64M-token passkey retrieval task despite being trained on only a 4K context.", "questions": "1.\tRegarding the model’s architectural complexity and its new hyperparameters (S, K, G): could the authors comment on their sensitivity? How were the reported values chosen, and what is the expected impact of varying them?\n2.\tRegarding the memory reset mechanism: could the authors provide an ablation study against simpler strategies (e.g., resetting to zero or to the last state) to better justify the choice of a randomized approach?\n3.\tTo validate the true advantage of the HSA hierarchical design, we suggest the authors provide results for an NSA baseline augmented with the same chunk encoder. This experiment is crucial for a fair comparison.", "limitations": "yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1.\tThe paper’s core contribution is a novel, well-motivated Hierarchical Sparse Attention (HSA) mechanism that directly addresses the critical trade-off between efficiency and random-access capability in long-context models.\n2.\tThe claims are backed by extensive and compelling experiments across a wide range of tasks. The standout performance on a 64M-token passkey retrieval task strongly validates the method’s length extrapolation capabilities.\n3.\tThe work extends beyond theory to a practical, system-level solution, including a hardware-aligned kernel design and an efficient inference strategy, which significantly enhances its real-world applicability.\nWeaknesses:\n1.\tWhile powerful, the RAMba architecture introduces significant complexity and multiple new hyperparameters (e.g., chunk size S, selected chunks K, HSA layer frequency G). The paper lacks ablation studies on these crucial design choices, making it difficult to assess their individual impact or the model’s sensitivity.\n2.\tThe choice of a randomized segment for the memory reset mechanism is not sufficiently justified. Although shown to be effective, the lack of comparison against simpler strategies (like resetting to zero) makes it unclear whether the randomization itself is a key ingredient or an unnecessary complexity.\n3.\tThe comparison against the NSA baseline is confounded. As shown in Fig. 3, RAMba has a significant advantage over RAMba-w/o-enc, which proves the chunk encoder is critical. Given this, the comparative experiments should have equipped the NSA baseline with the same component to fairly test the merits of the hierarchical attention design itself. Without this control, the source of the performance gain remains ambiguous.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your insightful feedback and recognition of our work.\n\nW1,Q1. **Regarding the model’s architectural complexity and its new hyperparameters (S, K, G): could the authors comment on their sensitivity? How were the reported values chosen, and what is the expected impact of varying them?**\n\nBased on previous research on hybrid models ([1], Fig.4), the optimal ratio between attention and Mamba is generally around 8.25\\%~20\\%. Therefore, we inserted 3 HSA layers into the 48-layer Mamba model. Generally, S is set to 64 to better align with hardware.\n\nThe value of K primarily impacts computational cost—larger K values typically result in better performance but also higher computational load and lower sparsity. The optimal value of K is worth conducting an empirical study on. However, since this work requires comparison with other baselines to verify its effectiveness, we have not delved into this aspect here. This will be left for future work.\n\nW2,Q2. **Regarding the memory reset mechanism: could the authors provide an ablation study against simpler strategies (e.g., resetting to zero or to the last state) to better justify the choice of a randomized approach?**\n\nThank you for your insightful question regarding the memory reset mechanism. To address your concern, we would like to clarify that our approach involves initializing the memory state using the last hidden state from a randomly selected segment.\n\nIn fact, a recent study [2] has conducted an empirical analysis specifically on this topic.\nThey refer to using the last hidden state from a random segment as the initial state as \"state passing,\" and using the last hidden state from the preceding sentence as the initial state as \"truncated backpropagation through time (BPTT or TBTT). They compared different strategies, including setting to zero, TBTT, and state passing.\nTheir results ([2], Fig. 5) demonstrate that initializing the initial state with the last state from a random segment (as employed in our paper) consistently outperforms other alternatives when extrapolating, such as resetting to zero or initializing with a randomized vector. This evidence supports the effectiveness of our chosen approach.\n\nW3, Q3. **We suggest the authors provide results for an NSA baseline augmented with the same chunk encoder**\n\nWe sincerely appreciate your suggestion. We did not conduct the comparison for two reasons. First, we believe that RAMba (w/o enc) is a fair comparison with Mamba w/ NSA. Secondly, NSA is inherently incompatible with the chunk encoder, as it is unique to the hierarchical attention framework. Please allow us to elaborate further.\n\nOne purpose of introducing RAMba (w\\o enc) is to enable a fair comparison with NSA. Aside from the differences in the attention module, RAMba (w\\o enc) and Mamba (w\\ NSA) are identical in every other aspect. However, RAMba (w\\o enc) significantly outperforms Mamba (w\\ NSA) on RULER (Table. 3), and this comparison validates that the performance improvement mainly stems from HSA. \n\nIn addition, NSA itself is not compatible with the encoder. Let's take an example to see the difference between how NSA and HSA work.\n\nAssuming every 2 tokens form a chunk, separated by |:\n```\nchunk ids             |       0      |      1    |      2    |     3     | ...\ntoken positions       |   0   |  1   |  2  |  3  |  4  |  5  |  6  |  7  | ... , i <- the currrent token\nattn logits[i, *]     | -20.0 | 20.0 | 1.0 | 0.0 | 2.0 | 1.0 | 1.0 | 1.0 | \nattn weights (softmax)|   0.0 |  1.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |  // O(L$^2$), L: sequence length\nchunk attn logits sum |      0.0     |    0.5    |    1.5    |    1.0    |  // O(L*(L/S), S: chunk size\n```\nNSA selects chunks according to \"chunk attn logits sum\" for efficiency.\nIf only the top-2 chunks can be selected, NSA would choose chunks 2 and 3, thus missing the chunk 0 with the highest sum of attention weights, **leading to inaccurate chunk selection**.\n\nAssuming NSA selects the top-2 chunks for sparse attention, the following process would be:\n``` \nchunk ids         |      2      |      3      | \ntoken positions   |   4  |   5  |   6  |   7  |  // sentence-level RoPE applied here (0 ~ L)\nattn logits       |  2.0 |  1.0 |  1.0 |  1.0 | \n                  |<-   enc?  ->|<-   enc?  ->|  \n                  // Chunk-wise encoding here will disrupt global positional information.\n                  |<-----     softmax   ----->| \nattn weights      | 0.49 | 0.17 | 0.17 | 0.17 |\n``` \nIf chunk encoding is applied at the chunk level, it will destroy the global positional information.\n\nFor the same example, HSA works as follows:\n```\nchunk ids          |       0      |      1     |      2    |     3     | ...\ntoken positions    |   0   |   1  |  2   |  3  |  4  |  5  |  6  |  7  | ...\nattn logits        | -20.0 | 20.0 |  1.0 | 0.0 | 1.0 | 1.0 | 1.0 | 1.0 | ...\ns[i,c]             |    s [i,0]   |  s [i,1]   |  s [i,2]  |  s [i,3]  | ...\n```\nWhere s [i,c] represents learnable token-to-chunk relevance scores as described in Line 155 of the paper. For example, HSA may also select chunks 2 and 3 when parameters are randomly initialized, we have:\n```\nchunk ids            |           2           |          3          |\n                     |<---      enc      --->|<---      enc    --->| // chunk-wised encoding\nattn logits          |   2.0     |    1.0    |    1.0   |    1.0   | // intra-chunk RoPE only (0~S for each chunk)\nintra-chunk weighting|<--     softmax      ->|<--    softmax     ->|\ntoken-level weights  |   0.73    |    0.27   |    0.5   |    0.5   |\ns[i,c]               |        s[i,2]         |        s[i,3]       |\ninter-chunk weighting|<---            stick-breaking           --->| // with positional inductive bias\nchunk-level weights  |        w[i,2]         |        w[i,3]       |   // s.t. w [i,2} + w [i,3] <= 1\nfinal weights        |0.73×w[i,2]|0.27×w[i,2]|0.5×w[i,3]|0.5×w[i,3]|\n```\nDue to the positional inductive bias inherent in the stick-breaking attention mechanism, it becomes feasible to employ a bidirectional encoder at the chunk level.\n\nMoreover, in this process, w[i,c], derived from s[i,c], participates in the final attention weight computation and the whole forward pass. This allows it to receive gradients and allocate higher weights to chunks with more important tokens. Even if random initialization initially misses the most relevant chunk 0, continuous training enables the model to gradually learn to select the most relevant chunks.\n\n\n[1] An Empirical Study of Mamba-based Language Models\n\n[2] Understanding and Improving Length Generalization in Recurrent Models"}, {"review_id": "mqTb4ZVVy7", "reviewer": "Reviewer_eEPG", "summary": "This paper introduces RAMba, which combines Mamba (a Linear RNN) with a novel attention method, called Hierarchical Sparse Attention (HSA), whose goal is to optimize the tradeoffs between efficiency, length generalization, and random access in long-context modeling. \nThe key innovation is HSA, a two-stage hierarchical attention mechanism that learns chunk-aware token-to-chunk relevance. \nThe authors implement hardware-aligned kernels for efficiency and introduce a memory reset mechanism during training to improve length generalization.\nRamba achieves perfect accuracy on context passkey retrieval up to 64M context length despite pre-training on only 4K contexts. \nOn language modeling tasks, it performs slightly outperforms with Transformers and Mamba on IID data (4K length), and performs on par with Mamba on OOD samples (16K and 64K context length).", "questions": "- Can you provide concrete evidence (citations, experiments) supporting the claim that existing chunk-based methods \"suffer from inaccurate chunk selection\"? I really think this is crucial because without this evidence, the fundamental motivation for HSA remains somewhat obscure.\n\n- Do the PPUs used in efficiency evaluation support tensor cores? If not, I don't see how you can validly compare against FlashAttention-2, which is optimized specifically for tensor core utilization. \n\n- Are the perplexity results in Table 1 on validation or test sets? \n\n- How and where is the sorting operation applied? If it's inside the Triton kernel, how is it implemented?\n\n- Can you provide a detailed breakdown of the FLOPs required for HSA and its arithmetic intensity? Specifically, I'd like to understand the computational cost per token including chunk selection, token-level attention, and chunk-level fusion, and whether your method is memory-bound or compute-bound.", "limitations": "I think the authors provide a limitations section in the Appendix acknowledging the focus on Mamba and computational resource constraints. However, I believe some limitations could be addressed in order to improve the paper:\n\n- Discussion of how findings might change at larger model/context scales\n\n- I don't see any acknowledgment of the experimental inconsistencies or hardware evaluation issues that concern me (see weakness above)\n\n- I believe there's limited theoretical analysis of when/why the approach might fail.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n- I believe the hierarchical sparse attention mechanism represents a meaningful advance over existing chunk-based methods. The distinction between \"chunk-aware\" and \"chunk-unaware\" learning seems intuitive and well-motivated to me, with the two-stage attention design enabling end-to-end optimization of chunk selection.\n\n- The 64M context passkey retrieval with perfect accuracy is quite impressive. That's 16000x beyond pre-training length, which I think demonstrates exceptional length extrapolation. The consistent improvements across language modeling benchmarks within their evaluated scope also support the method's effectiveness.\n\n- I appreciate that the authors provided hardware-aligned kernel design with detailed algorithmic descriptions (Algorithms 1-3), which is very important for reproducibility and practical deployment. \n\n- The efficiency-generalization-access trilemma is well-articulated and represents a fundamental challenge in long-context modeling.\n\n**Weakness**\n\nWhile I largely enjoying reading this paper, I believe it still has room for improvement. Here are some of them:\n\n- The paper states that existing methods \"often suffer from inaccurate chunk selection\" (Line 37) without citing any supporting evidence. I think this unverified claim serves as the primary motivation for their approach, which I consider a significant methodological flaw.\n\n- I believe the multiple interacting components (Mamba + HSA + chunk encoder + memory reset) make it difficult to isolate individual contributions and understand failure modes. A more objective presentation and evaluation of HSA sounds more compelling.\n\n- The paper reports NSA results for language modeling (Table 1), passkey retrieval (Figure 3), and other tasks using Mamba with NSA, but later states \"NSA's implementation currently does not support generation.\" Since language modeling IS generation (next-token prediction), I think this reveals either implementation inconsistencies or conceptual confusion (possibly on my side). \n\n- I find it hard to accept the claim that Mamba \"excels\" (Line 292) at RULER tasks based on 20-40% accuracy rates. I'd say, instead, that all models perform poorly after a context length of 64K.\n\n- I believe that a critical issue lies in the efficiency comparison using \"Physics Processing Units (PPUs)\", which, to the best of my knowledge, lack tensor cores. In turn, this makes me think the comparisons with FlashAttention-2 (specifically optimized for tensor core utilization) are misleading, making the efficiency claims invalid or overstretched.\n\n- Finally, the paper lacks detailed computational complexity analysis. Based on my calculations (disregarding sorting), for the settings from section 4.1 (K=8 chunks, S=64 chunk size), HSA requires approximately $O(n \\times \\lceil n/S \\rceil + n \\times K \\times S^2)$ FLOPs per sequence. For a 16K sequence with $d=1024$, we get ~540M FLOPs per token, and an arithmetic intensity of ~33 FLOPs/byte. This indicates that HSA is likely compute-bound rather than memory-bound on modern hardware (e.g., GPUs), which could better elucidate how memory offloading can provide empirical benefits.", "was_revised": true, "final_justification": "Revised score after rebuttal: clarifications regarding the functionality of the method and its evaluation.", "rebuttal": "We sincerely appreciate your insightful feedback and are delighted to hear that you largely enjoyed reading our paper. \n\nW1, Q1: **Can you provide concrete evidence supporting the claim that existing chunk-based methods \"suffer from inaccurate chunk selection\"?**\n\nApologies for the confusion. \n\nIn line 38, \"a closer examination reveals a critical weakness\" refers to the evidence observed in our experiments. The chunk-based methods mentioned in our introduction mainly refer to NSA, which we consistently compare against in all experiments. Notably, in Table 3, Mamba(w/ NSA) performs significantly worse than RAMba(w/o enc) on the more challenging MQ-NIAH task, for in-domain sequence lengths. The gap becomes even more pronounced in length extrapolation. Please note these two models only differ in the attention module (NSA vs HSA).\n\nThe results of RAMba(w/ NSA, w/ m.r.) further support the claim. Without RNN's memory, where retrieval relies solely on NSA, there is a significant performance drop across all tasks. In contrast, RAMba(w/ m.r.) exhibits no such issue. \n\nHere, **please allow us to further explain through theory** and specific examples why NSA may suffer from inaccurate chunk selection, as well as how HSA helps mitigate this problem.\n\nThe core principle of NSA's learnable chunk selection lies in approximating the token-to-chunk relevance using unnormalized attention scores. In self-attention, the relevance of token j to token i is defined as:\n\np[i,j] = $\\frac{e^{logits_{i,j}}}{Z_i}, \\quad Z_i = \\sum_{j < i}{e^{logits_{i,j}}}$,\n\nwhere logits[i,j] = $q_i^T k_j$ is the dot product between the query $q_i$ of token i and the key $k_j$ of token j.\n\nThe relevance of chunk c to token i is ideally the sum of the relevance of all tokens within that chunk:\n\nr[i,c] = $\\sum_{j \\in \\mathcal{C}}{p_{i,j}}= \\frac{1}{Z_i} \\sum_{j \\in \\mathcal{C}} e^{logits_{i,j}}$, \n\nwhere $\\mathcal{C}$ denotes tokens in chunk c.\nHowever, calculating r[i,c] requires computing $Z_i$ (the full softmax normalization across all tokens), which would necessitate full $ Q^TK$ computation and thus undermine the computational efficiency.\n\nTo ensure efficiency, NSA approximates chunk relevance using the mean-pooled key representation $K_c$ of chunk c:\n\nr'[i,c] = $q_i^T K_c= q_i^T \\frac{1}{S} \\sum_{j \\in \\mathcal{C}}{k_j} = \\frac{1}{S} \\sum_{j \\in \\mathcal{C}} q_i^T k_j = \\frac{1}{S} \\sum_{j \\in \\mathcal{C}}{logits_{i,j}}$,\n\nwhere $K_c=\\frac{1}{S} \\sum_{j \\in \\mathcal{C}}{k_j}$. While this approximation bypasses the need for normalized softmax scores across all tokens, it introduces a discrepancy between r'[i,c] (the approximation) and r[i,c] (the ideal one).\n\nLet's consider the following example, assuming every 2 tokens form a chunk, separated by |:\n```\nchunk ids       |       0      |      1    |      2    |     3     | ...\ntoken positions |   0   |  1   |  2  |  3  |  4  |  5  |  6  |  7  | ... , i\nlogits[i,j]     | -20.0 | 20.0 | 1.0 | 0.0 | 2.0 | 1.0 | 1.0 | 1.0 |\np[i,j]          |   0.0 |  1.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |\nr[i,c]          |      1.0     |    0.0    |    0.0    |    0.0    |\nr'[i,c]         |      0.0     |    0.5    |    1.5    |    1.0    |\n```\n\nIf only the top-2 chunks can be selected, NSA would choose chunks 2 and 3 according to r' [i,c], thus missing the chunk 0 with the highest sum of attention weights. \n\nAssuming NSA selects the top-2 chunks for sparse attention, the following process would be:\n``` \nchunk ids         |      2      |      3      | \ntoken positions   |   4  |   5  |   6  |   7  |\nlogits[i,j]       |  2.0 |  1.0 |  1.0 |  1.0 |\n                  |<-----     softmax   ----->|\np'[i,j]           | 0.49 | 0.17 | 0.17 | 0.17 |\n``` \n\nThroughout this process, r'[i,c] only participates in chunk selection but not in the forward computation, nor does it receive gradients. \n\n**The NSA's inaccurate chunk selection issue stems from using unnormalized attention logits to estimate chunk importance, with r'[i,c] not learnable, making the inaccuracy of chunk selection unavoidable.**\n\nFor the same example, HSA works as follows:\n```\nchunk ids          |       0      |      1     |      2    |     3     | ...\ntoken positions    |   0   |   1  |  2   |  3  |  4  |  5  |  6  |  7  | ...\nlogits[i,j]        | -20.0 | 20.0 |  1.0 | 0.0 | 1.0 | 1.0 | 1.0 | 1.0 | ...\np[i,j]             |  0.0  |  1.0 |  0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | ...\ns[i,c]             |    s [i,0]   |  s [i,1]   |  s [i,2]  |  s [i,3]  | ...\n```\nWhere s [i,c] represents learnable token-to-chunk relevance scores as described in Line 155 of the paper. For example, HSA may also select chunks 2 and 3 when parameters are randomly initialized, we have:\n```\nchunk ids            |           2           |          3          |\ntoken positions      |     4     |     5     |     6    |     7    |\nattn_logits [i,j]    |   2.0     |    1.0    |    1.0   |    1.0   |\nintra-chunk weighting|<--     softmax      ->|<--    softmax     ->|\ntoken-level weights  |   0.73    |    0.27   |    0.5   |    0.5   |\ns[i,c]               |        s[i,2]         |        s[i,3]       |\ninter-chunk weighting|<---            stick-breaking           --->|\nchunk-level weights  |        w[i,2]         |        w[i,3]       |   // s.t. w [i,2} + w [i,3] <= 1\nfinal weights        |0.73×w[i,2]|0.27×w[i,2]|0.5×w[i,3]|0.5×w[i,3]|\n```\n\nIn this process, w[i,c], derived from s[i,c], participates in the final attention weight computation and the whole forward pass. This allows it to receive gradients and allocate higher weights to chunks with more important tokens. Even if random initialization initially misses the most relevant chunk 0, continuous training enables the model to gradually learn to select the most relevant chunks.\n\nThis point can be further validated through a pure NSA Transformer.\nRULER results:\n| Models | 4K | S-N | MQ-N | VT | FWE | \n|-|-|-|-|-|-|\n|Transformer full-attn| | 95.08 | 88.59 | 97.12 | 60.0 | \n| Transformer w/ NSA | | 46.85 | 5.66 | 19.48 | 51.95 |\n\n\nW2. **The multiple interacting components make it difficult to isolate individual contributions**\n\nWe fully agree with your perspective and have conducted ablation studies on each component to observe their individual impacts.\n\nTable 3 includes all ablation experiments, such as removing the encoder (w/o enc), memory reset (w/o m.r.), and replacing HSA with NSA or purely MAMBA, reflecting the influence of different modules on the overall system.\n\nThe most straightforward comparison is between RAMba(w/o enc), Mamba, and Mamba(w/ NSA). The only difference lies in HSA, which shows significant improvement for inputs of 64K and below. The other variants, such as memory reset and encoder, primarily serve to further enhance HSA's extrapolation capability.  \n\nW3. **NSA's implementation currently does not support generation.**\n\nAs there is no official implementation for NSA thus we use an implementation from the open-source community (line 230).\n\nNSA's batch forward mode (used for training) calculates next-token perplexity for full sequences without needing KV cache management for new tokens, as no new tokens are generated. In contrast, generation mode requires handling a variable number of newly generated tokens, complex kv cache management, and positional encoding, especially with different attention mechanisms. This complexity is why vLLM performs well for inference but not training, while Megatron is optimized for the reverse. The KV management for generation was incomplete at the time of the paper's submission.\n\nEven so, this does not undermine our conclusions. The majority of experiments can be evaluated via batch forward passes, and the performance on PPL and RULER is sufficient to support our claims.\n\nW4. **I find it hard to accept the claim that Mamba \"excels\".**\n\nThanks for pointing that out. The original phrasing was imprecise. We will revise it to: \"Mamba performs better in long sequential statistical tasks.\"\n\nW6/Q5. **How memory offloading can provide empirical benefits. Regarding computational complexity analysis.**\n\nSorry for the confusion, the \"memory offloading\" refers to offloading the KV cache from **GPU to the CPU**, not HBM to SRAM; Memory offloading does not provide computational benefits; in fact, it slows down generation but reduces the GPU's memory footprint, as reflected in Figure 4(3) and Table 7.\nDuring long-sequence inference, the KV cache could overwhelm GPU memory. Offloading the KV cache to CPU and just reloading the KV cache of selected chunks can largely save the memory footprint.\nTherefore, memory offloading has no relationship with whether HSA is memory-bound or compute-bound.\n\nThe complexity of HSA is as follows:\n\nL: sequence length, d: hidden size, K: number of selected chunks, S: number of tokens in a chunk\nRetrieval: $O(L \\times \\frac{L}{S} \\times d)$\n\nAttention: $O(L \\times K \\times S \\times d)$\nCompared to your version, there is no $S^2$, each token attends to $K \\times S$ tokens.\n\nChunk-level fusion: $O(L \\times K \\times d)$\n\nQ2: **Do the PPUs support tensor cores?**\n\nYes, they do. \nHere, the \"PPU\" is just another name for GPUs not produced by NVIDIA. It also supports Torch and Triton. \nRecently, we have access to the H800, and the throughput of 8*H800 is twice the results shown in Table 6, which aligns with Figure 4. Without tensor cores, the PPU could not achieve such throughput.\n\nQ3. **Are the perplexity results on validation or test sets?**\n\nOn validation sets.\n\nQ4. **How and where is the sorting operation applied?**\n\nWhen computing token-to-chunk weighting in line 155. It's outside the HSA triton kernel; weights and indices are already computed before conducting HSA.\n\nYour suggestion on discussing the chunk selection issue is very helpful, and we will include it in the Appendix. We hope our replies address your concerns, and we would appreciate it very much if you could consider raising the score."}, {"review_id": "EBvi0ARutr", "reviewer": "Reviewer_XZCu", "summary": "This paper claims that current RNNs fail to randomly access historical context and questions whether the efficiency advantages still hold when attention mechanisms are integrated into RNNs. To this end, the authors propose Hierarchical Sparse Attention and combine it with Mamba2, leading to a hybrid RAMba architecture. HSA consists of two stages, including a bidirectional encoder within each chunk and a weighted summation for the chunks selected by top-k mechanisms. Moreover, they design a hardware-aligned Triton kernel to efficiently implement HSA. RAMba achieves remarkable performance in length generalization, performance, and efficiency, pushing ahead with the boundaries of long-context modeling.", "questions": "1. Why does HSA adopt the stick-breaking attention mechanism for constructing chunk weights, rather than other forms of attention with positional encoding? I would like to see a theoretical justification or supporting ablation studies. (Table 1 presents RAMba w/o m.r. w/o s.b., which uses softmax without positional encoding, but it does not include ablation related to positional encoding)\n\n2. Line 174: Does the equation for the HSA layer in this line omit the Mamba layers? I ask because the formula shows H_t^(l+1)=H_t^l+O_t^l, but in the architecture, there are components other than the HSA layer between H_t^(l+1) and H_t^l, including Mamba layers.\n\n3. Line 273: Why does NSA's implementation currently not support generation?", "limitations": "I agree that experiments with RNNs besides Mamba can be explored further.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n1. A sparse attention design with a two-stage hierarchical structure. Chunk-level information is captured in the first stage and then fused in the second stage, which helps the length generalization. The working memory and long-term memory from a psychological perspective are interesting.\n\n2. A hardware-aligned algorithm to achieve efficient parallel computation.\n\n3. Extensive experiments on long-range language modeling, Passkey retrieval, LongBench, RULER, and fine-tunable datasets. Results of 2.7B models on LongBench and RULER validate the scalability. The result that the running time and memory consumption outperform NSA is encouraging, thanks to the shared chunk memory design. Additionally, the ablation studies regarding memory reset and sticking weights are thorough.\n\nweakness\n\n1. The motivation and the solution of the paper are not tightly aligned. My main concern is that the authors use \"Random Long-context Access for Mamba\" as the central motivation, while the main focus of the paper is actually on proposing the HSA method. In fact, the authors do not design any sparse mechanism specifically for the Mamba block itself, nor do they provide an in-depth discussion on the relationship between HSA and Mamba. Instead, HSA is simply applied within a hybrid architecture. As a result, this narrative may dilute the core contribution of the work. Moreover, I believe that using \"random access\" to describe chunk selection may be misleading, although the phrase contributes to the formation of the acronym \"RAMba\".\n\n2. Line 29-30: Is there any literature supporting the claim that augmenting RNNs with attention mechanisms leads to poor length extrapolation? The authors appear to assume that hybrid models extrapolate worse than pure RNNs, but this assumption is questionable. To the best of my knowledge, hybrid architectures do not necessarily hinder length extrapolation [1].\n\n3. Lines 139 and 157: Does the introduction of the transformer-based bi-directional encoder and the mean-pooled representation compromise the model's causality? MoBA [2] points out: \"The routing to the current block could also violate causality, since mean pooling across the entire block can inadvertently include information from future tokens.\" To address this, they ensure that each token is routed to its respective current block and apply a causal mask during the current block attention. I would be glad to see the authors' perspective on this issue.\n\n4. The results in Figure 3 are excellent. However, I have one question: What mechanism enables RAMba to achieve better extrapolation performance than both Mamba-2 and NSA? Even when the encoder is removed, the results remain strong. If the authors could provide a reasonable explanation for this significant performance gain, it might offer useful insights for future architecture design.\n\n5. Can HSA operate independently of the Mamba blocks? If a model were designed using only HSA, would it still be faster than NSA?\n\n[1] MiniMax-01: Scaling Foundation Models with Lightning Attention.\n[2] MoBA: Mixture of Block Attention", "was_revised": true, "final_justification": "The authors have addressed the concerns I raised. Overall, to enhance the ability of RNNs to accurately retrieve historical information, the authors introduce a new sparse attention mechanism that facilitates precise chunk selection. Through the design of a hybrid architecture, they demonstrate that the proposed method achieves a balance among extrapolation ability, performance, and efficiency. The remaining issues for the authors are limited to refinements in methodology and some wording.", "rebuttal": "We sincerely appreciate your insightful feedback. \n\nW1. **Regarding Motivation, the authors do not design any sparse mechanism specifically for the Mamba block itself**\n\nWe believe that \"Enabling Random Long-Context Access for Mamba\" is a goal, and the sparse mechanism for Mamba block could be one way to achieve it, but maybe not the only way. This paper, \"via Hierarchical Sparse Attention,\" presents another approach to achieving this goal through a hybrid architecture.\n\nAs mentioned in our abstract and introduction, there is an information bottleneck in RNNs due to compressing variable-length sequences into fixed-dimensional representations, which inevitably leads to information loss. In contrast, the kv cache in attention grows with sequence length, which has no information bottleneck.\n\nThus, the current mainstream approach to enhancing Mamba's long-range random access ability is integrating Mamba with full attention [1]. As stated in lines 32-33 of the introduction, this leads to a trilemma. Consequently, our motivation lies in replacing full attention with sparse attention, as reflected in *\"Enabling Random Long-Context Access for Mamba via HSA.\"* \n\n**Regarding Random Access**\n\n\"Random access\" refers to the ability to access any arbitrary past token, not just selecting chunks. For example, in an open-book exam, you locate a page (chunk) and then pinpoint specific details (tokens). Similarly, HSA enables access to any past token by first identifying the chunk and then the token-level details.\n\nW4. **What mechanism enables RAMba to achieve better extrapolation performance than both Mamba-2 and NSA?**\n\nWe believe this is an interesting and important topic, so let us discuss it in advance.\n\nFirst of all, the underperformance of Mamba-2 has been discussed in many studies [5]. The main reason is the information bottleneck issue, as we mentioned above.\n\nTo know why HSA achieves better extrapolation performance, let's take a close look at how NAS and HSA work in the following example.\nAssuming every 2 tokens form a chunk, separated by |:\n```\nchunk ids             |       0      |      1    |      2    |     3     | ...\ntoken positions       |   0   |  1   |  2  |  3  |  4  |  5  |  6  |  7  | ... , i <- the currrent token\nattn logits[i, *]     | -20.0 | 20.0 | 1.0 | 0.0 | 2.0 | 1.0 | 1.0 | 1.0 |\nattn weights (softmax)|   0.0 |  1.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |  // O(L$^2$), L: sequence length\nchunk attn logits sum |      0.0     |    0.5    |    1.5    |    1.0    |  // O(L*(L/S), S: chunk size\n```\nNSA selects chunks according to \"chunk attn logits sum\" for efficiency.\nIf only the top-2 chunks can be selected, NSA would choose chunks 2 and 3, thus missing the chunk 0 with the highest sum of attention weights, **leading to inaccurate chunk selection**.\n\nAssuming NSA selects the top-2 chunks for sparse attention, the following process would be:\n``` \nchunk ids         |      2      |      3      | \ntoken positions   |   4  |   5  |   6  |   7  |\nattn logits       |  2.0 |  1.0 |  1.0 |  1.0 |\n                  |<-----     softmax   ----->|\nattn weights      | 0.49 | 0.17 | 0.17 | 0.17 |\n``` \n\nThroughout this process,  chunk selection can not be explicitly learned.\n\nFor the same example, HSA works as follows:\n```\nchunk ids          |       0      |      1     |      2    |     3     | ...\ntoken positions    |   0   |   1  |  2   |  3  |  4  |  5  |  6  |  7  | ...\nattn logits        | -20.0 | 20.0 |  1.0 | 0.0 | 1.0 | 1.0 | 1.0 | 1.0 | ...\ns[i,c]             |    s [i,0]   |  s [i,1]   |  s [i,2]  |  s [i,3]  | ...\n```\nWhere s [i,c] represents learnable token-to-chunk relevance scores as described in Line 155 of the paper. For example, HSA may also select chunks 2 and 3 when parameters are randomly initialized, we have:\n```\nchunk ids            |           2           |          3          |\ntoken positions      |     4     |     5     |     6    |     7    |\nattn logits          |   2.0     |    1.0    |    1.0   |    1.0   |\nintra-chunk weighting|<--     softmax      ->|<--    softmax     ->|\ntoken-level weights  |   0.73    |    0.27   |    0.5   |    0.5   |\ns[i,c]               |        s[i,2]         |        s[i,3]       |\ninter-chunk weighting|<---            stick-breaking           --->|\nchunk-level weights  |        w[i,2]         |        w[i,3]       |   // s.t. w [i,2} + w [i,3] <= 1\nfinal weights        |0.73×w[i,2]|0.27×w[i,2]|0.5×w[i,3]|0.5×w[i,3]|\n```\n\nIn this process, w[i,c], derived from s[i,c], participates in the final attention weight computation and the whole forward pass. This allows it to receive gradients and allocate higher weights to chunks with more important tokens. Even if random initialization initially misses the most relevant chunk 0, continuous training enables the model to gradually learn to select the most relevant chunks.\n\nSuch a learnable chunk selection mechanism is essentially a kind of token-to-chunk retrieval, which may lead to better extrapolation performance\n\nW2. **Is there any literature supporting the claim that augmenting RNNs with attention mechanisms leads to poor length extrapolation?**\n\nYes.\nSpecifically, lines 29-30 refer to hybrid architectures like Mamba with full attention. Mamba with sliding window attention is not included in this discussion. The issue stems from the poor length extrapolation ability of full-attention, which has never been fundamentally solved.\n\nFor instance, Figure 1 of YaRN [4] shows how perplexity (ppl) increases when full-attention-based models with RoPE extrapolate to unseen context lengths.\n\nTable 2 in [7] also reveals that, even with temperature scaling and discard positional encodings, perplexity still dramatically increases beyond a certain length, indicating poor extrapolation.\n\nThe existing full-attention-based models with 1 million context lengths primarily rely on techniques like ring-attention to brute-force scale the pretraining length, which requires huge computational resources.\n\nAs minimax[6] mentioned in their abstract, \"The context window of MiniMax-Text-01 can reach up to 1 million tokens **during training** and extrapolate to 4 million tokens during inference at an affordable cost. ” They also acknowledge the limitation of only being able to extrapolate up to 4 times the training length. The 1 million pre-training context window is also achieved via ring-attention.\n\nOur work achieves extrapolation from 4K to 64M (16000×) in passkey retrieval and maintains high performance in both training and inference.\n\nW3. **Does the introduction of the transformer-based bi-directional encoder and the mean-pooled representation compromise the model's causality?**\n\nNo, it doesn't. **Each token cannot access the chunk it belongs to**, it can only access the previous chunks. We will make further clarification in the revised version.\n\n\nW5. **Can HSA operate independently of the Mamba blocks? If a model were designed using only HSA, would it still be faster than NSA?**\n\nWe believe so. It can work together with sliding window attention. We believe HSA would still be faster than NSA because we do not require token compression attention, thus the overall FLOPs are lower.\n\nQ1. **Why does HSA adopt the stick-breaking attention mechanism for constructing chunk weights, rather than other forms of attention with positional encoding? I would like to see a theoretical justification or supporting ablation studies.**\n\nIf we use RoPE, a straightforward issue would be the out-of-domain problem during extrapolation (when the input length exceeds the pretraining length), as discussed in [4] and [7].\n\nHere is an ablation study to replace the stick-breaking mechanism with softmax+RoPE:\n\nppl with context length=64K  (the lower the better):\n| Models | pg19 | arxiv | code |\n|-|-|-|-|\nRAMba stick-breaking | 17.01| 3.65 | 3.07 |\nRAMba rope + softmax | 17.80 | 4.19 | 3.32 |\n\nRULER: (context len=64K) (the higher the better)\n| Models | S-N | MQ-N | VT | FWE|\n|-|-|-|-|-|\nRAMba stick-breaking | 85.07 | 55.22 | 55.22 | 53.73 |\nRAMba rope + softmax | 43.29 | 26.87 | 29.85 | 37.31 |\n\n\nQ2. **Does the equation for the HSA layer in this line omit the Mamba layers?**\n\nNo, sorry for the confusion.\nIt doesn't omit the Mamba layers.\nThe layer indices here include both Mamba and HSA layers, and we will clarify this point further later.\n\nQ3. **Why does NSA's implementation currently not support generation?**\n\nAs there is no official implementation for NSA thus we use an implementation from the open-source community (line 230).\n\nNSA's batch forward mode (used for training) calculates next-token perplexity for full sequences without needing KV cache management for new tokens, as no new tokens are generated. In contrast, generation mode requires handling a variable number of newly generated tokens, complex kv cache management, and positional encoding, especially with different attention mechanisms. This complexity is why vLLM performs well for inference but not training, while Megatron is optimized for the reverse. The KV management for generation was incomplete at the time of the paper's submission.\n\nEven so, this does not undermine our conclusions. The majority of experiments can be evaluated via batch forward passes, and the performance on PPL and RULER is sufficient to support our claims.\n\n[1] Jamba: Hybrid Transformer-Mamba Language Models\n\n[2] SAMBA: SIMPLE HYBRID STATE SPACE MODELS FOR EFFICIENT UNLIMITED CONTEXT LANGUAGE MODELING\n\n[3] Understanding and Improving Length Generalization in Recurrent Models\n\n[4] YaRN: Efficient Context Window Extension of Large Language Models\n\n[5] Repeat After Me: Transformers are Better than State Space Models at Copying\nTransformers are Better than State Space Models at Copying\n\n[6] MiniMax-01: Scaling Foundation Models with Lightning Attention\n\n[7] Length Generalization of Causal Transformers without Position Encoding\n\nWe hope our replies address your concerns, and we would appreciate it very much if you could give us more support."}]

Author rebuttal:

Meta-review: The paper introduces Hierarchical Sparse Attention (HSA), designed to augment RNN architecture, with efficient long-range random access capabilities while preserving linear computational complexity and length generalization. The core claim is that HSA addresses a "trilemma" in long-context sequence modeling: balancing efficiency (linear time/space), random-access flexibility (ability to retrieve arbitrary historical tokens), and length extrapolation. HSA divides sequences into chunks, learns via a two-stage hierarchical process (intra-chunk encoding and inter-chunk aggregation using stick-breaking attention), and integrates hardware-aligned kernels, including KV cache offloading to CPU. When combined with Mamba-2, the resulting RAMba model is pretrained on 4K-length contexts but extrapolates to 64M tokens, achieving perfect accuracy on passkey retrieval tasks up to 64M, near-constant memory footprint during inference, and  comparable performance on benchmarks like RULER and LongBench. Reviewers characterize the work as a meaningful advance in chunk-based sparse attention. 

Strengths:
1. A learnable two-stage sparse attention that directly addresses the efficiency–random access–length generalization trilemma. Overall the method is well-motivated and concretely presented.
2. Experimental results are solid.
3. Hardware-aligned kernels, and KV offload strategy makes it appliable

Weakness:
1. Generally, the proposed architectural is heavy in design, with HSA + chunk encoder + state passing + Mamba.
2. The focus is limited to Mamba, with no exploration of other RNNs or larger scales

I recommend acceptance as a poster. The key reasons include the paper's strong technical contributions to a high-impact area (long-context modeling in efficient architectures), backed by rigorous experiments and ablations.

The majority of reviewers (three out of four) rated it as accept or borderline accept, with high confidence, and the rebuttals effectively resolved most concerns through additional results.
Decision: Accept (poster)