Overcoming Long Context Limitations of State Space Models via Context Dependent Sparse Attention
Authors: Zhihao Zhan
Jianan Zhao
Zhaocheng Zhu
Jian Tang
Venue: NeurIPS 2025 poster ()
Keywords: long-context modeling
state-space models
sparse attention
novel architectures design
Primary area: deep_learning
TL;DR: We propose joint recall, a novel synthetic task, and hybrid sparse attention with context-dependent sparsity for better sub-quadratic long-context modeling.

Abstract: Efficient long-context modeling remains a critical challenge for natural language processing (NLP), as the time complexity of the predominant Transformer architecture scales quadratically with the sequence length. While state-space models (SSMs) offer alternative sub-quadratic solutions, they struggle to capture long-range dependencies effectively. In this work, we focus on analyzing and improving the long-context modeling capabilities of SSMs. We show that the widely used synthetic task, associative recall, which requires a model to recall a value associated with a single key without context, insufficiently represents the complexities of real-world long-context modeling. To address this limitation, we extend the associative recall to a novel synthetic task, joint recall, which requires a model to recall the value associated with a key given in a specified context. Theoretically, we prove that SSMs do not have the expressiveness to solve multi-query joint recall in sub-quadratic time complexity. To resolve this issue, we propose a solution based on integrating SSMs with Context-Dependent Sparse Attention (CDSA), which has the expressiveness to solve multi-query joint recall with sub-quadratic computation. To bridge the gap between theoretical analysis and real-world applications, we propose locality-sensitive Hashing Attention with sparse Key Selection (HAX), which instantiates the theoretical solution and is further tailored to natural language domains. Extensive experiments on both synthetic and real-world long-context benchmarks show that HAX consistently outperforms SSM baselines and SSMs integrated with context-independent sparse attention (CISA). Our code is available at: https://github.com/DeepGraphLearning/HAX.

Overcoming Long-Context Limitations of State-Space Models via Context-Dependent Sparse Attention
Zhihao Zhan 12., Jianan Zhao Zhaocheng Zhu Jian Tang
Mila Québee Al Institute, Univeessity of Montréal 'HEC Montréal, *CIFAR Al Chair Equal contribution "Correspondence: tangjian langian@mmll..uee mila.quebec
Abstract
Efficient long-context modeling remains a critical challenge for natural language processing (NLP), as the time complexity the predominant Transformer archi- tecture scales quadratically with the sequence length. While state-space models (SSMs) offer alternative oub-quadratic solutions, they struggle to capture long- range dependencies effectively. In this work, we focus analyzing and improving the long-context modeling capabilities of SSMs. We show that the widely used synthetic task, associative recall, which requires model to recall value associ- ated with a single key without context, insufficiently represents the complexities of real-world long-context modeling. To address this limitation, associative recall to a novel synthetic task, joint recall, which requires a model to recall the value associated with key given in specified context. Theoretically, we prove that SSMs do not have the expressiveness to solve multi-query joint recall in sub-quadratic time complexity. To resolve this issue, we propose solution based on integrating SSMs with Ceneext-Dependent Sparse Attention (CDSA), which has the expressiveness to solve multi-query joint recall with sub-quadratic computation. To bridge the gap between theoretical analysis and real-world applications, we propose locality-sensicee
extend the
Hashing Attention with sparse Key Selection (HAX), which instantiates the theoretical solution and is further tailored to natural language domains. Extensive experiments on both synthetic and real- real-world long-conted benchmarks show that HAX consistently outperforms SSM baselines and SSMs integrated with context-independent sparse attention (CISA). Our code available at: http://hhh..ccoo//eeearaaaaaaaaa
Introduction
Long-context modeling a central challenge natural language processing (NLP), which underpins variety of applications, such as document summarization, question answering, and machine translation [3]. Recent advances in large language models (LLMs) have broadened the landscape of long-context modeling, enabling new capabilities such as autonomous agents, retrieval- -atreeal- generation, dialogue systems, and long-context reasoning [22]. This growing demand has spurred intensive research into algorithms that can efficiently and effectively capture long- long-range dependencies [37].
Currently, the Transformer architecture (38 is the dominant paradigm in sequence modeling. How- ever, its applicability to long sequences is fundamentally constrained by the required computation that grows quadratically with the sequence length. This motivates the research direction for the invention of efficient architectures.
Recently, state-space models (SSMs) 10 16 31 have emerged as potential alternative solution, offering sub-quadratic time complexity as well as comparable performance to Transformers on short-context NLP tasks [1 However, empirical evidence suggests that SSMs are less effective
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Associative Recall does not specify context.
The occans are blue, and snow is white.
White!
What color is snow?
Joint Recall specifies context.
the U.S., people pay with dollars and drive cars on the right side of the road. In the U.K. people pay with pounds and drive cars on the left side of the road.
Which side roud do people drive in the U.S.?
Right!
Figure 1: Comparison of joint recall and associative recall. Associative recall does not account for context. Joint recall extends associative recall by incorporating context-dependency into key-value associations, For example, while associative recall may map "pay with" to either "dollar" or "pound", joint recall allows map "dollar" in the U.S. and "pound" in the U.K., depending on context. This makes joint recall more realistic and rigorous synthetic task for both theoretical analysis and empirical benchmark for long-context modeling.
than Transformers in capturing long-range dependencics [40]. Furthermore, theoretical analysis by 19 demonstrates that SSMs much less capable of handling long-context copying, due to the limitations of architecture representation capacity.
In this work, we aim to better understand and improve the long-context modeling abilities of SSMs. We first show that previous studies based on the widely used synthetic task, associative recall (2] might be constrained by its limited capability to simulate natural language in-context dependencies. To be specific, associative recall assumes that each key uniquely associated with a value, regardless of the surrounding context. However, natural language often defies this assumption: the same key can correspond to different values depending on its context. Consider the example in Figure [] when asked on which side of the road people drive, the correct answer should depend on the country being referenced. Without specifying whether the context is the US or the UK, the question becomes ambiguous. This example highlights& critical shortcoming of associative recall: it lacks the capacity to simulate context-dependent key-value association, which is very common in natural language.
To address this limitation, we extend the associative recall to more general synthetic task, joint recall. Unlike associative recall, joint recall requires the model to retrieve a value corresponding to key conditioned on specified context. Theoretically, we prove that standard SSMs lack the representational capacity to solve multi-query joint recall under sub-quadratic time complexity.
To overcome this expressiveness bottleneck, we propose to augment SSMs with Context-Depeneee Sparse Attention (CDSA), a class of sparse attention with sparse attention patterns that are conditioned on the context representations. Locality-sensitiv hashing (LSH) attention [20] exemplifies CDSA, while context- inndependen sparse attention (CISA) includes sliding window attention, A-shaped attention, and dilated attention [6] Compared to CISA, CDSA enables dynamic content-dependent routing of information, which essential to efficiently solve the multi-query joint recall task. We theoretically show that there exists CDSA which, when integrated with SSMs, enables solving the multi-query joint recall task in sub-quadratic time with respect to sequence length. Moreover, we establish an expressiveness gap between SSMs integrated with CDSA and SSMs integrated with CISA on multi-query joint recall.
Building upon this insight and to bridge the gap between theory and practice, we propose novel architecture: senssity- Hashing Attention with sparse Key Selection (HAX), HAX improves the expressiveness of LSH attention by incorporating our proposed Key Selection (KS) attention, and is further integrated with state-of-the-ar SSMs, Mamba and Mamba2 [44 4 instantiating the theoretically grounded solution. We validate the effectiveness of HAX through extensive experiments on both synthetic and real-wor wal-world long-context modeling benchmarks. The experiment results show that HAX consistently outperforms SSM baselines as well as SSMs augmented with CISA. These findings demonstrate that CDSA, when carefully integrated with SSMs, a critical component in unlocking their potential for long-context modeling.

---

Our main contributions are summarized as follows:
1. We introduce joint recall, novel synthetic task that extends associative recall context- dependent key-value association, which offers a new perspective for both theoretical analysis and empirical benchmark for long-context modeling.
2. Through theoretical analysis on joint recall, we demonstrate that integrating state-space models (SSMs) with context-dependent sparse attention (CDSA) has the expressiveness to solve multi- query joint recall with sub-quadratic computation.
3. Guided by this theoretical insight, we propose a novel architecture, HAX, based on SSM integrated with CDSA, which consistently outperforms SSMs and SSMs integrated with context- independent sparse attention (CISA) on both synthetic and real-world long-conted context benchmarks.
Preliminaries
In this section, introduce two prominent approaches for efficient architecture design: sparse attention in Sec. 2.1 (exemplified by LSH attention in Sec. (.2) and SSMs in Sec. [23 We also introduce associative recall, widely-used synthetic benchmark for long-context modeling, Sec. 2.4
2.1 Sparse Attention
For a sequence of length 4, we denote the attention scores auto-regressive sequence modeling as:
A Softmax (M
QK")
where M {0, 1 is the auto-regressive mask, Q, K E Rixd, d the hidden dimension. this work, we define the attention scores for sparse attention
A Softmax & M O QK
(2)
with S E {0, 1 representing the sparse attention pattern. We consider the sparsity constraint:
| < (2
(3)
To ensure per-step computational efficiency, we further tighten this constraint by requiring:
Vi, S <!
(4)
where S; denotes the i-th row of S. This implies that each query attends most keys, < I.
2.2 Locality-Sensine Hashing Attention
Given that locality-sensinee hashing (LSH) attention represents one of the most effective input- dependent sparse attention for auto-regressive modeling (20), reformulate simple algorithm to generate the sparse attention pattern LSH. This algorithm accepts the query and key matrices Q and K as input and outputs binary sparse attention pattern SESH- At each forward pass, Q and K first centralized and normalized to and K, respectively:
Q normalize(Q2 Q; Qi). K; normalize(K, K;)
Next, random projection matrix H i.d. N/ 1) Rdxh is sampled to project the normalized vectors Q and K into the hash space. Then, we consider binning rules which assign each vector Q and K, hash bin: the argmax binning rule (20 35 and the sign-bit binning rule {8 58.
The argmax binning rule assigns each vector to the index of its most aligned column in H:
bing, = argmax Q.H). bink; = argmax (K,H)
(6)
The sign-bit binning rule constructs binary hash code by computing the signs of the projected values and interpreting as binary number:
h
bing. Σ I(QHH) 2h bink, Σ
I(K,H) > 0] 2
j=1
j=1
3

---

The argmax binning rule assigns vectors to h bins, while the sign-bit binning rule assigns vectors to 2 bins. We will further discuss the relationship between these two binning strategies Appendix A Based on the assigned bins, a preliminary sparse pattern SLSu is constructed by allowing each query to attend to all preceding keys within the same bin:
SLSH, = [binq. =bink]]
(8)
Finally, satisfy the sparsity constraint defined in Eq. 4 per-bin sliding window mask MISH is applied, so that each query only attends to at most is nearest keys in the same bin:
SISH M @ SiSH
2.3 Generalized State-Space Model
Following the definitions introduced by 19 we formulate generalized state-space models as sequence models defined by an update rule u U V U and an output function r U + V. where V denotes the token vocabulary and 2/ represents the recurrent state. Let (o(Ø) E U denote the initial state. Given an imput sequence 01, 2/ E for in (1...z), the state (1 ... u¹) E and its corresponding output (",,,.... (U,,,..,
U/) E V are defined recursively as:
U/ 02) = u(U;-1 (1), ... 1/-1), v.) (u v4) = (U)(())..,,e))
(10)
(11)
2.4 Associative Recall
The associative recall task was originally introduced in (2] (27 found that the LLM performance on this task is strongly correlated with their in-context learning abilities. [ extended associative recall to the multi-query setting: a model is first given a sequence of associated key-value pairs, and then required to recall each value when queried with the corresponding key. Associative recall has been widely adopted synthetic benchmark for long-context modeling 4 18).
3 Joint Recall
We discuss the motivation and formulation of joint recall in Sec. finally provide the theoretical results in Sec. 3.3
3.1 and Sec. 3.2 respectively, and
3.1 Motivation
The motivation behind proposing joint recall is to overcome key limitation of the setup of associa- tive recall: each key corresponds to single fixed value. While this setup well-suited for studying the tasks that emphasize capturing stable lexical patterns, such as sub-word units or fixed multi- word expressions, it falls short in representing the context-sensitive nature of meaning natural lan- guage. Consider the following examples:
Multi-Query Associative Recall recalls each value associated with single key without context:
Multi-Query Associative Rocall
Input a5b2e3d (36.3d1e14c22230 :5b2c3dle4jc77227dd77
Output 34512
Multi- Multi-Query Joint Recall recalls each value associated with key given in specified context:
The legislative branch of the U.S. govern-
Multi-Query Joint Recall
ment is called Congress.
Input
Bb 4 B Ba?bbAb?a?
On Monday mornings. Alice studies math.
Output
1423
From a philosophical perspective, definitions are often constructed using genus keys and differentia context. In the first example, the value "Congress" is identified with the genus key "the legislative branch" and the differentia context "of the U.S. gov- ernment". The second example reflects a more daily scenario, where the value "math" is identifiable only when all contextual elements, "Monday" and "morning", are considered together with the key "Alice". These cases illustrate that accurate semantic interpretation in natural language often requires integrating context and keys, suggesting that models must move beyond the simplistic one-to-one mappings of associative recall capture the compositional and context-dependen nature of meaning. This observation motivates our introduction of novel synthetic task, which we refer to joint recall.
Figure 2: Comparison between synthetic multi-querr joint recall and associative recall.

---

3.2 Formulation
Associative recall requires a model to memorize ng associated key-value pairs. Joint recall generalizes this task: the model is required to memorize an ne X table of context-specific key-value associations, in which ny keys are associated with different values each of the n, contexts. Inspired by (. we also extend joint recall to a multi-query setting, requiring the model to recover the entire table instead of a specific entry the table.
Figure [ illustrates multi-query joint recall with n. = 2 and ny = 2. Following the structure of natural language, the sequentialized table input consists of nhe context blocks, each beginning with a context token (e.g. uppercase letter in Figure 2). followed by ny key-value pairs specific to that context (e.g, lowercase-lette-diigt pairs in Figure 2) Then, the model is asked to recall the associated values given each context-key pair, under arbitrary permutations of the context and key ordering.
Appendix C further extends the joint recall formulation to multi-level context: for example, in the sentence OOn Monday mornings, Alice studies math.", "Monday" and "morning" are contexts at hierarchical levels. It also provides theoretical analyses grounded in this extended formulation.
3.3 Theoretical Analysis
3.3.1
Categorization of sparse attention
The sparse attention pattern S defined in Sec. 2.1 can be categorized as context-dependent or context- independent, depending on whether it is predetermined or dynamically inferred from the context representations. Context-independent sparse attention (CISA) patterns, such as sliding window attention, A-shaped attention, and dilated attention [6], are fixed regardless of context. In contrast, context-dependent sparse attention (CDSA) patterns, exemplified by LSH attention [20), adapt according to the context representation. Appendix Figure 8 provides an illustration of both categories.
3.3.2 Limited expressiveness of SSMs
we demonstrate that solving the multi multi-query joint recall task imposes linear growth requirement on the state dimension of SSMs with respect to the number of entries n the joint recall table, = ne x ng. Let [ be the number of distinct representations that the recurrent state space U can encode, for a state with b bits of capacity, [u4] 2 We define the uniform multi-query joint recall distribution as the distribution in which all values are sampled i.i.d. from the uniform distribution over the token vocabulary V. In this setting, we obtain the following:
As corollary of Theorem 2.7 [9].
Corollary 3.1 (Limited expressiveness of SSMs). Under the uniform multi-query joint recall distri- bution, for any n, generalized state-space model defined in Sec. 2.3 incurs an error rate of at least -
1
Remark 3.1. To guarantee Prjerr] = 0, it is necessary that the number of representable states satisfies [24 z [V]". Taking the logarithm of both sides yields the condition 6 log V/. This implies that the state-space dimension of the model must grow linearly with the number of entries in the joint recall table, highlighting fundamental limitation of the representation capacity of SSMs.
3.3.3 Improved expressiveness of SSMs integrated with CDSA
For SSMs integrated with sparse attention, we establish the following results:
Proposition 3.2 (Improved expressiveness of SSMs integrated with CDSA). There exists 2-layer auto-regressive hybrid model consisting of an SSM layer followed by an LSH attention layer, which can solve multi-query joint recall in O log" n) time complexity with Oglog n) SSM state dimensions.
Proposition 3.3 (Limited expressiveness of SSMs integrated with CISA), There does not exist 2-layer auto nuto-regressive hybrid model consisting of an SSM layer followed by a CISA layer, which can solve multi-query joint recall with o(n²) time complexity, since requires at least ( SSM state dimensions, l is the maximum number of keys that each query allowed to attend to in the sparse attention module, as defined in Eq.
we see a clear representation capacity gap between the SSMs integrated with CDSA and the SSMs integrated with CISA.
Remark 3.2. Comparing Proposition 3.2 with Proposition 3.3
5

---

LSH atlention
Key Selection (KS) attention
Assign each query/ key hash bin. Calculate attention within each bin.
1. Calculate seore for each key.
Queries attend top-scored keys.
LSH e (04))
hin
SKS (0.1p
Conv
1
kdd
Kidk
Handles interactions context.
Memorizes important keys.
SHAX = max{SLsH Sks)
Integrating HAX with Mamba
Figure Illustration of the HAX architecture, and Mamba 14 integrated with HAX.
Remark 3.3. In practice, with appropriate constant k integrating CISA with SSMs still provides an advantage: unlike SSMs, which only have access to the last state representation, CISA layers can attend & different state representations simultaneously, a cost of k times of computation budget.
Complete proofs are provided in the Appendix
Method
Guided by the theoretical analysis Sec. 3.3 we propose new architecture, locality-sensitive Hashing Attention with sparse Key Selection (HAX). HAX improves the expressiveness of LSH attention by incorporating our proposed Key Selection (KS) attention, and further integrated with tftteoff SSMs, serving as instantiation of context-dependent sparse attention (CDSA) integrated with SSMs, thereby benefiting from the theoretical advantages discussed Sec. 3.3
In this section, we first discuss the expressiveness limitations of LSH attention in Sec. 4.1 and then address these limitations by introducing our proposed key selection (KS) attention Sec. Finally, Sec. 4.3 details the architecture of HAX as well as how HAX is integrated with state-of-the SSM architectures, Mamba and Mamba2 14
4.1 Limitation LSH Attention
In LLMs, certain keys (particularly those at the beginning of sequence) often receive attention from most queries, forming distinctive "vertical-stripe" attention patterns (39], as illustrated Appendix Figure 9 These globally attended keys play essential role in instruction following, where the model is expected to focus its attention on the instruction tokens (23).
Although LSH instantiates CDSA, as discussed in Sec. 2.2 it suffers from a key limitation: difficulty in capturing "vertical-stripe" attention patterns. This arises because in each hashing round, every key is mapped to single bucket, and attention is constrained to occur only between queries and keys within the same bucket. As result, when many queries are forced to attend to a limited set of key buckets, those buckets become overloaded, diminishing representation diversity and ultimately degrading attention quality.
4.2 Key Selection (KS) Attention
Goals. To address the limitation of LSH attention in capturing "vertical-stripe" attention patterns, propose to augment LSH attention by integrating it with a novel key selection (KS) attention module. This module designed to satisfy the following desirable properties:
1. "Vertical-stripe" capability: KS attention can express "vertical-stripe" attention patterns.
2. Auto-regressive compatibility: The computation of KS attention for the current token does not depend future queries or keys.
3. dontext-dependen sparsity: KS sparse attention pattern conditioned on the query and key representations in context and satisfies Eq. 4

---

Modeling, Taking the query and the key matrices Q and K as input, KS attention operates in two phases. The first phase is key scoring, where an scoring module computes an importance score for each key based on the key itself and all previous queries:
2i - (K; QI...))
(12)
The second phase is key selection: each query attends to the previous highest-scoring keys,
SKS;, [2y E Top-k(
(13)
With an ideal key scoring module that assigns the highest scores the globally important keys, KS attention effectively covers k "vertical-stripe" within attention patterns.
For simplicity, multilayer perceptron (MLP) the key scoring network:
fo Q1..s MLP normalize
Qp)) (O2))
(14)
ISPSi
Training. To train the scoring MLP, at each layer, we randomly sample & candidate keys. Their indices are denoted by Z. We compute the reference attention weights via a simple linear attention module, and calculate a pairwise ranking loss between these reference weights and the predicted scores. To be specific, we compute:
A' = Q K
y= y= (A') O M
(15)
where K[Z] is the selected key representations, M is the auto-regressive mask restricted to those positions, and (-)) is the sigmoid function. With predicted scores x E R& and target scores E IR*, we construct pairwise logits and targets:
if y Vj- if y Fj
Pı)()) = Ii Ij.
0.5
(16)
if yi y3
and define the ranking loss:
Lacare [x, 1)
BCE (P (e))),TT(())))
(17)
12
where BCE(..) denotes binary cross-entropy. This objective encourages the scoring network to assign higher scores to informative keys that receive higher attention weights. The final training objective sums the ranking loss across layers with the auto-regressive language modeling loss LLM:
= LLM + a Σ Lesore
(18)
layers
where a is a scalar hyperparameter that balances the contribution the ranking loss.
4.3 Hybrid Block Design
Finally, we propose combines LSH and KS attention patterns:
sensitive Hashing Attention with sparse Key Selection (HAX), which
SHAX max (SLSH, Sks} E {0, 13
(19)
When Vi, SLSn, llo S 5. 1/1|, llo - 2: it satisfies
Vi, SSA llo I k
(20)
Intuitively, LSH and KS attention are complementary, each addressing the other's limitations. LSH attention routes each query to semantically similar keys through randomized hashing, offering flexible, content-based interactions that KS attention alone lacks. In contrast, KS attention introduces broadcast connections to a small set of globally important keys, such as instructions or formatting markers, thereby recovering the "vertical-stripe" patterns that LSH attention struggles express. While LSH attention promotes diverse contextual representations, mitigating risks representation collapse, KS attention sharpens focus by allocating attention weights to the most informative positions, enabling stronger long-range control. Importantly, both mechanisms are inherently sparse, so their combination introduces sub-quadratic computational cost.
Figure 3 illustrates Mamba-based and Mamba2-base HAX layer. The proposed hybrid sparse attention layer mitigates the representation capacity limitations of SSMs by coupling them with parallel sparse attention branch. A parameterized gate rescales the sparse attention output before fusion, which promotes stable optimization.
7

---

Validation Loss vs Training Steps
Table 1: Results of multi- query joint recall. Integrating Mamba or Mamba2 with HAX achieves the best performance. which is in bold.
2.164
Kanta
Kamiu+ 1D
Mumbo LSH
Mumbu A.A.
MomntoeK3 +25
2.162
XaXX
Validition las
2,160
2.150
Mamba
Mamba2
36.6
2.156
+D
7.8
19.6
+SW
18.7
70.6
2.154
+SW+D
16.6
48.6
+A +LSH
16.4
49.3
10000
20000 30000
60000 50000
11.6
13.5
Training Stops
KS (owrs)
36.6
60.1
Figure 4: Validation LAM during continual training on Mamba 130M models. Mamba integrated with HAX is the only architecture that consistently exhibits decreasing vali- dation loss over the entire training process.
HAX (owrs)
Base (2 (2xdimm)
34.5
59.2
Samba
6.3
5
Experiments
We conduct experiments from two perspectives. First, we empirically verify the theoretical findings presented in Sec. on multi-query joint recall. Then, we demonstrate the effectiveness of the proposed hybrid sparse architecture across synthetic and real-world NLP benchmarks.
5.1 Empirical Verification on Multi-Query Joint Recall
Data. We construct a multi-query joint recall dataset in which the number of context blocks and the number of keys per context independently sampled from the range [5, 16], and the size of the vocabulary is fixed at [V] = 16. The dataset consists of 1.4M training examples, along with 14.4K samples each for validation and testing.
Baselines. We adopt Mamba 14 and Mamba2 as base architectures and evaluate various hybrid
sparse attention models built upon them, as illustrated in Figure These include dilated attention (D),
sliding window attention (SW), combination of sliding window and dilated attention (SW+D) [6]. A-shaped attention (A), attention, and HAX (a combination of LSH and KS attention). We also consider a Samba baseline.
hashing (LSH) attention, our proposed key selection (KS)
Setup. For fair comparison, we fix the hidden size to 64 and fix k = 64 defined in Eq. 4 (the maximum number of keys each query can attend to) across all hybrid architectures. For variants that integrate multiple sparse attention components, namely SW+D, A (consists of a SW component and a sink attention component, sink attention only attends the first tokens in the sequence), and HAX (LSH+KSS, we allocate k = 32 to each component, in order to maintain global h: = 64. We also consider a baseline that doubles the hidden size of SSMs to compensate for the absence of hybrid sparse attention. We fix the number of layers to 2 for all models. We conduct experiments using 3 different random seeds for each learning rate {3e-3, 1e-3, 3e-4), and report the average performance corresponding to the best-performing learning rate. For evaluation, we calculate mean accuracy par-sample and average over the testset.
Results. As summarized in Table [ compared to the Mamba and Mamba2 base architectures, most hybrid sparse attention models improve performance. In particular, our proposed HAX model consistently achieves the best performance, which underscores the effectiveness of our approach.
5.2 Continual Pre-training on Natural Language
Setup. To evaluate our method on real-world long-context natural language data, we perform continual pre-training based on the publicly released Mamba checkpoints [14. As in section 5.1 we augment the Mamba architecture with sparse attention, as illustrated Figure 3 We include SSMs integrated with CISA baselines, along with ablated variants of our HAX model. We fix the sparsity parameter k = 128. For variants that integrate multiple sparse attention components, namely SW+D A, and HAX (LSH+KS), we set & 64 to each component, in order to maintain global & 128.
8

---

Table 2: Ruler benchmark at 2K context length. We compare different sparse attention integrated with Mamba 790M model, including CISA methods: dilated attention (D), sliding window attention (SW), and their combination (SW+D), and A-shaped attention (A), and CDSA methods: LSH attention, and our proposed key selection (KS) attention and HAX. The best average performance in bold.
Model
NIAISI
NIAH52
NIARS3
NIAIMKI
NIAHMK2
NIAHMK3
NIAHMV
NIAHMO
47
CWE
FWE
Q
Average
Mamba
100
95.2
86.8
34.4
2.6
32.4
28.35
2.76
1.4
100
25.2
42.43
174
9.78
-SW
100
93.6
83.6
36.2
344
31.05
31.45
5.24
1.26
100
24.4
42.52
-SW+D
100
95
82.4
34
2.6
32.9
29.45
4.72
1.22
100
24.8
42.26
-LSH
0
0
0)
0
0
0.2
0.06
100
8.36
-KS Tours)
100
94.4
84.4
33.6
1.8
37.25
33.85
4.04
43.10
-HAX (owzz)
Table 3: LongBench English tasks. We compare different sparse attention integrated with Mamba
790M model, including CISA and CDSA methods. The best average performance is in bold.
Model
2WiLMQAA
GivRopert
HolfoQA
LCC
MullNews
MaliFQA
MuSiQue
NQA
PsgCnt
PaoRet
Qumpar
QM55m
Repetench
SamSum
Tree
Trivio(A
Average
Mamha
13.39
20.3
8.37
39.51
20.36
27.62
3.89
8.35
1.44
10.69
20.36
38.47
20.15
37.59
19.79
+D
0.21
2.87
0.26
0
0.72
D
0.51
+SW
13.81
20.58
9.25
38.06
20.44
27.55
5.15
8.12
1.2
1.5
10.55
19.57
38
22.42
36.55
19:77
+SW+D
13.6
20.39
8.55
39
20.74
27.55
4.35
7,98
0.83
1,56
10.17
19.73
38.84
22.49
37.02
19.74
27.31
+ISH
1.66
1.49
1.5
7.74
0.8
2.32
0.96
0.7
0.1
0
5.08
7.13
3,89
0.9
2.31
KKS (aurs)
12.93
20.93
9.34
38.22
20.83
26.56
4.34
7.46
0.27
L.77
11 11.65
19.38
38.
38.17
22.97
45
37.52
19.83
(ours)
19.95
Table 4: Performance of extrapolation from 2K to 4K context length on Ruler benchmark. We compare different sparse attention integrated with Mamba 790M model, including CISA and CDSA methods as in Table The best average performance is in bold.
Model
NIAISI
NAAS8
NIAIS1
NAARRIKI
NIARMIK
NIAHMK3
NAAHRY
NEAHMMQ
VT
CWE
FWE
0.11
Average
Mamba
100
19.2
0.65
1.85
5.72
43
100
26.2
78.6
27.01
+D
0
0
0
0
0
65.2
100
0
73
18:32
+SW
100
18.2
5
566
0.2
0.2
1.75
2.1
4.4
4.78
100
22.4
788
26.42
+SW-D
19.8
+1.SH
0
100
0.2
13.33
+KS (ours)
99.8
20.2
7.8
0.2
1.2
1.85
9.32
2.85
100
22.6
78.6
27.50
4HAX (owx))
100
11
9
0.8
3.4
2.75
6.36
2.54
100
29
78.8
28.20
Overall, our HAX adds only about 1 additional parameters to the released Mamba architecture. All models continually pre-trained for 50K steps with context length of 2K.
Validation Loss during Continual Pre-Training. Figure shows the validation loss CLM of 130M model continual pre-training The Pile [T]. As observed, the Mamba base architecture and all baseline variants exhibit either training instability or plateau early in the training process. In contrast, proposed HAX model is the only architecture that shows consistent decline in validation loss throughout the training process, indicating improved stability and sustained learning.
Ruler and LongBench Evaluation. Tables 2 and 3 present the downstream performance for
the Mamba 790M model after continual pre-training for 50K steps on TxT-360 (36 followed by
instruction fine-tuning for 10K steps on UltraChat 7 following Mamba-Chat 24 Ruler [T8] is
synthetic long-context NLP benchmark designed to assess model performance on tasks including retrieval, multi-hop reasoning, aggregation, and question answering. LongBench [ comprises real-world NLP long-context tasks, including question answering, summarization. few-shot learning. retrieval, aggregation, and code completion. The evaluation results, summarized in Tables and show that among all hybrid sparse attention variants, our proposed HAX model is the only one that outperforms the Mamba baseline by significant margin average.
Extrapolation Evaluation. We additionally evaluate the Mamba 790M models on the Ruler 18 benchmark at 4K context length, despite being continually pre-trained and instruction fine-tuned solely with 2K context length. As summarized in Table our proposed HAX model consistently outperforms all baselines, underscoring its robustness for context length extrapolation.

---

Related Work
6.1 State-Space Models
State-space models (SSMs) originated in control theory, exemplified by damped mass-spring systems HiPPO [13] was of the first efforts to adapt SSMs for machine learning applications. LSSL unified convolutional neural networks (CNNs), recurrent neural networks (RNNs), and ordinary differential equations (ODEs) under the SSM framework, enabling their implementation within deep neural networks. H3 [10] integrated SSM layers with short convolutional filters to enhance sequence modeling. Mamba (14] advanced this line of work by making all SSM parameters input-dependent, which significantly increases the representation capacity of SSMs, Mamba2 4 further improved the architecture and established connections between SSMs and Transformer attention.
29
6.2 Analysis on State-Space Models
Recent empirical studies have shown that SSM long-context modeling performance often lags behind that of Transformer architectures [40]. In theory, [19] demonstrated that even solving simple tasks like copying requires SSM state dimensions grow linearly with the sequence length. Furthermore, 25 established that the expressiveness of linear or diagonal SSMs is bounded by the complexity class TC 64 further showed that SSMs and Transformers capture overlapping yet distinct subsets of TC', providing theoretical basis for developing hybrid models that combine the strengths of both architectures.
6.3 Context-Dependeed Sparse Attention
I Locality-sensitiv hashing (L.HH) attention 20 stands as one the most widely adopted forms of context- depeext-dependent sparse attention. Follow-up LSH attention variants can be broadly categorized by their binning rules: the argmax binning rule adopted in 20 35 while the sign-bit binning rule
used in B 5. (28) introduces Triton kernel that accelerates hash-based sparse attention. Recently,
[3 proposed native sparse attention, a new context- context-dependent sparse attention that outperforms Transformers training from scratch, highlighting its potential for efficient long-context modeling.
6.4 Hybrid Architectures
Several works have explored architectures that mix a large proportion of SSM layers with small number of full attention layers, and have reported performance surpassing that of standard Transform- [40]. The effectiveness of such hybrid architectures has been further validated at billion-parametts parameter scale 211 In parallel, researchers have also investigated the design of hybrid sparse attention models
26 44 which offer sub-quadratic computational complexity, providing promising direction for efficient long-contex: modeling.
33 32
Conclusion
In this work, we introduce joint recall, novel synthetic task that generalizes associative recall to context-dependent key-value retrieval. Theoretically, we show that both SSMs and SSMs integrated with context-indes indeeendeen sparse attention (CISA) could not solve multi-query joint recall within sub- quadratic time complexity, while integrating SSMs with context- overcomes this limitation. Guided by this insight, we propose to integrate state-of-the-art SSMs with novel CDSA, locality-sensitivy Hashing Attention with sparse Key Selection (HAX). Experiment results confirm that HAX achieves improved training stability and consistent performance gains across synthetic and real-world long-context NLP benchmarks. The joint recall task therefore offers unified theoretical lens and empirical yardstick for long-context modeling, while HAX demonstrates the power of theory-driven architecture design. These results highlight the importance of aligning model design with expressiveness improvements, and demonstrate that combining efficient sequence models with CDSA is promising direction for scalable long-context modeling.
sparse attention (CDSA)
10

---

Acknowledge
We would like to thank Jingyang Yuan, Wanru Zhao, Meng Qu Zewen Chi for their insightful discussion on language model pre-training. We also thank the anonymous reviewers, AC, SAC, and PC for their time, thoughtful feedback, and constructive engagement throughout the review process. We also acknowledge funding from the Canada CIFAR Al Chair Program and the Intel-Mila partnership program. The computation resource this project is supported by Mila and the Digital Research Alliance of Canada

---

References
[1] Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. In The Twelveth International Conference on Learning Representations, 2024.
[2] Jimmy Ba, Geoffrey E. Hinton, Volodymyr Mnih, Joel Z. Leibo, and Catalin lonescu. Using fast weights to attend to the recent past. In The Twenty-ninth Annual Conference on Neural Information Processing Systems, 2016.
[3] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang. Zhengxiao Du, Xiao Liu, Aohan Zeng. Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics, 2024.
[4] Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv: 2405.21060, 2024.
[5] Aditya Desai, Shuo Yang, Alejandro Cuadron, Matei Zaharia, Joseph E. Gonzalez, and Ion Stoica. Hashattention: Semantic sparsity for faster inference. In The Forty-second International Conference on Machine Learning, 2025.
[6] Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang. Shaohan Huang, Wenhui Wang, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens, 2023.
[7] Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin. Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. Enhancing chat language models by scaling high figh-quality instructional conversations, 2023.
[8] Ning Ding, Yehui Tang, Haochen Qin, Zhenli Zhou, Chao Xu, Lin Li, Kai Han, Liao Heng, and Yunhe Wang. Memoryformer. Minimize transformer computation by removing fully-connecter layers. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[9] Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, ZUJIA CHEN, Ameya Sunil Mahabalesh- warkar, Shih- Liu, Matthijs Van keirsbilck, Min-Hung Chen, Yoshi Suhara, Yingyan Celine Lin, Jan Kautz, and Pavlo Molchanov. Hymba: A hybrid-head architecture for small language models. In The Thirteenth International Conference on Learning Representations, 2025.
[10] Daniel Y. Fu, Tri Dao, Khaled Kamal Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations, 2023.
[11] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace Hc, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling, 2021.
[12] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 2024.
[13] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ah- mad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi. Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, et al. The llama 3 herd models, 2024.
[14) Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In The First Conference on Language Modeling, 2024.
12

---

[15] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher RRé. Hippo: Recurrent memory with optimal polynomial projections. In The Thirty-fourth Anmual Conference on Neural Information Processing Systems, 2020.
[16] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The Tenth International Conference on Learning Representations, 2022.
[17] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christo- pher RChristopher RRé. Combining recurrent, convolutional, and continuous-time models with linear state-space layers. In The Thirty-fitth Annual Conference on Neural Information Processing Systems, 2021.
[18] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? In The First Conference on Language Modeling, 2024.
[19] Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and Eran Malach. Repeat after me: Transformers are better than state space models copying. In The Forty-first International Conference on Machine Learning, 2024.
[20]
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The Eighth International Conference on Learning Representations, 2020.
[21] Barak Lenz, Opher Lieber, Alan Arazi, Amir Bergman, Avshalom Manevich, Barak Peleg, Ben Aviram, Chen Almagor, Clara Fridman, Dan Padnos, Daniel Gissin, Daniel Jannai, Dor Muhlgay, Dor Zimberg, Edden M. Gerber, Elad Dolev, Eran Krakovsky, Erez Safahi, Erez Schwartz, Gal Cohen, Gal Shachaf, Haim Rozenblum, Hofit Bata, Ido Blass, Inbal Magar, Itay Dalmedigos, Jhonathan Osin, Julie Fadlon, Maria Rozman, Matan Danos, Michael Gokhman, Mor Zusman, Naama Gidron, Nir Ratner, Noam Gat, Noam Rozen, et al. Jamba: Hybrid transformier-mmbaa language models. In The Thirteenth International Conference on Learning Representations, 2025.
[22] Jiaheng Liu, Dawei Zhu, Zhiqi Bai, Yancheng He, Huanxuan Liao, Haoran Que, Zekun Wang, Chenchen Zhang, Ge Zhang, Jiebin Zhang, Yuanxing Zhang, Zhuo Chen, Hangyu Guo, Shilong Li, Ziqiang Liu, Yong Shan, Yifan Song, Jiayi Tian, Wenhao Wu, Zhejian Zhou, Ruijie Zhu, Junlan Feng, Yang Gao, Shizhu He, Zhoujun Li, Tianyu Liu, Fanyu Meng, Wenbo Su, Yingshui Tan, Zili Wang. Jian Yang, Wei Ye, Bo Zheng, Wangchunshu Zhou, Wenhao Huang, Sujian Li, and Zhaoxiang Zhang. A comprehensive survey on long context language modeling, 2025.
[23] Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 2024.
[24] Justus Mattern and Konstantin Hohr. Manba-chat. GitHub, 2023.
[25] William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. In The Forty-first International Conference on Machine Learning, 2024.
[26] Elvis Nunez, Luca Zancato, Benjamin Bowman, Aditya Golatkar, Wei Xia, and Stefano Soatto. Expansion span: Combining fading memory and retrieval in hybrid state space models, 2024.
[27] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodee Danny Hernandez. Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022.
(28) Matteo Pagliardini, Daniele Paliotta, Martin Jaggi, and François Fleuret. Faster causal attention over large sequences through sparse flash attention. In The Thirty-seventh Annual Conference on Neural Information Processing Systems, 2023.
[29] Badri Narayana Patro and Vijay Srinivas Agneeswaran. Mamba-360: Survey of state space models as transformer alternative for long sequence modelling: Methods, applications, and challenges, 2024.
13

---

[30] Saurav Pawar, S. M. Towhidul Islam Tonmoy, S M Mehedi Zaman, Vinija Jain, Aman Chadha, and Amitava Das. The what, why, and how of context length extension techniques large language models a detailed survey, 2024.
[31] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao. Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In The Fortieth International Conference on Machine Learning, 2023.
[32] 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.
[33] Liliang Ren, Yang Liu, Shuohang Wang. Yichong Xu, Chenguang Zhu, and Cheng Xiang Zhai. Sparse modular activation for efficient sequence modeling. In The Thirty Annual Conference on Neural Information Processing Systems, 2023.
[34] Yash Sarrof, Yana Veitsman, and Michael Hahn. The expressive capacity of state space models:
A formal language perspective. In The eighth Annual Conference on Neural Information Processing Systems, 2024.
[35] Zhiqing Sun, Yiming Yang, and Shinjae Yoo. Sparse attention with learning to hash. In The Tenth International Conference on Learning Representations, 2022.
[36] Liping Tang, Nikhil Ranjan, Omkar Pangarkar, Xuezhi Liang, Zhen Wang, Li An, Bhaskar Rao, Linghao Jin, Huijuan Wang, Zhoujun Cheng, Suqi Sun, Cun Mu, Victor Miller, Xuezhe Ma, Yue Peng, Zhengzhong Liu, and Eric Xing. Txt360: A top-quality Ilm pre-training dataset requires the perfect blend, 2024.
[37] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, 2022.
[38] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In The Thirtieth Annual Conference on Neural Information Processing Systems, 2017.
[39] Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in a transformer language model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, 2019.
[40] 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 Shoeybi, and Bryan Catanzaro. An empirical study of mamba-base language models, 2024.
[41] Wikipedia. Harry Potter. https: www.eeewwwiieddiarreeee.ee httppppddaa.oo///iiiiaaaaaraaaoa 2025.
[42] Songlin Yang and Yu Zhang. Fla: A triton-based library for hardware-efficient implementations of linear attention mechanism, 2024.
[43] Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo. Liang Zhao, Zhengyan Zhang, Zhenda Xie, Yuxing Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wen- feng Liang, and Wangding Zeng. Native sparse attention: Hardware-aligned and natively trainable sparse attention. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics, 2025.
[44] Luca Zancato, Arjun Seshadri, Yonatan Dukler, Aditya Golatkar, Yantao Shen, Benjamin Bowman, Matthew Trager, Alessandro Achille, and Stefano Soatto. B'MOJO: Hybrid state space realizations foundation models with eidetic and fading memory. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
14

---

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: We confirm that the abstract and introduction accurately reflect the paper's contributions and scope.
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 in 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 not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Yes, the limitations are discussed in the Technical Appendix A.
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-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
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 on 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
15

---

Justification:
We provide assumptions and proof in Technical Appendix B and C.
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 in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide the experimental settings Technical Appendix D. The code and data to reproduce the results can be found at:
https://githuu com CoomDDeeGrraae whpepprrapheeannnngeg... HAX
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 are 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 is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to 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) the contribution is primarily a new algorithm, the paper should make 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 (e-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 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 (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
16

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We provide the experimental settings in Technical Appendix D. The code and data reproduce the results can be found at:
httpp://gitthn com/ Seeem/Depprraaa http:/Derrraaeeaannngggeeeeeeeee
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/
cublic/guides. /ttp///eeeaaaaiiiiooooooiiieeiee
for more details.
While we encourage the release of code and data, we understand that this might 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 //ippscc/pppp wwcppuulcc//aeeecccccecccecceeee htaliicgguies//// httpc//guuddde/dooddubiiisssssss
for more details.
The authors should provide instructions on 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 much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide these information in Technical Appendix D. Full details can be found in the code.
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 a 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
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Main results (Ruler and Longbench) are obtained from the well established benchmarks.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoo bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
17

---

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).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CL, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that 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 or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide these information Technical Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer:
[Yes]
Justification: Yes, carefully respect the NeurIPS Code of Ethics.
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 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:
The paper would not raise any potential positive or negative social impacts.
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.
18

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 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 unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper would not pose any risk for misuse.
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 not 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: [NA]
Justification:
The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does 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 used and, possible, include URL.
The name of the license (e.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.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
19

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if 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 new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include 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 as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not 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.
According the NeurIPS Code of Bthies, 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/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: All results in-silico and do not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on 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 NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
20

---

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, originality of the research, declaration is not required.
Answer: [NA]
Justification: LLM is used only for writing and editing purpose during the research.
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.LM policy (https: for what should or should not be described.
21

---

C. Kuy
Kig
Information
Kg
Kag
Auto-Regressive Model
Input-
C
Kag
Via
Ky
Va
Vs
V2
Inquiry
Kag
Kay
Output
Kig
Figure 5: Input and output components of the auto-regressive multi-query joint recall task. The input sequence is further divided into an information component and an inquiry component.
A Relationship Between the Argmax and Sign-Bit LSH Binning Rules
In this section, show how the sign-bi bit LSH binning rule (Eq. 7 can be interpreted argmax LSH bimning rule (Eq. (6) applied to an expanded projection matrix with 2 columns. We first construct the expanded matrix, and then prove the equivalence.
Expanding the projection matrix. Let the original random projection be H HH RdxA H,
H.]
Lid NN00, 1). Define codebook of 2 signed prototypes
B = Bb = Σラ H (b.....,bby bn) E f-11, +11%) c Rd.
(21)
Stacking all B, as columns yields the implicit matrix H E Rdxx"
Equivalence of the two binning rules. For a normalized query vector Q we define its sign projection as = sign(Q) Q H) E +11++13 The inner product Q; and prototype B,
QQ. Bb)
Σ by QQ H,).
(22)
Because every term with bj * sy flips the sign of the positive quantity [Q,, H2)|, we have the strict inequality (Q., B&) (Q,, Bb) for all b * s. Hence
brgmax (Q., Bb) = -bin bin
(sign)
(23)
The sign-bit assignment is exactly the argmax rule applied to H. An identical argument holds for keys Ky. Thus, the sign-bit method equals the argmax method with 2 (expanded) columns.
B Theoretical Proof
Multi-query joint recall requires models to recall an ne x ny table of context-specific key-value associations, in which ni keys are associated with different values in each of the n. contexts, with n. x being the total number of entries. For clarity, we introduce some additional notations for multi-query joint recall in the auto-regressive setting, as illustrated Appendix Figure 3 The input sequence is divided into an information component and an inquiry component. The information component provides the context-specific key-value associations. The inquiry component permutes the order of context and keys in the information component, and the model is required to predict the corresponding values given each key under every specified context.
B.1 Proof of Corollary 3.1
Corollary 3.1 (Limited expressiveness of SSMs). Under the uniform multi-query joint recall distri- bution, for any n, generalized state-space model defined in Sec. 2.3 incurs an error rate of at least
M F
1
22

---

Information
Inquiry
SSM
Monsay
Allaz
physics
English
Allix
Chinoe
Monday
Alizz
Hybrid Dilated Attention
Miooday
Africe
moth
Bob
piyesccs
Thenday
English
Alloe
Chinose
Monday
Alice
Hybrid LSH Attention
Monslay
Alias
physis
Allox
Monsly
Alizz
Figure 6: Comparison between SSM, hybrid dilated attention model and hybrid locality- hashing (LSH) attention model on joint recall. By selectively bypassing irrelevant context, sparse attention alleviates memory overload in SSM layers and enhances the hybrid model's capability to retrieve relevant information.
The intuition behind the proof of Corollary 3.1 data instances that a model can accurately represent is fundamentally limited by the capacity of its recurrent state. Since all information from the input sequence must be encoded into fixed recurrent state following the context during processing, the total number of distinguishable output is bounded by the representation capacity of the state space [2]]. Consequently, if the size of the output space [V["" exceeds (41, the model inevitably incurs non-negligible error.
straightforward: the number of possible joint recall
As direct consequence of Theorem 2.7 in [9] adopt their proof strategy. We reformulate Lemma D.1 in [19] as the following Lemma B.I Let m denote the index of the last token in the information component. Then, for any fixed permutation P of the context and keys in the inquiry component, the following Lemma B.1 holds:
Lemma B.1. Let M be a fixed-state generalized SSM that maps the joint recall input space X' the output space V under any fixed permutation P of the context and keys in the inquiry component. Then there exists a function
+ V such that for all inputs x E X, the model output satisfies M(x) = G(Um(x)) Um is defined in Eq 10
Following [9]. bound the error of the model by comparing the number of possible model states to the number of distinct input instances.
Proof.
1 Prferr] Pr[M(x)= =ly y V"]
(24)
1 Σ 1[((x= =
(25)
yEv-
1
Σ Σ
I[G(u) y) ( ul
(26)
|V/"" yEv" ueld
1
Σ
1[Um(x) =
(27)
1 uEW
(28)
To guarantee Pr[err] 0, it is necessary that the number of representable states satisfies [2|] > |
Taking the logarithm of both sides yields the condition b 2 n log [V]. This implies that the state-space dimension of the model must grow linearly with the number of entries in the joint recall table, highlighting a fundamental limitation of the representation capacity of SSMs.
In contrast, as illustrated in Figure 6 hybrid sparse attention models mitigate this limitation by enabling information to propagate through multiple parallel paths, thereby alleviating the bottleneck imposed by sequential state updates.
23

---

B.2 Proof of Proposition 3.2
Proposition 3.2 (Improved expressiveness of SSMs integrated with CDSA). There exists a 2-layer auto-regressive hybrid model consisting of SSM layer followed by LSH atlention layer, which can solve multi-query joint recall in Onnlog" n) time complexity with (log n) SSM state dimensions.
Proof. We prove by construction. In the first layer, expect the SSM state concatenates each value token representation with its associated key token representation and context token representation. To be specific, we expect the SSM state representation at each value token to be
[c, k, v, is_ul Eu
where с is the representation of the current associated context token, k is the representation of the current associated key token, and v is the representation of the nearest value token. is_ U sign indicator (-1 or 1) that specifies whether the current token is value token.
To achieve this, we first construct each vector c, k and be distinct b-dimensional vector with unit norm without zero entries, i.e., ||c||2 = 1, ||k||| = 1, ||v||2 = 1, Vj, ej # 0, o O, v of 0. Since the number of distinct vectors that can be drawn from the unit sphere grows exponentially with dimensionality, Ollog n) embedding dimensions are sufficient to ensure that all representations are distinguishable. Then we define an embedding space in which context and value tokens are mapped to structured representations. Specifically, a context token is embedded as
|c, 0,
where с is the constructed representation of this context token on the unit sphere, and the final coordinate is set -1 to indicate that the current token is not a value. Similarly, key token is embedded as
(0, k. 0, -1]
and a value token is embedded as
[0, v.
where k and v are the key and value representations from the unit sphere, respectively, and the final coordinate is set only when the current token is value token. Following Eq. [0 we define the update rule of the generalized SSM as follows:
U, u(U,-,, U, o 1 0] + o lfe, +
(29)
R; = -(,
(30)
where is the current input embedding and ey refers to its j-th dimension. The update rule operates as a conditional overwrite: if a position does not carry information (i.e., the corresponding dimension in is 0), the previous state is preserved; otherwise, it updated with the current embedding. Following this update rule, the SSM state at each value token in the information component takes the form
[c, k, v]
while the SSM state at each key token takes the form
(c, k, ?. -1]
In the second layer, LSH attention operates on the SSM state (c. k, v,is_r] E U4 using (c, k, 0, is as the LSH attention key representation, [c, k, 0, 1] as the LSH attention query representation, and (0, 0, v, 1] as the LSS attention value representation. This design ensures that value tokens in the information component and key tokens in the inquiry component that share the same context and key (i.c., identical e and representations in the SSM state) will always be assigned to the same hash bin. With a sufficient number of, e.g. (n) hash bins, which can be efficiently constructed using sign-bit binning rule with Oloo n) random projection vectors, values associated with each key in every specified context are reliably retrievable by LSH attention. This step is the bottleneck of computation with (n log n) time complexity.
24

---

1st-level
2nd-level context
3rd-level
4th-level context
value
context
context
Alice studies math. Bob studies physics.
mornings,
on Monday
afternoons.
the Spring Semester,
on Tuesday on Wednesday
Monday
In the Fall Semester,
Figure 7: An example multi-level context natural language.
B.3 Proof of Proposition 3.3
Proposition 3.3 (Limited expressiveness of SSMs integrated with CISA). There does not exist 2-layer auto-regressive hybrid model consisting of an SSM layer followed by a CISA layer, which can solve multi-query joint recall with o(n²) time complexity, since requires at least ( SSM state dimensions, k is the maximum number of keys that each query allowed to attend to in the sparse attention module, as defined in Eq.
Proof. Consider a key given in the inquiry component of the auto-regressive joint recall task. The model is required to output the associated value token when this key token is provided as input. Taking the query representation from this key token, the sparse attention can attend to at most k key representations from previous tokens, where the key representations are calculated based on the SSM state representations of the first layer. To solve the joint recall task, these k key representations being attended must collectively encode the full information component. Since the full information component length is O(n), by Corollary 3.1 & state representations of the generalized SSM in the first layer must use at least ( dimensions to collectively store the information component. Thus, the per-key computational cost required by the second-layer sparse attention is O(k- _t ((n).
and therefore the total time complexity (n??).
Comparing Proposition 3.2 with Proposition 3.3 we see clear representation capacity gap between the SSMs integrated with CDSA and the SSMs integrated with CISA. In practice, however, with an appropriate constant k integrating CISA with SSMs still provides an advantage: unlike SSMs, which only have access the last state representation, CISA layers attend different state representations simultaneously, at cost times of computation budget.
C Extending Joint Recall to Multi-level Context
As illustrated in Figure 7 in many cases, natural language contexts exhibit hierarchical dependencies. This motivates us to extend joint recall the multi-level context setting, in which we regard the keys the last level of context.
C.1 Formulation
Given different levels of context vocabulary C1, C2, Cn, and token vocabulary V, multi-level context joint recall requires a model recover the mapping C1 C2 ... x C.p + The context of multi-level context joint recall is hierarchically structured analogously to natural languages. It is
25

---

divided into [[] first-level blocks, where each first-level block begins with a token from the first-level context vocabulary C₁. Each first-level block is further divided into sub-blocks beginning with a second-level context token from C2, and this recursive sub-division continues up to the w-th level. The last-level block consists of w-th level context token followed by value token from V. Note that associative recall is a special case of multi-level joint recall with w = 1, and joint recall is a special case of multi-level joint recall with w = 2. We similarly define multi-query multi-level context joint recall, where the model is required to recall all n [C] x |C2| x |Cwl entries in the full context-value supertable.
C.2 Expressiveness of SSMs Integrated with CDSA on Multi-Level Context Joint Recall
On multi-query multi-level context joint recall, both Corollary 3.1 and Proposition 3.3 continue to
hold under the same assumptions. We now extend Proposition 3.2 to the following Proposition C.T which demonstrates that SSMs integrated with CDSA remain expressive even in the presence of w levels of hierarchical contexts.
Proposition C.1 (Expressiveness of SSMs integrated with CDSA on multi-level context joint recall). There exists a 2-layer auto-regressive hybrid model consisting of an SSM followedd by an LSH attention layer, which can solve multi-query multi-level context joint recall in O(wn log? n) time complexity with O(w log n.) SSM state dimensions.
Proof. Similar to Proposition 3.2 we prove by construction. In the first layer, we hope the SSM state to consist of the context and value representations
[2 ...,..... Zwo, v, vj E EU
where z; denotes the representation of the nearest i-th level context token, v represents the nearest value token, and is_ e is a sign indicator (-1 or 1) that specifies if the current token is a value token.
To achieve this, we similarly construct each vector z, and v to be a distinct b-dimensional vector with unit norm without zero entries, i.e., | 1, | 22 1, Vj. 2i 0, V * O. Since the number of distinct vectors that can be drawn from the unit sphere grows exponentially with dimensionality, Ollog n) embedding dimensions are sufficient ensure that all representations are distinguishable. Consequently, the total size of the SSM state is O(w log n).
Then we define an embedding space in which context and value tokens are mapped to structured representations. Specifically, a i-th level context token is embedded as
[0, ..., 0, Zi, 0, ..., -1]
where z; is the context token representation, and the final coordinate is set to -1 to indicate that the token not a value. Similarly, a value token embedded as
e e= 00,0, ... 0, 1]
0,,0,....
where is the value token representation and the final coordinate is set to 1 to mark it as value token. We keep the update rule of the generalized SSM as in Eq. 29
U u(U(-1, U e 1 = + 1le, + 0] R₂ r(U))= = U₁
(31)
(32)
which operates as a conditional overwrite, Following this update rule, the SSM state at each value token takes the form
Zz. Z2. Zw v,
where 21, Z2, .... zw are the w levels of context representations the current token, v is the value representation, and the final dimension is set to to indicate that the current token is value token.
In the second layer, LSH attention operates on the SSM state [zz, Z2, .... Zwo V. _s__), [21, 2..... Zu. O, iss_0 v] as the LSH attention key representation, [Z1, Z2, .... Zw 0. 1] as the LSH attention query representation, and [0, 0, ... v, 1] as the LSH attention value representation. This design ensures that tokens that share the same context (i.e., identical |2, 22,2,... 22, ..., 2nd] in the SSM state)
v). using
will always be assigned to the same hash bin. With a sufficient number of, e.g. O(n) hash bins, which can be efficiently constructed using sign-bit bit binning rule with Ollog n) random projection
26

---

vectors, value representations associated with identical key combinations in the context are reliably retrievable by LSH attention. This step is the bottleneck of computation with Oww log" n) time complexity.
This construction establishes that a 2-layer hybrid model consisting of generalized SSM followed by LSH attention can solve multi-query multi-level joint recall efficiently, with sub-quadratic time complexity and sub-linear state complexity with respect to the input sequence length.
D Experiment Details
D.1 Empirical Verification on Joint Recall
For all models, we fix the number of layers to 2, set the hidden size to 64, and use A: = 64. For variants that integrate multiple sparse attention components, namely SW+D, A (consists of a SW component and sink attention component, a sink attention always attends to the first A: tokens in the sequence only), and HAX (LSH+KS), we allocate l 32 each component, in order to maintain global k = 64. For both LSH and HAX (LSH+KS), we adopt the sign-bit binning strategy (Eq, 7 with h 8, and refresh the random hashing matrix at each training step. To account for the absence of sparse attention mechanisms, we additionally evaluate Mamba and Mamba2 baselines with their hidden size doubled to 128, ensuring fair comparison. Additionally, we include Samba baseline, consisting of 2 Mamba layers and 2 sliding window attention layers. For Samba, the hidden size and sliding window width (k) are both set to 64. We use A optimizer. All models are trained for 400,000 steps with batch size of 64. Our implementation is based on Beaah-LLiner-Attennie [42
D.2 Continual Pre-training on Natural Language
For all experiments, we fix the sparsity parameter 128. For variants that integrate multiple sparse attention components, namely SW+D, A, and HAX (LSH+KS), we set k = 64 to each component, in order to maintain global k 128. For both LSH and HAX (LSH+KS), we adopt the argmax binning strategy (Eq. (6) with h = k. We resample the random hashing matrix at each training step and fix a random hashing matrix during evaluation.
To enhance the long-context modeling capability, we filter samples to retain only those with tokenized lengths of least 2,048 tokens. At the beginning of continual pre-training, the K and Q projection weights are initialized using the parameters of the B and C projections from Mamba, respectively, based on state-space duality 4]. During continual pretraining, we use cosine schedule with maximum learning rate of 3e-4 for 130M models and 1.5c-4 for 790M models, and a minimum learning rate of le-5. A warm-up phase of 200 steps with learning rate of followed by 800 steps of linear warm-up precedes the cosine schedule. For instruction tuning, we also apply 200-step O learning rate phase followed by 800 steps of linear warm-up, up, after which the learning rate remains constant at 3e-6. For both continual pre-trat pre-training and instruction-tume -tniing we use Adam W optimizer and train with a global batch size of 64 and context length 2K.
E
Additional Experiments
E.1
Short-context Modeling Benchmark
We follow Mamba [4] to evaluate the zero-shot short context modeling performance of the con- tinually pre-trained models on the LM evaluation harness benchmark from EleutherAl [12). Our results in T show that continual pre- trainini on long sequences will not lead to a significant performance drop short context benchmarks, where the Mamba w/o continual pre- pre-training results are copied from the Mamba paper [4].
27

---

Table 5: LM evaluation harness benchmark for continually pre-trained Mamba models. We compare different sparse attention integrated with Mamba, including CISA and CDSA methods as in Table
LamhdaPPI.
WinoGrande
PIQA
LambdhAce
HellaSwag
ARC-E
ARC-C
AverageAce
Mamba 15.58
51.8
63.8
44.3
35,3
47.8
24.4
44,6
15.75
52.8
64.4
43.9
47.7
24.1
44.7
+SW
15.47
52.8
63.8
44.3
35.2
47.9
24.7
44.8
+SW+D
15.35
53.0
64.0
44.7
35.2
47.6
24.4
44,8
15.65
53.0
64.0
44.2
35.3
47.7
24.7
44,8
+LSII
15.73
52.6
64.3
43,8
35.3
47.7
24.0
44,6
(ours)
+HAX (owrs)
15.62
52.5
63.9
44.3
35,3
47.9
24.5
44.7
w/o traninng
16.07
51.9
64.5
44.3
35.3
48.0
24.3
44.7
Sliding Window Attention
A-shaped Attention
Dilated Attention
Socality-Sensiive Hashing Attention
Figure 8: Examples for input-dependent and input-independe sparse attention patterns.
25
d Attention Probability
Query Token Index
75
100
25
50
75
100
Key Token Index
Figure 9: An example of the "vertical-stripe" attention pattern in LLM. We input the first paragraph of the Wikipedia term "Harry Potter" [41] into the Llama 3.2 IB model [3 and visualize the log attention probabilities of the last head in the first layer. The input text is: "Harry Potter is a series of seven fantasy novels written by British author J. K. Rowling. The novels chronicle the lives a young wizard, Harry Potter, and his friends, Ron Weasley and Hermione Granger, all of whom are students Hogwarts School of Witcheraft and Wizardry. The main story are concerns Harry's conflict with Lord Voldemort, dark wizard who intends to become immortal, overthrow the wizard governing body known as the Ministry of Magic, and subjugate all wizards and Muggles (non-magica people)."
28


Reviews:
[{"review_id": "PxXJovvNKi", "reviewer": "Reviewer_MdxV", "summary": "This paper addresses the limitations of state-space models (SSMs) in long-context natural language processing (NLP) by introducing a novel hybrid sparse attention mechanism. The authors propose a new synthetic task, joint recall, to evaluate the model's ability to integrate multiple contextual cues. Through theoretical analysis and empirical evaluation, they demonstrate that their proposed hybrid sparse attention mechanism, combining locality-sensitive hashing (LSH) and key-centered (KC) attention, significantly improves the long-context modeling capacity of SSMs. The paper also validates these findings on real-world NLP tasks, showing enhanced performance on benchmarks such as Ruler and LongBench.", "questions": "Apart from the weaknesses mentioned above, there might be some questions that would benefit from this manuscript.\n\n* (Q1) **Broader Impacts.** Could the authors provide a discussion on the potential societal impacts of their work, including any negative consequences and mitigation strategies?\n\n* (Q2) **Statistical Significance.** How do the authors plan to address the lack of error bars or other measures of statistical significance in their experimental results?", "limitations": "The paper is technically sound, with a clear and innovative approach to addressing long-context limitations in SSMs. The theoretical and empirical evaluations are comprehensive, supporting the claims made. The authors have acknowledged some limitations in the appendix, but a more explicit discussion in the main paper would be beneficial. Additionally, the paper would benefit from a clearer discussion on the potential societal impacts and future directions of their work.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n\n* **(S1) Innovative Approach**: The paper introduces a novel hybrid sparse attention mechanism that addresses the limitations of SSMs in \ncapturing long-range dependencies, demonstrating both theoretical and empirical improvements.\n\n* **(S2) Synthetic Task**: The introduction of the joint recall task provides a unique framework for evaluating the model's ability to integrate multiple contextual cues, enhancing the interpretability of the results.\n\n* **(S3) Empirical Validation**: The paper provides comprehensive empirical evidence supporting the effectiveness of the proposed method on both synthetic and real-world datasets, showcasing its practical utility.\n\n* **(S4) Technical Rigor**: The theoretical analysis is thorough, with clear assumptions and proofs, ensuring the soundness of the proposed approach.\n\n### Weaknesses\n\n* **(W1) Limited Discussion on Broader Impacts**: The paper does not discuss potential societal impacts, which is a missed opportunity to address broader ethical considerations.\n\n* **(W2) Statistical Significance**: Firstly, the paper lacks reporting of error bars or other measures of statistical significance, which could strengthen the credibility of the experimental results. Secondly, expand to explicitly contrast joint recall with positional-stress protocols (e.g., ‘lost-in-the-middle’) and multi-key associative tasks.\n\n* **(W3) Reproducibility**: While the paper mentions the availability of code and data, it does not provide explicit instructions or URLs, potentially hindering full reproducibility.\n\n* **(W4) Missing hardware-aligned sparse baselines:** More comparisons to relevant previous works and reporting of throughput/latency and memory utilization to substantiate “sub-quadratic” *in practice* would enhance the manuscript with more convincing evidence of the proposed method.", "was_revised": true, "final_justification": "After considering the rebuttal from the authors and review comments from other reviews, the authors have tackled my concerns, and I decided to keep my original score.", "rebuttal": "We sincerely thank you for your thoughtful evaluation of our work. We greatly appreciate the positive feedback regarding the novelty of our hybrid sparse attention, the clarity of our theoretical analysis, and the solidity of the empirical results. We are also glad that the introduction of the joint recall task was seen as a meaningful contribution to theoretical analysis and empirical benchmark for long-context modeling. Below, we address your specific questions and concerns in detail.\n\n---\n\n**W1/Q1: On the potential societal impacts of our work**\n\nThank you for raising this important point regarding broader societal impacts.\n\nOur work focuses on fundamental architectural improvements for long-context sequence modeling, particularly by enhancing the theoretical and empirical understanding of hybrid sparse attention in state-space models. The goal is to advance the efficiency and expressiveness of machine learning models at the architectural level, a domain that is general-purpose and infrastructure-oriented rather than application-specific.\n\nAs such, our contributions do not involve direct interaction with human users, personal data, or deployment in potentially risky domains (e.g., medical, legal, or surveillance applications). Consequently, we believe our work does not introduce significant ethical risks or societal concerns in its current scope.\n\nMeanwhile, we agree that architectural advances may eventually be used in downstream applications. Therefore, we fully support the community’s efforts to consider responsible deployment and fairness at later stages of the ML pipeline. We have updated the Conclusion section to explicitly clarify the scope and neutrality of our contributions with respect to societal impacts.\n\n---\n\n**W2/Q2: Statistical significance report**\n\nWe appreciate your suggestion regarding the inclusion of statistical significance measures.\n\nFor the Ruler [14] and LongBench [3] benchmarks, we follow the established evaluation protocol used in prior work, which typically reports single-run results without error bars. We adopted the same reporting standard to ensure comparability and consistency.\n\nFor our proposed joint recall task, we conducted experiments using three random training seeds and summarize the average results with standard deviations in the table below:\n\n||Mamba|Mamba2|\n|-|-|-|\n|Base|16.3 $\\pm\\$ 0.7|36.6 $\\pm\\$ 2.7|\n|+D|7.8 $\\pm\\$ 0.8|19.6 $\\pm\\$ 16.3|\n|+SW|18.7 $\\pm\\$ 4.5|70.6 $\\pm\\$ 46.8|\n|+SW+D|16.6 $\\pm\\$ 1.2|48.6 $\\pm\\$ 46.2|\n|+SW+SK|16.4 $\\pm\\$ 0.1|49.3 $\\pm\\$ 42.7|\n|+LSH|11.6 $\\pm\\$ 2.5|13.5 $\\pm\\$ 12.6|\n|+LSH+SW|16.5 $\\pm\\$ 0.2|50.7 $\\pm\\$ 43.2|\n|+KC|36.6 $\\pm\\$ 35.5|60.1 $\\pm\\$ 34.1|\n|+LSH+KC|38.0 $\\pm\\$ 37.3|74.3 $\\pm\\$ 35.0|\n|+FullAttn|48.4 $\\pm\\$ 45.2|100.0 $\\pm\\$ 0.0|\n|Samba|6.3 $\\pm\\$ 0.1|/|\n\nTwo new baselines are added to this table: hybrid LSH + sliding window (SW) attention and hybrid full attention. As observed, the variance in performance can be relatively high. This behavior is inherent to the nature of the synthetic task: the model either learns the underlying rule and achieves near-perfect accuracy, or fails to do so and performs poorly when being stuck at a local minimum. This leads to a bimodal outcome distribution, which inflates the standard deviation.\n\n---\n\n**W3: Data availability and code reproducibility**\n\nWe have provided the link to the full codebase at multiple locations in the paper (Lines 248, 479, and 518), and we will further highlight these links in the final version for better visibility.\n\nBoth the joint recall synthetic task and the continual pre-training experiments are fully reproducible by running the scripts in our repository. For the continual pre-training experiments in the paper, we use only publicly available datasets:\n\n- Continual pre-training follows Mamba [10] using The Pile dataset.\n- Instruction tuning follows Mamba-Chat [20] using UltraChat.\n- Evaluation is performed on Ruler [14] and LongBench [3], both of which are publicly accessible.\n\n\nWe provide a detailed description of training setups in Appendix D. We hope this clarifies the reproducibility of our experiments.\n\n---\n\nOnce again, we sincerely thank you for the helpful comments and suggestions. We sincerely hope that our rebuttal resolves the concern you highlighted, and we are grateful for your time and insights."}, {"review_id": "Rv353VAo2P", "reviewer": "Reviewer_YZT8", "summary": "This paper addresses the limitations of State-Space Models (SSMs) in capturing complex long-range dependencies that require integrating multiple contextual cues. The authors introduce a new synthetic task, \"joint recall,\" to demonstrate this weakness and propose a hybrid sparse attention mechanism combining Locality-Sensitive Hashing (LSH) and a Key-Centered (KC) attention. Empirical results on both synthetic and real-world NLP benchmarks show that augmenting Mamba with this hybrid attention significantly improves its long-context modeling performance.", "questions": "1. The QA2 task results are missing for RULER benchmark. Could you add them?\n\n2. Could you add a baseline that use dense attention instead of sparse attention for hybridization?\n\n3. Which dataset did you use for continual pretraining?\n\n4. How sensitive are the continual training results on the scalers of the auxilary loss in KC attention?\n\n5. How many parameters doe the MLP have for key scoring?\n\n6. Could you also add a setting of SW+LSH, so that we can better understand the effectiveness of KC?\n\n7. How many test samples are in the joint recall synthetic dataset? Could you include the standard deviation for the evaluation on synthetic tasks?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The proposed LSH+KC attention is thoughtfully designed. It combines the content-based clustering of LSH with a key selection mechanism (KC attention) specifically created to handle \"global hub\" tokens (like instructions), addressing a known weakness of simple LSH.\n\n2. The idea of using both key representation and historical query representaitons for key importance scoring is novel. And using ranking loss on downsampled number of keys to learn the scoring module drastically reduces the training computation costs, and is well motivated.\n\nWeaknesses:\n\n1. The paper seems rushly written. Experiment results are not complete, and the experiment setup and evaluation settings are missing. Technical details on KC attention are not clear. The results on LongBench seem not significant.\n\n2. The proposed KC attention has similar performance as sliding window attention on RULER. It is unclear how effective KC it is in real-world scenario.\n\n3. The scale of the experiments are small. The model trained on synthetic data are only two layers and the 130M sized Mamba model are continually pretrained for the results on real world data.\n\n4. Missing Citations. The authors should discuss with previous works on dynamic sparse attention, especially those also select QK representations [1,2], to better position their works.\n\n\n---\n\n[1] Faster Causal Attention Over Large Sequences Through Sparse Flash Attention\n\n[2] Sparse Modular Activation for Efficient Sequence Modeling", "was_revised": true, "final_justification": "The added results mostly resolved my concerns.", "rebuttal": "We sincerely thank you for the thoughtful and encouraging feedback. We are glad that you found the design of our LSH+KC attention mechanism to be novel and well-motivated.\n\nBelow, we address your specific questions and concerns in detail.\n\n---\n\n**W3: The scale of the experiments are small.**\n\nWe have newly conducted continual pre-training based on the 790M Mamba official checkpoint. These larger-scale experiments confirm the effectiveness of our proposed LSH+KC method on real-world long-context benchmarks. For these experiments, we use the TxT-360 corpus (not the same as The Pile used in our paper) as the continual pre-training dataset. We continual pre-train for 50,000 steps with a context length of 2K, which is followed with instruction tuning using UltraChat (the same as in the paper) for 10,000 steps with a context length of 2K. We also train new 130M models in this setting. The new results on Ruler and LongBench are presented below:\n\n**Table R.1: 4K Ruler results based on 790M Mamba**\n\n|Model|NIAHS1|NIAHS2|NIAHS3|NIAHMK1|NIAHMK2|NIAHMK3|NIAHMV|NIAHMQ|VT|CWE|FEW|QA1|QA2|Average|\n|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|\n|Mamba|100|19.2|9.2|5.2|0.6|0|0.65|1.85|5.72|4.3|100|26.2|78.6|27.0|\n|+D|0|0|0|0|0|0|0|0|65.2|0|100|0|73|18.3|\n|+SW|100|18.2|5|5.6|0.2|0.2|1.75|2.1|4.4|4.78|100|22.4|78.8|26.4|\n|+SW+D|100|20.2|8.2|6.2|0.4|0.2|1.45|1.25|4.36|4.32|100|24.6|79|26.9|\n|+SW+SK|100|19.8|7.2|6.6|1.2|0.2|1.3|0.8|4.36|3.24|100|24.2|79|26.8|\n|+LSH|0|0|0|0|0|0|0|0|0|0.04|100|0.2|73|13.3|\n|+KC|99.8|20.2|7.8|8.2|1|0.2|1.2|1.85|9.32|2.86|100|22.6|78.6|27.2|\n|+LSH+KC|100|23|11|9|0.8|0|3.4|2.75|6.36|2.54|100|29|78.8|**28.2**|\n\n**Table R.2: LongBench results based on 790M Mamba**\n\n|Model|2WikiMQA|GovReport|HotPotQA|LCC|MultiNews|MultiFQA|MuSiQue|NQA|PsgCnt|PsgRet|Qasper|QMSum|Repobench|SamSum|Trec|TriviaQA|Average|\n|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|\n|Mamba|13.39|20.3|8.37|39.51|20.36|27.62|3.89|8.35|1.04|1.5|10.69|20.36|38.47|20.15|45|37.59|19.79|\n|+D|0|0.21|0|2.87|0.26|0|0|0|0|0|0|0.72|3.46|0.7|0|0|0.51|\n|+SW|13.81|20.58|9.25|38.06|20.64|27.55|5.15|8.12|1.2|1.5|10.55|19.57|38.38|22.42|43|36.55|19.77|\n|+SW+D|13.6|20.39|8.55|39|20.74|27.55|4.33|7.98|0.83|1.56|10.17|19.73|38.84|22.49|43|37.02|19.74|\n|+SW+SK|13.63|20.77|8.55|38.59|20.2|27.31|4.56|7.66|0.27|1.74|11.32|20.09|38.5|20.73|44.5|37.16|19.72|\n|+LSH|1.66|1.49|1.5|7.74|0.98|2.32|0.96|0.87|0.1|0|2.3|5.05|7.13|3.89|0|0.9|2.31|\n|+KC|12.93|20.93|9.34|38.22|20.83|26.56|4.34|7.46|0.27|1.67|11.65|19.35|38.17|22.97|45|37.52|19.83|\n|+LSH+KC|12.73|20.64|8.79|39.3|21.79|28.31|3.95|7.64|0.83|2|11.94|19.63|39.07|21.04|43|38.47|**19.95**|\n\n**Table R.3: LongBench results based on new 130M Mamba**\n\n|Model|Average|\n|-|-|\n|Mamba|11.40|\n|+D|0.82|\n|+SW|11.52|\n|+SW+D|10.98|\n|+SW+SK|11.28|\n|+LSH|2.73|\n|+KC|10.99|\n|+LSH+KC|**11.58**|\n\n---\n\n**W1: The paper seems rushly written.**\n\nWe sincerely apologize for that. After submission, we have substantially revised the manuscript to improve clarity, organization, and technical precision.\n\n**W1: Experiment results are not complete; Q1: Add the QA2 task in Ruler benchmark.**\n\nThank you for pointing this out. The results are provided at Table R.1. These 790M models were pre-trained with a 2K context length, so the 4K evaluation tests their extrapolation ability. Notably, our LSH+KC architecture still outperforms all baselines on average Ruler performance (including QA2), highlighting its robustness to longer contexts.\n\n**W1: Experiment results are not complete; Q2: Add the hybrid dense attention baseline. Q6: Add the SW+LSH baseline. Q8: Include the standard deviation on joint recall.**\n\nAcross all benchmarks, our proposed Mamba+LSH+KC consistently outperforms Mamba+LSH+SW and narrows the gap between SSM and hybrid full attention. For the joint recall experiments, due to the rebuttal character limitations, please refer to our rebuttal of W2/Q2 to the reviewer MdxV. The Ruler results are provided in the following table:\n\n|Model|NIAHS1|NIAHS2|NIAHS3|NIAHMK1|NIAHMK2|NIAHMK3|NIAHMV|NIAHMQ|VT|CWE|FEW|QA1|Average|\n|-|-|-|-|-|-|-|-|-|-|-|-|-|-|\n|Mamba|99.6|53.4|7.4|20.8|0|0|17.8|4.8|1.5|1.3|100|14|26.7|\n|Mamba+LSH+SW|99.8|79.2|23.6|23|0|0|18.4|8.35|2.96|1.3|100|11.8|30.7|\n|Mamba+LSH+KC|100|92.4|34.6|24|0.2|0|20.4|3.8|4.9|1.7|100|12.8|32.9|\n|Mamba+FullAttn|100|100|57.4|35.8|1.2|5.4|31.35|10.55|2.04|4.88|100|15.6|38.69|\n\nAnd the LongBench results are provided in the following table:\n\n|Model|2WikiMQA|GovReport|HotPotQA|LCC|MultiNews|MultiFQA|MuSiQue|NQA|PsgCnt|PsgRet|Qasper|QMSum|Repobench|SamSum|Trec|TriviaQA|Average|\n|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|\n|Mamba|6.11|15.16|3.35|34.57|16.73|12.72|2.51|3.02|0.85|0.5|4.97|16.49|35.76|1.67|10.5|12.76|11.10|\n|Mamba+LSH+SW|6|14.17|2.95|34.65|15.58|13.65|2.13|2.28|0.41|0.75|5.01|13.71|36.24|1.91|14.5|12.82|11.04|\n|Mamba+LSH+KC|6.71|15.47|3.34|34.81|14.95|13.64|1.76|2.6|1.64|0.3|5.54|14.22|35.47|1.89|14|15.17|11.34|\n|Mamba+FullAttn|5.46|16.99|3.16|34.26|17.85|13.28|2.2|2.65|1.14|0|5.74|16.23|35.96|1.73|14.5|14.27|11.58|\n\n**W1: The experiment setup and evaluation settings are missing.**\n\nWe apologize if the presentation was unclear and would like to clarify that:\n- We describe the baseline setup and evaluation procedure in Section 5, with additional training and implementation details provided in Appendix D.\n- For the joint recall task, evaluation settings are detailed in Section 5.1.\n- For Ruler [14] and LongBench [3], we follow their standard evaluation settings.\n\n**W1: Technical details on KC attention are not clear.**\n\nWe provide the technical description of KC attention in Section 4.2, including the method, mathematical formulation, and training objective. The MLP is the simplest 3-layer MLP with ReLU activation.\n\n**W1: The results on LongBench seem not significant.**\n\nWe emphasize that LongBench poses significant challenges for long-context modeling and demands substantial training to perform well. Our largest continual pretraining run, a 790M model trained on approximately 8B tokens, remains far from convergence. Nevertheless, our proposed LSH+KC architecture consistently outperforms all baselines on LongBench across three evaluation settings (Table 3, R.2, R.3), underscoring its practical effectiveness even under limited academic training resources.\n\n---\n\n**W2: It is unclear how effective KC it is in real-world scenario.**\n\nWe would like to clarify that KC attention is not intended to serve as a stand-alone module, but rather as a complementary component to enhance the expressiveness of LSH attention. As shown in our experiments (Tables 2, 3, R.1, R.2, R.3), our hybrid LSH+KC attention consistently outperforms both the base SSMs and all other SSM variants augmented with input-independent sparse attention, including sliding window attention, A-shaped attention (sliding window + sink) and sliding window + dilated attention. This advantage holds across both synthetic and real-world long-context benchmarks, validating the effectiveness of KC as part of a larger input-dependent attention design.\n\n---\n\n**W4: Missing Citations.**\n\nThe two referenced papers contribute to sparse attention from an efficiency perspective. The first presents a Triton-based kernel that accelerates both QK-sparse and hash-sparse attention. The second introduces Sparse Modular Activation, which dynamically gates attention modules per token in hybrid SSM-attention models.\n\nIn contrast, our work is motivated by the goal of improving long-context modeling through a deeper understanding of its *expressiveness* constraints. First, building on a theoretical analysis of the joint-recall task, we identify a fundamental limitation in the *representational capacity* of SSMs and motivate the integration of LSH attention to address this challenge. Second, we further propose KC attention to address the *expressiveness* limitation of LSH. Finally, we integrate our method into modern state-space models, specifically Mamba and Mamba-2, demonstrating its empirical benefits at scale.\n\nWe appreciate the reviewer’s pointers to these works and will add a subsection in Related Work to discuss related sparse attentions.\n\n---\n\n**Q3: Which dataset did you use for continual pretraining?**\n\nWe followed the original Mamba setup and used The Pile for continual pretraining, as noted in Appendix D.2.\n\n---\n\n**Q4: How sensitive are the continual training results on the scalers of the auxiliary loss in KC attention?**\n\nThank you for bringing this up. While we did not extensively tune this scaling factor, we observed that choosing a suitable value is important for achieving good downstream performance. Here are the LongBench results based on different scaler values:\n\n|$\\alpha$|2WikiMQA|GovReport|HotPotQA|LCC|MultiNews|MultiFQA|MuSiQue|NQA|PsgCnt|PsgRet|Qasper|QMSum|Repobench|SamSum|Trec|TriviaQA|Average|\n|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|\n|0.05|5.3|14.62|2.63|35.62|13.4|11.42|1.82|1.92|0.66|1.08|4.73|12.5|36.73|1.63|15.5|14.11|10.85|\n|0.1|6.71|15.47|3.34|34.81|14.95|13.64|1.76|2.6|1.64|0.3|5.54|14.22|35.47|1.89|14|15.17|11.34|\n|0.2|5.56|15.23|3.18|35.78|14.01|11.98|1.78|2.03|0.66|1.17|4.83|12.89|36.29|1.66|15|14.6|11.04|\n\n---\n\n**Q5: How many parameters does the MLP have for key scoring?**\n\nThe MLP used for key scoring in KC attention is lightweight relative to the base model. We provide the comparison of the number of parameters in the following table:\n\n|Officially Released Mamba|MLP for KC attention|Mamba+LSH+KC|\n|-|-|-|\n|129.1M (100%)|0.5M (0.39%)|130.9M (101.33%)|\n|793.2M (100%)|1.0M (0.13%)| 799.1M (100.74%)|\n\nThis demonstrates that LSH+KC attention introduces minimal overhead while delivering meaningful performance improvements.\n\n---\n\n**Q7: How many test samples are in the joint recall synthetic dataset?**\n\nAs noted in Line 232, the joint recall synthetic dataset contains 10,000 test samples.\n\n---\n\nWe sincerely thank the reviewer once again for the constructive feedback and insightful questions, which have helped us strengthen the clarity, depth, and presentation of our work."}, {"review_id": "LLYm3UWbSc", "reviewer": "Reviewer_Q4tT", "summary": "The paper presents a novel synthetic task to measure recall capabilities of sequence models at varying complexity levels. The paper then studies how such a simple synthetic recall task can be hardly solved by models based on State Space Models (SSMs) while it can be easily solved by Hybrid models that combine SSMs with the Attention mechanism. The paper also presents an alternative Sparse Attention mechanism inspired by the proposed synthetic recall task and is shown to outperform other variants on small/mid scale experiments.", "questions": "1. What is the intuition for using the proposed LSH in section 4.1? The two proposed methods seem arbitrary, why is this the right way to implement the idea of having local hashing?  \n2. Why using the random matrix H? Do we need to use it at every forward pass or is it initialized and fixed when the model is initialized?\n3. Table 1 contains many ablations of the proposed method and it misses comparison with other SoTA methods (like SE-Attention and NSA in [2]). The only comparison is with Samba, however its results are pretty low, why is it the case?\n4. Table 2 and 3 is missing baseline numbers from the literature too.\n\nMinor: Line164, broken reference", "limitations": "Yes", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strength**:\n1. The paper tackles an important and timely problem: design of optimal subquadratic sequence layers.\n2. The introduced synthetic task is simple and could provide complementary information to other commonly used synthetic ones\n\n\n**Weaknesses**:\n1. The synthetic task introduced in this paper should be motivated more and better discussed in light of other commonly used benchmarks. Why do other associative recall tasks fail short? Which properties of trained models cannot be observed on other synthetic recall tasks that can only be observed in this new one? In its current form, it seems to be rather arbitrary, poorly motivated and it misses discussions with previously proposed tasks. For example, why isn’t it possible to use multiple keys with different values per key? This would still test the capability of the model to store all of the information in its state. \n2. The local hashing groups keys into bins and lets future queries attend to all keys in the same bin. However this would force the Attention to only look at similar “clusters” with little extra information, this has been shown to reduce expressivity as groups can be aggregated and merged together at reduced information loss. Why is that stage necessary? It’s been shown in SE-Attention (reference [22] in the manuscript) and [2] that it is possible to compress keys while still performing well in downstream tasks, what is the main intuition for not compressing information that is redundant? This is an important modelling assumption which deserves to be precisely motivated and validated.\n3. The manuscript misses some key references and comparisons, for example on Hybrid models the paper does not discuss the results in [1] in which authors propose a single unified Hybrid layer capable of compressing redundant information from the past and performing effective recall of old information, their main motivation is similar to the authors’ main ideas and the resulting model is similar to Figure 3, as prior art it should be included in the discussion and compared against. Furthermore, recent work like [2] proposes a similar hardware aware implementation to SE-Attention proposed in [22] (ref in the manuscript), the connection with both these works has not been thoroughly discussed in the manuscript and should be added in the method and the related work section. For example, can the authors comment more on the relationship of SE-Attention and their proposed method? \n4. How is the input dependent sparsity pattern implemented? An efficient and scalable implementation requires hardware awareness, but modern hardware cannot efficiently process and move to and from SRAM to HBM memory arbitrary access patterns. See SE-Attn and NSA for some hardware aware implementations, can the authors compare with those in their manuscript? \n\n[1] Zancato L., et al. “B’MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory”\n\n[2] Yuan J., et al. “Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention”", "was_revised": true, "final_justification": "The authors have been very receptive and engaged during the rebuttal, once the feedback is incorporated, the work will be much more solid and it'll be more impactful for the community.", "rebuttal": "We sincerely thank the reviewer for the thoughtful and positive assessment of our work. We greatly appreciate your recognition of the significance of our contributions, including the relevance of our study to subquadratic sequence modeling, the value of the joint recall task as a synthetic benchmark, and the motivation behind our proposed hybrid sparse attention mechanism. Your feedback encourages our continued exploration in this important direction. Below, we address your specific questions and concerns in detail.\n\n---\n\n**W1: On the motivation for proposing joint recall**\n\nThank you for raising this important point about the motivation of our proposed joint recall task.\n\nThe core motivation for joint recall lies in testing *context-dependent* memory, a key aspect of real-world language understanding that is largely absent in prior synthetic tasks such as associative recall. While associative recall benchmarks evaluate whether a model can retrieve a value when queried with a corresponding key, they assume a context-independent one-to-one key-value mapping. This design does not reflect the nature of natural language, where context disambiguates meaning.\n\nTo illustrate this, we provide the following example:\n\n- Associative Recall:\n  - Context: “The oceans are blue, and snow is white.”\n  - Query: “What color is snow?”\n  - Answer: “White.”\n\n- Joint Recall:\n  - Context:\n    - “In the U.S., people pay with dollars and drive on the right side of the road.”\n    - “In the U.K., people pay with pounds and drive on the left side of the road.”\n  - Query: “Which side of the road do people drive on in the U.S.?”\n  - Answer: “Right.”\n\n\nFrom this example, we can see that associative recall does not account for context. Joint recall extends associative recall by incorporating *context-dependency* into key-value associations. For example, while associative recall may map \"pay with\" to either \"dollar\" or \"pound\", joint recall allows it to map to \"dollar\" in the U.S. and \"pound\" in the U.K., depending on context. This makes joint recall a more realistic and rigorous synthetic task for both theoretical analysis and empirical benchmark for long-context modeling.\n\nRegarding the suggestion of using multiple keys with different values per key, we agree that such configurations have been explored in benchmarks like Ruler [14]. For reference, we have already tested the Ruler benchmark in our main experiment, including those tasks as well. However, those designs do not evaluate context-sensitive resolution, where the same query token (e.g., “pay with”) may map to different values depending on contextual modifiers (e.g., country). This is precisely the ability we aim to isolate and test with joint recall.\n\nThus, our task uniquely highlights the limitation of SSMs in capturing input-conditioned compositional associations, and it provides a controlled yet challenging environment for studying improvements from hybrid architectures like LSH+KC.\n\n---\n\n**W2: On the motivation of introducing LSH**\n\nThank you for this insightful question. We agree that compressing or merging keys can be beneficial as indicated in related works (e.g., SE-Attention, Native Sparse Attention). Below, we clarify our motivation for introducing LSH attention and discuss its relationship to key compression.\n\n1. *Theoretical necessity of input-dependent sparsity for hybrid sparse attention.*\n\nWe introduce the following *new proposition 4*:\n*There does not exist a 2-layer auto-regressive hybrid model consisting of an SSM layer followed by an input-independent sparse attention layer, which can solve the joint recall task with $o(L^2)$ time complexity, since it requires at least $O(L/k)$ SSM state dimensions, $k$ is the maximum number of keys that each query allowed to attend to in the sparse attention module, as defined in Eq.4.*\n\nThe *proof* is provided at the end of the response to this point. This establishes that any sub-quadratic hybrid sparse attention must employ *input-dependent sparse attention*, i.e. sparse attention with a pattern that is conditioned on the query and key representations at inference time.\n\n\n2. *LSH as an instance of input-dependent sparse attention which solves joint recall.*\n\nWe do not claim LSH attention to be the unique or universally optimal input-dependent mechanism. Rather, we show in Proposition 3 that LSH as a concrete instantiation of input-dependent sparse attention, provably suffices to solve joint recall when integrated with SSMs. By choosing LSH, we verify our theoretical result with a concrete, well-studied algorithm. Any other input-dependent scheme with similar properties would also satisfy the theory, and exploring those is an exciting avenue for future work.\n\n\n3. *LSH’s orthogonality to key compression.*\n\nIt is true that compressing keys (via aggregation or learned pooling) can preserve downstream performance in many real-world tasks. In our synthetic joint recall construction, however, no redundancy is present by our construction (please refer to our proof of Proposition 3 in Appendix C.2). For clarity and analytical tractability, we therefore omit key compression from both our proofs and our empirical study. That said, we do not rule out combining LSH+KC with key-compression techniques on real-world data; such hybrid approaches might yield additional efficiency gains.\n\nWe hope this clarifies our choice of standard LSH attention: it serves as a concrete and well-established example of input-dependent sparse attention, and when integrated with SSMs, it provides a provable solution to the joint recall task.\n\n*proof of proposition 4.* Consider a key given in the inquiry component of the auto-regressive joint recall task. The model is required to output the associated value token when this key token is provided as input. Taking the query representation from this key token, the sparse attention can attend to at most $k$ key representations from previous tokens, where the key representations are calculated based on the SSM state representations of the first layer. To solve the joint recall task, these $k$ key representations being attended must collectively encode the full context. Since the full context length is $O(L)$, by Corollary 1, $k$ state representations of the generalized SSM in the first layer must use at least $O(L/k)$ dimensions to collectively store the full context. Thus, the per-key computational cost required by the second-layer sparse attention is $O(k \\cdot (L/k)) = O(L)$, and therefore the total time complexity is $O(L^2)$.\n\n---\n\n**W3 / Q3 / Q4: On the related work: B’MOJO, SE-Attn, and NSA**\n\nWe would like to clarify that our motivation differs in key ways from the works mentioned. NSA targets hardware efficiency via compression-based sparse attention. SE-Attn improves retrieval in hybrid SSMs through relevance-based memory access. B’MOJO proposes a realization-theoretic architecture for transductive inference. In contrast, our method is motivated by an *expressiveness* analysis grounded in the joint recall task, a synthetic benchmark designed to capture key challenges in *context-dependency* modeling. This analysis reveals a fundamental limitation of SSMs under sub-quadratic constraints. To address it, we introduce a method which combines SSM with input-dependent sparse attention, instantiated by LSH attention, and we present both the theoretical analysis and empirical verification on the advantage of this method.\n\nWe also note that none of NSA, SE-Attention, or B’MOJO have released their code. As a result, a direct empirical comparison is currently not feasible. However, we will include a more detailed discussion of these works in the related work section.\n\n---\n\n**W4: On the hardware-aware implementation**\n\nWhile we agree that hardware-aware implementation is crucial for scalability and plan to explore it in future work, our current focus is on establishing the theoretical necessity and empirical effectiveness of input-dependent sparse attention. We believe this contribution is significant in its own right, as it identifies a fundamental limitation of SSMs and provides a provably expressive and empirically validated solution.\n\n---\n\n**Q1: On the hashing implementation**\n\nOur choice to implement both the argmax and sign-bit binning rules is by no means arbitrary, but follows directly from the two most widely adopted LSH-attention schemes in the literature: argmax binning (represented by Reformer [16], Learning-to-hash attention [C.1]), and sign-bit binning (represented by MemoryFormer [C.2], HashAttention [C.3]).\n\n---\n\n**Q2: On the random hash projection matrix**\n\nFollowing Reformer [16], we use a random matrix for hash projection. In our implementation, we randomly sample a matrix at each step during training, and fix a random matrix during inference. We do not perform any tuning or adaptation of the random matrix at inference time.\n\n---\n\n**Q3: On Samba performance**\n\nThank you for the question. While we have not fully understood the root cause of this discrepancy, we ensured that the comparison was fair by using consistent training settings under the FLA implementation [C.4]. The details of our setup and efforts to ensure fairness are provided in Appendix D.1. We will continue investigating this behavior to better understand the performance gap in future work.\n\n---\n\n**Minor: broken reference**\n\nWe’ll correct it in our future versions.\n\n---\n\nWe once again thank you for your thoughtful and constructive feedback. Your comments have truly helped us clarify our motivation, contributions, and distinctions from prior work, and we appreciate the opportunity to improve our future research as well as the manuscript through your insights.\n\n---\n\n[C.1] Sparse Attention with Learning to Hash. ICLR 2022.\n\n[C.2] MemoryFormer: Minimize Transformer Computation by Removing Fully-Connected Layers. NeurIPS 2024.\n\n[C.3] HashAttention: Semantic Sparsity for Faster Inference. ICML 2025.\n\n[C.4] fla-org/flash-linear-attention. github."}]

Author rebuttal:

Meta-review: The paper argues that to achieve sub-quadratic long-context modeling with State-Space Models (SSMs), it needs input-dependent sparse attention. It introduces "joint recall," a novel synthetic task testing context-dependent key-value associations. And based on this task, It formalizes the assumption (SSMs need needs input-dependent sparse attention) via an expressiveness argument (e.g., a proposition showing input-independent sparsity is insufficient under sub-quadratic constraints).

Guided by this analysis, the authors propose a hybrid mechanism LSH + Key-Centered attention:
LSH attention for content-addressable retrieval with query/key-conditioned sparsity.
Key-Centered attention to emphasize “global hub” tokens (e.g., instructions), with a lightweight scoring MLP trained via a ranking objective.

Empirical evaluation is done on Joint Recall, RULER, and LongBench, including continual pre-training of Mamba/Mamba-2 at 130M and 790M scales. Reported results show consistent gains for LSH+KC over SSM baselines and SSMs with input-independent sparsity.

Strengths
1. The expressiveness argument well motivates the necessary of input-dependent sparsity
2. Joint Recall is a useful task that isolates context-dependent composition
3. Across synthetic and real benchmarks, LSH+KC outperforms pure SSMs and SSMs with input-independent sparse patterns

Weaknesses:
1. While LongBench and RULER are appropriate, coverage of widely used tests is buried in the appendix and could be surfaced
2. Claims of “sub-quadratic” would be more convincing with throughput/latency/memory numbers
Decision: Accept (poster)