Probing Length Generalization in Mamba via Image Reconstruction

TL;DR

This paper investigates how Mamba, a sequence model, struggles with longer inference sequences in image reconstruction tasks, revealing adaptive behaviors and proposing a length-aware variant to enhance generalization.

Contribution

The study provides a detailed analysis of length generalization in Mamba and introduces a length-adaptive version to improve performance beyond training sequence lengths.

Findings

Mamba's behavior adapts to training sequence lengths but fails to generalize.

A length-adaptive Mamba improves performance across various sequence lengths.

Analysis offers insights into architectural improvements for better length generalization.

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.