Short Window Attention Enables Long-Term Memorization
Authors: Loïc Cabannes
Maximilian Beck
Maria Lomeli
Gergely Szilvasy
Matthijs Douze
Jade Copet
Pierre-Emmanuel Mazare
Gabriel Synnaeve
Herve Jegou
Venue: ICLR 2026 Poster (2026)
Keywords: hybrids
xLSTM
SWA
memory
attention
long-context
architecture-design
LLM
stochastic
Primary area: foundation or frontier models, including LLMs
TL;DR: We evidence the counter-intuitive fact that, in hybrid architectures, shorter sliding windows lead to better length-extrapolation. We introduce a training procedure that allows the hybrids to perform well on both short and long-context tasks.

Abstract: Recent works show that hybrid architectures combining sliding window softmax attention layers with linear recurrent neural network (RNN) layers outperform both of these architectures taken separately. However, the impact of the window length and the interplay between softmax attention and linear RNN layers remain under-studied. In this work, we introduce SWAX, a hybrid architecture consisting of sliding-window attention and xLSTM linear RNN layers. 

A counter-intuitive finding with SWAX is that larger sliding windows do not improve the long-context performance. In fact, short window attention encourages the model to better train the long-term memory of the xLSTM, by relying less on the softmax attention mechanism for long context-retrieval. 

The issue with small sliding windows is that they are detrimental for short-context tasks, which could be solved with information from moderately larger sliding windows otherwise. Therefore, we train SWAX by stochastically changing the sliding window size, forcing the model to leverage both a longer context window and the xLSTM memory. SWAX trained with stochastic window sizes significantly outperforms regular window attention both on short and long-context problems.

Published as a conference paper at ICLR 2026
SHORT WINDOW ATTENTION ENABLES LONG-TERM MEMORIZATION
Loie Cabannes!.. Maximilian Beck Gergely Szilvasy', Matthiis Douze', Maria Lomeli', Jade Copet', Pierre-Emmanuel Mazaré', Gabriel Synnaeve', Hervé Jégou'
Meta FAIR "Johannes Kepler University Linz 3 Ecole Normale Supérieure Paris Saclay, Paris Cité University, Inria Paris
"Work done during an internship at Meta FAIR
ABSTRACT
Recent works show that hybrid architectures combining sliding window softmax attention layers with linear recurrent neural network (RNN) layers outperform both of these architectures taken separately. However, the impact of the window length and the interplay between softmax attention and linear RNN layers remain under-studied. In this work, we introduce SWAX, hybrid architecture consisting Blding--innoon attention and xLSTM linear RNN layers.
A counter-intuitive finding with SWAX is that larger sliding windows do not im- prove the long-context performance. fact, short window attention encourages the model to better train the long-term memory of the xLSTM, by relying less on the softmax attention mechanism for long context-retrieval
The issue with small sliding windows is that they detrimental for short-context tasks, which could be solved with information from moderately larger sliding win- dows otherwise. Therefore, window size, forcing the model to leverage both longer context window and the XLSTM memory. SWAX trained with stochastic window sizes significantly out- performs regular window attention both on short and long-context problems.
train SWAX by stochastically changing the sliding
SWAX:stochastic
35
SWAX:126
stoxsstochattic
60
SWAX:2048
Long Context Performance
SWAXI512
50
accuracy
40
30
20
SWAX-2048
10
SWA:2048
15
39.0
39.5
40.0
40.5
41.0
4096
8192 16384 32768 65536 131072 sequence length
Short Context Performance
bench- Figure/ Short (average marks) vs long context performance for 1.4B XLSTM, SWA (sliding window attention) and SWAX with different sliding window sizes.
Figure 2: RULER Needle-In-A-HHann Haysttka racy 1.4B SWAX model with fixed slid- ing window size of 2048 vs our method using a stochastic window size of 128/2048.

---

Published as a conference paper at ICLR 2026
INTRODUCTION
Memory is core concept in Neural Network Architectures (Zhong et al., 2025). Modern LLMs based on softmax attention have a working memory in the form of the key-value (KV) Cache and yield state-of-the-ar long-context performance. This working memory expands indefinitely as the sequence length grows, incurring linear growth in both compute and memory to generate each new token. With such an unbounded compute cost, current models become prohibitively expensive for in-context learning on long sequences such as codebases and long reasoning traces.
On the other hand, recurrent neural networks (RNNs) like State Space Models (SSMs) (Gu & Dao, 2024) or variants of linear attention (LA) (Katharopoulos et al., 2020) maintain and iteratively update a hidden state. Through input-dependent update rules, RNNs manage to decide whether to keep previous or add new information (Hochreiter & Schmidhuber, 1997; Chung al., 2014). In this way, the compute and memory cost is constant and independent of the sequence length. This allows models to learn test time from large sequence lengths and reason without a specific token limit. Recently, these linear RNNs have been generalized in the context of online learning (Liu et al., 2024; Sun al., 2025). However, the recall ability of linear RNNs remains inferior to that of Transformers (Fu et al., 2023). This shortfall has hindered their adoption in favor of global attention- based architectures, which remain the current state-of fate-of-thee of-the-ar art architecture for language and code models (Jain et al., 2024).
Recent works like those of De et al. (2024), Ren et al. (2025). Dong et al. (2024) and Arora et al. (2025) have aimed at combining the advantages of softmax attention and Linear Attention into hybrid architectures (Wang al., 2025). Following this line of research, in this paper, we study hybrid architectures, which combine linear RNNs and sliding window attention both components with fixed maximum state size and thus fixed compute cost per token.
In this context, we make the following contributions:
1. We investigate the impact of the sliding window length wide range of tasks, encom- passing validation perplexity, short-context reasoning, common sense benchmarks, and long-context modeling tasks;
2. We show that, contrary previous belief, for hybrids architecture interleaving sliding win- dow attention and linear RNNs, longer sliding windows actually hurl performance long- context retrieval tasks compared to using shorter windows;
3. We present a training strategy based on a stochastic window size that provides the long- context performance enabled by short windows and the short-context and reasoning perfor- mance of longer windows.
2 BACKGROUND
The attention mechanism handles sequences of key vector k, E Rddu and value vectors v, E Rdd. A fundamental perspective proposed by Katharopoulos et al. (2020) is that all forms of attention update matrix memory by adding to the outer product of key vector k, E RR and value vectors V, Rd This holds, provided that can apply vector mapping to each of these vectors. Then, to read from the memory, a query qr E Rd is compared to the previous keys using a similarity metric, usually the inner product (q. k). In order improve the accuracy of subsequent retrieval operations, a pre- pre-procsssing feature mapping @ is applied to the keys and queries. Defining the memory tensor as
H. vT. Rdaad
2, Σ o(k))
Rdd,
normalized read is performed as
y ((q.) H.
(is)) o(ks)) Vi
(2)
Φ(q.) Zt
(
2

---

Published as a conference paper at ICLR 2026
Linear attention. Equation (1) shows that if the kernel & is diite-dimensional mapping, then the saauue-mmapedd keys as well as the memory tensor are also diiee-dimensionll and can be mate- rialized and cached for future retrievals (H,, Zz) in constant memory:
H + H + (k,) v
Ze - - + o(k)).
(3)
All the keys and values are thus stored constant memory. The per-token read cost Oddys x de). Importantly, it does not depend on the sequence length S.
Softmax attention (SA). Katharopoulos et al. (2020) show that softmax attention can be seen as performing the attention operation defined in Equations I and 2, i.e., as writing outer products between keys and values in memory. In such a case, the keys and queries undergo an infinite- dimensional feature mapping induced by the exponential kernel in softmax attention. Compared to exponential feature map reduces the interference between
linear attention, an infinite-dimensionnn the stored keys and yields an improved retrieval accuracy. Another consequence is that the memory He cannot be materialized and cached. Instead, one needs to maintain all the previous keys and queries in memory in order to compute the exponential of the dot products of the keys and queries, also referred to as the "KV Cache". A well known issue inherent to the self-attention mechanism, is that the KV Cache size (and per token computation) increases linearly with the sequence length.
Gated linear attention. Another way to limit interference between keys in the sequence is to learn when to "forget" information and remove from the memory. This is the idea behind Gated Linear Attention (Yang et al., 2024), which improves stability and long-context performance through selective retention/forgetinng of the information. Let Ott, Gr, A E R2 be write-, read-, and decay- gates or, equivalently, broadcastable vectors. Gating is often implemented by learned affine maps and element-wiss sigmoids. The update and reading rule are as follows:
H& diag (x)) H + diag(a; Φ(k,) vT.
(4)
yr (dag((3)) ((q:)) H.
(5)
Gated Linear Attention as well as other modern RNNs remove the normalizing constant and, instead, rely normalizing layers such as LayerNorm (Ba al., 2016) and RMSNorm (Zhang & Sennrich, 2019) in the network to stabilize training (Beck et al., 2025a).
Sliding window attention (SWA). Softmax attention maintains all past (ki, 0z) pairs, producing linear growth in memory and compute with / due to the KV cache. Variants with a sliding window of size w restrict the attention process only the previous w tokens, changing the memory and time complexity per-token from ((S) to O(w) complexity (Beltagy et al., 2020). This theoretically allows SWA architectures to handle arbitrarily large input sequences. In practice the receptive field of the model is limited to (Iw) where I is the number of SWA layers in the model. Moreover, it is unlikely that the theoretical receptive field is fully utilized in practice (Xiao, 2025).
Hybrids between local attention and global softmax attention. Through multi-tun interactions (Gehring et al., 2025), tool-use or long Chain-Ora Of-ii--O-Thought reasoning (Wei et al., 2023; DeepSeek- AI, 2025), the length which models have to process has grown from few thousands of tokens to tens or hundreds of thousands of tokens. This motivates several recent works (OpenAI, 2025; Dong al., 2024; NVIDIA, 2025; Ren et al., 2025) to consider new architectures whose computational cost grow less rapidly relative to sequence length than that of global softmax attention, while still performing well on long-context tasks. One such type of architectures are hybrids, for which most layers have a fixed state size like sliding-windon or Linear Attention Layers, and the rest of the layers are global softmax attention layers. However, because those architectures still keep some global attention layers to remain competitive on long-context tasks, they also keep the (SS) scaling in state size and FL.OPs per token.
Hybrids between linear attention and sliding window softmax attention. Another kind of hy- bridization involves only component with a fixed state size, as considered by De et al. (2024) and Ren et al. (2025), who hybridize linear attention variants with sliding window attention. SWA paired with linear attention provides a natural split: the linear path maintains a compressed working mem- ory with an unlimited receptive field; the windowed softmax path offers high-fidelity local reasoning.
3

---

Published as a conference paper at ICLR 2026
Figure We compare 4 different types of archi- tectures, including 3 hybrid architectures:
Unbounded memory
Bounded memory
transformer
local/global
XLSTM
SWAX
(1) The transformer with vanilla self-attention (SA). Its complexity is prohibitive for long con- texts lengths.
FFN
FFN
FFN
FFN
SA
SA
mLSTM
MLSTM
(2) This is circumvented by replacing SA layers by sliding window attention (SWA) layers (Gemma Team, 2025; OpenAI, 2025).
FFN
FFN
FFN
FFN
SWA
mLSTM
SWA
(3) xLSTM (Beck et al., 2024) offers a memory with unbounded time horizon, albeit not as pre- cise SA for handling the recent context.
FFN
FFN
FFN
(4) SWAX is hybrid architecture that includes both SWA layers and long-term memories layers, implemented with mLSTM memory cells.
SA
SWA
mLSTM
SWA
Moreover, they demonstrate that, despite having fewer LA layers which are the only ones with an unlimited receptive field, the long-context performance of such hybrids is actually higher than that of a purely Linear Attention architecture. the size of the sliding window on validation perplexity. They found that longer windows yield better performance, making the choice of window size purely trade-off between performance and com- pute. However, they did not investigate the impact of the sliding window length on the long-context performance of the models.
particular, (De et al., 2024) investigated the impact of
3 HYBRID ARCHITECTURE DESIGN
In this work, we focus on hybrid architectures that alternate sliding window attention and linear RNNs. As candidate for the linear RNN component, we choose the xLSTM (Beck et al., 2024), as this architecture has been scaled to models having to 7B parameters and has shown strong per- formance wide variety of tasks. Importantly, fast and efficient Triton kernels are available (Beck et al., 2025b;a). xLSTM introduced two novel memory cells: the sLSTM with scalar memory and the mLSTM with matrix memory. However, on language tasks the mLSTM cell shows superior per- formance over the sLSTM, which has been abandoned in the latest 7B xLSTM model. We follow this choice and rely solely on the mLSTM cell in our hybrid architecture. Subsequently, we use xLSTM and mLSTM interchangeably to refer to the same architecture.
There exist many ways to hybridize these two components ( (Wang et al., 2025). In our case, we adopt the simple design of inter-layer hybridization, which alternates between SWA layers and layers of linear attention. For the sake simplicity, we adopt a 1:1 ratio meaning, that for every xLSTM layer there is one Sliding Window Attention layer. Figure illustrates how the layers in pure architectures and our SWA-xLSTM hybrid architecture. Most hybrids use window sizes between 128 in OpenAl (2025) and 2048 in De et al. (2024). We evaluate at intermediate lengths with sliding attention windows of lengths 128, 256, 512, 1024 and 2048. Finally, we evaluate stochastic training procedure that aims at improving length-ettrapolation. This training procedure stochastically chooses for each new batch either short or long window. In our experiments, we sample window size of either 128 or 2048 with probability 0.5 for each length. A similar strategy was proposed by Zhang & Bottou (2025) in the context of the Memory Mosaic architecture. However, in their case, the stochastic attention mask was applied to a long-term memory layer. In constrast, apply it to Sliding Window Attention layer with the explicit goal of reducing over- reliance on the SWA layers for long-context recall.
interleaved
EXPERIMENTS
4.1 EXPERIMENTAL SETUP
Our experiments focus language modeling, with emphasis on understanding the compromise between conrt- and long-context recall performance. In particular, we investigate the impact of the SWA window size on long-context retrieval. For this purpose, mainly rely the needle-
4

---

Published as a conference paper at ICLR 2026
SWAX:12
SWA:128
50
XLSTM
accuracy
40
30
20
10
4096
8192
16384
32768
65536
131072
sequence length
Figure 4: RULER needle-in-a-haystacc average performance on varying sequence lengths for 1.4B models. For SWA and SWAX we indicate the sliding window size after the colon.
in-a-haystack tasks of the RULER benchmark (Hsieh et al., 2024). A common practice for models using global attention is to pre-train them on shorter sequence lengths like 4k or 8k to reduce the cost of the attention operation, and then fine-tuned in a second training stage longer sequences to improve their long-context ability (Peng et al., 2023). In our case, we mainly focus on fixed- memory, fixed-compute architectures. Therefore, longer training sequences do not increase the required compute to attain a total training tokens target. We choose to train our models on a 16k sequence length from the start. Since are interested in the capabilities of the model after standard pre-training, we do not perform any task-specific fine-tuning on long-context tasks. Except stated otherwise, our experiments use model with 1.4 billion parameters. From our observations, is at this size that models become able to perform recall on sequence lengths in the tens of thousands of tokens. However, to validate our method of stochastic window size at larger scale, we also evaluate models at the 7B parameter scale.
The 1.4B models have 24 blocks and a model dimension of 2048 while the 7B models have 32 blocks and model dimension of 4096. In each block, the FFN is gated MLP Liu et al. (2021) with Silu activation (Elfwing et al., 2017). For the SWA layers of the hybrids, we use Rotary Positional Embedding (RoPE) (Su et al., 2023) with frequency 0 of 10000, and 16 attention heads. All models are trained on 150 billion tokens following a warmup-cosine learning rate schedule with peak learning rate of 3 10-4 and a minimum learning rate of 10-6 6 The batch size is 10t tokens. Our training data mix consists mostly of web-data and code. Since we are most interested in the performance of models long sequences, and code data has, average, 10 times longer documents than our web data, we report the validation perplexity on the code data subset of our data mix. For the MBPP (Austin al., 2021) and HumannvalPPLus Lus (Liu et al., 2023) pass @ 10 results, we use a sampling temperature of 0.8.
4.2 HYBRID vs. PURE ARCHITECTURES
We start investigation by reproducing the finding from (De et al., 2024) whose hybridization of Local Attention and Linear Attention variants improve performance across the board on both short-term reasoning and long-term recall tasks.
Long-context performance. Figure shows that the pure SWA architecture performs poorly on long-context recall. This is expected because of its limited receptive field of 128 24 = 3072 tokens. More importantly, confirms the counterintuitive finding from (De et al., 2024) that hybrids, despite having fewer global receptive field layers, outperform the pure variants in long-context recall. Intuitively, this explained by the fact that, although the SWA layers have a limited receptive field, the softmax feature mapping allows them to better model local dependencies than equivalent number of linear attention layers. Since most of the information necessary to predict the next token comes from local dependencies (Ruiz & Gu, 2025), fully linear attention model dedicates most of its layers to modeling local dependencies and few layers to model long-term dependencies.
On the other hand, in SWA-LA hybrids, local dependencies attention layers, which are more precise because of their direct access to recent history. As a con- sequence the linear attention layers specialize in modeling long-term dependencies, which the SWA
rather routed to the softmax local
5

---

Published as a conference paper at ICLR 2026
Model
Transformer xLSTM
SWA
SWAX
window length
n/a
128
256
512
1024
2048
FLOPP/token (x10%)
6.174
2.978
3.029
3.004
3,016
3.041
3.092
3.192
J
2,431
2,602
3,036
2.551
2,540
2.546
2.538
2.523
HEplus/pass@10
14.63
12.80
13.41
12.20
12.80
14.02
14.63
15.24
ARC-o
30.90
28.93
31.25
29.27
30.13
31.76
30.30
29.79
ARC-e
66.43
65.79
65.58
65.75
67.82
67.82
66.89
67.15
Hellaswag
46.18
44.68
44.45
45.37
45.37
45.47
45.47
45,91
MBPP/pass 1
31.40
22.60
23.80
24.20
26.80
28.40
29.20
28.80
NaturalQuestions
13.45
12.49
12.37
13.51
13.45
13,40
12.31
13.19
PIQA f
73.99
73.39
73.99
73.83
74.05
74.48
73.67
73.99
RACE.high
37.19
33.65
33,25
33.56
35.71
35.71
35.11
35,99
RACE.mid
50.63
46.59
45.54
46.52
49.09
48.89
48.40
49.30
SIQA f
42.48
40.23
41.61
42.53
42.02
41.71
41.91
41.71
TriviaQA f
30.11
27.96
28.15
28.95
29.59
28.26
28.98
29,95
Winogrande 1
61.41
58.01
62.12
62.04
59.91
58.33
59.27
59.51
average
41.57
38.93
39.63
39.81
40.56
40.69
40.51
40.88
Table 1: Validation perplexity and accuracy on short-context reasoning and commonsense tasks. All models
have 1.4B parameters. To compute the transformer FL.OPs use the training sequence length of 16384.
layers cannot model due to the limited window size. This highlights the impact that the window size can have on how much supervision the linear attention layers receive.
Short-context performance Table shows that the performance of hybrid models on short- context reasoning benchmarks higher than that of xLSTM and also slightly higher than that of a pure SWA architecture. This further highlights the fact that for short contexts, hybrid models leverage the high precision of the softmax sliding window attention layers. Hybrid models therefore take the strong short-conter performance of softmax attention, and the improved long-context recall ability of the Linear Attention layers.
4.3 IN SEARCH OF AN OPTIMAL WINDOW SIZE FOR HYBRIDS
Hereafter, establish that windows that are too long actually hinder the Linear Attention layers from learning to model long-term dependencies during training. We hypothesize that this degrada- tion is due under-training of the linear attention layers on the long-context recall task. To validate this hypothesis, we train SWAX with varying window sizes in {128, 256, 512, 1024, 2048) and test the models on both short context reasoning tasks and, more importantly, also on long-context recall tasks like RULER NIAH.
Short-context performance. De et al. (2024) experimented with different window sizes to find the optimal sliding window size. However, they only evaluated the different window sizes using the validation perplexity. Table shows that, indeed, the hybrid with the largest softmax attention window (SWAX:2048) has the best performance from the validation perplexity point of view.
However, raw perplexity not sufficient to accurately predict performance downstream tasks, and especially not in long-context modeling tasks (Fang et al., 2024). Thus, we also evaluate the impact of the window size on short-context reasoning and common sense benchmarks and on long-context retrieval tasks from the RULER benchmark. Table shows that on short-context reasoning bench- marks all window sizes except the shortest one, 128, give similar results, with the best performing hybrid being the one trained with the longest window size of 2048. The worse performance of the shortest window of size 128 is understandable as most prompts, even from those relatively short reasoning benchmarks, do not fit within a sliding window of 128 tokens.
Long-context performance. Figure 5 shows that once tested on longer sequences, the perfor- mance of the hybrid trained with a window size of 2048 drops the most. On the other hand, the SWAX models trained with shorter window sizes like 128, 256. and 512, maintain better perfor- mance even up to sequence lengths of 65k and 131k tokens. On the NIAH single task, SWAX models with shorter window have around 30% recall accuracy at 131k sequence length, while the

---

Published as a conference paper at ICLR 2026
miah_sngle__avergg
niah multikyy aeaemuttikkkyaeerageeee
100.0
25.0
80.0
accuracy
accuracy
20.0
60.0
40.0
10.0
20.0
5.0
0.0
4096
8192
16384
32768
65536
131072
4096
8192
16384
32768
65536
131072
sequence length
sequence length
niah_ miah_multiquery_aveeer average
niah_ httpmmuuliivluueaveraaa
80.0
80.0
accuracy
60.0
accuracy
40.0
40.0
20.0
20.0
4096
8192
16384
32768
65536
131072
4096
8192
16384
32768
65536
131072
sequence length
sequence length
SWAX:2048 SWAX:1024
SWAX:512 SWAX:256
SWAX:128
Figure 5: RULER NIAH subtasks accuracy for 1.4B SWAX models with different window sizes
SWAX with a window of 2048 has near 0% recall. Even the shortest sliding window of size 128, which consistently underperformed the longest ones in terms of PPL and short-context reasoning, significantly outperforms the model with the 2048 window length on all RULER NIAH tasks.
As shown in Figure 6, averaging over all sequence lengths and NIAH tasks, the SWAX model with a window size of 128 actually performs the best out of all the window sizes we tested. In particular, it outperforms the 2048 window size by 16 accuracy points. In other words, the SWAX with the shortest window has a recall 88.9% higher than the SWAX with the longest window. The most likely cause for this phenomenon is that during training, most of the dependencies model fall inside the 2048 tokens window.
Therefore, during pretraining. it was advantageous for the model with window of 2048 to the more precise softmax attention from the sliding window rather than having to rely on the less precise Linear Attention layers model most dependencies. However, once tested longer sequence length where the dependencies are outside of the window length, the model does not extrapolate since it never learned to rely on the Linear Attention layers do long-context modeling.
On the other hand, the models with shorter windows had to rely on the Linear Attention layers to propagate information since many dependencies fell outside of the sliding window. We give further evidence in Appendix A which further indicates that this is indeed the reason for the poor long- context performance of hybrids with long sliding windows.
All these results show that, contrary to previous belief, longer sliding windows do not always provide better performance and can even have a negative impact when extrapolating to tasks beyond the sliding window size and training sequence length. On the contrary, shorter window sizes push the Linear Attention layers, that have global receptive field, to receive more supervisory signal and specialize in long-context dependencies. Overall, shorter sliding windows allow the model to better extrapolate to tasks beyond the sliding window size and even far beyond the training sequence just way of reducing computational cost or
length. This also means that shorter windows maximize hardware utilization as was often thought to be the case, as in Arora et al. (2025) and De al. (2024).
7

---

Published as a conference paper at ICLR 2026
Table 2: Validation Perplexity and accuracy downstream tasks. Stochastic
models use w 2048 by default and switch w 128 with probability p = 0.5 at 1. 1.4B scale and with probability p = 0.75 at 7B scale.
model parameters
SWAX 1.4B
SWAX 7B
120
0.00
0.00
train window length window length
stochastic 2048
128
stochastic 2048
Pain-time window Length
128
2048
2048
128
2048
2048
256
0,00
val. PPL.
2.551
2.502
2.523
2.291
2.272
2.283
HEplus/pass@G 10
12.20
12,80
15.24
24.39
24.39
26.83
512
0.25
0:30
0.33
0.00
0.00
ARC-c f
29.27
30.82
29.79
40.77
40.86
41.55
ARC-e
65.75
68.71
67.15
75.05
74.46
74.80
1024
0.23
MBPP/pass @ 10
Hellaswag f
45.37
24.20
30.60
45.51
28.80
45.91
44.60
53.34
47.20
53.59
45,40
53.69
NaturalQuestions
13.51
12.47
13.19
22.21
23.45
23.04
2048
0.14
PIQA f
73.83
74.43
73.99
77.26
77.97
76.55
RACE.high f
33,56
35.91
35.99
37.65
38.99
39.88
128
256
512
1024
2048
RACE.mid
46.52
48.33
49.30
52.65
54.32
54.80
Test-time window length
SIQA f
42.53
41,86
41.71
43.55
44.22
42.78
TriviaQA t
28.95
28.99
29.95
46,20
47.15
46.80
Figure 6: average NIAH aecuracy of 1.4B SWAX models depending their train and test time window sizes.
Winogrande
62.04
59.27
59.51
67.88
67.64
65.75
average 1
39.81
40.81
40.88
48.80
49.52
49.32
4.4 DIFFERENT WINDOW SIZES AT TRAIN AND TEST TIME
We now explore training strategy allowing for a large window size test time, to have the best reasoning performance possible, while still being trained such that the Linear Attention layers for long-term dependencies and extrapolate to longer sequences.
Length extrapolation. As preliminary analysis, we first evaluate the performance of models when tested with different window size than the one used at training time. Figure 6 shows that, as expected, naively extending the window size beyond its training length results catastrophic collapse. This is a common phenomenon in softmax attention with RoPE which is used in the SWA layers (Peng et al., 2023). On the other hand, windows of size 1024 and less show little degradation when reducing their train-time window size by half. Overall, models need to be trained windows sizes to be able to use large windows during testing. At the same time, we cannot allow the models to over-rely on the long softmax attention windows since those do not perform well on very long-context tasks.
large
Stochastic window size. To solve this dilemma, we introduce training procedure that, throughout the training. stochastically alternates between large window size and small window size. Our hypothesis is that this will prevent the model from over-relying on the SWA layers, while still making the model capable of using the larger window size at test time.
Moreover, to validate our experiments at a larger scale, we also train 7B parameter models using the same experimental setup as for the 1.4B models. For the 1.4B experiments, at each new batch of data, we set the window size to 128 with probability 7B scale, we use slightly higher probability p 0.75 of sampling the short window. We provide an ablation for the value of in Appendix B. Finally, to foree the model to make better use of the larger test-time window of size 2048, we anneal the stochastic training procedure by not sampling the smaller window size anymore for the last 10% of training. We find that this short period of fixed windows at the end of the training significantly helps short-context performance without degrading long-context performance. We provide ablation of the annealing in Appendix B.
0.5 or leave the default window size of 2048. At
Table shows how training with a stochastic window size alternating between 128 and 2048 and annealing gives a short-context performance comparable or even better than training with fixed window size of 2048. In particular, at both 1.4B and 7B scales, stochastic training gives considerably better short-context performance than fixed-sizer window of 128. From validation perplexity perspective, the stochastic window size outperforms all models trained with fixed window sizes at all parameter scales. Therefore, training with stochastic window size and testing with longer
8

---

Published as a conference paper at ICLR 2026
1.4B, SWAX:126
70
78, SWAX:128
1.48. SWAKistochastic
78. SWAKistochastic
60
1.48, 5WAX:2018
60
78, 5WAX:2048
50
50
accuracy
accuracy
40
40
30
10
10
0
4096
8192
16384
32768
65536
131072
4096
8192
16338
32768
65536
131072
sequence length
sequence length
Figure 7: Average RULER NIAH accuracy of 1.4B SWAX models with different window sizes.
Figure 8: Average RULER NIAH accuracy 7B
SWAX models with different window sizes.
window yields better results on short-context tasks compared to a short window at both train and test time.
Compared to training with a fixed long sliding window of size 2048, stochastic training gives com- parable performance at the 1.4B scale and even slightly superior performance at the 7B scale on short-context reasoning tasks. This indicates that indeed, even if during training the model has seen the longer window size only part of the time, it is still able to take advantage of the longer window size for short-medium context reasoning tasks.
Long-context performance of the stochastic training. We evaluate the stochastically trained SWAX models with annealing on the last 10% of the training. This is to ascertain this strategy gives a performance on long-context tasks as good as short-window variants.
Figures 7 and 8 show that on RULER's long- congecontext recall tasks, the models trained with stochastic window size and annealing perform on par or better than the model trained with the short window of 128, and drastically better than the hybrid model trained with window of 2048 tokens. For instance, at 1.4B parameter scale depicted in Figure 7, the stochastic training SWAX model performs similarly to the context performance of models at 7B scale. At 7B scale, compared to using a fixed sliding window of size 2048, the stochastic training gives much better retrieval accuracy at all sequence lengths. Furthermore, just as at 1.4B scale, stochastic training gives similar or even superior long-context performance compared using short window throughout training. Overall, stochastic window size at training time maintains all the benefits of having short window for long-context recall, and most if not all of the benefits of longer window for short/medium-coni reasoning tasks.
SWAX. Figure 8 shows that stochastic training also improves the long-
Moreover, this result further confirms that the poor performance of hybrids trained on long windows not intrinsically due to long sliding window at test time. Instead, these results show that the poor length-gen gengtaagnncrliiitionnn of hybrids with long sliding windows is due to the training procedure. Indeed, if the model allowed to use the long sliding window throughout training, will over- rely on the more precise softmax attention of the sliding window even for recall tasks, which will not extrapolate to longer sequence lengths. It will under-utilize the Linear Attention Layers for the long-term recall task. On the contrary, if during training the model is not allowed to rely the long window, and is instead stochastically forced to use a shorter window, then the linear attention will be required the medium/long term recall tasks. Since, essentially, this amounts to stochastically reducing the capacity of the model to make it more robust, this can be seen as a form of dropout (Srivastava al., 2014) on the attention mechanism.
More long-context benchmarks. The Ruler tasks are artificial long context benchmarks, built to be sensible to model variants. We also evaluated SWAX on more realistic tasks from several families of benchmarks:

---

Published as a conference paper at ICLR 2026
Training SWA size
Table 3: Other long-conters benchmarks
Task family
128
stochastic
2048
SWAX architecture with 1.4B parameters. The results are averaged over all tasks within fam- ily. Top: xLSTM linear attention memory. hot- tom: Gated DeltaNet linear attention. For Gated DeltaNet also report the averaged performance of Ruler multiquery NIAH, averaged over se-
Longbench
10.67
11.79
11.14
Longbench2
22.99
27.26
22.65
Babilong
4.68
9.35
7.29
Gated DeltaNet
Ruler multiquery NIAH
40.58
40.90
28.80
quence sizes 4096, 8192 and 16384.
Longbench 9.14
9.94
11.42
Babilong
3.74
7.21
8.95
LongBench (Bai et al., 2024a) is a bilingual, multitask, and comprehensive benchmark evaluate the long context understanding capabilities of LLMs. It is composed of six categories and twenty one different tasks. It evaluates long-text application scenarios such single-document QA, multi-document QA, summarization, few-shot learning, synthetic tasks and code completion.
Babilong (Kuratov et al., 2024) This benchmark extends the Babi tasks (Weston et al., 2015), designed evaluate long-term memory and reasoning capabilities of LLMs. Babilong uses long input sequences, making it suitable for evaluating models with ad- vanced memory architectures. The main categories are: single book reasoning, memory and retrieval and temporal and spatial reasoning.
LongBench2 (Bai et al., 2024b) builds on the original LongBench, expanding both the scale and diversity of tasks to better stress-test LLMs for real-world long-context scenarios.
The results are summarized in Table 3 (bottom). They show that the tasks are difficult on aver- age for such small models. Stochastic training outperforms fixed-size training in some settings, but for others, the long context (2048) performs better. Table 5, Appendix we present the aver- age performance of the LongBench summarization and question-answering tasks, where the same conclusion holds.
4.5 EXPERIMENTS WITH GATED DELTANET
The Gated DeltaNet (Yang et al., 2025) a state-of-the-art hybrid architecture that combines gating and delta-update rule to update the linear attention memory. To assess whether our observations apply to other linear attention architectures, we replace the xLSTM component in the previous ex- periments with Gated DeltaNet layer (we call this combination SWAX/GDN). We did not change any training setting, and ran the same long-context benchmarks. Table 3 (top) shows that hybrid architectures with Gated DeltaNets also benefit from the stochastic training.
CONCLUSION
Through an empirical analysis of hybrid architectures, we evidence the counter-intuitive fact that shorter sliding windows lead to better length-extrapolltion retrieval tasks. Moreover, we intro- duce training procedure that stochastically changes the window size throughout training. This training procedure offers a strong performance on short-context tasks (enabled by longer sliding windows) and the length-extrapollatio ability Linear Attention layers, enabled by shorter win- dows at training time.
REFERENCES
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall- thrallthroouhhut tradeoff, 2025, URL https::/ ttte://arxiii org/ alos 24022402..2 18668.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski. David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, and Charles Sutton. Program synthesis with large language models, 2021. URL https: /arxiiv www.riivvoovooaeeeeeeee abs 08 07732.
10

---

Published as a conference paper at ICLR 2026
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016. URL pettps//arxerr wwtp/oggg///ooooooooooo...oo..oo wbs/1607.064.44
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. LongBench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meet- ing of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 3119-3137, Bangkok, Thailand, August 2024a. Association for Computational Linguistics. doi: 10.18653/ v v1/2024.acl-long.J 172. URL https: / a aawpaaannannooooo.aaaaaaaao org/ 2024.0022.22 acl-long......l 172.
Yushi Bai, Shangqing Tu, Jiajie Zhang, Hao Peng, Xiaozhi Wang, Xin Lv, Shulin Cao, Jiazheng Xu, Lei Hou, Yuxiao Dong. Jie Tang, and Juanzi Li. Longbench v2: Towards deeper understanding and reasoning on realistic long-context multitasks. arXiv preprint arXiv:2412. 2024b.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory, 2024. URL https: ptttps://arxiv.orggg/ abs 2405. 04517.
Maximilian Beck, Korbinian Pöppel, Phillip Lippe, and Sepp Hochreiter. Tiled Flash Linear Attention: More efficient linear and xlstm kernels. arXivv, 2503. 14376, 2025a. URL
https:/ //tp//aavvvv////////////e//aaav //taaaivvrggg//.e/.////....o.olt
wtt./bs/////oooooooeo/o.../..aer 2203..503.. 2503.14376.
Maximilian Beck, Korbinian Pöppel, Phillip Lippe. Richard Kurle, Patrick M. Blies, Gümter Klam- bauer, Sebastian Böck, and Sepp Hochreiter. xistm 7b: A recurrent Ilm for fast and efficient inference, 2025b. URL https: /tt////axivvv// org abs///b///////// 2503. 13427.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020. URL wwt.t:tps//t/...tm arxiv. http:///wwwww..waaee/eeesee 200.b.000000000000 05150.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piga: Reasoning about physical commonsense natural language. 2019. URL https: http::///axxvv.rrrr 11641.
aaaaaaiivvv..a orw. absalbbss/99 19911.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fo- tios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji. Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob Mc- Grew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating large language models trained code. 2021.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks sequence modeling, 2014. URL https: /atrsv/////vv../// orgg/ abbs/412. 3555.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018. URL https: arxii.iiv...........aaaaaaaaa.va http:/bs//00000 05457.
Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murary Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Des- jardins, Arnadd 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 https: ptttps://arxiie /ttpa/iivv./////////////////aata org/ absg///bb/aaae 2202224220222 19427.
DeepSeek-AL. Deepseck-rl: Incentivizing reasoning capability in Ilms via reinforcement learning, 2025. URL. http: htt.ttttt///xiiiiee wwtpxiiiwog/as////o/// 22501 12948.

---

Published as a conference paper at ICLR 2026
Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih- Yang Liu, Matthijs Van Keirsbilek, Min-Hung Chen, Yoshi Suhara, Yingyan Lin, Jan Kautz, and Pavlo Molchanov. Hymba: A hybrid-heee head architecture for small language models, 2024.
URL https:tp///// http:///xxvvvrr////e// absp//aa///// 2411. 13676.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Nigmoid-weighte linear units for neural network function approximation in reinforcement learning, 2017. URL https: /arxiaaxiiv/...////i//aaaaaaaail org/ abs/ 1702. 03118.
Lizhe Fang, Yifei Wang, Zhaoyang Liu, Chenheng Zhang, Stefanie Jegelka, Jinyang Gao, Bolin Ding, and Yisen Wang. What is wrong with perplexity for long-context language modeling?, 2024. URL https://arxiae
www.aas////..ooooooo///........e /0000./211..230 23771.
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models, 2023. URL https://
arxiv. org/ abs/ 22222212222222.
Jonas Gehring, Kunhao Zheng, Jade Copet, Vegard Mella, Quentin Carbonneaux, Taco Cohen, and Gabriel Synnaeve. Rlef: Grounding code lims in execution feedback with reinforcement learning, 2025. URL https: arrxiv. org/ abs / 2400.02089.
Google DeepMind Gemma Team. Gemma 3 technical report, 2025. URL https://arxix:
org/abs/ 2503 19786.
Albert Gu and Tri Dao. Mamba: linear-time sequence modeling with selective state spaces, 2024. 231 2. 00752.
URL hwtphttp////. http:htts///aaae http://wwwwvo//////ee.//
Sepp Hochreiter and Jürgen Schmidhuber: Long short-term memory. Neural Computation, 9:1735- 1780, 11 1997. doi: (91.1//2eecco.999.98...
Cheng-Ping Hsieh, Simeng Sun. Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models?, 2024. URL htttt:::: htapps:///rxxiy //twxvxxvvv//////////////////avl org /abs/////////aaa//a/////////// 2404. 0 06654.
Naman Jain, King Han, Alex Gu, Wen-Ding Li, Fanjia Yan, Tianjun Zhang, Sida Wang, Armando Solar-Lezama, Koushik Sen, and lon Stoica. Livecodebench: Holistic and contamination free evaluation of large language models for code. arXiv preprint, 2024.
Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Regina Barzilay and Min-Yen Kan (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp.
Vancouver, Canada, July 2017. Association for Com- putational Linguistics. doi: 10. (0.11663/11/PP1--1144 URL ht /att:ss:/aclaathoollo P17-1147/.
/
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns: Fast autoregressive transformers with linear attention, 2020. URL https://arxis /arxiv.. org/abs/ 2006 16236.
Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Igorevich Sorokin, Artyom Sorokin, and Mikhail Burtsev. BABILong: Testing the limits of LLMs with long context -n--ningg------hsstakk.. In The eight Conference on Neural Information Processing Sys- tems Datasets and Benchmarks Track, 2024. URL https: hwwp///eeeeeeeeeeeeeeeeee net forum?
id=u7m2CC84BQ.
Tom Kwiatkowski. Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming- Ming-WWi Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452 466, 2019. doi: 16.1162/tacl 00276. URL https:/ /aaa:///cllao antho wog..org/go.. 9- 1026/.
12

---

Published as a conference paper at ICLR 2026
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. Race: Large-scale reading comprehension dataset from examinations, 2017. URL 04683.
1104
Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, and Qiang Liu. Longhom: State space models are amortized online learners, 2024. URL https://arxiig http://////a////////// abs/2407. 14207.
Hanxiao Liu, Zihang Dai, David R. So, and Quoe Le. Pay attention to mlps, 2021. URL https: //arxiv. org. http//b//210005 08050.
Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. Is your code generated by chatGPT really correct? rigorous evaluation of large language models for code generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https:// openreview. httr://ooommnndqqxx_e
Team NVIDIA. Nemotron-h: A family of accurate and efficient hybrid mamba-transformer models, 2025. URL https: arxiv. http 03624.
OpenAl. gpt-oss-120b & gpt-oss-20b model card, 2025. URL http:://ersiv..rr//bbbb 2508. 10925.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models, 2023. URL wtttttsss///aiii/ http://aratttoorgaa/////3000 00000...000....
00071.
Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling, 2025. URL https: //arxiv..rr org/ abs/ 2406. 07522.
Ricardo Buitrago Ruiz and Albert Gu. Understanding and improving length generalization in recur- rent models, 2025. URL htwpt:tts/sceeccm //apivaaavvvvvv/av www.vooorgggooo. abtgg//b/////// 02782.
Keisuke Sakaguchi, Ronan Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale, 2019. URL https://arxiv org/abs/1907 10641.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Common- sense reasoning about social interactions, 2019. URL nttps::/arxiii wwwp:xiwi.orra/....aoooeoe haspgg/bb/19904 09728.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, ((56):1929-1953 2014. URL http:/ http://jmlz.ont /t////mmll..rr/.///e./////../o.e wwtpppppppp//ss///.../// sri srivastaval .
Jianlin Su. Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding, 2023. 2104. 09864.
Yu Sun, Xinhao Li. Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, Tatsunori Hashimoto, and Carlos Guestrin. Learning to (learn test time): Rnns with expressive hidden states, 2025. URL https: /ttp:::///axvvvo/r/// abs/2401 04 620.
Dustin Wang, Rui-J Zhu, Steven Abreu, Yong Shan, Taylor Kergan, Yuqi Pan, Yuhong Chou, Zheng Li, Ge Zhang, Wenhao Huang, and Jason Eshraghian. A systematic analysis of hybrid linear attention, 2025. URL https:ttp........... / /arxiv.. hatgg//bb/2500// 250.7066777.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia. Ed Chi, Quoc Le, and Denny Zhou. Chain-of-thougth prompting elicits reasoning in large language models, 2023. URL httpt: org. abs ab2/ 22001 111220111991 11903.
13

---

Published as a conference paper at ICLR 2026
Jason Weston, Antoine Bordes, Sumit Chopra, and Tomas Mikolov. Towards ai-complete question answering: A set prerequisite toy tasks. arXiv: Artificial Intelligence, 2015. URL https: //api. wea.pissetttttcccoooaa...oo 3
Guangxuan Xiao. Why stacking sliding windows can't see very far. httpp::/gganggant http//logsstackinng swa. htm., 2025.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear atten- tion transformers with hardware-efficient training, 2024. URL https://arxiv.... wwtpeiivorr///aaee abs/ 2312.06633
Songlin Yang. Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule, 2025. httpsttt////ax.tt org/ abs/ 2412. 06464.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence?, 2019. URL https: //arxiaa...e///////// .rg. aws/1905.............. 07830.
Biao Zhang and Rico Sennrich. Root mean square layer normalization, 2019. URL https:// arxi v. orgoorg//aa abs /1 1910 07467.
Jianyu Zhang and Léon Bottou. Memory mosaies scale. arXiv preprint arXivv2507..3225, 2025.
Shu Zhong, Mingyu Xu, Tenglong Ao, and Guang Shi. Understanding transformer from the per- spective of associative memory, 2025, URL https: //tp/aaaavvv///a///aa//aaa/aawal org/ abs/2505 11.05.99440 19488.
14

---

Published as a conference paper at ICLR 2026
SUPPLEMENTARY MATERIAL
A RESULTS OF PURE SWA MODELS
In section 4.3 hypothesize that the worse performance of SWAX models with long windows comes from the model utilizing the SWA layers instead of the xLSTM layers. To further confirm this hypothesis, we train 1.4B pure SWA model with a window size of 2048 and compare its performance to the SWAX model with the same window size. If the hypothesis that the SWAX model relies on the SWA layers for recall valid, then we expect its performance to be similar to that of a pure SWA architecture.
xLSTM
SWAX:128
accuracy
SWA-2048
SWAX 2048
10
4096
8192
16384
32768
055.36
131072
sequence length
Figure 9: Average accuracy of 1.4B parameter models over all RULER NIAH tasks.
Figure shows that indeed, the SWAX model trained with window of 2048 performs very similarly to the pure SWA architecture. On the other hand, the accuracy the SWAX model trained with window of 128 is dissimilar from that of the pure SWA model with a window of 128. This further evidences that low long-context performance of hybrid models with long windows comes from the model over-relying on the SWA layers for long-context recall instead using the xLSTM layers.
B STOCHASTIC SAMPLING PROBABILITY AND ANNEALING OF STOCHASTICITY
In this experiment we perform an ablation on the stochastic sampling probability p its schedule during training for the two model sizes 1.4B and 7B we consider in our study. Table 4 shows that, at 7B scale, a higher probability of sampling the small window during training is necessary to significantly improve short and long context performance compared the probability 0.5 which worked at 1.4B scale. Looking at the impact of annealing, i.c. using a stochastic window size for the first X% of training, at both 1.4B and 7B scale, annealing improves short-context performance compared to keeping the stochasticity until the end of training. At 7B scale, the annealed SWAX model even performs better on short-context than the SWAX model trained with a fixed window size of 2048. In terms of long-context performance, compared keeping the stochasticity until the end of training, annealing slightly degrades the long-context performance at 1.4B scale but keeps or even slightly improves long-context performance 7B. We believe that exploring different annealing procedures might provide even better short-context performance improvements while at the same time keeping good long-context performance.
C
BENCHMARKS
Code generation We use two benchmarks that evaluate the code generation capabilities AI models: HumanEval+ and MBPP.
The HumanEval+ (Liu et al., 2023) benchmark is extension of HumanEval (Chen al., 2021), which designed to evaluate the functional correctness code generated by Al models.
15

---

Published as a conference paper at ICLR 2026
model
xLSTM
SWAX
SWAX
parameters
7B
7B
1.4B
train-time window
NA
128
p=0.9
p=0.75 2048
p =0775
p=0.5
2048
p=0.5 2048
p =0.5
test-time window
NA
128
2048
2048
2048
2048
2048
niah single
61.20
62.43
58.99
63.36
63.20
55.27
53,46
62.61
61.59
niah multiquery
44.18
34.78
35.95
32.72
32.23
17.94
21.62
30.85
27.42
niah multikey
10.14
9.96
13.32
17.23
17.86
14.34
12.29
14.52
12.19
niah mullivalie
39.28
26.11
23.55
27.32
27.56
19.48
17.30
30.63
27.63
niah average
37.19
34.76
34.55
37.73
37.87
30.78
29.52
36.61
34.55
HEplus/passee 10
25.00
24.39
25.61
23.17
24.39
21,95
26.83
13.41
12.80
37.42
40.77
40.00
40.86
40.86
40.09
41.55
32.45
30.82
arc-e
73.32
75.05
74.76
74.50
74.46
74.63
74.80
67.23
28.71
hella
52.95
53.34
53.34
53.48
53.59
53.47
53.69
45.61
45.51
mbpp/pass@11
43.80
44.60
43.60
42.80
47.20
45.80
45.40
28.00
30.60
nq
22.12
22.21
23.12
23.16
23.45
22.58
23.04
12.95
12.47
piqa
76.93
77.26
76.71
78.07
77.97
77.31
76.55
74.32
74.43
race.high
37.02
37.65
37.71
38.74
38.99
38.74
39.88
34.82
35,91
race.mid
52.85
52.65
54.11
54.32
54.32
53,90
54.80
48.26
48.33
siqa
43.96
43.55
44.06
44.01
44.22
44.37
42.78
41.91
41.86
tqa
46.39
46,20
47.44
46.90
47.15
46.55
46,80
29.35
28.99
wino
66.06
67.88
68.51
67.32
67.64
66.77
65.75
59.20
59.27
short-average
48.15
48.80
49.08
48.95
49.52
48.85
49.32
40.62
40.81
Table 4: NIAH and downstream tasks accuracy for 7B models. p indicates the probability of using a window of 128 for batch, otherwise using window of 2048. 700% indicates annealing, i.e., only doing the stochastic window size first 90% the training and then using a fixed window size of 2048 for rest training. NIAH single and multikey results are the average overall all sub-tasks for each.
MBPP (Austin et al., 2021) designed evaluate the code generation abilities of AI models, particularly for Python programming tasks.
Common sense and general reasoning. We use benchmarks consisting of question-answer or multiple-choice questions designed to evaluate the common sense reasoning abilities of AI models, particularly the context of natural language understanding: HellaSWAG (Zellers et al., 2019), ARC (Clark et al., 2018), PIQA (Bisk et al., 2019). SIQA (Sap et al., 2019), Winogrande (Sakaguchi et al., 2019), NaturalQuestions (Kwiatkowski et al., 2019), RACE (Lai et al., 2017) and TQA (Joshi et al., 2017).
D LONGBENCH SUMMARISATION AND QA TASKS PERFORMANCE
We report the averages for all the summarization and question answering tasks, observe that training with stochastic SWA size outperforms the fixed-size training windows.
Task
SWAX:128
SWAX:2048
SWAX:stochastic
Longbench summarization tasks (average)
10.570
I1.152
12.429
Longbench QA tasks (average)
5.203
6.112
7.257
Table 5: Average performance for summarization and question-answering LongBench tasks
The LongBench summarization tasks are: bench multi news and longbench.vesuu
long-
The LongBench QA tasks, which include single-QA and multi-QA tasks, are: long- bench hothooo Longbench_2wkking longbench.musique, musique, longbench_duuraader, ddureader, longbench narrattveeaa longbench.qaspe longbench. hatpheiehcmttiifella..eeen and long- bench shhcc..hltiiieldqqa..h.
longbench.
16


Reviews:
[{"review_id": "B8lgNr14oZ", "reviewer": "Reviewer_4vk4", "summary": "The paper proposes **SWAX**, a hybrid that interleaves sliding-window softmax attention (SWA) and xLSTM layers. Core claims: (i) **shorter** SWA windows surprisingly yield **better long-context recall** because they force the RNN path to learn long-term memory; (ii) a **stochastic window** training schedule (occasionally using a small window such as 128, otherwise 2048, with an anneal in the last 10% of training) recovers short-context performance while preserving long-context extrapolation. Evaluations use a 1.4B and 7B model trained on 150B tokens at 16k sequence length; long-context performance focuses on RULER NIAH, with downstream short-context benchmarks reported.", "questions": "* Line 075: “**hybrids architecture**” → “**hybrid architectures**.” \n* Line 198: “In **constrast**” → “In **contrast**.” \n* Line 233:  “we also **train evaluate** models at the 7B parameter scale.” → “train **and** evaluate.” \n* Table/figure captions could be clearer (e.g., units/averaging in Fig. 6 heatmap; spell out metrics). \n* **Add broader long-context tasks** (e.g., LongBench-style QA/summarization/code comprehension) to test whether the NIAH gains translate to practical workloads. \n* **Tighter causal tests**: (i) equalize total compute across window settings; (ii) vary document-length distribution; (iii) report how gradients/attention mass shift between SWA and xLSTM under different windows; (iv) include ablations on the anneal schedule (you mention some in Appendix B—bring the most informative ones into the main paper).", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The stochastic window schedule (sampling 128 vs 2048; annealing off in the last 10% of training) is simple, inexpensive to try, and consistently improves length extrapolation without sacrificing short-context quality at both 1.4B and 7B scales. This is well demonstrated in Figs. 7–8 and Table 2.", "weaknesses": "1. **Limited and unbalanced evaluation of “long-context” ability.**\n   The central claim rests mostly on **RULER Needle-in-a-Haystack (NIAH)** accuracy; other long-context behaviors (QA, summarization, codebase-level reasoning, tool-use traces, multi-hop retrieval) are not evaluated. As a result, it’s unclear whether the observed gains generalize beyond synthetic token recall. The paper itself emphasizes RULER NIAH and heatmaps across sequence lengths but lacks a broader long-context suite. This weakens external validity of the main claim.  \n\n2. **Confounds in the “short vs long window” story and insufficient controls.**\n   The paper argues that **longer** windows under-train the RNN path and therefore **hurt** extrapolation; however, compute is not cleanly equalized across variants (Table 1 shows **FLOPs/token** rising with window size), and the training schedule/data mix may interact with window length (e.g., document length distribution skewed to code). Stronger controls (fixed total FLOPs, matched optimization schedules, varied data mixes, and tests at matched effective receptive fields) are needed to isolate causality. As is, the causal interpretation (under-training the RNN) is plausible but not conclusively established.  \n\n3. **Novelty is incremental relative to prior hybrid work; framing oversells.**\n   Alternating xLSTM and local attention is now common; the **new ingredient** here is mainly the **stochastic window schedule** and the empirical observation that small windows train the long-term path better. The paper acknowledges closely related hybrids and prior window-size analyses (e.g., De et al.), and even positions the method as akin to **dropout** on the attention mechanism. Without broader/better-controlled evaluations or theoretical support, the contribution feels like a useful training tip rather than a substantial architectural advance.", "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": "4shqrUL5ZH", "reviewer": "Reviewer_pAFQ", "summary": "The paper studies hybrid models that interleave sliding-window softmax attention (SWA) with linear RNN layers (xLSTM/mLSTM) in a 1:1 pattern. The core, counter-intuitive claim is that short SWA windows improve long-context recall because they push the model to train its recurrent memory instead of relying on local softmax; conversely, long windows hurt length extrapolation. The authors also propose stochastic window training—randomly switching between short and long windows during pretraining with an anneal near the end—to regain short-context quality without giving up the long-context gains. Empirically, they evaluate mainly on RULER “needle-in-a-haystack” tasks and a suite of short-context benchmarks;", "questions": "1) How many random seeds were run per configuration? Please report mean ± standard deviation across seeds for the key metrics in Table 1 and Figures 5–8 (and include the per-seed results in an appendix or supplementary file). If only a single seed was used for any result, please state this explicitly.\n\n2) Could you specify the full optimization setup: the optimizer used (e.g., AdamW), exact β₁/β₂ (and ε) values, weight-decay coefficient, gradient-clipping norm and threshold, all dropout rates?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) Clear, testable capacity-allocation hypothesis. Long windows allow SWA “absorb” the learning signal; short windows force the xLSTM to learn long-term memory. The paper probes this across window sizes and model scales.\n\n2) Simple, practical training recipe. Stochastic window training with a short end-of-training anneal; task-agnostic and effective at 1.4B and 7B.\n\n3) Comprehensive comparisons. Side-by-side results for pure (xLSTM, SWA) vs. hybrid models, compute-aware reporting (FLOPs/token; Table 1), and stringent length-extrapolation stress tests (RULER NIAH up to 131k).", "weaknesses": "1) The pretraining mix is described only at a high level (“mostly web and code”); perplexity is reported on a code subset, but sources, filters, licenses, and exact proportions are not disclosed. This reduces reproducibility and external validity.\n\n2) RULER NIAH—primarily a needle-retrieval test—is the centerpiece, but broader long-context reasoning is underrepresented (multi-hop QA over long documents, extended chain-of-thought, codebase-level tasks). The generality of the findings remains unclear. A stronger case would include evaluations across varied tasks, e.g., Babilong and long open-QA benchmarks (HotpotQA and MuSiQue in full-wiki mode, LongBench).\n\n3) The experiments are exclusively based on the xLSTM as the linear RNN component. A brief discussion or experiment on whether this phenomenon holds for other popular models, including neural memory (for example, Titans[1]), would significantly increase the generalizability and impact of the findings.\n\n[1] Ali Behrouz and Peilin Zhong and Vahab S. Mirrokni. Titans: Learning to Memorize at Test Time, arXiv:2501.00663", "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": "L17halcx5j", "reviewer": "Reviewer_SVZm", "summary": "The paper proposes SWAX, a hybrid model that alternates sliding‑window softmax attention (SWA) with xLSTM (linear RNN) layers, arguing that short attention windows during training force the model to use the RNN path as true long‑term memory. Empirically, for 1.4B and 7B models trained on 150B tokens at 16k sequence length, the authors find: (i) shorter windows (e.g., 128) yield better long‑context retrieval on RULER Needle‑in‑a‑Haystack (NIAH) than longer windows (e.g., 2048) when extrapolating up to 131k tokens, and (ii) a stochastic window schedule preserves long‑context gains while recovering short‑context quality at test time with a 2048 window.", "questions": "N/A", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Clear, counter‑intuitive empirical result: In hybrids, short windows train better long‑term memory than long windows. The heat‑map in Fig. 6 (p.8) plus the degradation curves in Fig. 5 (p.7) make the effect easy to verify.\n-  The stochastic window + brief annealing yields near‑best short‑context results while keeping long‑context robustness. This is a practical recipe others can adopt with minimal code churn.", "weaknesses": "- Long‑context evaluation is narrow: RULER NIAH is synthetic. The paper lacks real‑task long‑context evaluations.\n- On the short‑context suite, the Transformer baseline still leads on average. The paper should contextualize this gap with wall‑clock throughput and memory at inference.\n- Design space under‑explored: The hybrid is fixed to a 1:1 SWA:xLSTM interleave and xLSTM only. It would be more valuable if more results on other setup could be added.", "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": "IaF1WrGNUy", "reviewer": "Reviewer_bBjL", "summary": "This paper empirically evaluates how the window size of sliding window attention affects the long-context performance, when mixing sliding window attention with xLSTM linear RNN layers in a hybrid architecture. The authors find that short window attention encourages the model to better train xLSTM, leading to better long-context performance. However, short window attention may lead to worse short-context performance. The authors then propose a simple strategy to randomize the window size in training, achieving the best of both worlds.", "questions": "1. Could the authors comment on the generalizability of the observation to other RNN architectures?\n\n2. Could the authors comment on whether the results continue to hold if the models are further fine-tuned?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well-written, clear, and easy to follow. The paper provides an interesting and practically insightful observation that short window attention leads to better long-context performance in hybrid models mixing window attentions and RNN layers.\n\n2. The authors provide a simple adaption to train a model with good short-context and long-context performance.", "weaknesses": "1. The paper would be stronger if it evaluates more than one types of RNN layers to show that the findings are robust against different RNN architectures.\n\n2. The findings are based on models without fine-tuning with long-context data. It is unclear whether the results continue to hold if the models are further fine-tuned.", "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)