Skip to yearly menu bar Skip to main content


Poster
in
Workshop: XAI4Science: From Understanding Model Behavior to Discovering New Scientific Knowledge

From Markov to Laplace: How Mamba In-Context Learns Markov Chains

Marco Bondaschi · Nived Rajaraman · Xiuying Wei · Kannan Ramchandran · Razvan Pascanu · Caglar Gulcehre · Michael Gastpar · Ashok Makkuva


Abstract: 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 (S$6$) 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.

Chat is not available.