From Markov to Laplace: How Mamba In-Context Learns Markov Chains
This work addresses the knowledge gap in understanding the fundamental learning capabilities of Mamba, which is significant for researchers and practitioners interested in developing efficient and effective alternatives to transformer-based language models.
The authors investigated the in-context learning capabilities of Mamba, a structured state space sequence model, and found that it can efficiently learn the optimal Laplacian smoothing estimator for Markov chains, achieving Bayes and minimax optimality. This result was observed even with a single-layer Mamba.
While transformer-based language models have driven the AI revolution thus far, their computational complexity has spurred growing interest in viable alternatives, such as structured state space sequence models (SSMs) and Selective SSMs. Among these, Mamba (S6) and its variant Mamba-2 have shown remarkable inference speed ups over transformers while achieving comparable or superior performance on complex language modeling tasks. However, despite these architectural innovations and empirical successes, the fundamental learning capabilities of Mamba remain poorly understood. In this paper, we address this gap by studying in-context learning (ICL) on Markov chains and uncovering a surprising phenomenon: unlike transformers, even a single-layer Mamba efficiently learns the in-context Laplacian smoothing estimator, which is both Bayes and minimax optimal, for all Markovian orders. To explain this, we theoretically characterize the representation capacity of Mamba and reveal the fundamental role of convolution in enabling it to represent the optimal Laplacian smoothing. These theoretical insights align strongly with empirical results and, to the best of our knowledge, represent the first formal connection between Mamba and optimal statistical estimators. Finally, we outline promising research directions inspired by these findings.