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.