Probing Length Generalization in Mamba via Image Reconstruction
Authors: Jan Rathjens, Robin Schiewer, Laurenz Wiskott, Anand Subramoney
Organizations: Institute for Neural Computation, Faculty of Computer Science, Ruhr University Bochum, Bochum, Germany · CITEC, Bielefeld University, Bielefeld, Germany · Department of Computer Science, Royal Holloway, University of London, Egham, UK
Mamba has attracted widespread interest as a general-purpose sequence model due to its low computational complexity and competitive performance relative to transformers. However, its performance can degrade when inference sequence lengths exceed those seen during training. We study this phenomenon using a controlled vision task in which Mamba reconstructs images from sequences of image patches. By analyzing reconstructions at different stages of sequence processing, we reveal that Mamba qualitatively adapts its behavior to the distribution of sequence lengths encountered during training, resulting in strategies that fail to generalize beyond this range. To support our analysis, we introduce a length-adaptive variant of Mamba that improves performance across training sequence lengths. Our results provide an intuitive perspective on length generalization in Mamba and suggest directions for improving the architecture.
Figures & tables
Figure 1: Model overview. An input sequence of image patches augmented with spatial information, a separator token, and query tokens specifying spatial locations is processed by a Mamba model. The model produces reconstruction tokens, each predicting an image patch for the queried location. When query tokens densely cover the image, the reconstruction tokens can be assembled into a visualization of the input image.
Figure 2: Performance of Mamba (Left) and Transformers (Right) as a function of image and query token sequence lengths under different architectural settings. Curves show mean MSE on the Omniglot test set with 95% confidence intervals. Avg denotes a baseline that predicts the dataset-wide mean image patch.
Figure 3: Reconstructions across models. Column 1 shows the original image. Columns 2–8 show reconstructions given the indicated number of randomly sampled image tokens. Rows correspond to models trained with different TI values.
Figure 4: Scanning an image in different orders. The left panel shows the original image. The middle and right panels display reconstructions of the image quadrants with different processing orders. Rows correspond to models with different TI values.
Figure 5: ( Left). Sequential processing of two images. Columns 1–2 show the original images. Columns 3–4 show reconstructions after processing tokens from the first image. Columns 5–6 show reconstructions after processing tokens from the first image, followed by those of the second image. Rows correspond to models trained with different TI . (Right). Mean MSE of reconstruction tokens queried at spatial locations corresponding to the provided VI tokens on the Omniglot test set (95% confidence intervals). Lines correspond to models trained with different TI values and a sequence-length–adaptive variant (prepend Mamba).
Figure 6: (Left). Mean MSE of adapted Mamba variants across varying image-token sequence lengths, evaluated with 1024 query tokens on the Omniglot test set, with 95% confidence intervals. (Right). Each column corresponds to a representative block of the model. The top row shows the default model, and the bottom row shows the prepended Mamba variant. The curves depict the evolution of ∥Δt∥2 over the first 128 tokens, indexed by t , for different VI values, averaged across test sequences.
Recurrent linear attention models (RLAs) such as Mamba offer efficient linear-time sequence modeling as an alternative to Transformers, yet their fixed-capacity recurrent states limit long-sequence modeling. Drawing inspiration from hierarchical human memory, we propose Hierarchical Memory Mamba (HMM) to address this limitation. Building upon a pre-trained Mamba backbone, HMM integrates a lightweight working memory that extracts slow paragraph-level semantics (PLS) from the fast sensory memory embedded in the backbone's hidden states. The PLS is subsequently compressed into persistent long-term memory for task-relevant retrieval. The hierarchical processing of semantic information overcomes the representation bottleneck of RLAs and endows HMM cross-task generalization through parametric learning, which is not observed in other long-context enhanced Mamba variants. Evaluations on Passkey Retrieval and LongBench-E tasks demonstrate that HMM improves retrieval success by 34.3--37.1% and reasoning accuracy by 1.6--14.2% over strong Mamba-based models, while adding only 2% extra parameters and with minimal training overhead.
Qinwen Wang, Jieping Luo, Aoxiang Qin +5
The Chinese University of Hong Kong, Shenzhen · Hainan Bielefeld University of Applied Sciences · BrainGalaxy
Length generalization is the ability of sequential models to perform well on context lengths unseen during training. In this work, we show that the challenge of achieving length generalization is related not only to architectural choices such as positional encoding and the attention mechanism but also to the training procedure itself. We study how regularization affects length generalization and find that weight decay can hinder extrapolation. In contrast, dropout improves extrapolation when its placement within the architecture is reconsidered. In that regard, we show that the standard placement of dropout before layer normalization introduces a systematic distributional mismatch, and that applying dropout just before the linear projection resolves this issue. For example, a modified SmolLM3 with sliding window attention, continually pre-trained with dropout, can extrapolate perfectly to 64× on Needle-in-a-Haystack and far beyond the pre-training context size on RULER and HELMET. Mamba2 also benefits from dropout, suggesting an architecture-agnostic nature of the problem. We further propose Variance-Preserving Affine Dropout (VPAD), a new dropout strategy that substantially reduces the resulting pre-activation variance mismatch, leading to further extrapolation improvements in transformers.
Pavlo Vasylenko, Matthias Lindemann, André F. T. Martins +1
Instituto Superior Técnico, Universidade de Lisboa · Instituto de Telecomunicações · ELLIS Unit Lisbon +1
The state space model Mamba has recently emerged as a promising paradigm in computer vision, attracting considerable attention for its efficient handling of long-sequence tasks. Its inherent causal structure makes it particularly well suited for autoregressive pretraining. However, existing autoregressive pretraining methods in vision are largely limited to short-sequence settings and may not fully use Mamba's capacity to model longer contexts. To investigate this setting, we introduce SeparaTors for AutoRegressive pretraining (STAR), a new autoregressive pretraining method for Vision Mamba that explicitly marks the boundaries between different images. STAR increases the patch-token sequence length from 144 to 640 by packing four images and four separator clusters. This is approximately 4.4× the ARM patch-token sequence length. The increase is achieved without changing the resolution of any individual image: we use 192×192 inputs for autoregressive pretraining and 224×224 inputs for downstream classification fine-tuning. With this long-sequence pretraining scheme, STAR-B achieves 83.5% EMA top-1 accuracy on ImageNet-1K after 1,600 epochs of pretraining. The learned representation also transfers beyond in-distribution classification: compared with ARM, STAR-B improves COCO box AP from 46.11 to 46.84 and mask AP from 40.74 to 41.45, while raising the mean top-1 accuracy across five ImageNet robustness benchmarks from 55.1% to 56.8%. Under the evaluated four-image setting, these results indicate that separator-based long-sequence pretraining improves recognition robustness and dense visual prediction relative to ARM.