Understanding and Improving Length Generalization in Hierarchical Sparse Attention Models
Authors: Jiaqi Leng
Xiang Hu
Junxiong Wang
Jianguo Li
Wei Wu
Yucheng Lu
Venue: ICLR 2026 Poster (2026)
Keywords: long-context modeling
length generalization
length extrapolation
sparse attention
language modeling
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We demonstrate that extreme length generalization in hierarchical sparse attention is enabled by the interplay of an expressive chunking, a stable bypassing residual path, and enforced retrieval sparsity.

Abstract: Effectively processing long contexts is a critical challenge for language models. While standard Transformers are limited by quadratic complexity and poor length extrapolation, alternative architectures like sliding window attention and state space models sacrifice the ability to effectively utilize the full context due to their fixed-size memory. Chunk-based sparse attention has emerged as a promising paradigm for extreme length generalization, yet the key architectural principles underpinning its success are not yet fully understood. In this work, we present a systematic dissection of these models to identify the core components driving their performance. Through a unified framework and comprehensive ablation studies, we demonstrate that a combination of three design principles is critical: (1) an expressive, non-linear Chunk Encoder with a dedicated CLS token to produce   representations for retrieval; (2) a Bypassing Residual Path to stably integrate retrieved global information without it being overridden by the local residual stream; and (3) enforced selection sparsity during pre-training to bridge the train-test distribution gap. We provide a theoretical motivation for intra-chunk information processing and landmark generation. By combining these principles, we establish a new state-of-the-art for training-free length extrapolation, successfully generalizing models trained on a 4K context to 32 million tokens on RULER and BABILong. Our findings provide a clear and empirically-grounded set of design principles for developing future, highly-capable long-context language models.

Published as a conference paper at ICLR 2026
UNDERSTANDING AND IMPROVING LENGTH GENER- ALIZATION IN HIERARCHICAL SPARSE ATTENTION
MODELS
Jiaqi Leng'* L+
Xiang 2+
Junxiong Wang
Jianguo Li*
Wei Wu
Yucheng Lu
Fudan University
Tencent AI Lab
*Cornell University
Ant Group
Ant International
NYUU Shanghai
jqleng22em. Fudan.ddu.cm cn, shammxxxhuetencent.ccom lu. www.yucheenenyu.com edu
Q github.coomm http://wwww.o..oo///////c//////.
ABSTRACT
Effectively processing long contexts is critical challenge for language mod- els. While standard Transformers are limited by quadratic complexity and poor length extrapolation, alternative architectures like sliding window attention and state space models sacrifice the ability to effectively utilize the full context due to their fixed-size memory. Chunk-basee based sparse attention has emerged as promis- ing paradigm for extreme length generalization, yet the key architectural prin- ciples underpinning its success are not yet fully understood. In this work, we present a systematic dissection of these models to identify the core components driving their performance. Through a unified framework and comprehensive abla- tion studies, we demonstrate that a combination of three design principles is crit- ical: (1) an expressive, non-linear Chunk Encoder with a dedicated CLSS token produce representations for retrieval: (2) Bypassing Residual Path stably inte- grate retrieved global information without it being overridden by the local residual stream; and (3) enforced selection sparsity during pre-training to bridge the train- test distribution gap. We provide theoretical motivation for intra-chunk infor- mation processing and landmark generation. By combining these principles, we establish a new state-of-the-an for training-free length extrapolation, successfully generalizing models trained 4K context 32 million tokens RULER and BABILong. Our findings provide a clear and empirieally-grounded grounded set of design principles for developing future, highly-capable long-context language models.
INTRODUCTION
The recent proliferation of Large Language Models (LLMs) has catalyzed an urgent demand for models capable of processing long contexts (Brown et al., 2020) Touvron et al.) 2023 OpenAl,
2024). However, the performance of standard Transformer architectures (Vaswani et al. 2023), which serve as the universal backhone, is well-documented to degrade sharply on ofitt--distributioon context lengths (Hsich et al. 2024, Xiao et al.) 2024 Lu et al. 2024; Kuratov et al.) 20244. Concur- rently, naively increasing the context window during training is often infeasible due to the quadratic complexity of the attention mechanism, which leads prohibitive computational costs. This has spurred the community to investigate novel architectures designed explicitly for length extrapola- tion.
We posit that ideal length extrapolation necessitates two key properties: (1) the ability to train on
"Equal contribution
I Work done at Ant Group.
1 "Corrsponniig authors

---

Published as a conference paper at ICLR 2026
shorter sequences while achieving stable or even lower perplexity on longer sequences during infer- ence, and (2) the capacity to effectively utilize the entire context. Significant research has focused on the first property. For instance, architectures like Sliding Window Attention (Beltagy al. 2020) and recurrent models can maintain a nearly constant perplexity as the context extends beyond the training length (Buitrago & Gu 2025). However, these methods fall short the second property. The former is constrained by a fixed-size local receptive field, while the latter compresses the entire history into fixed-size state, creating an information bottleneck. As revealed by in-context re- trieval tasks like SULER(Hsieh. et al. 2024), the inability to access and utilize arbitrary information from long contexts remains a critical limitation for many applications, from multi-round chatbots to agent-based systems (Liu et al. 2024 Pan et al. 2025 Guo et al. 2023).
Fortunately, recent methods based on retrieval and sparse attention (Mohtashami & Jaggi 2023 Hu et al. 2025b]a) a) can maintain a near-constant perplexity on out-of-domain long sequences and achieve high accuracy on simple retrieval tasks like passkey retrieval. Their approach involves dividing the entire sequence into chunks, selecting the K most similar chunks and attending to the tokens within those chunks. However, on more complex retrieval tasks, their accuracy still degrades as context length increases. Moreover, comprehensive analysis of why these models extrapolate so effectively- and what role their different designs play in this capability is largely missing.
To address this gap. we present systematic dissection of chunk- based sparse attention mechanisms, identifying the core architectural principles, theoretical motivations, and sparsity configurations that enable extreme generalization. Our main contributions are:
1. Unified Framework and Systematic Ablation. We present a unified framework for chunk- based sparse attention and conduct systematic ablation identify the critical architectural compo- nents for generalization. We empirically demonstrate that an optimal combination of three design choices is essential: a non- non-linear Chunk Encoder for expressive representations, a dedicated CLS token for disentangling retrieval and content information, and a Bypassing Residual Path for stable information integration.
2. Theoretical Motivation and Diagnostic Analysis. We provide a theoretical motivation for the Chunk Encoder, framing it as necessary non-linear approximator for full attention scores. We sup- port our architectural claims with in-depth diagnostic analysis that correlates intermediate retrieval accuracy with final task performance, revealing how each component improves generalization by enhancing either retrieval prominence or information integration.
3. State-of-the-Art Extrapolation and Insights on Sparsity. By combining these optimal design principles, we establish a new state-of-the- art, generalizing model trained on 4K context to 32M tokens. We further provide a detailed analysis on sparsity, showing that both a large context for contrastive learning and enforced selection sparsity during pre-training are crucial for generalization.
2 RELATED WORK
Approaches to long-context modeling can be broadly categorized. The first comprises local methods like Sliding Window Attention (Beltagy et al. 2020 Xiao et al., 2024). While effective at maintain- ing perplexity, they are fundamentally limited by fixed-size local receptive field, preventing access to information beyond their local window. A second category includes recurrent architectures (Gul & Dao 2024 Dao & Gu 2020; Sun et al.) 2023 Qin et al. 2023 Peng et al.) 2023 Yang et al. 2024). Recent work shows that zero-initialized recurrent models struggle to generalize to longer sequences; it attributes the failure to models entering unexplored states and introduces State Pass- ing and Truncated Backpropagation Through Time (Sutskever 2013) Williams Peng 1990) to improve length generalization, achieving stable perplexity on cut-of-distribution lengths (Buitrago & Gu 2025). However, on challenging retrieval tasks (Hsieh et al., 2024), performance still drops significantly beyond the training length. This is because the entire context history is compressed into a fixed-size state, sacrificing the ability to recall specific, distant information losslessly. The third category seeks extend full attention by modifying positional encodings (Press et al.) 2022 Peng
eet al. (2024] Chen et al.] 2023) Tan et al. 2005). Their training-fre extrapolation gains are modest, typically limited to small multiple of training length.
In pursuit of more scalable extrapolation, chunk-based sparse attention has emerged as a distinct and highly effective paradigm. Landmark Attention (Mohtashami & Jaggi 2023) pioneered this by
2

---

Published as a conference paper at ICLR 2026
using dedicated landmark tokens to represent chunks for retrieval. While innovative, its length exp- trapolation capability is limited to about 64x training length. NSA (Yuan et al.) 2025) proposed an efficient hardware-alignes implementation of chunk selection but has very limited length extrapola- tion (Hu et al. 2025a). State-of-the-an training-free frre extrapolation has been achieved by Differen- tiable Retrieval- bbased Transformers (DRT) (Hu al.) 2025b) and RAMba (Hu al. 2025a). which can generalize to 1000x their training length on simple retrieval tasks by introducing a learnable hierarchical chunkwise sparse attention for retrieval. These methods also differ in their retrieval fre- quency, with DRT operating per-chunk while RAMba, NSA, and Landmark retrieve per- per-token. Our work builds upon the SWA+Hierarchical Sparse Attention (HSA) architecture with per-token re- trieval, providing a systematic analysis to dissect and identify the specific architectural components responsible for the extreme generalization capabilities of this emerging paradigm.
PRELIMINARY
3.1 RANDOM CONTEXT ACCESS: THE KEY ABILITY TO UTILIZE FULL CONTEXT
Effective long-context reasoning hinges on the ability to flexibly retrieve information from any part of the preceding context. This capability, which we term Random Context Access (RCA), critical for models to excel at tasks requiring the synthesis of distal information. We define it as follows:
Definition:Al any given step t, Random Context Access is the ability of a model to attend to any subset of tokens from the preceding context {zo..... -z-1 } in data-dependent and lossless manner.
Under this definition, different architectures exist on a spectrum of capability. The standard Trans- former, with its full-attention mechanism, perfectly exemplifies RCA by granting unrestricted access to the entire KV cache. In stark contrast, many efficient architectures trade this property for com- putational gains by imposing structural constraints. These include: locality constraints, in which models like Sliding Window Attention are confined to a local neighborhood: information bottle- necks, created by State Space Models (SSMs) that compress history into lossy, fixed-size state; and static access patterns (e.g.. [Chil et al. (2019); [Beltagy et al. (2020), where data-independent sparse attention predetermines information pathways, with more details in Apx. [B.] These archi- tectural bottlenecks fundamentally limit model's capacity for true long-range reasoning.
3.2 DYNAMIC CHUNKWISE SPARSE ATTENTION
Chunkwise sparse attention offers a practical and scalable approximation of RCA. The principle of dynamic chunkwise sparse attention mirrors an efficient strategy for open-book exam. Rather than exhaustively reading the entire source material (akin to full attention), one first identifies key pages using high-level bookmarks (landmarks). Queries then resolved by first consulting these landmarks to locate the relevant pages (chunks) and only then reading their detailed contents (to- kens). This hierarchical retrieval drastically reduces the search space, with its success contingent the quality of the landmarks.
Following this principle, our method provides access the global context by dynamically selecting
SWA+HSA Architecture. Our SWA+HSA model is primarily composed of Trans- formers with Sliding Window Attention (SWA), with layers divided into upper layers and lower layers. Lower decoders are composed of SWA followed by an FFN. At the architectural midpoint, chunking layer splits representations from lower decoders into chunks, and encode them by chunk build global memory, each chunk of which contains a landmark vectors (Imk) and encoded chunks (E). Upper de- coders each contains a local self-attention layer for local context, followed by an HSA to attend selected global memory and then an FFN.
Figure 1:
chunked memory shared by HSA
Uspar decoder
link,
Lmkka
FFN
WOWTOOO: Chunking
HSA
SWA
H2
Chunk Selection
FFN
SWA
24
3/S
3

---

Published as a conference paper at ICLR 2026
fixed number of chunks. While this imposes chunkwise access pattern, it aligns with empirical findings that attention scores often exhibit a blockwise structure (Yuan et al.) (025) and that most queries only require a local subset of the context. This approach thereby offers a compelling trade- off between unrestricted context access and computational feasibility.
3.3 HIERARCHICAL SPARSE ATTENTION AND MODEL. DESIGN
As shown in Figure D our SWA+HSA model treats the global memory in chunkwise approach. During decoding. the chunking layer gathers hidden representations into a buffer and encodes them complete chunks. To maintain causality, the chunk containing the current token, well as all subsequent chunks, are not accessible for retrieval. Following the notation indices with [-] to indicate that it is indexed by chunk rather than by token. For instance, Hi == +,.0+1((++)CC E C is the chunked hidden states. At chunk selection layer, a query vector qt RRdee computes a dot-product similarity score on landmark Imke E TR formulated as St 9: 1mk E R. The indices of the top-N chunks are then selected, forming set Ię-
[un et al.) (2023),
A critical design choice lies in how these selected chunks are weighted, as this determines how the retrieval signal backpropagated. The weighting function Wi, is zero all non-selected chunks. For the selected chunks, the weights are assigned according to distinct strategies for NSA (Eq. Ta GCA (Eq. lb), and HSA (Eq. Ic Let S (e,,)
(Ic) This comparison highlights clear progression: from the implicit binary weighting of NSA, to the proportional softmax weighting in GCA (Hu et al., (0255), and finally to the recency-biased
Wr,
(la)
W1,i = Softmax(S:): (lb)
Wr, = StickBreak(Ss);
allocation via stick-breaking proces: Tan et al. 2025) used in HSA (our work and Hu et al. 2025a).
In the second stage, the selected chunk indices I, and their corresponding weights we are used to compute the final attention output. The HSA layer takes the query Q, and the selected keys and values, (K(z,) Vizi]l and performs a weighted-sum over standard attention operations:
Q:Ki
O,
Σ Wt, Attention (Q. VEI
Qg, K VI.
Σ SEZ
Wh,s Softmax
VHI M
(2)
Vdle
The direct participation of the weights Wt, in the forward computation (Eq. 2 creates differ- entiable pathway where chunks deemed more important contribute more significantly to the final output. Crucially, in the backward pass, the gradients flowing to each chunk's parameters are scaled by their weight Wryy. This enables the model to perform effective credit assignment, concentrating learning the most relevant retrived chunks.
SYSTEMATICAL STUDY OF RETRIEVAL AND INTRA-CHUNK PROCESSING
4.1 CHUNK INFORMATION PREPROCESS
Imkk
Mean Pool
[muk
As shown in Fig. h serves dual purposes, (1) passed to upper decoder layer as input, (2) gath- ered with adjacent hidden states to form the chun- ked hidden states H to generate global mem- ory (chunked KV and landmarks). We define these operations through a pair of functions: f(-) for landmark aggregation and g(-) for intra-chunk processing. The final outputs are then derived via linear projections:
Chunking
Chunking
H2
16]
(b)
cts
Figure 2: Design of Encoder: (a): Encoder
w/o CLS (b): Encoder with a learnable CLS token.
'Specifically, let I sort(I,) be the selected chunk indices sorted by their scores. then wi.k ass,, I[ (s where is the sigmoid function.
4

---

Published as a conference paper at ICLR 2026
Table 1: Unified Comparison of Chunk Processing: The Encoder is bidirectional Transformer encoder shared by f()) and g(-). For the "w/ CLS" variant, [0] selects the output corresponding to the prepended XCLS token, and [1 :] selects the outputs for the original chunk hidden states H[].
Configuration
{(H
(1)
NSA
Pool (1)
H
HSA w/o Encoder
MeanPool (N
RMSNorm(HH (Hp))
R
HSA w/ Encoder w/o CLS
MeanPool (Encoder (eenoodelEn H (1))
Encoder(H(g)
HSA w/ Encoder w/ CLS
Encoder([xes (xcrsr H (1)
Encodor( [xis; H [1
1mkg) = Linear(f(H(g)) E R and (i). V(n] Linear(g(H()) E T gge +d.)
The different architectural configurations we investigate, summarized in Table D can be expressed as joint definitions of (f+g). In the "w/ CLS" variant, we prepend learnable token, XCLS, to the input chunk [:), as shown in Fig. [ The Encoder processes this combined sequence, and its output corresponding to the XCLS position is used form the landmark, while the remaining outputs form the KV cache.
Chunked Hidden States: Disentangling Representations for Retrieval and Prediction. The hidden state h must be immediately useful for autoregressive next token prediction, while also serving as durable memory for future retrieval. Considering the dual purposes of h, we posit that simply operating on these raw hidden states conflates these two distinct functional objectives. The introduction of non-linear layers aims disentangle these two roles. By processing the raw states, these components can produce refined representations specifically optimized for retrieval, separate from the immediate demands of next-token prediction. A RMSNorm provides simple mechanism to introduce oon-iinerrtty, while the Encoder takes this step further by explicitly learning to re- encode the chunk's information.
Landmark: Approximating Full Attention via Learnable Summarization. We frame HSA as a structured approximation of full attention. The goal for the HSA weight (r, 3b faithful approximation of the full attention weight ar (Eq, Bab.
exp(q: k,)
exp(q: k2)
C1,i
ビj=1 1-1 exp(q: kj)
(3a)
6,, = Wtee ESec, exp(q: ky)
(3b)
Formally, the ideal chunk weight we,c, should approximate the ratio of its chunk's 's attention mass to the global attention mass, as shown in Eq. 4a However, since landmark is generated using only local chunk information, it is impossible to compute the global denominator. Therefore, the model must learn a practical proxy: the unnormalized selection score St,e; must be made proportional to the chunk's unnormalized attention mass, relationship defined in Eq. 4b
Ejecr exp(g: kj) Wh.cc -1=1 exp(q: k;)
q. Imk [j & Σ exp((: ky) (4b)
(4a)
JEC.
The core challenge is thus to generate a landmark Imkij that satisfies the relationship in Eq. 4b The right-hand side of this equation is highly nonlinear function of the keys within the chunk. A simple linear operator like MeanPool is insufficient for this task. A multilayer Chunk Encoder, as powerful learnable function, essential for learning this complex relationship. Adding a CLS token and using its output landmark, further improve the nonlinearity and learnailiity compared with mean- mean-pool seaalnpooing the Chunk Encoder outputs.
4.2 SKIP CONNECTION DESIGN FOR CROSS-LAYER INFORMATION FUSION
Some HSA layers in model can retrieve information to the shared long rang-range memory, by attend- ing to the key-value states at the output of the chunking layer. This design introduces challenge: how to effectively fuse the more abstract information of the current upper layer, X with the more
5

---

Published as a conference paper at ICLR 2026
literal, lower-level context retrieved by the HSA module, H(xx))
We investigate two alternative skip connection strategies for integrating the HSA layer within the upper decoder blocks. In both designs, the output of the HSA layer is added to the imput residual to form an intermediate representation. x' (Eq. 5). This representation is then processed by the sub- sequent MLP block, MM-). The key distinction between the two strategies lies in the final residual connection, The Standard Sequential Path (Eq. [aa follows the conventional Transformer design, adding the MLP's output back the intermediate representation x'. In contrast, the Bypassing Residual Path (Eq. 6bb provides more controlled mechanism, ensuring the direct output from the HSA module bypasses the final residual addition.
Xou = x + +Mxx) = Xin + +M(x) + H(xx)
(6a)
x' Xin H((x)))
(5)
Xout Xin + M((')
(6b)
We hypothesize that the Bypassing Residual Path is particularly suited for our cross-layer attention architecture. The direct addition of cross-layer information H(xx)) the highly-processed residual stream, as in the Standard Sequential Path may disrupt the information flow. On the contrary, the Bypassing Residual Path tasks the MLP with explicitly learning how to resolve the information gap between network depths.
EXPERIMENTS
We organize experiments follows. Section 5. describes the experimental setup, including model architecture, training details, benchmarks, and baselines. Section presents length extrap- olation results BabiLong and RULER, demonstrating the superiority of our SWA+HSA architec- ture. Section 5.3 provides an in-depth analysis of the retrieval mechanism, disentangling retrieval accuracy from information integration. Section 5.4 investigates the role of sparsity in enabling extreme-length generalization. Finally, Section 5.5 validates the scalability and general capabilities of our approach.
Model Performance vs Sequence Length (Average across all tasks)
0.8
w/ encoder w/ CLS, w/ bypass ours w/ encoder w/ CLS, w/o bypass ours encoder w/o CLS, w/ bypass ours w/ encoder w/o CLS, w/o bypass ours w/o encoder, w/o CLS, w/ bypass ours w/o encoder, w/o CLS, w/o bypas
0.6.
Average
LLaMA with YaRN
Stick-Breaking Transformers
Landmark Attention
0.2
Sliding Window Attention
Mamba2 w/ State Passing
8.1
4K
200k
Sequence Length (Tokens)
The best SWA+HSA model (ours) demonstrates robust long-context extrapolation, maintaining high accuracy up to 8M tokens. In stark contrast, baseline models exhibit clear limitations. Full-attention methods collapse near-zero accuracy shortly after their training length. Prior sparse methods like Landmark Attention show better generalization but also fail decisively at longer contexts (264k). Sliding Window Attention plateaus at a low accuracy consistent with random guessing, while Mamba2's performance degrades substantially.
Figure 3: BabiLLong Evaluation Results:

---

Published as a conference paper at ICLR 2026
Table 2: RULER evaluation results. Our best SWA+HSA model maintains high accuracy up to 32M tokens, significantly outperforming all baselines in long-context scenarios. In Model Specs, "Enc" denotes the number of encoder layers, 'CLS' indicates the use of CLS token, and 'Bypass' refers to the Bypassing Residual Path.
Model
Enn
CLS
Bypuss
S-N
MQ-N
VT
Avg.
S.N
MQ.N
VT
Avg.
S-N
MQ.N
VT
Avg.
Model Specs
ctx-len=4K
cIx-lem=32K
Mambu2
0
99,54
15.96
80.80
65.43
1.48
1.49
0.40
1.12
0.00
0.00
0.00
0.90
NSA
22.82
3.15
12.43
12.80
0.00
0.00
4.48
0.00
0.00
0.00
0.00
Llama Yarn
92.39
69.11
93.51
85.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.90
Stick Break
0
76.72
22.17
37.20
45. 36
73.88
24.63
31.34
43.28
33.33
18.18
30,30
27.27
Landmark
37.00
40.00
83.00
0.00
SWA-HSA
0
no
nO
84.40
61.30 30
37.02
60.91
78.97
38.80
12.14
43.30
77.49
35. 14
7,80
40.14
SWA-ISA
81,20
68.40
87
79.13
$1.00
66.20
87.40
78.20
77.20
62.60
86,60
75.47
SWA-HSA
I
yes
B00
79. 80
70:40
90.00
80.07
81.00
86.40
77.60
75.60
62.20
84.80
74.20
SWA-HSA
2
no
nO
77.39
67.52
89.24
78.05
78.58
65.75
87.66
77.33
74.83
64.76
87.27
75.62
SWA-HSA
2
yes
R0
86.28
71.96
93.98
82.82
69
82.16
SWA-HSA
0
nn
yes
90.33
71:77
72.26
78. 12
88.06
71.08
69.20
76.11
84.21
69.00
65.75
72, 72.99
SWA-HSA
I
89,20
77.00
93.80
86.67
89.40
75.60
90.80
85.27
84.40
73.40
90,20
82.67
SWA-HSA
I
yes
yes
91.20
75.20
89.20
72.80
83.93
20
74.00
87.80
83.33
SWA-HSA
2
100
yes
91.71
72.56
93.68
85,98
88.15
72.36
92.20
84.24
86.97
71.17
92.99
83.71
SWA-HSA
2
yes
yes
91.51
78.78
87.92
89.14
73.84
91.61
84.86
88.94
76.21
91.81
85.65
Models
Enc
CLS
Rypass
S-N
MQ-N
VI
Avg.
S-N
MQ-N
VT
Avg.
S-N
MQ-N
VT
Avg.
Model Specs
ctx-len=IM
ctx: SM
ctx-len=52M
SWA-HSA
0
no
n0
72.28
21.78
3.96
32.67
61.39
6.93
1.988
23.43
61.80
1.12
0.00
20.97
SWA-HSA
I
80
72.73
58. 18
72.73
67.88
69.09
34.55
58.18
53.94
70.91
47.27
43.64
53.94
SWA-ISA
yes
B00
40.00
54.55
38.18
52.12
35.76
SWA-HSA
2
D0
B0
63.37
63.37
72.28
66.34
58.42
48.51
66.34
57.76
67.42
42.70
70.79
60.30
SWA-IISA
2
yes
83,17
46.53
$2.18
70.63
78.22
28.71
80.20
62.38
82.02
24.72
61,80
56.18
SWA-HSA
0
BO
yes
80.20
64.36
45.54
63.37
71.29
46.53
13.84
44.55
82.02
31.46
3.37
38.95
SWA-IISA
I
nO
69.09
56.36
45.45
56.97
58.18
30.91
10.91
33.33
63.64
36.56
12.73
37.58
SWA-HSA
yes
yes
69.09
80.00
72.73
78.18
75.76
SWA+ISA
2
nO
yes
85,15
74.26
74.26
77.89
78.22
65.35
65.35
69.64
82.02
66.29
41.57
63.29
SWA+HIA
2
yes
yes
88.12
76.24
93,07
85.81
79.21
77.23
90.10
82.18
82.02
68.54
88.76
79.77
5.1 EXPERIMENTAL SETUP
Our experiments are conducted on series of models based on the SWA+HSA architecture detailed in Figure [ All model variants contain approximately 240M parameters and are trained with context length of 4K. Unlike the original DRT (Hu et al. 2025b) which uses Grouped Cross Atten- tion (GCA) for per-chunk retrieval, our architecture employs HSA with stick-breaking weighting strategy (Eq. Icp to enable per-token retrieval, providing greater flexibility for dynamic access to long-range information.
Key hyperparameters for this setup are fixed across experiments. The sliding window size for the lower layers is 512 tokens, and the HSA chunk size is 64. Unless otherwise noted, we use Top- K of 8 for chunk selection, giving the HSA module an effective receptive field of 512 tokens. A comprehensive description of our training recipes and further hyperparameter details is provided in Appendix C
Benchmark. We evaluate our models on two long-context benchmarks: BabiL.nng (Kuratov al., 2024) and RULER (Hsich et al.) 2024). BabiLong is designed to simulate complex question- answering scenarios over long documents. It extends the classic bAbl dataset (Weston et al. 2015) by embedding the original short stories and factual sentences into a long distractor context. This design challenges a model to first perform robust, long-range information retrieval to locate all rele- vant facts, and then apply the reasoning required by the original task. Consequently. BabiLong tests retrieval and reasoning capabilities simultaneously.
To complement this and more directly assess model's core retrieval ability in isolation, we also use RULER. From RULER, we select three representative sub-tasks: Single-Needle (S-N); Multi- Query Needle (MQ-N), where 2 of 6 key-value pairs are queried; and Variable Tracking (VT), which tests transitive assignment chains. Collectively, these tasks provide a targeted evaluation of model's capacity for precise, long-range information retrieval and tracking simple logical chains. As RULER requires accurately retrieving several tokens, it provides a strict test of models' random context access ability.
Baselines. We compare our method against several baselines. Our primary full-attention baseline is Llama with YaRN and Mamba2 with State Passing (Buitrago & Gu 2025). which is trained
7

---

Published as a conference paper at ICLR 2026
SWA HSA Model:: Retr eval Recall Top-K
SWAA HSA Models: Ratr loval Recall @ Top-K
SWAk HSA Model:: Marginal Integration Precision
Cumulytire Recil Rore
Margina Predron
Comulettve
SNA+HSA Model Setup wi ercoder CLS. of typass
-0- wi encoder CLS wit bysass
wi encoder CLS, bypass
0: wi encoder wis CLS we bycess
bypass
(a) S-N Top-K Retrieval Accuracy
(b) MQ-N Top-K Retrieval Accuracy
(c) S-N Integration Precision
Figure 4: Analysis of Retrieval and Utilization Mechanisms. Subplots (a) and (b) show retrieval accuracy (Recall @ Top-K) K) on the S-N and MQ-N tasks, respectively. The results highlight that the Bypass mechanism is the most critical component for successful retrieval, consistently outperform- ing variants without it. The learnable encoder and CLS token also provide noticeable improvements. Subplot (c) plots the conditional answer accuracy on the S-N task, given that the correct chunk was retrieved at a specific rank (by weight). Across all models, integration precision drops sharply as the rank of the correct chunk decreases. This demonstrates stricter requirement: a chunk must not only be successfully retrieved (recall) but also be assigned high-ranking weight be effectively utilized by the model.
on 4k context length. To ensure fair comparison with other efficient attention mechanisms, we include Landmark Attention, configured to match our model's training FLOPs during training by using sequence length of 756. Furthermore, benchmark against methods with similar computational budget, we include Stick-breaking Attention and Sliding Window Attention. Both of these are trained with Ik context window keep the receptive field same to our approach.
5.2 LENGTH EXTRAPOLATION RESULTS
As shown in Figure the results on BabiLong highlight the stark differences in extrapolation ca- pability across architectures. The benchmark's classification-like nature establishes a random-guess baseline around 0.3 accuracy, level to which Sliding Window Attention predictably converges due to its limited context. Other baselines also fail to generalize, with full-attention and Landmark At- tention collapsing various points beyond their training regimes. Crucially, the figure also reveals the importance of our key design choices within the SWA+HSA architecture. The model variants with learnable chunk encoder consistently outperform those without.
To further validate these findings and provide a more granular analysis of retrieval performance, we The trends observed here strongly
turn to the RULER benchmark, with detailed results in Table [ corroborate our conclusions from BabiLong. First, comparing our SWA+HSA models against the baselines reveals significant performance gap in extrapolation. Llama with Yarn (Llama-Yand despite its strong in-domain performance at 4K, fails catastrophically at 32K, with their scores drop- ping zero. Stick-breaking Attention (Stick- BreckBBrekk and Landmark Attention demonstrate some extrapolation ability but exhibit significant performance decay and are consistently outperformed by our models on longer contexts.
Second, while all our model variants showcase robust long-context performance, the results also highlight a clear performance hicrarchy among our different architectural contigurations. Our full model (Enc=2, CLS=yes, Bypass=yes) consistently achieves state-of-the-art results. It not only secures the highest average score of 87.92% at the 4K training length but also degrades remarkably gracefully. maintaining average score of 79.77% even at extreme context length of 32M. This indicates that the specific combination of our proposed components critical for unlocking maximum performance, interaction will analyze in detail in the subsequent section.
8

---

Published as a conference paper at ICLR 2026
Table 3: Ablation study on sparsity using the RULER benchmark. The interplay between train- ing context length and sparsity is critical for extreme extrapolation. A model trained on a 4K context with high selection sparsity (Top-K ( ultimately achieves the best performance at 32M tokens.
Model Spees
etx-lenssK
ctx-lon=32K
ete-lenm 128K
TPR
Top-K
Train-Length
S-N
MQ-N
VT
Avg.
S-N
MQ-N
VT
Avg.
S-N
MQQN
VI
Avg.
SWA+HISA
1024
38.68
2.50
11.60
17.59
$1.49
0.75
0.00
17.41
36.36
0.00
0.00
12.12
SWA+HSA
2048
95.27
67.90
56.86
73.34
96.27
71.64
10.45
59.45
87.88
63.64
18.18
56.57
SWA+HIA
78.78
84.86
SWA+HSA
8192
98.70
90.54
96.25
95.33
95.52
92.54
97.01
95.02
96.97
87.88
96.97
93,94
SWA+HSA
64
4096
94.97
$8.85
97.73
93.85
95.56
$6.38
97.63
93.19
91.51
87.96
98.72
92.73
SWA+HSA
4076
93.69
31.82
71.80
97.288
24.63
88.06
06.06
93,94
6.06
63.64
54.55
Model
Spees
orx-len-1M
cnx-lan-MM
crx-Ien-32M
TPR
Top-K
Tmin-l ength
S-N
MQQN
VT
Avg.
S-N
MQQN
VT
Avg.
S-N
MQ-N
VI
Avg.
SWA+HSA SWA+HSA
1096
88.12
76.21
93.07
85.81
79.21
77.23
90.10
82.18
82.02
68.54
88.76
79.77
$4.00
79.17
95.00
SWA+HSA
64
4096
78.18
90.91
96.36
88.48
51.31
76.36
94.55
74.07
50.91
69.09
89.09
69.70
5.3 IN-DEPTII ANALYSIS OF RULER RETRIEVAL
To further understand the mechanisms behind our architectural designs, we move beyond final task scores to conduct a diagnostic analysis of the intermediate retrieval process. Specifically, we investi- gate the relationship between the accuracy of chunk selection and the correctmess of the final model output on the S-N and MQ-N tasks. In these retrieval-centric tasks, we can unambiguously identify the critical "needle" chunk(s) containing the information required answer the prompt. We then measure the retrieval recall for these essential chunks and correlate it with final answer accuracy. This allows us to disentangle two primary failure modes: failure to retrieve the relevant context versus failure to utilize the context once it has been retrieved. Given that the evaluation context length far exceeds the local sliding window, successful task performance becomes contingent on accurate onggrange retrieval in most cases.
Our analysis reveals that even at an extreme context length of 8M, most model configurations, with the notable exception of the poorly performing baseline (without encoder bypass path), suc- cessfully retrieve the correct "needle" chunk within their top-k(k=8) selections. The significant performance gap, therefore, does not originate from complete failure of retrieval, but rather from two more subtle factors: Retrieval Prominence: Superior models assign a significantly higher rank and selection weight the correct chunk within the top-k set, indicating more precise and confident retrieval score, as indicated in Fig. 4a and Fig. 4b As shown in Fig. 4c the model can integrate information more accurately if the target information is assigned higher importance among the selected chunks. Information Integration: After successfully retrieving the relevant chunk(s), better-performing models are generally more effective at integrating this information to produce the correct final answer.
Our experiments empirically validate the role of our proposed architectural components in improv- ing both these aspects: 1. The Bypassing Residual Path is erucial for effective information inte- gration. We observe that the bypassing path ((Bypass-yes') consistently and significantly improves final answer accuracy, even when the correct chunks are among selected chunks. This provides strong evidence for our hypothesis that retrieved information requires dedicated modulation before being integrated into the main residual stream. The bypassing design proves to be highly effective mechanism for this modulation. 2. The Chunk Encoder and CLSS token improve retrieval promi- nence. Our results show direct positive correlation between using the encoder and CLS token and the rank of the correct chunk. These models have the same total parameter count; Models with one more encoder layer have one less lower-decoder layer. The fact that this architectural trade-off leads to superior performance on ectrieval-nteessive tasks demonstrates that the encoder is a more carameter-efficiee mechanism for learning robust retrieval representations and adding a CLS token to generate landmark can improve the parameter efficiency of the chunk encoders. Based on these findings, we also significantly improve length generalization RAMba, detailed in Appendix. D
5.4 THE ROLE OF SPARSITY
The effectiveness of our model critically depends on the careful orchestration of sparsity across multiple dimensions. We conduct systematic ablation study to analyze three key dimensions of

---

Published as a conference paper at ICLR 2026
sparsity using our best best-performing configuration, with detailed results presented in Table
Sufficient Training Length. This dimension specifies the context window size used during train- ing, which determines the pool of chunk candidates available for the sparse attention mechanism to learn from. A longer training context exposes the model to more diverse retrieval scenarios, en- abling it to develop a more robust retrieval policy through richer contrastive signals. The empirical evidence strongly supports this: models trained on 4K contexts successfully extrapolate up to x their training length, while those trained on only 1K contexts fail to generalize beyond the training regime. Further extending training length to 8K brings additional gains, confirming clear positive correlation between training context length and extrapolation capability.
Selection Sparsity (Top-K). This parameter controls the number of chunks the model can attend to at each decoding step, directly regulating the sparsity of the retrieval mechanism. While a large Top-K value (e.g., 64) permits nearly dense global attention, small Top-K (e.g., 8) enforces strict sparsity, compelling the model to learn highly discriminative retrieval. This enforced sparsity during training proves essential for extreme-length generalization. At 32M tokens, the model trained with Top-K=8 substantially outperforms its Top-K=64 counterpart, demonstrating that a sparse retrieval policy learned during training generalizes more effectively to the resource constraints and scale demands of inference at unprecedented sequence lengths.
Retrieval Frequency (Tokens Per Retrieval). This dimension defines the retrieval interval: how many tokens are generated before the HSA mechanism refreshes its selected chunks. In other words, specifies whether chunk selection is updated every single token (per-token retrieval) or once per block of multiple tokens (per-chunk retrieval). Two competing strategies exist: per-token retrieval offers maximal flexibility through fine-grained, dynamic context access at every step, while per- chunk retrieval potentially provides more stable training signals by reducing update frequency. To empirically determine the optimal strategy, we compare per-token retrieval against per-8 per-8-token re- trieval. The results show a clear tradeoff: per-8-token retrieval achieves marginal gains on simple sengle-needle tasks (S-N) but suffers severe degradation on complex multi-query tasks (MQ-N, VT) requiring multiple targeted retrievals. This decisively favors per-token retrieval as the more robust choice for our final architecture.
5.5 SCALABILITY AND GENERAL CAPABILITIES
To further validate the scalability and general capabilities of our proposed architecture, we evaluate scaled-up Lixture-of-Expertt (MoE) variant (8B parameters, 1B active). The detailed training recipe and model specifications follow the setup in Hu et al. (2025c). As shown in Appendix [A] our architecture achieves competitive performance on general reasoning benchmarks (c.g., GSM8K, MMLU) compared to standard Transformers. confirming that our length-extrapolatie modifications do not compromise short-context capabilities.
CONCLUSION
In this work, we conducted a systematic investigation into the architectural principles underpinning extreme length generalization in chunk-based sparse attention models. Through unified framework and comprehensive ablation studies, we demonstrated that a combination of three design principles is critical: (1) an expressive, non-linear Chunk Encoder with dedicated CLS token to generate disentangled representations for retrieval and content; (2) a Bypassing Residual Path to stably inte- grate retrieved global information without disrupting the local processing stream; and (3) enforced selection sparsity during pre- pre-training to align the model's learned retrieval policy with the demands of inference unseen, longer contexts. Our in-depth analysis revealed precisely how these com- ponents improve performance by enhancing either retrieval prominence or information integration. By combining these principles, we developed model that generalizes from a 4K training context 32 million tokens without fine-tuning. Ultimately, our findings distill the complex design space into set of guidelines for developing future, highly-capable long-context
clear, language models.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
We thank the amonymous reviewers for their insightful comments. This work was supported by the Ant Group Research Intern Program.
ETHICS STATEMENT
The authors have read and adhere to the ICLR Code of Ethics. This work is foundational, focusing on the architectural principles of language models rather than a specific application. To uphold the prin- ciples of scientific transparency and reproducibility, our experiments were conducted exclusively on the publicly available datasets, and we commit releasing our source code upon de-anonymization to allow for verification and future research. We believe that transparent and reproducible research crucial for the responsible development future AI systems.
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of our findings, we have provided comprehensive details of our exper- imental setup. The proposed architectural modifications are described thoroughly, with detailed architectural diagram and conceptual overview in Figure All hyperparameters and training proce- dures, including the specific recipes for pre- pre-training and supervised fine-tuning on the RULER and BabiLong benchmarks, available, and our evaluation setup described Section Furthermore, we commit to releasing source code upon de-anonymization to allow for full verification of our results and support future research in this area.
thoroughly documented in Appendix [ The datasets used are publicly
REFERENCES
Joshua Ainslic, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. Etc: Encoding long and structured inputs transformers, 2020. URL http:///aaxvv..oo/bbssssss 200p///0040088883333
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long- dong-document transformer. arXiv preprint arXiv:2004.051 .00000....55,,,... 2020.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Ho- race He, Connor Leahy, Kyle McDonell, Jason Phang. Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. GPT- NeoX-20B: An open-source autoregressive language model. In Proceedings of the ACL Work- shop on Challenges & Perspectives in Creating Large Language Models, 2022. URL https: wwrrxiv.org/bba.. 2204 06745
Aleksandar Botev, Soham De, Samuel L Smith, Anushan Fernando, George-Cristian Muraru, Ruba Haroun, al. Recurrentgemma: Moving past transformers for efficient open language models, 2024. URL httphhttss//axiiiorrrt hwrgg//b//20 ab2/2404 2404. 07839
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL https://aara wttpor/g///bb////////e 14165
Ricardo Buitrago and Albert Gu. Understanding and improving length generalization recurrent models. In Forty-second International Conference on Machine Learning, 2025. URL https:
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation, 2023. URL htttps///arxiv http:////axxxoo.oaabsssss
2306. 15595

---

Published as a conference paper at ICLR 2026
Rewon Child, Scott Gray, Alee Radford, and llya Sutskever. Generating long sequences with sparse 18509
transformers, 2019. httl:ttps::/axx..orr///////
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transforee-xl:: Attentive language models beyond fixed-length context. In Proceedings of the 57th annual meeting of the association for computational linguistics, pp. 1978-2988, 2019.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference Machine Learning, 2020.
Soham De, Samuel L. Smith, Anushan Fermand,, Aleksandar Botev, George Cristian Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Des- jardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, and Caglar Gulcehre. Griffin: Mixing gated linear recurrences with local attention for efficient lan- guage models, 2024. URL http:: wttp./wwwwwrr///////..... 19427
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. 2020. URL btttp:://arxiv.org/aa 2101.00027
Yizhao Gao, Zhichen Zeng, Ya Du, Shijie Cao, Peiyuan Zhou, Jiaxing Qi, Junjie Lai, Hayden Kwok-Hay So, Ting Cao, Fan Yang, and Mao Yang. Secrattention: Self-distileed attention gating for efficient long-context prefilling. In The Thirty- Anmual Conference on Neural Information Processing Systems, 2025. URL https://openrevieee wtwpeerreiieee.eeeooommmmiiiiooo 8y
Paolo Glorioso, Quentin Anthony, Yury Tokpanov, Anna Golubeva, Vasudev Shyam, James Whit- tington, Jonathan Pilault, and Beren Millidge. The zamba2 suite: Technical report, 2024a. URL
https://arxiv. www.oog/g/bbaoo ab///4b/211111 11242
Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. Zamba: A compact 7b ssm hybrid model, 2024b. URL (htps:
//arxiv.org// www.orggg//////eeeeee.. 2405. 16712
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First conference on language modeling, 2024.
Daya Guo, Canwen Xu, Nan Duan, Jian Yin, and Julian McAuley. Longcoder: A long-range pre- trained language model for code completion. In International Conference on Machine Learning. pp. 18099-12109) PMLR, 2023.
Haowen Hou, Zhiyi Huang, Kaifeng Tan, Rongchang Lu, and Fei Richard Yu, Rwkv-x: A linear complexity hybrid language model, 2025.
21463
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. RULER: What's the real context size of your long-context language models? In First Conference on Language Modeling, 2024. URL httpss///openreiieweeeeeeeeeeeee kloBbc76Sy
Xiang Hu, Jiaqi Leng, Jun Zhao, Kewei Tu, and Wei Wu. Hardware-aligned hierarchical sparse attention for efficient long-term memory access. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025a. URL htttp:////penreiii..eete//rroroe
dIHSZTx9Lu
Xiang Hu, Zhihao Teng, Jun Zhao, Wei Wu, and Kewei Tu. Efficient gengghh-geeeallzaalle attention via causal retrieval for long-context language modeling. In Forty-second International Conference on Machine Learning, 2025b.
net/
Xiang Hu, Zhanchao Zhou, Ruiqi Liang, Zehuan Li, Wei Wu, and Jianguo Li. Every token counts: Generalizing 16m ultra-long context in large language models, 2025c. URL https://axxi..
org/abs/2511. 223.b./39911.19
Jamba Team. Jamba-1.5 Hybrid transformer mamba models at scale, 2024. URL httph::tp//// //arxiv. org/abs/24g0) at/aab//248... 2225) 12570
12

---

Published as a conference paper at ICLR 2026
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer, 2020.
URL httpss://axxi www.:/orgg/a///eeeeeeee 2001.44411
Yury Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. Babilong: Testing the limits of Ilms with long context Advances in Neural Information Processing Systems, 37:106519-10655 2024.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, and Yoav Shoham. Jamba: A hybrid transformer-mant language model, 2024. URL htttt:://axxiii www.g////////aae 19887
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, 12:157-173, 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- ence on Learning Reprexentations, 2019.
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang. Zhiqi Huang, Huan Yuan, Suting Xu, Xinran Xu, Guokun Lai, Yanru Chen, Huabin Zheng, Junjie Yan, Jianlin Su, Yuxin Wu, Neo Zhang, Zhilin Yang. Xinyu Zhou, Mingxing Zhang. and Jiezhong Qiu. Moba: Mixture of block attention for long-context Ilms, 2025. URLhhtpss///arxt hwwp:/iwworg/bb///eeeee
13189
Yi Lu, Jing Nathan Yan, Songlin Yang, Justin T. Chiu, Siyu Ren, Fei Yuan, Wenting Zhao, Zhiyong Wu, and Alexander M. Rush. A controlled study on long context extension and generalization in 12181
Ilms, 2024. URL httptttpss//aa..oogggroeeeee/e/e
MiniMax. Minimax-01: Scaling foundation models with lightning attention, 2025. URL https:
08313
Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. In Workshop on Efficient Systems for Foundation Models @ ICML2023, 2023. wttp:///t//////eeoooe//eee/ee/ee
OpenAl. Gpt-4 technical report, 2024. URL
Zhuoshi Pan, Qianhui Wu, Huiqiang Jiang, Xufang Luo, Hao Cheng, Dongsheng Li, Yuqing Yang, Yew Lin, H. Vicky Zhao, Lili Qiu, and Jianfeng Gao. Secom: On memory construction and retrieval for personalized conversational agents. In The Thirteenth International Conference Learning Representations, 2025. URL
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong. Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Jie Zhu. Rwkv: Reinventing rmns for the transformer era, 2023. URL https://arrat htt.://rgg////////e//o/oo/ssss//
abs/ 2305. 13048
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. YaRN: Efficient context win- dow extension of large language models. In The Twelfth International Conference on Learning Representations, 2024. URL
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for se-
quence modeling. Advances in Neural Information Processing Systems, 36:33202-33222 2023.
13

---

Published as a conference paper at ICLR 2026
Liliang Ren, Yang Liu, Yadong Lu, Chen Liang, Weizhu Chen, al. Samba: Simple hybrid state space models for efficient unlimited context language modeling. In The Thirteenth International Conference on Learning Representations, 2025.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content- sparse attention with routing transformers. Transactions of the Association for Computational Linguis- fics, 9:53-68, 2021.
Karthik Soman, Andrew Langdon, Catalina Villouta. Chinmay Agrawal, Lashaw Salta, Braian Pee- toom, Gianmarco Bellucci, and Orion Buske. Zebra-llama: A context-aware large language model for democratizing rare disease knowledge, 2024. URL http:ps//aaatt..rgg/aa////////// 02657
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma. Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models, 2023. URL
https://arxiv arg/abs/230/ .86221
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang. Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. In The eighth Annual Conference on Neural Information Processing Systems, 2024. URL
https://openreviiw. hatppperreiieeeettooouuuoiiiiie nett/foruum1d==550x5555666
Ilya Sutskever. Training recurvent neural networks. PhD thesis, University of Toronto, 2013.
Shawn Tan, Songlin Yang, Aaron Courville, Rameswar Panda, and Yikang Shen. Scaling stick- breaking attention: An efficient implementation and in-depth study. In The Thirteenth Interna- fional Conference on Learning Representations, 2025. URL http:///ppenreieeeeee/eoooo
8d=r8J3DSD5kF
Tencent Hunyuan Team. Hunyuan-turbon turboss: Advancing large language models through mamba- transformer synergy and adaptive chain-of-thought, 2025.
15431
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, et al. Llama 2: Open foundation and fine-tuned chat models, 2023. Ittps://////rrxi /rgg/bb//3// 09288
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023. URL https://arxiv.. org/abs/1706. 03762
Junxiong Wang, Daniele Paliotta, Avner May, Alexander Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models. Advances in Neural Information Processing Systems, 37:62432-62457, 2024.
Jason Weston, Antoine Bordes, Sumit Chopra, Alexander M. Rush, Bart van Merriënboer, Armand Joulin, and Tomas Mikolov. Towards ai-complete question answering: A set of prerequisite toy tasks, 2015.
wtt.k//.org/bb///02.000 056988
Ronald J. Williams and Jing Peng. An efficient gradient-based algorithm for on-line training of
recurrent network trajectories. Neural Computation, 2(4):490-501, 1990. doi: 10.1 1622nc2oo 1990.2.4.490.
Yuhuai Wu, Markus N. Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers, 2022. URL httpht:///s:/xii .orllloooorgoooooo.o.o htrgg/bbb/20.... 08913
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference Learning Rep- resentations, 2024. URLlhttps. htttps:://ppeeeeiieeeee wwwpreerrrieeeneeeeerooeeuee h
Songlin Yang, Bailin Wang. Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- ers with the delta rule over sequence length. Advances in neural information processing systems, 37:1 19.115991-115522 2024.
14

---

Published as a conference paper at ICLR 2026
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, Wenfeng Liang, and Wangding Zeng. Native sparse attention: Hardware-aligned and natively trainable sparse attention. In (C), pp. 23078-23097, 2025. URtht/ps://acc htattp:://aclanthooog..oooo anthoologyoorraooooooooooo org//255.
acl-long. 1126/
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-177977, 2020.
15

---

Published as a conference paper at ICLR 2026
APPENDIX
A EVALUATION ON GENERAL CAPABILITIES AT SCALE
critical concern for specialized long-context architectures is whether the modifications introduced the hierarchical sparse attention and chunk
to support extreme length sacoding-negativee impact performance on standard short-context tasks.
To empirically verify the scalability of our approach and ensure no performance regression occurs in general domains, we evaluate a scaled-up version our model, denoted as SWA-HSA-MoE.
Experimental Setup. We conduct controlled comparison between our method and standard Transformer baseline in the large-scale regime:
Model Scale: Both models are 8B-parameter Mixture-of-Experts (MoE) architectures with approximately 1B active parameters per token.
Training Data: Both models were pre-trained on 8 trillion tokens using identical data mixtures and training recipes.
Architecture: The Transformer-MoE baseline uses standard dense attention with MoE Feed-Forwanl Networks (FFNs). The SWA-HSA-MoE replaces the attention mechanism with our proposed Hierarchical Sparse Attention (HSA) framework (including the non- linear chunk encoder and bypassing residual path) while retaining the identical MoE FFN structure.
We refer readers to Hu et al. 2025c) for comprehensive details regarding the training infrastructure, MoE implementation, and dataset composition used in these experiments.
Results. We evaluate both models on a diverse set of standard short-context benchmarks covering knowledge, reasoning, math, and coding. As shown in Table SWA-HSA-MoE achieves competi- tive superior performance across the majority of tasks:
1. No Regression:
The overall average score of SWA-HSA-Mol (57.69) surpasses the stan- dard Traanfforme--Mo baseline (55.62), confirming that our long-context mechanisms do not compromise core language modeling capabilities.
2. Reasoning & Coding Gains: Notably, model shows significant improvements in com- plex reasoning tasks, including GSM8K (+0.61), MATH (+4.02), and coding benchmarks like MBPP+ (+5.56). This suggests that the hierarchical selection capability may offer benefits even in non-retrieval-hent contexts by efficiently filtering information.
Table 4: General Capabilities Evaluation (8B MoE). Comparison between standard Transformer-Mal Moo and proposed SWA-HSA-MoE. Both models have 8B total parameters (IB active) and are trained on 8T tokens. Our method demonstrates robust performance on standard short-text benchmarks, outperforming the baseline 6 out of tasks.
Benchmark
Category
Transformer-MoE
SWA-HSA-MoE (Ours)
MMLU
Knowledge & Reasoning
58.74
57.83
PIQA
Knowledge & Reasoning
77.48
78.84
BBH
Knowledge & Reasoning
50.34
51.70
GSM8K
Math Reasoning
66.41
67.02
MATH
Math Reasoning
37.96
41.98
HumanEval+
Coding
48.17
50.61
MBPP+
Coding
50.26
55.82
Average
Overall
55.62
57.69
These results serve as strong empirical evidence that the architectural principles proposed in this
16

---

Published as a conference paper at ICLR 2026
work Stee---ciiiicllyy the synergy between the chunk encoder and bypassing residual path scal- able and universally effective, extending context length significantly without sacrificing the model's general utility.
B ADDITIONAL RELATED WORK
In the main body of our paper, focus dynamic sparse attention methods specifically designed for length extrapolation. Here, we provide broader overview of the rich literature on sparse at- tention, which has historically focused on different goal: reducing the quadratic complexity of attention for long, but still in-distribution, sequences. We categorize these methods into static and dynamic approaches, and for each, we explain why their design principles do not typically lend themselves to training-free free length extrapolation.
B.1 STATIC SPARSE ATTENTION
Static sparse attention methods reduce computational complexity by restricting each token to a pre- defined, fixed subset of other tokens. These data-independent patterns are highly efficient but lack flexibility required for length extrapolation. Early work like Sparse Transformer (Child et al.] 2019) introduced fixed strided and dilated patterns, foundational concept that influenced many subsequent models. This idea evolved into more complex topologies that combine local windowed attention with some form of global connectivity. For instance, Longformer (Beltagy et al. 2020) and ETC (Ainslie et al. 2020) designated few "global" tokens to act as information hubs for the entire sequence. BigBird (Zaheer et al. 2020) generalized this by combining local, global, and ran- dom connections to approximate the small-world properties of a dense graph. These methods, while successfully extending the manageable context length for tasks like long-document processing, fundamentally tied to the positional relationships learned during training. Their fixed patterns and reliance on learned positional biases prevent them from generalizing sequence lengths orders of magnitude beyond what they were trained on.
B.2 DYNAMIC SPARSE ATTENTION
Unlike static methods, dynamic sparse attention mechanisms determine attention patterns adaptively based on the input content, aiming to approximate the expressiveness of full attention by focusing computation on relevant tokens. Early approaches often relied on heuristic-baser bbsed grouping or clus- tering. Reformer (Kitaev et al. 2020) pioneered the use of locality-sensitie
hashing (L.H)) to bucket tokens for attention, while the Routing Transformer (Roy et al.) (02T) used online k-means clustering. These methods were a significant step towards content-aware sparsity. Another line of work focused on memory, with Transformer-XL. (Dai et al. (2019) introducing a recurrent memory and Memorizing Transformers (Wu et al.) 2022 using kNN index for retrieval from an exter- nal long-term memory. While these methods are "dynamic", their selection mechanisms are still parameterized and optimized for the context lengths seen during training. The learned LSH func- tions, cluster centroids, or memory management policies are not inherently length-agnostic and thus do not extrapolate well. More recent hardware-aware methods like NSA (Yuan et al.) 2025) and MoBA (Lu et al.) (2055) have achieved significant practical speedups, but their design goal remains computational efficiency within known context distribution, not training-free generalization ex- treme lengths. Our work, in contrast, focuses specifically on learning retrieval policy that is robust drastic changes sequence length, distinct and complementary objective.
Learned Sparse Attention and KV Cache Sharing. Two concurrent lines of work share surface- level similarities with our approach. SeerAttention (Gao et al. 2025) uses learned block-level gat- ing to approximate dense attention efficiently, and YOCO (Snn et al. 2024) introduces decoder- decoder architecture with a shared global KV cache. However, both focus on efficiency within existing long-conter models rather than training-free free length generalization. Our work investigates a distinct goal: identifying the architectural components that enable robust extrapolation from short training contexts to extreme lengths.
17

---

Published as a conference paper at ICLR 2026
C TRAINING RECIPE
C.I PRE-TRAINING RECIPE
Our pre-training recipe was designed for our models of approximately 240M parameters and is summarized in Table 5 All models were pre-trained a tokenized version of the deduplicated Pile dataset (Gao et al, 2020) using the GPT-NeoX-20BB tokenizer (Black et al. 2022). The models trained for approximately 8.05 trillion tokens using global batch size of 1.05 million tokens. We used the Adam W optimizer (Loshchilov & Hutter 2019) with a cosine learning rate schedule, a peak learning rate of 10-3, and minimum learning rate of 2 10-1
C.2 SUPERVISED TUNING RECIPES
Following pre-training, separate fine-tuning runs were conducted for the RULER and BabiLong benchmarks. All SFT runs were initialized from their respective pre-trained checkpoints.
SFT for RULER. The SFT data was constructed dynamically. For each sample, one of the three RULER tasks (S-N, MQ-N, or VT) was randomly selected. The task's "needle" was then in- serted at random position within background context sampled from the Pile. The model was trained to generate the correct answer. This stage ran for approximately 1.0 trillion tokens with peak learning rate of X 10-4
SFT for BabiLong. In a similar fashion, the SFT data for BabiLong was created by dynamically embedding tasks from the original bAbl dataset (Tasks 1-20) into background context from the Pile. For each training sample, one of the 20 bAbl tasks was randomly chosen, and its factual sentences were embedded into a context. The model was then trained on the associated question- answering pair. The hyperparameters were identical to the RULER SFT, with the primary difference being slightly longer training schedule of approximately 1.07 trillion tokens.
Hyperparameter
Pre-training
SFT (RULER) SFT (BabiLong)
Common Parameters
Model Parameters
~ 240M AdamW W Cosine with 2% warmup 1.05 Million tokens
Optimizer
LR Schedule
Global Batch Size
Stage-Specific Parameters
Training Tokens
8.05 Trillion
1.0 Trillion
1.07 Trillion
Peak Learning Rate
10
2 10-
2 10TH
Minimum Learning Rate
2 x 10
4 x 10-5
4 x 10
Table 5: Key hyperparameters for Pre-training and SFT stages.
D HYBRID ARCHITECTURES WITH SSM
The recent success of hybrid architectures, which synergistically combine different components to balance performance, efficiency, and long-context reasoning. marks a significant trend in sequence modeling. This paradigm has been prominently validated in models built upon the Mamba archi- tecture, such as Jamba (Lieber al., 2024). 7amba/7amba2 (Glorioso et al. 2024b a), Samba (Ren) et al. 2025). and others (Tencent Hunyuan Team 2025 [Soman et al. [2024] Wang al. (024). The principle extends to hybrids combining attention with linear RNNs or other mechanisms, as seen in Griffin (De et al.] (0244). RWKV-X (Hou et al. 2025). RecurrentGemma (Botev al. 2024). and
large-scale MoE models (MiniMax 2025| Jamba Team 2024).
A notable model in this space is RAMba (Hu et al. (202aa), which integrates an SSM with HSA and demonstrates 64x length generalization on RULER. We hypothesize that our architectural principles further enhance the performance of such hybrid models. To test this, we apply our components
18

---

Published as a conference paper at ICLR 2026
Upper decoder
Upper decoder
Mamba
FFN
HSA
Mamba
HSA
Chunk Selection
Lower decoder
Chunk Selection
Mamba2
Lower decoder
Mamba2
RAMba
RAMba-MLP
Figure 5: RAMba Architectures, An MLP is inserted into upper decoder layer in RAMB-MBL as Mamba does not have explicit MLP layer.
Model
Enc Layers
CLS
Bypuss
S:N
MQN
VI
Arg.
S.N
MQ.N
VI
Avg.
S-N
MQ-N
VT
Avg.
Model Spees
cts Jek-lee4KK
ctty Jen-32K
CCV los- 128K
Mambo-2
99.54
15.96
90.80
65.45
1.48
1,49
0,40
1.12
0.000
0,000
0.00
0.30
Mamba wd NSA
35.62
34,01
0.00
15.15
5.05
RAMba foriginal)
RAMba
yos
78.20
84.85
yes
RAMEa-MLP
mu
yes
08
79.04
89. 70
85.82
100.00
45.45
75.76
RAMba MLP
yes
200
93332
$1.35
96.66
90.44
89.33
78.35
96.27
86.32
84.85
06.87
96.97
82.83
RAMIs-MMP (ours)
yes
yos
95.36
83.95
96.64
91.99
97.76
88.06
97.76
94.53
100.0
87.88
100.0
95.96
Models
Enc
CLS
Hypass
S:N
MQQ.N
VI
Aug.
S:N
MQ
VI
Avg.
Model Spees
ctx Icl-co--MMM
let-SM
RAMba coriginal)
2
00
100
53.66
51.49
43.56
42.90
6,533
24.75
22.77
18.15
RAMba
yes
57.41
52.48
32.67
47.53
17.83
22.77
17.82
19:47
yes
3.54
RAMba-MLP
yes
RAMPa-MLP
yes
80
RAMES-ME.P (owrs)
3es
yes
160.00
$4.00
Table 6: RULER results for the RAMba hybrid architecture. The results confirm that the de- sign principles identified for DRT generalize effectively to an SSM-based model. Our full RAMba configuration (w/ Encoder, CLS, Bypass) once again achieves the best performance, particularly extreme context lengths (2 1M), validating the critical and synergistic role of these components across different local architectures.
to the RAMba architecture. The original RAMba design does not include a distinct MLP block due to Mamba's block design, which is necessary for controlled study our Bypassing Residual Path, as shown in Figure 5 We therefore introduce an MLP after each HSA block, enabling to isolate the effects of our proposed encoder, CLS token, and bypass mechanism on a state-of-the-art SSM- based foundation. We systematically investigate RULER, detailed results in Table
EFFICIENCY ANALYSIS
In this section, we provide a detailed analysis of the computational efficiency of our proposed Hier- archical Sparse Attention (HSA) mechanism, including theoretical complexity bounds and empirical measurements of inference latency.
E.1 COMPLEXITY ANALYSIS
Consider model with HSA operating on a context of length L, with chunk size C, top-K selection, and hidden dimension d. We analyze the complexity of each component:
19

---

Published as a conference paper at ICLR 2026
Chunk Encoder. The chunk encoder processes each chunk of size C' independently. With L/C chunks total, the complexity is:
o ( (4Cd + O (L (4d 2Cd))
which scales linearly with context length L.
Top-K Chunk Selection. The top-K selection computes scores between the query and each chunk's landmark representation. For single decoding step, the complexity is:
o (
(8)
The cumulative complexity across all steps is:
o L d -0 (42
L3d
Hierarchical Sparse Attention. size C. For single step:
HSA computes attention over the selected K chunks, each of
O (K (4Cd + 2))
(10)
The cumulative complexity for all L steps is:
O ( K- K-CC& +2)
(11)
Discussion. Only the chunk selection component involves a quadratic term with respect to L. However, compared full attention's (I.2d) complexity, this term is reduced by factor of 1/C. Furthermore, chunk selection is performed only once in the entire model (at single layer), making its quadratic contribution minimal in practice.
E.2 INFERENCE EFFICIENCY
To empirically validate the efficiency of our proposed architectures, we measure inference latency for generating 100 tokens across varying prompt lengths. All experiments use a batch size of 16, and we report the time cost excluding the initial prefilling phase.
Table 7: Inference Time Comparison. Time cost (in seconds) for generating 100 tokens different prompt lengths. RAMba (Mamba-2 + HSA) maintains near-constant latency regardless of context length, while full-attention Transformers exhibit substantial degradation.
Model
4K +
16K
64K +
Transformer (Full Attention)
2.26
8.90
32.12
Mamba-2
2.92
2.82
2.84
RAMba (Mamba-2 + HSA)
3.14
3.05
2.76
As shown in Table 7 both Mamba-2 and RAMba exhibit nearly constant inference time across context lengths from 4K to 64K, whereas the full-attention Transformer's 's latency increases by over 14x. Notably, RAMba incurs only a marginal overhead compared to the base Mamba-2 model while providing the additional capability of lossless long-range memory retrieval through HSA. Similarly, SWA+HSA benefits from the linear scaling properties of sliding window attention, with HSA adding only acceptable computational overhead. These results demonstrate that our proposed architectures are highly efficient for practical deployment in long-context scenarios.
F LLM USAGE STATEMENT
In accordance with the disclosure made during the submission process, large language models (LLM) were utilized writing assistant to aid in polishing the text of this manuscript. The LLM's
20

---

Published as a conference paper at ICLR 2026
role was strictly limited to improving the clarity, conciseness, and narrative flow of author-written content. All scientific contributions, including the core research ideas, experimental design, and analysis of results, are the sole work of the human authors. The authors have reviewed and edited all text and take full responsibility for the final content of this paper.
21


Reviews:
[{"review_id": "9rmyRKRcLe", "reviewer": "Reviewer_Nj58", "summary": "The paper try to interpret why chunk-based sparse attention models, for examples DRT and RAMbe, tend to generalize well to longer contexts than they were trained on. The authors argue that these models implicitly approximate a property they call Random Context Access (RCA)—the ability to flexibly retrieve any relevant part of the past without full attention. They identify three key design choices that make this work. 1. A non-linear chunk encoder with a CLS token, which learns better chunk summaries than simple mean pooling. 2. Bypassing residual path, which helps integrate retrieved global information without it being drowned out by local residuals. 3. Training with small Top-K retrieval, which avoids train–test mismatch.\nThey build on a 240M-parameter DRT backbone trained at 4k tokens and show strong length extrapolation—up to 32M tokens—on RULER and BabiLong, outperforming full attention and other long-context baselines.", "questions": "see weakness section", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.The paper offers a clear conceptual story. They provide a nice way to unify many previous chunk-based designs with the proposed concept, and the proposed concept is simple and effective.\n2. The empirical results are strong: models trained on 4k contexts can extrapolate to tens of millions of tokens. \n3.The ablation study is especially thorough—it covers encoder depth, presence of CLS, bypass design, Top-K size, and training length, giving a convincing picture of what really matters.", "weaknesses": "1.The theory part is more intuition-driven than rigorous; the RCA argument isn’t formally analyzed.\n2. The experiments more rely on synthetic or controlled tasks. The experiment on realistic long-context applications is limited.\n3. The detailed efficiency analysis of the method is not discussed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "m9o8JmA6Pt", "reviewer": "Reviewer_diBz", "summary": "Paper summary\n\nThe paper studies why hierarchical, chunk-based sparse attention models can extrapolate to extreme context lengths and how to make them more reliable. Using a unified view of “dynamic chunkwise sparse attention,” it decomposes these systems into (i) how chunks are encoded and landmarks are formed for retrieval, (ii) how retrieved global information is fused back into the model, and (iii) how sparsity is enforced during training. A theoretical argument shows landmarks must approximate a chunk’s total attention mass, which requires a nonlinear chunk encoder rather than simple pooling. Empirically, models trained on 4K tokens generalize, without further tuning, to contexts up to 32M on RULER and to 8M on BabiLong. Figure 1 (page 3) diagrams the architecture; Figure 3 (page 6) shows BabiLong accuracy remaining high as length grows; Table 2 (page 7) reports strong RULER results at 32M; Table 3 (page 9) details sparsity ablations. \n\nKey contributions\n\n* Unified framework + ablations: Formalizes chunk processing as functions over hidden states and compares weighting schemes (NSA, softmax, stick‑breaking) and encoder designs (with/without CLS), isolating which parts drive extrapolation. \n* Theory for landmark design: Shows that effective landmarks must approximate a chunk’s total attention mass, motivating a nonlinear Chunk Encoder and a dedicated CLS token for learnable summarization instead of mean pooling. \n* Cross‑layer fusion mechanism: Introduces a Bypassing Residual Path so retrieved global context is integrated via the MLP rather than overwritten by the residual stream; diagnostics separate retrieval recall from integration precision. \n* State‑of‑the‑art training‑free length extrapolation: 4K‑trained models reach 32M tokens on RULER with graceful degradation (e.g., ~80% average at 32M in Table 2) and maintain high BabiLong accuracy up to 8M tokens (Figure 3). \n* Principled sparsity insights: Finds that selection sparsity during pre‑training (small Top‑K like 8), sufficient training length (≥4K), and per‑token retrieval are critical for extreme generalization (Table 3). \n* Generality to hybrids: Applying the same principles to RAMba (Mamba + HSA) yields consistent gains across long lengths (Appendix C, Table 5).", "questions": "1. Precisely delineate what is new.\n   The paper combines chunk encoders, a CLS landmark, and the bypass residual within a unified HSA view, but it is not crystal-clear what is truly novel relative to DRT, RAMba, and NSA beyond the specific encoder+CLS objective and the bypass wiring. Please provide a one‑page “delta table” contrasting your equations and learning signals with DRT’s retrieval and RAMba’s stick‑breaking setup, and quantify the marginal gains when retrofitting only your two advertised contributions (encoder+CLS and bypass) into prior baselines without other recipe changes. A tight novelty delineation would materially strengthen the paper. See the unified comparison in Table 1 (page 5), Eq. 6a/6b (page 5), and the baseline list in Section 5.1 (page 6). \n\n2. Landmark theory: can you give a formal guarantee or falsifiable test?\n   Eq. 4a/4b (page 5) proposes that the landmark score should track a chunk’s total attention mass. Could you provide either: (i) a bound on how well a finite encoder can approximate this mass under simple distributions (e.g., mixture of von Mises–Fisher keys), or (ii) a calibration plot showing correlation between your selection score and the true within‑chunk mass across lengths? A short, synthetic study or correlation analysis would make the claim more than heuristic. \n\n3. Bypassing residual: mechanism and gradients.\n   The bypass changes the residual addition point (Eq. 6a vs 6b, page 5), and Figure 4 suggests improved utilization. Can you show gradient‑flow or activation‑scale diagnostics demonstrating that the bypass prevents residual stream “overwrites”? A simple ablation with a learned gate or scalar on the bypass path and a rank‑1 sensitivity analysis would help attribute the gain to integration rather than retrieval side effects. Cite Figure 4 (page 8) and Eq. 6a/6b (page 5). \n\n4. Baseline fairness and apples‑to‑apples choices.\n   Landmark Attention is trained on shorter sequences to match FLOPs, whereas your models train at 4k; retrieval frequency and Top‑K also differ across methods. Please report a “length‑matched” Landmark baseline at 4k and harmonize retrieval frequency (per‑token) and Top‑K across baselines. Adding one strong positional baseline beyond YaRN (e.g., ALiBi/PI) at 4k would reduce the chance that your positional baseline is a soft target. The request is small at 32k/128k (not 32M). See Section 5.1 (page 6) and Table 2 (page 7). \n\n5. Compute and efficiency at long lengths.\n   The results emphasize accuracy up to 32M tokens, but the paper lacks wall‑clock, memory, and throughput numbers for chunk selection vs token attention. Please add a figure with end‑to‑end latency and peak memory versus length for Top‑K in {4, 8, 16}, chunk size 64, and the hardware used. Even a 32k→128k→1M sweep would make the practical story credible. Table 4 (page 17) shows large training/SFT budgets; efficiency data would help separate algorithmic benefits from sheer scale. \n\n6. Scale and backbone generality.\n   Most main results use ~240M parameters. Appendix C shows encouraging RAMba results, but still at small scale and with a RAMba‑MLP variant. Can you share any existing runs or partial results at ≥1B parameters, or at least trends when moving the chunking boundary earlier/later in depth in the Transformer backbone? Figure 1 (page 3) and Table 5 (page 18) are the relevant anchors. Even partial curves at 32k/128k would be informative.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Originality.\n\n* Introduces a clean formalization of *Random Context Access* as the capability a long‑context model actually needs, which sharpens the objective beyond “long perplexity is stable” and motivates the architectural choices thereafter (Section 3.1). \n* Provides a unified lens for *dynamic chunkwise sparse attention* that decomposes chunk selection, weighting, and intra‑chunk processing; the side‑by‑side comparison of weighting rules (NSA, softmax, stick‑breaking) in Eq. 1a–1c and the f/g functions for landmark vs KV formation is a helpful synthesis that organizes a messy literature (Section 3.3; Table 1 on page 5). \n* Offers a simple but convincing theoretical target for landmark design: landmarks should approximate a chunk’s total attention mass, which in turn argues for a nonlinear chunk encoder and a CLS‑based summarizer rather than linear pooling (Eq. 4a–4b on page 5). \n* Proposes a *Bypassing Residual Path* that changes how retrieved global information is fused, with a concrete formulation (Eq. 6b on page 5) and a diagnostic that separates retrieval recall from integration precision (Figure 4 on page 8). \n* Demonstrates transfer of the principles to an SSM hybrid (RAMba) by inserting an explicit MLP to test the bypass mechanism and encoder/CLS design in a different backbone (Appendix C; Table 5 on page 18). \n\nQuality.\n\n* Careful experimental design: all DRT variants hold parameter count (~240M) fixed and standardize key knobs (window 512, chunk size 64, Top‑K 8) so ablations isolate the proposed components rather than incidental capacity changes (Section 5.1). \n* Strong and transparent baselines: full‑attention with RoPE scaling (Llama‑YaRN), Landmark Attention, sliding‑window, stick‑breaking, and Mamba2 with State Passing are included; training FLOPs or receptive fields are matched where relevant (page 6). \n* Evaluations cover both retrieval‑only skills (RULER: S‑N, MQ‑N, VT) and retrieval‑plus‑reasoning (BabiLong), reducing the chance of overfitting to a single metric. The curves on BabiLong reach up to 8M tokens (Figure 3 on page 6), and RULER tables extend to 32M (Table 2 on page 7). \n* Diagnostics move beyond end accuracy: retrieval recall@Top‑K and *marginal integration precision* clarify that the bypass path mostly helps utilization while the encoder/CLS improves ranking prominence (Figure 4 on page 8). This is rare and genuinely informative. \n* Thoughtful sparsity study: clean ablations over training length, Top‑K, and tokens‑per‑retrieval show why selection sparsity during pre‑training and per‑token retrieval matter at extreme lengths (Table 3 on page 9). \n\nClarity.\n\n* The architectural flow is easy to follow: the DRT diagram highlights lower vs upper decoders and the chunking layer that produces landmarks and encoded chunks (Figure 1 on page 3). The two encoder variants are visually contrasted (Figure 2 on page 5). \n* Mathematical definitions are compact and targeted: Eq. 2 stitches chunk weights into standard attention; Eq. 4b states the landmark learning objective in one line; Eq. 6a–6b precisely defines the bypass change. The paper ties each equation to an experiment that tests it, which keeps the narrative crisp. \n* Results are presented with the right granularity: a single multi‑length RULER table makes degradation profiles legible (Table 2 on page 7), BabiLong’s averaged curve makes plateau/failure modes obvious (Figure 3 on page 6), and Appendix B lays out hyperparameters and training schedules in a reproducible table (Table 4 on page 17). \n\nSignificance.\n\n* Establishes training‑free extrapolation far beyond common practice: models trained at 4K hold up to 32M on RULER with graceful decay and maintain strong BabiLong accuracy to 8M, while RoPE‑style baselines collapse shortly after the training length (Table 2 on page 7; Figure 3 on page 6). This pushes the empirical frontier and sets a new reference point for “extreme” long‑context capability. \n* Distills actionable design rules that others can adopt: use a nonlinear chunk encoder with a CLS landmark, fuse via a bypass that lets the MLP mediate cross‑layer information, enforce selection sparsity during pre‑training, and retrieve per token. These prescriptions are validated again in a distinct hybrid (RAMba), which broadens the paper’s impact beyond a single architecture (Table 5 on page 18). \n* Clarifies what it means to “use the whole context” by centering Random Context Access and measuring retrieval prominence vs integration. That conceptual shift is likely to influence how future long‑context work frames objectives and ablations, not just how it reports numbers (Sections 3–5). \n\nOverall, the paper combines a tidy theoretical target, disciplined ablations, and record‑scale evidence to produce design guidance that is both novel and broadly useful to the long‑context community.", "weaknesses": "1. Positioning and novelty are under‑specified.\n   The paper’s core design elements (chunkwise retrieval with landmarks, per‑token retrieval, and hierarchical sparse attention) are largely inherited from DRT/RAMba/NSA; the main additions are the specific *bypass* residual and the non‑linear chunk encoder with a CLS landmark. The text acknowledges these antecedents but doesn’t clearly delimit what is genuinely new vs. what is a careful re‑combination. Tighten the novelty claim by: (i) explicitly contrasting your equations and learning signals with DRT’s GCA and RAMba’s stick‑breaking weighting; (ii) running a head‑to‑head where you *retrofit* prior baselines with your two contributions (encoder+CLS and bypass) without other recipe changes, isolating marginal gains. Table 5 partially does this for RAMba‑MLP; extend the same “surgery” to Landmark and DRT to make the incremental contribution undeniable. \n\n2. Theory is insightful but stops short of guarantees.\n   The “landmark should approximate a chunk’s total attention mass” target (Eq. 4a–4b) is a compelling heuristic, but it is not accompanied by identifiability or approximation bounds. As written, it is unclear when a finite‑depth encoder can satisfy Eq. 4b under distribution shift or adversarial distractors. Strengthen the section by: (i) providing an upper bound on the KL or total‑variation gap between α and α̂ (Eq. 3a vs 3b) under assumptions on within‑chunk key dispersion; (ii) proving that mean pooling fails in a controlled setting (e.g., mixture of von Mises–Fisher keys) and that a 2‑layer encoder with CLS can represent the log‑sum‑exp proxy; (iii) stress‑testing Eq. 4b with synthetic counter‑examples where per‑chunk keys are multi‑modal. Figure 2 and the argument around Eq. 4b are the right hooks; add formal statements and failure cases. \n\n3. Evaluation is narrow; external validity is uncertain.\n   The study leans heavily on RULER S‑N/MQ‑N/VT and BabiLong. These are excellent for retrieval diagnostics, but they are synthetic and cue‑rich; they do not stress discourse coherence, paraphrase invariance, or multi‑hop compositionality across heterogeneous documents. Broaden the evidence by adding: (i) long‑document QA with paraphrase and entity aliasing; (ii) multi‑file code navigation where identifiers are obfuscated; (iii) book‑length summarization with contradiction traps; and (iv) “lost‑in‑the‑middle”‑style probes that deliberately place relevant spans at low prior positions. Figure 3’s averaged BabiLong curve and Table 2’s RULER tables are strong; add at least one real‑world, non‑templated task to bolster claims of “utilize the entire context.” \n\n4. Comparative fairness has confounds.\n   Several baseline choices can be contested: Landmark is trained with shorter sequences to match FLOPs, which handicaps its retrieval policy; LLaMA‑YaRN is known to degrade beyond train length, so a single positional‑scaling baseline is a soft target. Since Section 5.1 already modifies DRT to use stick‑breaking for per‑token retrieval, some of the gains vs “DRT 0” could be due to weighting, not just encoder/CLS/bypass. Make the comparison harder to dispute by: (i) giving baselines the same retrieval frequency (per‑token) and Top‑K; (ii) matching *training sequence length* rather than FLOPs for Landmark; and (iii) adding a second strong positional baseline (e.g., an ALiBi/PI variant) trained with the same 4k context and SFT recipe. See the setup on page 6 and Table 2 on page 7. \n\n5. Compute and efficiency claims are underspecified.\n   The paper emphasizes accuracy at 1M–32M tokens but gives no end‑to‑end latency, memory, or throughput numbers for these lengths, and no complexity constants for chunk selection. Appendix Table 4 signals substantial training/SFT budgets (≈8.05T pre‑training tokens; ≈1T tokens per SFT), which can swamp architectural effects. Add: (i) wall‑clock and peak‑memory curves vs. sequence length for Top‑K∈{4,8,16} and chunk size 64 on a standard accelerator; (ii) ablate retrieval index cost (score computation on landmarks) separately from token‑level attention; (iii) a compute‑normalized comparison (accuracy vs. TFLOPs) at long lengths. This will make the “practical” story credible. \n\n6. Scale and backbone generality remain open.\n   All primary models are ~240M parameters. The transfer to RAMba is helpful, but results are still at small scale and with a custom RAMba‑MLP that inserts an MLP not present in the original block. Demonstrate robustness by: (i) repeating key ablations at ≥1B parameters; (ii) applying the encoder+CLS and bypass to a pure Transformer baseline with different local windows; and (iii) showing the trend holds when the chunking boundary is moved earlier/later in depth (Figure 1 shows a fixed mid‑point). Table 5 provides a good template; broaden it to other backbones and sizes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RKSdiPhfQc", "reviewer": "Reviewer_UbFu", "summary": "This paper analyzes chunk-based hierarchical sparse attention models for extreme length extrapolation, specifically calls out three design choices: a non-linear chunk encoder, a bypassing residual path, and high selection sparsity. By combining these elements, the authors achieve state-of-the-art training-free length extrapolation, generalizing models trained on a 4K window to 32M tokens on the RULER and BABILong benchmarks. The work also provides theoretical motivation for the Chunk Encoder and diagnostic analysis correlating retrieval accuracy with task performance.", "questions": "1. How do you initialize the CLS token, and does it matter? Have you experimented with multiple CLS tokens per chunk? \n2. Are there tasks where you find that your method underperforms full attention at shorter contexts, which would demonstrate a fundamental limitation of chunk-based retrieval?\n3. Is your method sensitive to the chunk size hyperparameter, and does this vary with context length or task? An ablation on this would be valuable.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper has a clear presentation and nicely unifies different sparse attention approaches under a single framework. \n2. This method achieves significant training-free length extrapolation gains and retains them consistently on RULER and BABILong, marking strong empirical evidence.\n3. The justification for the chunk encoder as approximating full attention scores is compelling.\n4. The result in Appendix C showing that this method generalizes to hybrid SSMs is exciting and empirically sound.", "weaknesses": "1. The theoretical analysis is informal, and is not derived — there is no formal analysis of when and why this approximation is reasonable to make, although it is intuitive and stated as such. \n2. It would be valuable to have an analysis of computational costs, and any impact on inference speed (e.g. TPS and TTFT statistics) through this method.\n3. Evaluation primarily focuses on mostly retrieval tasks and synthetic reasoning, some real-world tasks that are a bit more diverse (e.g. summarization over long-documents) would be valuable to show general utility. \n4. The model scale is limited, and it is unclear whether this extends to larger models; I am willing to overlook this given the results are quite compelling, but for practical utility, this is important to study.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uoRPnr8sAh", "reviewer": "Reviewer_BsZ4", "summary": "This paper explores a method combining chunk-wise attention and retrieval, enabling models trained on short texts to extrapolate to ultra-long contexts. The framework, an enhancement of DRT, investigates improvements like adding encoder for global KV representation, a CLS token for predicting landmark, and bypass residual connection. Experiments on RULER and babilong demonstrate that these incremental modifications yield significantly superior performance.", "questions": "1. The authors state that sliding-window attention cannot access information outside its window. Given this limitation, why is its performance on the babilong benchmark at 8M context nearly identical to its performance at 4K? Does this imply the task itself does not require genuine long-range dependencies?\n\n2. Do these long-context modifications harm the model's short-text capabilities? It would be beneficial to evaluate the model on standard short-text benchmarks (e.g., MMLU, Wikitext-2) to verify that the proposed changes do not cause performance regression in non-long-context scenarios.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The proposed strategies—such as adding an encoder for global KV representations, using a CLS token as a landmark, and incorporating a bypass residual connection—are simple, yet deliver significant and immediate gains on downstream tasks. This makes the approach highly practical.\n\n\n2. The paper provides detailed ablation studies that convincingly demonstrate the necessity of each component and the synergistic benefits of their combination.\n\n\n3. The model is evaluated on well-known long-text benchmarks (RULER, babilong). The choice of benchmarks is reasonable and effectively characterizes the model's long-context capabilities.\n\n4. The authors attempt to provide a rationale for their modifications, such as explaining how the encoder-processed CLS token landmark can better approximate the highly non-linear attention denominator.", "weaknesses": "Despite the strong empirical results, the paper's primary weakness lies in its insufficient theoretical justification for why the proposed modifications are effective. The framework is fundamentally similar to DRT, yet the analysis of the new components is cursory.\n\n1. Ambiguous Role of the Encoder: The distinction in performance with and without the encoder is not substantial. It is unclear if the limited improvement stems from its purported \"decoupling\" function or simply from increased model capacity via extra parameters. The paper fails to clarify if this addition removes a specific bottleneck, or if adding the same parameter count elsewhere (e.g., in the base model) would yield similar gains.\n\n2. Unclear Mechanism of the Bypass Connection: The large performance boost from the simple bypass residual connection is surprising. The authors' explanation—that \"the Standard Sequential Path may disrupt the information flow\"—is vague and unconvincing. The paper lacks a sufficient explanation for the mechanism by which this simple \"re-wiring\" dramatically improves retrieval performance.\n\nOverall, the theoretical arguments for these modifications feel brief and tenuous, which undermines the reliability of the empirical findings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)