Probing Length Generalization in Mamba via Image Reconstruction

2026-03-12

AI Summary

This work addresses the significant length generalization bottleneck of Mamba models, which struggle to generalize during inference to input sequences longer than those seen during training. Through a controlled visual experiment based on image patch reconstruction, the study reveals for the first time that this performance degradation stems from the model’s strong reliance on the training sequence length distribution. To mitigate this limitation, the authors propose a length-adaptive variant of Mamba that dynamically adjusts its internal state modeling mechanism to accommodate varying input lengths. Experimental results demonstrate that the proposed method consistently improves reconstruction performance across diverse training length configurations, effectively enhancing the model’s ability to generalize to unseen sequence lengths.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: (Large) Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

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.

Problem

• length generalization
• Mamba
• sequence modeling
• out-of-distribution
• image reconstruction

Innovation

• length generalization
• Mamba
• sequence modeling
• image reconstruction
• length-adaptive architecture