ArXiv: 2410.07145
🎯 Pitch
Mamba models fail on long contexts not because they forget too much, but because they never learned to forget in the first place—their recurrent states are so oversized for typical training lengths that simply storing everything minimizes loss, bypassing the need for a functional forgetting mechanism. The minimum training length required to induce forgetting scales linearly with state size, but once that threshold is crossed, the maximum recallable context length scales exponentially with state size, unlocking dramatically longer effective contexts.
1. Executive Summary
This paper analyzes why Mamba-based recurrent language models degrade sharply when processing contexts longer than their training length, using Mamba-2 checkpoints evaluated on language modeling and passkey retrieval. The authors identify the root cause as the inability to forget—the model's failure to learn robust memory decay mechanisms (quantified by per-token retention strength α_t and retrieval error accumulation in the recurrent state)—which leads to interference between token representations and incoherent outputs. They attribute this to state overparameterization, where the recurrent state is excessively large relative to the training context length, enabling the model to minimize loss by retaining all information rather than learning to discard irrelevant tokens; they establish that a forget threshold exists and scales linearly with state size (), that forgetting is learned only when training length exceeds this threshold, and that the maximum recall context length for passkey retrieval scales exponentially with state size—with a 370M-parameter Mamba-2 achieving near-perfect retrieval at 256K tokens after continued pre-training, outperforming similarly sized transformers. The findings establish that robust long-context performance in RNNs emerges only when training contexts are long enough to force the model to learn forgetting rather than simply storing everything in oversized states.
2. Context and Motivation
The Core Problem: Recurrent Models Cannot Generalize Beyond Their Training Length
The fundamental question this paper tackles is: why do state-of-the-art recurrent language models like Mamba-2 catastrophically fail when confronted with sequences longer than those seen during training? This is not an incremental performance degradation — the models produce incoherent outputs, lose the ability to recall information from any position in the context, and exhibit language modeling loss spikes that render them unusable (Figure 1, Figure 2). For a 370M-parameter Mamba-2 trained on 8K-token sequences, the loss at position 24K roughly doubles compared to positions within the training window (Figure 1), and passkey retrieval accuracy drops from near-perfect at ≤8K to essentially zero at 16K+ contexts (Figure 2).
This gap is significant because recurrent architectures have been widely promoted as the efficient alternative to transformers for long-context processing. Unlike self-attention, which scales quadratically with sequence length (), RNNs maintain a fixed-size state that encodes all historical context, giving them constant per-token time and space complexity during inference. This makes them dramatically more efficient for streaming applications, on-device deployment, and processing extremely long documents. However, this efficiency is worthless if the models cannot reliably process sequences beyond their training length — which in practice is often quite short. Most production Mamba checkpoints (130M through 2.8B parameters) are trained on only 8K-token contexts (Table 1, Figure 12), far shorter than the 128K–1M token contexts that transformer-based models routinely handle.
The practical implications are immediate: any application that requires processing long documents, maintaining extended conversations, or performing retrieval over large contexts — essentially any deployment scenario beyond short-form generation — cannot rely on these RNNs without either expensive retraining on longer sequences or a fundamental architectural fix. The paper's results suggest that the training context length required to achieve robust generalization scales linearly with the state size (), meaning that scaling model capacity (and thus state size) makes the problem worse, not better — exactly the opposite of what we expect from scaling laws for transformer models.
Why the Problem Matters: The Tension Between Efficiency and Capability
The gap this paper addresses sits at a critical intersection in the current LLM landscape. Transformer-based models (GPT-4, Gemini, Llama 3) have demonstrated impressive long-context capabilities — Gemini 1.5 processes up to 10M tokens, MiniMax-01 handles 4M tokens — but they pay a steep computational price. The KV cache for a single transformer layer at sequence length with hidden dimension requires floating-point numbers in memory, and the attention computation itself scales as in the naive case. For a 7B-parameter model processing a 1M-token sequence, this becomes prohibitively expensive in both memory and FLOPs.
Recurrent models promise to resolve this tension. Their state size is fixed: a Mamba-2 model with hidden dimension has a total state size of parameters across all layers (Appendix A.1), which is equivalent to a transformer's KV cache for 128 tokens. This means that at long sequence lengths, the recurrent model's memory footprint and per-token cost remain constant while the transformer's grow without bound. The theoretical advantage is enormous — but only if the recurrent model can actually utilize long contexts effectively.
The paper's findings cut to the heart of this tradeoff. The issue is not simply that Mamba-2 is undertrained or poorly tuned; it is that the architecture's core mechanism — the recurrent state update rule with learned memory decay — does not learn to forget when the state capacity exceeds the information content of the training context. This is a fundamental design tension: we want large states to store more information, but larger states allow the model to "cheat" during training by retaining everything rather than learning a meaningful forgetting policy. The result is that scaling up model size (which increases state size) actually degrades length generalization (Figure 1, comparing 130M through 1.3B models), creating a perverse scaling dynamic where larger recurrent models become worse at the very task they are designed for.
Prior Approaches and Where They Fall Short
The paper identifies several categories of prior work that touch on this problem, none of which have directly addressed the forgetting mechanism as the root cause.
Length Generalization in Transformers. The vast majority of recent length generalization research has focused on transformer architectures, specifically on positional encoding schemes. Transformer models with relative positional encodings (RoPE, ALiBi) can theoretically process arbitrarily long sequences, but they exhibit significant performance degradation beyond the training length due to out-of-distribution positional values. Solutions have included modifying the positional encoding to extrapolate beyond training — YaRN (Peng et al., 2024b) scales frequencies, LongRoPE (Ding et al., 2024) searches for optimal rescaling factors, and SelfExtend (Jin et al., 2024) uses grouped attention — with some approaches achieving training-free length generalization to moderate extents. However, these approaches are fundamentally architectural solutions for the attention mechanism; they do not transfer to recurrent models, which have no positional encoding and no KV cache. The failure mode in RNNs is categorically different: it arises from the state dynamics, not from positional information.
Prior Investigations of Mamba's Length Generalization. Several works have empirically documented Mamba's performance drop beyond training length but have attributed it to different causes. Jelassi et al. (2024) showed that Mamba struggles to copy information from context unless its state size grows linearly with context length — a capacity limitation framing. Arora et al. (2024a) analyzed the associative recall capabilities of transformers versus RNNs, finding that RNNs require explicit state capacity for each token pair they need to recall. These works correctly identify that copying and associative recall are challenging for fixed-size states, but they frame the problem as one of insufficient capacity — the state is too small to hold everything. The current paper's insight inverts this: the problem is that the state is too large relative to the training context, allowing the model to avoid learning forgetting and thus failing when the context exceeds training length. This is a qualitatively different diagnosis with fundamentally different implications (lengthening training is the solution, not enlarging states).
Wang et al. (2025) and Over-Smoothing. The most closely related prior work discusses over-smoothing in state space models caused by the memory decay term. In this framing, the exponential decay in the recurrent update causes token representations to blend together, losing the distinctiveness needed for precise recall. However, as the paper notes in Section 6, this analysis explores a regime where recency effects should dominate — recent tokens should be clearly distinguishable from earlier ones — yet the retrieval failure occurs even for very recent tokens when the total context is long (Figure 2 shows zero accuracy at the rightmost positions for 16K+ contexts). This cannot be explained by over-smoothing alone; it points to active interference from earlier tokens that should have been forgotten but were not.
Engineering Heuristics for Mamba Length Extension. Some practitioners have developed ad-hoc fixes for Mamba's length generalization. LongMamba (Zhang, 2023) divides the discretization term by a constant (e.g., 0.5), which makes the memory decay closer to 1 — effectively making the model retain more information by decaying less. DeciMamba (Ben-Kish et al., 2024) proposes heuristics for identifying which heads to modify. However, as the paper demonstrates in Figure 4, these approaches either compromise short-context performance (LongMamba degrades within the training window) or simply fail to address the root cause. The paper's sliding window and RRI (Reduced Memory Retention and Insertion) interventions show that the opposite approach — inducing more forgetting — can improve length generalization, strongly suggesting that the field has been looking at the problem backward.
Training Length and State Capacity. The only prior work that directly connects training length to recurrent model behavior is the compute-optimal context size analysis by Buckman & Gelada (2024), which argues that training context length should be chosen based on the data's long-range dependency structure. However, this work does not connect training length to the recurrent state size or to the learning of forgetting mechanisms. Yang et al. (2024a)'s work on GLA (Gated Linear Attention) mentions using truncated backpropagation through time to extend effective context length for RNNs, but does not analyze why this helps or how it interacts with state capacity.
How This Paper Positions Itself
The paper's central claim is a unified explanatory framework that connects three previously separate phenomena: (1) the sharp performance degradation of Mamba-2 beyond training length, (2) the model's failure to learn forgetting despite having a built-in decay mechanism, and (3) the linear relationship between state size and the training length required for robust forgetting. The key conceptual move is to reframe the problem from one of state capacity (the state cannot hold enough information) to one of state overparameterization (the state is so large that the model never needs to learn how to forget during training).
This reframing is anchored in a specific mechanistic hypothesis about how the recurrent state accumulates and how forgetting interacts with retrieval. The state at time is a weighted sum over all past insertions:
When retrieving the memory inserted at time by querying with , the output includes the target signal scaled by plus an error term from all other tokens:
As the context length grows, the vectors cannot remain mutually orthogonal (the state dimension is finite), so the error term accumulates. The model has a knob to control this: produce small values (strong decay) to quickly diminish the contribution of earlier tokens to the state. The paper's key empirical finding is that the model does not learn to turn this knob appropriately when trained on contexts shorter than the state's information capacity — it keeps extremely close to 1 (Figure 3 shows some heads with , meaning per-step decay factors are essentially 1.0), attempting to retain everything and consequently suffering catastrophic interference when the context extends beyond training.
The paper positions this not as a failure of the Mamba-2 architecture per se, but as a failure of current training practices — specifically, training on contexts that are too short relative to the model's state capacity. The linear relationship (Figure 11) provides a concrete, actionable guideline: to train a Mamba-2 model that generalizes to arbitrary-length contexts, one must use training sequences whose length exceeds this threshold. The exponential scaling of recall capacity with state size (Figure 9: ) provides an optimistic upper bound: once the forgetting mechanism is properly learned, these models can theoretically recall information from extremely long contexts, since the amount of distinct information in repetitive "filler" text is constant and the state can discriminate between an exponential number of patterns.
This positions the paper at the intersection of architecture design (the recurrent update rule), training methodology (context length selection, data filtering), and empirical scaling laws (the linear and exponential relationships). It does not propose a new architecture or a new training algorithm; instead, it provides the diagnostic framework and empirical evidence needed to understand when and why current approaches fail, and what the path forward looks like. The contribution is primarily analytical — establishing the forgetting mechanism as the bottleneck, the state overparameterization hypothesis as the explanation, and the linear scaling law as the practical guideline — with the implication that future RNN designs should account for the interplay between state size, training length, and forgetting as a first-class design constraint rather than an afterthought.
3. Technical Approach
This is primarily an empirical analysis (mechanistic investigation) paper whose core idea is that Mamba-2's catastrophic failure on sequences longer than its training length is caused by the model's inability to learn a robust forgetting mechanism, which in turn stems from training on contexts that are too short relative to the recurrent state's capacity—a phenomenon the authors term state overparameterization.
3.1 Reader Orientation
The "system" being analyzed is a pre-trained Mamba-2 language model—a recurrent neural network that processes tokens sequentially, maintaining a fixed-size hidden state that encodes all past context. The problem it solves (or rather, the problem the paper diagnoses) is that this model produces incoherent, high-loss outputs and loses all retrieval capability when the input sequence exceeds the length it was trained on (typically 8K tokens). The solution's "shape" is not a new architecture or training algorithm, but rather a diagnostic framework and empirical scaling law: by identifying that the root cause is the model's failure to learn memory decay (forgetting) during training, and by establishing that the training length must exceed a threshold that scales linearly with the state size, the paper provides both an explanation and a concrete prescription for fixing the problem—train on longer sequences.
3.2 Big-Picture Architecture (Diagram in Words)
The analysis framework has four major components that interact as follows:
-
Pre-trained Mamba-2 Language Models (130M–1.3B parameters): The base objects of study. These are recurrent models trained on 8K-token sequences with next-token prediction. Each model maintains a fixed-size hidden state
$h_t \in \mathbb{R}^{N \times P}$per head that accumulates information across time steps via a learned update rule involving memory decay$\alpha_t$and new information insertion$B_t x_t$. The key observable is that these models exhibit sharp loss increases and retrieval failures when$t$exceeds the 8K training length. -
Diagnostic Probes (retention strength measurement, state statistics collection, intervention experiments): A suite of analytical tools for inspecting what the model has learned. These include:
- Computing cumulative retention strength
$\alpha_{1:t}$of the first token over time (Section 3.2.1) - Collecting mean and variance statistics of recurrent states
$h_t$and convolutional states when processing controlled prompts like repeated newline characters (Section 3.3, Appendix G) - Artificially modifying the update rule at inference time to induce forgetting (RRI and sliding window interventions, Section 3.2.2)
- Measuring passkey retrieval accuracy as a function of context length and answer position (Section 2.2, Appendix B)
- Computing cumulative retention strength
-
Controlled Training Experiments (varying state sizes, training lengths, and data budgets): To establish causal relationships, the authors train models from scratch or continue pre-training from official checkpoints while systematically varying:
- Model size (and thus state size
$N_S$, ranging from 0.8M to 19.3M parameters) - Training context length
$T_{\text{train}}$(from 512 to 256K tokens) - Amount of training data (10B, 20B, 40B tokens for the 512-length experiments) This enables measuring how the forgetting behavior changes with these variables.
- Model size (and thus state size
-
Empirical Scaling Law Extraction: From the controlled experiments, the authors fit two mathematical relationships:
$T_{\text{forget}}$as a function of state size$N_S$: the minimum training length needed for the model to learn robust forgetting (Section 5.1–5.2)$T_{\text{recall}}$as a function of state size$N_S$: the maximum context length from which the model can accurately retrieve a 5-digit passkey after being trained with$T_{\text{train}} > T_{\text{forget}}$(Section 5.3)
The flow of analysis is: start with pre-trained models → observe failure → hypothesize inability to forget → measure retention strength and state statistics → confirm via intervention experiments → propose state overparameterization as explanation → test via controlled training with varying state sizes and training lengths → extract quantitative scaling laws → validate that models trained above the forget threshold indeed exhibit robust length generalization and long-range recall.
3.3 Roadmap for the Deep Dive
- First, the Mamba-2 recurrent update rule in full detail (Equations 1–5, 11), because every subsequent diagnostic and intervention operates on this mechanism—we need to understand what
$\alpha_t$,$B_t$,$h_t$, and$\Delta_t$are and how they interact before we can discuss what goes wrong. - Second, the formal decomposition of the state as a weighted sum and the retrieval error analysis (Equations 6–7), since this provides the mathematical framing for why forgetting matters—it shows precisely how interference from insufficiently decayed earlier tokens corrupts recall.
- Third, the two diagnostic methods for measuring whether the model has learned forgetting: cumulative retention strength
$\alpha_{1:t}$(Section 3.2.1) and state distribution statistics (Section 3.3), along with what "normal" versus "pathological" values look like. - Fourth, the two intervention experiments that artificially induce forgetting—RRI scaling and the sliding window mechanism (Section 3.2.2)—including their specific hyperparameter choices and what their success/failure implies causally.
- Fifth, the controlled training setup for establishing the forget threshold and maximum recall length (Section 4, Section 5), including data processing, model configurations, hyperparameter sweeps, and the evaluation protocols for determining
$T_{\text{forget}}$and$T_{\text{recall}}$. - Sixth, the passkey retrieval evaluation protocol (Section 2.2, Appendix B), which serves as the primary downstream task for assessing whether forgetting has been learned and how far recall extends.
3.4 Detailed, Sentence-Based Technical Breakdown
The Mamba-2 Recurrent Update Rule
The Mamba-2 architecture consists of $L$ layers, each containing $H$ heads computed in parallel. The layer output is the sum of head outputs. For a single head, the computation at time step $t$ is governed by two core equations (simplified notation from Section 2.1; full details in Appendix A, Equations 9–15):
Query rule:
where $y_t$ is the head's output vector at time $t$, $C_t \in \mathbb{R}^{P \times N}$ is a query matrix that reads from the state, $h_t \in \mathbb{R}^{N \times P}$ is the recurrent hidden state at time $t$, and $P$ is the head dimension (always 64 in Mamba-2). This equation determines how memory is read: the query matrix $C_t$ selects which stored information to extract from the state.
Update rule:
where $h_{t-1}$ is the state from the previous time step, $\alpha_t \in \mathbb{R}$ is a scalar memory decay (retention strength) at time $t$ that controls how much of the past state is preserved, $B_t \in \mathbb{R}^{N \times 1}$ is an insertion vector that determines where in the state space new information is written, and $x_t \in \mathbb{R}^{1 \times P}$ is the new content to be inserted at time $t$.
What it computes: The update rule is a first-order linear recurrence. At each time step, the previous state $h_{t-1}$ is scaled by $\alpha_t$ (decaying old memories) and then a new outer product $B_t x_t$ is added (inserting new information). The state $h_t$ grows as a running weighted sum of all past insertions. The output $y_t$ is produced by multiplying the query matrix $C_t$ against the current state.
Why this form: The multiplicative decay $\alpha_t$ (rather than additive or gated) ensures that the contribution of very old tokens is exponentially smaller than recent ones, which is the standard solution to the vanishing/exploding gradient problem in RNNs. The additive insertion via outer product $B_t x_t$ (rather than a nonlinear transformation of $h_{t-1}$ and the input) makes the state a linear function of past inputs, which enables the model to be expressed as a weighted sum (Equation 6) and makes the analysis of memory retention tractable. The separation into a scalar decay per head (rather than a full matrix decay) keeps the parameter count manageable while giving each head independent control over its forgetting rate.
The key parameters in the update rule are derived from the input $u_t \in \mathbb{R}^d$ through learned projections:
Discretization step (gating):
where $W_\Delta \in \mathbb{R}^{d \times 1}$ and $b_\Delta \in \mathbb{R}$ are trainable parameters, and Softplus is a smooth approximation to ReLU that ensures $\Delta_t > 0$. This scalar acts as a gating mechanism: it controls how much the model "pays attention" to the current input, since both the decay and insertion strength depend on it.
Memory retention strength (decay):
where $A \in \mathbb{R}$ is a trainable parameter (one per head). Since $\Delta_t > 0$ and $\exp(A) > 0$, we have $\alpha_t \in (0, 1)$, with $\alpha_t \to 1$ meaning "retain everything" and $\alpha_t \to 0$ meaning "completely forget the past." The double exponential form comes from the continuous-time state space model derivation: $A$ is the continuous-time decay rate, and $\Delta_t$ is the discretization step size.
Insertion vector:
where $\tilde{B}_t = \sigma(\text{Conv}(W_B u_t)) \in \mathbb{R}^{N \times 1}$ is computed via a short 1D convolution (kernel size 4) followed by SiLU activation, and $W_B \in \mathbb{R}^{d \times N}$ is a trainable projection. The multiplication by $\Delta_t$ couples insertion strength to the gating: when the model decides to "pay attention" (large $\Delta_t$), it both decays past memory more (smaller $\alpha_t$) and inserts new information more strongly (larger $B_t$). This coupling is a design choice that links reading and writing—it means the model cannot simultaneously strongly retain the past AND strongly insert new information, creating an explicit tradeoff.
Query matrix:
where $W_C \in \mathbb{R}^{d \times N}$ is a trainable projection, and the same convolution + SiLU pattern is used as for $B_t$. The query vector $C_t$ is what reads from the state to produce the output.
Content embedding:
where $W_x \in \mathbb{R}^{d \times P}$. This is the actual information being stored in the state via the outer product with $B_t$.
Full output (per head):
where $D \in \mathbb{R}^{1 \times P}$ is a trainable skip-connection parameter. This allows the model to bypass the state entirely for some dimensions, which is useful for representing information that should not be subject to the recurrent dynamics. The head output is then multiplied element-wise by a gated version of the input and projected back to dimension $d$:
where $W_o \in \mathbb{R}^{P \times d}$, $W_{\text{gate}} \in \mathbb{R}^{d \times P}$, and Norm is RMS normalization. The gating $W_{\text{gate}} u_t$ acts as a nonlinearity that controls which dimensions of the recurrent output $o_t$ are used.
Key structural facts (Appendix A.1): Mamba-2 always sets $P = 64$ and $N = 128$. The total number of heads is $H = 2d / P$. This means the total state size across all heads and layers is $H P N \cdot L = 2d N L$. Since $N = 128$, the total state size is $256 d L$. In comparison, a standard transformer of hidden dimension $d$ with $L$ layers has a KV cache size of $2 d L T$. Therefore, the Mamba-2 state is equivalent to a transformer's KV cache at sequence length $T = 128$. The number of layers $L$ in Mamba-2 is roughly twice the number of attention layers in a comparably-sized transformer, so the per-layer state is even smaller relative to transformers. This structural fact—that the state size is fixed and relatively compact compared to transformers' growing KV caches—is what creates both the efficiency advantage and the overparameterization problem.
State Decomposition and Retrieval Error Analysis
The linearity of the update rule (Equation 2) allows the state $h_t$ to be unrolled as an explicit weighted sum over all past insertions:
where $\alpha_{i:t}$ is the cumulative decay applied to the insertion at timestep $i$ by timestep $t$, defined as:
where $\alpha_j$ is the per-step retention strength at step $j$, and the product runs from the step after insertion $i$ through the current step $t$.
What it computes: The state at time $t$ is a superposition (weighted sum) of all $t$ insertions $B_i x_i$ made at previous time steps. The weight on the $i$-th insertion is the product of all per-step retention factors from step $i$ to $t$. Because each $\alpha_j \in (0,1)$, earlier insertions have smaller weights (they have been multiplied by more $\alpha_j < 1$ terms), creating a temporal decay. This form reveals that the state's information content at any time is a mixture of all past inputs, with recent inputs weighted more heavily.
Why this form: This representation is what makes the analysis of forgetting mathematically tractable. Instead of treating the state as a black box, we can explicitly see each past token's contribution and its decay. The cumulative product $\alpha_{i:t}$ quantifies exactly how much of the $i$-th token "remains" in the state at time $t$. If the model produces $\alpha_j \approx 1$ for all $j$, then $\alpha_{i:t} \approx 1$ even for very early $i$, meaning nothing is forgotten. If the model produces small $\alpha_j$, then $\alpha_{i:t}$ decays rapidly, and early tokens contribute negligibly.
Now consider what happens when the model tries to retrieve the information inserted at a specific past timestep $s$. To query for that information, the model produces $C_t = B_s$ (the query vector matches the insertion vector for the target token). The retrieved output is:
where the first term $\alpha_{s:t} (C_t B_s) x_s$ is the desired signal (the $s$-th token's content $x_s$, scaled by its retention $\alpha_{s:t}$ and the query-insertion match $C_t B_s$), and the second term $\sum_{i \neq s} \alpha_{i:t} C_t B_i x_i$ is the retrieval error—interference from all other tokens.
What it computes: The retrieval operation isolates one token's contribution from the summed state by using a query vector $C_t$ that correlates with the target's insertion vector $B_s$ and ideally is orthogonal to all other $B_i$ for $i \neq s$. The dot product $C_t B_s$ amplifies the target, while $C_t B_i$ for $i \neq s$ should ideally be zero (perfect orthogonality). In practice, with finite state dimension $N$ and $t$ potentially much larger than $N$, the $B_i$ vectors cannot all be mutually orthogonal, so the error term is non-zero. The severity of interference depends on two factors: (1) how much the earlier tokens have been decayed (controlled by $\alpha_{i:t}$), and (2) how well-separated the $B_i$ vectors are in the $N$-dimensional space.
Why this form: This decomposition makes explicit the tradeoff that the forgetting mechanism must navigate. If $\alpha_{i:t}$ is large for many $i$ (the model retains everything), then the error sum contains many non-negligible terms, and when the total number of terms $t$ exceeds the training length, the accumulated error can overwhelm the signal term, causing retrieval failure for tokens at any position (not just early ones). The model has a learned mechanism to control this—produce small $\alpha_t$ values to rapidly decay early tokens—but the paper's central claim is that it does not learn to use this mechanism when trained on short sequences, because within the training length, the error from $T_{\text{train}}$ tokens is still tolerable, so there is no gradient pressure to develop aggressive forgetting.
Diagnostic Method 1: Measuring Cumulative Retention Strength
To determine whether the model has learned to forget, the authors compute the cumulative retention strength of the first token—$\alpha_{1:t}$—as a function of position $t$ for each head in each layer (Section 3.2.1, Figure 3).
Procedure: For a given input sequence, the model computes $\alpha_t = \exp(-\Delta_t \exp(A))$ at each position. The cumulative product from position 1 to the current position $t$ is then:
This value represents the fraction of the first token's insertion that remains in the state at position $t$. It is computed separately for each head (since each head has its own $A$ parameter and $\Delta_t$ values).
What it measures: If $\alpha_{1:t}$ decays toward zero as $t$ increases, the model is forgetting the first token—its contribution is being progressively diminished by small multiplicative decay factors. If $\alpha_{1:t}$ stays close to 1, the first token is being retained almost perfectly, meaning the model is NOT forgetting.
Key finding (Figure 3): In Mamba-2 370M (trained on 8K sequences), three out of the first eight heads in layer 38 have $\alpha_{1:8000} > 0.997$. This means that after 8,000 time steps, the first token retains over 99.7% of its initial contribution to the state. Since this is a cumulative product, the per-step decay must be even closer to 1—the model is essentially not decaying at all within the training window. Similar patterns are observed in other heads and layers.
Why this is diagnostic: If the model produces $\alpha_t \approx 1$ throughout the training window, it has learned a strategy of retaining everything rather than selectively forgetting. This strategy works fine within 8K tokens because the state size (12.9M parameters for the 370M model) is large enough to store information about 8K tokens without catastrophic interference. But the strategy fails catastrophically when the context extends to 16K or 24K tokens: the model has never experienced a situation where interference becomes problematic, so it has never learned to produce smaller $\alpha_t$ values that would prevent the error accumulation in Equation 7 from overwhelming the signal.
Diagnostic Method 2: State Distribution Statistics
The second diagnostic approach examines how the recurrent state's statistical distribution changes as a function of context length, particularly when the length exceeds the training window (Section 3.3).
Procedure: The authors feed a controlled prompt—the "newlines" prompt, consisting solely of repeated newline characters "\n\n\n..."—to Mamba-2 and record the hidden state $h_t$ (the recurrent state from Equation 2, not the full output) for every head in every layer at every position. For each head, they compute the mean and variance of the state values across its $N \times P$ elements at each position, then plot these statistics as a function of $t$. They also examine the per-channel distribution at specific positions.
Why the newlines prompt: The newlines prompt is chosen because (1) it can be arbitrarily long, (2) it contains no varying linguistic content, so any changes in state statistics reflect the recurrent dynamics rather than the input, (3) it produces the "most consistent and smooth layer statistics" (Appendix I), making it easier to detect systematic patterns. The paper notes that similar state distribution changes are observed on real pre-training data and passkey retrieval prompts, but those have variable length and content that obscure the systematic trends.
Key finding (Figure 5, Figure 16 in Appendix G): For many heads in the Mamba-2 370M model, the mean and variance of the state exhibit a sharp change—often an explosion—when the context length $t$ exceeds the training length of 8K tokens. In the specific example of layer 38, heads 0–7 (Figure 5), some heads show the mean jumping from near-zero to noticeably non-zero values at positions beyond 8K, while the variance increases dramatically (from roughly 10–50 to hundreds or thousands). Figure 6 shows the per-channel distribution for one head at $t = 8K$ (within training) versus $t = 20K$ (beyond training): the majority of channels have small, stable values, but a few "outlier channels" have values that have grown to the range of ±100, driving the variance explosion.
Full-layer statistics (Figure 16, Appendix G): The pattern is not uniform across all layers. Early layers (0–7) show relatively stable statistics out to 20K+ positions. Middle layers (16–31) show the most dramatic explosions. Late layers (40–47) show more modest changes. This suggests that the pathological state accumulation—the failure to forget—propagates through the network in a layer-dependent manner.
Contrast with convolutional states (Figure 17, Appendix G): The short convolution layers (kernel size 4) that produce $B_t$, $C_t$, and $x_t$ have their own state (the previous 4 inputs to the convolution), and the statistics of this convolutional state remain stable far beyond the training length. This is expected because the convolutional state only stores 4 tokens of history, so it cannot accumulate error. The contrast highlights that the explosion is specific to the recurrent state with its unbounded temporal accumulation.
What this measures: The explosion in state statistics beyond the training length is the manifestation of the model's inability to forget. When the context exceeds the training length, the model has never seen states with this many accumulated insertions. Because it has not learned to produce small $\alpha_t$ values (it retains everything during training), the state continues to accumulate information linearly, and the values grow without bound. The "outlier channels" that explode are the dimensions where the accumulated interference from past tokens is most concentrated—they are the channels that happen to have strong projections onto many $B_i$ vectors, causing their values to blow up when too many tokens are stored.
Why this form: The mean-and-variance diagnostic is simple to compute and visualize across all layers and heads, making it a practical tool for identifying which components of the model are suffering from over-retention. The per-channel analysis (Figure 6) reveals that the problem is not a uniform degradation—most channels remain stable—but rather a concentration of the accumulated error into a few dimensions, which then corrupt the output when those dimensions are read by the query mechanism.
Intervention Experiment 1: Reduced Memory Retention and Insertion (RRI)
To establish a causal link between over-retention and length generalization failure, the authors introduce an inference-time intervention that artificially induces more forgetting: Reduced Memory Retention and Insertion (RRI) (Section 3.2.2).
Procedure: At inference time, after the model computes $\alpha_t$ and $B_t$ normally, they are scaled by constant multipliers smaller than 1:
where the specific multipliers used are $c_\alpha = 0.9999$ (for the retention/decay factor) and $c_B = 0.75$ (for the insertion vector). These values are chosen by validation using the average loss on pre-training data with 32K context length.
What it computes: This intervention forces each head to decay past information slightly more aggressively at every step (multiplying $\alpha_t$ by a value less than 1 reduces the cumulative product $\alpha_{i:t}$ faster) and to insert new information less strongly (scaling $B_t$ by 0.75 reduces the magnitude of new insertions). The net effect is that earlier tokens contribute less to the state at any given time, reducing the interference term in Equation 7.
Why these multipliers: The value $c_\alpha = 0.9999$ is extremely close to 1—this is not an aggressive intervention. The authors are making a minimal change that nudges the cumulative decay to be slightly faster. The fact that such a tiny multiplier produces a visible improvement in length generalization (Figure 4) demonstrates how sensitive the system is to the exact decay rate and how the model's natural $\alpha_t$ values are essentially at the edge of stability. The insertion scaling $c_B = 0.75$ is more substantial because the insertion vectors $B_t$ must be large enough to be distinguishable in the state space; reducing them helps prevent new insertions from dominating the state.
Key result (Figure 4): The RRI-modified model (green curve) shows lower language modeling loss at positions beyond 8K compared to the original model (blue curve). However, it also shows slightly higher loss within the first 8K positions (visible from the initial offset in Figure 4), because the reduced insertion strength means less information is stored per token, which degrades short-context performance. This is the expected tradeoff: more aggressive forgetting helps long contexts but hurts short contexts, where retaining all information is actually optimal. The original model over-optimizes for short-context retention at the expense of long-context stability.
Why this is causal evidence: The model's parameters are unchanged—the only difference is the scaling multipliers applied to the computed $\alpha_t$ and $B_t$. By showing that simply reducing memory retention and insertion (without any training) improves length generalization, the experiment demonstrates that the failure is NOT due to the model's inability to represent long contexts, but rather due to its failure to learn the appropriate decay rates during training. The model has the architectural capacity to handle long contexts; it just has not learned the parameter settings that would enable this.
Intervention Experiment 2: Sliding Window State
The second causal intervention is more principled: instead of heuristically scaling the decay, the authors implement an exact sliding window mechanism using the algebraic properties of the weighted-sum state representation (Section 3.2.2).
Procedure: Let $w$ be the window size (exact value not specified in the main text, but described as a hyperparameter). Instead of using the full state $h_t$ (which contains information from all tokens 1 through $t$), the model should use a windowed state $h_t^{(w)}$ that contains only the last $w$ tokens:
Using the unrolled representation (Equation 6), this can be computed exactly as the difference between two full states:
where $h_t$ is the full state at time $t$, $h_{t-w}$ is the full state at time $t-w$, and $\alpha_{t-w+1:t} = \prod_{j=t-w+1}^{t} \alpha_j$ is the cumulative decay applied to the old state over the window period.
Maintaining the sliding window efficiently: To compute $h_t^{(w)}$ at each step during streaming generation, the system must maintain three quantities:
$h_{t-1}$: the full state at the previous step (needed for normal state update)$h_{t-w}$: the full state from$w$steps ago (needed for the subtraction)$\alpha_{t-w+1:t}$: the cumulative decay over the current window
The third quantity is computed indirectly to avoid floating-point instability from multiplying many $\alpha_j$ values. Instead, the system maintains $\Delta_{t-w:t} = \sum_{i=t-w}^{t} \Delta_i$ (the sum of discretization steps), and recomputes $\alpha_{t-w:t} = \exp(-\Delta_{t-w:t} \exp(A))$ at each step, which is numerically stable and computationally negligible.
What it computes: The sliding window state $h_t^{(w)}$ is mathematically identical to what the model would produce if it had only ever seen the last $w$ tokens, with all earlier tokens completely removed. The subtraction $h_t - \alpha_{t-w+1:t} h_{t-w}$ works because the contribution of tokens 1 through $t-w$ to $h_t$ is exactly $\alpha_{t-w+1:t}$ times their contribution to $h_{t-w}$ (they have been further decayed by the window's worth of $\alpha$ factors), so subtracting the correctly scaled old state eliminates them entirely.
Why this works algebraically: The state update is linear and the decay is multiplicative—there is no nonlinear state transformation that would make the subtraction invalid. This is a special property of the Mamba-2 update rule (and other gated linear RNNs) that enables exact windowing without recomputation. The key insight is that the state's weighted-sum representation (Equation 6) is explicit: we can always identify which terms came from which time interval, because the decay factors $\alpha_{i:t}$ are known and purely multiplicative.
Computational cost: Maintaining the sliding window requires storing an additional state vector $h_{t-w}$ and the scalar $\Delta_{t-w:t}$. The extra memory is $O(N \cdot P)$ per head—the same size as the state itself, roughly doubling the state memory. The extra computation per step is one scalar-vector multiplication (for $\alpha_{t-w+1:t} \cdot h_{t-w}$) and one vector subtraction, which is negligible compared to the cost of computing the full state update.
Key result (Figure 4): The sliding window intervention (red curve in Figure 4) provides the best length generalization among all methods compared—better than the original model, better than LongMamba, and better than RRI. The loss remains stable far beyond the training length. This is because the sliding window guarantees that the state never contains information from more than $w$ tokens ago, completely eliminating the interference from very early tokens that causes the state explosion.
Why this is causal evidence: Unlike RRI, which makes a heuristic modification to the decay rate, the sliding window intervention makes a structural change: it enforces an absolute bound on how far back the state can remember. The fact that this bound prevents the loss explosion confirms that the explosion is caused by tokens beyond the window contributing to the state—i.e., by the failure to forget. Furthermore, the sliding window does not require modifying the model's learned parameters; it only changes how the state is used at inference time, proving that the model's short-context representations are still valid when isolated from long-range interference.
Controlled Training: Establishing the Forgetting Threshold
The core empirical contribution of the paper is establishing that there exists a training length threshold $T_{\text{forget}}$ above which Mamba-2 learns robust forgetting, and that this threshold scales linearly with the state size (Section 4, Section 5).
Core hypothesis (Section 4): The state is "overparameterized" relative to typical training lengths. The state size $N_S$ (total number of parameters across all recurrent states) determines the information capacity—roughly speaking, how many tokens' worth of distinct information can be stored before interference becomes severe. If the training length $T_{\text{train}}$ is less than this capacity, the model can simply retain everything (by keeping $\alpha_t \approx 1$) and still perform well on the language modeling objective, because the accumulated interference from $T_{\text{train}}$ tokens is within the state's tolerance. The model receives no gradient signal encouraging it to forget, because forgetting would only reduce the amount of stored information without improving predictions within the training window. When $T_{\text{train}}$ exceeds the capacity, however, retaining everything leads to noticeable interference even within the training window, creating gradient pressure to produce smaller $\alpha_t$ values and learn selective forgetting.
The hypothesis predicts: (1) for any state size, there exists a threshold $T_{\text{forget}}(N_S)$ such that the model learns to forget if and only if $T_{\text{train}} > T_{\text{forget}}$, and (2) $T_{\text{forget}}$ scales monotonically with $N_S$ (larger states can store more before needing to forget).
Evidence for "more training leads to less forgetting" (Section 4.1, Figure 8): The authors pre-train Mamba-2 370M from scratch with $T_{\text{train}} = 512$ tokens using the RedPajama corpus (a short training length, far below the state capacity). They evaluate intermediate checkpoints at 10B, 20B, and 40B tokens on passkey retrieval. The results show that:
- At 10B tokens, the model achieves reasonable retrieval accuracy (~80–90%) for answer positions within the last few hundred tokens, even when the total context is 8K tokens.
- At 20B tokens, accuracy within the recent window improves, but accuracy for the 8K context length drops.
- At 40B tokens, the model has near-perfect accuracy for contexts ≤512 tokens, but essentially zero accuracy for 8K contexts.
The model's in-distribution performance improves (better retrieval within the training length), but its out-of-distribution performance degrades (worse retrieval at longer lengths). The interpretation: the model is learning to retain more and forget less as training progresses, because within the 512-token window, retention is strictly beneficial. This is analogous to overfitting—the model's state dynamics are optimized for the short-context distribution and fail to generalize.
Controlled sweeping of training lengths by model size (Section 5.1–5.2): To establish the quantitative relationship, the authors train multiple Mamba-2 variants with different state sizes and sweep training lengths for each. The models are:
| Model Size | State Size $N_S$ | # Layers | Hidden $d$ | # Heads |
|---|---|---|---|---|
| 36.4M | 0.8M | 6 | 512 | 16 |
| 47.0M | 1.6M | 12 | 512 | 16 |
| 84.6M | 2.4M | 12 | 768 | 24 |
| 130M (official) | 4.8M | 24 | 768 | 24 |
| 370M (official) | 12.9M | 48 | 1024 | 32 |
| 780M (official) | 19.3M | 48 | 1536 | 48 |
The three smallest models are trained from scratch. The three larger ones are continued from official Mamba-2 checkpoints (all originally trained on 8K sequences). For each model, the authors sweep training lengths $T_{\text{train}}$ up to 256K tokens, continuing pre-training on the RedPajama-V2 dataset.
Data processing for long-context training (Section 5, Appendix F): The RedPajama-V2 corpus is deduplicated and filtered to remove documents shorter than 4K tokens (removing ~97.6% of the data, since most web documents are short). This filtering ensures that the training data contains genuine long-range dependencies rather than concatenations of unrelated short documents. For training lengths longer than the longest available document, sequences are concatenated and delimited with a special EOS token. During evaluation, documents longer than 16K tokens are sampled and concatenated if needed.
Truncated backpropagation through time (TBPTT): States are initialized as the final state of the previous sequence rather than zeros, effectively concatenating sequences but stopping gradient backpropagation at intervals. This makes the distribution of initial states $h_0$ more varied (they reflect actual intermediate states rather than always being zero), which Yang et al. (2024a) showed helps extend the effective context length. The default is to concatenate 12 sequences with this technique.
Training hyperparameters (Appendix F.1):
- Learning rate scheduler: WSD (Warmup-Stable-Decay) with 10% decay steps, chosen because it allows simple resumption from intermediate checkpoints (unlike cosine schedules, which require restarting from scratch). The warmup phase is 1000 steps of linear increase, the stable phase maintains the peak learning rate, and the decay phase is 50K steps of linear decrease.
- Learning rate sweep:
$\{1 \times 10^{-5}, 2 \times 10^{-5}, 5 \times 10^{-5}, 1 \times 10^{-4}, 2 \times 10^{-4}, 5 \times 10^{-4}, 1 \times 10^{-3}\}$, with the best value selected by validation on passkey retrieval (not by language modeling loss, since "the loss of many checkpoints was highly similar [but] their performance in passkey retrieval can differ a lot"). - Batch size: 0.5M tokens (half the original Mamba-2 batch size of 1M, chosen empirically for more stable continual pre-training).
- Precision: BF16 for most computations, with FP32 for some activations (matching the official implementation).
- Optimizer: AdamW with weight decay 0.1 and gradient clipping at 1.0.
Determining whether forgetting has been learned (Section 4.2 criterion): For each trained model, the authors feed prompts with 1M tokens and check whether the language modeling loss at any position exceeds $2 \times$ the maximum loss observed within $T_{\text{train}}$ tokens. The loss is averaged over 128 prompts. If the loss stays bounded (never exceeds this 2× threshold), the model is judged to have learned robust forgetting—it generalizes to arbitrarily long contexts without catastrophic degradation. If the loss explodes, the model has not learned to forget and still suffers from interference accumulation.
This criterion is operational rather than theoretical: it directly tests whether the model's state dynamics remain stable when the context length is orders of magnitude beyond the training length. The 2× multiplier is a conservative threshold that allows for some increase in loss (longer contexts are inherently more difficult) while flagging the catastrophic mode where loss doubles or more.
Key result (Figure 10, Figure 11): For each model size, there is indeed a training length threshold above which the loss remains stable at 1M tokens and below which it explodes. For Mamba-2 130M ($N_S = 4.8\text{M}$), the threshold is approximately 16K–32K tokens. For Mamba-2 370M ($N_S = 12.9\text{M}$), it is approximately 64K–128K tokens. The 780M model ($N_S = 19.3\text{M}$) still exhibits poor length generalization at training lengths up to 128K (the maximum the authors could afford to train), confirming that the threshold continues to grow with state size.
Fitting the linear relationship (Figure 11): Plotting $T_{\text{forget}}$ (the minimum training length that produces stable loss at 1M tokens) against state size $N_S$ (in millions of parameters) yields a strong linear fit:
with $R^2 > 0.999$. The state size $N_S$ is measured in millions of parameters (so for the 370M model, $N_S = 12.9$). $T_{\text{forget}}$ is in thousands of tokens.
What it computes: Given a Mamba-2 model's state size $N_S$ (which is determined by the architectural parameters: $N_S = H \cdot P \cdot N \cdot L = 2d \cdot N \cdot L$, where $N = 128$ for Mamba-2, $d$ is the hidden dimension, and $L$ is the number of layers), this equation predicts the minimum training context length required for the model to learn robust forgetting. For $N_S = 12.9$ (370M model), the prediction is $5.172 \times 12.9 - 4.469 \approx 62.2$K tokens, which matches the empirical observation that the threshold lies between 32K and 64K. For $N_S = 19.3$ (780M model), the prediction is $5.172 \times 19.3 - 4.469 \approx 95.5$K, consistent with the observation that 128K training is still insufficient for this model.
Why this form is linear: The linear relationship suggests that the information capacity of the state scales linearly with the number of parameters in the state. Each state parameter can store approximately $1/5.172 \approx 0.193$ tokens' worth of information (the inverse slope). Alternatively, approximately 5.17 state parameters are needed per token of training context to reach the point where the state is "full" enough that forgetting becomes necessary. This linear scaling has a clear intuitive interpretation: the state is a matrix of size $N \times P$ per head, and the total number of elements across all heads and layers is $H \cdot N \cdot P \cdot L$. Each element stores a scalar value that is a weighted sum of past insertions; the total information content is proportional to the number of such scalars, hence linear.
Controlled Training: Establishing the Maximum Recall Context Length
The forgetting threshold $T_{\text{forget}}$ represents the point where the training context contains enough information to force the model to learn decay. However, once forgetting is learned, the model is not limited to recalling only the last $T_{\text{forget}}$ tokens—it should be able to selectively retain and recall specific pieces of information from much longer contexts, as long as the total amount of distinct information in the context is bounded.
The passkey retrieval task as a probe (Section 5.3, Appendix B): Passkey retrieval is a synthetic task where a model must find and recall a 5-digit number (the "passkey") embedded in a long context of repetitive filler text ("The grass is green. The sky is blue..."). The prompt structure (Appendix B.1) is:
- A brief instruction: "There is important info hidden inside a lot of irrelevant text. Find it and memorize it."
- Filler text
- The passkey insertion: "The passkey is 34847. Remember it. 34847 is the passkey."
- More filler text
- The query: "What is the passkey? The passkey is"
The passkey can be placed at any position in the context (parameterized by "Answer Depth" = passkey position / context length, as a percentage). The filler text contains very little distinct information—it is highly repetitive—so the amount of information the model needs to store is essentially just the 5-digit passkey itself. The challenge is preserving this information across potentially very long filler sequences while ignoring the filler.
Why this task isolates recall from capacity: In language modeling, a 256K-token document may contain thousands of distinct facts, entities, and narrative elements—the total information content grows roughly linearly with length. In passkey retrieval, the total information content is constant (5 digits + the positional context needed to know "where" the passkey is). This means the state's information capacity is not the bottleneck; the bottleneck is the model's ability to selectively retain the passkey while forgetting the filler. If the model has learned robust forgetting (the $\alpha_t$ mechanism appropriately decays irrelevant tokens), it should be able to handle arbitrarily long filler sequences—the state capacity required is constant, and the discrimination ability (how many distinct patterns the state can distinguish) grows exponentially with state size.
Evaluation protocol (Section 5, Appendix F): For each trained model, the authors sweep context lengths $T \in \{1\text{K}, 2\text{K}, 4\text{K}, \dots, 256\text{K}\}$ and passkey positions evenly distributed from 0% to 100% of the context length. Passkey placement uses $n$ evenly spaced positions, with the $i$-th passkey inserted at position $\lfloor T \times i / n \rfloor$. The exact prompt template and the specific filler text are provided in Appendix B.1. Decoding is greedy (temperature 0) for reproducibility, and both model parameters and activations use FP32 to eliminate precision errors as a confounding factor (BF16 introduces non-negligible errors in $\Delta_t$ and $\alpha_t$ computation).
Determining recall success: A passkey is considered correctly retrieved if the model outputs the exact 5-digit sequence. For a given context length $T$, the model is evaluated at all position depths, and the accuracy is reported as a heatmap (Figure 2, Figure 8, Figure 9). The "maximum recall context length" $T_{\text{recall}}$ for a model is defined as the largest $T$ for which the model achieves over 95% accuracy across all answer depths.
Key result for models trained above $T_{\text{forget}}$ (Figure 9): The maximum recall context length $T_{\text{recall}}$ scales exponentially with state size:
with $R^2 > 0.999$. $T_{\text{recall}}$ is in thousands of tokens, and $N_S$ is in millions of parameters.
What it computes: For a given state size $N_S$, this equation predicts the longest context from which the model can successfully retrieve a 5-digit passkey. For $N_S = 12.9$ (370M model), the prediction is:
This matches the empirical result: the Mamba-2 370M model, after continued pre-training at lengths above $T_{\text{forget}}$, achieves near-perfect passkey retrieval at 256K context length—a result the authors claim outperforms similarly sized transformer models.
Why the exponential form: The exponential scaling arises from the combinatorics of the state space. The passkey retrieval task requires the model to discriminate between the passkey-bearing position and all other positions. The state is a vector in a high-dimensional space (dimension $N \cdot P \cdot H \cdot L$ total). The number of nearly-orthogonal directions in this space—and hence the number of distinct positions that can be uniquely addressed—grows exponentially with the dimensionality. Each additional unit of state capacity multiplies the number of distinguishable patterns by a constant factor (here, 1.365). This is analogous to the capacity of associative memory models (Hopfield networks, modern Hopfield models), where the number of storable patterns grows exponentially with dimension.
Why this is critically different from the forgetting threshold: $T_{\text{forget}}$ is about training dynamics—when does the model receive sufficient gradient signal to learn the forgetting mechanism? It scales linearly with state size because it depends on how much information can be crammed into the state before interference becomes noticeable. $T_{\text{recall}}$ is about inference-time capability—once forgetting is learned, how far can selective recall extend? It scales exponentially because it depends on the combinatorial capacity of the high-dimensional state space to distinguish patterns. The two thresholds answer different questions and obey different scaling laws, but they are causally linked: you cannot achieve any $T_{\text{recall}}$ benefits until you first train with $T_{\text{train}} > T_{\text{forget}}$.
Training cost caveat: The exponential recall capacity is only realized when the model is trained on contexts long enough to learn forgetting. The current practice of training Mamba-2 on 8K sequences means that for the 370M model ($T_{\text{forget}} \approx 62$K), the model never learns to forget, and $T_{\text{recall}}$ collapses to essentially zero for lengths beyond 8K. The paper's scaling laws provide a concrete recipe: to unlock the exponential recall capacity, one must first pay the linear training length cost.
Summary of Design Choices and Their Justifications
- Using existing pre-trained checkpoints for the largest models (130M, 370M, 780M) rather than training from scratch: saves enormous computational cost while still allowing exploration of the forgetting threshold through continued pre-training at longer lengths. The fact that the linear
$T_{\text{forget}}$relationship holds across both from-scratch and continued models validates this choice. - Using passkey retrieval as the primary downstream task rather than more complex benchmarks: provides a clean, controllable measure of pure recall ability without confounding factors like reasoning or knowledge. The simplicity of the task means failure can be confidently attributed to the forgetting mechanism rather than to other capabilities.
- The "newlines" prompt for state statistics collection rather than real text: eliminates input-driven variance in the state trajectories, making the explosion pattern clearly visible. The paper verifies that similar patterns occur on real data, but uses newlines for clean visualization.
- The 2× loss threshold for determining
$T_{\text{forget}}$rather than a more continuous measure: provides a clear binary criterion for whether forgetting has been learned. The threshold is conservative (allowing some loss increase) while clearly separating models that catastrophically degrade from those that remain stable. - The WSD learning rate scheduler rather than cosine: enables resumption from intermediate checkpoints, which is critical for the extensive sweep of training lengths—if each training length required a full separate training run from scratch, the computational cost would be prohibitive.
- Truncated backpropagation through time with state carryover rather than independent sequences: improves the diversity of initial state distributions, making the training regime closer to the long-context inference regime where states are always "warm" (carrying information from prior context).
- The 4K-token minimum document length filter for training data: ensures that training sequences contain genuine long-range structure, preventing the model from learning that it can simply reset/ignore context at document boundaries.
4. Key Insights and Innovations
Innovation 1: Reframing the Problem from Capacity Insufficiency to Overparameterization
The paper's most fundamental conceptual move is inverting the dominant framing in prior work on RNN length generalization. Before this paper, the standard explanation for why recurrent models fail on long contexts was insufficient capacity: the fixed-size state cannot hold enough information, so performance degrades as the context grows beyond what the state can store. This framing appears in Jelassi et al. (2024), who show that Mamba struggles to copy unless its state size grows linearly with context length, and in Arora et al. (2024a), who analyze the associative recall capacity limits of various architectures. Under this view, the solution is to make states larger—scale the state dimension with the desired context length.
This paper inverts that narrative entirely. The problem is not that the state is too small; it is that the state is too large relative to the training context length, creating a situation where the model can achieve low training loss by simply retaining everything (keeping α_t ≈ 1) rather than learning a selective forgetting policy. The evidence for this reframing is direct and multi-pronged:
- Figure 3 shows that within the 8K training window, some heads have cumulative first-token retention α_{1:8000} > 0.997—meaning the model retains essentially everything it has ever seen, with no evidence of learned forgetting.
- Figure 1 demonstrates that larger models (with larger states) have worse length generalization—exactly the opposite of what a capacity-limitation hypothesis would predict. If the state were too small, larger states should generalize better; the fact that they generalize worse is strong evidence that excess capacity is the culprit.
- Figure 4 shows that artificially inducing more forgetting (via RRI scaling or sliding window) at inference time improves length generalization without any additional training. If insufficient capacity were the problem, reducing the model's ability to retain information should hurt performance, not help.
This reframing is more than a semantic distinction—it fundamentally changes the prescription for fixing the problem. Under the capacity-limitation view, the solution is architectural (bigger states, more expressive update rules). Under the overparameterization view, the solution is about training methodology: train on longer sequences, specifically sequences whose length exceeds the state's capacity to store everything without interference. The paper's linear scaling law (T_forget = 5.172 · N_S − 4.469, Figure 11) operationalizes this prescription: it tells you exactly how long your training sequences must be for a given state size.
Significance: This is a fundamental reframing with direct practical consequences. It explains why larger Mamba-2 models have paradoxically worse length generalization (a phenomenon visible in Figure 1 that had no prior explanation), it reconciles the apparent contradiction between Mamba's impressive short-context performance and catastrophic long-context failure, and it provides a clear, actionable guideline for practitioners that does not require architectural changes. The reframing also connects to a broader theme in deep learning—overparameterization leading to undesirable implicit biases—that has been extensively studied in the context of generalization in feedforward networks but has not previously been applied to the temporal dynamics of recurrent architectures.
Relationship to prior work: The paper explicitly contrasts its diagnosis with the "state capacity" framing of Jelassi et al. (2024) and Arora et al. (2024a), and with the "over-smoothing" analysis of Wang et al. (2025). It also distinguishes its findings from the engineering heuristics of LongMamba (Zhang, 2023) and DeciMamba (Ben-Kish et al., 2024), which attempted to improve length generalization by making decay factors closer to 1 (retaining more)—exactly the opposite of what the overparameterization diagnosis would prescribe. The paper's sliding window and RRI interventions (Figure 4) demonstrate that inducing more forgetting is the correct direction, directly contradicting these prior approaches.
Innovation 2: The Memory Decay as a Learned Behavior That Can Be "Overfit" Away
The paper's second distinctive contribution is demonstrating that the memory decay mechanism in Mamba-2—the parameter α_t that controls how much past information is retained—is not a fixed architectural property but a learned behavior that changes over the course of training and can be pathologically optimized for the training distribution at the expense of generalization. This insight emerges from the training dynamics shown in Figure 8: as the model sees more data (10B → 20B → 40B tokens) when trained on short (512-token) sequences, its in-distribution retrieval accuracy improves while its out-of-distribution retrieval accuracy degrades. The model converges toward a strategy of retaining more information and forgetting less, because within the 512-token window, every retained token is potentially useful for prediction, and there is never enough interference to create negative gradient signal.
This behavior—"more training leads to less forgetting"—is counterintuitive from the perspective of standard deep learning, where more training data generally improves generalization. It is, however, precisely what one would expect from an overfitting dynamic: the model's state dynamics become increasingly specialized to the short-context distribution, and the implicit bias that favors simple solutions (in this case, "retain everything") becomes dominant as training proceeds. The model never encounters the failure mode (state explosion from accumulated interference) during training, so it never receives gradient signals that would push it toward learning a forgetting policy.
This finding has important implications for how we think about temporal dynamics in learned recurrent systems. In classical RNNs (LSTMs, GRUs), the gating mechanisms that control memory retention and forgetting are often thought of as solutions to the vanishing gradient problem—they are architectural features that enable gradient flow across long time spans. Under this view, the forget gate "works" because of how it is structured (sigmoid gating, additive updates), not because of what it learns from data. The Mamba-2 results challenge this assumption: having a forgetting mechanism in the architecture is necessary but not sufficient. The mechanism must also receive the right training signal to learn when to forget, and that signal only exists when the training context is long enough to create interference that forgetting can resolve.
Significance: This is a conceptual advance that bridges the gap between architectural design and training methodology. It implies that evaluating recurrent architectures solely on their in-distribution performance (e.g., perplexity on validation sets with the same context length as training) may systematically miss the most important property of the learned dynamics—whether the model has learned to forget appropriately for long-context generalization. The practical implication is that training procedures for recurrent models should include explicit out-of-distribution evaluation (e.g., loss curves at extended context lengths, as in Figure 10) as part of the model selection criterion, rather than relying on in-distribution validation loss, which the paper notes can be "highly similar" across checkpoints with very different length generalization behavior (Appendix F.1).
Relationship to prior work: This finding reframes the results of prior work that studied Mamba's length generalization as a static property of trained checkpoints. Jelassi et al. (2024) evaluated final models; Waleffe et al. (2024) benchmarked existing checkpoints; Ben-Kish et al. (2024) and Zhang (2023) proposed inference-time fixes for trained models. None of these works examined how the forgetting behavior evolves during training, which is where the paper's key insight lies. The training progression in Figure 8 reveals that the length generalization failure is not an architectural limitation but a training artifact—a finding that opens the door to fixing the problem through better training recipes rather than architectural changes.
Innovation 3: Empirical Scaling Laws for Forgetting and Recall in Recurrent Models
The paper's third major contribution is establishing quantitative scaling laws that relate the recurrent state size to two critical quantities: the minimum training length needed to learn robust forgetting (T_forget), and the maximum context length from which the model can accurately recall specific information after forgetting is learned (T_recall). These two relationships are shown to follow fundamentally different functional forms—linear for T_forget, exponential for T_recall—reflecting different underlying mechanisms.
The linear relationship T_forget = 5.172 · N_S − 4.469 (Figure 11, R² > 0.999) is derived from controlled experiments sweeping state sizes from 0.8M to 12.9M parameters and training lengths from 4K to 256K tokens. The experimental design is noteworthy for its systematic nature: six model sizes, multiple training lengths per size, a clear operational criterion for determining whether forgetting has been learned (loss at 1M tokens stays below 2× the maximum in-training loss), and careful control of confounding factors (data filtering for long documents, TBPTT for state diversity, WSD scheduler for checkpoint resumption). The near-perfect linear fit is striking and provides strong evidence that the underlying relationship is genuine rather than an artifact of the specific experimental choices.
The exponential relationship T_recall = 4.756 · (1.365^{N_S} − 1) − 0.742 (Figure 9, R² > 0.999) is derived from passkey retrieval evaluation on the same set of models after they have been trained above their T_forget thresholds. The exponential form has a deep theoretical interpretation: the passkey retrieval task requires discriminating between the passkey-bearing position and all other (repetitive, information-poor) positions. The state is a vector in a space of dimension proportional to N_S, and the number of distinguishable patterns in a high-dimensional space grows exponentially with dimension—this is the same capacity scaling that underlies associative memory models. The fact that the empirical fit is exponential (base 1.365 per million state parameters) means that modest increases in state size can yield dramatic increases in maximum recall length: doubling the state size from 5M to 10M increases the predicted T_recall from approximately 19K to approximately 105K tokens.
Significance: These scaling laws provide the first principled guidance for how to configure training length and state size in recurrent language models. Before this work, the choice of training context length for Mamba models was largely ad-hoc—the standard was 8K tokens, inherited from the original Mamba-2 paper, with no theoretical or empirical justification for why this length was chosen or whether it was sufficient. The linear T_forget law provides a clear answer: for the 370M model with state size 12.9M, 8K training is far below the required ~62K threshold, explaining why the model catastrophically fails beyond 8K. For the 130M model (state size 4.8M), the threshold is ~20K, which is closer to 8K but still insufficient, consistent with the observation that the 130M model shows somewhat better retrieval at 16K (Figure 2a) than larger models.
The two scaling laws together reveal an important asymmetry in the cost-benefit tradeoff of scaling state size. To unlock the exponential recall benefits (T_recall), you must first pay the linear training cost (T_forget). This means that scaling state size is a double-edged sword: larger states can theoretically achieve much longer recall, but they also require much longer training sequences to learn forgetting in the first place. If training length is held fixed (as in current practice), scaling state size makes things worse, not better (Figure 1). If training length can be scaled proportionally, the exponential payoff in recall capability eventually dominates.
Relationship to prior work: These scaling laws are conceptually analogous to the compute-optimal scaling laws for transformers (Hoffmann et al., 2022) in that they provide quantitative relationships between model hyperparameters (state size, training length) and downstream capabilities (length generalization, recall distance). However, they address a fundamentally different phenomenon: transformer scaling laws relate training compute to loss, while these laws relate inference-time behavior (the ability to process long contexts) to training-time choices (context length). This is a novel category of scaling relationship that has no direct precedent in the literature. The only prior work connecting training length to recurrent model behavior is Buckman & Gelada (2024), who argue for choosing training length based on data statistics, but without connecting to state size or forgetting dynamics.
Innovation 4: The Retrieval Error Analysis as a Unifying Diagnostic Framework
The paper's fourth distinctive contribution is the decomposition of the recurrent state as a weighted sum (Equation 6) and the associated retrieval error analysis (Equation 7), which together provide a mechanistically grounded diagnostic framework for understanding why forgetting matters and how its failure manifests. While the weighted-sum representation itself is not novel—it follows directly from the linearity of the update rule—the paper's use of this representation to link three previously separate observations into a single causal chain is a significant conceptual contribution.
The three observations that this framework unifies are:
-
High retention strength (α_{1:t} ≈ 1, Figure 3): The model is not producing small enough per-step decay factors to meaningfully reduce the contribution of early tokens to the state.
-
State distribution explosion (Figures 5, 6, 16): When the context exceeds the training length, specific channels in the recurrent state experience sharp increases in variance, with a small number of "outlier channels" growing to dominate the state norm.
-
Catastrophic retrieval failure at all positions (Figure 2): When the context length significantly exceeds the training length, the model cannot retrieve tokens from any position—including very recent ones—rather than experiencing a graceful degradation that affects only early positions.
The retrieval error analysis connects these observations causally. When α_{i:t} is close to 1 for all i (observation 1), the error term ∑{i≠s} α{i:t} C_t B_i x_i in Equation 7 accumulates linearly with the number of tokens. Since the B_i vectors cannot be mutually orthogonal in N-dimensional space when the number of tokens exceeds N, the error term grows without bound as the context lengthens. This error concentrates in specific state dimensions—the channels with the strongest projections onto many B_i vectors—causing the outlier channel explosions (observation 2). When these outlier channels dominate the state, the query operation C_t h_t returns values corrupted by the accumulated interference, destroying retrieval accuracy for tokens at any position—even very recent ones that should be dominant by recency (observation 3).
What makes this framework distinctive: Prior work had observed some of these phenomena in isolation. Wang et al. (2025) discussed over-smoothing and noted that the decay term can cause token representations to blend. Jelassi et al. (2024) observed that Mamba struggles with copying beyond certain lengths. But these observations were treated as separate problems with separate explanations (over-smoothing vs. capacity limits), and neither explained the non-graceful nature of the failure—specifically, why retrieval fails for very recent tokens when the total context is long, rather than only for early tokens. The retrieval error analysis provides a single mechanism that explains all three observations simultaneously and predicts the non-graceful failure mode: the error accumulation is a property of the total accumulated state, not of individual token decay, so when it becomes severe, it corrupts retrieval for all tokens, not just early ones.
Significance beyond this paper: This diagnostic framework is not specific to Mamba-2. Any recurrent architecture whose state can be expressed as a weighted sum over past insertions—which includes GLA (Yang et al., 2024a), RWKV (Peng et al., 2024a), RetNet (Sun et al., 2023), HGRN-2 (Qin et al., 2024), and many others—is subject to the same interference dynamics. The framework thus provides a general tool for diagnosing and predicting length generalization failures in the broader class of linear recurrent models. The paper's finding that RWKV-6 shows less severe failure (Appendix D, Figures 13, 14) is consistent with this framework: RWKV has a smaller state size (Figure 12), so its state reaches the interference threshold at a longer context length for a given training length, or equivalently, its T_forget is lower for the same state size.
The framework also suggests diagnostic procedures that practitioners can apply to their own models: measure α_{1:t} to check whether the model is decaying at all; inspect state statistics at extended lengths to identify which layers are accumulating interference; compute retrieval accuracy as a function of total context length with the target token at a fixed recent position to specifically test the "non-graceful failure" prediction. These diagnostics require no additional training, only forward passes on long sequences.
Relationship to prior work: No prior work on RNN length generalization has provided a mechanistic explanation at this level of specificity. The closest is Arora et al. (2024a), who analyze the associative recall capabilities of various architectures in terms of state capacity, but their analysis focuses on what the state can store in principle (capacity limits), not on how the state dynamics evolve when capacity is exceeded. The retrieval error analysis complements this by addressing the dynamics—what happens to the state values and retrieval accuracy when the number of stored tokens exceeds what the state can represent without interference.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The RedPajama-V2 corpus (Computer, 2023), an open dataset with 30T tokens from the Internet, is used for both pre-training from scratch and continued pre-training of Mamba-2 models. Documents shorter than 4K tokens are filtered out, removing approximately 97.6% of the data to ensure training sequences contain genuine long-range structure (Appendix F). For evaluation of language modeling loss as a function of position, documents longer than 16K tokens are sampled and concatenated if needed. The passkey retrieval task uses a synthetic prompt with repetitive filler text ("The grass is green. The sky is blue...") and a 5-digit passkey inserted at controlled positions; the exact template is provided in Appendix B.1. For state statistics collection, a "newlines" prompt consisting solely of repeated newline characters is used because it produces the most consistent and smooth layer statistics while eliminating input-driven variance (Section 3.3, Appendix I).
-
Base model(s). The primary subjects are Mamba-2 checkpoints at four scales: 130M, 370M, 780M, and 1.3B parameters, all officially released and pre-trained on 8K-token sequences (Table 1). These are chosen because Mamba-2 has "shown strong capabilities on several tasks and has publicly available checkpoints of multiple sizes, allowing us to explore the relationship between state sizes and length limits" (Section 2). Additional smaller models (36.4M, 47.0M, 84.6M parameters) are trained from scratch with configurations reported in Table 2 to extend the state size sweep downward. For comparative evaluation, Mamba-1 (130M, 370M, 790M, 1.4B, 2.8B), RWKV-5 (1.5B, 3B), RWKV-6 (1.6B, 3B), and HGRN-2 (1.3B) are also evaluated on passkey retrieval or length generalization (Appendix D, H). The relationship between state size and model size for these architectures is shown in Figure 12.
-
Metrics. Three primary metrics are used. Language modeling loss as a function of token position (Figures 1, 4, 10): cross-entropy loss computed at each position in a long sequence, revealing whether the model degrades beyond its training length. This is evaluated on RedPajama documents or the newlines prompt. Passkey retrieval accuracy (Figures 2, 8, 9, 13, 14, 15): the fraction of trials where the model outputs the exact 5-digit passkey, evaluated as a function of context length and answer position (parameterized as "Answer Depth" = passkey position / context length, as a percentage). Accuracy is reported as heatmaps. Greedy decoding with FP32 precision is used for reproducibility and to eliminate precision errors as a confounding factor (Appendix B). State statistics (Figures 5, 6, 16, 17): the mean and variance of the recurrent hidden state h_t and convolutional state values are computed across all elements for each head at each position, plotted as a function of token position to detect distributional shifts when context exceeds training length.
-
Baselines. The paper does not compare against alternative training recipes as baselines in the traditional sense—there is no "competing method" for learning to forget. Instead, the primary comparisons are: (1) The original Mamba-2 model evaluated at lengths within versus beyond its 8K training length (Figures 1, 2), serving as the baseline that exhibits catastrophic degradation. (2) LongMamba (Zhang, 2023), a heuristic that divides the discretization term ∆_t by a constant (evaluated with multiplier 0.5 and also swept across values, Appendix C), which makes α_t closer to 1 and thus retains more information—evaluated in Figure 4 as a comparison point for the intervention experiments. (3) Larger Mamba-2 models (370M, 780M, 1.3B) compared against the 130M model within Figure 1, establishing the counterintuitive result that larger models have worse length generalization. (4) For the passkey retrieval evaluation, Mamba-1, RWKV-5, and RWKV-6 official checkpoints serve as architectural comparisons (Figures 13, 14, 15). (5) For the FLOPs-matched comparison implicit in the recall scaling analysis, the paper notes that the 370M Mamba-2 after continued pre-training at lengths above T_forget "achieves near-perfect retrieval at 256K context length, outperforming similarly sized transformer models" (Section 1, Section 5.3), though no specific transformer baseline is quantitatively reported in the paper.
-
Generation budget / compute accounting. The paper does not use a FLOPs-based compute budget since the primary analysis is about length generalization behavior rather than compute-optimal allocation. Training cost is measured in tokens processed (10B, 20B, 40B tokens for the training dynamics experiment; unspecified total token counts for the threshold-finding experiments). Inference cost for passkey retrieval is measured in context length (1K to 256K tokens) with greedy decoding. For the sliding window intervention, the additional memory required is one extra state vector h_{t-w} plus a scalar ∆_{t-w:t} per head, roughly doubling the state memory with negligible additional computation (Section 3.2.2). The paper explicitly notes that training cost constrained the maximum sweeps: "we do not have enough resources to train the [780M] model beyond [128K training length]" (Section 5.2).
-
Cross-validation / statistical protocol. For the controlled training experiments, the learning rate is selected by validation on passkey retrieval performance, not by language modeling loss, because "the loss of many checkpoints was highly similar [but] their performance in passkey retrieval can differ a lot" (Appendix F.1). The best checkpoint is selected using this validation criterion. For the forgetting threshold determination (Section 4.2), the criterion is whether loss on 1M-token prompts exceeds 2× the maximum loss within T_train tokens, averaged over 128 prompts—this provides a statistical basis for the binary classification of whether forgetting has been learned. For passkey retrieval evaluation, each context length T is evaluated with n evenly spaced needle positions (n not explicitly specified for all experiments, but described as "evenly distributed" in Appendix B.1), and accuracy is aggregated across positions. No formal confidence intervals or standard errors are reported for any of the scaling law fits, though R² values > 0.999 are provided (Figures 9, 11).
Main Quantitative Results
The Forgetting Threshold Scales Linearly with State Size
The paper's central quantitative finding is the linear relationship between minimum training length required for robust forgetting (T_forget) and total recurrent state size (N_S). The fitted relationship (Figure 11) is:
with R² > 0.999, where N_S is measured in millions of parameters and T_forget in thousands of tokens. The data points span state sizes from 0.8M (36.4M model) to 12.9M (370M model), with the 780M model (N_S = 19.3M) confirmed to still exhibit poor length generalization at training lengths up to 128K, serving as an additional point above the fitted line.
For each model size, the determination of whether forgetting has been learned follows an operational criterion defined in Section 4.2: the model is evaluated on 1M-token prompts, and forgetting is judged to be learned if the loss at any position never exceeds 2× the maximum loss observed within T_train tokens (averaged over 128 prompts). This criterion is designed to be conservative—allowing moderate loss increases that naturally accompany longer contexts—while clearly separating models that catastrophically degrade from those that remain stable.
The specific thresholds for each model (read from Figure 10, with quantitative values implied by the fit):
-
36.4M model (N_S = 0.8M): The fitted line predicts T_forget ≈ 5.172 × 0.8 − 4.469 ≈ −0.33K, meaning this model should learn forgetting even at very short training lengths. This is consistent with the observation that smaller models have better length generalization in Figure 1.
-
47.0M model (N_S = 1.6M): Predicted T_forget ≈ 5.172 × 1.6 − 4.469 ≈ 3.81K. Training at 4K should be near the threshold.
-
84.6M model (N_S = 2.4M): Predicted T_forget ≈ 5.172 × 2.4 − 4.469 ≈ 7.94K. This is approximately the standard 8K training length, suggesting this model is at the boundary.
-
130M model (N_S = 4.8M): Predicted T_forget ≈ 5.172 × 4.8 − 4.469 ≈ 20.4K. Figure 10(a) shows that at T_train = 4K and 8K, loss explodes beyond the training length; at T_train = 16K, the model shows "much better length extrapolation"—placing the threshold between 8K and 16K, roughly consistent with the prediction.
-
370M model (N_S = 12.9M): Predicted T_forget ≈ 5.172 × 12.9 − 4.469 ≈ 62.2K. Figure 10(b) shows that at T_train = 8K, 16K, and 32K, loss explodes; at T_train = 64K, the loss remains stable well beyond the training length. The threshold is between 32K and 64K, consistent with the prediction.
-
780M model (N_S = 19.3M): Predicted T_forget ≈ 5.172 × 19.3 − 4.469 ≈ 95.5K. The paper reports that this model "also has poor length generalization at training lengths below 128K" (Section 5.2), with the authors unable to train beyond 128K due to resource constraints. This is consistent with the prediction that the threshold is above 95K and likely near 100K, which would require >128K training to clearly exceed.
The fact that continued pre-training from official 8K checkpoints (for the 130M, 370M, and 780M models) yields thresholds consistent with the from-scratch models (36.4M, 47.0M, 84.6M) validates that the linear relationship holds regardless of the pre-training history—the threshold depends only on state size and final training length.
Computation of state size: For Mamba-2 with P = 64, N = 128, and H = 2d/P heads per layer, the state size per layer is H × N × P = (2d/64) × 128 × 64 = 256d parameters. Across L layers, N_S = 256dL. For the 370M model with d = 1024 and L = 48, N_S = 256 × 1024 × 48 ≈ 12.6M, closely matching the 12.9M figure reported in Table 2 (the small discrepancy may arise from rounding or inclusion of the convolutional state, which the paper notes contains "much less contextual information" and is "largely ignored" in state size calculations; Appendix A.2).
The Maximum Recall Context Length Scales Exponentially with State Size
When Mamba-2 models are trained with T_train > T_forget (i.e., they have learned robust forgetting), their maximum passkey retrieval context length T_recall scales exponentially with state size (Figure 9):
with R² > 0.999, where N_S is in millions of parameters and T_recall in thousands of tokens.
The procedure for determining T_recall: For each model trained above its T_forget threshold, passkey retrieval accuracy is evaluated at context lengths sweeping from 1K to 256K, with passkey positions evenly distributed from 0% to 100% of the context. The maximum context length with accuracy exceeding 95% across all answer depths is recorded as T_recall.
Specific data points (read from Figure 9 and the fitted curve):
-
For N_S = 0.8M (36.4M model): Predicted T_recall ≈ 4.756 × (1.365^0.8 − 1) − 0.742 ≈ 4.756 × (1.282 − 1) − 0.742 ≈ 0.60K, suggesting very limited recall even after learning to forget—consistent with a small state having limited discrimination capacity.
-
For N_S = 1.6M: Predicted T_recall ≈ 4.756 × (1.365^1.6 − 1) − 0.742 ≈ 4.756 × (1.645 − 1) − 0.742 ≈ 2.33K.
-
For N_S = 2.4M: Predicted T_recall ≈ 4.756 × (1.365^2.4 − 1) − 0.742 ≈ 4.756 × (2.110 − 1) − 0.742 ≈ 4.54K.
-
For N_S = 4.8M (130M model): Predicted T_recall ≈ 4.756 × (1.365^4.8 − 1) − 0.742 ≈ 4.756 × (4.47 − 1) − 0.742 ≈ 15.8K.
-
For N_S = 12.9M (370M model): Predicted T_recall ≈ 4.756 × (1.365^12.9 − 1) − 0.742 ≈ 4.756 × (54.8 − 1) − 0.742 ≈ 255K. This is the headline result: the 370M Mamba-2, after continued pre-training above its T_forget threshold (~62K), achieves near-perfect passkey retrieval at 256K context length. The paper emphasizes that "to the best of our knowledge, no previous models with less than 1B model parameters have near-perfect accuracy at this length in this task" (Section 5.3).
-
For N_S = 19.3M (780M model): Predicted T_recall ≈ 4.756 × (1.365^19.3 − 1) − 0.742, which extrapolates to extremely large values (1.365^19.3 ≈ 410), predicting T_recall well into the millions of tokens. However, this extrapolation is untested since the 780M model could not be trained above its T_forget threshold with available resources (Section 5.2).
The exponential scaling is interpreted as arising from the combinatorial capacity of the high-dimensional state space to discriminate between distinct token positions: "the number of combinations of the state grows exponentially with the number of elements" (Section 5.3). Since passkey retrieval involves a constant amount of information (5 digits) embedded in repetitive filler text, the total information load is independent of context length; the bottleneck is the state's ability to maintain a distinct "address" for the passkey position among many possible positions, which scales exponentially with dimensionality.
Training Dynamics: More Data Leads to Less Forgetting
When Mamba-2 370M is trained from scratch on short (512-token) sequences, the passkey retrieval behavior evolves in a counterintuitive direction as training progresses (Figure 8, Section 4.1).
At 10B tokens (Figure 8a): The model achieves moderate retrieval accuracy (~80–90%) for passkeys placed within the last few hundred tokens, even when the total context is much longer (up to 8K tokens). There is some ability to recall recent information from longer contexts.
At 20B tokens (Figure 8b): Accuracy for very recent positions (within the last ~1K tokens) improves, but accuracy for the extended 8K context length drops. The model is becoming more specialized to the short-context distribution.
At 40B tokens (Figure 8c): The model has near-perfect accuracy for contexts ≤512 tokens (the training length) but essentially zero accuracy at 8K contexts regardless of answer position. The in-distribution retrieval performance has improved to near-ceiling, while out-of-distribution performance has collapsed to floor.
This progression—in-distribution improvement coinciding with out-of-distribution degradation—is the core evidence for the claim that the model is "overfitting" its forgetting behavior to the training distribution (Section 4.1). Within 512 tokens, retaining all information (keeping α_t ≈ 1) is strictly beneficial because the accumulated interference from 512 tokens is within the state's tolerance. As training proceeds, gradient descent amplifies this strategy, producing decay factors closer and closer to 1. The model never experiences the failure mode (state explosion from interference) during training, so there is no gradient signal to push toward learning selective forgetting. The result is a model that is locally optimal for short contexts but brittle to any length extension.
This finding is operationalized as a diagnostic: "language modeling loss is only computed for tokens within the training length, [so] this behavior is induced by minimizing loss" (Section 4.1). The implication is that training loss alone is insufficient to detect whether a recurrent model has learned appropriate forgetting dynamics; out-of-distribution evaluation (loss at extended lengths, passkey retrieval) is necessary.
Artificially Inducing Forgetting Improves Length Generalization
Figure 4 shows language modeling loss as a function of token position for Mamba-2 370M under four conditions: the original model, LongMamba (divide ∆_t by 0.5), RRI (scale α_t by 0.9999 and B_t by 0.75), and sliding window (exact subtraction of tokens beyond a window of unspecified size).
Original model (blue curve in Figure 4): Loss is lowest within the 8K training length but begins to climb sharply afterward, reaching approximately 2× the in-training loss by position 24K and continuing to degrade to position 32K. This is the baseline catastrophic failure.
LongMamba (purple curve): The loss degradation beyond 8K is less severe than the original model—the curve rises more gradually and plateaus at a lower level—but in-training loss is slightly higher because reducing ∆_t amplifies the memory decay (makes α_t closer to 1) on all tokens indiscriminately, which "unnecessarily diminishes the inserted information on all tokens" (Section 6). This illustrates the tradeoff: more uniform retention helps long contexts but hurts short-context performance.
RRI (green curve): Shows a similar pattern to LongMamba but with a different tradeoff profile. The reduced insertion strength (B_t scaled by 0.75) raises in-training loss slightly (less information stored per token), but the reduced retention (α_t scaled by 0.9999) prevents the catastrophic loss increase beyond 8K. The fact that c_α = 0.9999—an extremely subtle change—produces a visible improvement demonstrates how close the model's natural decay factors are to 1 and how sensitive the system is to the exact decay rate.
Sliding window (red curve): Provides the best length generalization of all methods. The loss remains stable well beyond the training length, with no visible explosion. This is because the sliding window enforces an absolute bound: tokens beyond the window size contribute literally nothing to the state, eliminating the interference accumulation entirely. The cost is additional memory (one extra state vector and scalar per head) but no architectural change or retraining.
The success of these interventions—particularly the sliding window, which is a structural guarantee of forgetting rather than a heuristic adjustment—confirms causally that the failure mode is over-retention leading to interference, not a fundamental representational limitation.
State Distribution Explosion Beyond Training Length
The manifestation of the forgetting failure in the state values is documented through mean and variance statistics computed on the "newlines" prompt for Mamba-2 370M (Section 3.3, Appendix G).
Figure 5 shows the mean and variance of the first 8 heads in layer 38 as a function of token position out to ~35K tokens. Several heads exhibit a clear explosion when t exceeds the 8K training length: the variance jumps from values of roughly 10–100 within the training window to values of hundreds or thousands at extended positions. The mean also shifts from near-zero to noticeably non-zero values. The explosion is not instantaneous at position 8001—it develops gradually as the context lengthens, consistent with the accumulation of interference from an increasing number of insufficiently decayed tokens.
Figure 6 examines the per-channel distribution of one head (head 2, layer 38) at two specific positions: t = 8K (within training) and t = 20K (far beyond training). At 8K, the distribution of channel values is concentrated around zero, with counts peaking at ~10³–10⁴. At 20K, the distribution has changed dramatically: a small number of channels have values in the range ±50–100, while the majority remain near zero. The variance explosion is "largely attributed to a few outlier channels while most channels are relatively stable" (Section 3.3). This is exactly what the retrieval error analysis (Equation 7) predicts: the accumulated interference concentrates in specific dimensions of the state space—those with the strongest projections onto many B_i vectors—rather than being uniformly distributed.
Figure 16 (Appendix G) provides the full set of mean-and-variance plots for all 48 layers (grouped in panels of 8 layers each). The patterns vary by layer depth: early layers (0–7) show relatively stable statistics even at extended positions, suggesting they process more local information that doesn't accumulate across long ranges. Middle layers (16–31) show the most dramatic explosions, consistent with these layers being responsible for integrating information across longer temporal spans. Late layers (40–47) show more modest changes, possibly because their role is to read out from the state rather than to accumulate it further.
Figure 17 (Appendix G) shows the same statistics for the convolutional states (the short 1D convolutions with kernel size 4 that produce B_t, C_t, and x_t). These remain stable throughout—means stay within ±0.04 and variances within 0–2, even at 25K+ positions. The contrast with the recurrent state explosion is stark and confirms that the instability is specific to the unbounded temporal accumulation in the recurrent state, not to any numerical precision issue or general architectural property.
Comparative Architecture Evaluation
The paper evaluates several other recurrent architectures on passkey retrieval to determine whether the forgetting failure is specific to Mamba-2 or more general (Appendix D).
Mamba-1 (Figure 15): The 790M, 1.4B, and 2.8B checkpoints all show clear passkey retrieval failure patterns similar to Mamba-2. At context lengths near the training length (likely 2K or 4K based on the checkpoint release notes), retrieval accuracy is high for recent positions. As context length increases to 16K–32K, accuracy drops to near-zero across all answer depths. The failure mode is the same non-graceful pattern: when the context exceeds a threshold, retrieval fails for all positions, not just early ones. This is consistent with Mamba-1 having the same overparameterization dynamics, though its state size is 8× smaller than Mamba-2 for comparable model sizes (Mamba-1 uses N = 16 vs. Mamba-2's N = 128), which would predict a proportionally lower T_forget and better length generalization at a given training length—a prediction partially borne out by the observation that Mamba-1 failure occurs at longer absolute lengths than Mamba-2, though the training lengths also differ.
RWKV-5 (Figure 13): The 1.5B and 3B checkpoints show less severe failure than Mamba variants. At 1K–8K context lengths, retrieval accuracy for passkeys in the last 1K–8K tokens (depending on the curve) remains reasonably high. The 1.5B model in particular maintains accuracy above 50% for passkeys in the last 8K tokens even at 16K total context. The 3B model shows some degradation at 16K but still above zero. The paper hypothesizes that "this difference is a result of architectural differences and [the fact] that the state size is smaller in RWKV-5 and RWKV-6" (Appendix D), as shown in Figure 12, where RWKV models have significantly smaller state sizes than comparably-sized Mamba-2 models.
RWKV-6 (Figure 14): The 1.6B checkpoint shows passkey retrieval accuracy patterns that are qualitatively better than Mamba-2 but still exhibit degradation at 16K context length. Accuracy for very recent passkeys (within the last 1K–2K tokens) remains high, but drops off for earlier positions. The failure is more graceful than Mamba-2's catastrophic collapse—accuracy degrades with answer depth rather than dropping to zero uniformly.
HGRN-2 (Figure in Appendix H): The 1.3B model shows perplexity increase on the newlines prompt starting considerably before the training length is reached—"perhaps surprisingly, the increase in perplexity happens considerably before the context length reaches the training length" (Appendix H). This is noted as potentially distinct from the Mamba pattern, hypothesized to result from training distribution differences.
The comparative results broadly support the state size hypothesis: models with smaller states (RWKV series) exhibit less catastrophic failure at a given context length, consistent with their T_forget being lower and thus more likely to be exceeded by their training length (assuming similar training lengths across architectures).
Ablation Studies and Robustness Checks
Training length sweep for fixed model sizes (determining T_forget existence): Figure 10 shows the language modeling loss curves for Mamba-2 130M and 370M when continued pre-training at different training lengths (4K, 8K, 16K, 32K, 64K). For the 130M model (Figure 10a), training at 4K produces loss that begins increasing immediately after 4K; training at 8K delays the increase to ~8K; training at 16K produces stable loss well beyond 16K, confirming the threshold exists between 8K and 16K. For the 370M model (Figure 10b), training at 8K, 16K, and 32K all produce eventual loss explosions; training at 64K produces stable loss far beyond 64K, confirming the threshold between 32K and 64K. This ablation establishes that T_forget is a real threshold—a qualitative change in behavior—rather than a continuous improvement with longer training.
State size sweep for fixed training length (the overparameterization evidence): Figure 1 implicitly shows this ablation: at a fixed 8K training length, the 130M model degrades noticeably but less severely than the 370M model, which degrades less severely than the 780M model, which degrades less severely than the 1.3B model. Larger models (with larger states) have worse length generalization when training length is held constant, directly supporting the overparameterization hypothesis. If the problem were simply insufficient training, all models would degrade similarly; the size-dependence confirms that larger states make the problem worse because they increase the capacity that must be exceeded for forgetting to be learned.
Training data quantity (forgetting as a function of optimization progress): The three checkpoints in Figure 8 (10B, 20B, 40B tokens) constitute an ablation over optimization steps at fixed training length (512 tokens). The progressive degradation of out-of-distribution retrieval accuracy with more training demonstrates that the forgetting failure is not simply a result of undertraining—more training makes the problem worse, not better, because gradient descent increasingly optimizes for the short-context distribution. This is a critical robustness check: it rules out the hypothesis that the models simply need more training data at their current sequence length.
Decay rate modification magnitude (sensitivity of length generalization): The RRI intervention uses c_α = 0.9999, a multiplicative factor extremely close to 1. The fact that such a tiny adjustment produces a visible improvement in Figure 4 demonstrates that the model's learned decay rates are poised at a critical point where even minuscule changes to the cumulative decay substantially affect long-context stability. While not presented as a formal sweep over c_α values, the comparison between RRI (c_α = 0.9999), LongMamba (divide ∆_t by 0.5, equivalent to making α_t closer to 1, the opposite direction), and the sliding window (complete elimination of old tokens) spans the spectrum from "retain more" to "retain slightly less" to "retain only the last w," establishing that the direction of improvement is toward more forgetting, with the optimal amount being task-dependent (LongMamba hurts short-context performance, RRI hurts short-context slightly, sliding window doesn't affect short-context at all since it only eliminates tokens beyond the window).
Convolutional vs. recurrent state stability: Figures 16 and 17 (Appendix G) together constitute an ablation over which component of the state accumulates the interference. The recurrent state (Figure 16) exhibits explosions in many layers; the convolutional state (Figure 17) remains stable in all layers. Since the convolutional state only aggregates information over the last 4 tokens (kernel size 4), while the recurrent state aggregates over the entire sequence, this contrast isolates the temporal extent of accumulation as the necessary condition for the explosion. It also rules out numerical precision issues (which would affect both state types similarly) and input-driven artifacts (since both states process the same input).
Architecture variants (Mamba-1 vs. Mamba-2 vs. RWKV): The comparison across architectures (Figures 13, 14, 15) serves as an ablation over specific design choices—Mamba-1's smaller N (16 vs. 128), Mamba-2's larger state, RWKV's different gating mechanisms. The pattern that architectures with smaller states (RWKV) show less catastrophic failure, while those with larger states (Mamba-2) show worse failure, supports the state-size-centric explanation and suggests that the phenomenon is not tied to specific parameterizations of the update rule. However, the paper notes that RWKV-7 and Gated DeltaNet "have gone beyond a gating-based memory decay mechanism and are out of the scope of this paper" (Section 6), so the findings may not transfer to those architectures.
Precision effects: Appendix B notes that preliminary evaluations were conducted with BF16, FP16, and FP32 precision. BF16 introduces errors around 1e-3 in computing ∆_t and α_t, but "the explosion of channels in the states is consistently observed despite this precision error." FP16 shows no noticeable differences from FP32. The robustness of the qualitative finding (state explosion) across precisions rules out floating-point artifacts as the explanation, though FP32 is used for all reported results to eliminate precision as a confounding variable.
Greedy vs. non-greedy decoding for passkey retrieval: Appendix B notes that "our preliminary results show that other decoding parameters give noticeably worse performance on passkey retrieval." This is not explored in depth, but it suggests that the retrieval failure is sensitive to decoding strategy, which is consistent with the interference mechanism: if the retrieved signal is corrupted by accumulated error, stochastic decoding (which adds noise) would further degrade accuracy.
Critical Assessment
Claim 1: The inability to forget causes performance degradation. The evidence for this claim comes in three forms: (a) high retention strength values (α_{1:8000} > 0.997 in some heads, Figure 3), (b) state distribution explosions when context exceeds training length (Figures 5, 6, 16), and (c) improved length generalization when forgetting is artificially induced (RRI, sliding window; Figure 4). This evidence is mutually reinforcing and, taken together, makes a compelling causal case. The sliding window intervention is particularly strong because it is a structural guarantee (tokens beyond the window are mathematically eliminated from the state), not a heuristic tuning, and it produces the best length generalization. The high retention values are correlational—they show the model is retaining everything, but don't by themselves prove this retention causes the failure. The intervention experiments close this gap by demonstrating that reducing retention (causally) improves generalization.
A limitation worth noting: the α_{1:t} measurement in Figure 3 is shown for only 8 heads in one layer (layer 38) of one model (370M). The text states "similar observations can be found in other heads and in other layers as well," but the quantitative evidence is narrow. A systematic survey of retention strengths across all layers and model sizes would strengthen the claim, particularly to establish whether the heads with high retention are specifically those whose state statistics explode (Figure 5), which would directly link the two diagnostic methods.
Claim 2: State overparameterization explains the inability to forget. This is the paper's most important conceptual claim, and the evidence is substantial but circumstantial. The two main supports are: (a) training on short sequences (512 tokens) progressively eliminates forgetting as optimization proceeds (Figure 8)—the model converges toward a "retain everything" strategy because it works within the training window; (b) larger models have worse length generalization at fixed training length (Figure 1), and the training length needed for forgetting scales linearly with state size (Figure 11).
The "overparameterization" framing—that the state is too large relative to training length—is a compelling interpretation, but the paper does not provide direct evidence that state capacity is the mechanism. An alternative explanation could be that longer training sequences simply provide more diverse state initial conditions (via TBPTT), and it is this diversity, not capacity pressure, that enables the model to learn forgetting. The paper uses TBPTT with state carryover specifically to increase initial state diversity (Appendix F), but the role of this technique versus raw sequence length is not ablated. Would a model trained on 64K-token sequences but with states always reset to zero (no TBPTT) still learn forgetting? If yes, the capacity-pressure explanation is supported. If no, the diversity-of-initial-conditions explanation would gain traction. This ablation is not performed.
Additionally, the claim that "state capacity is exceeded at T_forget" is inferred from the fact that forgetting is learned when training length exceeds that threshold, but there is no direct measurement of information-theoretic state capacity (e.g., by measuring reconstruction accuracy of stored tokens as a function of state size and number of stored tokens). The linear fit itself is strong (R² > 0.999), but with only six data points (five used for fitting, one out-of-range for validation), the functional form—while clearly linear in the tested range—might not hold at much larger state sizes. The 780M model's failure at 128K training is consistent with linear extrapolation but does not confirm it; training this model at 256K or 512K to observe whether forgetting emerges at the predicted ~100K threshold would substantially strengthen the claim.
Claim 3: The minimum training length scales linearly with state size, and the maximum recall length scales exponentially. The linear fit (Figure 11) is based on five data points with state sizes from 0.8M to 12.9M, plus one validation point (19.3M) that is consistent with the extrapolation but not confirmed to lie on the line since the actual threshold for the 780M model is unknown (it is only bounded below by 128K). Five points are sufficient to establish linearity in the tested range, but the extrapolation to larger models (e.g., a hypothetical 7B Mamba-2) rests on the assumption that the relationship remains linear, which is theoretically plausible (capacity ∝ number of parameters) but empirically untested.
The exponential fit (Figure 9) is based on the same models, evaluated after being trained above their respective T_forget thresholds. The R² > 0.999 is impressive, but the data range is small: T_recall spans from ~1K to ~256K over state sizes from 0.8M to 12.9M. This is roughly 8 doublings of T_recall and 4 doublings of state size. An exponential fit to four doublings is plausible but not strongly constrained—other functional forms (e.g., a high-degree polynomial) might fit equally well. The theoretical justification (combinatorial capacity of high-dimensional spaces) supports an exponential relationship, but alternative mechanisms (e.g., the state learning to use a small number of dimensions for positional addressing, with capacity growing as a power law) cannot be ruled out from the data alone.
A more significant concern is circularity in the data pipeline: the models used for the T_recall measurement are the same models (or continuations thereof) used to establish T_forget. Since T_recall is only measured for models trained above T_forget, and T_forget is determined from the same training runs, there is no independent test of whether the linear and exponential relationships generalize to models trained with entirely different recipes (different data, different hyperparameters, different architectures within the Mamba-2 family). The paper's use of continued pre-training from official checkpoints (for the 130M, 370M, and 780M models) and from-scratch training (for the smaller models) provides some robustness, since these have different initialization and optimization trajectories, but all models share the same architecture and training data.
Claim 4: The 370M Mamba-2 with continued pre-training achieves near-perfect passkey retrieval at 256K, outperforming similarly sized transformers. The 256K retrieval result for the 370M model is the paper's most striking claim but also the least substantiated in terms of comparison. The paper states this outperforms "similarly sized transformer models" (Section 5.3) and "no previous models with less than 1B model parameters have near-perfect accuracy at this length" (Section 5.3), but provides no head-to-head comparison with a specific transformer baseline at comparable scale. The passkey retrieval literature includes results for transformers at various scales, but the paper does not cite or reproduce specific numbers for a 370M-parameter transformer on this exact task with this exact prompt format. This makes the claim difficult to evaluate independently.
More fundamentally, the claim conflates two distinct achievements: (a) that the 370M Mamba-2 can retrieve from 256K after continued pre-training above T_forget, and (b) that this is better than what transformers can do at this scale. Achievement (a) is directly demonstrated by the experiments. Achievement (b) is an informal comparison without rigorous benchmarking. The paper would be stronger if it either provided the transformer comparison data or limited the claim to the demonstrated fact: that Mamba-2's recall capacity, when properly trained, extends to at least 256K, far beyond the 8K limit of the standard checkpoints.
Missing experiments that would strengthen the paper:
-
Direct measurement of information capacity: Train models to store and recall specific numbers of random tokens, measuring retrieval accuracy as a function of tokens stored and state size. This would provide an independent estimate of state capacity in bits and directly test whether T_forget corresponds to the point where the information load exceeds this capacity.
-
Ablation of TBPTT state carryover: Train models at fixed long sequence lengths but with and without state carryover (i.e., states reset to zero vs. continued from previous sequence). This would distinguish between the "capacity pressure" explanation (longer sequences force forgetting because they contain more information) and the "state diversity" explanation (state carryover provides more varied initial conditions for training).
-
Transformer baseline at matched scale: Train a ~370M transformer on the same data with the same sequence lengths and evaluate on the same passkey retrieval task with the same prompt format, allowing a direct comparison.
-
Sweep of a single model to very long training lengths: Train, for example, the 370M model at 128K, 256K, and 512K sequence lengths to determine whether T_recall continues to scale or saturates. Currently, the exponential fit extrapolates far beyond the measured range.
-
Measurement of α_t statistics across training: Track how the distribution of α_t values evolves during training at different training lengths, providing direct evidence for the claim that short-sequence training drives α_t toward 1 while long-sequence training drives it to smaller values.
-
Evaluation on non-synthetic long-context tasks: Passkey retrieval is intentionally simple. Evaluating models trained above T_forget on realistic long-context benchmarks (e.g., ∞Bench, Zhang et al., 2024a) would demonstrate whether the improved forgetting translates to practical capabilities.
-
RWKV with extended training: The paper observes that RWKV shows less catastrophic failure and hypothesizes this is due to smaller state size. Training a RWKV model at progressively longer lengths to find its T_forget and comparing against the Mamba-2 linear fit would test the generality of the scaling law.
Conditions under which the claims hold: The paper's claims are explicitly conditioned on the Mamba-2 architecture (with the caveat that similar principles may apply to other gated linear RNNs), the RedPajama training data distribution, and the specific hyperparameter choices (batch size, learning rate, TBPTT length). The linear and exponential scaling laws are extrapolations from models with ≤780M parameters and ≤12.9M state parameters; their validity at larger scales is plausible but untested. The finding that forgetting is learned only when training length exceeds state capacity is likely robust for any architecture with a similar linear recurrent update rule, but the specific coefficient (5.172) is likely architecture-specific and would differ for other state sizes, head dimensions, or gating mechanisms.
6. Limitations and Trade-offs
Difficulty Estimation Cost Is Unaccounted for in the Headline Claims
The assumption or constraint: The paper's central practical prescription—train Mamba-2 models with context length exceeding T_forget to achieve robust length generalization—depends crucially on knowing what T_forget is for a given model. Determining T_forget requires an expensive sweep: training the model at multiple context lengths and evaluating whether loss explodes on 1M-token prompts at each length. The paper does not propose a method for estimating T_forget without performing this sweep. The authors acknowledge the training cost implicitly: "we do not have enough resources to train the [780M] model beyond [128K training length]" (Section 5.2), confirming that even the authors could not fully characterize the threshold for their larger model.
The consequence: A practitioner wanting to deploy a Mamba-2 variant at a specific scale cannot simply plug their state size into the linear formula T_forget = 5.172 · N_S − 4.469 with confidence, because this formula was derived from only 5–6 data points on one architecture with one training recipe. The coefficient 5.172 is likely sensitive to model architecture (head dimension, number of layers, gating design), training data distribution, and hyperparameters (batch size, learning rate, TBPTT configuration). Using this formula for a different Mamba variant, a different dataset, or a different optimizer configuration could substantially misestimate the required training length. Underestimating leads to models that still fail catastrophically on long contexts; overestimating wastes compute on unnecessarily long training sequences. The paper provides no method for cheaply estimating T_forget—no proxy metric computed during training, no scaling law with additional conditioning variables (data distribution, architecture details), and no early-stopping criterion based on in-training statistics. The cost of the characterization sweep itself is comparable to training the model itself and is emphatically not included in any efficiency or cost analysis. A practitioner cannot answer "how long should I train?" without effectively doing the experiment first.
What evidence exists in the paper: The paper explicitly acknowledges this gap in Section 5.2: the 780M model at 128K training length still exhibits poor length generalization, and the authors state they lack the resources to train beyond this point to confirm the predicted ~95K threshold. The linear fit relies on only 5 data points with R^2 > 0.999 (Figure 11), but this tight fit is within a specific experimental setup and does not guarantee transferability. No ablation over data distribution, optimizer, batch size, or TBPTT configuration is performed to test the robustness of the 5.172 coefficient. The exponential recall fit (Figure 9) relies on the same models and same training procedure, inheriting the same fragility.
Mitigation status: The paper does not address this limitation. There is no attempt to develop a cheap proxy for T_forget (e.g., measuring the rate of state variance growth during training at different sequence lengths, or evaluating retrieval accuracy at 2× training length as an early indicator). The authors do not propose a method for estimating the coefficient without a full sweep. Future work on predictive models of T_forget given architectural and data characteristics is implicitly suggested by the existence of the scaling law, but not explicitly called out as a necessary next step.
Single Architecture Family Evaluated; Generality to Other Recurrent Models Is Unproven
The assumption or constraint: The paper's quantitative findings—the T_forget linear scaling law, the T_recall exponential scaling law, the state explosion diagnostics, and the effectiveness of the sliding window and RRI interventions—are all demonstrated exclusively on Mamba-2 models. Mamba-1 and RWKV are evaluated only in inference mode on passkey retrieval using their official checkpoints (Appendix D), with no controlled training sweeps and no measurement of T_forget for these architectures. The paper explicitly restricts scope: "RWKV-7, xLSTM, and Titans have gone beyond a gating-based memory decay mechanism and are out of the scope of this paper" (Section 6). Even the Mamba-2 results are from a single training recipe (RedPajama-V2 with specific filtering, WSD scheduler, 0.5M batch size, TBPTT with 12 sequences).
The consequence: The claim that "some conclusions/insights may apply to other architectures" (Section 2.1) is speculative. The weighted-sum state decomposition (Equation 6) applies to any linear recurrent model, suggesting the interference mechanism is general. However, the specific relationship T_forget ∝ N_S could depend critically on how the state is structured and parameterized. Mamba-2 allocates N = 128 dimensions per head with P = 64, meaning its state is factorized as H × N × P. Mamba-1 uses N = 16, RWKV uses a completely different gating mechanism, and GLA uses matrix-valued decays—each architecture may have a different relationship between nominal state size (parameter count) and effective information capacity. A practitioner working with RWKV, GLA, RetNet, or HGRN-2 cannot use the 5.172 coefficient; they cannot even be certain that a linear relationship exists for their architecture, or that extending training length will improve length generalization rather than simply wasting compute. The paper's comparative passkey evaluation (Figures 13–15, Appendix H) shows that RWKV-5 and RWKV-6 exhibit less catastrophic failure than Mamba variants, but provides no causal evidence that this is due to state size rather than other architectural differences (token mixing, feedforward components, time mixing design).
What evidence exists in the paper: The passkey retrieval evaluations for Mamba-1 (Figure 15) show similar catastrophic failure patterns to Mamba-2, which is consistent with the overparameterization hypothesis since Mamba-1 also has a large state (though 8× smaller than Mamba-2 at comparable model sizes; Appendix A.1). RWKV-5 and RWKV-6 (Figures 13, 14) show more graceful degradation. The state size comparison across architectures (Figure 12) shows RWKV models have much smaller states than comparably-sized Mamba-2 models. However, these are purely correlational observations on pre-existing checkpoints with unknown training lengths and procedures—no controlled experiments vary training length for non-Mamba architectures to test whether their behavior follows the same T_forget dynamics.
Mitigation status: The paper partially acknowledges this scope limitation by stating that exhaustive ablation studies across architectures are "left for future work" (Section 2.1) and by noting which architectures are out of scope (Section 6). However, the paper's title ("Stuffed Mamba") and abstract's framing as revealing "a critical limitation in current RNN architectures" (emphasis added) overclaims relative to the evidence, which is specific to Mamba-2 with suggestive extensions to Mamba-1 and RWKV. A practitioner reading the abstract would reasonably believe the findings apply to RNNs broadly, but the empirical support is narrow.
No Demonstration That Learned Forgetting Translates to Realistic Long-Context Tasks
The assumption or constraint: The paper's evaluation of whether forgetting has been successfully learned relies on two metrics: (1) whether language modeling loss remains bounded on 1M-token prompts (the T_forget criterion, Section 4.2), and (2) passkey retrieval accuracy as a function of context length (Section 5.3). Passkey retrieval is an intentionally simple synthetic task: the model must find and recall a 5-digit number embedded in highly repetitive filler text. The paper explicitly acknowledges this simplicity: "in this task, the noisy context is repetitive, thus, the amount of contextual information is largely independent of the context length" (Section 4.3). The model is evaluated exclusively on next-token prediction loss and the passkey task—no realistic long-context benchmarks (e.g., long-document QA, multi-turn conversation, summarization, code completion over long files) are tested on models trained above T_forget.
The consequence: There is a substantial gap between "the model can stably process 1M tokens without exploding" and "the model can usefully leverage information across 256K tokens for practical tasks." The T_recall exponential scaling law is derived from passkey retrieval, which requires storing essentially one piece of information (5 digits) and discriminating it from massive redundancy. Realistic long-context tasks involve storing, integrating, and reasoning over thousands of distinct facts distributed across the context—a fundamentally different information load. A model that perfectly retrieves a passkey at position 200K might fail to answer a question that requires synthesizing facts from positions 50K, 120K, and 180K because the cumulative information exceeds what the state can maintain without interference even after forgetting is learned. The paper's claim that the 370M Mamba-2 "outperform[s] similarly sized transformer models" (Section 5.3) at 256K retrieval is specifically about passkey retrieval, not about general long-context capability—but this distinction is easily missed given the prominence of the claim. A deploying practitioner cannot assume that a Mamba-2 model trained at T_train > T_forget will match a transformer on their specific long-context task just because it passes the passkey test.
What evidence exists in the paper: The paper acknowledges the limited nature of the passkey task in Section 4.3, noting that the information content is "largely independent of the context length" and that this is why exponential recall scaling is possible. No experiments test the models on realistic benchmarks (∞Bench, LongBench, needle-in-haystack with multiple needles, summarization, QA over book-length texts). The language modeling loss curves (Figures 1, 4, 10) show that loss stabilizes when models are trained above T_forget, but stable loss does not guarantee useful downstream performance—the model might still fail to attend to relevant context, confuse facts, or produce incoherent outputs even with bounded loss.
Mitigation status: Not addressed. The paper does not evaluate on any realistic long-context task beyond language modeling loss and passkey retrieval. The authors do not acknowledge this as a limitation or suggest what realistic benchmarks would be appropriate for future evaluation. This is a significant omission given the paper's practical framing: the abstract promises "valuable insights for improving long-context modeling" and "a promising foundation for improving Mamba's performance in long-context modeling," but the gap between the demonstrated capability (passkey retrieval) and useful long-context modeling is never discussed.
The Sliding Window Intervention Proves Causality but Is Not a Deployable Solution
The assumption or constraint: The sliding window intervention (Section 3.2.2) is presented as the strongest causal evidence that over-retention causes the length generalization failure: by mathematically eliminating tokens beyond a window of size w, the model achieves stable loss far beyond its training length (Figure 4, red curve). This intervention works by maintaining two full states (h_t and h_{t-w}) and computing the windowed state as h_t^{(w)} = h_t - α_{t-w+1:t} · h_{t-w}. The paper frames this as an efficient mechanism requiring "one extra state vector and scalar per head."
The consequence: While the sliding window demonstrates causality cleanly, it is not a practical solution for deployment because it fundamentally limits the model's effective context length to the window size w. For the 370M model trained on 8K sequences, a sliding window that prevents explosion must use w ≤ 8K (the training length, since the model has not learned to handle more). This means the model can only ever access the last w tokens—exactly the same limitation as the original catastrophic failure, but now enforced by design rather than discovered by accident. The model cannot leverage information from tokens 1 through t-w for any purpose, even if that information would be useful for the current prediction. This defeats the purpose of long-context modeling entirely. The paper does not evaluate at what window sizes the model's in-training performance degrades—setting w too small would hurt even short-context quality because useful information within the training window would be prematurely discarded.
Furthermore, the sliding window is presented as an inference-time intervention, but it could also be viewed as a training objective: if the goal is a model that naturally learns to window its state, one could train with a loss that penalizes long-range retention. The paper does not explore whether the window size could be made adaptive (learned per head, per token) or whether the sliding window mechanism could be used during training to shape the learned decay behavior.
What evidence exists in the paper: Figure 4 shows the sliding window achieving the best length generalization among all interventions, but the window size w is never specified in the main text or appendices. The paper does not evaluate how short-context performance varies with w, does not measure retrieval accuracy under sliding window, and does not discuss the information loss imposed by the window. The sliding window is treated purely as a diagnostic tool, not as a candidate solution.
Mitigation status: The paper does not address this tradeoff. The sliding window is correctly presented as evidence for the causal role of over-retention, not as a deployment strategy. However, the paper does not discuss what a real solution would look like—for example, training with a loss that encourages learned forgetting (e.g., a penalty on α_{1:t} for large t), or using the sliding window during training to teach the model that information beyond w is unrecoverable, thereby incentivizing it to learn its own forgetting schedule. The gap between "we can force forgetting" and "we can make the model learn to forget appropriately" is left unaddressed.
No Guidance on Optimal Training Length Selection Beyond the Linear Fit
The assumption or constraint: The paper establishes that T_train must exceed T_forget for the model to learn robust forgetting, and provides a linear relationship T_forget = 5.172 · N_S − 4.469 to estimate this threshold. However, the paper does not investigate what happens when T_train substantially exceeds T_forget. Is there a point of diminishing returns? Does training much longer than necessary improve T_recall further, or does it saturate? Does excessively long training hurt short-context performance or introduce other pathologies? The paper sweeps training lengths (4K, 8K, 16K, 32K, 64K) only up to the point where the loss stops exploding (Figure 10), but does not measure performance at lengths far above the threshold.
The consequence: A practitioner who determines that T_forget for their model is approximately 64K faces a practical question: should they train at 64K, or 128K, or 256K? Training at longer lengths is more expensive—the cost scales approximately linearly with T_train for next-token prediction. But the paper provides no information about the benefit curve. The exponential T_recall scaling (Figure 9) suggests that larger state sizes unlock dramatically longer recall, but this is measured at a single training length (presumably just above T_forget for each model). It is unknown whether training the 370M model at 256K instead of 64K would push its T_recall beyond 256K, or whether recall capability is fundamentally bounded by state size regardless of how much training length exceeds the threshold. The absence of a "training length vs. recall length" scaling law means practitioners cannot make cost-benefit tradeoffs about how far above T_forget to train.
What evidence exists in the paper: The paper provides no sweep of T_train for a fixed state size to measure how T_recall scales with T_train once above T_forget. All T_recall measurements are taken at a single T_train per model—the minimum length that achieves stable loss. The loss curves in Figure 10 show that at T_train = 16K (just above the threshold for the 130M model), loss is stable but not identically flat; at T_train = 64K for the 370M, the loss curve is similarly stable. But passkey retrieval is not evaluated at these different training lengths to see whether recall capability improves further.
Mitigation status: Not addressed. The paper does not discuss optimal training length selection beyond the threshold, does not measure the benefit of exceeding the threshold, and does not propose any criterion for choosing T_train given computational constraints. The linear fit is presented as the key finding, but it answers only the minimum requirement, not the optimal allocation. This is a significant omission for a paper that positions itself as providing "valuable insights for improving long-context modeling" and "a promising foundation."
Hard Problems (Informationally Dense Long Contexts) Are Not Addressed
The assumption or constraint: The paper's entire analysis of length generalization and the forgetting mechanism assumes that the challenge of long-context processing is temporal—maintaining stable state dynamics as the number of time steps grows. The passkey retrieval task operationalizes this by keeping information content constant regardless of length: the filler text is highly repetitive, so the model only needs to store ~5 digits of actual information. The paper explicitly notes: "If we train Mamba-2 on passkey retrieval data, the model can theoretically handle infinitely long contexts" (Section 4.3, footnote). This reveals that the paper is studying a specific regime where context length grows but information content does not.
The consequence: Many—arguably most—real-world long-context tasks violate this assumption. Processing a 200K-token legal document, a book chapter, or a long conversation involves contexts where the total information content grows with length: more entities, more facts, more narrative threads, more relationships. The paper's framework does not address this regime. When both context length and information content grow, the model faces two simultaneous pressures: (1) the temporal interference problem analyzed in the paper (many tokens in the weighted sum), and (2) the capacity saturation problem (the state simply cannot store all the distinct information). The paper's finding that T_forget scales linearly with state size suggests that to handle informationally dense long contexts, one would need to scale state size with the information content, which grows with length—exactly the same scaling that transformers face with their KV caches, eliminating the efficiency advantage that motivates recurrent architectures in the first place.
The paper's exponential T_recall result is specifically predicated on "the amount of information in the context [being] largely independent of the context length" (Section 4.3). When information grows with length, the recall capacity would not scale exponentially—it might not scale at all, or might scale sublinearly with state size. The paper provides no analysis, theoretical or empirical, of what happens when information density is high.
What evidence exists in the paper: The passkey task results (Figures 2, 9, 13–15, Appendices B, D) all use the same filler text template with a single 5-digit passkey. No experiment varies the amount of distinct information in the context while holding length constant. No experiment embeds multiple distinct facts or requires the model to integrate information across the context. The language modeling loss curves (Figures 1, 4, 10) are evaluated on RedPajama documents, which contain genuine long-range structure, but the loss measurement alone does not indicate whether the model is correctly tracking entities, resolving references, or maintaining factual consistency across the context—it only measures next-token prediction quality, which can remain reasonable even if the model has "forgotten" specific facts from hundreds of tokens ago in favor of local linguistic coherence.
Mitigation status: The paper does not acknowledge this limitation. The abstract claims the work "provides valuable insights for improving long-context modeling" without distinguishing between the temporally-extended low-information regime (where the findings apply directly) and the informationally-dense long-context regime (where the findings may not transfer). The scaling laws are presented as general properties of the architecture rather than as conditional on low information density. A practitioner seeking to use Mamba-2 for document understanding or multi-turn dialogue would find no guidance on whether training above T_forget is sufficient or whether state capacity fundamentally limits performance on these tasks.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper fundamentally reframes the conversation around recurrent language models and long-context processing. Before this work, the dominant narrative was that RNNs fail on long contexts because their fixed-size states lack sufficient capacity—the state simply cannot hold enough information. The prescription that followed from this diagnosis was architectural: make states larger, design more expressive update rules, add external memory. This paper inverts that diagnosis entirely. The problem is not that the state is too small; it is that the state is too large relative to the training context length, creating a perverse training dynamic where the model learns to retain everything rather than learning to forget selectively. The prescription that follows is methodological, not architectural: train on longer sequences, specifically sequences whose length exceeds the state's capacity to store everything without interference.
This reframing is significant for several reasons. First, it explains a paradox that had no prior resolution: why larger Mamba-2 models (with larger states) have worse length generalization than smaller ones (Figure 1). Under the capacity-limitation hypothesis, larger states should generalize better. Under the overparameterization hypothesis, larger states make the problem worse because they raise the threshold at which forgetting becomes necessary, and when training length is held fixed (as in current practice at 8K), larger models are further below their threshold. The paper provides a quantitative explanation: the 130M model (N_S = 4.8M) has T_forget ≈ 20K, placing 8K training at roughly 40% of the required length, while the 370M model (N_S = 12.9M) has T_forget ≈ 62K, placing 8K training at only 13% of the required length. The degradation is proportional to this gap.
Second, the paper reconciles a tension in the literature between works that documented Mamba's length generalization failure (Jelassi et al., 2024; Waleffe et al., 2024; Ben-Kish et al., 2024) and works that proposed engineering fixes (LongMamba, DeciMamba). The prior fixes attempted to improve generalization by making decay factors closer to 1—retaining more information. The paper's intervention experiments (Figure 4) demonstrate that the opposite direction—inducing more forgetting via RRI scaling or sliding windows—actually improves generalization. The prior fixes were addressing a symptom (loss of information at distance) rather than the cause (interference from insufficiently decayed information), which is why they produced mixed results. The paper's diagnostic framework (retrieval error analysis, Equations 6–7) provides the mechanistic understanding needed to design interventions that target the correct mechanism.
Third, the paper introduces a new category of scaling law—relating state size to training length requirements and recall capacity—that has no direct precedent in the literature. The Chinchilla scaling laws (Hoffmann et al., 2022) relate model size and training data quantity to loss. The current paper's scaling laws relate architecture design choices (state size) to training methodology choices (context length) and inference-time capabilities (maximum recall length). This is a fundamentally different axis of scaling that addresses the temporal dynamics of information processing rather than static loss optimization. The finding that T_forget scales linearly with state size while T_recall scales exponentially (when information density is low) reveals an asymmetry in the cost-benefit tradeoff: scaling state size imposes a linear training cost but unlocks exponential recall benefits, making larger models potentially much more efficient for long-context tasks if one is willing to pay the upfront training cost.
Conceptually, the paper establishes that the memory decay mechanism in learned recurrent systems is not a fixed architectural property but a learned behavior that can be pathologically optimized for the training distribution. Figure 8 demonstrates this directly: as training progresses on short sequences, the model converges toward retaining more and forgetting less, because within the training window, retention is strictly beneficial and the model never experiences the failure mode (interference explosion) that long contexts induce. This is a form of overfitting specific to recurrent architectures—overfitting the temporal dynamics rather than the input-output mapping. This insight has implications beyond Mamba-2: any learned recurrent system trained on sequences shorter than its effective capacity will develop brittle temporal dynamics that fail under distribution shift in the time dimension.
The paper also redirects research attention in the RNN space. Before this work, significant effort was invested in designing more sophisticated state update mechanisms—gated delta networks, matrix-valued states, dynamic recurrence, Titans-style learned memorization. The paper's findings suggest that the bottleneck for current architectures is not the expressiveness of the update rule but the training recipe: simply training existing architectures on longer sequences (above T_forget) may unlock capabilities that more complex architectures cannot achieve if also trained too short. This does not mean architectural innovation is unnecessary—the exponential T_recall scaling suggests that state size is a genuine capacity lever—but it does mean that training methodology is the binding constraint for current models, and that architectural comparisons that do not control for training length relative to state size are potentially misleading. A fair comparison between Mamba-2 and a new architecture must ensure both are trained above their respective T_forget thresholds.
Finally, the work changes how practitioners should evaluate recurrent models. The paper demonstrates that language modeling loss—the standard metric for training and model selection—can be "highly similar" across checkpoints with dramatically different length generalization behavior (Appendix F.1). The loss is only computed on tokens within the training length, so it is blind to the temporal dynamics that govern out-of-distribution behavior. The paper implicitly argues for a new evaluation protocol: any recurrent model intended for long-context use should be evaluated on loss curves at extended lengths (as in Figure 10) and on passkey-style retrieval tasks that specifically probe whether interference is accumulating. This is not a minor methodological tweak—it is a fundamental change in what "good performance" means for recurrent architectures, shifting from "low loss within the training window" to "stable dynamics at arbitrary length."
Follow-Up Research This Work Enables
Cheap estimation of T_forget during training. The paper's most immediate gap is the cost of determining T_forget: it currently requires training the model at multiple context lengths and evaluating loss on 1M-token prompts at each length. This sweep is as expensive as training the model itself. A strong follow-up would develop a proxy metric computable during a single training run that predicts whether the current training length exceeds T_forget. One candidate: monitor the variance of the recurrent state h_t at positions near the training length. The paper shows (Figures 5, 16) that when T_train < T_forget, state variance explodes beyond the training length. If this explosion begins to manifest within the training window as T_train approaches T_forget—perhaps as a subtle increase in variance or the emergence of outlier channels—then tracking these statistics during training could provide an early warning that the training length is insufficient. A practical experiment: train the 370M model at 16K, 32K, 48K, and 64K, and at each length, measure the maximum state variance within the training window. If the variance begins to increase detectably as training length approaches T_forget ≈ 62K, then a simple threshold on state variance could replace the expensive sweep. Alternatively, one could directly measure the cumulative retention strength α_{1:t} at the end of training sequences; if it drops below a threshold (e.g., α_{1:T_train} < 0.5 for some fraction of heads), the model is learning to forget.
Direct measurement of state information capacity. The paper infers that T_forget corresponds to the point where information content exceeds state capacity, but never directly measures capacity. A follow-up could design a controlled experiment: train Mamba-2 variants to store and recall random token sequences of varying length, with retrieval tested immediately after the sequence. By sweeping state size and sequence length, one can measure the probability of correct recall as a function of stored tokens and state parameters, yielding an empirical estimate of information capacity in bits per parameter. This would test whether the linear relationship T_forget ∝ N_S arises from a constant bits-per-parameter capacity (the slope 5.172 would then represent tokens per unit capacity times the average information per token in RedPajama data). If the measured capacity is significantly different from what the linear fit implies, alternative explanations for T_forget (such as gradient signal quality rather than information saturation) would need to be considered. This experiment is newly tractable because the paper provides a clear operational definition of the phenomenon to be explained.
Training length sweeps for non-Mamba architectures. The paper demonstrates the T_forget phenomenon for Mamba-2 and provides suggestive passkey retrieval evidence for Mamba-1 and RWKV, but no controlled training length sweeps for non-Mamba architectures. A direct follow-up would replicate the Figure 10 analysis for RWKV-6, GLA, and RetNet: take each architecture at a fixed scale (e.g., matching the 370M Mamba-2 in parameter count), train at multiple context lengths (8K, 16K, 32K, 64K, 128K), and determine whether a similar T_forget threshold exists, whether it scales with each architecture's state size, and whether the relationship follows the same linear coefficient (5.172) or a different one. RWKV's smaller state size for comparable model scale (Figure 12) predicts a lower T_forget and thus better length generalization at standard training lengths—a prediction the paper's passkey data partially supports but does not causally confirm. This experiment would establish whether "state overparameterization" is a universal phenomenon in gated linear RNNs or specific to Mamba's design choices (large head count, N = 128, coupling of ∆_t to both decay and insertion). A negative result—finding that RWKV learns forgetting even at short training lengths despite having a nominally smaller state—would suggest that state structure (how the state is organized across heads and layers) matters more than raw parameter count, refining the overparameterization hypothesis.
Information-dense long-context evaluation. The paper's passkey retrieval results establish that when information content is constant, recall length scales exponentially with state size. But realistic long-context tasks involve contexts where information content grows with length. A critical follow-up would evaluate Mamba-2 models trained above T_forget on tasks with controlled information density. One design: the multi-needle retrieval task, where K distinct passkeys are embedded in a context of length T, and the model must recall a specified one. By varying K (information content) independently of T (context length), one can map out the two-dimensional capability surface (K, T) as a function of state size and training length. Does the model's performance follow a total-information budget (K × information-per-needle roughly constant for a given state size), or does the temporal interference analyzed in the paper dominate even when K is small? If the former, the exponential T_recall result would not generalize to information-dense settings; if the latter, Mamba-2 might be genuinely competitive with transformers for document understanding. This experiment would also test the paper's implicit claim that the forgetting mechanism, once learned, enables the model to selectively retain important information while discarding filler—a claim that passkey retrieval partially supports but that needs verification with multiple distinct pieces of information.
Training above T_forget: how much is enough? The paper establishes that T_train must exceed T_forget but provides no guidance on how far above the threshold one should train. A natural follow-up would sweep T_train / T_forget ratios (1×, 2×, 4×, 8×) for a fixed state size and measure T_recall at each ratio. This would produce a "training length scaling law" complementary to the state size scaling law: for a given model, how does maximum recall length improve as training length increases beyond the minimum threshold? Does T_recall saturate once T_train exceeds T_forget by some factor, or does it continue to improve? Does the exponential relationship with state size hold at all training length ratios, or does the base of the exponential (1.365) depend on how far above T_forget one trains? This experiment has direct practical implications: a practitioner who can afford to train at 128K for their 370M model (2× the ~62K threshold) needs to know whether the additional cost buys proportionally more recall capability or primarily provides a safety margin. The paper's current data suggests that models at exactly T_forget achieve the reported T_recall values, but this is not explicitly tested by evaluating the same model at multiple training lengths above threshold.
Combining the sliding window mechanism with training. The sliding window intervention (Section 3.2.2) proves that eliminating old tokens prevents interference, but it does so destructively—information beyond the window is permanently lost. A research direction suggested but not explored by the paper is using the sliding window as a training objective to teach the model to learn its own forgetting schedule. During training, one could apply a sliding window of size w to the state before computing the output, meaning the model's predictions at position t can only use tokens t−w+1 through t. This forces the model to learn that information beyond w is unrecoverable, which should create gradient pressure to (a) store important information aggressively within the window, and (b) learn to produce small α_t values that rapidly decay tokens once they exit the window, since retaining them only contributes to interference without being usable for prediction. At inference time, the window could be removed entirely—if the model has learned robust forgetting during training, the state should remain stable without the window. This is a form of "curriculum learning" for temporal dynamics: start with a small w and gradually increase it during training, teaching the model to handle progressively longer contexts. A concrete experiment: train Mamba-2 370M at 64K sequence length with a sliding window of w = 64K (equivalent to no window, since the sequence never exceeds the window), and compare against training with w = 16K initially, increasing to 64K over the course of training. If the curriculum-trained model achieves better T_recall or more stable loss at 1M tokens, it suggests that explicit forgetting pressure during training produces better temporal dynamics than simply exposing the model to long sequences.
Practical Applications and Downstream Use Cases
Cost-efficient long-document processing with smaller models. The paper's most directly actionable finding for practitioners is that a properly trained 370M-parameter Mamba-2 can perform passkey retrieval at 256K context length, and the exponential scaling law (Figure 9) suggests that modestly larger models could extend much further. For applications that involve searching for specific information in very long documents—contract review, e-discovery, academic literature search, regulatory compliance checking—this opens the possibility of using sub-billion-parameter recurrent models instead of much larger transformers. The key practical requirement is that the model must be trained with context length exceeding its T_forget threshold; for the 370M model, this means at least ~64K tokens of training length. The computational advantage during inference is substantial: a 370M Mamba-2 processes each token in constant O(1) time and memory regardless of context length, while a transformer of comparable capability at 256K context length would need either a much larger model or significant KV cache compression, both of which increase cost. The paper's finding that even the 130M model (N_S = 4.8M, T_forget ≈ 20K) can achieve T_recall ≈ 16K when properly trained means that very small recurrent models can handle moderately long contexts—16K tokens covers most practical document lengths—at a fraction of the cost of transformer alternatives. The caveat, as discussed in Section 6, is that these results are demonstrated only on passkey retrieval with low information density; deployment on realistic document tasks would require validation on appropriate benchmarks first.
Streaming applications with unbounded context. The sliding window mechanism described in Section 3.2.2, while presented as a diagnostic tool, is directly deployable for streaming applications where only recent context matters. The implementation requires maintaining two state vectors (h_t and h_{t-w}) and one scalar per head—roughly doubling the state memory—with negligible additional per-step computation. For applications like real-time transcription, live captioning, or continuous speech recognition, where the model must process an unbounded stream of input but typically only needs the last few thousand tokens of context, the sliding window provides a mathematically exact guarantee that old tokens contribute zero to the current state. This eliminates the need for complex state reset heuristics or periodic context truncation that can introduce artifacts. The window size w can be chosen based on the task requirements; the paper's analysis suggests w should not exceed the training length for models trained below T_forget, but for models trained above T_forget, w could potentially be much larger. A practical deployment might train at T_train = 64K (above T_forget for the 370M model), set w = 64K at inference, and process arbitrarily long streams with stable loss and constant memory. This is a concrete engineering solution enabled by the paper's diagnostic framework even without solving the general long-context problem.
Checkpoint selection and model evaluation for recurrent models. The paper's finding that validation loss is a poor predictor of length generalization (Appendix F.1: "the loss of many checkpoints was highly similar [but] their performance in passkey retrieval can differ a lot") has immediate practical implications for model development pipelines. Teams training recurrent models should add two diagnostics to their standard evaluation suite: (1) a "length stress test" that measures loss as a function of position on 2× to 4× training length prompts, flagging checkpoints where loss begins to increase sharply beyond the training window, and (2) a passkey-style retrieval task at multiple context lengths to directly probe whether interference is accumulating. These diagnostics are cheap—they require only forward passes on synthetic data—and can be run periodically during training to detect when the model is converging toward a "retain everything" strategy (Figure 8) before it becomes pathological. The paper's training dynamics experiment (Figure 8) provides a concrete timeline: for 512-token training, the degradation becomes visible between 10B and 20B tokens. For longer training lengths, the timeline would shift, but the principle—that retrieval accuracy at extended lengths degrades before in-training loss shows any problem—is likely robust. Incorporating these diagnostics into the training loop would allow early stopping before length generalization collapses, or trigger a switch to longer training sequences.
Guiding architecture design through state-size-aware training length selection. For teams designing new recurrent architectures, the paper's linear scaling law provides a back-of-the-envelope design constraint: the training length must grow proportionally with the total state size. When comparing architectural variants, one should either match training length as a multiple of T_forget (e.g., train both models at 2× their respective T_forget) or explicitly account for the different training costs imposed by different state sizes. A proposed architecture that achieves lower loss than Mamba-2 at the same model size and training length might simply have a smaller state (and thus a lower T_forget, meaning it learns forgetting at the given training length while Mamba-2 does not). This is not a fair comparison—it confuses architectural quality with training adequacy. The paper's framework enables fair comparisons by providing the T_forget concept: architectures should be compared at training lengths that are equal multiples of their respective forgetting thresholds, or the comparison should explicitly report T_forget and note how the fixed training length relates to it. This is analogous to how the Chinchilla scaling laws enabled fair comparisons of transformer architectures by controlling for compute-optimal training.
When to Prefer This Method
This paper does not propose a new method or architecture that competes with alternatives; it provides a diagnostic framework and training prescription (train above T_forget) for existing Mamba-2 models. The paper does position its recommended training approach against current standard practice (training at 8K regardless of model size) and against prior length-extension heuristics (LongMamba, DeciMamba), but these are not "alternatives" in the sense of competing methods—the standard practice is simply undertraining, and the heuristics are shown to be suboptimal. The paper also does not explicitly compare Mamba-2 trained above T_forget against transformers at matched scale on realistic tasks, so there is no decision rule for choosing between architectures. A "prefer A when" framework is therefore not applicable to this work. The paper's contribution is diagnostic and prescriptive within the Mamba-2 family: given a Mamba-2 model, you must train it with T_train > T_forget = 5.172 · N_S − 4.469 to achieve robust length generalization. Whether the resulting model outperforms a comparable transformer on a specific task remains an open empirical question that the paper does not resolve.
