ICLR 2026 · Oral

From Markov
to Laplace

How Mamba in-context learns Markov chains

Presented by

Shengbo Gong

Authors

Bondaschi, Rajaraman, Wei, Ramchandran, Pascanu, Gulcehre, Gastpar, Makkuva

arXiv:2502.10178 / EPFL · UC Berkeley · Google DeepMind · Télécom Paris

01

The question underneath

What counts as
intelligence?

One workable answer: infer a rule from a handful of examples and apply it immediately — no retraining, no gradient step.

01

In-context learning

this talk

02

Early stopping before failure

open

03

Lifelong learning

open

04

Zero-shot transfer

open

05

Deploy-time learning

open

Inferring an unseen task from examples in the prompt, and applying it in the same forward pass — with the weights frozen.

Example · mathematical

A rule drawn at random

Example · natural language

A label map drawn at random

Account A

Implicit gradient descent

von Oswald+ 2022 · Dai+ 2023

Account B

Implicit Bayesian inference

Xie+ 2021

Sequential modeling / literal HMM

States move; observations follow.

Infer a transition law from one trajectory, then predict the next observation. In the paper, the emission is identity: z_t = x_t.

LLM demonstrations / HMM-style hypothesis

Meaning is hidden; text and answers reveal it.

Infer the shared task mechanism from demonstrations, then complete the missing answer. Tokens are treated as emissions from a latent task state.

Shared inference template

p(next | context) = ∫ p(next | θ, context) p(θ | context) dθ

02

Two architectures, one estimator

Background

What Mamba is, what Laplacian smoothing is, and why a Markov chain is where they meet.

Autoregressive interface

Prefix in.
Next-token distribution out.

01

Embed every observed token x_t.

02

Pass embeddings through Mamba; its recurrent state moves left to right.

03

MLP + linear + softmax produce P(x_t+1 | x_≤t).

Paper Figure 2. Input tokens are embedded, passed through recurrent Mamba cells and an MLP, projected to logits, and converted to a next-token distribution at every position.

Horizontal arrows between Mamba cells carry the recurrent state. Vertical arrows carry the token representation toward the language-model head.

Mamba-2 1 layer · 1 head — the exact model analysed in this talk

Counts for one fixed context

n — times that context occurred.
n_j — times token j followed that context.

Naive: raw frequency

n_jn

Assigns probability zero to anything unseen — and a single zero sends cross-entropy to infinity.

FIX: ADD-𝛽 BEFORE DIVIDING · LAPLACE, 1814

P_β(x_t+1=j|ctx) = n_j + βn + 2β

Why this exact estimator

The Bayes-optimal closed form

For the unknown binary transition row θ ∼ Dirichlet(𝛽, 𝛽), the posterior predictive is exactly

E[θ_j | n₀, n₁] = n_j + βn + 2β

That closed form is add-𝛽. It is the exact Bayes-optimal predictor for next-token cross-entropy, not a heuristic correction.

So the loss has a known floor — ask not "does it learn?" but "does it reach the optimum?"

03

How this field actually works

Toy examples

You cannot prove anything about a frontier model. So you build the smallest synthetic task that still contains the phenomenon.

01

Linear & sparse regression

Fresh weight vector per sequence — compare against least squares. Garg+ 2022

02

Trees & two-layer nets

Labels from a random tree or random network — tests non-linear classes. Garg+ 2022

03

Random Markov chains

Fresh transition matrix per sequence — the only setting used from here on. Edelman+ 2024

From here on: random binary Markov chains and their known add-𝛽 optimum.

Architecture

Mamba-2, 1 layer, 1 head

Baselines

GPT-2 1–2L, linear attn 1–2L

Optimizer

AdamW, lr 1e-3, cosine

Iterations

10 000

Embedding dim d

grid { 2, 4, 8, 16, 32 }

SSM state dim N

N = d

Expansion factor e

e = 1 ⇒ inner width ed = d

State matrix H_t

H_t ∈ ℝ^(ed)×N = ℝ^d×d

Conv window w

2 … 6

Compute

1× A100, 10–60 min per run

Objective

L(θ) = −Σ_t E [ log f_θ(x_t+1 | x₁^t) ]

averaged over positions t

Plain next-token cross-entropy — nothing bespoke. Its minimizer is provably add-β, so the loss floor is known in advance.

Markov order k

Dirichlet β 1.00

Seed

Batch of 5 one fresh P per row

Five sequences, five different transition matrices. At order k, the recurrent state H_t must keep one transition-count row for each of the 2^k contexts.

Markov order k

Dirichlet β 1.00

Sequence t = 0 / 96

Recurrent state — transition counts

Charted context —

add-β

—

true kernel

—

| error |

—

raw freq.

—

Estimated P(next = 1) for one fixed context

add-β estimate true kernel (hidden) raw frequency

The solid line is computed from counts alone; the dashed line is the kernel the estimator was never told. They converge — that convergence is the in-context learning.

04

The finding

One layer is
enough.

A single-layer Mamba tracks the Bayes-optimal estimator. A single-layer transformer cannot get off the ground.

The 1-layer Mamba curve sits on the optimum. The 1-layer transformer never leaves the marginal; two layers only track it loosely.

Paper Figure 1. Left: predicted next-token probability across a fixed test sequence — 1-layer Mamba overlays the optimal estimator, 2-layer Transformer tracks it loosely, 1-layer Transformer stays flat near 0.5. Right: L1 distance from optimal across Markov orders 1 to 4.

Left: the predicted probability along one test sequence. Right: the L1 gap from optimal across Markov orders 1–4. Mamba's line and the optimum are indistinguishable.

The depth gap comes from where binding happens — not from attention.

05

Ablations

Which part is
load-bearing?

Take the architecture apart one component at a time. The answer is not the part anyone expected.

Bolt a convolution onto a transformer's Q, K, V and one layer suddenly suffices — the curve climbs onto the optimum. The depth requirement was about binding, not attention.

Paper Figure 11. Predicted probability and test loss comparing 1-layer Mamba, 1-layer Transformer, and 1-layer Transformer with convolution added to K, Q, V. Adding convolution lets the single-layer transformer match Mamba and the optimal predictor.

The convolution-augmented single-layer transformer lands on Mamba and the optimum; the plain single-layer transformer stays flat, well above the loss floor.

The state transition factor a_t across 500 positions at convergence is nearly constant at one; an inset zoom shows values around 0.996.

The learned gate stays at a_t ≈ 0.996 everywhere.

The optimal estimator needs every historical count, so the model learns almost no forgetting. Selectivity is available, but this stationary task does not use it.

An order-k transition contains two objects: a context of k tokens and the one token observed after it.

This is the local requirement. It says what feature must enter the recurrence before any counting can be correct.

Paper Fig. 3b · learned, not constructed

Paper Figure 3b. Test loss by Mamba convolution window for Markov orders one through four, collapsing sharply once the window reaches k plus one.

The threshold moves one step for every increase in k.

Below w = k+1 the gap is large; at and above it, the gap collapses to zero. The staircase is the signature of a representational threshold.

The same k creates the global requirement

An order-k binary chain has 2^k contexts. The recurrent state H must preserve a separate predictive count for every one of them.

H_t := all recurrent states, flattened ∈ ℝ^d

For fixed accuracy and precision:
dim(H) = Ω(2^k).

This is an information bound on the total recurrent state, not a quirk of one Mamba parameterization.

Test loss by Markov order for hidden dimensions 2, 4, 8, 16, 32. Each dimension solves orders up to log2(d) and degrades beyond it, confirming d = 2^k is sufficient.

Measured, not just argued: d = 4 handles k = 2 and fails at k = 3. Each curve gives out one order after the last.

So order controls both resources: locally, w ≥ k+1; globally, the state capacity grows exponentially in k.

Test loss gap from optimal by Markov order for one to four layers. Mamba is uniformly low and flat in depth because one layer already has enough capacity at these orders; transformers improve sharply from one to two layers.

Mamba is flat in depth here because one layer already saturates the task.

Transformer / local binding

Layer 1 binds.
Layer 2 retrieves and counts.

Depth supplies the missing local step — which is why a transformer needs two layers.

SSM / Mamba / global memory

Depth and capacity are interchangeable.

Stacking recurrent layers adds capacity linearly through L · C, where C = dim(H) is the recurrent-state capacity per layer.

Gap to add-β collapses to zero exactly when L · C ≥ 2^k.

One layer · grow state capacity C

k \ C

2

4

8

16

1

.000

2

.234

.000

3

.319

.243

.000

4

.374

.326

.209

.000

Zero exactly on and past C = 2^k. Capacity alone must grow exponentially in k.

Fixed C = 4 per layer · grow depth L

k \ L

1

2

4

8

1

.000

2

.000

3

.243

.000

4

.326

.209

.000

Same staircase — depth and capacity enter symmetrically, as the product L · C. Depth does add capacity, but only linearly, while the demand grows as 2^k — so the layers needed still double each order.

06

Consequences

At LLM scale

Does any of this survive on real text — and what did the frontier labs ship once they had to choose?

Same ablation, real language. Mamba-2 at 14.5 M parameters on WikiText-103 — remove one component at a time and read the perplexity.

Model

PPL

Δ

Mamba-2, full

27.55

—

without non-linearities

28.98

5%

without convolution

30.68

11%

without gating factor

32.16

17%

The two components the Markov analysis singled out are the two that hurt most on real text.

An honest caveat

At 12 layers the
advantage nearly
vanishes.

Mamba-2, 110 M

21.38

without convolution

21.46

Likely because other layers learn to approximate the convolution themselves. The mechanism is still needed — it just stops needing its own module.

01

A 1-layer Mamba implements the exact Bayes-optimal add-β estimator.

02

The convolution binds context to successor; recurrent state H accumulates the counts.

03

Order k sets both costs: w ≥ k+1 locally and L × C ≳ 2^k globally.

04

Markov prediction is not language, and in-context learning is only one capability of an LLM.

Thank you

Questions?

Paper

From Markov to Laplace: How Mamba In-Context Learns Markov Chains. ICLR 2026, Oral. arXiv:2502.10178

Architecture

Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752 · Dao & Gu, Transformers are SSMs (Mamba-2 / SSD). arXiv:2405.21060

Code

github.com/Bond1995/Markov-Mamba — and the training-free reproduction shown here, in toy_reproduce.py

Presenter

Shengbo Gong

From Markov
to Laplace

What counts as
intelligence?

Five capabilities

In-context learning

Two accounts

One lens

Background

Complexity

Mamba LM

Mamba block

The data

The estimator

Worked example

Toy examples

Lineage

Setup

Batch lab

Convergence

One layer is
enough.

On the optimum

The measurement

Two layers

One block

Which part is
load-bearing?

Ablation

Convolution

The measurement

Gate open

Order sets the window

Order sets the state

What depth changes

The capacity law

Trade-offs

At LLM scale

Real text

Kimi K3

Take away

Questions?

From Markovto Laplace

What counts asintelligence?

Five capabilities

In-context learning

Two accounts

One lens

Background

Complexity

Mamba LM

Mamba block

The data

The estimator

Worked example

Toy examples

Lineage

Setup

Batch lab

Convergence

One layer isenough.

On the optimum

The measurement

Two layers

One block

Which part isload-bearing?

Ablation

Convolution

The measurement

Gate open

Order sets the window

Order sets the state

What depth changes

The capacity law

Trade-offs

At LLM scale

Real text

Kimi K3

Take away

Questions?

From Markov
to Laplace

What counts as
intelligence?

One layer is
enough.

Which part is
load-bearing?