Understanding and Improving Length Generalization in Recurrent Models
Authors: Ricardo Buitrago
Albert Gu
Venue: ICML 2025 poster (2025)
Keywords: length generalization
length extrapolation
mamba
long context
language models
Primary area: deep_learning->sequential_models_time_series

Abstract: Recently, recurrent models such as state space models and linear attention have become popular due to their linear complexity in the sequence length. Thanks to their recurrent nature, in principle they can process arbitrarily long sequences, but their performance sometimes drops considerably beyond their training context lengths---i.e. they fail to length generalize.
In this work, we provide comprehensive empirical and theoretical analysis to support the \textit{unexplored states hypothesis}, which posits that models fail to length generalize when during training they are only exposed to a limited subset of the distribution of all \textit{attainable} states (i.e. states that would be attained if the recurrence was applied to long sequences).
Furthermore, we investigate simple training interventions that aim to increase the coverage of the states that the model is trained on, e.g. by initializing the state with Gaussian noise or with the final state of a different input sequence. With only 500 post-training steps ($\sim 0.1\%$ of the pre-training budget), these interventions enable length generalization for sequences that are orders of magnitude longer than the training context (e.g. $2k\longrightarrow 128k$) and show improved performance in long context tasks, thus presenting a simple and efficient way to enable robust length generalization in general recurrent models.

Understanding and Improving Length Generalization in Recurrent Models
Ricardo Buitrago Ruiz 12
Albert Gu
Abstract
tion to matching the performance of Transformers (Vaswani et al., 2017) across many tasks, the recurrent mechanism brings two benefits: the ability to efficiently process long sequences thanks to its linear complexity, and the capacity to easily process tokens beyond their training context by simply rolling out the state. Nevertheless, in practice these benefits are often unrealized, given that their performance can drop considerably when the sequence length exceeds their training context (Waleffe et al., 2024; Ben-Kish et al., 2024; Ye et al., 2025; Yuan et al., 2024). This naturally leads to two questions: (1) why do these models fail to length generalize? and (2) how can we efficiently enable length generalization across several recurrent models?
Recently, recurrent models such as state space models and linear attention have become popu- lar due to their linear complexity in the sequence length. Thanks to their recurrent nature, in prin- ciple they can process arbitrarily long sequences, but their performance sometimes drops consider- ably beyond their training context lengths-i.e. they fail to length generalize. In this work, we provide comprehensive empirical and theoretical analysis to support the unexplored states hypothe- sis, which posits that models fail to length gener- alize when during training they are only exposed to a limited subset of the distribution of all ar- tainable states (i.e, states that would be attained if the recurrence was applied long sequences). Furthermore, we investigate simple training in- terventions that aim to increase the coverage of the states that the model is trained on, e.g. by initializing the state with Gaussian noise or with the final state of a different input sequence. With only 500 post-training steps (~~ 0.1% of the pre- training budget), these interventions enable length generalization for sequences that are orders of magnitude longer than the training context (e.g. 2k + 128k) and show improved performance in long context tasks, thus presenting simple and efficient way to enable robust length generaliza- tion in general recurrent models.
Recently, some works have studied the length generalization of Mamba (Dao & Gu, 2024) and have proposed solutions such as forcing the model to forget previous context (Chen et al., 2024b) or skipping tokens in the state update to reduce the effective context of the processed sequence (Ye et al., 2025; Ben-Kish et al., 2024). However, these methods re- quire changing the internal mechanism of Mamba and might not be easily transferable to other architectures. Other works have linked length generalization to state capacity and over- fitting (Wang, 2024; Chen et al., 2024b), proposing training on longer sequences and with Truncated Backpropagation Through Time (TBTT) (Williams & Peng. 1990; Sutskever, 2013) as a way to enable length generalization. In this work, we reason about the distribution of states that the model is trained on to introduce a precise hypothesis that explains why recurrent models fail to length generalize. Moreover, we perform comprehensive interventions that elucidate on what distributions recurrent models need to be trained to enable length generalization.
1. Introduction
More concretely, we propose the unexplored states hypothe- sis, which suggests that models f'ail to length generalize when their recurrence applied long sequences pro- duces state distributions that have not been explored during training. We support this hypothesis through sev- eral experiments and additionally introduce Effective Re- membrance, a novel metric that measures the impact of previous parts of the context in the output of the model. We find out that models that fail to length generalize are dispro- portionately impacted by the initial tokens of the sequence, suggesting that these models overfit to states produced early in the sequence when they are trained on short contexts with
Recurrent architectures like state space models (Gu et al., 2022; Smith et al., 2023; Gu & Dao, 2023; Dao & Gu, 2024) and linear attention variants (Katharopoulos et al., 2020; Peng et al., 2023; Yang et al., 2024b;c) compress the pre- vious context of sequence into state, with each output depending on previous tokens only through the state. In addi-
'Carnegic Mellon University CCartesia AL. Correspondence to: Ricardo Buitrago Ruiz <ricardobuitrgooriiee comm".
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

Understanding and Improving Length Generalization in Recurrent Models
Mamba-1 Position-Wis Wise Perplexity
Mamba-2 Position-wss wise Perplexity
GLA Position-wise Perplexity
30
30
20-
Perplesity
Perplextty
20
19
102
10a
20
101
10a
16a
10
30
10'
1D1
10f
Sequence Position
Saquanco Postton
Sequerca Position (1)
130m
130m
(a) Mamba-1 Perplexity.
(b) Mamba-2 Perplexity.
(c) GLA Perplexity.
Mamba-1 EffRlom
Mamba-2 EffRam
GLA EffRem
10a
ETFamett
10-1
20a1
2000
4000
6000
2000
2000
4003
5000
5000
500
100G
1500
2000
Sequencce Position
Sequence Position
Sequence Position (1)
330m
110m
310m
Litke
(d) Mamba-1 EffRemr (1) (1' = 8192).
(e) Mamba-2 EffReer (1) (T' = 8192).
(1) GLA EfReem (T' 2048).
Figure 1. (Top) Perplexity as function of token position on the Pile validation dataset (Gao et al., 2020) for the official Mamba-1 and Mamba-2 checkpoints trained with context T = 2048, well as for Gated Linear Attention (GLA) models trained with context T = 512. In dashed lines, we show the same models post-trained with State Passing (SP), which is an intervention that initializes the state with the final state of different sequence (see Section 4.1). State Passing is a simple technique that enables length generalization across several recurrent architectures. Mamba-2 and GLA are post-trained for 500 steps and Mamba-1 I is post-trained for 1000 steps. A similar plot for RWKV-v6 (Peng al., 2024a) is shown in Figure 10. (Bottom) Effective Remembrance for recurrent models and their State Passing post-trained counterparts. Effective Remembrance at time i roughly measures the impact of the tokens at positions [0. t) on the output of the model at later position T, with 0 indicating no impact (no "effective remembrance" of tokens (0, 4)) and indicating maximal impact (see Section 3.4 for a precise definition). The baseline models are disproportionately affected by tokens that are very far away in the past, indicating that they are not correctly handling the recent context. State Passing fixes this behavior.
a zero-initialized fixed state.
terventions enable length extrapolation in long context tasks such as BABILong (Kuratov al., 2024), passkey retrieval (Mohtashami & Jaggi, 2023) and synthetic copying (Jelassi et al., 2024); and also change the curves of the Effective Remembrance metric so that the model no longer overfits to the states of early parts of the sequence.
Based on these findings, we explore training interventions that modify the initial state to expose the models to wider range of state distributions (which is equivalent to starting from some given initial context). We investigate techniques such as TBTT and additionally propose new ones that are directly motivated by our hypothesis, such as initializing the state with some type of Gaussian noise. Through com- prehensive comparison of these interventions, we conclude that the key length generalization is training on initial states that are similar to the states that the model attains when processing long sequences- in particular, training on long sequences (directly or indirectly through TBTT) is not always necessary. In our experiments, with as little as 500 post-training steps (~ 0.1% of the pre-training budget) we enable length generalization from 2k training contexts to sequences of length 1 128k validation. Furthermore, the in-
Our study argues that length generalization in recurrent models is expected to be readily achievable through sim- ple training interventions. This simplifies the process of comparing new recurrent architectures by allowing re- searchers to primarily focus on their in-length performance, which we consider particularly significant in light of the re- cent proliferation of newly proposed recurrent architectures (Peng et al., 2023; Sun et al., 2023; Beck et al., 2024; Katsch, 2023; Ma et al., 2024; Liu et al., 2024; De et al., 2024; Arora et al., 2024; Sun et al., 2024; Yang et al., 2024b;c;a). As whole, contributions bring new theoretical insights
2

---

Understanding and Improving Length Generalization in Recurrent Models
on the behavior of stateful models and simple empirical techniques to improve recurrent models large.
3.1. Definition of Length Generalization
First, we provide a concrete definition for position-wise perplexity-iee the average perplexity that the model achieves at each position in the sequence.
2. Preliminaries
2.1. Notation
Definition 3.1 (Position-wise Perplexity). Let a(-lc) be the next token probabilities of an autoregressive sequen- tial model given context c, and let D be distribution over sequences. We define the position -wise perplexity at position i the average perplexity that the autoregressive model achieves at position t:
Given a sequence x = (ro, F1, (T, we will use Tss to refer to the subset of the sequence (I', Ta+1s..,, zz). Addi- tionally, for any sequence x we will use xx to denote the multiplication of elements of x, i.e. xx o Zi. For convenience, if we will define z 1.
Position-wise Perplexity()) = Fen-plexp(((zz//ooo
2.2. Recurrent Models
Position-wise perplexities can easily be computed in a dataset simply by averaging the typical perplexity by se- quence position. These values serve to define length gener- alization:
The core of both Mamba and Mamba-2 is state space model (SSM), which is transformation of 1-dimensional sequence E RºT' y I through an implicit latent state h E (TT,N), In its general form, it can be written as:
Definition 3.2 (Length Generalization). Let M be au- toregressive sequential model trained with context Tuain [0, 7] an interval with 7 > Tuain. and D a distribution over sequences. Let p" be the minimum position-wis perplex- ity that .M achieves on [0, Train] and let t* be the position where is achieved. Then, M said to length generalize in [0, T] if the position-wist perplexities for t ; are smaller equal than p".
hr- A,h1-1 + Be
(1a)
yt =2=C( hr
(Ib)
where A, E R(N.N), B, E R(M.)), and C. E R(N.)) The output y/ depends on the past history To: only through the state h4- Thus, in autoregressive models h, maintains a compressed representation of the past Fo.: which is useful to predict the next token. Equation la is valid for t E [o, T]. For / = 0, it is also necessary to define h-1. which we denote by initial state. In the standard implementations of Mamba, the state h-1 is initialized with zeros h-1 = O. Note that this is equivalent to starting to predict the sequence yo:r without any previous context.
3.2. Short Training Contexts Impede Length Generalization
In this subsection we show the results for models trained from scratch with different training contexts. In Figure 2 it can be seen that when the training context length is too short, the models' perplexities diverge for positions after the context length. Additionally, the larger the model, the longer the context needed to length generalize. Based on these results, we conclude that the longer the training context, the better the length generalization. Details on the training recipe and model configurations are given in Section B.
Throughout this work, we use the term "recurrent models" to refer to architectures that accept a formulation like the one presented in Equation which includes most modern recurrent models like Linear Attention (Katharopoulos et al., 2020), RWKV (Peng et al., 2023), Retnet (Sun et al., 2023) and Gated Linear Attention (Yang et al., 2024b), to name a few. The difference between the architectures mostly resides in the parametrization of A, B, and C,; and in how they are obtained from the inputs. We note that some modern recur- rent architectures do not accept a formulation like Equation 1 (e.g. XLSTM (Beck et al., 2024)); we hypothesize that they would exhibit similar behavior but we leave them outside the scope of this work. We refer to Yang et al. (2024c) for an overview of modern recurrent models.
3.3. Training for Many Tokens Impedes Length Generalization
In the previous subsection we trained models for several times what Chinchilla scaling laws dictate, which is typical to maximize their performance. Now, we will study whether the failure to length generalize also occurs when training for fewer tokens. Figure shows the position wise perplexity of Mamba-2 model at different checkpoints during training, showing that the model needs to be trained for at least 7.5x Chinchilla scaling laws to fail to length generalize. Thus, extended training hinders length generalization.
3. Analyzing the Length Generalization Failure of Recurrent Models
In this section, we identify training conditions under which models fail to length generalize and explain this behavior through our proposed unexplored states hypothesis.
We note that, in addition to short contexts and extended training, other hyperparameters also influence length gener- alization in our experiments. For example, in Section E we
3

---

Understanding and Improving Length Generalization in Recurrent Models
Mambe2 45m trained with context
Mambe2 trained with context T
Mamba2 125m trained with cortext T
trained connext
T-256
T-3024
T=107d
T-256
T=512
T-2048
1-204d
T=512
T-1024
T-4056
T=6096
T-1024
T=2048
T=8192
T-2048
T-4056
Perplenks
Perstesits
121 102
121
30
121
102 3>
301
307
327
309
321
102
37'
Sequence
Sequence Prsition O2
Sequerice Position
Sequence Position 0t
(a) Mamba-2 45m
(b) Mamba-2 85m
Mamba-2 125m
(d) GLA 70m
Figure 2. Position-wise perplexities for Mamba-2 (Dao & Gu, 2024) and GLA (Yang et al., 2024b) trained from scratch with different context lengths T on the Pilc (Gao et al., 2020). The longer the training context, the better the length generalization. The 45m model trained for 2B.5B tokens (25x Chinchilla laws), the 70m and 85m trained for 34B tokens (20x Chinchilla Laws), and the 125 model is trained for 25B tokens (10x Chinchilla Laws).
observe that having an extended learning rate warm-up pe- riod hinders length generalization. This indicates that there a complex relationship between length generalization and training recipes. In Section 4 we will show that there are inexpensive training interventions that enable length gener- alization, removing the need to carefully tune the training recipe to achieve generalization.
straightforward way to eliminate previous parts of the se- quence from the state). Consequently, their failure to length generalize arises because their handling of recent context compromised by having already processed earlier parts of the sequence.
To analyze this phenomenon further, we propose compar- ing the outputs of the model when it is given the full con- text, and the output when it is given only a recent context. Concretely, given a context To:2', an autoregressive model predicts vector of probabilities for the next token position, q(- ((-(Fo:T) E I where |V] is the vocabulary size. We propose Effective Remembrance as measure understand how the output of different partial contexts Tu. differs from the output of the full context.
Mamba-2 85m at intermediate checkpoints
40
Perpexity
5x
10x
20
12
17 3x
Definition 3.3 (Effective Remembrance). Given autore- gressive model which outputs probabilities q(- context) over vocabulary V. an input sequence No.T and distance be- tween probability distributions d, we define the Effective Remembrance for E [0, T] as:
102
10a
104
Sequence Postion
Figure 3. Position-wise perplexities of a Mamba-2 85m model trained for different number of tokens with context of 1024 on the Pile (Gao et al., 2020). 2.5x means that model trained 2.5 times what Chinchilla scaling laws dictate for that model size. Thus, the checkpoints correspond to range between 4.25B and 34B tokens. In this case, the failure to length generalize occurs after training for more than 7.5x Chinchilla laws.
EffRemy d(g(-1x0.)) q (rer))
(2)
Unless otherwise stated, we will use total variation' as a distance between distributions:
TV
1 Σ lq(x) - q(x')l ZEV
3.4. Effective Remembrance: a Metric to Understand How Sequence Models Process Context
(V((()),
(3)
Sequence models would perform reasonably well on long sequences if they only used sliding context window of length equal to the training context to predict each output, as this would be equivalent to processing sequences of the same length as those seen during training. This is the mech- anism of some architectures like Sliding Window Attention (Beltagy et al., 2020). However, recurrent models use the the full context for their predictions (indeed, there is not a
Additionally, we will present the values of Effective Re- membrance averaged over a dataset. We note that Effective Remembrance is applicable to all autoregressive sequence models, not just recurrent ones.
II Section G we perform an ablation on the choice of the distance metric for Effective Remembrance and conclude that the shape of the curve EffRemy (t) is very similar when using total variation, the Jensen-Shanme distance or cosine similarity.

---

Understanding and Improving Length Generalization in Recurrent Models
Effective Remembrance can be understood as how much the past tokens 10---1 influence the output at time T. If EfReem (t) = 0, this means that the predictions using FeT and using Zo:T are the same, meaning that the model does not "effectively remember" any of the past tokens Tu:--1. Conversely, if (t) is high, the model is substantially influenced by the tokens I0 since remov- ing them from the context changes the prediction signifi- cantly. Naturally, for distribution such text we would expect (f) R 0 for i < T (tokens that are very far away from 7 barely have an effect on the prediction), and EfRRem (t) + 1 when t ++++. T'?
Standard deviation of full state
255
2.0
Std Dev
1.5
Base
State Passing
2000 4000
6000
0000
Figure 4. Standard deviation of the full state of the Mamba-2 130m official checkpoints versus the sequence position t (he in the no- tation of Section 2.2). The standard deviation taken across all elements of the state in all layers. The Mamba-2 2 130m post tsaintr with State Passing (Section 4.1) produces states whose standard deviation do not significantly change after the training context T = 20418. In contrast, the official Mamba-2 checkpoint reaches a standard deviation almost twice as large at position 8192 than the one at position = 2048.
Figure 1 shows the Effective Remembrance of several re- current models, using 7 four times larger than the training context (8192 for Mamba and 2048 for GLA). Models that fail to length generalize are substantially affected by far away tokens. Indeed, the 1. 1.3b Mamba-2 model has high values of (f for small f. indicating that its outputs strongly depend on the initial tokens (in other words, if the initial tokens were not included for the prediction, the out- put would be very different). Therefore, the failure to length generalize is linked the model depending too much on the initial tokens and overfitting the states that arise when processing the early parts of the sequence without previous context (i.e. with a fixed zero-initialized state).
ization based on the distribution of the state. At time t, recur- rent models take an input Vg and the recurrent state ht to out- put the prediction yr. That is, we can write yr = ( ht) for some function f. Therefore, the key to length general- ization is understanding which distributions of states and inputs the function f has good performance.
3.5. The Distribution of States Changes Over Time
We can gain further insights into why recurrent models fail to length generalize by studying statistics of the distribution of the state over time. Concretely, if we have a distribution of sequences of arbitrary length X, we can study the distri- bution of the state at time t, De := (Xo.c), where (z) refers to rolling out the state recurrence (e.g. Equation la)
For many tasks, we can assume that the distribution of inputs at time t does not change much after the initial tokens (for example, in text, we expect that the marginal distribution of tokens at two different positions and t' should not be too different). Thus, can focus on the distribution of the state. In particular, we are interested in the distribution of attainable states, that is, the ones that the model would produce when processing a sequence.
on x. In Figure 4 we show that the standard deviation of the state of Mamba-2 130m increases over time. In particular, after the training context the model encounters states with a distribution different from training.
Again, we denote by D. := hx(Xoz) the distribution of the state at time t. When a model is trained with context length of T, it is exposed to state distributions Dr for 0 S < T, however it is never exposed to distributions De for t' > T. If the distribution Dr is different to the ones seen during training, the model is not guaranteed to have good performance. Based on this insight, we present our hypothesis for why models fail to length generalize:
We already note that the problem with length generalization is not necessarily related to having large norms in the state; rather, is due to having states whose distribution change after the training context: the State Passing intervention (Section 4. achieves length generalization but has higher standard deviation in the state at almost all positions (Figure 1). In Section C we show that this change in distribution is due to specific specific layers and heads of the state.
Unexplored states hypothesis: Recurrent models fail to length generalize when they are trained only on subset of all attainable state distributions-i. i.e. on subset of the stutes that would be attained if the state recurrence was rolled out indefinitely. When trained for long enough, the model overfits to this subset and performs poorly on long sequences because it encounters unexplored state distribu- tions.
3.6. Unexplored States Hypothesis
Inspired by our previous empirical findings, in this subsec- tion we arrive at a hypothesis to understand length general-
NNote that because total variation is bounded between 0 and 1, 0 < (t) < 1.
5

---

Understanding and Improving Length Generalization in Recurrent Models
This hypothesis explains the previous behavior we have observed. When the training context T is short, the model overfits to limited number of distributions D; for 0 < t < T, which can be different to the distribution of the states attained after rolling out the recurrence on long sequences. Conversely, if the training context T' is increased, the model trained on a wider range on distributions and is more likely to perform well on the distribution of attainable states.
model is very different from an IID Gaussian with same mean and variance in all layers, especially for large model sizes. Thus, even though this non-zero initialization exposes the model to more state distributions. they are not realistic and do not cover the distribution of attainable states, which explains why it does not fix length generalization. More details on the post-training procedure are given in Section B.2.
An ablation the impact of is shown in Section F. We also note that the random initial state is applied only in training, while a zero-initialized state is used for validation. Introducing noise in validation significantly degrades the model's performance on the first tokens.
4. Training Interventions to Enable Length Generalization
In the previous section, we verified that training with long sequences enables length generalization, which is in agree- ment with the unexplored states hypothesis. In this section, we further support this hypothesis by showing that length generalization is also achieved when the model is exposed to initial states that are similar to attainable states-in in other words, training with long sequences is sufficient but not necessary for generalization. To do so, we propose four different interventions to initialize the state and study the generalization performance of recurrent models post-trained with those interventions.
4.3. Intervention 2: Fitted Noise
For the second intervention, we initialize the state from a distribution that resembles attainable states more closely. To do so, we sample from a Gaussian distribution with mean and standard deviation fitted to the final states seen during training. More concretely, during training for each layer I and head l we dynamically compute the mean and variance of the final states using a moving average with 3 = 0.1:
1000 + (1 B) M)Mnn (10) + 3p
4.1. SSM Equations For Non-Zero Initial State
(5)
201) + (1 3) Variance (1) + 30
First of all, we will derive the equations for SSM whose initial state is not zero initialized. Assuming h-1 of O, we can unroll Equation la to have:
(6)
where the Mean and Variance are taken across the head dimension and state expansion dimension N. Then, at each training step the values of the initial state of the layer and head h are sampled independently from Gaussian with the updated parameters N (1(0)) 2(1))
14 =
A 1 Badg + 1
(4a)
==
31- = C Σ
A + C. h- ==0
(4b)
We apply this intervention by post-training the official Mamba-2 models for 100 steps (~ 0.02% of the pre-training budget) and present the results in Figure 5 (dashed orange line). This approach significantly enhances the generaliza- tion of the 370m and 780m models but still fails to correct the 1.3b model. We hypothesize that this is because indepen- dent sampling does not generate realistic states for the 1.3b model: as the model size increases, the states likely become more complex, with stronger dependencies between values. As in the previous intervention, we note that Gaussian state initialization is used only during training.
Thus, an SSM that is initialized with h-1 F 0 differs from zero-initialized SSM only in additive factors of h- and CT A A in the final state and output, respectively. As we mentioned in Section 2.2, most recurrent models can be formulated using Equation 1, and thus Equation is valid for them. An explicit derivation of the formula given in Section H.
4.2. Intervention 1: Random Noise
For the first intervention, we initialize all the values of the state by sampling independently from a Gaussian (0, d²). We post-train the official Mamba-2 checkpoints for 100 steps with this technique (~ 0.02% of the pre-training bud- get) and show the results in Figure 5 (dashed red line). This simple method slightly improves the generalization of the 370m model, but does not work for the 780m and 1.3b mod- els. We believe the reason why this method does not work that the distribution of attainable states of the Mamba-2
4.4. Intervention 3: State Passing
In this intervention, we initialize the state the final state of different sequence (for convenience, shuffle all the se- quences and use the final states of the previous batch during training). Note that this is equivalent to sampling an initial state from the distribution of attainable states. Additionally, we want the model to perform well on sequences with no prior context. To achieve this, we implement a dropout
6

---

Understanding and Improving Length Generalization in Recurrent Models
Post training interventions on Mamba 370m
Post training interventions on Marnas-2 780m
Post-irainnng nterventions on Momba 2 1.3b
11-
Aarplaiity
Perpleaty
(Base)
Random Noise
ritted Neise
Passing
TRIT
Sequence Position (t)
Sequence Position F-)
Sequence Posit (t)
(a) Mamba-2 370m.
(b) Mamba-2 780m.
(c) Mamba-2 1.3b.
Figure 5. Position-wise perplexity of official Mamba-2 models (BASE) and four interventions that applied these models with 100 post-training steps. For the STATE PASSING and TBTT interventions, 100 steps is enough to enable length generalization in 32k length sequences for models. The interventions modify the initial state the recurrent models, thus facilitating the exploration of wider range of states. They sample an initial state from distributions that progressively get closer to the true distribution of attainable states: (1) RANDOM NOISE samples from a Gaussian distribution with fixed variance; (2) FITTED NOISE samples from a Gaussian distribution with mean and variance calibrated to the final states seen during training; (3) STATE PASSING uses the final state of different sequence as initial state; and (4) TBTT splits a sequence into several chunks and uses the final state of the previous chunk as initial state. STATE PASSING directly samples from the distribution of attainable states and together with TBTT has the best performance, supporting the unexplored states hypothesis. For STATE PASSING, the results for Mamba-1 I and Gated Lincar Attention are also shown in Figure 1.
mechanism that randomly zero-initializes the state with a probability of 0.1.
4.5. Intervention 4: Truncated Backpropagation Through Time (TBTT)
The results for this intervention are shown in Figure (dashed blue line). With only 100 steps (~ 0.02% of the pre-training budget), this intervention enables length gener- alization, which supports our hypothesis that the failure to length generalize is due to the model not being trained on attainable state distributions.
For the fourth intervention we use Truncated Backprop- agation Through Time (TBTT) (Williams Peng, 1990; Sutskever, 2013). This method consists in splitting a long sequence into smaller chunks, and for each chunk using the final state of the previous chunk as initial state. Even though the gradient propagation is stopped between chunks, the model still learns to model sequence based on some previous context (initial state).
In Figure 1, we also apply this techniques to enable length generalization in both Mamba-1 and GLA. Additionally, show the results for Effective Remembrance in these intervened models. The Effective Remembrance values are smaller than the baseline, especially for 1 < T, suggesting that the intervened models are not substantially affected by far away tokens and correctly model the recent context. Interestingly, the Effective Remembrance curves for the Mamba models and GLA models have different shapes. In GLA, the tokens that are more than 512 positions away from T have small impact on the output, but the model is extremely affected by tokens that are just 512 positions apart (we recall that 512 is the training context length for the GLA models we are evaluating).
A comparison with the rest of the interventions is shown in Figure 5 (purple dashed line). It is worth noting that this method achieves very similar performance to State Pass- ing, thus indicating that they key for length generalization is exploring attainable state distributions, not necessarily processing (chunked) long sequences.
5. Performance of Interventions on Long Context Tasks
In the previous section we showed that the fitted noise, State Passing and TBTT interventions help length generalization on the Pile (Gao et al., 2020). As we mentioned in Sec- tion 3.4, length generalization would also be possible if the model only used the last T tokens of sequence to output
Additionally, in Figure 4 we show that a model post-trained with State Passing produces states whose distributions do not seem to significantly change over time, which sheds light into how this intervention achieves length generaliza- tion: the state update mechanism is modified so that the distribution of attainable states does not change much after the training context.
'From an implementation point of view, the difference with State Passing is that TBTT does not shuffle the sequences and requires carefully setting up the dataloader sampler so that the final state of chunk is used as initial state for the next chunk.
7

---

Understanding and Improving Length Generalization in Recurrent Models
BABILong Mamba-2 780m: Few Shot
BABILong Mamba-2 780m: Finetuned
Base
0.4
Accuracy
0.3
State Passing
Accuracy
0.3
TBTT
Base
0.2
0.2
State Passing
0.1
0.1
TBTT
0.0
0.0
1K
2k
4K
16k
32k
64k
128k
256k
0k
1k
2k
8k
16k
32k
64k
128k
256k
Sequence Length
Sequence Length
(a) Few Shot.
(b) Finetuned.
Figure 6. Few shot and finetuned performance for the BABILong benchmark (Kuratov et al., 2024) of Mamba-2 780m model under three settings, BASE corresponds to the official checkpoint; STATE PASSING is the official checkpoint post-trained with State Passing on the Pile (see Section 4.1); and TBTT corresponds the official checkpoint finetuned with Truncated Backpropagation Through Time on the Pile (see Section 4.5). The STATE PASSING and TBTT interventions significantly improve the baseline in both the few shots and finetuned settings, achieving reasonable performance in sequences up to length 256k despite only heing trained with a context of 2k. More finetuning details on Section B.3
its prediction (where T' is the training context). However, this not a desirable behavior, as the model would fail in tasks that require reasoning over tokens that more than T positions apart. Additionally, length generalization was measured using perplexity, but no other metrics were explored. Thus, an important question remains: Are these interventions truly effective in tasks that require reasoning over sequences longer than the training context T? In this section, we answer affirmatively by showing results on three long context tasks.
fitted noise are capable of handling relationships between to- kens that are much more than 2k positions apart. Moreover, they achieve better performance much harder tasks; in particular the 780m achieves perfect accuracy sequences of length 256k. We provide more details on the tasks and finetuning recipe on Section B.4.
5.3. Synthetic Copying
The synthetic copying task (Jelassi et al., 2024) consists in copying an arbitrary sequence of tokens. In Table 1 we present the results for a Mamba-2 model that is trained from scratch, and we show that using State Passing during training greatly improves length generalization sequences more than three times longer. Thus, State Passing helps the model solve long context tasks that harder than those seen during training.
5.1. Long Context Reasoning: BABIL.ng benchmark
BABILong (Kuratov et al., 2024) is a challenging bench- mark which tests both the common sense understanding of model as well its ability to capture long range depen- dencies in text. In Figure 6, it can be observed that State Passing enhances the length extrapolation capabilities of the model in both the few shot and finetuned settings (we recall that the model is trained and finetuned on sequences of length 2048). Therefore, State Passing is not only useful in fixing the diverging perplexity of established language models, but also in enhancing their ability to solve long context reasoning tasks.
ARCHITECTURE
CHARACTER ACCURACY
Mamba-2
0.27 + 0.03
Mamba-2 State Passing
0.47 + 0.11
Table 1. Length generalization results for the synthetic copying task (Jelassi et al., 2024), which consists of copying arbitrary sequence of tokens. The models have 45 million parameters, are initialized randomly and trained on I million sequences of length between 50 and 100. They are evaluated on sequences of length 300. The State Passing intervention greatly improves length gener- alization.
5.2. Long Context Passkey Retrieval
The passkey retrieval task (Mohtashami & Jaggi, 2023) re- quires the model to retrieve a 5-digit passkey inserted at given depth of a long context. In Figure 8 we present results for two Mamba-2 model sizes under three different settings; zero shot evaluation, standard finetuning with context length 2048 and finetuning with the fitted noise intervention, also with 2048 (see Section B.4 for why we use the fitted intervention in this task). The models finetuned with
6. Related Work
Length generalization in Mamba through changes in the state update mechanism. Some works have proposed changing the internal computation of Mamba, for exam-
8

---

Understanding and Improving Length Generalization in Recurrent Models
Mamba2-370M Zero Shot
Mamba2-370M Finetuned
Mamba2-377 Finetuned (Fitted Noise)
Passkey Depth
0.1
Passkey Depth
0.1
Passkey Depth
0.1
0.75
0.5
0.5
0.5
0.50
0.9
0.5
0.9
0.25
0.00
64k
128k
25ck
64k
128k
236k
64k
128k
256k
Sequence Length
Sequence Length
Sequence Length
(a) 370m shot.
(h) 370m finetuned.
(c) 370m finetuned with fitted noise.
Mamba2-780M Zero Shot
Mamba2-780M Finetuned
Mamba2-780M Finetuned (Fitted Noise)
1.00
Passkey Depth
0.9 0.5 0.1
Posskey Depth
0.1
Parskey Depth
0.1
0.5
0.3
0.50
0.9
0.9
0.25
64k
128k
256k
64k
128k
256k
64k
128k
256k
Sequence Length
Sequence Length
Sequence Length
(d) 780m zero shot.
(e) 780m finetuned.
(1) 780m finetuned with fitted noise.
Figure 7. Performance on the passkey retrieval task (Mohtashami & Jaggi, 2023) of the official Mamba-2 checkpoints under three scttings: (left) zero shot. (middle) standard finetuning, (right) linetuning with fitted noise. Finetuning with fitted noise enables solving the passkey task on sequences at least twice as long than with standard finetuning. Additionally, although the models are finetuned with sequences of length 2048, the fitted noise intervention solves tasks that require leveraging dependencies between tokens that are more than 2048 positions apart. More details on finetuning are given in Section B.4.
ple by updating the state Mamba only on some selected tokens-which equivalent to shortening the effective con- text (Ye al., 2025; Ben-Kish et al., 2024; Yuan et al., 2024), or by foreing the model forget previous context (Chen al., 2024b). While these methods show increased performance in tasks like passkey retrieval and document re- trieval (Mohtashami & Jaggi, 2023; Rajpurkar et al., 2018 8), a downside is that they change the internal state update mechanism and might be hard to transfer to architectures different to Mamba. In contrast, our interventions do not change the implementation of the architecture and are appli- cable several recurrent models.
model needs to be trained to enable length generalization namely, distributions that are close to the distribution of final states staching the conclusion that training on long sequences or with TBTT not always necessary.
7. Conclusion and Future Work
This work presents an empirical and theoretical study of the length generalization of recurrent models. We introduce the unexplored states hypothesis, which states that the failure to length generalize is due to the model not being trained the distribution of the states that would be attained if its recurrence applied to long sequences. Besides bringing important insights about the states and the behavior of re- current models through tools like Effective Remembrance, this works also presents simple and general tools to enable length generalization for general recurrent architectures.
Length generalization in SSMs and overfitting to short sequences. Some other works link the failure to length generalize of SSMs with state overparametrization and over- fitting to short sequences. Wang (2024) identifies that zero- initialized recurrent models struggle with length generaliza- tion, provides& theoretical analysis that links length gen- eralization with polynomial extrapolation, and proposes pre-training models with State Passing and TBTT (Williams & Peng, 1990; Sutskever, 2013) to enable length general- ization. Chen et al. (2024b) shows that training models for longer sequences enables length generalization and argues that the failure to length generalize is due to the model being overparametrized for its training context and not learning to forget past tokens. We build upon these insights to propose the unexplored states hypothesis as precise explanation for why recurrent models fail to generalize. We focus on the dis- tribution of the states and specify on which distributions the
We mention several directions for future work. Firstly, while our work mostly focuses length generalization in tasks related to text modeling, further study of length extrapo- lation in other distributions with long range dependencies would be insightful. Secondly, we chose Mamba, Mamba-2, Gated Linear Attention and RWKV-v6 as a representative subset modern recurrent architectures, but would be interesting analyze the training interventions on other re- current architectures, for example time invariant models like S4 (Gu et al., 2022) or hybrid models (Lieber et al., 2024; Ren et al., 2024; De et al., 2024; Arora et al., 2024)).
9

---

Understanding and Improving Length Generalization in Recurrent Models
Acknowledgements
lives in Creating Large Language Models, 2022. URL
fetpp://arxivvorgge /204 06745.
RBR was supported by the "la Caixa" Founda- tion (ID 100010434). The fellowship code is LCF/BQ/EU2//1190090 We thank Aakash Lahoti for assistance with the synthetic copying experiment.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020, URL https::/ppooeeddnggs neurips. cc httppaper_fiieee ccc/:ppper_LLleeeeppppr//////L 1 145700600fa.00004418199eel222222 pdf.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
References
Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S., Zinsley, D., Zou, J., Rudra, A., and Ré, C. Sim- ple linear attention language models balance the recall- throughput tradeoff. arXiv:2402. 18668, 2024.
Chen, S., Wong, S., Chen, L.., and Tian, Y. Extending context window of large language models via positional interpolation, 2023a. URL https: hatp:s//arxxv.org/ abs/ 2306. 15595.
Bansal, A., Schwarzschild, A., Borgmia, E., Emam, Z., Huang, F., Goldblum, M., and Goldstein, T. End-to-end algorithm synthesis with recurrent networks: Extrapola- tion without overthinking. In Advances in Neural Infor- mation Processing Systems, volume 35, pp. 20232-20242, 2022.
Chen, S., Wong, S., Chen, I.., and Tian, Y. Extending context window of large language models via positional interpolation, 2023b. URL 2ttp/2306.155993
Beck, M., Pöppel, K., Spanring, M., Auer, A., Prudnikova, O., Kopp. M. K., Klambauer, G., Brandstetter, J., and Hochreiter, S. XLSTM: Extended long short-term mem- ory. In The Thirty-eight Anmual Conference on Neural Information Processing Systems, 2024. URL https: //openreview net/ forum? S-AAAAAAPPPAAAA
Chen, Y., Qian, S., Tang, H., Lai, X., Liu, Z., Han, S., and Jia, J. LongloRA: Efficient fine-tuning of long- context large language models. In The Twelfth In- lernational Conference on Learning Representations, 2024a. URL. htw.t::t//////.... /ttp:///oppeeeeeeeeeeeeeeeeeeeee net/ forum? id- 6PJooRfaaK
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. arXiv:2004.055 05100, 2020.
Chen, Y., Zhang, X., Hu, S., Han, X., Liu, Z., and Sun, M. Stuffed mamba: State collapse and state capacity of rnn-based long-context modeling, 2024b. URL https:
Ben-Kish, A., Zimerman, I., Abu-Hussein, S., Cohen, N., Globerson, A., Wolf, L., and Giryes, R. Decimamba: Ex- ploring the length extrapolation potential of mamba, 2024. URL httph:tt////// /ttp/aaaivvv/v.////v///.../a..va org/ abs /ass/244 6 14528.
// //arxiv..co.//////////////aa org/ httpa/ab////.0.1...
Dao, T. and Gu, A. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML), 2024.
Biderman, S., Schoelkopf, H., Anthony, Q. G., Bradley, H., O' Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E., et al. Pythia: A suite for ana- lyzing large language models across training and scaling. In International Conference Machine Learning, pp- PMLR, 2023.
De, S., Smith, S. L., Fernando, A., Botev, A., Cristian- Muraru, G., Gu, A., Haroun, R., Berrada, I., Chen, Y., Srinivasan, S., Desjardins, G., Doucet, A., Budden, D.. Teh, Y. W., Pascanu, R., De Freitas, N., and Gulcehre, C. Griffin: Mixing gated linear recurrences with local atten- tion for efficient language models. arXiv [cs.LG), Febru- ary 2024. URL http:/ btttp://arxive /ttp/aiivvvv/////////ee////aaeva org/ abs/ 2402. 19427.
Black, S., Biderman, S., Hallahan, E., Anthony, Q., Gao, L., Golding, I.,, He, H., Leahy, C., McDonell, K., Phang, J., Pieler, M., Prashanth, U. S., Purohit, S., Reynolds, L., Tow, J., Wang. B., and Weinbach, S. GPT-NeoX-2055 An open-sonnee source autoregressive language model. In Pro- ceedings of the ACL Workshop on Challenges Perspec-
Ding, Y., Zhang, L. L., Zhang, C., Xu, Y., Shang, N., Xu, J., Yang, F., and Yang, M. Longrope: Extending lim
10

---

Understanding and Improving Length Generalization in Recurrent Models
context window beyond million tokens, 2024. URL https:://rxiv. org. abs /2002.22222. 13753.
Information Processing Systems Datasets and Bench- marks Track, 2024. URL https:/ / /pwnrevie...eeeeeeeeeeeeeeeeee.r net/ fer//oorm?? 7
Gao, L., Biderman, S., Black, S., Golding, L.. Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 2020.
Li, T., Zhang. G., Do, D., Yue, X., and Chen, W. Long- context Ilms struggle with long in-context learning, 2024. URL https:
: / /ett://axxxttt wwwaaivvvvooo www.voogg/aaa. /a/g//bb//// 24444. 02060.
Lieber, O., Lenz, B., Bata, H.,, Cohen, G., Osin, J., Dalmedi- gos, I., Safahi, E., Meirom, S., Belinkov, Y., Shalev- Shwartz, S., Abend, O., Alon, R., Asida, T., Bergman, A., Glozman, R., Gokhman, M., Manevich, A., Rat- ner, N., Rozen, N., Shwartz, E., Zusman, M., and Shoham, Y. Jamba: A hybrid transformer-mamba lan- guage model, 2024. URL a abs/24033 1987..
Grattafiori, A. et al. The llama 3 herd of models, 2024, URL https: //axivv. org/ abs/ 2407 21783.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.000752 2023.
Gu, A., Goel, K., and Ré, C. Efficiently modeling long se- quences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022.
Liu, B., Wang, R., Wu, L., Feng, Y., Stone, P., and Liu, Q. Longhorn: State space models are amortized online learners, 2024. URL https: /ttp/vixvvv./////////////ee orgg /bss/ 2407.14207.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Vinyals, O., Rae, J. W., and Sifre, L. Training compute-optimal large lan- guage models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS '22, Red Hook, NY, USA, 2024. Curran Associates Inc. ISBN
Ma, X., Yang. X., Xiong, W., Chen, B., Yu, L., Zhang, H., May, J., Zettlemoyer, L., Levy, O., and Zhou, C. Mega- lodon: Efficient Ilm pretraining and inference with un- limited context length. arXiv preprint 20X0v:2404.08801, 2024.
Mohtashami, A. and Jaggi. M. Random-access infi- nite context length for transformers. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023. URL https:://///.......... httpt:://ppeereeiieeteee//// wwwpr/weeeeeeeeeeeeepeeeeeeeeeee net forum? 64-77Hnn6400V
Jelassi, S., Brandfonbrener, D., Kakade, S. M., and Malach, better than state space models at copying. In Salakhutdinov, R., Kolter, Z., Heller, K. Weller, A., Oliver, N., Scar- lett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learn- ing, volume 235 of Proceedings of Machine Learn- ing Research, pp. 21502-21521. PMLR, 21-27 Jul 2024. URL StRLhttps://proceeee http///proeeeeng.....e www.mnn..mmm...m.mm............. mlr press// v235/jelassi24a.ttt
E. Repeat after me: Transformers
Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H.. Cheng, X., Chung, M.. Grella, M., GV, K. K., He, X., Hou, H., Lin, J., Kazienko, P., Kocon, J., Kong. J., Koptyra, B., Lau, H., Mantri, K. S. I., Mom, F., Saito, A., Song, G., Tang, X., Wang, B., Wind, J. S., Wozniak, S., Zhang, R., Zhang, Z., Zhao, Q., Zhou, P., Zhou, Q., Zhu, J., and Zhu, R.-J. Rwkv: Reinventing rnns for the transformer era, 2023. URL attps://arx:: httpss////xxvvvoo//bbeee httg:/abs/3300..1 2205.
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on
Machine Learning, pp. 5156-5165. PMLR, 2020.
Peng, B., Goldstein, D., Anthony, Q., Albalak, A., Alcaide, E., Biderman, S., Cheah, E., Du, X., Ferdinan, T., Hou, H., Kazienko, P., GV, K. K., Kocoñ, J., Koptyra, B., Krishna, S., Jr., R. M., Lin, J., Muennighoff, N., Obeid, F., Saito, A., Song, G., Tu, H., Wirawan, C., Wožniak, S., Zhang, R., Zhao, B., Zhao, Q., Zhou, P., Zhu, J., and Zhu, R.-J. Fagle and finch: Rwkv with matrix-valued states and dynamic recurrence, 2024a. URL https: / //arxiv. org/ abs 2444 .08002.559....
Katsch, T. Gateloop: Fully data-controlled linear recurrence for sequence modeling. arXiv, abs/23 1ttp/2311.000 1.01927, 2023.
Kazemnejad, A., Padhi, Ramamurthy, K. N., Das, P., and Reddy, S. The impact of positional encoding length generalization in transformers, 2023. URL https:// arxiv. org. abs/ 2205 11.00..9966 19466.
Kuratov, Y. Bulatov, A., Anokhin, P., Rodkin, I., Sorokin, D. I., Sorokin, A., and Burtsev, M. BABILong: Test- ing the limits of LLMS with long context reasoning-in a-haystack. In The Thirty-eight Conference on Neural
Peng, B., Quesnelle, J., Fan, H., and Shippole, E. YaRN: Efficient context window extension of large language models. In The Twelfth International Conference on

---

Understanding and Improving Length Generalization in Recurrent Models
Learning Representations, 2024b. URL https:// openreview. ntt/ htt.:/oormmiiidwwHee 3=wHHfxhhulu.
Waleffe, R., Byeon, W., Riach, D., Norick, B., Kor- thikanti, V., Dao, T., Gu, A., Hatamizadeh, A., Singh, S., Narayanan, D., Kulshreshtha, G., Singh, V., Casper, J., Kautz, J., Shoeybi, M., and Catanzaro, B. An empirical study of mamba-based language models. arXiv [es.LG), June 2024. URL http: p /aaxivv..v./../////////////a/iv/ org/ wttgabs/////ooooooeot...o...sere 2406. 07887.
Press, O., Smith, N., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapo- lation. In International Conference on Learning Represen- tations, 2022. URL https: hatps:///pppeeeee http//pppneeewweceee/e net/ forum? RBpppPCCC
Wang. S. Longssm: On the length extension of state-space models in language modelling, 2024. URL https:// arxiv. wwxiv.or///.... org/abs///////ooooooeo/...aaaaar 2406. 02080.
Rajpurkar, P., Jia, R., and Liang. Know what you don't know: Unanswerable questions for SQuAD. In Gurevych, and Miyao, Y. (eds.), Proceedings of the 56th Annual Meeting of the Association for Compula- tional Linguistics (Volume 2: Short Papers), pp. 784-789, Melbourne, Australia, July 2018. Association for Compu- tational Linguisties. doi: 10. 1/11866///1111882224 URL https: / acl antho wttp://-1/11111111
Weston, J., Bordes, A., Chopra, S., and Mikolov, T. To- wards ai-complete question answering: set of prereq- uisite toy tasks. In Bengio, Y. and LeCun, Y. (eds.), 4th International Conference on Learning Represen- tations, ICLR 2016, San Juan, Puerto Rico, May 2-
4, 2016, Conference Track Proceedings, 2016. URL http: / //aaxxxxvvoorg/ org/ abs 11////12222 05698.
Ren, I.,, Liu, Y., Lu, Y., Shen, Y., Liang, C., and Chen, W. Samba: Simple hybrid state space models for effi- cient unlimited context language modeling, 2024, URL ht https: //apsi///avvv..///a org/ wwgg/bb//.............geeee abs/2405 07522.
Williams, R. J. and Peng. J. An efficient gradient-based
algorithm for on-line training of recurrent network trajec- tories. Neural Computation, 2(4):490- 1990. doi: 10. 162/neco. 1990.2.4.490.
Smith, J. I., Warrington, A., and Linderman, S. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023. URL https:://// wpwpr/peeeeeeeeeeeeeeeeeeeee.e.. net/ forum? (d=Ai8Hw 3WXckss
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef- ficient streaming language models with attention sinks. In The Twelfth International Conference Learning Representationss, 2024. URL https:tt////// http:///pppeeviieeeeee net het/forummid-aa Nd
Sun, Y. Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. Retentive network: A successor to transformer for large language models. arXiv preprint 17Xiv22307.08621, 2023.
Yang, S., Kautz, J., and Hatamizadeh, A. Gated delta net- works: Improving mamba2 with delta rule, 2024a. URL https: /tts:///arxivv arwiv. org abs / 22/22211222 04464.
Sun, Y., Li, X., Dalal, K., Xu, J.. Vikram, A., Zhang, G., Dubois, Y., Chen, X., Wang, X., Koyejo, S., Hashimoto, T., and Guestrin, C. Learning to (leam test time): Rnns with expressive hidden states, 2024. URL https: //arxiv. org/ abs 04620.
Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated linear attention transformers with hardware-efficinn train- ing. Proceedings of ICML, 2024b.
Yang. S., Wang. B., Zhang, Y., Shen, Y., and Kim, Y. Parallelizing linear transformers with the delta rule over sequence length. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024c. URL https: http::///peeeeeeeeeeeeee net/ forum? 1d=y8Rm4VNRBS
Sutskever, I. Training Recurrent Neural Networks. Ph.d. thesis, University of Toronto, Toronto, Ontario, Canada, 2013.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In Proceedings of the 31 International Conference on Neural Information Processing Systems, NIPS' 17, pp. 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Ye, Z., Xia, K., Fu, Y., Dong, X., Hong, J., Yuan, X., Diao, S., Kautz, J., Molchanov, P., and Lin, Y. C. Long- mamba: Enhancing mamba's long-contex. capabilities via training-fre receptive field enlargement. In Pruceed- ings of the International Conference on Learning Repre- sentations (ICLR), 2025. 2025.
Veerabadran, V., Ravishankar, S., Tang, Y., Raina, R., and de Sa. V. Adaptive recurrent vision performs zero-shot computation scaling to unseen difficulty levels. Ad- vances in Neural Information Processing Systems, vol- ume 36, pp. 18132- 183132-181444 18145, 2023.
Yuan, D.,, Liu, J., Li, B., Zhang, H., Wang. J., Cai, X., and Zhao, D. ReMamba: Equip mamba with effective long- sequence modeling. arXiv [cs. [cs.CL], August 2024. URL http:://axxv..org// 2408 15496.
12

---

Understanding and Improving Length Generalization in Recurrent Models
Zhang, B. and Sennrich, R. Root Mean Square Layer Normalization. In Advances in Neural Information Processing Systems 32, Vancouver, Canada, 2019. URL httpt::////// wwweeneeeiww.neteeeeeeeeee. pdf2id=S1qBAffrr.
Zhao, L., Feng, X., Feng, X., Zhong, W., Xu, D., Yang. Q. Liu, H., Qin, B., and Liu, T. Length extrapolation of transformers: A survey from the perspective of positional encoding. In AI-Onaizan, Y., Bansal, M., and Chen, Y.-N. (ods.), Findings of the Association for Computational Lin- guistics: EMNLP 2024, pp. 9959-9977, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10. 10.186555w120224fnddin----mmm 582. URL https:tt////// httpt:/ttacaaaeeeeee http:///acaaattoooooooor//////// 582/.
Zhu, D., Yang, N., Wang. L., Song, Y., Wu, W., Wei, F., and Li, S. Pose: Efficient context window extension of ilms via positional skip-wise training, 2024. URL https: /tte:///arre /awrivvvvvv/va///////ee./aaaaav.
13

---

Understanding and Improving Length Generalization in Recurrent Models
Table Model configurations and training hyperparameters for our experiments. The learning rates follow the values of previous works
(Brown et al., 2020; Biderman et al.. 2023).
ARCHITECTURE
PARAMS
n_layers
d_model
n_heads / /dheedd
d_state
LEARNING RATE
BATCH SIZE
130m
24
768
16
6e-4
0.5M tokens
Mamba-1
350m
48
1024
16
3e-4
0.5M tokens
45m
12
512
8 64
128
le-3
0.5M tokens
85m
12
768
12 64
128
1e-3
0.5M tokens
130m
24
12/64 64
6e-4
0.5M tokens
Mamba-2
350m
48
1024
16/ 64
128
3e-4
0.5M tokens
780m
48
1536
16/96
128
2.5e-4
0.5M tokens
1.3b
48
2048
32/64
128
2e.4
0.5M tokens
70m
512
4/ 128
le-3
0.5M tokens
GLA
120m
768
4/ 182
6e-4
0.5M tokens
360m
20
1024
4/6
3e-4
0.5M tokens
A. Extended Related Work
Length generalization in Transformers. In order length generalize, Transformers face the added difficulty of handling positional encoding, which is not easily extendable beyond the training context (Zhao al., 2024). To deal with this issue, some works have proposed several position interpolation techniques (Chen et al., 2023b; Ding et al., 2024; Kazemnejad et al., 2023; Zhu et al., 2024; Peng et al., 2024b). Other works have proposed finetuning with sparse attention to allow more efficient length generalization (Chen et al., 2024a; Mohtashami & Jaggi, 2023), or they have developed methods to maintain a fixed-size sliding window for the KV cache (Xiao et al., 2024). These methods are bespoke to the usage and architecture and not always achieve good performance in long context tasks (Press al., 2022; Li al., 2024; Chen et al., 2023a). In contrast, recurrent models can principle extend beyond their training context by simply rolling out the state recurrence. In this work, we show that there simple and general interventions to enable length generalization in recurrent models, thus fully leveraging their benefits.
Algorithmic extrapolation. Some works have evaluated recurrent models on algorithmic tasks that are harder to solve than the ones the model has been trained on, and they have proposed training modifications that prevent the model from overfitting to the simpler training tasks (Bansal et al., 2022: Veerabadran et al., 2023). particular, Bansal et al. (2022) proposes Incremental Progress Training, which reuses states from previous iterations of the task discourage the model from learning iteration-specific behaviours, which simlar to why the TBTT intervention works text. While there are some differences in the tasks (e.g. in these algorithmic extrapolation tasks the state should converge fix point), some of the techniques used are similar. This work presents the unexplored stales hypothesis and framework to reason about length generalization based on the distribution of states, which helps understand why these interventions work. We leave it as future work to study them in a broader set of cases.
B. Training recipes
B.1. Pre-training Language Models
For the experiments Section 3, we train on the Pile (Gao al., 2020) with the EleutherAI/gpt -nexx- 200+ tokenizer (Black et al., 2022). For the learning rate, peak learning rate given by Table and decay to le - 5. The gradients are clipped to 1.0 and no dropout is used. Additionally, we also follow the improved training recipe of Grattafiori et al. (2024), with an Adam optimizer with 31 = 0.9 and 833 = 0.95, weight decay scheduling with peak of 0.01, RMSNorm (Zhang & Sennrich, 2019) instead LayerNorm and no linear biases. For Mamba-1 I and Mamba-2 we use a training context of 2048 and for GLA we use a context of 512. Depending on the experiment, the models are trained for several times what Chinchilla scaling laws dictate (Hoffmann et al., 2024), see Figure and Figure 3.
cosine scheduling with warmup in the 10% first training steps,
"httpe://hugginngface. co/ http://wwgggnnga....eeeeeeeeeeea
14

---

Understanding and Improving Length Generalization in Recurrent Models
B.2. Post-training Interventions on the Pile
For the post-training interventions of Section 4, we use the same recipe as for model pre-training (section B.1) yet using a peak learning rate that is ten times smaller than the one given in Table 2. We post-trate pest-train for 100 steps.
B.3. Babilong Finetuning
The models are evaluated on the tasks [qal, qa2, qa3, qa4, qa5, qa6, qa7, qa8, qa9, qa10] of the benchmark, and finetuned on such tasks for one epoch using facts and questions from the BABI training dataset (Weston et al., 2016). In the finetuned setting, all the models are finetuned on BABILong without State Passing nor TBTT, thus the benefits of having State Passing or TBTT finetuned checkpoint are not lost when finetuning again for this task.
B.4. Passkey Retrieval Finetuning
Contrary to typical language modeling datasets, the distribution of tokens in the passkey task is not stationary. This is why we show results for the fitted noise intervention, as it does not require using the final state of the sequence (i.e., right after revealing the passkey), which might not be appropriate as the initial state.
We finetune the official Mamba-2 checkpoints on the passkey retrieval task using the same procedure as section B.1, this time for 1000 steps and using a peak learning rate ten times smaller than the one given in Table 2. In order to finetune, we mask all the tokens the sequence that do not correspond the passkey, trim the filler sentence to have uniform batch size, and sample different passkey depth between 0 and for each sample. Additionally, we add a period after the passkey (""") because we observed that some models failed because they repeated the passkey several times consecutively (i.e. for passkey of 12345, the output contained 123451 Figure 8.
An example of a sample from the dataset is shown in
There is important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.
green. The sky is blue. The sun is yellow. Here we go. There and back again. The grass green. The sky blue. The sun is yellow. Here we go. There and back again. The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. The The pass key is 3327. Remember it. 3327 is the pass key.
The grass is green. The sky blue. The sun yellow. Here we go. There and back again. The grass is green. The sky is blue. The sun is yellow, Here we go. There and back again. The grass green. The sky is blue. The sun yellow. Here we go. There and back again. The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. What is the pass key? The pass key is 3327.
Figure 8. Sample of the passkey retrieval task (Mohtashami & Jaggi, 2023) with length 256 and depth 0.5.
C. Distribution of the State at Given Heads and Layers Over Time
Figure shows the standard deviation of certain heads of the Mamba-2 state as a function of time. In the official checkpoint, the standard deviations increase after the context length, thus producing distribution shift which explains why the model fails to length generalize. In contrast, the State Passing intervention fixes this issue by producing states whose standard deviations are more stable after the training context.
D. Position-wise Perplexity for RWKV-v6
Figure 10 shows the position-wise perplexity of a RWKV-v6 model (Peng et al., 2024a) trained from scratch with a context of T 256 and post-trained with State Passing, When the training context too short, the RWKV-v6 model fails to length generalize (as discussed in Section 3.2), whereas the State Passing intervention (Section 4.1) achieves length generalization.
E. Impact of the Warm-Up Period on Length Generalization
In Figure 1I we show the results of pre-traue neeitraining a Mamba-2 45m model from scratch with two different learning rate warm-up periods and a context length of T 1024. Having an extended warm-up period hinders length generalization
15

---

Understanding and Improving Length Generalization in Recurrent Models
Standard deviation (Layer 4, Head 2)
Standard deviation (Layer 6, Head 4)
Standard deviation (Layer 16, Head 1)
17.5
30-
14
25-
12.5
10
10.0
2.0-
Std Dev
Base
15
State Passing
50-
101
Base
Base
State Passing
State Passing
Standard deviation (Layer 16, Head 9)
Standard deviation (Layer 18, Head 7)
Standard deviation [Layer 20, Head 11)
10-
Base State Passing
Dev
Std Dev
Base
Base
State Passing
State Passing
2000
4000
6000
8000
4000
0050
8000
2000
4000
6000
ro20
Sequence Position
Figure Standard deviation of parts of the state of the Mamba-2 130m official checkpoint versus sequence position (he in the notation of Section 2.2). As in the case of Figure 4, the model post-trained with State Passing produces states whose standard deviations are more stable than the baseline across time. The layers and heads are selected based on which had more variation across time, the majority of other heads in the baseline model do not exhibit this abnormal behavior. Thus, the distribution shift that is observed in Figure 4 only due to specific heads in the state. Sometimes, the State Passing intervention generates states with a higher standard deviation than the pre-trained model (see Layer 4, Head 2 and Layer 6, Head 4). Thus, note that the solution to length generalization is not avoiding large standard deviations in the state; but rather having states whose distribution do not significantly change after the training context.
RWKV-V6 70m
50
40
Perplexity
20
10a
10a
10a
Sequence Position (t)
Base
State Passing
Figure 10. Position-wise perplexity of RWKV-v6 70m model (Peng et al., 2024a) pre-trained from scratch for 35B tokens and post- trained with State Passing for 250m tokens on The Pile (Gao al., 2020) (see B details). It exhibits a similar behavior to the models shown in Figure I: when pre-trained in short context for many tokens fails to length generalize, but post-training with small amount of steps with State Passing enables length generalization.
16

---

Understanding and Improving Length Generalization in Recurrent Models
Mamba-2 45m (two warm-up periods)
20.0
19.5
19.0
18.5
Perplexity
18.0
17.5
17.0
16.5
16.0
10a
10a
104
Sequence Position (t)
10% warm-up
25% warm-up
Figure Mamba-2 45m trained on context 25% refer to having a warm-up period that lasts 10% and 25% of the full training. respectively). An increased warm-up period hinders performance beyond training context, demonstrating impact training recipe length generalization. The models are trained for 9B tokens (10x Chinchilla scaling laws) with the recipe given in Section B.1.
1024 with different warm-up periods of the learning rate scheduler (10% and
in this case, suggesting that length generalization can be affected by training hyperparameters. Our work removes the need to carefully analyze and tune these hyperparameters, since length generalization is expected to be achievable with the interventions we propose in Section
F. Impact of the Standard Deviation on the Random Gaussian Initial State Intervention
Mamba-2 125m post-trained with different random noise a
15
13:
Perplexity
12
10
102
10a
Sequence Position
Base
a= 0.125 ==3 2
0= 32
== 0.03125
o= 0.5
g= 8 0= 128
Figure 12. Position-wise perplexity of Mamba-2 125m post-trained for 100 steps with random initial state at different standard deviations a. This post-training technique enables length generalization, the noise level too high hurts performance.
In Section 4.2, we proposed sampling the initial state from a Gaussian distribution of standard deviation o (see Section 4.2), which is parameter that needs to tuned for each model. In Figure we show the position-wiss wise cross entropy of a 125m Mamba-2 model post-trained with different standard deviations in the random initial state. This post-training technique
17

---

Understanding and Improving Length Generalization in Recurrent Models
Mamba-2 EfiRem (Total Variat on)
Mamba-2 EffRem (jensen Shannon)
Mamiba-Z 2 EtfRem (Cosine Similanty)
30"
Enfoamytt
10.22
10
10
10
2000
4000
6030
2000
2000
4003
5000
5000
2000
4008
5000
5000
Sequencce Pos costionn
Sequence Position
Sequence
130m
(a) Total Variation.
(b) Jensen- Shaen-Shannon Divergence.
(c) 1 Cosine Similarity.
Figure Comparison of different distance metrics for the computation of Effective Remembrance (Section 3.4). The models shown are the Mamba-2 official checkpoints (solid lines), and their post-trained counterparts with State Passing (Section 4.1). They are evaluated on the valildation dataset the Pile (Gao et al., 2020), and the Effective Remembrancce value 8192. these models, the shape of the Effective Remembrance is robust to the choice of distance.
facilitates length generalization; however, an excessively high noise level degrades performance. On the other hand, we have observed that if the noise standard deviation is too small, this intervention does not enable length generalization.
G. Ablation on the Choice of Distance for Effective Remembrance
In Section 3.4 we introduced Effective Remembrance as a metric to understand the impact of certain parts of the context in the output an autoregressive model. Since Effective Remembrance compares two different outputs of the autoregressive model, which are probabilities, we used total variation as a distance between them. In this subsection, we explore the impact of using different distance (the Jensen- distance) and also using the cosine similarity between the probability vectors. The results for the Mamba-2 baseline models and the same models post-trained with State Passing is shown in Figure 13. Although the absolute values vary, it can be seen that the shape of the Effective Remembrance does not change much depending on the distance.
H. Derivation of SSM Equations for Non-Zero Initial State
Taking equation la and unrolling it:
he Arhi-1 + B,3
A (Ar-1hi--2 + B BL-IIA-1) + BAs A-1 + B1-1ART1-1 + B
-- + Σ A
==0
State Passing Pytorch Pseudocode
In section 4.1 we provide the equations to compute the output and final state recurrent model when the initial state not pero-initialized. Figure 14 we provide pseudocode to compute the contribution of the initial state to the final state and output.
18

---

Understanding and Improving Length Generalization in Recurrent Models
import
import cunch.i...nnniiiin
def it torch. Tencar, Pensor,
Tenaos,
C: torch.
state: torch, Fensor,
tuple (torch. torch, Tensor):
8: Batch size
3t Bequenoe Dimension
Midden disenvion
dimsnsion
Topular
Japut (aquation
(8,
B: (9, H, 2) à (Squet lon Ia]
C1 (8, H,
C (Squation 1h)
Initiei
Outputa:
ya (8, 8, 2) Coatonbution of the initial state to the outouc Nj Finel stacH
stabe_
n.*
Generate A 10:1), Allli
(T:7)
Rece_rev Lorch. ma www.nn... (lip(-)), dine 1).[1ip(-11
exp padia, finel_state torch.einaunl fhet, hnt, bint obhn", X, B, Aookrrels, f f 1:13
a Final scate cont contz iuution From initial state
Final. state final, state
revis,
Calouiate conce cbotribtioon fon
state on outpata
An torch torch.esg mep storehh torch,eiing ael (fb,,, biint bit. Sbht", state, Agea, C)
111 on 7)
y, tinel_state
Figure Pseudocode to compute final state (Equation 4a) and contribution to output from initial state (second term of the right hand side of Equation 4b) for recurrent model when the initial state is not zero-initialized. This code assumes that the matrix Aç diagonal.
19


Reviews:
[{"review_id": "SIFfUqpEgw", "reviewer": "Reviewer_jMRg", "summary": "The authors in this paper propose a framework to analyze the problem of length generalization in recurrent networks. The authors primarily focus on State Space Models and how they behave when test sequences are significantly longer than training sequences by studying their response to 4 training interventions: (1) initializing the SSM with a random initial state, (2) learning the distribution from which the initial state from (1) is sampled, (3) initializing the SSM with the final state of a different sequence and (4) using truncated back propagation through time. The authors claim that the proposed simple interventions lead to significant improvement in length generalization when compared to vanilla Mamba SSMs trained without these interventions.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback and helpful discussion around length generalization in recurrent models more broadly. We are encouraged to see that the reviewer believes this type of work is relevant and that our training intervention to the initial states of SSMs are sound. We provide answers to their questions:\n\n> RE: Experiments limited to SSMs and Gated Linear Attention\n\nWe understand the reviewer's concern about our results being specific to SSMs and Gated Linear Attention (GLA). **For that reason, we have performed a new experiment in another recurrent architecture, RWKV-v6 [1], and have found a similar phenomenon**: the model trained on short contexts and zero-initialized states diverges, but post-training with state passing enables length generalization:\n\nhttps://postimg.cc/zLkhvF2x\n\nWe will include this result in the final version of the paper to increase the scope of studied recurrent models. We also tried using DeltaNet and Gated DeltaNet [2], but the open source implementations seem to have a bug when using a non-zero initial state (the output of the forward applied to a long context versus substituting part of the context with an initial state do not match), so we could not test our interventions. We hypothesize that the benefits of training with non-zero initialized states will also apply to other modern recurrent networks, but we agree with the reviewer's remark and in the final version of the paper **we will make our claims more specific to the scope of architectures investigated.**\n\n[1] Peng, B. et al. (2024). _Eagle and Finch: RWKV with matrix-valued states and dynamic recurrence_. In _First Conference on Language Modeling_.\n\n[2] Yang, S. et al. _Gated Delta Networks: Improving Mamba2 with Delta Rule_. In _The Thirteenth International Conference on Learning Representations.\n\n> RE: Evaluation on other long context tasks besides passkey retrieval and perplexity\n\nWe recognize that passkey retrieval on its own might not be enough to assess the length generalization capabilities of our interventions. Therefore, **we have evaluated the models on two additional long context tasks related to language modeling, which go beyond passkey retrieval in their complexity**. The initial state interventions bring significant benefits in these tasks. Due to limited space, we refer to the response to Reviewer b5dM for more details on the tasks.\n\nResults for the BABILong task: https://postimg.cc/6TjDhBZ0\n\nResults for the synthetic copying task: https://postimg.cc/ctwryLQV\n\n> RE: References related to length extrapolation in recurrent networks and TBTT / state passing\n\nWe thank the reviewer for providing these algorithmic extrapolation references. The references propose architecture modifications to enable length extrapolation in recurrent models, and also propose training them on a more diverse distribution of states (Incremental Progress Training). While related, our work is focuses on language modeling, where established models exhibit a diverging perplexity after the context length - which is surprising as this does not require length extrapolation, given that predicting a language token at say position 2500 should be roughly as hard as predicting it at position 2000. Several recent works in language modeling attempt to solve this length generalization issue by changing the inner mechanism of the models (see Section 6 for specific references), so **one of our key contributions resides in showing that length generalization in established language models is achievable through simple (yet underutilized) interventions on the initial state.** \n\nWe also agree that some of these interventions were used in prior works, but we believe our main contribution lies in establishing a framework for understanding why they help in recurrent models and providing empirical experiments to support it. In particular, our work: (1) introduces the **\"unexplored states hypothesis\", which explains the poor generalization performance of recurrent models by reasoning about the distribution of the states attained when processing tokens beyond the training context**; **(2) systematically evaluates a range of interventions on the initial state, including random and fitted noise, which are naturally motivated by the \"unexplored states hypothesis\"**; and (3) **provides a deeper understanding of the state and how recurrent models processlong context through metrics like Effective Remembrance and the analysis of the results of the interventions**.\n\nAdditionally, **we also show that these interventions enable length extrapolation in some tasks related to core aspects of language modeling, like passkey retrieval, synthetic copying and BABILong.** We agree that the algorithmic length extrapolation tasks of the provided references are important for improving machine learning models at large and some of the methods used are similar, so in the final version we will include them both as related work and as avenues for future work.", "claims_and_evidence": "- I believe that the claims about state space models benefitting from the abovementioned training interventions is validated with sound experimentation. \n- The authors however claim that they have developed a theoretical and empirical framework to understand length generalization broadly in recurrent models. This claim is not validated as the class of recurrent architectures evaluated in this work are (mostly just) SSMs and in addition, a specific type of hardware-efficient transformer (GLA).", "methods_and_evaluation_criteria": "- I believe the proposed methods and evaluation criteria are necessary and make sense for the problem of length generalization. However, the tasks used in this paper seem quite simplistic in nature and don't stress test the training interventions on more challenging length generalization problems.\n- Using perplexity to measure length generalization I believe only measures the stability of the recurrent state space at test sequence lengths beyond what was experimented with during training. However, it doesn't study whether the extrapolation state space retains the same high expressivity (seen during training) in this unseen domain of longer recurrent sequences. \n- I appreciate the authors evaluating their interventions on the passkey retrieval task, however, this task seems quite simplistic and would request the authors to consider tasks such as Maze Solving task in [1, 2] and the incremental grouping task in [3]. Evaluating on these more challenging tasks makes the work more impactful and practical.\n\nReferences:\n1. Bansal, A., Schwarzschild, A., Borgnia, E., Emam, Z., Huang, F., Goldblum, M., & Goldstein, T. (2022). End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking. Advances in Neural Information Processing Systems, 35, 20232-20242.\n2. Veerabadran, V., Ravishankar, S., Tang, Y., Raina, R., & de Sa, V. (2023). Adaptive recurrent vision performs zero-shot computation scaling to unseen difficulty levels. Advances in Neural Information Processing Systems, 36, 18132-18145.\n3. Goetschalckx, L., Govindarajan, L. N., Karkada Ashok, A., Ahuja, A., Sheinberg, D., & Serre, T. (2023). Computing a human-like reaction time metric from stable recurrent vision models. Advances in neural information processing systems, 36, 14338-14365.", "theoretical_claims": "NA.", "experimental_designs_or_analyses": "The authors have evaluated the effect of different training context lengths, number of trainable parameters and post-training interventions on a series of SSM architectures (Mamba-1, Mamba-2) and a hardware-efficient transformer architecture (GLA) in the context of length generalization. This experimental design does not necessarily reflect any lack of rigor, but like I mentioned, I find the task of passkey-retrieval to be quite simplistic and the choice of architectures to be incoherent with the overall phrasing of the contributions applying generally to all types of RNNs.", "relation_to_broader_scientific_literature": "The work is very relevant to current art in language modeling (LM) with recurrent architectures. Developing LMs with stable hidden state spaces is crucial to efficiently scaling performance at long sequence lengths during inference time. Hence, I find this work to be well situated to broader scientific literature on sequence modeling.", "essential_references_not_discussed": "I believe the two main interventions studied here (state passing and T-BPTT) are studied in other RNN models in prior art. So I don't agree that these interventions are novel contributions in this work. I believe that the authors have applied these interventions found in prior RNN length generalization literature on SSMs and GLA, which I believe is the contribution of the current work.\n\nPrior findings list:\n1. State passing: This is essentially what is studied under the name of Incremental Progress Training in the following paper on recurrent architectures.\nBansal, A., Schwarzschild, A., Borgnia, E., Emam, Z., Huang, F., Goldblum, M., & Goldstein, T. (2022). End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking. Advances in Neural Information Processing Systems, 35, 20232-20242.\n2. T-BPTT: Truncated back propagation through time is a well-known prior approach to learning stable recurrent models without gradient collapse as the authors have included in their references. This also is not a novel contribution of this paper.\n3. There is a line of recent works in length generalization in recurrent models that is not mentioned in the related work. I request the authors to please look at this list of prior work in length generalization in RNNs, and add comment on how their proposed work is related to these highly related prior art.\ni) Bansal, A., Schwarzschild, A., Borgnia, E., Emam, Z., Huang, F., Goldblum, M., & Goldstein, T. (2022). End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking. Advances in Neural Information Processing Systems, 35, 20232-20242.\nii) Veerabadran, V., Ravishankar, S., Tang, Y., Raina, R., & de Sa, V. (2023). Adaptive recurrent vision performs zero-shot computation scaling to unseen difficulty levels. Advances in Neural Information Processing Systems, 36, 18132-18145.\niii) Goetschalckx, L., Govindarajan, L. N., Karkada Ashok, A., Ahuja, A., Sheinberg, D., & Serre, T. (2023). Computing a human-like reaction time metric from stable recurrent vision models. Advances in neural information processing systems, 36, 14338-14365.", "other_strengths_and_weaknesses": "NA. Please refer to my above review.", "other_comments_or_suggestions": "NA.", "questions_for_authors": "I have already conveyed my questions and concerns in other parts of the review. I thank the authors for their submission and look forward to reading their rebuttal.", "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": "IAFUHSUWdc", "reviewer": "Reviewer_8LQW", "summary": "The paper explores the reasons for limited length generalization of recurrent (mainly modern linear ones like mamba, GLA) neural networks.\nThe paper explores the hypothesis that this is due to \"unexplored states\" - i.e. the kind of states that occur after long context tend to be unfamiliar to models trained on shorter context (as such states are not attainable). The authors perform several empirical investigations surrounding this. The paper shows that the \"effective remembrance\" (a proposed measure that quantifies the influence of earlier tokens) is high for RNNs - and models struggling towards length generalization are highly affected by distant tokens. It shows training on short contexts can impede generalization further. All these suggests that the models struggle due to distribution shift of states that it encounter for longer sequence. \nTo counteract it the paper explores various post-training interventions:\n\n1. Random Gaussian noise to initialize states  to increase the diversity of states a model encounters.\n2. Fitted noise - from mean and variance of states occurring in longer contexts.\n3. State passing - Passing some state of a different (randomly chosen) sequence end point as an initial state for another sequence.\n4. Truncated Backpropagation Through Time.\n\nThe paper shows that random Gaussian and Fitted noise does enhance generalization to an extent for certain models but still limited because they don't necessary simulate realistic long-sequence states. \n\nHowever, state passing and truncated backpropagation do enable effective generalization.\n\nLater the paper also demonstrates maintenance of long-term dependence with Fitted noise for passkey retrieval task.\n\n## Update after rebuttal\n\nThe authors address most of my concerns. I increased the score accordingly.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed response and we appreciate that they find most our claims well supported, including our analysis on how state passing and TBTT enable length generalization because they simulate realistic states. We provide answers to their questions and observations:\n\n> RE: Does passkey involve length generalization?\n\nYes, passkey retrieval involves length generalization, as the models have only been finetuned on samples of length $T=2048$, yet solve the task at length up to $256k$.\n\n> RE: Effectiveness of state passing when finetuning in passkey, and zero shot performance in passkey of models post-trained with interventions\n\nOriginally, we decided to run the experiment with the fitted noise intervention, as in our opinion it is the more suited to this setting. As the reviewer mentions, it is somewhat unnatural to do state passing in this task, as it would be equivalent to providing a context containing a previous passkey. But we recognize that it is natural to wonder about the performance of state passing and TBTT, so here we also provide additional results for these two settings, where **overall state passing and TBTT bring some benefits to length generalization:**\n\nState passing finetuning in passkey: https://postimg.cc/vxwhrDXB\n\nZero shot of TBTT post-trained checkpoints: https://postimg.cc/rKkcmJG6\n\n> RE: Do the interventions achieve length generalization for other long context tasks?\n\nWe agree with the reviewer that perplexity and passkey retrieval might not be enough to assess the long context capabilities of the models, **so we have run experiments on two more long context tasks that are related to core aspects of language modeling**. Due to limited space, we refer to the response to Reviewer b5dM for more details on the tasks.\n\nResults for BABILong: https://postimg.cc/6TjDhBZ0\n\nResults for the synthetic copying task: https://postimg.cc/ctwryLQV\n\n**It can be seen that state passing and TBTT bring significant benefits in length extrapolation for these tasks**, for example by improving length extrapolation on sequences up to length 256k in BABILong. We also thank the reviewer for providing more related references on length generalization, which we will include in the paper as related work and avenues for future work.\n\n> RE: Scope of term \"recurrent models\"\n\n**We have added new results for another architecture, RWKV-v6, and we have observed the same phenomenon where the models trained on short context with zero-initialized states diverge, whereas post-training with state passing enables length generalization**\n\nhttps://postimg.cc/zLkhvF2x\n\nWe hypothesize that other modern recurrent models will exhibit a similar behavior, but we agree with the author that all the studied models fall under the category of linear recurrent architectures that accept a formulation in terms of SSMs. Thus, we will change the wording in the final version to more clearly set the scope of the claims to the linear recurrent models for which we show experiments, and cite xLSTM as a different recurrent model which might have other behavior.\n\n> RE: the \"unexplored state hypothesis\" being obvious\n\nIndeed we agree with the reviewer that the hypothesis feels very intuitive in hindsight, but we believe our work is very relevant to the community for two reasons. \n\nFirstly, many recent works deal with length generalization by intervening on the architecture mechanism (see specific examples in Section 6), or by proposing new architectures and comparing their performance beyond the training context (for example see [1] or [2]). In contrast, our work: **(1) proposes the \"unexplored states hypothesis\" as a framework to understand the length generalization of recurrent models by reasoning about the distribution of states, and not the architecture mechanism**; **(2) systematically analyzes four training interventions which are motivated by the \"unexplored states hypothesis**\"; and **(3) shows that the interventions enable length generalization without changes to the architecture.**\n\nSecondly, even though it is natural to think that models fail to length generalize because they are out of distribution somehow, **we believe that it is a valuable contribution to propose training with several types of non-zero initialized states to fix this issue**. More specifically, we show that length generalization is enabled when the model is trained with initial states that resemble the states that are attained when processing long sequences. **We believe this finding is both intuitive and interesting, which we consider a strength of our work and hope makes a shift in the community towards simple (yet underutilized) training interventions with non-zero initialized states**.\n\n\n[1] Yang, S. et al. (2025). _Gated Delta Networks: Improving Mamba2 with Delta Rule_. In _The Thirteenth International Conference on Learning Representations_.\n\n[2] Sun, Y. et al (2024). _Learning to (Learn at Test Time): RNNs with Expressive Hidden States", "claims_and_evidence": "Mostly the claims are supported. \n\nFew points I am a bit critical of:\n\n1. The passkey retrieval task seems to be shown only for Fitted noise - which did not appear as an effective solution in Pile. So I would be curious:\n     1. How does other methods particularly BPTT and state passing do on the task? If we are passing the state of different passkey-retrieval context - wouldn't that confuse the model? \n     1. Does passkey retrieval also involve length generalization tests? If not, it's also desirable to see if the models can generalize in situations where actual long range dependency exists and is enforced (unlike general language modeling). \n\nThe paper also uses the general term \"recurrent\" to scope out the kind of models it explores instead of a more specific term like \"linear recurrent\". It doesn't seem to consider non-linear recurrent models (which have modern variants as well such as xLSTM [1]). Could be good to limit the scope more explicitly by specifying this if that's the intention. \n\n\n[1] xLSTM: Extended Long Short-Term Memory - Beck et al. Neurips 2024", "methods_and_evaluation_criteria": "Yes. Some potential misses are discussed above.", "theoretical_claims": "As far as I have checked, no fundamental issues.", "experimental_designs_or_analyses": "No strong issues stand out to me besides the one I already discussed. \n\nOne issue I may mention (which is broader than this paper - so I will not particularly penalize this paper for this) is as follows: \n\nOne point is that the context of length generalization matters. Simply length generalization in language modeling may not give the full picture because in the dominant distribution of natural language contexts -- it could be very well possible long-range dependencies and state tracking is not required. Even some popular solutions like AliBi and such seem to work by having more of a \"forgetting\" bias towards the past. \n\nIt's unclear, how well length generalization can be achieved in other contexts where the final result is very sensitive to simple changes in the past. Passkey retrieval can be one example, but there are other contexts as well one could explore like logical inference and ListOps length generalization - examples of such exploration include [1,2,3,4] among others.\n\nEarlier SSMs like S4D were shown poor generalization there [4].\n\n[1] Modeling Hierarchical Structures with Continuous Recursive Neural Networks - Ray Chowdhury et al. ICML 2021\n[2] Ordered Memory - Shen et al. NeurIPS 2019\n[3] The Neural Data Router: Adaptive Control Flow in Transformers Improves Systematic Generalization - Csordas et al. ICLR 2022\n[4] Recursion in Recursion: Two-Level Nested Recursion for Length Generalization with Scalability - Ray Chowdhury et al. NeurIPS 2023", "relation_to_broader_scientific_literature": "The paper is connected to the lineage of linear recurrent models like RWKV, GLA, Mamba. Prior papers on linear recurrent models have already attempted different investigations regarding length generalization. This papers continues that direction but focusing more on the unexplored state hypothesis and non-architectural interventions - like BPTT and passing states from other models. \n\nKey contributions seem to be theoretical formulation of a hypothesis why RNNs struggle to generalize and some initial empirical investigations both to support the hypothesis and explore some initial countermeasures to enable length generation.", "essential_references_not_discussed": "Not as familiar with this specific direction of length generalization in linear RNN.\n\nHowever, this paper:\n\nThe Illusion of State in State-Space Models - Merill et al. ICML 2024\n\nSeems relevant to a degree. This may suggest some root issues in SSMs - in state tracking; and that could also limit length generalization in certain contexts. Simply hacking the initial state may not be the resolution in this case. \n\nxLSTM paper suggests better state-tracking - and should count as a model to further investigate if the title and abstract is kept general - i.e. making claims about \"recurrent\" neural networks and not \"linear recurrent\" neural networks more specifically.", "other_strengths_and_weaknesses": "**Strengths**\n\nOverall, the paper is well-written. Despite any limitations it is a decent exploration of an idea. Exploration of the different post-training intervention techniques are interesting and can work as \"quick fixes\" for better generalization in the at least some contexts in the interim.\n \n\n**Weakness**\n\n1. Besides the things mentioned above, to an extent the hypothesis also feels a bit \"obvious\". I am not sure how much scientific weight to give to it. \nSince recurrent models have the same parameters at every time step, in a sense, obviously, it should obviously generalize if the states are similar. And the corollary of that would be that if generalization is failing - the issue would be overfitting in the distribution of seen states. \n\n1. It's not clear if we are getting an effective single intervention strategy that works both in general language modeling, and also in length generalization in contexts with long-range dependencies - passkey, listops etc. (part of which - especially for tasks requiring state-tracking - this might be also simply impossible for linear RNNs - given Merill et al's work cited above)", "other_comments_or_suggestions": "n/a", "questions_for_authors": "1. Does passkey retrieval involve length generalization tests?\n1. Are the other strategies BPTT/State-passing tested on passkey retrieval?\n1. How about the zero-shot performance of the post-training intervention models on passkey?", "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": "J8LSAb3C1H", "reviewer": "Reviewer_b5dM", "summary": "The authors study length generalization in recurrent models. They begin with an empirical analysis of length generalization failures for Mamba v1 / v2 and gated linear attention. They define a new metric, \"Effective Remembrance\", to quantify the influence of the context prefix on a model's predictions, and show that Mamba v2 has high Effective Remembrance for early tokens. They hypothesize that length generalization failures are due to incomplete exploration of the recurrent state distribution during training and propose four state initializations to mitigate this. They provide empirical results showing that post-training using these initialization schemes substantially improves length generalization measured in terms of perplexity.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for its thoughtful analysis of the paper and concrete suggestions, and we are encouraged to see that they found the experiments sounds and the paper easy to follow. We provide responses to their discussion:\n\n> RE: Perplexity as an evaluation metric\n\nWe understand the reviewer's concern that perplexity on its own might not be sufficient to asses the capabilities of the models. Thus, we have run two additional evaluations.\n\n(1) **BABILong** [1] : BABILong is a benchmark that tests the ability of the model to process a long context and respond to a natural language understanding question based on some facts in the context. The results are shown in the following link:\n\nhttps://postimg.cc/6TjDhBZ0\n\n**It can be seen that TBTT and especially state passing greatly help the model achieve a better performance than the baseline, improving from ~9% to ~24% in a context length of $256k$**.\n\n(2) **Synthetic copying** [2] : We have also evaluated length generalization on the synthetic copying task, which consists in copying an arbitrary sequence of tokens. The results are shown in the following link:\n\nhttps://postimg.cc/ctwryLQV\n\n**It can be observed that state passing greatly helps length extrapolation, improving the accuracy in sequences three times longer that those seen during training from 27% to 47%.** We believe these new experiments highlight the length extrapolation capabilities of the interventions on the initial state and hope they help address the reviewer's concern on perplexity as an evaluation metric.\n\n[1] Kuratov, Y. et al.(2024). _BABILong: Testing the limits of LLMs with long context reasoning-in-a-haystack_. In _The Thirty-eighth Conference on Neural Information Processing Systems Datasets and Benchmarks Track_.\n\n[2] Jelassi, S. et al. (2024). Repeat After Me: Transformers are Better than State Space Models at Copying. <i>Proceedings of the 41st International Conference on Machine Learning</i>\n\n\n> RE: Supporting the unexplored states hypothesis by comparing state distributions in short and long contexts\n\nWe thank the reviewer for this suggestion, which we have added to the final version of the paper. **We measured the distribution of the state depending on the sequence position, and found that the standard deviation (across all elements of the state) increases as the sequence position increases. Thus, when processing longer sequences the model encounter states from a distribution that it has not seen during training. The results are in the following link:**\n\nhttps://postimg.cc/HcpL3yWP\n\n**Moreover, in the figure we show that our state passing intervention fixes this issue by producing states whose distributions do not vary so much based on the sequence position.** This effect can also be observed in specific heads of layers:\n\nhttps://postimg.cc/LJ0vZsZd\n\nThis shift in distribution is correlated with performance: the models' perplexity diverges after the training context because it encounters a state that has not been seen during training (see Figure 1a). **Thus, we believe these new findings reinforce our unexplored states hypothesis.**\n\n> RE: Is lower Effective Remembrance always better?\n\nYes, we absolutely agree that the \"right\" level of Effective Remembrance is not necessarily as low as possible. A certain level of Rffective Temembrance is needed for the model to remember tokens from the past (if it were zero at time $t$, it would be \"effectively\" ignoring the tokens in positions $[0,t]$, which is undesirable when long range dependencies are needed). However, the Rffective Remembrance curves of the baseline models shown in Figure 1 are clearly wrong, given that they are too influenced for tokens that are far away in the past. The state passing intervention yields more reasonable Effective Remembrance curves, where the model places more focus on the recent context, which is correlated with better performance.\n\nMore generally, we believe Effective Remembrance can be a useful tool when applied to other settings and we are eager to use it in future work. Several recent works point out that linear recurrent architectures may fail to achieve tasks like associate recall or copying because they cannot remember past tokens or store previous context into the state (compared to transformers, which do not compress the past context). This could be verified by showing that Effective Remembrance is too low in these tasks (which would be undesirable here, as the reviewer suggested). However, in the context of our work (language models failing to length generalize), the opposite is actually the case: Effective Remembrance is too large. We hypothesize this is due to the states being overparametrized and due to the models overfitting to the states that arise when processing short sequences (which motivate our unexplored state hypothesis and the the use of non-zero initial states to train on a wider distribution of states and fix the issue).", "claims_and_evidence": "The authors claim that:\n1. Models fail to generalize when they reach states that were not explored during training (the \"unexplored state hypothesis\")\n2. Length generalization in recurrent models can easily be improved by post-training with simple interventions on the state initialization.\n\nThe evidence for 2 is quite good given their experimental results, as long as you're comfortable characterizing length generalization in terms of perplexity. The evidence for 1 does not seem as strong. There is no direct investigation of the unexplored state hypothesis, even though this would seem to be straightforward - e.g by comparing the distribution of states under short vs long contexts and correlating their discrepancy with some measure of performance.", "methods_and_evaluation_criteria": "Effective remembrance is an intuitive way to measure the influence of a context prefix. The interventions described for the state are clear, though there is not much effort to distinguish why you'd prefer one over the other, apart from the empirical results.\n\nIn terms of evaluation, length generalization is defined and evaluated almost entirely in terms of perplexity. While this is standard for the literature, it's worth noting that there meaningful debate as to whether perplexity is really informative for long-context performance (see below). The authors also study performance on a passkey task, but there are many other instances of long-context tasks, or even whole benchmarks, that could provide an alternative set of evaluation criteria for length generalization.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The experiments are sound overall. It might have been useful to include a \"negative control\" for the post-training, perhaps by post-training on the same amount of data with the standard zero-initialization for the state.", "relation_to_broader_scientific_literature": "Length generalization is a long-running topic in sequence modeling, including for recent architectures like Mamba. Recurrent models have long been studied in terms of the statistical properties or dynamics of the state, including aspects that may lead to better performance over long sequences (\"long memory\") or more stable behavior.", "essential_references_not_discussed": "There is an active debate on whether perplexity is informative for long-context performance. Recent works include [1, 2], both of which observe low correlation between perplexity and performance on benchmarks such as LongBench [3]. It would be useful to cite this discussion and to hear the authors' thoughts given their approach in this paper.\n\n[1] Fang, L., et al. (2024). What is Wrong with Perplexity for Long-context Language Modeling?. arXiv preprint arXiv:2410.23771.\n\n[2] Hu, et al. (2024). Can Perplexity Reflect Large Language Model's Ability in Long Text Understanding?. arXiv preprint arXiv:2405.06105.\n\n[3] Bai, Y., et al. (2023). Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508.", "other_strengths_and_weaknesses": "The paper is well written and easy to follow. The claims are stated clearly and the methods are well defined.", "other_comments_or_suggestions": "Typos:\n- $x_t$, not $x_T$, in line 114?\n- I'd suggest using \"et al.\" for the 1.5-page Llama-3 citation", "questions_for_authors": "- Could the authors share their thoughts on using perplexity to evaluate length generalization? \n- Why not test the unexplored state hypothesis directly by comparing state distributions in short and long context scenarios?\n- The authors' effective remembrance analysis essentially concludes that the recurrent models (and Mamba v2 in particular) are too strongly influenced by observations in the distant past. This seems at odds to some extent with the main objective of modern recurrent architecture development, which is to produce models that are able to retain information over very long sequences. Moreover, it would seem that the \"right\" level of effective remembrance depends on the data generating process, so that less is not always better. Is there a balance to be struck here?", "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": "Uv9ExsurEm", "reviewer": "Reviewer_5H8P", "summary": "The paper investigates why recurrent models fail to generalize to sequence lengths beyond their training context and proposes methods to improve length generalization. They also propose a metric (Effective Remembrance) which basically captures the difference in a model's next token distribution at a point, when considering two different context lengths (starting at t=0 or t=some other index in the past >0).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are encouraged to see that the reviewer thinks that our unexplored states hypothesis is well supported and that our interventions on the initial state are useful to achieve length generalization. We provide responses to their questions:\n\n> RE: Results on long-context generalization benchmarks besides passkey retrieval to strengthen the paper\n\nWe followed the suggestion of the reviewer and **present new results for two long context benchmarks.** \n\nThe first one is BABILong [1], a challenging benchmark which tests both the common sense understanding of the models as well as its ability to capture long range dependencies in text. We present results for the baseline and our interventions in the following link:\n\nhttps://postimg.cc/6TjDhBZ0\n\n**It can be observed that state passing enhances the length extrapolation capabilities of the model, improving performance from ~9% to ~24% in sequences of length $256k$** (we recall that the model is trained and finetuned on sequences of length $2K$). This reinforces our claim that state passing is not only useful to fix the diverging perplexity of established language models, but also improves helps length generalization by exposing them to more initial states during training and thus enhancing their ability to handle long contexts.\n\nThe second benchmark is Synthetic Copying [2], which consists in copying an arbitrary sequence of tokens. The results are shown in the following link:\n\nhttps://postimg.cc/ctwryLQV\n\n**It can be seen that state passing greatly helps length extrapolation, improving the accuracy in sequences three times longer than those seen during training from 27% to 47%.** We believe these new experiments highlight the length extrapolation capabilities of the interventions on benchmarks that are closer to real world data, as the reviewer suggested.\n\nFinally, regarding the Longbench benchmark specifically, unfortunately we lack the computational resources to train and evaluate models that can achieve decent performance on this benchmark. Even models of more than 7B parameters are barely better than random guessing (see Table 2 in [3]), and we have focused our work on models in the 1B scale and below. However, we find this benchmark very relevant to our methods and will include it as an avenue for future work in the final version of the paper.\n\n[1] Kuratov, Y. et al.(2024). _BABILong: Testing the limits of LLMs with long context reasoning-in-a-haystack_. In _The Thirty-eighth Conference on Neural Information Processing Systems Datasets and Benchmarks Track_.\n\n[2] Jelassi, S. et al. (2024). Repeat After Me: Transformers are Better than State Space Models at Copying. <i>Proceedings of the 41st International Conference on Machine Learning</i>\n\n[2] Bai, Y. et al. (2024). LongBench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics. Association for Computational Linguistics.\n\n> RE: The difference between state passing and TBTT resides in whether the initial state is related to the context of the current sequence\n\nThat is correct: the difference between state passing and TBTT is that in state passing the initial state is unrelated to the given sequence, whereas in TBTT the initial state corresponds to prior context of the sequence. From an implementation point of view, the only difference is that in state passing the dataset is shuffled, whereas in TBTT the dataloader carefully orders the samples so that when passing a final state to the next sequence, such final state corresponds to the prior context of the sequence. As the reviewer points out, in principle it could be possible that in state passing the final state would correspond to the correct context, but the probability negligible given that the datasets are so large.\n\nEven though in practice the implementation differences are small, in theory the methods are doing very different things. Both methods can be seen as sampling initial states from a given distribution. State passing samples initial states from the distribution of final states (for convenience, we simply take the final state of the previous batch, which is a good enough approximation because the batches are unrelated). In contrast, TBTT samples states from a degenerate distribution, which specifically consists of the final state of the previous context of the sequence. A very interesting result of our paper is that TBTT is sufficient but not necessary for length generalization, given that state passing also works (or even sampling from random or fitted noise can be good enough). Besides being a practical and efficient way of enabling length generalization, we believe our results for these interventions shed light on the core issue with length generalization: models fail to length generalize when they are not trained on the distribution of states attained when processing a sequence (our \"unexplored states hypothesis\").", "claims_and_evidence": "Most claims are well-supported, particularly the core hypothesis that unexplored states lead to generalization failure and that state-based training interventions can fix the issue.", "methods_and_evaluation_criteria": "The paper presents comparisons using only the perplexity metric on The Pile dataset (so the task is next token prediction), and on a very synthetic long context task (passkey retrieval). It would be good to show the results on a long context task derived from real world data (maybe LongBench or even something simpler that requires 2 or 3 hop reasoning).", "theoretical_claims": "No.", "experimental_designs_or_analyses": "-", "relation_to_broader_scientific_literature": "The paper shows that initialising the state of a sequence model with noise that has been fit to the typical distribution of how the states are at the end of a context window (say N), helps to train/finetune models with context length N, but genralise to 2N (since now the model has been trained while seeing what state distributions sort of look like at the end of 2N steps of context even if the gradient was not back propagated 2N steps). This tries to mimic what was effectively done when other works trained with TBTT and carried forward the state into a new rollout starting from the end point of the previous rollout of a sequence model - but this time with fitted noise or state passing - fitted noise intervention works well for the small models but not a bigger model, state passing works well for both.", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "I think the proposed approach of trying to understand train-time randomisations/methods to encourage length generalization makes sense and is useful. Presenting results on a long-context generalization benchmark that is not very toy or synthetic would significantly strengthen the paper.", "other_comments_or_suggestions": "-", "questions_for_authors": "Am I correct in assuming that state passing is different form TBTT since the state value (say s) used to intialise a new rollout at training iteration i, for batch element say 0 - is not related in any way to previous states that might have been seen in the prior context that corresponds to batch element 0?\nSo if we have a context of length 2N, but we are only able to train for context lengths of N, it is not the case that we are using the states at the end of N to initialize the rollout when we go from N+1 to 2N in the next batch (because the states at the end of the batch are shuffled - so while in some cases they might actually correspond to the correct context/trajectory - it many cases it will not?).", "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:

Meta-review: The paper investigates length generalization in recurrent networks. The authors showed that performance of RNNs drops significantly beyond training context length, and proposed several interventions leading to significant improvement in length generalization. Reviewers appreciate the simple intervention for practical reasons and the empirical validation. Some raised their scores after the rebuttal as their main concerns are addressed. There are some remaining questions, e.g., using perplexity as the measurement, missing discussion/comparison to prior works, but overall reviewers agreed on a more positive outcome, considering the merits outweigh the limitations.
Decision: Accept (poster)