Mamba Modulation: On the Length Generalization of Mamba

The quadratic complexity of the attention mechanism in Transformer models has motivated the development of alternative architectures with sub-quadratic scaling, such as state-space models. Among these, Mamba has emerged as a leading architecture, achieving state-of-the-art results across a range of language modeling tasks. However, Mamba's performance significantly deteriorates when applied to contexts longer than those seen during pre-training, revealing a sharp sensitivity to context length extension. Through detailed analysis, we attribute this limitation to the out-of-distribution behaviour of its state-space dynamics, particularly within the parameterization of the state transition matrix A. Unlike recent works which attribute this sensitivity to the vanished accumulation of discretization time steps, exp(-sumt=1NDeltat), we establish a connection between state convergence behavior as the input length approaches infinity and the spectrum of the transition matrix A, offering a well-founded explanation of its role in length extension. Next, to overcome this challenge, we propose an approach that applies spectrum scaling to pre-trained Mamba models to enable robust long-context generalization by selectively modulating the spectrum of A matrices in each layer. We show that this can significantly improve performance in settings where simply modulating Deltat fails, validating our insights and providing avenues for better length generalization of state-space models with structured transition matrices.