From Markov to Laplace: How Mamba In-Context Learns Markov Chains
Authors: Marco Bondaschi
Nived Rajaraman
Xiuying Wei
Razvan Pascanu
Caglar Gulcehre
Michael Gastpar
Ashok Vardhan Makkuva
Venue: ICLR 2026 Oral (2026)
Keywords: State-space models
Markov chains
In-context learning
Laplacian smoothing
Primary area: interpretability and explainable AI
TL;DR: We uncover an interesting phenomenon where a single-layer Mamba represents the Bayes optimal Laplacian smoothing estimator when trained on Markov chains and we demonstrate it theoretically and empirically.

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 (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 an interesting phenomenon: even a single-layer Mamba efficiently learns the in-context Laplacian smoothing estimator, which is both Bayes and minimax optimal. 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.

Published as a conference paper at ICLR 2026
FROM MARKOV TO LAPLACE: HOW MAMBA IN-CONTEXT LEARNS MARKOV CHAINS
Marco Bondaschi EPFL
Nived Rajaraman UC Berkeley
Xiuying Wei EPFL
Razvan Pascanu Google DeepMind
Caglar Gulcehre
Michael Gastpar EPFL
Ashok Vardhan Makkuva Télécom Paris
EPFL
ABSTRACT
While transformer-based language models have driven the Al revolution thus far, their computational complexity has spurred growing interest in viable al- ternatives, such as structured state space sequence models (SSMs) and Selec- tive 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 learn- ing capabilities of Mamba remain poorly understood. In this paper, we address this gap by studying in-context learning (ICL) on Markov chains and uncover- ing an interesting phenomenon: even a single-layer Mamba efficiently learns the in-context Laplacian smoothing estimator, which is both Bayes and minimax op- timal. To explain this, we theoretically characterize the representation capacity of Mamba and reveal the fundamental role of convolution in enabling to repre- sent 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. Code is available http::///itthu...ooooo.oo Markoro Mamba
INTRODUCTION
Transformers have been the forefront recent breakthroughs in language modeling, driving the AI revolution (Vaswani et al. 2017: Radford & Narasimhan 2018 Devlin et al.) 2018). Despite their empirical success, transformers suffer from high computational complexity, such as quadratic scaling in sequence length during training and linear cache size at inference (Gu & Dao (023a). To address these limitations, there is a growing interest in designing alternative efficient architectures among which structured state space models (SSMs) are the most prominent, In particular, Selective SSMs such as Mamba and Mamba-2, have achieved state-of-the-are results in various language modeling tasks, while greatly improving the inference throughput (Cirone al. 2025)..
Motivated by this success, there is tremendous interest in understanding the sequential modeling abilities of SSMs, especially that of Mamba. In particular, mirroring a theme that has been successful in unraveling fundamental mechanisms (e.g. induction heads) behind transformers (Makkuva et al.
2025 2024 Rajaraman al.) 2024 Nichani et al., [024 Edelman et al. 2024). growing body of research explores Mamba through its in-context learning (ICL) capabilities (Grazzi et al. 2024 Halloran et al. 2024: Akyürek
2024 Park et al. 2024). While these works reveal interesting
insights about Mamba's s ICL abilities vis-a-vis transformers, they are largely empirical in nature, and we currently lack a fundamental theoretical understanding of Mamba and its underlying learning mechanisms. We are thus motivated to ask:
Can we systematically characterize the ICL capabilities of Mamba?
"oorrespnndence rco. bondaschi @epf1 ch.

---

Published as a conference paper at ICLR 2026
0.35
0.9
-laayer Momba
1 1-laeer Mamba
-layeer Transformer
0.30
1-layer Transformer
0.8
2-layer Transformer
2-layer Transformer
Optimal estimator
0.25
0.7
Prodictal probability Pe
0.6
0.20
0.15
0.5
0.10
0.4
0.05
0.3
0.00
0
10
20
30
40
50
60
70
2
8:
Markov order
(a) Next-token probability estimation
(b) J1 distance of model estimation from optimal
Figure Single-layer Mamba learns the optimal Laplacian estimator when trained on random Markov chains, exhibiting ICL. (a) shows the predicted probability distribution
fixed test sequence for models trained on binary first-order Markov sources. (b) quantifies the L1 deviation from the optimal estimator for random sequences and various Markov orders. The error intervals show the standard deviation across 5 runs. Sec. 4.3 and Fig. 10 further discuss Mamba vs. Transformers.
In this paper, we approach this question from the point of view of representation power, and char- acterize Mamba's ICL capabilities Markov processes, building upon the Markov-ICL framework originally introduced for transformers (Edelman et al.) 2024). As opposed to a Mamba vs. Trans- formers comparison, here leverage this framework for a detailed study of Mamba, and uncover an interesting phenomenon: even single-layer Mamba efficiently learns the in-context I Laplacian smoothing estimator, which is both Bayes and minimax optimal, for all Markov orders (Figs. [ and 10a). Towards explaining this, we theoretically characterize the representation capacity of Mamba and demonstrate that the convolution mechanism, together with selectivity and recurrence, plays a fun- damental role in realizing the Laplacian smoothing. Importantly, we showcase that these theoretical insights align strongly with empirical results, even outside the realm of Markovian data. To the best of knowledge, this is the first result of its kind connecting Mamba and optimal statistical estimators.
In summary, we make the following contributions:
Leveraging the Markov-ICL. framework, uncover the surprising fact that even single-layer Mamba learns the optimal in-context estimator for all Markov orders (Fig. [ Intriguingly, convolution plays pivotal role, more so than gating and non-linear activation, in this learning ability (Sec. (3.
Towards explaining this phenomenon, we characterize the representational capacity of single-layer Mamba and show, both theoretically and empirically, how it represents the optimal in-context estimator for any finite-state first-order processes, through an intricate interplay of convolution, selectivity and recurrence. Further, provide fundamental limits for higher-order processes (Sec.
We demonstrate the generality of our findings on non-Markovian data and illustrate the fundamental role of convolution even complex language language-modeling tasks (Sec.) 5
1.1 RELATED WORK
SSMs (Gu et al. 2020 2021) have been recently introduced as an alternative recurrent architecture
aimed at rivaling the well established transformer backbone Vaswani et al. 2017). The model was
originally introduced as a discretized linear dynamical system (Gu et al. 2022 Recent works tried to
re-frame the architecture from a linear recurrent perspective (Orvieto et al. 2023b). However, there
are still many gaps in understanding this family of models (Team et al. 2024), such as questions around expressivity (Orvieto al. 2023a). This is particularly important given the proliferation of inspired architectures that have emerged since its introduction (Qin & Liu 2024 Csordás
al. 2024 Zhu al. 2024 Gu & Dao 2023h De et al. 2024 Beck al. 2024)
2

---

Published as a conference paper at ICLR 2026
To this end, our work squarely focuses understanding the representation power of Mamba, and in particular its ICL capability, which, while extensively studied for transformers (Xie et al. 2021
Hendel et al., 2023; Bai et al.) 2023). remains largely unexplored for SSMs. In this space, recent studies such as Sushmma et al. 2024). have shown that SSMs can perform gradient-based learning for in-context adaptation similar transformers. There is conflicting evidence whether Mamba's s ICL
abilities are better (Grazii et al. 2024) or worse (Halloran et al. 2024 |024]|Akyinn Akkürek et al.) (024) compared transformers. Nonetheless, SSMs have demonstrated promising results in in-context reinforcement learning tasks (Lu et al., 2024). well as in next-state prediction for dynamical models (Joseph) al., (204)), highlighting the potential of SSMs as efficient alternatives transformers for ICL tasks. Motivated by this, as opposed to an architectural comparison Jelassi et al. (2024); Bhattamishra et al. (2024); Merrill et al. 2022)); ); Sarrof et al. 2024), here we solely focus on Mamba's 's ICL capabilities, specifically, through the lens of random Markov processes. This framework has been successfully applicd to transformers (Edelman et al. 2024; Makkuva et al., 2025 2024 Rajaraman et al., [024]] Nichani al. 2024), where it helped unveil fundamental learning mechanisms of transformers such as induction heads. Ours is the first work that employs this framework for Mamba and SSMs.
2 PROBLEM SETUP
We formally define the problem setting and provide necessary background. We the following notation: scalars are denoted by such italic lower case letters as I, y. Euclidean vectors by bold x, y. and matrices by upper case X, etc. 1 refers the all-one vector. For T N. [T] A {1...., T), and for sequence (#4)121, define æt (Ik...., 24). For E R, sigmoid(2) & 1/(1 + e=²), RzLU(z) & max(0, 2). and softplus(z) & log(1 + e*). Unif(S) denotes the uniform distribution over a set S and Dir(3) denotes the Dirichlet distribution with parameter O. DKL (PIIQ) denotes the KL divergence between distributions / and Q.
2.1 INPUT DATA: RANDOM MARKOV CHAINS
To investigate the ICL. capabilities of Mamba, we build upon the Markov-ICL. framework of Edelman (et al. (0024). In particular, we let the input tokens to be stochastic and drawn from random Markov chain order k. That is, the token sequence I = (zt) X on the state space (vocabulary) X follows the transition dynamics:
IP (0++1=-1 = x (N++1 1
almost surely for all f E [2], and the h"horder Markov kernels, IP (I1+1 = 2 k+1 = (+11))
are sampled independently for each tuple (ir-kt1,.. ix) from the Dirichlet prior Dir(3 1). with When 9 1, this corresponds to the uniform distribution on the S-dimensional simplex A where size S = |
The transition matrix p (P:t) E [0, 1 encapsulates the set of all S conditional probabilities of the chain, each row corresponding to one of them. While this transition matrix governs the generation of each token I, for t > k, the first k-tokens &1,..., are drawn i.i.d. from Uni This constitutes the joint law of the random variables (P, 2), termed random Markoy distribution henceforth. More succinctly.
Data generation (Random Markov sequences).
1. Draw P with each row sampled i.i.d. from Dir(3 1).
2. For t 1,...,.. k sample x, ~ (t).
3. For t = k k,..., sample I++1 ~ Ps 1
4. Return the input (x:)
5. Repeat the above steps generate batch (2 )E[B)-
Why Random Markov is good testbed for ICL. As consequence of the generation process, every sequence follows a different Markov distribution. Therefore, at inference, a model trained on this random Markovian data has to estimate the next-token distribution in-context for every test sequence. Hence, this data class serves as a good sandbox to gauge the ICL capabilities of Mamba, which was also used in similar context for transformers (Nichani al. 2024 Rajaraman al. 2024
3

---

Mamba block is then incorporated into a full-fledged language Published as a conference paper at ICLR 2026
2.2 MAMBA ARCHITECTURE
Selective SSMs such as Mamba and Mamba-2 are a class of sequence-to-sequenc mod- els that are closely related RNNs and classical state space models (Gu & Dao A key feature underpinning these models is the selectivity mechanism, enabling them to selectively choose inputs every timestep, as opposed to linear time-invariant (LTT) systems. While tures the behavior of all selective SSMs, we will specifically focus on the state-of-the- art Mamba-2 model to simplify exposition. By slight abuse of terminology, henceforth we will also refer to this model simply as Mamba. Mathematically speaking, Mamba implements the sequence-to-sequene ping Mamba Raxı + Rdxt, where given sequence of input embeddings a = (xx)?=1 E pdxT of dimension d, outputs the corresponding output beddings = (ot) E RdxT of the same dimension with
2023b).
fo(x))
So(a))
Jo(xF))
x
believe our work cap-
logit
logit;
logitr
Linear
Linear
Linear
1
Pr
MLP
MLP
MLP
map-
Wy
Mamba
Marnuu
Mamba
Embedding
Eimbedding
Einbedding
Mamba(ar). More precisely, fix t E [T]. Then the output On at time is computed o, Mamba(z() using the following recurrence equations (Dao Gu. 2024):
1
1]
Figure 2: Mamba-base language model.
- exp(-a A.) E (0, 1),
He = a, H + 3e b E RedxN
A. A softplus((wa)) xx) 5) E R,
yt = H, E RRdd
3, - ReLUUconv x (Wx x/-w+1) A.s
2. y: @ ReLU(W, x,) E Red
b. - ReLU(conv 8 (Wo x
O, = W, 2, E Rd.
- ReLU(conve (Wc w1-4)))
(Mamba)
(Input selectivity)
where the initial state Hg = 0, W- E Redxd W. E RRdxed a 2 0, wA E Rd,ōE R, Wx E Redxd WB Nxd and Wo N are all learnable parameters, and conv(2) time-wise convolution of window to E N with distinct kernels per dimension. Here e E N is the feature expansion factor, typically 2. Let OMamba denote the set of all these parameters.
Intuition behind Mamba. The underlying intuition behind the update equations in Mambalis simple: given sequence of input embeddings (xx,), we first capture their local temporal information using separate convolutions to compute ã,, b, and [Input selectivity). Equipped with this local memory, we perform linear state update compute the current state H, from the past H,-1. weighed by an input-dependent decay factor a, E (0, 1). and (æt, b.). Subsequently, we compute the state projection ye. modulate it with an selut-selective term yield z,, and finally project it down get the output embedding o,, which is a function of the entire input sequence until then, æf. i.e., o, = Mamba (mr).
Mamba-basel language model. model as follows:
Tt. E 10
Embedding
iP
Manbuu
Ut
MLP
V7c
L
Lineur
logit,
Prediction
Sa(x1),
(2)
where f0 A Fe (I1+1 = | softmax(logit,) E |0, 1 is the probability estimation for the next symbol rz+1 conditioned on the past zί. We omit the layer norm here for simplicity. We compactly denote the set of all model parameters as 0 E DRP. We refer to $ for more details.
4

---

Published as a conference paper at ICLR 2026
2.3 LEARNING TASK: NEXT-TOKEN PREDICTION
With the objective of auto-regressively estimating the next token, we train the model parameters 0 minimize the cross-entropy loss between the next-token predicted probability fo(xf) and the corresponding ground-truth symbol I++1 across all the positions t E [
L(Ø) -
Σ EE[T]
E + log p (xf)]
where (x)) & Po (Iz+1 = - for j E X. and the expectation is both over the transition kernels P and the Markov sequences x (=1)(=1 sampled from P. In practice, it is replaced by empirical average across a finite set of batches, sampled according to the random Markov distribution in Sec. 2.1 For our experiments we use the AdamW optimizer (Kingma Ba 2015).
2.4 OPTIMAL ESTIMATOR: LAPLACIAN SMOOTHINNG
Given the Bayesian prediction loss in Eq. (3. it is natural to ask: what is the optimal I minimizing it? follows from a classical result in statistics (Rissanen (1984). {(A) that this minimum is achieved when the corresponding model prediction matches the (average) ground-truth predictive distribution, =f)]. for all I. Given the joint distribution of the pair (P, Sec. 2.1 where the kernel P ~ Dir(B 1), can be shown ($ A that the conditional expectation above simplifies to the well-known Laplacian smoothing, also known as the add-B
i.e. IPo (Iz+1 = xrf = pin [P (Ia+1
estimator (see e.g. Merhav & Feder (998)):
(kk (Ie+1 j 2 Epls [PP (#++11 -i-1=jlo [ n + 3
(Laplacian smoothing)
2,3
where ny is the number of times token j follows the current -oddr context si-x+1 1+1 in the sequence af, i.e. ny = Hí: (xf=)==)) #x) = (+i- and the frequency of this context, i.e. n = Ki / = x Adjusting these counts by 8 plays the role of additive smoothing, which avoids assigning zero probabilities to unseen events, an idea dating back to Laplace (Laplace 1814). is also known that the add-3 estimator is asymptotically minimax optimal, as T + 0o (Xie & Barron 1997
Hao et al. 2018).
How Laplacian smoothing implies ICL. If Mamba realizes this smoothing estimator, i.e. Ie (2k) automatically implies its ICL abilities: given fresh test sequence at inference, in order to optimally predict the next token, it has process the input tokens in-context compute the relevant counts, as in the Laplacian smoothing] But does Mamba realize this optimal counting estimator in practice?
3 DOES MAMBA LEARN IN-CONTEXT ESTIMATORS?
To investigate the ICL capabilities of Mamba, we consider the problem setup described above and train Mamba and transformer models using AdamW on the next-token prediction loss in Eq. (3) on random Markov chains (we refer to $ For more experimental details). These experiments reveal interesting and rather surprising insights about Mamba:
Mamba learns the optimal Laplacian smoothing estimator on the Markov prediction task, even with a single layer (Fig. lap.
2. Convolution mechanism plays fundamental role in Mamba, more so than gating and non-linear activations, aiding its learning abilities (Fig. Bap.
In the sequel, we expand upon these observations in detail.
I) Mamba learns the Laplacian smoothing. After training, we evaluate Mamba and transformers on the same test sequence fixed beforehand and compare their performance to that of the optimal Laplacian smoothing estimator. Specifically,
compare their next-token prediction probabilities with those of the add-3 estimator, Fig. F Iillstrrates these results for various Markoy orders, which uncovers surprising phenomenon: even a single-layer Mamba sharply matches the optimal
5

---

Published as a conference paper at ICLR 2026
0.40
0.35
Mamba
Order
0.35
MambaZero
0.30
Order 2
0.30
Mamba without convolution
Order
0.25
Order
Test loss
0.20
Test
0.15
0.10
0.05
0.05
0.00
0.00
10
15
20
25
30
35
Iteration (x200))
Convolution window of Mamba
(a) Importance of convolution
(b) Relation between window size and Markov order
Figure 3: (a) illustrates the fundamental role of convolution, without which the model fails to learn the task. In contrast, simplified variant with just the convolution (MambaZero) matches the performance the full model. (b) highlights the relation between the Markov order A: and the window size of Mamba. It is required that w + 1 for the model to learn the order-k prediction task.
estimator on the whole sequence. The same conclusion holds for larger state spaces and deeper models (Fig. [oa), well in wverpparaeetrized settings (Fig. 99. and even when part of the dataset is held out (see & E]. For transformers, we observe that a two-layer model also matches the predictor, albeit less sharply, whereas single layer fails solve the task. This aligns with recent theoretical results (Sanford et al. 20244] Ekbote et al., 2025). that show that two layers are required for transformers implement an induction head (realizing the counting estimator) efficiently, We also observe that linear attention performs similarly to softmax attention in this setting (cf. Fig. 7 Sec. 4. further discusses Mamba vs. Transformers.
2) Convolution is the key. To decipher the key architectural component behind Mamba's 's success Markov prediction task, we do ablation study on its three main features: convolution in Input selectivity (ii) ReL.U non-linearity in Input selectivity and (iii) the gating mechanism in Mamba and MI.P| Amongst them, interestingly. convolution plays fundamental role the model's performance, illustrated in Fig. Here we compare the full Mamba architecture from Sec. 2.2 Mamba with just the convolution in Input selectivity removed, and a simplified Mamba architecture with only convolution (MambaZero in Sec. 4.1p. Further experiments on adding/removing convolution to Mamba and transformers, as well as experiments with varying width, are shown in $ (E.6) while experiments on natural language closeness of each of these models' losses L(O)) to that of the optimal add-3 estimator Lg, i.e. i.e. [T(0)) T.gl. The closer this metric is zero, the better the model's 's performance is. Remarkably, the simplified Mamba with just the convolution succeeds on the Markov prediction task, while the full model without convolution fails, highlighting its fundamental importance. This raises natural question: how does convolution help Mamba to implement the optimal Laplacian estimator?
deferred Sec. 5.2 As metric of comparison,
the
4 HOW MAMBA IMPLEMENTS THE LAPLACIAN ESTIMATOR
Motivated by its success in learning the optimal estimator, here we study how Mamba represents Laplacian smoothing. Specifically, we provide a concrete theoretical construction backed by empirical results, that illustrates the mechanism Mamba uses to implement the estimator in practice.
4.1 MAMBAZERO: SIMPLIFIED MODEL
Building upon the insight that Mamba with just the convolution achieves the same performance as that of the full model (Fig. Ba), we consider its simplified version: MambaZero. MambaZero retains only the essential elements of the full model in Sec. 2.2 the Ehheddddd layer, the convolution

---

Published as a conference paper at ICLR 2026
inside the Mamba block in Input selectivity and the Linear layer.
More formally, it is given by:
zt = ex. E R
(Embedding)
H₂ a H
-r bj @Rex
3y + MambaZero(2)),
(MambaZero)
y: H. c2 E R
(MambaZero)
logit, W TRS
(Linear)
O. = W. Yt E R
[a(11) (logit,/ ||loiit/))))
(Prediction)
where Cr is the token embedding for x X, and the input-selective terms dg, Ht, by and c₁ are computed as in Input selectivity) without ReLU and just the convolution. Here we use the L normalization instead of the softmax in the (heePiccitioon layer to ease theoretical analysis, similar to Nichani et al. (2024); Rajaraman et al. (2024). Let ((ei)iex, OMambaZero We) E IP denote the full set of parameters for appropriate D 221.
4.2 MAIN THEOREM: MAMBA REPRESENTS THE LAPLACIAN ESTIMATOR
We now present our main theorem that MambaZero can represent Laplacian smoothing for any Emite-state first-order Markov process. key defining feature of our constructive proof is that it aligns with the structures empirically learned by the model, shedding light on the fundamental learning mechanisms of Mamba.
S} size |X| S, there a choice of parameters for the canonical MambaZero model, with dimensions N = S. 28, = and convolution window 2. such that its output prediction exactly matches that of the Laplacian estimator, for first-order Markov chains on X. More formally, for any $ > O, there exists a set of parameters O such that, for all sequences (14))21 and all
Theorem For state space X'
DKL. (P 1 P (- (.121) = 0.
Remark. The KL divergence above is precisely the penalty paid in the cross-entropy loss in Eq. (3) at time when using the predictor Pe instead of the optimal P(". In other words, the result implies that the loss of MambaZero can be made exactly equal to the optimal.
4.2.1 KEY MECHANISM AND PROOF SKETCH
Main idea. build our intuition towards how MambaZero can realize the add-;9 counting estimator for first-order Markov sequences, let's focus on the core MambaZero block. The key observation here is the following: if the state H;-1 can capture all the transition counts + till 2 the new state Ht, can be updated to account for the current transition x;-1 Iz top of the existing counts, by suitable choice of a, Tt. and b, Then the relevant count information corresponding to the current prefix 2; could be read off from the state projection y: H, and be modified to account for 3-smoothing via the (inear and Predicctinn] layers. Buttressing this idea are two key empirical facts, which in fact hold for any A: 2 underpinning
construction:
(i) State-to-state transition factor at RS for all 2 1. We empirically observe that when the MembaZero model is trained on random first-order Markov data, convergence we have a, ss 1 for all é 2 1 (Fig. 4). Since modulates how much past information flows into the present, a, = is required for the state H, to store all previous transition counts. Note that this can easily achieved by setting either a or A to be zero in Input selectivity which we empirically observe well.
(ii) Convolution window 2 l: 1. Recalling that à the Markov order, we empirically observe that the window size = k + 1 is sufficient for the full Mamba to learn the Laplacian smoothing on order Markov chains (Fig. 3b). To understand why, note that in the MambaZero architecture above, apart from the Maabaa/ero//blon
all remaining equations operate on the current token at time In the MaabaaZro block, the dependency of the output Yt on the previous tokens due to that of the state HI, on (æt, b,) in the update equation, and of in the state projection. Since (Ø,, b., c,) depend on the past through the convolutions, a window of size k + enables them to keep track of the current token well its length-A prefix, which is necessary to compute the counts needed in Laplacian smoothing On the other hand, if wxx Wp & k. then one can find confusable sequences, i.e. sequences that share the number of occurrences of all length-k prefixes, but whose counts

---

Published as a conference paper at ICLR 2026
of the tokens following each prefix is different, resulting in the model's estimate to deviate from that of the optimal add-3. We refer to for more details. While having all the window sizes wx, w, w k + 1 is sufficient, it can be further strengthened to w, & ($
We now detail our construction for the first-order case, capitalizing on these insights.
Construction. Let us fix 1 + = 2. Then, T. and b, only depend on the current token 2; and the previous one 24-1, while C2 only depends on Ht. Thus, Ty and b, can only take 5 possible values depending the last transition in the sequence, whereas C₁ only S. To ease the notation, denote these values by x(5) b(i)), and c respectively, for i E X. Additionally, at = 1. these terms depend only on the current symbol, taking two additional values each, denoted by お(0), b()) Let Raj denote the number of transitions + j in the input sequence æ'. Then, unfolding the state update recursion in MambaZero we get that the output of the MambaZero block is
will
0 = W. Tuby ct + Σ nej Wo x (( C..
While the output in Eq. (4) depends on all the transition counts, in view of Laplacian smoothing) we ideally want only those counts pertaining to relevant transitions, i.e. if x4 = 0, the counts Roj- for j E X', and similarly for other values of zz. To this end, we empirically observe that at convergence, the model's parameters are such that ((5)) C2 ~ O whether i & 24 (cf. Fig.5). Due to this property, only the counts that are involved in the computation of the Laplacian estimator for the current token Te appear in the output og. Stitching these facts, the final logits in the Linear| layer depend on the first and current token via
logit, We + WoW, Zoby Ce + Σ no, WeW ( ((e)) Ct.
(5)
The final step is to then show that for properly chosen parameters, one can make the two vectors associated with the counts to be orthogonal, and the other vectors, independent of the counts, to sum up to the vector 31 (which we also empirically verified, cf. Fig. 5). Subsequently, the L normalization in inPrrdictionn layer will give next-token probability estimate, matching that of the add-B estimator. We defer the full proof and additional details to $[C
Dimension reduction for binary state space. Interestingly, for the binary case X {0, it is possible to further reduce the hidden dimension = 25 in Thm. to d = S = 2 by leveraging the relationship between the transition counts. The key theoretical insight is that the transition counts in binary sequences are strongly correlated. Specifically, no1 and n0 are at most one apart: every time a transition > 1 occurs, either the sequence is followed by only I's until the end, or subsequent transition 1 + 0 also occurs. Therefore, the dependency of the output on noi is in fact a dependency on 1110- One can leverage this property to help MambaZero realize Laplacian smoothing with just two-dimensional embeddings, with arbitrarily small error. We refer to Thm. in C.3 for full details.
4.3 LOWER BOUND: FUNDAMENTAL LIMIT ON THE REPRESENTATION POWER OF MAMBA
We now provide a fundamental limit on the representation power of recurrent architectures like Mamba, in the form of a lower bound on the hidden dimension that is required to represent the optimal estimator. In particular, our result establishes that with finite bit precision, irrespective of depth, for any recurrent architecture implement the Laplacian estimator, the hidden dimension has to at least scale as (222²).
Theorem 2. Consider a recurrent model of the form
H. h4 (He
yt Po x1) = g.
with transformations (he, 9t), where H. E Rdd and the model has a bit precision of Suppose that
the 1hh -oddr Markov kernel P sampled from the Dirichlet prior with 3 1, Dir(1 1).
Suppose also that the recurrent architecture satisfies the following point-wise guarantee: for any sufficiently large t, almost surely over P and a ~ P.
IPo (-121) - p(0) xt) Se,
(6)
8

---

Published as a conference paper at ICLR 2026
where pp(+( a4) is the Laplacian estimator for B 1. Then, the recurrent architecture must satisfy d p.p 2 2 - log(1/E)
We defer the full proof and additional details to App. D
Depth. We note that Thm. 2 does not assume depth one, and holds for recurrent models of any depth.
Mamba vs. Transformers. As Thm. demonstrates, to capture a -oddr Markov process, Mamba requires the hidden dimension to scale exponentially in k; whereas the best known result for transformers needs a three layer model with the hidden dimension growing linearly in & (Rajaraman et al. 2024). On the other hand, for first-order sources we empirically observe from Fig. that 1-layer Mamba tracks the optimal estimator more sharply than a transformer (see E for additional comparative results). While these comparisons are meant to provide more detailed context for Mamba, we would like to emphasize that the main focus of our paper is not a comparative study but rather fandamental understanding of Mamba's ICL abilities.
Higher orders and learning dynamics. While Thm. demonstrates that Mamba can represent the optimal estimator for finite-state first-order processes, our empirical results in Fig. Tb strongly suggest that a similar conclusion holds for higher-order sources. In a similar vein, analyzing Mamba's learning dynamics in its convergence this smoothing estimator an interesting topic of future research, but outside the scope of this paper, whose focus is on representation power.
5 BEYOND MARKOV
5.1 SWITCHING MARKOV MODEL
A key component of Mamba enabling selectivity is the state-transition factor ae, that controls the flow of information from the past state H to the current H₃: if a, 1, the past information is fully utilized in computing the current state, and hence the output, whereas a₂ = O completely ignores the past. In the Markovian setting considered so far, the role of a, has largely been dormant: 0 Rs 1 for all t 2 1, as the optimal Laplacian predictor requires counts of all transitions, demanding the use of full past (Sec. 4.2)p. p. To better highlight this selectivity mechanism, we consider a Markovno process, where the role of a, becomes fundamental.
Specifically, we focus on the switching Markov prucess, where we add switch token to the binary alphabet, i.e. we consider X' = {0, I,.). The key difference here compared to the random Markov generation in Sec. 2. is that until we hit switch token, we follow the same binary Markov sequence generation the former, but once the switch state is reached,
sample new Markov kernel and then generate a new Markov sequence. The switch tokens are sampled according to a parallel ii.d. Bernoulli process with probability Pswitch (0.01 in our experiments). The sampling process is described in detail in $ [E.7 With this data model, the optimal prediction strategy is to use the add-3 estimator in between two switch tokens, and reset the transition counts every time a switch occurs. We provide empirical evidence in & E.7 Indeed, Fig. 13 illustrates that Mamba implements precisely this strategy, closely tracking the switching events via the transition factor ay: it sets aj be zero whenever I22 = S and to one otherwise.
5.2 NATURAL LANGUAGE MODELING
To test the generality of our finding that convolution plays a key role Markovian data (Fig. 3}. we conduct experiments on language modeling using the WikiText- 103 dataset. Details on the experimental setup can be found in $ F By adding or re- moving convolution in both these models, we obtain the results in Table 1 The re- sults illustrate that convolution enhances the performance of the two architectures, in particular for Mamba (11% vs. 2%).
Table 1: Perplexity results on the WikiText- 103 dataset.
Model
Params.
Perplexity
Mamba-2 (w/o conv)
14.53 M
30.68
Mamba-2 (w/ conv)
14.54 M
27.55
Transformer (w/o conv)
14.46 M
29.28
Transformer (w/ conv)
14.46 M
28.67

---

Published as a conference paper at ICLR 2026
highlighting its saliency. Further ablation studies on this task show that together with convolution, gating also plays a central role (17% change, cf. $ $EE.8 Table 4). Furthermore, additional experiments with deeper models show that the relative importance of convolution seems to decrease as the number of layers increases (ef. $ | TTble 5). This may be due to the fact that the role of convolution is taken over by other Mamba layers, which are known successfully approximate convolution (Wang) & Xue 2023).
6 CONCLUSION
Structured state space sequence models (SSMs) and Selective SSMs such as Mamba have shown re- markable inference speed-ups over transformers while achieving comparable or superior performance on complex language modeling tasks. In this paper, we studied in-context learning (ICL) capabilities of Mamba on random Markov chains and show that, unlike transformers, even single-layer Mamba efficiently learns the in-context Laplacian smoothing estimator. To explain this, we theoretically and empirically characterized the representation capacity of Mamba, which revealed the fundamental role of convolution, together with selectivity and recurrence, in enabling it. We further provided additional empirical results on non-Markovian data, showing the generality of our insights. Extending results to deeper Mamba models, as well as investigating Mamba's 's learning dynamics, are some interesting future directions.
ACKNOWLEDGMENTS
The work this manuscript was partially supported by the Swiss National Science Foundation under Grant 200364.
REFERENCES
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Arhitec- tures and algorithms. arXiv preprint arXiv.2401.12973, 2024.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Workshop on Efficient Syslems for Foundation Models ICML2023, 2023.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint arXiv:2405.04511 2024.
Satwik Bhattamishra, Michael Hahn, Phil Blunsom, and Varun Kanade. Separations in the represen- tational capabilities of transformers and recurrent architectures. Advances in Neural Information Processing Systems, 37:36002-36045, 2024.
Marco Bondaschi and Michael Gastpar. Batch universal prediction. In 2024 IEEE International Symposium on Information Theory (ISIT), pp. 3552-3557, 2024. doi: 10.1199/SIT77764220244 10619270.
Marco Bondaschi and Michael Gastpar. Alpha-NML universal predictors. IEEE Transactions on Information Theory, 111)):177-1118, 2025. doi: 10.11099TT.2024.3322211
N. Desa-Bianchi and G. Lugosi. Prediction, Learning, and Games. Cambridge University Press, 2006.
Nicola Muca Cirone, Antonio Orvieto, Benjamin Walker, Cristopher Salvi, and Terry Lyons. Theo- retical foundations of deep selective state-space models, 2025. URL ptttps://arxiv.org
abs/2402. 19447 1
Róbert Csordás, Kazuki Iric, Jürgen Schmidhuber, Christopher Potts, and Christopher D Manning. Moeut: Mixture-of-experts universal transformers. arXiv preprint arXi::2405.16033 2024.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. arXiv preprint arXiv:2405.2 2024.
10

---

Published as a conference paper at ICLR 2026
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Pristian-Murary Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models. arXiv preprint arXiv:2402. 19427, 2024.
Jacob Devlin, Ming- Weig-Weiile Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding, 2018. URL https://arxiv.org/
abs abs/1810. 04805
Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The Evolution Statistical Induction Heads: In-Context Learning Markov Chains, 2024.
Chanakya Ekbote, Marco Bondaschi, Nived Rajaraman, Jason D. Lee, Michael Gastpar, Ashok Vard- han Makkuva, and Paul Pu Liang. What cannot, two Two-layer transformers provably represent induction heads on any-order markov chains, 2025. URL
abs/2508 07208
Riccardo Grazzi, Julien Siems, Simon Schrodi, Thomas Brox, and Frank Hutter. Is mamba capable of in-context learning? arXiv preprint arXiv:2402.03170, 2024.
Albert Gu and Tri Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv preprint arXiv: 23/2.00752, 2023a.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint ar arXiv:2312.00752, 2023b.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Advances in Neural Information Processing Systems, volume 33, pp. 1474-1487, 2020.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-tim models with linear state space layers. In Advances in Neural Information Processing Systems, volume 34, pp. 572-585, 2021.
John Halloran, Manbir Gulati, and Paul F Roysdon. Mamba state-space models can be strong downstream learners. arXiv preprint arXiv:2406. 00209, 2024.
Yi Hao, Alon Orlitsky, and Venkatadheerai Pichapati. On learning markov chains. In Advances in Neural Information Processing Systems, volume 31, pp. 646-655, 2018.
Roee Hendel, Mor Geva, and Amir Globerson. In-context learning creates task vectors. arXiv preprint arXiv:2310.15916. 2023.
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Trans- formers are better than state space models at copying. arXiv preprint arXiv:2402.01032, 2024.
Federico Arangath Joseph, Kilian Konstantin Haefeli, Noah Liniger, and Caglar Gulcehre. Hippo- prophecy: State-space models can provably learn dynamical systems in context. arXiv preprint arXiv:2407.09375. 2024.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Pierre Simon Laplace. Essai philosophique sur les probabilités. Courcier, Paris, France, 1814. Reprinted by Cambridge University Press, 2009. In the reprint, the estimator appears on page 23.
Chris Lu. Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Hyeji Kim, Michael Gastpar, and Chanakya Ekbote. Local to Global: Learning Dynamics and Effect of Initialization for Transformers. In The eiiht Annual Conference on Neural Information Processing Systems, 2024.
11

---

Published as a conference paper at ICLR 2026
Ashok Vardhan Makkuva, Marco Bondaschi, Alliot Nagle, Adway Girish, Hyeji Kim, Martin Jaggi, and Michael Gastpar. Attention with Markov: curious case of single-layer transformers. The Thirteenth International Conference Learning Representations, 2025.
N. Merhav and M. Feder. Universal prediction. IEEE Transactions on Information Theory, 44(6): 2124-2147, 1998. doi: 00.1109/18.205544
William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. arXiv preprint arXivv2404.08819, 2024.
Eshaan Nichani, Alex Damian, and Jason D Lee. How Transformers Learn Causal Structure with Gradient Descent. arXiv preprint arXiv:2402. 14735, 2024.
Antonio Orvieto, Soham De, Caglar Gulcehre, Razvan Pascanu, and Samuel L Smith. On the univer- sality of linear recurrences followed by nonlinear projections. arXiv preprint vrXiv:2307.11885 2023a.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. arXiv preprint wrXiv:2303.06349 2023b..
Matteo Pagliardini. GPT-2 modular codebase implementation.
11m baselines Accessed: Jan. 2025.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jacwoong Cho, Samet Oymak, Kang- wook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? comparative study on in-context learning tasks. In Proceedings of the 41 International Conference on Machine Learning, 2024.
Jiahao Qin and Feng Liu, Mamba-spike: Enhancing the mamba architecture with a spiking front-end for efficient temporal data processing. arXiv preprint arXiv:2408. 2024.
Alec Radford and Karthik Narasimhan. Improving language understanding by generative pre-training. 2018.
semant http://ph...aoaaaooooooooorrrrpr
Nived Rajaraman, Marco Bondaschi, Ashok Vardhan Makkuva, Kannan Ramchandran, and Michael Gastpar. Transformers on Markov data: Constant depth suffices. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
Rissanen. Universal coding, information, prediction, and estimation. IEEE Transactions on 1056936.
Information Theory, 30(4))629-636, 1984. doi:
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. One-layer transformers fail solve the induction heads task, 2024. URL httppp::/aaxxvv.ooo wttp:///wwwwwr/////.....oo 23/5//.abb/220..1332 14332
Yash Sarrof. Yana Veitsman, and Michael Hahn. The Expressive Capacity of State Space Models: A Formal Language Perspective. arXiv preprint arXiv: 2405.17394. 17394, 2024.
Neeraj Mohan Sushma, Yudou Tian, Harshvardhan Mestha, Nicolo Colombo, David Kappel, and Anand Subramoney. State-space models can learn in-context by gradient descent. arXiv preprint 2tXiv22410.11687. 2024.
Jamba Team, Barak Lenz, Alan Arazi, Amir Bergman, Avshalom Manevich, Barak Peleg, Ben Aviram, Chen Almagor, Clara Fridman, Dan Padnos, et al. Jamba-1.5: Hybrid transformer-mant models at scale. arXiv preprint arXiv:2408. 12570, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, pp. 5998-6008, 2017.
Shida Wang and Beichen Xue. State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. Advances in Neural Information Processing Systems, Systems,36:74022 36:77021--44007333 2023.
12

---

Published as a conference paper at ICLR 2026
Qun Xie and A.R. Barron. Minimax redundancy for the class of memoryless sources. 1EEE Transactions on Information Theory, 43(2):646 657, 1997.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference. arXiv preprint arxi::2111.00000 2021.
I. Zhu, Liao, Q Zhang, X Wang. W Liu, and X Wang. Vision Mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint vxXiv2241.094111 2024.
13

---

Published as a conference paper at ICLR 2026
A PRELIMINARIES ON LAPLACIAN SMOOTHING
Laplacian smoothing is a mature and well understood topic. An account can be found, c.g., in Merhav & Feder (1 [998); Desa-Bianchi & Lugosi 2006), with some recent updates in Bondaschi & Gastpar (2024] 2025). For the sake of completeness, we provide a brief outline of how it applies to our context. For k-th order Markov data, at every time instant f. the Laplacian add-ß estimator applied to the subsequence of tokens with the same context i E X as the current one is the predictor that minimizes the Bayesian cross-entropy loss in Eq. (3), when the Markov kernel sampled according to the product Dirichlet distribution Dir(3 1). We first give an intuition of why this is the case, and we provide full proof at the end of the section. We consider the binary case X' {00 1]. but the results can be extended to arbitrary finite alphabets.
Consider a given sequence (a (== For every length-k context i X let (z1) byf be the subse- quence of tokens preceded by 2t. Note that, since each sequence (ay) is generated by k-th order Markov chain, all the tokens in the sequence with the same length-& prefix share the same conditional probability distribution. Furthermore, since each of the conditional distributions of the chain is randomly chosen independently from the others, the subsequence (æz)læy is a sufficient statistic to
estimate the probability distribution of all the tokens with the same prefix if. Therefore, the optimal prediction for a sequence (xz)²= is given by employing the optimal predictor for each i.id. subse- quence (dz))ss= for every af E X Since each conditional distribution is sampled from Dirichlet distribution with parameter 3, it is well known that the optimal predictor for such subsequences is the add-constant estimator, with constant equal to /3. More specifically, if xr x = it, then the optimal estimation for 2z is
Tp! (I141 = j
n2 +
72 n 233
where ny is the number of times token j appears in the subsequence aili = (xx x -----t) and n is the length of the subsequence.
We now provide a formal proof of this fact. Theorem 3. Consider the class of all k-th order Markov kernels P (Psş) (Ptt)ifeee where each P.ę-
P(I ( is a probability distribution on X = 10. Let each P. be sampled i.i.d. from Dir(B 1), and let at ~ (n* and x141 (1 ~ Pai Then, the predictor (() = P(x141 = j at).
for j E {0, 1}, that minimizes the loss
L
EpE. St Ap [10+1 log (xf + -1 (1-1)+))- log (x1)]
(8)
T [E[T]
is the add-B estimator in Eq.
7 ie. the minimizer flo = -P(/²) (IL+1 = i x1), for all t 12 k.
Proof. First note that
L T Σ E EpE, Ep [re41 log (() + (1 2(41) log (( (z1)
J T #e [1e41 log (z) + +11 - -1411) log (x1))
T
Σ
R E Be (1441) log f (x (1 Bryzlle (12+1)) log f (x1)).
-
Let us define the distribution f (xf) e Enssi [1] and p. (af) & 1 12 (af). Then, we can rewrite the loss as
L= T Σ
Eeg ((9 log ( - 11 (xf) log f (x
For every t E [T] and every of E X", the term inside the expectation is minimized by picking
(() = f(()) (at). In fact, note that it can be rewritten as
-1!2 log (x _!")(()) (af) log (0) (x1)
14

---

Published as a conference paper at ICLR 2026
12 (x1) + (xt) log
(xf) (x1)
= (of) log (x))
1 (x1) log ( fz (z+) log ( (x4)
(x)
DRa. (L-(EIII((E) H(L(x)))
which is minimized when DRL (f. (z((a ie., when f(af) = fo (21). We will now show that f- (x11 is precisely the add-3 estimator. Consider any context if and any sequence of such that
x/-R+1 if. Let also p - (1) P(1 if). Then,
f (xf) - Emc lnf [Te+1]
= Eppp |ot E Pe [4441]
= -Bra Int [P)
= -Bpgletl (1)
where the last equation we used the fact that, when ot ~ Unif(x**), the subsequence æilış is sufficient statistic for Pq. Hence,
(x+) = Bry 1 [2+(1]
h/' p0-1(1 (011-p)8-300 100-11nn111111 p)) 1000
pp
So (1-) 1)
gmi (1 - n0 dg
1 p) - dp
Jo 1mm m
[(m + 3+ 1 + +3
I(n 25)
1'n +(=+28- 28 + 1) 1'(n + (no + 3)
n1 +
n + 2
where we used the fact that P. ~ Dir(B 1), that Jo q01-(( - 2)00 = (2)))T + =0). and that [[2 1) 21)
The proof above is for o ~ ) However, note that the same proof would also work for a distributed according to any distribution that independent of the Markov kernel instead the distribution depends on P (e.g., the stationary distribution of the Markov chain), then the proof would fail the step where wf lep is a sufficient statistic for Pa--
Remark.
Remark. It is important note that, to be able to implement such predictor requires in-context capabilities: at inference, in order to optimally predict the next token, the model must be able to look
into the previous tokens test sequence, and count the tokens with the correct prefix.
15

---

Published as a conference paper at ICLR 2026
B MAMBA-BASED LANGUAGE MODELING ARCHITECTURE
Mamba block can be incorporated into a full-fledged language model as follows: let a = (21, 22, ................................ ,IT) X be an input token-sequence over the alphabet X; here X {0, 1] as explained in Sec. 2.1 Then, at every & E [T', the output of the language model 0 is given by the following sequence of equations (Dao Gu 2024
3t = Es, IR
(Embedding)
U + Mamba(æ1) E Rd
(Mamba)
+ (W))) @ Wou,) e R",
(MLP)
logit, W E ER:
(Linear)
[o(a)) - Po (Ie+1 = x4) soltma E [0, 1]
(Prediction)
where the parameters E Rd Wi E lixd W2 E Rdx4d and W₁ E Sxd are learnable, and so(xf) the probability law for the next symbol It-1 conditioned on the past x1. We omit the layer norm here for simplicity. We compactly denote the set of all model parameters as 0, i.e. (ee) Mamba. 21223, Wi) E R..
16

---

Published as a conference paper at ICLR 2026
C PRELIMINARIES AND PROOF OF THM. AND THM. 4
C.I EMPIRICAL INSIGHTS
Here expand upon empirical observations in 4.2.1 which form the basis of our proof.
State-to-state transition factor at =3 for all t 2 We empirical evidence supporting this
observation in Fig.
1.0
1.000
0.906
0.986
0.6
0.984
0.982
0.4
0.980
0.2
0.0
100
200
300
400
500
Position
Figure 4: Value of az across positions at convergence.
Convolution window 2 A: + 1. Recalling that & is the Markov order, we empirically observe that the window that w = k + sufficient for the full Mamba learn the Laplacian smoothing on hth Markov chains. To understand why, note that in the MambaZero architecture above, apart from the MambaZZeroblood the MambaZero block, same the Mamba block except ReLU, the dependency of the output yı on the previous tokens is due to that of the state H, on (zt, by) in the update equation, and c2 in the state projection. Since (z., b,, c,) depend on the past through the convolutions, window of size k 4- enables them to keep track of the current token as well as its length-k prefix, which is necessary compute the counts needed in Laplacian smoothing] On the other hand, if S k, then can find confusable sequences, i.c. sequences that share the same number of occurrences of all length-k prefixes, but whose counts of the tokens following each prefix is differrntt
all remaining equations operate on the current token at time t. In
For such sequences, the state H, is the same, and so are the predicted probabilities by the Mamba model: however, the optimal estimator, depending on the transition counts, would give very different probability estimates, allowing Mamba's 's prediction loss to deviate from that of the optimal. For example, consider
1. If u: 1, then (i:,, b,, C,) depend only on the current token F,. Then, consider the two sequences x = (0, 0, 1, 1) and f = (0, 0, 0, 1, At time / 6. these two sequences would give the same state H, and the same output yt. since they share the same number of tokens 0 and 1. Therefore, the estimated probability given by the model would be the same in both cases. However, the optimal add add-constant estimator (with 8 = 1) would estimate the probability of It-1 1 to be 1/4 for æ, and 3/4 for T.
Further, it is sufficient that the convolution for e, has window we: = k. That is, the convolution conve involved in the computation of C₁ have window size equal to the Markov order & (i.c., less than conv and conv л) without affecting the model's capability of learning the task (or, equivalently, be taken be zero). Intuitively, this is because the
the left-most kernel coefficients of conve role of C in the state projection is to select the correct transition counts for the computation of the estimator, distilled into yt. In order to do so, it sufficient to know the length-k context of the current symbol which can be encoded by convolution with window size k.
Orthogonal count-dependent vectors. The inner products ( (*) corresponding to * l go to at convergence: only the correct counts are kept in the final logit.
17

---

Published as a conference paper at ICLR 2026
Convergence to the optimal / = 1. The count-independent part of the final logit converges to 3 1, corresponding to the optimal Laplacian estimator.
1.75
1.25
product
1,00
0.75
0.50
0.25
0.25
0
25
50
73
100
125
150
175
200
Iterations
Itorations (x10)
(a) Counts-rclated inner products
(b) Counts-independent vector
Figure 5: (a) Counts-related inner products across interations. Only the correct counts corresponding to ( (k) for i have a non-zero inner product at convergence. (b) Binary coordinates of the counts-independent vector across iterations. Both coordinates converge to the optimal / = 1.
C.2 PROOF OF THM.
Let be the constant of the considered add-constant estimator. Let us fix 0 and A, so that a₃ = 1, for all t 2 1. This can be done by picking, e.g., Wa = 0 and 8 such that softplus(O) 1. Note that the application of convolution to given sequence of vectors z can be rewritten as linear matrix-form operation. For example, for conv one has that
x D 2t-1
D
(9)
where D and D are diagonal matrices. The same holds for conv n and conve, with correspond- ing diagonal matrices 1 (D and D
Let us take the embedding vectors e,, E X. to be the one-hot encoding vectors for the alphabet X. interleaved with zeros, i.c., let e; be such that 64 1 and otherwise. Furthermore, take Wx to be the 25 x 25 matrix such that ji = for = 2 and j = 2h- 1, for < l < S, and 0 otherwise. (The role of Wx shift each coordinate of the embedding vectors by one.) Take now conv to be such that its output is simply equal to the current vector, i.c.. take D = 0 and D I25x25- Take also W Isxs and conv n so that the output is equal to the second-to-le lastt vector, i.e., take (( Isxs and D 0. Finally, take Wc Isxs and conve such that (1()) and ((1) Tsxs-
The final logit vector is in general equal to logit, = W,,, + W.W, a ((=)) (())T c +
Σ naj W.W, (00
(10)
Using the matrices chosen above, we can simplify the formula as follows. Firstly, note that ((m) = 0. as the b vectors only depend on the second-to-lase vectors. Furthermore, since the embedding vectors are orthogonal to each other, we have that B(6)) is, only the correct counts are kept in the logit computation.) With these simplifications, the logit formula becomes
whether i 2z, and 0 otherwise. (That
logit, = Wex, + Σ Du.. W,W, 50)
(11)
Finally, take W, = 12sx2s and take We such that We(1) 31 for all odd i, W2()i, for 1 < i < S, and 0 otherwise. With this choice, we get, for all >
logit, 31 + Σ n
(12)
18

---

Published as a conference paper at ICLR 2026
where e; is the one-hot vector for symbol
X. After the normalization, we finally get
Maj + 3
fo fo(x)); = Σx 7i + S3
(13)
my for i E X. This is precisely the required add-ß Laplacian estimator.
C.3 THIM. 4 DIMENSIONALITY REDUCTION FOR THE BINARY CASE
Theorem 4. For the canonical MambaZero model with dimensions d = 2, and convolution window w = 2, there is a choice of parameters such that the model prediction is arbitrarily close to the Laplacian estimator for random first-order Markov chains. More formally, for any 3 > 0 and e (0, 1), there exists a set of parameters such that, for all sequences (2t)121 and all 2 L,
DEL ( (-12(11))1 a (- (-121) S.
Proof. Fix e > and let B > 0 be the constant of the considered add-constant estimator. Let us fix = 0 and Å, 1, that a, 1, for all 2 This can be done by picking, e.g., wa and such that softplus(6) = 1. Let us compactly denote the convolution kernels as
[100
0011
Yo0 701 n10
conv =
conv =
(14)
0:10
011
711
where each row corresponds to the kernel weights applied time-wise each coordinate of the input sequence (#+))21. Since the window for conve is we 1, we can simply assume w.l.o.g. that C. Wa
Let us denote the embedding vectors to be (Con, co1 T and (e10, e11) and assume that vectors are not collinear. Take also Wx = WB such that
0
Wx eo
Wx
(15)
and take We such that
c0 )
0
Wc eg
Wee
(16)
c
Let us also take the kernels of convx and conv to be the same across coordinates, i.e.,
(xo on
Tro 71
conv =
conv =
(17)
70
71
such that the following conditions are satisfied:
oo + 0171
0o + 0/10 > on * an 0o'1 anntann to
(18)
Be
Note that, with such a choice of parameters, we have
X
01
(11)
B(0)
71
B(1)
(19)
0
71
(00)
cn
(1))
(20)
(00)
[00 + a1
(1)
(10)
a1
(11)
0
000 + a1
(21)
B(00)
(10 + 71
0(1)
Yo
B(10)
71
B(1)
Yo + 71 )
(22)
m
70
:
19

---

Published as a conference paper at ICLR 2026
(We replaced the vector notation of Sec. with matrix notation, so that ( has be intended as x00), and so on.) Take also W. = Wt T. With this choice of parameters, the final logit vector becomes, using Eq. (5)
logit, = We 2t + Wo ((x)) B(x)) c(+)) + 1 (m1 Wo Wox(way B(ar)))((T)))T)) c(
+ Rasa aa ((rv)) Blu Ce + nza, Wo BleasT + x((()))) B(Rese) C₁
Blu
(23)
+ [mm=11 = =(s0 + coa171)
0 (coa0on) + n00 (100 tan) + m)o
601
(107) + (
[1000] + amm)e
+ no1
(24)
= (000 + coam7))
-(c00
0 + noo (craoo
(aoo)t + v0
+ 1 (m1=1)
601
0
+ no1
(ao7 + ta10)co
(25)
if [r 0, and
C10 logit, = (en1 + feamm) [m)==)) (aau))) +
((007 + an))e (0070 + ( tamea
nu
(00 an)( an) + m n)a
(26)
= =( =(n e10
(c10071 ) n10 (1007) + an (anta)) (m1=0)
en
0
nn
(aom + airo)es
(27)
if x. 1. Take now
6oo c11 (ao71 + a171))aa a1
(28)
601 = 110 = (ao7) + 0170) Ao7/
(29)
With this choice of parameters, after the layer normalization, the final output probability vector
n00 + 8
Ru1 + (z1 0} Be
fo(x1) =
(30)
1100 + no1 + 23 + 1 (x1=0} Be noo + n01 + 2,3 + (m1=0) Be
if r 0, and
nuu + 1 1 (m 1) Be
n1 +
T
fo(a))=
(31)
1110 + 711 + 23 + 1 [2,=1) Be n10 + n11 + 23 + (m)=1) Be
if 2, = 1. Note that the resulting predicted probabilities exactly match the add-3 estimator when 21 * xr, but they are slightly different when # = I, due to the additional ,e factor. We now show that, when the additional factor is present, the two predictors nevertheless differ by at most in KL distance. We show it for the case 21 = It = 0, the other case follows in the same way. In fact, note that
Rot + + Be
no1 +3
+ noi+D
(32)
1000 + no1 + 2, + Be
noop + A41 + 2/3
1 +
9e
Roo +oo
+28
Now, since
Be
1 S 1 1+
11+c
(33)
Rot + B
20

---

Published as a conference paper at ICLR 2026
and
Be
1 S 1+
S 1 +
(34)
n00 Ro1 + 2
we have that
R01 + :
Ro1 + 3 3E
(1+
(35)
1000 + 111 + +33
1100 + no1 2B + 3e
but we also have
100o + 110 + +23 3
Be
$ 11+
S 11+., E,
(36)
Roo0 + 141 + 2/
Roo + nu1 + 2
so that
Pl (5++1 x1)
DKL (PS) (-124) IPo (. (-1x4) = =P/") (x141 = 01 log
Po (2141 21)
pll) (I++1 = 1 x+1
+ 1 (Je+1 1 2 log
Pp (r++1 = x
(37)
Rog
Rou + 3
log
+ +28
7100 + n01 2/9
+-8 YABC +-00 +234+9
f01 + 3
log
+22
1100 + n00 + 2,3
Roi +.3.3e
Aoo 1oo +
(38)
Rou +33
Mu1 +
log(1 + E) +
log (1 +
7100 + no1 + 2,
7400 + no1 + 2,3
(39)
< log(1 + +e)
(40)
Se
(41)
concluding the proof.
21

---

Published as a conference paper at ICLR 2026
D PROOF OF THM. 2
Consider a recurrent model of the form If, h(H 24) and g(He) for each 2 where H2 E Rd and the model has bit precision of p. In this proof, we will assume that the state space of the underlying Markov chain is f0, 1}. By the recurrent architecture, the predicted distribution over the next token T4+++1 is of the form,
V Fe (I1+R+1= 2 ) = g(2, Hetk).
(42)
Recall that the add-1 estimator is defined as,
n 1t. 1+h k) +1 1+1
(43)
8+R n m +
where n(=f) Σi I(x 24) indicates the number of times 2f appears in the sequence. This is the optimal estimator for sequences drawn from the product-Dirichlet prior: for every ih. P(. aft ~ Dir(1 1), which is the distribution we will assume for this proof. Fixing æf, we can write the add-1 estimator more explicitly as a function of x/+1 as,
ni 2) +1+1
p (5e+H+1 = 11, 2111 = =zf)
(44)
n==)) +
Now, fixing r1. correctness of the recurrent model means that, almost surely over P drawn from the prior, and z ~ P and 2 P(-[z)),
g( 0, H244) E IP(() (I1+R+1 0 wf # =1) + [-s,e]
(45)
where Hetk is function of a1 and 24. Ast oo, under the randomness of the draw of æ P. by the strong law large numbers RHS converges almost surely to the conditional distribution under P. almost surely over the choice of P from the product-Dirichlet prior. Here drawn from the product-Dirichlet prior, P(=|=*) > 0 almost surely, and so the resulting distributions exponentially mixing and ergodic. Namely, for each 2 E (0, 1 almost surely P drawn from the product-Dirichle prior,
the fact that for P
p (1248+1 = 1000/,00f0 +f, c+h =)) P(O)= > > >)))=
Pr sup t=000
(46)
for any > 0. Therefore, a necessary condition to satisfy Equation 45 is, for each 2t E x
g( = 0, H++x) E P(O + I- (t), E + (t)]
(47)
for some (t),)) which is a function of P satisfying lim 7P (t) = 0 almost surely over P drawn from the prior; note that H is implicitly a function of 2 and 2/. Divide the interval [0. 1] into 1/c disjoint intervals of size each. Recall that P(-1=f)
Dir(1 1). which implies that
the random variable P(0) = for each fixed = (randomness is over P) is distributed as,
Pr [P(]) -2af] Unif 1]).
(48)
Consider the buckets B, {0,,e,, 22)..... [1 E, 1]}. Define the function round(p) [0, 1] (0.--. [Bel - 1} to return the index of the bucket in B- such that p falls in that bucket.
Lemma 1. Consider any function f(z))= x* + {0, B& 1} such that, pointwise,
Iround (round(P(0))))) f(+)) s..
(49)
Then, when 2() idd Dir(1 1).
H Hshamoon ((f(2)) ot {0. (0.13*3) 2 ((1 (1- (e)-33)og(1//6 log(1/E) -log(2r 1))
(50)
where the randomness is over the draw of P and Hshannon the discrete Shannon entropy.
22

---

Published as a conference paper at ICLR 2026
Proof. Recall that P(O itd Unif([O, 1)) across of E x Then,
Pr(roun((P(012))) = - Pr(((()(=)) E [s(i - eij + 1/2)))
(51)
{e
if j S |B2l - -2,
(52)
3e/2
j = or j = B&l
This implies that, by independence of the P(O]zf)'s across zf E x
Hshannon ((P(O)=)) : o E x*)) 2 ( - 38) log(1/e).
(53)
Let c(zf) ) be the random variable P(O 2 f(zt). P() is measurable function of f(zt) and e(z1), ), and therefore,
Hshannon (1S(=)) o xx u fe(zt) o E x*)) > Shammon ((P(0)=)): E x**)) (54)
Note that c( bounded in the rate {-r,... r) and can take at most 2r + values. Therefore, H. (fe(z)) 2 x) s |211 log(2r 1). Since HA,, B) S H(A) + (B). have that,
H Shannon (15(2)) o E x²)) 2 Hshamnon ((P(0 o E x*}) - H (fe(z)) : =
(55)
2 |21 ((1 - (3) log(1 1/8) - log(2r 1))
(56)
Recall that we are guaranteed that g(5++*+1 O, He+s) E P(O)zt) + [-- +
This implies that the recurrent model is able to recover round(p) for p P(0)=f) up to an error
of r [tt(t/e]] for each 2f X* by computing round(p) where = g(x+4x+1 = 0, H.+s).
Informally, this just means that á is likely to fall in bucket close to p. In combination with Lemma for f(xx) = g( = 0, He+s) we have that,
H [Shnnnon (aaiinooon(ggge ((9( = 0, He+k)) 2f E (0, (0.1)*}) 2 2t ((1- (c) log (1/s) log(2 + 1))
(57)
Note however, that g(Zx+++1 = 0, H(+x) is a function of of implicitly, through Hz+* (which is also function of æ1). Since the dimensionality of Hetk is d and the model is implemented bits of precision,
H Shannon (g(nmnnn((((((+t11 = O, He+x): : of E (0,1)*}) < Hshannon (Hena) S dp
(58)
where all randomness here is induced by the random draw of the - Markov kernel P. Therefore, for the correctness guarantee Equation (45) to hold, we need,
dp 2 2* ((1-3) 38) log(1/E) - (t)/ +
(59)
in the limit + oo, and noting that lim sup (P(t) almost surely over P drawn from the prior, it is necessary that,
dp > 2 3s) log( 1/ee..
(60)
L
Fig. provides empirical evidence that the exponential dependency of the theorem on the order k is consistent.
Remark. The proof above assumees that the - Markov chain is on binary state space. However, the result can easily be extended to give the lower bound d 2 2([2*)()) for larger state spaces, well as similar scaling results for priors Dir( 1) for any / >>> 0. Furthermore, we helieve it should be possible to replace the Los егтоr guarantee in Equation (6) by the KL-divergence between the two distributions without significantly changing the conclusion (d = 20(k)
The intuition behind this result is in the manner in which the recurrent architecture carries out computation: by proceeding sequentially and compressing the information from the sequence has seen thus far at some time i. into small hidden vector, the model does not know what the next
Intuition.
23

---

Published as a conference paper at ICLR 2026
k tokens will be: the knowledge of this is vital to be able to compute the add-3 estimator time t + + with a small memory footprint. Indeed, when the identity of the next & tokens changes, the output of the model at time + k + 1 must look drastically different (as the add-3 estimator corresponds to approximately evaluating IP(- which are unrelated distributions under different choices of t. There when d and are small, the output of the model just cannot take so many values: it can realize at most 2dP possible sets. In other words, in order to succeed, the recurrent architecture is essentially forced to keep track of the number of occurrences each if E {0, 1] in the sequence at each time which costs exponential dependence & in the hidden dimension/precision.
2 possible values the set P = (P(-)+) of E 13 can take. But
0.35
d 2
0.30
d = 4
8
0.25
d == 16
d 32
Test loss
0.20
0.15
0.10
0.05
0.00
2
3 Order
Figure 6: Relation between the Markov order & and the hidden dimension d of the 1-layer Mamba model. The plot shows that d = 2 sufficient for the model to learn the k-th order Markov task. This corroborates the fact that Theorem 2 in the main paper has the correct order dependency.
24

---

Published as a conference paper at ICLR 2026
E ADDITIONAL RESULTS
E.1 LINEAR ATTENTION
0.35
1-ayst Marbe
Trandomer
Tonndorer
2-áper Translomer
0.30
0.8
layer Lincor Tronsforer
1-ayer Transbmer
Predicted probsbiitt
Optimal adiimater
layer Uneer mransforer
0.25
Translomer
0.20
0.15
0.4
0.05
0.00
Position
Orders
(a) Predicted probability
(b) Test loss
Figure 7: Comparison of predicted probability and test loss between 1 l-layer Mamba and the other baselines, including linear attention. Mamba outperforms all baselines. Linear attention and softmax attention Transformers perform similarly.
E.2
MULTIPLE STATES
Fig. shows that our results extend to larger number states.
0.5
0.4
gep relalive optinal
0.4
0.3
0.3
Test loss
0.2
0.2
0.1
0.1
0.0
0.0
5
10
15
10
15
20
25
Iteration 200)
Iteration (x
(a) Absolute test loss gap
(b) Relative test loss gap
Figure Test loss gap from the optimal for 1-layer Mamba and first-order Markov data, for different number of states. (a) shows the absolute gap T.()) L": (b) shows the relative gap (T.(0) TT" The loss gap consistently small for all state sizes.
E.3 DEEPER NETWORKS
Fig. 10 show that our results are consistent with larger number of layers.
E.4
PWER-PARAMETRIIED SETTINGS
Fig. 9 shows that observations hold also for larger convolution width and deeper networks.
25

---

Published as a conference paper at ICLR 2026
0.4
Lass from aptimal
0/2
gop
0.1
0.1
20
80
100
Iterations
Iterations
(a) Convolution width
(b) Depth
Figure 9: Test loss gap from optimal across iterations pee-pprametrized settings. (a) shows 1-layer Mamba with varying convolution width. (b) shows Mamba with convolution width 2 and varying number of layers. The models correctly learn the optimal predictor in all cases.
0.35
0.35
layer
layer
0.30
layers
0.30
2 layers
layers
3 layers
layers
layers
Test loss
0.20
0.20
Test
0.15
0.15
0.10
0.05
0.05
0.00
0.00
Order
Order
(a) Mamba
(b) Transformers
Figure 10: Test loss gap from the optimal for Mamba and Transformers, for different number of layers and Markov orders. Mamba has a smaller loss gap than transformers across all orders. Furthermore, adding more layers to Mamba does not significantly improve performance. As expected, 1-layer Transformers cannot solve the Markov task, while I-layer Mamba can.
E.5 HOLD-OUT EXPERIMENT
We also consider the following interesting experiment: we trained Mamba only on sequences from subset of the Markov simplex, specifically the interval j0, 0.5). and we tested it on sequences from its complement [0.5, 1]. Interestingly, optimal. However, by inspecting the actual estimated probabilities on a fixed test sequence, we see that the absolute difference between the model's estimation and the optimal Laplacian smoothing is high for the first samples, and gradually decreases as the sequence progresses. This is justified by the fact that the Laplacian estimator of the prior on the simplex) in the limit of long sequences (i.e., the gap from the optimal loss goes n 3o). Table 2 shows the absolute difference [P 1|x[-1 F* (x+ 1|24-1) for several t, showing how this difference gradually goes to zero as increases.
experiments show that the test loss still converges to the
only Bayes-optimal, but also minimax optimal (i.e., independently
E.6 CONVOLUTION
Fig. [[]shows that adding convolution to transformers (similarly to Mamba) makes the model solve the task with just one layer. On the contrary, Fig. 12 shows that Mamba needs two layers to solve the
26

---

Published as a conference paper at ICLR 2026
Table Results for the hold-out experiment.
Length
Abs. difference from optimal
0.378 0.050
50
0.098 = 0.039
100
0.007 0.003
300
0.001 = 0.001
task if convolution is removed. Table shows that one-layer Mamba without convolution cannot learn the optimal estimator, no matter how wide the model is.
1.0
0.40
1-layer Mamba
0.35
1-layer Transfomer
0.8
1-layer Transformer convolution
Predicted probabilty
0.30
0.6
0.25
0.4
1-ayer Mamba
0.2
1-aaer Transformer
1laayrr Transformer convolution
Optimal predictor
0.00
10
20
30
40
50
60
70
5
10
15
20
25
30
35
=
Iteration 200)
(a) Predicted prohability
(b) Test loss
Figure 11: Predicted probability and test loss for Transformers with and without convolution. Adding convolution the K, Q, V matrices of transformers makes the models succeed in learning the Markov task, similarly to 1-layer Mamba.
0.5
1-laer
1-layer Mamba without convolution
0.8
0.4
2-layer convolution
Predicted probabilty
0.7
0.3
Teat
0.2
1-layer Mamba
1-layer Mamba without convolution
0.1
2-layer Mamba without convolution
Optimal predictor
0.0
10
20
30
40
60
70
6
8
10
12
14
t:x
Iteration (x 200)
(a) Predicted probability
(b) Test loss
Figure Predicted probability and test loss for the full 1-layer Mamba and 2-layer Mamba without convolution. Similarly to transformers, Mamba needs two layers to solve the Markov task when convolution removed.
E.7 SWITCHING MARKOV
Here we detail the switching Markov process more formally:
1. Initialize /. 0.
27

---

Published as a conference paper at ICLR 2026
Table Experiments on one-layer Mamba without convolution, with varying width. The model does not learn the optimal estimator successfully.
Hidden dimension d Avg. test loss gap from optimal
10
0.113 + 0.032
100
0,158 + 0.055 0.140 + 0.072
1000
2. Draw a binary Markov kernel P with each row sampled i..d. from Dir(3 8 1).
3. Let = S with probability Pswitch: or sample 2t ~ Pa; #11 with probability Pswitch-
4. If = S, set & = + and to step 2; if Te * S, set t = =t + 1 and go to step 3.
Fig. 13 illustrates the behavior of Mamba on this process when Pswitch = 0.01.
0.8
1-layer Mamba
1.0
0.7
Optimal estimator
0.8
Prodictd probability Pe (R211
0.6
0.5
0.6
0.4
0.3
0.4
0.2
0.2
0.1
0.0
0.0
50
100
150
200
250
300
f: = 0
Position
(a) Predicted probabilities
Value a, across positions
Figure 13: One-layer Mamba on switching Markov data. Mamba is able to learn the optimal predictor by forgetting past counts every time switch token occurs. This achieved by setting a, = at every switch, and a, = otherwise.
E.8 FURTHER ABLATION FOR NATURAL LANGUAGE
The fundamental role played by convolution in modeling natural language is demonstrated in Table However, we find that other components, in particular gating factor o₁, play an essential role as well, when natural language is considered. In Table 4 we show the increase in test perplexity when convolution, gating factor a and non-linearities are individually removed from the full Mamba-2 architecture.
Table Perplexity results on the WikiText-103 dataset.
Model
Params.
Perplexity
Percentage increase
Mamba-2 (full)
14.54
27.55
Mamba-2 (w/o conv)
14.53 M
30.68
11%
Mamba-2 (w/o gating factor)
14.54 M
32.16
17%
Mamba-2 (w/o non-linearities)
14.54
28.98
5%
28

---

Published as a conference paper at ICLR 2026
Table 5: Perplexity results on Wik and PG-19 datasets for Mamba with 12 layers.
Dataset
Model
Params.
Perplexity
WikiText- 103
Mamba-2
110 M
21.38
WikiTexx-100 103
Mamba-2 w/o convolution
110 M
21.46
WikiText-100 103
Mamba-2 w/o gating
110 M
21.71
PG-19
Mamba-2
200 M
14.16
PG-19
Mamba-2 w/o convolution
200
14.28
PG-19
Mamba-2 w/o gating
200 M
14.66
29

---

Published as a conference paper at ICLR 2026
F MODEL ARCHITECTURES AND HYPER-PARAMETIE
The following tables discuss details on the architectures and hyperparameters used in all the paper's experiments. Each experiment was run on single Nvidia A100 GPU. The time taken by each experiment was between 10 to 60 minutes.
Table 6: Parameters in the Mamba architecture with their shape.
Parameter
Matrix shape
embedding
x
mamba.A mamba.dt
1
mamba.in_
(2ed + 2N + 1) x d (ed 2N)
mamba.conv ld
namba.out_prye
d x (2ed + 2N + + Ad x d
mlp.fel
mlp.fc2 Im head
d x 4d dx2
Table 7: Settings and parameters for the Mamba model used in the experiments.
Dataset
k-th order binary Markov source
Architecture
Based the Mamba architecture implemented in Dao & Gu 2024
Batch size
Grid-searched in {16, 32, 61, 128, 256}
Accumulation steps
Optimizer
AdamW (31 0.9, 82 0.95)
learning
0.001
Scheduler
Cosine
# Iterations
10000
Weight decay
1 x 10-3
Dropout
0
Sequence length
Grid-scarched in {128, 22,,25511
Embedding dimension
Grid-scarched (2, 4, 16,32) 32}
Mamba layers
Heads Convolution window
Between 2 and 6
Repetitions
5
Table 8: Parameters in the transformer architecture with their shape.
Parameter
Matrix shape
transformer.wte
2
transformer
N x d
transformerh.in (x()
d dxl 1 3d x d d dx 4d x d d d 4d
transformerh.attt.. (x
transformerhhatth. c_torrrc__ooooo (x
transformer...n. (x
transformer..mipp hwtp.c.lp.e....... fc (x
transformer.h.no proj (x
transformer.In
dxl
30

---

Published as a conference paper at ICLR 2026
Table Settings and parameters for the transformer model used in the experiments.
Dataset
k-th order binary Markov source
Architecture
Based GPT-2 architecture as implemented Pagliardini
Batch size Accumulation steps
Grid-searched {16, 32, 64, 128, 256}
Optimizer
AdamW (31 = 0.9, 82 = 0.95)
Learning rate
0.001
Scheduler # Iterations
Cosine 10000
Weight decay
1 x 10
Dropout
Sequence length
Grid-searched in (128, 256, 512, 1024]
Embedding dimension Transformer layers
Grid-searched in (4, 8, 16, 32}
Between 1 and 2 depending on the experiment
Attention heads
Repetitions
5
31


Reviews:
[{"review_id": "1mFrImUweT", "reviewer": "Reviewer_ntkq", "summary": "The paper studies in-context learning in Mamba-2 using Markov Chains, and shows an interesting finding of a single layer Mamba-2 model being able to very closely track the performance of the optimal in-context Laplacian smoothing estimator. The paper contains several experiments and ablations that focus on comparing a one layer Mamba-2 model to a 1 and 2 layer Transformer, specifically focusing on the Convolution part of the overall architecture as an important piece. The paper also presents several theoretical insights which map well to the observed results.\n\nMy review comes from viewing the paper from the lens of someone familar with the empirical results with larger scale Mamba-2 models and how they relate to Transformers at the whole model (not necessarily one or two layer wise) level.\n\nOverall, I found the paper technically sound and addressing a relevant question to the Mamba line of models. I am slightly learning towards accepting the paper, but do have certain comments and questions that withhold me from a stronger rating (see below).", "questions": "- Re: the importance of convolution, I found these two related comments by the Mamba-2 author on their official implementation:\nhttps://github.com/state-spaces/mamba/issues/525#issuecomment-2290739534\nand \nhttps://github.com/state-spaces/mamba/issues/433#issuecomment-2199557446\nSpecifically, “conv1d helps a bit on perplexity: e.g. at 360M on 7B tokens on the Pile, using GPT2 tokenizer, w conv1d we get around 8.6 perplexity and without conv1d is around 8.9 ppl.”. This represents a slight drop in overall model performance but not as significant as the one seen in the authors experiments in the paper. Do you have any theory about why this may be?\n\n- In the Mamba-2 paper, the original authors illustrate the duality between Transformers and SSM’s. The Mamba-2 layer itself can be seen as a special case of a Transfomer (and vice versa) under certain assumptions and has generally very similar training characteristics. Given this duality, how do the authors feel about the presented results showing that a 1-layer Transformer is significantly worse at ICL than a one layer Mamba?\n\n- Do you have any speculation about how your findings may change / evolve over multiple stacked layers as is the typical case in most models?\n\n- For larger scale practical deployments of Mamba models, do you think we should expect better in-context learning results generally compared to Transformers?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- I found the idea interesting! Thank you for sharing this from a perspective that is intriguing and different to the typical “we ran a bunch of experiments on the same benchmarks and boosted performance by X%” kind of LM study. It is positive to see Mamba models being connected with optimal statistical estimators.\n\n- The paper is well written. Even though it may not be easy to read on a cursory look (I needed to go back and forth to see symbols and notation frequently), it does contain all the information necessary to understand it. The paper also introduces the core idea behind Mamba-2 well.\n\n- In addition to the well motivated experiment operating on random Markov chains sampled from the vocabulary, the authors also repeat the same exercise on a true token distribution from a dataset like WikiText103 that would be more representative of the real world token transition probabilities. Modeling these “biased” next-token probabilities seems like the essence of language modeling and in-context learning to me, and I find it positive that we see perplexity results on real text too.\n\n- Overall, I found the experiments to be well designed and aligned well with the theoretical insights.", "weaknesses": "- For the non-Markovian setting in Section 5, (ppl on WikiText103 being the only result), I feel like the evaluation could have been a bit more diverse and rigorous. This is the setting that I assume most practitioners (for better or worse) would be interested in, and it may have been useful to do this over a slightly wider range of model sizes (number of layers, dimensionality etc) to help build better intuition.\n\n- There is a lot of related work on architectures that do sub-quadratic attention, e.g. the original Linear Attention work, or related recurrent/State Space style models like DeltaNet, RWKV etc. While I understand the specific focus on Mamba-2 in this work (which in itself needed a lot of analysis), the inclusion of 1-2 other related methods in the comparison (e.g. transformer with Linear Attn in addition to the 1 and 2 layer Transformer) would have strengthened my evaluation of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1L3Uw2vGeg", "reviewer": "Reviewer_5agm", "summary": "This paper studies the in-context learning abilities of the Mamba (S6) architecture, as a cheaper alternative to transformers (specifically, their quadratic attention cost in the sequence length). To study this problem, the authors consider finite-state k-order Markov Chains, and evaluate Mamba and Transformers in their ability to output the transition probabilities of the optimal Laplacian smoothing estimator in-context. Empirically, the findings include Mamba fitting the Laplacian smoothing estimator with only 1 layer, and convolution being the main ingredient in this ability. Furthermore, the authors show that a simplified Mamba variant can theoretically encode all the bits needed for storing the Laplacian smoothing operator with specific choices of its parameters.", "questions": "- Did the authors check that 2 or more layers Mamba achieves the same result as 1-layer Mamba in the considered experiments? Or in other words, what is the impact on performance if the model is over-parameterized with respect to the bits needed to encode the Laplacian smoothing estimator?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The paper is clearly written and well motivated.\n- The methodology considered in this study is sound and made explicit: Building on the empirical observation that the convolution is the key ingredient in the ICL abilities of Mamba, the authors construct a minimal architecture that allows a theoretical analysis all while keeping key similarities with the original architecture. This allows the derivation of sound theoretical results characterizing Mamba's ability to learn in-context, as function of the choices in parameters and the hidden dimension.\n- Given equivalences between language modeling, and k-order Markov Chains (e.g. n-gram models which previously had been the standard in language modeling), I believe that the insights included in this study are an important step towards building more efficient architectures in general, and in Language modelling in particular.", "weaknesses": "- Theorem 1 is insightful in proving the in-context ability of Mamba. However, in my understanding it only tackles the model misspecification problem, in the sense that it proves that the optimal solution can be represented by the model (zero approximation error). I think it would be interesting to discuss the intersections between the theoretically chosen parameters that achieve the result in Theorem 1, and those actually obtained by optimizing Eq 3 on a given training set (with or without distribution set) starting from standard random initializations of these parameters. This is partially done, with the observation that $a_t \\approx 1$ at convergence, but can also be extended to other parameters (for instance, the count-independent term in Eq 5 that need to sum up to $\\beta \\mathbf{1}$).\n- The natural language modeling experiments only validate the key role of the convolution in Mamba. I think that even the theoretical results (especially Th 2) can also be validated through experiments. For instance, by generating language from an n-gram model, one can show that a given hidden dimension size of Mamba is strictly needed in order to be able to perform in-context learning tasks in the generated language corpus. Of course this comes with challenges related to the exponential dependence on the Markov Chain order, but having a discussion around this would be a nice addition to the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0EtcU2071F", "reviewer": "Reviewer_6hFD", "summary": "This paper builds upon the previous Markov-ICL framework to empirically and theoretically investigate the ICL ability of Mamba architecture.\nSpecifically, they design experiments to train Mamba for next-token prediction loss on random Markov chains, and find one layer Mamba match the performance of the statistical optimal estimator, and convolution layer play an important role to ahcieve this performance. Theoretically, the authors provide a constructive proof showing how Mamba's components work in concert to implement the laplace smoothingy: convolution identifies the $k^{th}$-order transition, recurrence (with $a_t \\approx 1$) accumulates all transition counts into the hidden state, and selectivity reads out the specific counts needed for the current prediction. Finally, the paper proves a lower bound showing that for Mamba the hidden dimension $d$ must scale exponentially with the Markov order $k$.", "questions": "1. Does the result in Theorem 2 imply that the transformer is more powerful than Mamba in in-context learn the complex task (higher order Markov process)? Can we conclude transformer is more parameter-efficient than Mamba when solving complex tasks? Do authors have any insight into a way to design a Mamba block that sidesteps this trade-off?\n\n2. In the theoretical proof, the author omits the gating mechanism in Mamba, but later found in the NLP task that gating is an important component. If we want to add gating back to theoretical analysis, how does it interact with the \"counting\" mechanism described in Theorem 1?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This is the first work to investigate theoretically the ICL ability of the Mamba architecture.\n\n2. The paper is well-written; the logic and arrangement are easy to follow, and the author provides enough background in the main text to help readers understand the scientific question being studied. It starts from an empirical observation, finds convolution plays an important role, and then echoes this importance in the theoretical analysis.\n\n3. The paper is perfectly designed and written to shed light not only on the theoretical analysis but also on the real-world application and the Mamba architecture. The theory echoes the key component and mechanism in the original Mamba paper (convolution and the selective mechanism) and is supported by experiments on synthetic data.", "weaknesses": "1. One possible weakness is that, to the best of our current understanding, Mamba has not demonstrated the strong ICL ability seen in large transformer-based language models. While this is not a weakness of the theoretical analysis itself, Mamba's actual ICL capability in real-world applications will influence the impact of this theoretical work on the community.\n\n2. The real-world experiment is only on the language modeling task (e.g., perplexity on WikiText-103) instead of a dedicated, real-world ICL experiment, which limits the generality of the findings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IApC7wSCc2", "reviewer": "Reviewer_PRst", "summary": "In this paper, the authors intended to provide theoretical understanding on the learning capability of MAMBA on Markovian data. The authors are able to demonstrate through a constructive proof that Mamba is able to efficiently learn the\nin-context Laplacian smoothing estimator even with just one-layer, thus ensure both Bayes and minimax optimality through the properties of the Laplacian. In addition, the authors established the lower bound on the hidden dimension required for the representing the optimal estimator.", "questions": "My main question is related to the Dirichlet distribution assumption and its relation with Laplacian estimator, How much do these assumptions affect the generality of the results.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is very well-written, the basic concepts and main steps in the arguments are clearly presented and explained. \n\nThe results presented in the paper is valuable to the understanding of Mamba.", "weaknesses": "The optimality of the Laplacian estimator is based on the Dirichlet distribution assumption, it will be helpful to see the impact of the assumption. \n\nIt would be helpful that the quality of the lower bound presented in Theorem 2 can be addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)