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

In plain words

A language model reads a prompt and guesses what comes next. When the text follows a hidden pattern, a simple and strong strategy is to count: how often did each symbol follow the last few symbols so far? With enough counts the best guess is the most frequent follower, with a little smoothing so that nothing is ruled out completely. This is an old textbook estimator. The question here is whether Mamba, a recurrent alternative to transformers, can do this counting inside its running memory while it reads the prompt.

The authors answer with made-up sequences. Each sequence comes from its own random rule for what follows what, so the model has to work out the rule from the sequence it is reading. A single Mamba layer learns to match the counting estimator closely. On these sequences the authors find that a convolution step inside the block matters more than the gating and the non-linearity. They also build, by hand, a small Mamba that does the counting exactly for the simplest case, and they show that the memory needed grows quickly as the pattern gets longer.

It matters because it gives a concrete account of how a model without attention can learn from its prompt, and it points to which building block carries that ability.

TL;DR

A one-layer Mamba trained on random binary Markov chains learns, in context, the add-beta (Laplacian) estimator that the authors call Bayes and minimax optimal. Removing the convolution inside the block breaks this, while the simplified block with only the convolution still works. For first-order chains, a Mamba with embedding size 2S, state size S and convolution window 2 can reproduce the estimator exactly; for order k a recurrent model needs a hidden size that grows as 2^k. All numbers are author-reported.

Problem

Transformers learn in context on Markov data, and a two-layer transformer is known to realise the counting through induction-head-like attention §3. For Mamba, a selective state-space model, the in-context capacity on this task was not characterised: the authors ask whether a Mamba can represent and learn the counting predictor §1. The optimal predictor for a chain whose rows are Dirichlet(beta) is the add-beta estimate over the transitions seen so far in the prompt [§2.4, app. A]. The difficulty is that Mamba keeps only a fixed-size state, so it must compress those counts §4.3.

Scope and assumptions: binary chains of order k with a fresh transition table per sequence §2.1; a single Mamba-2 style block without layer norm in the write-up app. B. The theory uses a simplified model with no MLP and L1 normalisation instead of softmax §4.1.

Method

Data §2.1. Each sequence gets its own transition table. Its rows are drawn from a Dirichlet(beta) prior. The first k symbols are uniform and i.i.d.; after that each symbol is drawn from the row selected by the previous k symbols.

# Random binary Markov chains of order k  [§2.1]
# Output: x of shape (batch, T), symbols in {0, 1}
def sample_chains(batch, T, k, beta):
    # One transition table per sequence; each row ~ Dirichlet(beta)  [§2.1]
    P = dirichlet(beta, shape=(batch, 2**k, 2))
    # First k symbols are i.i.d. uniform  [§2.1, step 2]
    x = bernoulli(0.5, shape=(batch, k))
    for t in range(k, T):
        # Row index: last k symbols as a binary number
        # [code: src/optim/utils.py#L92]
        row = binary_index(x[:, t-k:t])
        # Next symbol drawn from that row  [§2.1, step 3]
        x_next = sample_categorical(P[batch_index, row])
        x = append_column(x, x_next)
    return x

Mamba block §2.2. The input-dependent step size is a softplus of a linear function of the token. Three causal convolutions over the last w tokens produce the write vector, the key and the read vector of a linear state recurrence. The state is decayed by a_t and read out by the read vector; the output is then gated by a ReLU branch.

# One Mamba block, one time step  [§2.2, eq. Mamba]
# x_t: (d,)  window: last w embeddings (w, d)  H: state (e*d, N)
def mamba_step(x_t, window, H, params):
    # Input-dependent step size, always positive  [§2.2, eq. Input selectivity]
    delta_t = softplus(dot(params.w_delta, x_t) + params.delta)
    # Causal convolutions over the window  [§2.2, eq. Input selectivity]
    x_tilde = relu(conv_X(params.W_X @ window)) * delta_t   # (e*d,)
    b_t = relu(conv_B(params.W_B @ window))                 # (N,)
    c_t = relu(conv_C(params.W_C @ window))                 # (N,)
    # Decay; a_t = 1 when the learned a is 0  [§2.2, §4.2]
    a_t = exp(-params.a * delta_t)
    # Linear state update: keep past counts, add the new transition  [§2.2]
    H = a_t * H + outer(x_tilde, b_t)                       # (e*d, N)
    # Read out, then gate with a ReLU branch  [§2.2, eq. Mamba]
    y_t = H @ c_t                                           # (e*d,)
    z_t = y_t * relu(params.W_z @ x_t)                      # (e*d,)
    return params.W_o @ z_t, H                              # (d,)

Full model §2.2, app. B. Embedding, the Mamba block with a residual, a gated MLP with a residual, and a linear head with softmax.

# Full language model over one sequence  [§2.2] [app. B]
# tokens: (T,) symbols; probs: (T, S) next-symbol distributions
def forward(tokens):
    # Look up the embedding of each token  [app. B, eq. Embedding]
    x = embedding[tokens]                                   # (T, d)
    # Residual Mamba output over the prefix  [app. B, eq. Mamba]
    u = x + mamba_over_sequence(x)                          # (T, d)
    # Gated MLP with residual  [app. B, eq. MLP]
    v = u + W2 @ (relu(W1 @ u) * (W3 @ u))                  # (T, d)
    # Logits, then softmax over symbols  [app. B, eq. Linear, eq. Prediction]
    logits = v @ W_l.T                                      # (T, S)
    return softmax(logits, axis=-1)                         # (T, S)

Loss §2.3. Cross-entropy of the next symbol at every position, averaged over positions and over sequences from the chain distribution.

# Next-symbol cross-entropy  [§2.3, eq. loss]
def loss(probs, tokens):
    # Probability the model gave to each true next symbol  [§2.3]
    picked = probs[:-1][arange(T - 1), tokens[1:]]          # (T-1,)
    return -mean(log(picked))

Target §2.4, app. A. The Bayes-optimal predictor is the add-beta estimate computed from the transitions seen so far in the same prompt.

# Laplacian (add-beta) predictor for a binary symbol  [§2.4] [app. A]
def laplace_estimate(past, context, beta):
    # n_j: how often symbol j followed this context so far  [§2.4]
    counts = count_followers(past, context)                 # (2,)
    # n: total number of occurrences of the context  [§2.4]
    total = sum(counts)
    return (counts + beta) / (total + 2 * beta)             # (2,)

Key equations (cited by the tag the paper prints, because equation numbers were not checked against a compiled PDF):

H_t = a_t H_{t-1} + \tilde{x}_t\, b_t^\top,\quad y_t = H_t\, c_t,\quad a_t = \exp(-a\,\Delta_t),\quad \Delta_t = \mathrm{softplus}(\langle w_\Delta, x_t\rangle + \delta) §2.2, eq. Mamba and eq. Input selectivity

\hat p_j = \frac{n_j + \beta}{n + 2\beta} §2.4, eq. Laplacian smoothing

What the trained model does §3, §4.2.1. Four observations hold at convergence: (i) a_t is close to 1 at every position, so the state keeps all past counts fig. 4; (ii) the block's convolution window must be at least k+1 for order k app. C.1; (iii) the read vector is orthogonal to the write vectors of every transition that does not start with the current symbol, so only the right counts reach the logits fig. 5; (iv) the count-independent part of the logit converges to beta = 1 fig. 5.

Construction for first-order chains §4.2, app. C.2. With w = 2, the embedding size d = 2S, state size N = S and e = 1, the authors give one choice of parameters whose output equals the add-beta predictor for every sequence and every t (Theorem 1). A binary-case variant uses d = N = 2 and matches the predictor up to a KL error epsilon §4.2, app. C.3.

Lower bound §4.3, app. D. For any recurrent model with hidden size d and precision p bits that approximates the beta = 1 predictor to within epsilon in L-infinity, the paper derives d * p >= 2^k (1 - 3 epsilon) log(1/epsilon) (Theorem 2). The argument counts how many distinct predictor tables the hidden state must separate.

Hyperparameters

Paper values come from the appendix tables tab. 7 and tab. 9. Values from the repository are cited to code where they differ.

Results

Author-reported. Setup: binary Markov chains of order k, one-layer Mamba-2 style model, trained with AdamW on random transition tables §3, app. F. WikiText-103 language-modelling runs (about 14.5M parameters unless noted) are described in §5.2; the paper does not give their training schedule.

Paper vs code

Official code: github.com/Bond1995/Markov-Mamba, read at 1c14026538ad413e04d64dac6f9f145ac8fcb8dc. The paper links it in the abstract. The card reads arXiv v2 (metadata updated 2026-06-22); the code commit is the one we cloned. Repository has no LICENSE file; the README credits state-spaces/mamba and nanoGPT README.md. Cite only; no code is quoted.

• Prior. Paper: rows are Dirichlet(beta) with beta > 0 §2.1. Code: each row is drawn uniform on [0, 1], which is beta = 1, and there is no beta option utils.py. The count estimator also defaults to beta = 1 utils.py.
• Reference loss. Paper: the gap |L(theta) - L_beta| is measured against the Laplacian loss §3. Code: the logged val/loss_gap is measured against optimal_est, which uses the true transition table P, so it is the oracle entropy, not the add-beta loss utils.py base.py base.py. The probability comparison uses the add-beta estimator with beta = 1 utils.py base.py.
• Output normalisation. Paper: softmax for the full model §2.2, eq. Prediction; L1 normalisation only for the simplified model §4.1. Code: flag --l1_loss, which the main Markov run script passes schedule.sh. With it, logits become (|logit| + 1e-4) / 100 and are then L1-normalised mamba_llm.py. The absolute value is not in the paper's equation. The paper does not say which normalisation produced each figure.
• Activation. Paper: ReLU after the convolutions and in the gate §2.2. Main run script: --conv_act --gate --activation silu, so SiLU is used schedule.sh mamba2.py mamba2.py.
• Gated MLP. Paper: relu(W1 u) times W3 u app. B, eq. MLP; the parameter table lists mlp.fc1 as 4d x d tab. 6. Code with --gate: one fc1 maps d to 8d, split into two halves, with the activation on the gate half mlp.py. The table shape does not match the gated code.
• Learning rate and length. Paper table: 0.001 and 10000 iterations tab. 7. Main run script: 0.002 and 2000 iterations schedule.sh schedule.sh. Schedule: start at lr/100, warm up over 2% of steps, then anneal to lr/5 (final_div_factor 0.05) main.py.
• Expansion. Text: e typically 2 §2.2. Script: 1 schedule.sh. Config default: 2 markov.py.
• Convolution window. Theory: w = k + 1 = 2 suffices for order 1 §4.2 app. C.1. Script: d_conv = 4 schedule.sh. Config default: 4 markov.py.
• Layer norm. Paper omits norms "for simplicity" app. B. Script uses --layernorm schedule.sh, which adds a norm before the mixer, before the MLP and at the output mamba_llm.py mamba_llm.py mamba_llm.py.
• Decay parameter. Paper: a >= 0 is learned §2.2. Code: A = -exp(A_log), with A drawn uniformly from [1, 16] at initialisation mamba2.py mamba2.py. A fixed-decay option (a_t = 1) exists as --fix_A mamba2.py; the paper's runs do not use it.
• Step size. Paper: softplus of a linear function plus delta §2.2. Code matches mamba2.py. The step-size bias is initialised log-uniformly between 0.001 and 0.1 mamba2.py; the paper does not state this.
• Simplified model. Paper: MambaZero, which drops the MLP and ReLU and uses L1 normalisation §4.1. Code: no class of that name in Markov-Mamba/src; flags have to be combined by hand (for example --no_mlp without --conv_act). The repository does not show which flags produced the MambaZero figures.
• No-convolution sweep. Paper: one-layer Mamba without convolution across widths 10, 100 and 1000 tab. 3. Script: --no_mlp, sequence length 512, 5000 iterations, state size equal to width schedule2.sh schedule2.sh schedule2.sh schedule2.sh.
• Unstated details in the code. Weight init: normal with std 0.02 for Linear and Embedding; output projections scaled by 1 / sqrt(2 * layers) mamba_llm.py. Gradient clipping at 1.0 base.py. Weight decay is excluded for norms, embeddings, convolutions, biases, A and the step-size bias mamba_llm.py mamba2.py mamba2.py.
• Switching experiments. Separate folder Markov-Mamba-Switch; it resamples the transition table after each switch token utils.py. Not checked line by line against the paper's appendix.
• Stale README. It lists a Markov-Mamba-L1 folder that is not in the repository; the L1 option is the --l1_loss flag README.md.

Pitfalls

• Step size is never zero. Section 4.2.1 says a_t = 1 can be reached by setting "a or Delta_t to be zero". Delta_t is a softplus, so it is strictly positive; a_t = 1 needs a = 0. The proof sets a = 0 and Delta_t = 1 by w_Delta = 0 and softplus(delta) = 1 app. C.2.
• a_t is trained, not imposed. The paper observes a_t close to 1 at convergence fig. 4. Only the proof forces it. An implementation that does not learn it will not reproduce the counts.
• Exact construction is fragile. Theorem 1 uses interleaved one-hot embeddings, a specific shift matrix and a specific output matrix app. C.2. Trained models do not have these weights; the paper checks the mechanism empirically fig. 5.
• Probable typo in the proof. The text says "D_X^{(1)} = 0" when setting the B convolution app. C.2. The intended term is probably D_B^{(1)} = 0.
• Window notation disagrees. Section 4.2 says w_X, w_B <= k would fail; appendix C.1 says w <= k. The reading that matches the proof is: the convolution window must be at least k + 1 for X and B, and may be k for C app. C.1.
• Binary construction needs sign conditions. The binary-case theorem requires specific kernel sign patterns and a relation between kernels and beta epsilon app. C.3. It gives KL <= epsilon, not equality.
• Lower bound scope. Theorem 2 assumes an L-infinity error bound. The authors believe a KL bound gives the same scaling but do not prove it app. D.
• Hold-out gap is large early. The model is far from the optimum for short prefixes (0.378 at t = 1) and converges later tab. 2. Short prompts are the worst case.
• Transformer comparisons are narrow. The "sharper than a transformer" claim is for first-order sources §4.3; higher-order comparisons are in the appendix fig. 10.
• WikiText-103 setup is missing. Section 5.2 points to the appendix app. F, which lists only the Markov setups. Tokeniser, learning rate and number of steps for language modelling are not stated.
• Compute. Each experiment ran on one A100 for 10 to 60 minutes app. F; the paper's settings use 5 repetitions each tab. 7.

Lineage

• Builds on: the Markov in-context framework of Edelman et al. (2024), first applied to transformers §2.1. Laplace's add-constant estimator and its optimality §2.4, app. A. The Mamba-2 block as implemented in Dao and Gu (2024) §2.2.
• Compared against: softmax transformers with one and two layers, and linear attention §3, app. E. Lower-bound comparison with the known three-layer transformer construction for order-k chains §4.3.
• Superseded by: not known from this wiki.