🤖 machine learning

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

This paper demonstrates that single-layer Mamba models can efficiently learn the optimal Laplacian smoothing estimator for Markov chains in-context, theoretically establishing the first formal connection between Mamba's convolution-based architecture and Bayes/minimax optimal statistical estimation.

Original authors: Marco Bondaschi, Nived Rajaraman, Xiuying Wei, Kannan Ramchandran, Razvan Pascanu, Caglar Gulcehre, Michael Gastpar, Ashok Vardhan Makkuva

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Marco Bondaschi, Nived Rajaraman, Xiuying Wei, Kannan Ramchandran, Razvan Pascanu, Caglar Gulcehre, Michael Gastpar, Ashok Vardhan Makkuva

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: A New Kind of AI Brain

Imagine the current superstar of Artificial Intelligence, the Transformer (the engine behind most chatbots), as a brilliant but heavy librarian. It can read a whole book and find connections instantly, but it gets tired and slow if the book is too long because it tries to remember every single word at once.

Enter Mamba, a newer type of AI model. It's like a nimble, fast-running courier. It processes information much faster and uses less memory, making it a promising alternative to the heavy librarian. But scientists didn't fully understand how this courier was so smart. They knew it worked well, but they didn't know the secret sauce.

This paper acts like a detective story, figuring out exactly how Mamba solves a specific type of puzzle called Markov Chains.

The Puzzle: Predicting the Next Step

To test Mamba, the researchers gave it a game called "Next Token Prediction."

The Setup: Imagine a sequence of events, like a string of colored beads (Red, Blue, Red, Red, Blue...).
The Rule: The color of the next bead depends on the colors of the beads that came just before it. This is a "Markov Chain."
The Challenge: The model sees a random string of beads and has to guess the next one. Crucially, the rules of the game (how likely Red follows Blue) change for every new string. The model has to figure out the rules on the fly just by looking at the current string. This is called In-Context Learning (ICL).

The Discovery: Mamba is a Perfect Statistician

The researchers found something surprising. Even a single-layer Mamba (a very simple version of the model) learned to become the perfect statistician for this game.

In the world of statistics, there is a "Gold Standard" way to guess the next bead when you don't know the rules perfectly. It's called Laplacian Smoothing (or the "Add-β" estimator).

The Analogy: Imagine you are guessing the next card in a deck. If you've seen 10 Aces and 0 Kings, you might guess the next one is an Ace. But a smart statistician knows, "Wait, I haven't seen a King yet, maybe it's just bad luck." So, they add a tiny "ghost" King to their count to avoid being too sure. This prevents them from saying "Zero chance!" for something they haven't seen yet.

The Paper's Claim: Mamba doesn't just guess; it mathematically learns to do this exact "ghost counting" perfectly. It calculates the counts of previous patterns and adds that tiny bit of "smoothing" automatically, just like the optimal statistical formula requires.

The Secret Ingredient: The Convolution "Flashlight"

The researchers asked: How does Mamba do this? Is it because of its complex gating mechanisms? Its non-linear activation functions?

They ran experiments where they removed parts of Mamba to see what broke.

The Finding: The most important part is the Convolution.
The Analogy: Think of the Convolution as a flashlight that Mamba shines on the recent past.
- To guess the next bead, Mamba needs to know: "How many times did 'Red' follow 'Blue' in the last few steps?"
- The convolution acts like a window that slides over the history, counting these patterns instantly.
- The researchers found that if you take away the flashlight (convolution), Mamba becomes blind and fails the task. If you only keep the flashlight (and remove the complex gating), Mamba still solves the puzzle perfectly.

Key Takeaway: The "flashlight" (convolution) is the hero here. It allows Mamba to look back, count the occurrences of patterns, and apply the optimal statistical smoothing without needing a deep, complex brain.

The Limits: How Big Does the Flashlight Need to Be?

The paper also looked at how hard the puzzle gets.

If the game depends on the last 1 bead (1st-order), a small flashlight works.
If the game depends on the last 5 beads (5th-order), the flashlight needs to be wider to see all 5 beads at once.
The Theorem: The paper proves that to handle a game depending on $k$ previous steps, the model's "memory size" (hidden dimension) needs to grow exponentially with $k$ . It's like trying to remember a password: the longer the password, the exponentially harder it is to keep it all in your head at once.

Comparison with Transformers

The paper compares Mamba to the Transformer (the heavy librarian):

Transformers: To solve this "counting" puzzle, a Transformer usually needs two layers (two brains working together) to build a mechanism called an "induction head" to count the patterns. A single-layer Transformer fails.
Mamba: A single-layer Mamba solves it immediately because its convolution mechanism is built-in and efficient at counting.

Summary

This paper reveals that Mamba's superpower in learning from context comes from a specific architectural feature: Convolution.

Mamba learns to act like a perfect statistician, using a method called Laplacian Smoothing to predict the next item in a sequence.
It achieves this by using a convolutional "flashlight" to count past patterns and apply the correct statistical adjustments.
This happens even in a very simple, single-layer model, whereas Transformers need more complexity to do the same job.

The authors conclude that this is the first time anyone has formally connected Mamba to these optimal statistical estimators, proving that Mamba isn't just a fast model, but one that fundamentally understands how to count and smooth data efficiently.

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

Problem Statement

While Transformer-based models have driven recent AI advancements, their computational complexity (quadratic scaling in sequence length) has spurred interest in efficient alternatives like Structured State Space Models (SSMs), specifically Selective SSMs such as Mamba and Mamba-2. Despite their empirical success and inference speed-ups, the fundamental learning capabilities of Mamba remain poorly understood. Specifically, there is a lack of theoretical understanding regarding Mamba's In-Context Learning (ICL) mechanisms compared to the growing body of work explaining Transformer ICL (e.g., induction heads).

The central question addressed is: Can we systematically characterize the ICL capabilities of Mamba? The authors investigate this by studying Mamba's performance on random Markov chains, a framework previously used to uncover fundamental mechanisms in Transformers.

Methodology

The paper employs the Markov-ICL framework, where input tokens are drawn from random Markov chains of order $k$ .

Data Generation: Transition kernels are sampled independently from a Dirichlet prior, ensuring every test sequence follows a unique Markov distribution. The model must estimate the next-token distribution in-context without explicit parameter updates.
Optimal Baseline: The theoretical optimal estimator for this task is the Laplacian smoothing (add- $\beta$ ) estimator, which is both Bayes and minimax optimal. This estimator requires counting token occurrences within specific contexts and applying additive smoothing.
Architectural Analysis: The authors analyze the Mamba architecture (specifically Mamba-2), focusing on its recurrence, selectivity, and convolution mechanisms. They introduce a simplified variant, MambaZero, which retains only the embedding, convolution-based input selectivity, and linear layers, removing gating and non-linearities to isolate the core learning mechanism.
Theoretical Characterization: The paper provides a constructive proof demonstrating how MambaZero can represent the optimal Laplacian estimator. It also establishes a lower bound on the hidden dimension required for recurrent architectures to represent such estimators.

Key Contributions

1. Mamba Learns the Optimal Laplacian Estimator

The authors demonstrate that even a single-layer Mamba efficiently learns the in-context Laplacian smoothing estimator for all Markov orders.

Empirical Evidence: Experiments show that a 1-layer Mamba matches the performance of the optimal estimator on binary and multi-state Markov chains, significantly outperforming a 1-layer Transformer (which fails the task) and matching a 2-layer Transformer.
Surprising Finding: Unlike Transformers, which typically require multiple layers to implement induction heads (counting mechanisms), Mamba achieves this optimal counting capability with a single layer.

2. The Pivotal Role of Convolution

Through ablation studies, the paper identifies convolution as the fundamental component enabling Mamba's ICL ability, more so than gating or non-linear activations.

Ablation Results: Removing convolution from Mamba causes it to fail the Markov prediction task, even with multiple layers. Conversely, a simplified model containing only convolution (MambaZero) successfully learns the task.
Mechanism: The convolution window size $w$ must satisfy $w \ge k + 1$ (where $k$ is the Markov order) to track the necessary context for counting. The convolution allows the model to maintain a state that effectively counts transitions $i \to j$ based on the current context.

3. Theoretical Characterization and Limits

Representation Capacity: The authors provide a constructive proof (Theorem 1) showing that MambaZero can exactly represent the Laplacian estimator for first-order Markov processes. The mechanism involves using the state $H_t$ to accumulate transition counts, where the convolution ensures the state updates depend on the correct context history.
Fundamental Limits: Theorem 2 establishes a lower bound for recurrent architectures: to represent the optimal estimator for a $k$ -th order Markov process with finite bit precision, the product of hidden dimension $d$ and precision $p$ must scale as $\Omega(2^k)$ . This contrasts with Transformers, which can represent these estimators with hidden dimensions growing linearly in $k$ (using 3 layers).
Selectivity: In non-Markovian "switching" Markov processes, the paper shows Mamba utilizes its selectivity mechanism (the decay factor $a_t$ ) to reset counts when a switch token occurs, demonstrating the adaptability of the learned mechanism.

Results

Markov Tasks: Single-layer Mamba achieves near-zero loss gap from the optimal Laplacian estimator across various Markov orders and state space sizes.
Comparison with Transformers: While 2-layer Transformers can solve the task, 1-layer Transformers cannot. Mamba solves it with 1 layer.
Natural Language: Experiments on the WikiText-103 dataset show that removing convolution from Mamba-2 increases perplexity by 11%, whereas removing it from a Transformer increases perplexity by only 2%. This suggests convolution plays a more critical role in Mamba's performance on complex language tasks than in Transformers.
Switching Processes: Mamba correctly implements a strategy of resetting transition counts upon encountering a "switch" token by setting the state decay factor $a_t$ to 0, effectively mimicking the optimal Bayesian strategy for non-stationary environments.

Significance and Claims

The paper claims to provide the first formal connection between Mamba and optimal statistical estimators.

Theoretical Insight: It reveals that Mamba's ICL capability is not merely an emergent property but is structurally grounded in its ability to implement optimal counting via convolution and recurrence.
Architectural Understanding: The work clarifies that for Mamba, convolution is the primary driver of its ability to learn context-dependent statistics, distinguishing it from the attention-based mechanisms of Transformers.
Generality: The findings extend beyond synthetic Markov data to natural language modeling, suggesting that the mechanisms identified (convolution-driven counting) are fundamental to Mamba's success in broader language tasks.

The authors maintain modesty regarding future directions, noting that while they characterize representation power, the learning dynamics of Mamba's convergence to these estimators remain an open area for future research. They also emphasize that their focus is on understanding Mamba's intrinsic capabilities rather than a direct architectural comparison with Transformers, though such comparisons provide necessary context.