Understanding Input Selectivity in Mamba: Impact on Approximation Power, Memorization, and Associative Recall Capacity
Authors: Ningyuan Teresa Huang
Miguel Sarabia
Abhinav Moudgil
Pau Rodriguez
Luca Zappella
Federico Danieli
Venue: ICML 2025 poster (2025)
Keywords: Mamba
State Space Model
Input selectivity
Function Approximation
Memorization
MQAR
Primary area: deep_learning->theory
TL;DR: In this work, we demystify the role of input selectivity in Mamba, investigating how it impacts its function approximation power, long-term memory, and associative recall capabilities.

Abstract: State-Space Models (SSMs), and particularly Mamba, have recently emerged as a promising alternative to Transformers. Mamba introduces input selectivity to its SSM layer (S6) and incorporates convolution and gating into its block definition. While these modifications do improve Mamba's performance over its SSM predecessors, it remains largely unclear how Mamba leverages the additional functionalities provided by input selectivity, and how these interact with the other operations in the Mamba architecture. In this work, we demystify the role of input selectivity in Mamba, investigating its impact on function approximation power, long-term memorization, and associative recall capabilities.
In particular: (i) we prove that the S6 layer of Mamba can represent projections onto *Haar wavelets*, providing an edge over its Diagonal SSM (S4D) predecessor in approximating discontinuous functions commonly arising in practice; (ii) we show how the S6 layer can dynamically counteract memory decay; (iii) we provide analytical solutions to the MQAR associative recall task using the Mamba architecture with different mixers --- Mamba, Mamba-2, and S4D. We demonstrate the tightness of our theoretical constructions with empirical results on concrete tasks. Our findings offer a mechanistic understanding of Mamba and reveal opportunities for improvement.

Understanding Input Selectivity in Mamba:
Impact on Approximation Power, Memorization, and Associative Recall Capacity
Ningynan Huang 12
Miguel Sarabia
Abhinay Moudgil 14
Pau Rodríguez
Luca Zappella 3 Federico Danieli 3
Abstract
ers, and parallelizability compared to (nonlinear) Recurrent Neural Networks (RNNs). In particular, Mamba (Gu & Dao, 2023; Dao & Gu, 2024) demonstrated state-of-the-art per- formance on various language modeling tasks, with smaller model size and faster inference than Transformers. The suc- cess of Mamba has largely been attributed to the fact that the parameters in its SSM layer (S6) are input-dependent, leading to improved expressivity compared to its SSM pre- decessors (Cirone et al., 2024). The goal of this work is to provide a more structural explanation for Mamba's 's superior performance. We do so by answering two questions: (i) how does the S6 layer's expressivity translate into its practical performance? (ii) how can the S6 layer interact with the rest of the Mamba block to solve concrete tasks?
State-Space Models (SSMs), and particularly Mamba, have recently emerged as a promising alternative to Transformers. Mamba introduces input selectivity to its SSM layer (S6) and in- corporates convolution and gating into its block definition. While these modifications do improve Mamba's 's performance over its SSM predeces- sors, it remains largely unclear how Mamba lever- ages the additional functionalities provided by input selectivity, and how these interact with the other operations in the Mamba architecture. In this work, we demystify the role of input selectiv- ity in Mamba, investigating its impact on function approximation power, long-term memorization, and associative recall capabilities. In particular: (i) we prove that the 56 layer of Mamba can repre- sent projections onto Haar wavelets, providing an edge over its Diagonal SSM (S4D) predecessor in approximating discontinuous functions com- monly arising in practice: (ii) we show how the S6 layer can dynamically counteract memory decay; (iii) we provide analytical solutions to the MQAR associative recall task using the Mamba architec- ture with different mixers - Mamba, Mamba-2, and S4D. We demonstrate the tightness of our the- oretical constructions with empirical results on concrete tasks. Our findings offer mechanistic understanding of Mamba and reveal opportunities for improvement.
We answer question (i) by providing fine-grained analysis of the S6 layer via the lens of function approximation and long-term memory. We prove that the S6 layer can represent projections onto Haar wavelets and thus efficiently model discontinuous signals, which is relevant for solving practical tasks. Moreover, we show that the S6 layer still suffers from exponential memory decay, but highlight mechanism which allows it to dynamically counteract such decay.
Building on the understanding of the S6 layer, we then investigate how S6 interacts with the other components of the Mamba block (which also encompasses a short- convolution and a gate branch, see Tab. 2) to tackle the Multiple-uuery Associative Recall (MQAR) task (Arora et al., 2024a). We prove that 1-layer Mamba and Mamba-2 models can both solve MQAR, even without gating; we describe how S6 and the convolution interact to achieve this, and how Mamba-2 leverages its independent convolutions to get more parameter-efficiee efficieet solutions. We also show that a I-layer Mamba can solve MQAR exactly even without input- dependence in its SSM: this (perhaps unexpected) re- sult helps cementing the importance of convolution and gat- ing in the Mamba architecture. This analysis further informs us how a variation to the functional form of Mamba, and particularly to how the SSM state matrix is affected by the input, can improve its performance on the INDUCTION HEADS task (Bietti et al., 2024; Sanford et al., 2024a).
1. Introduction
State Space Models (SSMs) have recently emerged as a promising approach for long-range sequence modeling, due to their computational efficiency compared to Transform-
'Wor done while Apple FFlatiron Institute, New York, USA 2 "Apple "Mila Research Institute, Montreal, Canada. Correspon- dence to: Ningyuan Huang <thuang @ flatironinstiuut.org..
Proceedings of the nd International Conference Machine Learning, Vancouver, Canada, PMLR 267, 2025. Copyright 2025 by the author(s).
We complement our theoretical findings with numerical ex- periments on synthetic sequence modeling tasks, For the S6

---

Understanding Selectivity in Mamba
layer, we demonstrate its approximation power on discon- tinuous functions, and its counteraction of memory decay, via the KEEP n-T1l task - generalization of KEEP FIRST from (Chiang & Cholak, 2022) requiring to memorize the n-th token in a sequence. Finally, for the full Mamba model, we confirm empirically that the model sizes prescribed the- oretically by our analytical solutions to the MQAR and INDUCTION HEADS tasks are tight in practice. Overall, our contributions can be summarized as follows:
language theory and approximation theory. Studies follow- ing the first approach investigate what formal languages can SSM recognize (Merrill et al., 2024; Sarrof et al., 2024; Grazzi et al., 2025). Studies ascribing to the second ap- proach - including this work - characterize what function classes can SSMs approximate: Li et al. (2022) showed that SSMs can approximate linear functionals with exponential memory decay; Wang et al. (2024) extended this analysis to nonlinear RNNs, showing however that adding nonlinearity (in the hidden state recurrence) does not fix the memory decay issues. Orvieto et al. (2024) proved that SSMs aug- mented with MLPs are universal approximators of regular functionals, but this improvement over Li et al. (2022) and Wang et al. (2024) requires the hidden state size to grow lin- early with sequence length. Cirone et al. (2024) extended the results from Li et al. (2022) input-eepenent SSMs such as Mamba, showing their universal approximation the class of nonlinear functionals arising from controlled differ- ential equations. While Cirone et al. (2024) highlighted how Mamba can approximate a larger class of functionals than S4D, we explore the consequences of this observation by characterizing specific function classes arising in practice.
We prove that the S6 layer of Mamba can represent projections onto Haar wavelets, providing an edge over the S4D layer in approximating discontinuous functions commonly arising in practice (Sec. 4.1).
We use sensitivity analysis to show the S6 layer gener- ally suffers from exponential decay of memory, and de- scribe how it can dynamically counteract (Sec. 4.2).
We show how the Mamba architecture can exactly solve the MQAR task using different SSM mixers, with an explicit characterization of the required model size, which helps explaining their performance differ- ence (Sec. 5.
Associative Recall Capability Associative Recall (AR) describes the ability of a model to retrieve information from its memory, based on the input context. Tasks for evaluating associative recall capabilities include INDUCTION HEADS (Olsson et al., 2022; Sanford et al., 2024a), k-HOP INDUC- TION HEADS (Sanford et al., 2024b), MULTIPLE- AULTIPLE-QUERY ASSOCIATIVE RECALL (MQAR) (Arora et al., 2024a), and NEEDLE-IN-TIE- STACK (Kamradt, 2023). 1mmpricmlii,,llly Mamba (Gu & Dao, 2023) and Mamba-2 (Dao & Gu, 2024) demonstrated performance competitive with Transformers on (a simple version of) INDUCTION HEADS and MQAR, respectively. Theoretical understanding of how language models perform associative recall begins to emerge: Bi- etti et al. (2024) constructed a 2-layer Transformer with positional encoding that solves the INDUCTION HEADS task; Sanford et al. (2024b) extended such construction to a log-depth Transformer that solves the k-H1P INDUC- TION HEADS task; Arora et al. (2024a) showed that a gated convolution model can solve MQAR. However, these con- structions all require the model size to scale with the input sequence length. In this work, we show that 1-layer Mamba models can solve INDUCTION HEADS and MQAR with model size independent of the sequence length, highlighting the role of the convolution operation and the gate branch, that has been previously overlooked in the literature.
Our findings reveal opportunities to further improve Mamba, such as by changing the way input dependence is incorporated into the SSM state matrix (Sec. 5.2).
2. Related Work
State-Space Models (SSMs) SSMs (i.e., linear RNNs) have recently emerged as a promising soquence-modeling approach, with faster training than (nonlinear) RNNs, and faster inference time than Transformers. To enable long- term memorization capability, Gu et al. (2020) designed SSMs as polynomial approximations of signals, by prescrib- ing non-normal HiPPO matrices as state matrices. To im- prove computational efficiency, Gu al. (2022b) proposed the $4 model: first by reparameterizing HiPPO sum of normal and low-rank matrices, then by further consider- ing a diagonal simplification in the S4D model (Gu et al., 2022a). All these SSMs are Linear Time-Invarian and thus computationally efficient, but consequently lack the ability to process information in a vimeinarying, dependde manner. Mamba (Gu& & Dao, 2023) overcomes this by us- ing input-dependent SSM parameters, without sacrificing computational efficiency, showing performance competitive with Transformers on long-range language modeling tasks. Mamba-2 (Dao & Gu, 2024) simplifies the state matrix to a scalar multiple of identity and modifies the Mamba model architecture, showing further empirical improvements.
3. Preliminaries: Linear RNNs as SSMs
The application of a Linear RNN can be interpreted as the discrete solution of a linear dynamical system, or SSMs,
Expressivity of SSMs The analyses on the expressivity of SSMs can be divided into two main approaches: formal
h(t) A (())((())(())))))
y(t) = C(()((). (1)
2

---

Understanding Selectivity in Mamba
Here, the input st() E I acts as forcing term for the hidden state h E TRN through the application of the input matrix B(t) E IR-x11 The natural evolution of the hidden state is dictated by the state matrix Att) E RNX,, Finally, the output y E IRdy is obtained by linearly transforming the hidden state via the output matrix C(t) E RR Col- lectively, we call A()), B(),, (()) the SSM parameters. For computational efficiency, modern Linear RNNs con- sider diagonal state matrices, with negative eigenvalues A(t) (() = diag( (t)....... (t()... This simpli- fies the solution of (1), and ensures stability in the evolution of the hidden state. Under this assumption, and considering an initial state (0) = O, we can explicitly write the hidden state solution as an integral function of the input (Dahleh et al., 10:
and Mamba (Gu & Dao, 2023), which models tame-narying (input-dependent) dynamies,
(t)) = (onneaaaa A = A,
(t) ((x(t)) = Lincar(æ(t)).
(7)
C(1) = (zx(() = Linear(zc())-
Note that both the Mamba-2 mixer and Linear Attention (Katharopoulos et al., 2020) can be interpreted as special- izations of Mamba (Dao & Gu, 2024), where A, = X, and A, = I, respectively.
4. Mamba SSM Mixer Layer Analysis
In this section, we focus on the core of the Mamba architec- ture: its SSM mixer layer, also referred to as S6. The goal is to understand how input-selectivity impacts its expres- sivity and long-range memorization eapacity. We show that an S6 layer can represent projections onto wavelets, thus efficiently modeling discontinuous signals. This useful in practice, e.g., for isolating specific tokens in a sequence. For long-range memorization, to show that, while S6 suffers from exponential decay of memory (akin to S4D), input-selectivity allows to decrease the rate of such memory decay by "freezing time". We vali- date our theoretical insights by experimenting on the KEEP n-TH task.
h(t)= = S eli dr) B(s) ((s) ds.
Generally, inputs to Linear RNNs are provided as (discrete) sequences of values [r=]]=1, and thus we consider a dis- cretization of system (1). This amounts to substituting the
sensitivity analysis
differential equation with an (approximating) recurrent one:
(1) = h = A + Z1,
yr = Ct)(t),
(3)
where Ac, Br. C2 are the discrete counterparts of A(), Bt)), C(t), respectively. Notice we can recover an explicit solution to (3) by unrolling the recurrence relation starting
4.1. Function Approximation Power: Expressing Wavelets via the S6 Layer
from hg = 0, to obtain
h
A, A,
B. Jz.
(4)
We begin by analyzing the expressivity of the S6 layer and the power of its input-dependent discretization in terms of function approximation capabilities. We prove that an S6 layer can approximate projections onto wavelets arbitrar- ily well (Thm. 1), while an S4D layer can at best project onto Fourier bases. Consequently, in approximating target functions with discontinuities, S6 achieves a faster approx- imation rate than S4D (Cor. 1). The advantage of S6 over S4D in approximating discontinuous functions translates into their performance differences in memorization tasks.
x=1 Arss1
The discretized SSM parameters are usually obtained follow- ing a Zeroth- Hold (ZOH) scheme (Tóth et al., 2008).
Given a discrete time-step (t) E ,R,,, ZOH prescribes
C, = C(t), B. = (( - (B()))()))
A, = AA0AAO))
(5)
Note that in the original Mamba formulation the authors use Forward Euler for (t) for simplicity, B, = (())
Linear RNNs Projectoon Onto Basis Functions The idea that Linear RNNs could perform projections onto specific sets of basis functions is not new, and indeed served as theoretical grounding for the HiPPO work (Gu et al., 2020). However, to our knowledge, this interpretation has not yet been leveraged explain the capabilities of modern Linear RNNs. In what follows, we take such interpretation to analyze their expressivity. To streamline the analysis and slim notation, we focus on 1D inputs and view them as continuous signals, x, = x(s) E IR, s E [0, 7', without loss of generality. For direct comparison with S4D, we con- sider a simplified version of S6 where the -epuuddeee only affects the state matrix via A A(r), and does not af-
In general, there is some flexibility in the choice of the func- tional form that the system parameters can take. The main requirements are that: (i) A, is diagonal(-zaable,, so that the product I 1 A, in (4) can be computed efficiently: (ii) the eigenvalues of A, are bounded in 1, 1]. to ensure stability; and that (iii) the SSM parameters do not depend on the state htt),, so to keep the system linear in (tt,, and allow to compute its solution in parallel along 6. The most relevant choices analyzed in this paper are: S4D (Gu al., 2022a), which models linear time-invarian dynamics,
(1) == 1, A(t)= = = A, B(I) = B, C(t) = CC:
(6)
3

---

Understanding Selectivity in Mamba
.., By]" independent of the input 244). Substituting this into (2), and recalling that A diag( (A,, Ax]), each component n = I,... N of the hidden state can be evaluated separately as an inner product between time-dependent functions:
feet (Ix) (i.e., B(xe)= [B,....,
Theorem 1. Consider a Haar wavelet $3 : [0, R. and the Mamba basis function (8) at t 1, -5.6 1, -3. St Asa(zz)dr B3 Let ã. == concat(z, s) be the in- put signal augmented with time positional encoding. For any > 0, there exist 3 Mamba basis functions 92,k ff.x Fzk such that the approximation error
(t) = 1/' e -- ((a)) de By x(s) ds = x)). (8)
(2n(s)) - (3: 1,5) (x -
We refer to (h; t, x) as the Mamba basis function, with the notation emphasizing its general dependency on the input signal x up to time I. For ease of comparison, we can recover an analogous formula to (8) also for S4D, by letting A(x)) = 1 Vt (see also (6)). This gives
is smaller than e, Vs E 10.
The proof relies on tweaking the input-dependent discretiza- tion (s)) to output (s)) -> oc or (ss) - 0 (note that A can directly depend on the time variable & instead of the input signal Is, due to time Positional Encoding (PE)). This effectively allows Mamba to represent Heaviside functions as bases: by linearly combining shifted Heaviside bases, one can immediately recover the required Haar wavelets (see Fig. I. middle-right subplots). The details of the proof are reported in App. A.2, where we also proceed to relax the inclusion of PE as an assumption for Thm. 1.
(t) I-- -dnlt-))B, Br x(5) ds (0 .x)
(9)
1i (s:t)
As we can see, S4D can provide only exponentials as basis functions. The approximation properties of these functions are limited: this is established in the literature, and ties back to the theory of Vandermonde matrices (Gautschi & Inglese, 1987), recently pointed out by Orvieto et al. (2024). Their poor performance are mainly due to: (i) the stability con- straint, Re(-X,) I O, which causes an exponentially fast decay to 0, de- de-facto limiting the effective support of said bases; (ii) the large degree of overlap between different bases (obtained by varying the only free parameter X,) The only way to curb these negative effects is by pushing the eigenvalues to be equispaced onto the complex unit disk, namely e An + c This is precisely the strategy rec- ommended by Orvieto et al. (2024), however it reduces the application of the S4D layer to simply performing a Fourier transform. In contrast, the additional flexibility provided by the uaput-selectivity in Mamba allows for a much richer va- riety of basis functions (8) to be employed in the projection, an example of which is shown next.
Theorem translates into practical advantages of Mamba over S4D, Haar wavelets are much better than Fourier bases for approximating discontinuous functions common in practice. This is formalized in the following corollary.
Corollary 1. For piecewise-constant function p(t) with m Z 1 discontinuities, there exist N Mamba functioos (8) such that the L?2 approximation error ||p - > Iül l1> is of order O(2 sm) ). On the other hand, S4D basis func- tions can achieve an approximation error of O(N
Corollary stems from Thm. 1, and from approximation results using Haar wavelets and Fourier bases available in the literature (Vetterli, 2001; Bekhoff, 1993); see proof in App. A.2. In the following, we illustrate how approximating wavelets translates into advantages on concrete tasks.
Mamba Bases Can Represent Haar Wavelets Here we provide the main theoretical result in this section, namely that the Mamba S6 layer can perform projections onto Haar wavelets. Due to their ability to capture local aspects of a function such as spikes and discontinuities, wavelets are generally better suited than Fourier bases in solving cer- tain signal processing tasks, (e.g., needles-in in-the-haystace haystack (Kamradt, 2023), transient signals (Mallat, 2012)). Recall the Haar wavelets are defined by dilation and translation,
Task KEEP n-TH The goal of the KEEP n-TH task is to recover the n-th element in a mondooll-ggeeratter sequence, y/ = Wn. This generalizes the KEEP FIRST task in (Chiang & Cholak, 2022) where n = 1. The solution of KEEP n-TH can be directly represented by combining the projections of diecewise-constar signal x²s)) Ei= (.))(s)))) onto two Heaviside functions H n., H(s - (n - 1)). As we have shown in Thm. 1, one S6 layer in Mamba can reproduce precisely this type of projections, provided the input is augmented with time-positional information. Thus, we arrive at the following Corollary:
(5) = (+) == 11 (+) - 1 (s) (s) = 23/ 22/2 k):
(10)
There exists an 56 layer that solves KEEP n-TH on input augmented with time Positional Encoding.
Corollary 2.
with E N denoting the dilation scale, and k = o,... 2 -1 the translation. Higher-orde wavelets correspond to local- ized and "spiky" bases; see Fig. A.1 (left) for an illustration.
The proof of Cor. 2 is reported in App. A.3; the results in Tab. verify empirically that Mamba with PE can perfectly
4

---

Understanding Selectivity in Mamba
where ö(A, Mas Buy depends on the imput subsequence Tgg independent of the sequence length t.
Table 1. KEEP FIFTH experimental results. Average accuracy across 3 seeds with T = 50 and - = 128. Standard error across 3 seeds 0.00 for models. Mamba and S4D models consist of embedding, SSM (h E I ) and linear layers (without convo- lution and gating, see Tab. 2). Positional Encoding (PE) encodes the position in the last clement of the embedding, resulting in [V] fower parameters.
Lemma 1 shows that both S4D and S6 have exponential decay of sensitivity when the sequence length increases. For S4D, the only mitigation strategy is set Xn - 0 (and thus e - 1). While S6 can implement the same strategy, it can also counteract the decay by adapting the input-dependent discretization, as formalized in Lem. 2.
PA
MAMIB
S4D
S4
TRANSPORMUR
Accuracy1
1.00
0.08
0.09
0.08
1.0
Lemma 2. Consider the discrete-time S6 in (7) where solve the task (same as Transformers), whereas Mamba without PE and S4D both fail, highlighting the advantage of Mamba over S4D in approximating discontinuous functions in practice. Additional ablations on model size, sequence length, and the role of PE are reported in App. D.2.
Parameters
9.3k
8.84
8.7k
6.42
B, = - (zz)...., Bn Bn(x))) E RN Suppose there exists a constant e 2 0 such that
lim A. Eal
(12)
r=1
Then the sensitivity of the n-th component of the state at time t with respect to any imput Tz is lower bounded by
4.2. Long-Range Modelling: Sensitivity Analysis
In this section, we examine the long-range memorization capacity of SSM layers, by performing a sensitivity analysis of the layer output with respect changes to the input, as the sequence length increases. To this end, we analyze the derivative of the SSM hidden state at time t with respect to the past input at time j. 1h, J. We argue that preserving sensitivity (i.e.,a non-zero derivative) is necessary for mem- orization: if the past input has no impact on the current state, one cannot hope for any information about it to be retained. With Lem. 1, we show how generally the sensitivity of an S6 layer decays exponentially fast, similarly to S4D. Thanks to input selectivity, however, the S6 layer can adjust the rate of this decay, thus dynamically tweaking the amount of in- formation to retain. In Lem. 2 we illustrate this mechanism, which proves to be useful for solving the task in Sec. 5.2.
Oh M
lim 1000
2 c
(x.( (x.) A(x)) dr,
(13)
Or,
This implies that, to retain sensitivity for longer sequences, Mamba must necessarily push ^A(z,) - 0 (or equivalently, e -Ad(2)) - 1). We illustrate this with the KEEP n-TII task shown in Fig. 1, where we report the distribution of the learned parameters e ->)(@.)) as we increase the sequence length. The shift towards 1 appears clear, validating the condition discussed in Lem. 2.
While in this section we mainly focus on the properties of the original Mamba mixer layer (Gu Dao, 2023), we note that the results proven in Thm. 1, Lem. I and Lem. 2 hold analogously for the Mamba-2 mixer (Dao & Gu, 2024). We remind that the Mamba-2 mixer layer can be seen as simplification of Mamba's 's, whereby the state matrix is parameterized by a single scalar, A = XI, ruther than its full diagonal. Nonetheless, we do not rely on Mamba's 's additional flexibility for our derivations; choosing a suitable sealar X suffices (see details in App. A.2).
Remark 1.
For simplicity, we consider ID inputs 2₁ E R. Given a generic, input-dependent recurrence relationship as in (4), we show in App. B that the sensitivity of the state with respect to its inputs at the j-th instant 2y E is given by
oh, 0 0z, dr,
(Eimi (I A. Boys)
=
= (II A.) (
Ju, (BJz))
(11)
5. Full Mamba Architecture Analysis
+ 87, Øzj
Σ
[Ii=14+1 A, B,N.).
In this section, we describe the full Manba architecture and study how its SSM mixer coordinates with the other com- ponents in the model to efficiently solve associative-recall tasks. full Mamba architecture includes an embedding layer, a number of Mamba mixer blocks, and an output layer. The mixer block is further composed of short convolution, an SSM, and gate we refer to Tab. 2 for details in the differences between Mamba and Mamba-2, but point out that Mamba-2 leverages three independent short convolu- tions (rather than a single common one) to compute its SSM parameters, as outlined in (14b).
Lemma 1. Consider the hidden states arising from the S4D and 56 SSMs defined in (6) and (7). The sensitivity of the n-th component of their states at time i with respect to the input at time j < is given by, respectively,
ah.
= o(\n, Bn, x2, No -xn(t-(i+1))
= =((,,>o An, Bn, zijj) e da Eiestt Alax),
e
0xy
5

---

Understanding Selectivity in Mamba
Propotioon of -Ad vouner hin
Progertion of vectoos bin
10075
S4D Musbo
0.006 > 00.05
73%
Proportion of
20%5
25%
10
20 30 40 50
015
Altionstes
Sequence
Figure 1. Distribution e -AA computed on test inputs by models trained to successfully solve KEEP n-TH tasks for various sequence lengths 7 (same setup Tab. 1). The left histogram confirms that the Mamba model must push e e^^A, + T increases, implied by Lem. 2, to retain information throughout the sequence. Also the S4D must behave similarly. Meanwhile, the right histogram shows that, compared S4D, Mamba has additional flexibility in forgetting irrelevant information (e AAx is mostly before timestep n) and memorizing information selectively (e xdz is mostly from timestep n onwards), while S4D is forced to memorize indiscriminately.
Table Comparison Mamba (including S4D as special case) and Mamba-2 (single-head mixers. We denote with & the Hadamard (clementwise) product, & the Kronecker (outer) product, conv(x))s) the 1-th output of the convolution, and the pointwise nonlinearity.
MAMBA
MAMBA-2
A E A 2. E R", to E NAKN
A = À E R. æ. E E lor E NdN
(14a)
(c)))(()))) E Rd
2, (conv (oone=((iiner(())))) E R
(14b)
SohtPlus(Lincar(H,, () gd
(4,t0 1)
(ottells(L E
(14c)
B, = Lān E
(B. B)
= (coav (Coeev((iinea((()))) E
(14d)
C, Linear(#:) E ENN
(C C)
C, =(comv((((nnar(())))) R
(14e)
h Ao(A11EN) @ be (A, ② dh) & B,
h1 AAA (A, & B,
(14f)
y = h: Ci 6 ER", 5 g(m:) o == (Linear(w:)) @ E xd
(14g)
While Mamba and Mamba-2 achieve performance competi- tive with Transformers and outperform their SSM predeces- sors in solving MQAR and INDUCTION HEADS, the details of how this solution can be assembled by the architecture remain elusive, with only lower bounds on the SSM mixer size available in the literature (Arora al., 2024b; Sanford al., 2024b), lacking the consideration of other compo- nents of Mamba such as convolution and gating. Here we close this gap by providing analytical constructions for 1 I-layer Mamba model that can exactly solve these tasks for any input. Perhaps counterintuitively, as we prove in Thm. 4, the components of a single Mamba mixer block are already powerful enough to solve MQAR exactly, even just using an S4D mixer layer (replacing the S6). With Thm. and Thm. 3 we further show how, thanks to the input selectivity of S6, Mamba and Mamba-2 can leverage leaner mecha- nisms to solve MQAR. Particularly, the S6 layer can use B, and C, to efficiently structure information within its hidden state, and retrieve when required. Finally, we show how the ability to structure the hidden state (used for MQAR) can be combined with the capacity to dynamically adjust the rate of memory decay (investigated for KEEP n-TH) to exactly solve INDUCTION HEADS with variant of S6. We name this variant Mamba-A T, and discuss it in Sec. 5.2.
We remark that the constructions described in this section are just possible solutions that the Mamba architectures can implement, and we do not exclude the existence of alternative ones. Nonetheless, in Fig. we verify empirically that our solutions are tight in terms of model size.
5.1. 1-Layer Mamba Can Solve MQAR
The MQAR task (Arora et al., 2024a) prescribes input and output sequences as follows
x k1, U1 knn Uxe
1,
koy-value pairs
shuffled keys, interwoven with noise
y [x, X, X,..,, X,X,
Vig .i...... ].
The keys k; are randomly chosen from a key set of size K, whereas the values U; and the noise are randomly taken from vocabulary of size [V]. The goal is to correctly predict the value associated with the corresponding key at the query positions, while other non-query positions (denoted with x) are ignored.
In this section we prove that the MQAR task can be ex- actly solved by three architectures: vanilla Mamba (Thm. 2), Mamba-2 (Thm. 3), and Mamba with an S4D mixer (Thm. 4), which we refer to as Mamba-S4D, Next we pro-
6

---

Understanding Selectivity in Mamba
Table Overview of exact solutions the MQAR task that can be implemented by the Mamha model with S4D-mixer in Thm. 4 (top), Mamba- mixer in Thm. 2 (middle), and Mamba-2-mixe in Thm. 3 (bottom). While S4D-miss mixer lacks input selectivity in its SSM layer, it can solve MQAR via gated-convolution mechanism on a larger embedding space (top). By contrast, both Mamba and Mamba-2 2 can solve MQAR with the same selective SSM layer construction without gating, and differ on the choice of convolutions (middle, bottom).
S4D
Embedding
Convolution
SSM
Gate
Encodings for any (key, value)
Combines pains with non linear size kemelss Filters out all except for (key, valacc-paiis.
Hidden state 2 xXX VV vector, summing Retrieves chunk of hiddees state over all the (key, value))pain.
combination can be uniquely identified.
corresgonding to key.
Classifier
Hidan modiiien
Sonizal
ypor
SSM
Convolation
Embaakingg
10 03
19
- bg
MAMBA
Embedding
Convolution
SSM
Orthogonal encodings. Keys have higher weights than values.
Combines pairs with non linear size kemelss Imppiamentsc - ReLU (1%,1 + 24) -
The columas are indexed the keys, storing associated values. Retrivaal seleste appreprisic column. denotes scalur 29.
Classifser
Dellddding Oupaa Jolve
Malles run encaass
Space
NIGY MR
2
Eimbedding
Pog D4] (1] [4]
P4 p9)
MAMBA-2
Embedding
Convolution
SSM
orthogonal encodings.
There linear convolutions kernel MennttyS)) 1 Stiti)))
The hiddes state columns are indesed by the keys, storing ssocisted values only (due to shift-1 convolition). Retrieval selects appropriate colume. 10M sppar
Classifier
Tolon
Dettedling Oupar
thate P4] P4
E NI NN
Eibrelding
A
vide: sketch of the proof', which we further illustrate in Tab. 3, and defer the full details to App. C.
Proof sketch. The construction is similar to Thm. 2 for Mamba, but with some simplifications due to the additional flexibility provided by three independent convolutions in Mamba-2. Specifically, we set conv = (1,0) as a shift- convolution kernel (i.e., shifting the input sequence to the right by one position), while conv, conve as the identity maps. Once again, JL Lemma allows us to reduce dimensionalities for both keys and values embeddings.
Theorem 2. There exists 1 1-layer Mamba model without gating that solves MQAR with k pairs using embedding size d = ( + log 1), and state size N = K.
Proof sketch. We first prove a construction using standard basis embeddings with d = + |V], and then apply Johnson- Linderstrass (JL) Lemma (see Lem, 4) to reduce dimension- ality. We specify the Mamba model as follows. A size-2 (nonlinear) convolution kernel combines token pairs, ex- tracting key-value pairs and removing non-informative ones (c.g., value-key, value-value). The SSM layer (S6) organizes the hidden state matrix (14f) via B, (14d) such that each column corresponds to a specific key, and holds the value embedding associated with said key (hence the state size N k). At query time, C, (I4e) uses the query(=key) embedding to retrieve the desired value from the correct column in the hidden state matrix.
Our construction of Mambu-2 similar the construction in (Li et al., 2025) that shows how convolution- augmented Transformers can solve MQAR. Indeed, in light of the convolution layer in Tab. 2, we can interpret Mamba-2 as a convolution -ugueette subquadratic Transformer.
Remark 2.
Remark 3. With the results from Thm. 2 and 3, we can infer that the extra convolutions included in the Mamba- 2 layer over vanilla Mamba in general allow for a more purameter-eficient solution of MQAR, by noting that d. = log + log [V| < k log [V), and N = =log < K.
There exists a 1-layer Mamba-2 model with- out gating that solves MQAR with n pairs using embedding size d = Ollg + log Vl), and state size N- log R.
Theorem 3.
Theorem 4. There exists a 1-layer Mamba model with an S4D mixer that solves MQAR with k pairs, using embed- ding size d = O( log [V),, and state size N 11
7

---

Understanding Selectivity in Mamba
Our Yordd
*
1000 securay
O+
5695 secuncy
<55% aooEEny
S4D
MAMBA
MAMBA-2
MAMBA AND MAMBA-A
128
Medel desemsio (ef)
Model dimension
12
*
*
X
16-
8
X Xx
X x
88
88
x X
X
88
88
x
20 40
Nomber Keys (e)
Vacabulary size (VD
Figure 2. Trained models accuracy on MQAR task (best of 7 seeds). varying and d. For
Figure 3. Trained models accuracy on INDUC- TION HEADS task (best of 5 seeds), varying [V] and d. N 4 4. Mamba-A 's per- formance (outlined) is equal or better than Mamba's 's (filled) and only hits 100% above the theoretical bound from Lem. 3 (black).
S4D N = 4, for Mamha N = 2x, and for Mamba-2 N = 8 In K. We use T = 100 and
| for all runs. The theoretical hounds model size for assembling the solutions proposed in Thm. 2 to 4 (black lines) separate reasonably well models that can achieve 100% accuracy (above black lines) from those that do not (below). In terms model size efficiency. Mamba-2 is better than Mamba, which in turn is better than S4D.
Proof sketch. We first present the construction using d = (xe[V)), and again apply JL Lemma to reduce to d (k log [V]) To overcome the limitation given by the lack of input selectivity in the SSM layer, we organize the hidden state along the embedding dimension only, and partition it so that each chunk holds the value associated with a specific key. At query time, the retrieval of the desired key-value chunk is done by leveraging the gating layer.
(F from a finite vocabulary 2 E V, the goal is report, for each 24, the token coming immediately after the latest previous occurrence in the input of token
Z4. That is, the output Yt must be Yt Iz(0)+1 where
j()) = max : =j < 1 = xz} (or a "blank" token, y4 = x, if ie appears for the first time).
Note the INDUCTION HEADS task is similar to MQAR, but with two significant differences. (i) There is no logical distinction between keys and values, so the model needs to identify the role of each token, but this can be handled by the short convolution, as we will show. More importantly, (ii) we need to retain information about only the latest previ- ous occurrence of a token, so the model should dynamically forget and remember information pertaining different to- kens. Since the latter memorization ability was investigated in Sec. 4.2, it is natural to leverage those insights in solution. To do so efficiently, we introduce a slight variation to the S6 layer: the Mamba-AT T SSM mixer, prescribing the following hidden state evolution
Remark 4. Our construction of the S4D-mixer relies both on gated convolution and the SSM recurrence. Arora et al. (2024a) provided a construction to solve MQAR based on gated convolution only. While such construction has better parallelizability, it requires the convolution kernel size to scale with the sequence length (in order to look back and find the matching keys). In contrast, our construction relies on size-2 kernels only, thanks to the SSM recurrence.
To validate our theoretical constructions, we train Mamba, Mamba-2, and Mamba-S4D models with varying embed- ding size d, on MQAR tasks with different number of key- value pairs K. The goal is to check how tight in practice are the theoretical bounds on model dimension derived above, and whether indeed trained models must respect them to solve the tasks. Results are reported in Fig. 2, where dashed curves denote our theoretical bounds and markers indicate empirical results. Notice our bounds in Thm. 2 to 4 are close to the empirical threshold between models sizes that can recover exact solutions or not, illustrating the tightness of our theoretical results. See additional details in App. D.3.
h, -e A O h, + 2, 2 B..
(15)
Comparing this to (14f), the only difference lies in the action of ((t:), which now varies along the state dimension N, rather than the embedding dimension d. While Gu & Dao (2023) hypothesized a similar performance for both versions, our findings reveal that the dependence along state dimen- sion is better suited to solving the INDUCTION HEADS task:
Lemma 3. There exists a 1-layer Mamba model with the Mambu-A SSM mixer (15) that solves INDUCTION HEADS with vocabulary V using embedding size 2 and state size N [V.
5.2. L-Layer Mamba-AT T Can Solve INDUCTION HEADS
The INDUCTION HEADS task was first introduced by Ols- son et al. (2022) to study Associative Recall capabilities in Transformers. Here we use the formulation from San- ford et al. (2024a): given an input sequence of tokens
Proof sketch The proof follows closely the MQAR con- struction, in that we leverage the matrix structure in the hidden state such that its columns are indexed by the keys
8

---

Understanding Selectivity in Mamba
and store the associated values. In the INDUCTION HEADS task, though, each token æ; acts as key in the (Ira, Fi-1) ) pair, and as value in the (2;-1, 4i) pair. To handle this distinction, we simply duplicate the embedding. and let the convolution layer correctly combine tokens information pairwise, so that after convolution each token encapsulates information both regarding a value and its preceding key. Moreover, we use (zx) to selectively erase outdated information, or to retain eurrently valid information, depending on the input observed. Pushing A (xx) + & flushes a previously mem- orized value, while (ax) + 0 preserves it. We refer to App. C.2 for the detailed proof.
Impact Statement
The goal of this paper is to improve the understanding of the Mamba architecture, specifically through analyzing the role of input selectivity. Our findings contribute to the ad- vancement of State Space Models, which may in turn further democratize access to Large Language Models, sharpening both the existing positive and negative aspects of LLMs. No additional societal impact is expected from this work.
Acknowledgements
We thank Pierre Ablin, Samy Bengio, Adam Goliński, Aryo Lotfi, Jason Ramapuram, Jonathan Sheaffer (in alphabetical order) for their helpful feedback and critical discussions throughout the process of writing this paper. We would also like to thank Alberto Bietti for pointers to the associative recall tasks.
Remark 5. To solve INDUCTION HEADS, Bietri et al. (2024): Sanford al. (2024b) constructed 2-layer Trans- formers relying length. On the other hund, we propose 1-layer Mamba composing a convolution and a variant SSM layer. Notably, this allows us to drop the PE and thus have the model size depend only on [Vi, and not the sequence length, improving upon the constructions for Transformers.
PE and with size scaling sequence
References
Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J., Rudra, A., and Re, C. Zoology: Measuring and Improving Recall in Efficient Language Models. In International Conference on Learning Representations, 2024a.
To demonstrate the efficiency of our Mamba-AT T variant on the INDUCTION HEADS task, we compare against the Mamba baseline, and report results in Fig. 3. For all model sizes considered, Mamba-A performs equally or better, demonstrating that selectivity along the state dimension (in the state matrix) improves Mamba's s ability to solve the INDUCTION HEADS task.
Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S., Zou, J., Rudra, A., and Re, C. Simple Linear Attention Language Models Balance the Recall- Throughput Trade- off. In International Conference on Machine Learning, volume 235, pp. 1763-1840, 2024b.
6. Conclusion and Future Work
In this work, we demystify the role of input selectivity in Mamba, showing its impact on approximation power, long- term memory, and associative recall capabilities. We prove that the S6 layer can efficiently represent discontinuous sig- nals and adaptively mitigate sensitivity decay. We also un- cover the role of other architectural components in Mamba, particularly convolution and gating. We present a mecha- nistic explanation of how Mamba solves memorization and associative recall tasks, with tight theoretical model size bounds matching empirical results. Our findings reveal op- portunities to further improve Mamba, such as alternative way to inject input dependence within the SSM state matrix.
Bietti, A., Cabannes, V., Bouchacourt, D., Jegou, H., and Bottou, L. Birth of Transformer: A Memory Viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
Chen, L.., Peng, B., and Wu, H. Theoretical Limitations of Multi-Layer Transformer. arXiv preprint, 2024. URL 02975.
https: http://xxxvv..ggabbs////
Chiang, D. and Cholak, P. Overcoming a Theoretical Lim- itation of Self-Attentio.. In Anmual Meeting of the As- sociation for Computational Linguistics, volume 1, pp. 7654-7664, 2022. doi: 10. 10.1866//1/22222el-oooggg.
Our current theory does not consider the aspects of optimiza- tion and generalization, both of which are interesting future directions to explore. Moreover, our analysis focuses on sim- ple associative recall tasks; extending it more complicated tasks such as --HOP INDUCTION HEADS (Sanford et al., 2024b). SEQUENTIAL FUNCTION COMPOSITION (Chen et al., 2024), and POINTER VALUE RETRIEVAL (Zhang et al., 2021) would be a natural next step. Overall, our work and proposed improvements add to the growing understand- ing of SSMs and could accelerate their development.
Cirone, N. M., Orvieto, A., Walker, B., Salvi, C.,, and Lyons, T. Theoretical Foundations of Deep Selective State-Space Models. In Advances in Neural Information Processing Systems, volume 37, 2024.
Dahleh, M., Dahleh, M. A., and Verghese, G. Lectures on Dynamic Systems and Control. MIT OpenCourseWare, 2011. URL https: //w////ww////////// mit edw.courses////////ooooooo 6- C-241j 1241i-dynamia e --w-s-st-eadd-oooeeeee----------

---

Understanding Selectivity in Mamba
griig=2001/990035555a00000001111 1185F9_MIT6_24111 htF_:___MT66_2211111eeeeetteoo
Li, Z., Han, J., Weinan, E., and Li, Q. Approximation and Optimization Theory for Linear Continuous-Time Re- current Neural Networks. Journal of Machine Learning Research, 23(42):1 2022.
Dao. T. and Gu, A. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. In International Conference Ma- chine Learning, volume 235, pp. 10041-10071, 2024.
Mallat, S. Group Invariant Scattering. Communications on Pure and Applied Mathematics, 65(10):1331-1398, 2012. doi: 10.1002/cpa.214111
Eckhoff, K. S. Accurate and Efficient Reconstruction of Discontinuous Functions from Truncated Series Expan- sions. Mathematics of Computation, 61 (204):745-76) 1993. doi: 10.1090/00025-577 1993- 195430-1.
Merrill, W., Petty, J., and Sabharwal, A. The Illusion of State State-Space Models. In International Conference on Machine Learning, 2024.
Gautschi, W. and Inglese, G. Lower Bounds for the Con- dition Number of Vandermonde Matrices. Numerische Mathematik, 52:241-250, 1987. doi: 10. (0.1007/BF01398 78.
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context Learning and Induction Heads. Transformer Circuits Thread, 2022. URL wttttttss///t/ttttttoottoo htt:://taanssoomseeeeeeee cuits ptiits.pb//2ue nd ndd-ndduttioo--headdd//
Grazzi, R., Siems, J., Franke, J. K., Zela, A., Hutter, F., and Pontil, M. Unlocking crace-Trackin in Linear RNNs Through Negative Eigenvalues. International Confer- ence on Learning Representations, 2025.
Gu, A. and Dao, T. Mamba: Linear-Time Sequence Mod- eling with Selective State Spaces. arXiv preprint, 2023. URL https::tp//......... /tte://rxxiii /aaiivivvvv. www.oorggaa... abs 2312 2, 00752.
html.
Orvieto, A., De, S., Gulcehre, C., Pascanu, R., and Smith, S. L. Universality of Linear Recurrences Followed by Non-linear Projections: Finite-Width Guarantees and Ben- efits of Complex Eigenvalues. In International Confer- ence Machine Learning, 2024.
Gu, A., Dao, T., Ermon, S., Rudra, A., and Ré, C. HiPPO: Recurrent Memory with Optimal Polynomial Projections. Advances in Neural Information Processing Systems, 33: 1474-1487, 2020.
Gu, A., Goel, K., Gupta, A., and Ré, On the Parameteri- zation and Initialization of Diagonal State Space Models. Advances in Neural Information Processing Systems, 35: 35971-35983, 2022a.
Sanford, C., Hsu, D., and Telgarsky, M. One-Layer Trans- formers Fail to Solve the Induction Heads Task. arXiv preprint, 2024a. URL https://///// htL.ttss///axxii.rrrg http://axxvvv..g/.. abs/ 2408 14332.
Gu, A., Goel, K., and Ré, C. Efficiently Modeling Long Sequences with Structured State Spaces. In International Conference on Leurning Representations, 2022b.
Sanford, C., Hsu, D., and Telgarsky, M. Transformers, Paral- lel Computation, and Logarithmic Depth. In International Conference on Machine Learning, 2024b.
Kamradt, G. Needle In A Haystack Pressure Testing LLMs. Github, 2023. URL https: htRps:/t //gtthub com/ gkamr [Accessed
Sarrof, Y., Veitsman, Y., and Hahn, M. The Expressive Capacity of State Space Models: Formal Language Perspective. Advances in Neural Information Processing Systems, 37:41202-41141 2024.
Needl halt:/LLMees__eeeelllrnnaaaaaa 2025-05-29].
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are RNNs: Fast Autoregressive Transform- ers with Linear Attention. In International Conference on Machine Learning, pp. 5156-5165, 2020.
Tóth, R., Felici, F., Heuberger, P., and Van den Hof, P. Cru- cial Aspects of Zero Orderrrrr Hold LPV State-Space Sys- tem Discretization. IFAC Proceedings Volumes, 41(2): 4952-4957, 2008
Kingma, D. P. and Ba, J. Adam: A Method for Stochastic Optimization. arXiv preprint, 2014. URL https://aa xiv. /tv/orr/aaaa abs/ 412....6666 6980.
Vershynin, R. Biigh-Dimnsiooal Probability: An Intro- duction with Applications in Data Science. 2018. doi: 10.1017/978 108231596.
Li, M., Zhang, X., Huang, Y., and Oymak, S. On the Power of Convolution Augmented Transformer. In Conference on Artifical Intelligence, number 17, pp. 18393-18402, 2025. doi: 10. 11109//aaa.vv//// v39i17..4409
Vetterli, M. Wavelets, Approximation, and Compression. IEEE Signal Processing Magazine, 18(5):59 2001. doi: 00111009999999900
10

---

Understanding Selectivity in Mamba
Wang, S., Li, Z.,, and Li, Q. Inverse Approximation Theory for Nonlinear Recurrent Neural Networks. In Interna- tional Conference on Learning Representations, 2024.
Zhang. C., Raghu, M., Kleinberg, J., and Bengio, Pointer Value Retrieval: A New Benchmark for Understanding the Limits of Neural Network Generalization. arXiv preprint, 2021. URL hattps://aaxttt / 2107. 12580.

---

Understanding Selectivity in Mamba
A. Approximation Power of Mamba
A.I. Notation
We typically use bold upper case A, B, C to denote matrices and bold lower case æ; to denote vectors or sequence. In Sec. and Sec. 4, the hidden state time is denoted as (t) (in the continuous setting) or h, (in the discrete setting). In Sec. 5, with a slight abuse of notation, the hidden state h2 E RdxN denotes matrix. We A, À denote the discretized versions of A. We use R. N to denote the reals and the natural number. The identity matrix is denoted as I, where the all-ones vector is denoted 1. We let dvag(v) be the diagonal matrix with diagonal filled with the vector We denote SoftPlus, ReLU, SILU as the corresponding pointwise nonlinearity and Linear as the linear layer. We let 1 be the indicator function, (s) be the heaviside function. We denote with @ the Hadamard (elementwise) product and & the Kronecker (outer) product. We use for embedding size, N for state size, and or for sequence length.
A.2. Mamba Approximates Haar Wavelets
In the following, we recall and outline the complete proof for Thm. 1.
Theorem 1. Consider a Haar wavelet vbj,k [0, R. and the Mamba basis function (8) at 1, 1, B,,a)dBB,.-.
Asa f (x,((, B5,k- Let it = concat (oncat(zg;: be the input signal augmented with time positional encoding. For any 0, there exist 3 Mamba basis functions 95.ks ff.k I7,k such that the approximation error
(2sa(s) -0 (s; + Ma (,) - (s;
smaller than e, Vs [0,1].
Proof. The Haar wavelet at scale with translation is is defined on the interval 8 E [0, 1] as
21/2
[0, (0,2-(+11)
())(s)
22/2
E (2-(1+1),27) 2-5)
(16)
otherwise
21/2
s S
(),((5) = -
-2/22
8 s 2- + 2 2- 1))
(17)
otherwise
See also Fig. A.l for an illustration. Notice that each wavelet can be represented linear combination of three shifted Heaviside functions, namely:
(s)(s) = 2 (s-22-1))) - -2H - (2-i2- (2--->+22-6tt)) (s
(s-
(18)
where H(s) = 1, 1 denotes the Heaviside function. The goal is then to show that Mamba can indeed reproduce a shifted Heaviside as one of its basis functions g5 by opportunely choosing its free parameters Aj. Aj and B₁ determining ghx
(m)) 95.6
(ma)
(mas
5000
9j.k
2//2
2//2
2
24/2
1
2-1
<huu
2 *
- + 2-0+1
2 (Ix
2//2
24/2
2//2
24/2
Figure A.I. (Left) Example of Haar Wavelets; (Middle-Right) Shape of the 3 Mamba basis 1 k 9y whose linear combination arbitrarily approximate the Haar wavelet y.
12

---

Understanding Selectivity in Mamba
To this end, we consider ID input. We set Ay,b and B, 24, which reduces the Mamba basis function
Hy (s; 1, = e S. d(
(19)
We then leave it to the last free parameter ((x()) perform most of the heavy-lifting, By extending the input signal æ to include the absolute time t as input, directly write A(x(t)) A(t) (this requirement will be relaxed in Prop. 1). Notice that we can use the Mamba basis function recover (scaled) Heaviside centered at any given instant l,, provided A (t) behaves as follows:
(s) (x(s) {
i < t
fo 0 if < in
(0;
- 24 2 (x).
(20)
else
else
Remember from (7) that ((s) = (s; wa) = SoftPlus (WaS - ba). Note that Softplus(x) O, lim Softplus(z) oc.
lim
(21)
Then, by choosing مa = wsst and pushing WA + -o, can get arbitrarily close to approximating (s)... Concretely, we have
(= 2 5 (r:hs -- - ) dr
2 H (s fx))
(22)
2- (k + 1)}. the resulting mamba basis functions 8y,k- k 9j.k 9 can approxi- mate arbitrarily well H(s 2- H(s (2- 2 2-(++))) H(s 2- (k 1)), respectively, which are precisely the shifted Heaviside functions appearing in (18). This allows us to finally write that Ve, Aie such that
By substituting & E
(1))
(5 + 9j 20 <c
(23)
Additionally, to remove the requirement of explicitly augmenting the Mamba input with time positional encoding, we show that Mamba can autonomously recover said time position by considering a constant input Ie = 1 Vt. This further relaxes the assumption in Thm. to using an additional Mamba layer receiving all-ones as input.
Proposition 1. A 1-layer Mamba can recover absolute time, given an all-ones input.
Proof. The proof follows by considering a Mamba layer (8) with 1, = Substituting this into (8), and providing an all-one input 2t 4I gives
1²(1)= /'- a(s)ds
(24)
This applies analogously to the discrete view, hE = - 2(3) =t
Before proving Cor. 1, recall the function approximation problem. Given function p E L²( 1]) (as function of time t e [0, 1)), and a set of bases parameterized by the SSM {gn. (e.g. Mamba or S4D), the goal is to project into the best size-N basis GN = (91....,99 } } that minimizes the approximation error.
S( (p(t) (rojjss())() dt,
argmin |p - projg, (p)llr= arg min GN
(25)
GN
where N denotes the hidden state size and
projex (0) = Σ (n,gn) gn) In-
(26)
In ENN
We note that the N basis GN is optimally chosen for the target function p (instead of using a fixed set of bases independent of p). The inner products (p. 9n) = f p(t)g dt) for the various determine the coefficients of the projection of the target function p onto the basis IN-
We are now ready prove Cor.
13

---

Understanding Selectivity in Mamba
Corollary 1.
For a piecewise-consiant function p(() with m 2 1 discontinaities, there exist N Mamba basis functions (8) such that the L² approximation error |pp - En=1 9 lillaa La is of order O 3m). On the other hand, S4D basis functions can achieve an upproximation error of ON
Proof. As discussed in Sec. 4.1, a S4D basis function (9) can represent any Fourier basis function fn = =f2an When approximating a target function with discontinuities, the Fourier coefficients decay very slowly (Eckhoff, 1993),
lim Ke, O(
(27)
This implics that using N Fourier bases (and thus S4D basis functions) for approximating discontinuous functions yields an approximation error O(N-1), proving the second part of Cor. 1.
By Thm. 1, N Mamba basis functions can approximate arbitrarily close any N/3 number of Haar wavelets. A set of N/3 Haar wavelets can achieve an approximation error of O(2 O(2-N/3m)) (Vetterli, 2001) when targeting a piecewise- constant function with 2 1 discontinuities. We provide concrete argument for completeness. Suppose the function is piecewise constant with m = 1 discontinuity. Then using an optimal set of N Haar wavelets H. C A.y.A.eeN..o.. -1, we can achieve an approximation error of
2/-1
min |pp - -projxy Rx
(p)ll1==
Σ
Σ
(p, 11x))
Σ
(p, (0))))))2
(28)
j=0 k=0
N-1
Σ
(p, )1MN)))
Σ
(e 23,
(29)
j=0
1=0
Σ ==N
(e, 03)
(30)
= -(
(31)
The second equality holds by noting that (p, }3,x) fo p(0)),, = if is constant in the interval [2- (k: 2)) Thus, for a piccewise-constant P with one discontinuity, only one wavelet at each scale j has (p, ()),(z))) # 0. Given that we can choose adaptively the best N wavelets, we pick the ones with nonzero coefficients across all = N scales. Then the approximation error is bounded the coefficient the highest scale, which has magnitude O(2-N/2 and squared error O(2-N), Similar analysis shows that for piecewise-constan functions with m discontinuities, the optimal adaptive wavelet basis with N wavelet functions achieves an approximation error of O(2-N/mn (2-N/m).
Remark 6. While the assumption of a piecewise-consten target function seems restrictive in Cor. 1, we note that this can be relaxed to a continuous target function by additionally assuming the input signals x are pracewise-constane The equivalence due to the intermediate value theorem. Note that piecewise-constant input signals are ubiquitous in language modelling tasks, where Ia takes values in a finite-sies size vocabulary.
A.3. Mamba can solve KEEP n-TH
Here we provide a direct application of the result from App. A.2, to explain the ability of Mamba to exactly solve the KEEP n-TH task when equipped with Positional Encoding.
Corollary 2.
There exists an 56 layer that solves KEEP n-TH on input augmented with time Positional Encoding.
Proof. To memorize the n-th position of piecewise-constane signal a(s) = Ei= 2; (= suffices to represent two Heaviside functions h1 = (s n2, b2 = H(s - (n - 1)) (recall H(s n) = 1g>n), by noting that hı - h2 is one the interval [n 1, n) and zero elsewhere, we have (rr, hn h2) = Tn, as desired.
14

---

Understanding Selectivity in Mamba
As shown in (22), the Mamba basis function can represent any Heaviside function. Specifically, let ws -00 and
00 t < n t2 n
(1([ SoftPlus(wst - wan)
(32a)
A2([xatt SoftPlus(wst- - wa(n 1))
x
I < n -1 12 1
(32b)
Suppose À = 1, B = 1 in the Mamba basis function (8). Using the above 1 A2, we can obtain the Mamba basis functions representing the desired Heavisides,
() = - S2 (r((r) -A-- H(s - 12).
(33a)
() =e S2 dz( -A--3, H(s - (n 1)).
(33b)
15

---

Understanding Selectivity in Mamba
B. Sensitivity of Mamba
In this section, we prove Lem. I and 2. These follow directly from the general sensitivity formula (11), which we derive in more details below: Given a generic, input- dependenenddnt recurrence relationship as in (4), we have that the sensitivity of the state time with respect its imput sequence the j-th instant 2 R is given by
@h a
Σ
II rrst1
A
B
0z1
dzz
==1
II A, \r=++1
Ba +
II Ar
B BBtz)
dz,
1'=++
JAj
II Ar mmo+l x
Batgdees +
A,
(Bg
dzz
prast1
(34)
B BBaz) + 0Λ,
A
B
dx,
dz,
r=j+1
s=1
rej
8 Jry
0A
II Ar P=1+1
(Bjr,)
+
II Ar
By
8zy
Since A, = - (Ai An (r)l) is diagonal and B; = (i)) Bn),),, we obtain from (34) the sensitivity of the n-th component of the hidden state he, = h, with respect to the imput 2y as
The dz,
II (r) r=j+1
0 (Ba(i)z,) 0 Oz,
dr, Σ II
In() B,((s)
(35)
r=s+1
With these ingredients, we are now ready prove the main results in this section. Lemma 1. Consider the hidden states arising from the S4D and S6 SSMs defined in (6) and (7). The sensitivity of the n-th component of their states at time I with respect to the inpul at time j < 1. is given by, respectively,
0h.4D
o(>n, Bay Izz) - -xa(t-(1-1))
0r,
0h
ö(D An, Buy &zs) e An Ei A(wx).
0x.
where ö(A, Xn, Bay Iz5s) depends on the input subsequence Tej, independent of the sequence length t.
Proof. Recall the Mamba discretization (5) chooses input-dependent SSM parameters as
X, (r) e An
Bn( Bn (zz))
(36)
Substituting (36) to (35), we see that the first factor becomes
II Xn(r) r=1+1
~ Minn A(xr)
(37)
while the second factor becomes
())());;)) + @m
j-1 II (r(rr) Armst1
Bn Bs)s)aa
dr,
sml
0 (z, ((x,) xz)) - -e- AA(1))) (1
Σ e da Mi A(xe) Br(x,) A(xs) Z..
(38)
Oz,
16

---

Understanding Selectivity in Mamba
We are interested in the behavior of the sensitivity for j fixed and 3o (i.e., the sensitivity of the current state with respect early input in long-range sequences). Notice that the second factor does not scale with (since is fixed), and can be bound in terms of the parameters defining the transformations (x) and A(x), as well as An and as such:
(388) S e(Bn, A, An, Tgj).
(39)
On the other hand, the first factor in (37) shows in general exponential dependence on t. Putting both together, we have
0h dr,
I E(A,, An; B,, - Eimsta A(x)).
(40)
The proof for S4D is immediate by taking A 1, (n())) B, for all t. resulting in
0hg
$ Si(An, Buy Bnjze -- (x((0(((+))
(41)
Lemma 2.
Consider the discrete-time 56 in (7) where B, [B,(E))...BB((TT)) E R.. Suppose there exists a
constant с 2 such that
lim An Σ (ar)) Sc.
(12)
Then the sensitivity of the n-th component of the state at time with respect to any input 2; is lower bounded by
Ohg
Ô dr,
lim
2 e
Bn (xx)) ((t.) x.
(13)
dz,
Proof. Recall for the scalar input sequence, each component of the Mamba hidden state given by
h
Σ e
Eim. A(xr) Br (x( A (xxx) Fa-
(42)
Then the condition in (12) implies
lim t0000
h
Bn (xx) A(x.)
(43)
Straightforward computation of lim
DM de,
completes the proof.
For comparison, we provide a similar sensitivity analysis for the softmax Attention layer as well.
Proof of Sensitivity of Softmax Attention.
Dyttea
0
Σ
T. W WKR.
WVX.
0x.
Re W WKST
as W. WKa (st Wo W Wva, Wv)
w. WK W
W WK Wyzs
Σ WJ Wxe.
w.W. WHE,
(44)
W. WKZ)
L8=1 W. WKE, Wvzs
W. Wx Wvz, Wv W. WK
W WKzr
WKK,
Dt
e W. WKR. WVT. )
W.+ Wx
Wvaz
r= 2h W Wx2,
Dr " W WK
17

---

Understanding Selectivity in Mamba
Even in this case, the second term be bound by constant C' in terms of [[ I: W
Wv and | (notice that
the rightmsst term, where fraction of two over appear, behaves weighted average of Wvs,, and hence does not scale with t). Only the denominator at the first factor remains function of t, providing
Ohgee 0x,
< C 1 t min c W Wiae)
(45)
18

---

Understanding Selectivity in Mamba
C. Proofs of Section 5
C.1. Proofs of Mamba Solving the MQAR Task
The MQAR Task
As a reminder, the MQAR task receives an input in the form
U1
U/
lis
Ugg
Ux
Ux ki, Ux
Ux
Kig Ux
Ux
kin
Ux
E
(46)
K kay-value pairs
shuflled queries (=keys), interwoven with noise
where the keys ,; are randomly taken from a key set of size x, and the values v, and the various noise tokens ex are randomly taken from vocabulary of size [V]. The goal to output sequence that, the location of each query, reports the value matching the corresponding key, that is
y
x
Vi
x
Vi
Vie x],
(47)
while the other components of the output (denoted with x) are ignored.
For the of the proofs, often make use of the ality. For completeness, we state and prove JL lemma below.
(JI.) Lemma reduce the embedding dimension-
ii.d. 1 Lemma 4 (JL Lemma). Given set of standard bases {ei...., eat and E (0,0.5), there exists random projection matrix M _ + IP where p O log d), Mij
vp Unif(#1} such that for all pairs (i,j).
KMe, Mej)l cc
(48)
Proof. Let /pMi. = Zi.j, where Zig is the symmetric Bernoulli variable. For any i,j E [d], we have
1 Zz Zij
1 p
(Mei, Mej) = (Ma.), Mbaal)
Z
P 1=1
where the last equality follows from the fact that the product of two independent Bernoulli variables is Bernoulli. In other words, the dot product of the two projected vectors is sum of A i..d. symmetric Bernoullis. By Hoeffding's incquality (e.g. Vershynin (2018, Theorem 2.2.2)).
P
pc²
P
21 2
exp
(49)
Thus
< 2exp ( pe
P[KMe; Meiil 2 c) 2P
7 2
(50)
p 1=1
Therefore, except with probability less than 2 2xxp ò E (0, 1). By a union bound, this holds for all (t) pairs of (e,, e,) except with probability
(- (-
pr"
2 ). it holds that (MM,, Me, I1 S e. Let = 3 In for some
d (9) exp (- (- < dl exp (- _ = exp(-21n
pE²
2xp
g =
(51)
Since there is a probability grater than - 52 that (Me; Mej)l S holds for all i, this guarantees the existence of such M by the probabilistic method.
Theorem 2.
There exists L-layer Mamba model without gating that solves MQAR with pairs using embedding size O(x + log [V), and state size = K.
Proof. The proof is divided into two steps: We first construct 1-layer Mamba model without gating that solves MQAR using standard basis vectors d = O(k + [V];; based on such construction, we then apply the JL Lemma (Lem. 4) together with shifting trick to complete the proof for O(n + log [V]).
19

---

Understanding Selectivity in Mamba
Step I: Construction With d = + O(VV). We consider one-hot encoding of keys and values. The main idea is (i) use the size-2 convolution to combine key-value pairs and filter out other uninformative pairs (particularly, preserve key-value pairs and discard value-value pairs); (ii) use appropriate input-dependent SSM matrices to store and retrieve the key-value information in the hidden state; (iii) use the output layer to project to the value embedding subspace. Crucially, we observe that the hidden state of Mamba at time & is a x N matrix (14f). Thus, we leverage this matrix structure such that each column of the state corresponds to a given key (N = k), and holds the value associated with it, that is:
h
jja
V/ J I Rd
(52)
We first describe how our proposed solution operates, and then prove that indeed our solution solves the MQAR task exactly,
Embedding The task of the embedding layer is to clearly distinguish values and keys. We achieve this by letting it map to orthogonal directions: let the embedding dimension be d = [ + K. and e; E IRd denote the standard basis vector. We impose
hgs + k; = k- Ca. and 02 J U2 = Cx++
(53)
for some parameters k v >> O.
We 2 convolution to combine information of key-value pair, as it is essential to the task solution. We remind that for a sequence (..., -t- ..). the action of the convolution layer with size-2 kernel with left padding to is given by
Convolution
c conv = (co Iz C1 ① Ct b)
Zg 0,
(54)
for certain parameters Co. c1, b E Rd and a nonlinearity' a = ReLU. For our construction, it suffices to pick Co =,1,, C1 1,,, and b = 61, with co, c1, 6 O. We want such nonlinear convolution to preserve information from the (k; fÿ) pairs, so we impose cak: > b and c b. We also need to prevent the value in (v; ky) from getting associated to the wrong key, but we want to preserve key information to be able to extract it at retrieval time, so we ask Cov & b and ei k >> b. The (k;, k;)) pair represents an edge-case, and we will show later how to deal with this. Finally, we want to ignore contributions from (v., Uz) pairs (as they refer to "noise" tokens), but we will delegate this task to ieput-selectivity in the SSM layer. To summarize, we need to satisfy
cok > b,
C > b,
kk>> b,
Cov S b.
(55)
A feasible parameter combination satisfying (55) is given by:
1, =
co = 1,
C =
and b- =
(56)
Let us show how such nonlinear convolution acts on the four different types of input pairs, under the assumptions (55):
conv ((k;; 12)) = ReLU(cuk; + quy b) = (cok - b)e, + (crv - b)ex+s
(57a)
conv ((v;, k2)) = ReLU(coU; + ca - b) = (cik - -bje,
(57b)
conv( (ak, k;)) = ReLUUccok; + tak, - b) = (cok - b)e; + (cik - b)e,
(57c)
conv((s5, uy)) ReeLU(Co + cr0y b) = (cir b)extj-
(57d)
Mamba SSM The SSM layer organizes a hidden state matrix where its columns are indexed by the keys and store information of the associated values. To this end, we let
A 0, 1, B(x)= = [ = W.
(58)
where We projects the input to the key embedding subspace. Consequently, the SSM layer covers three roles at once:
In the original Mamba definition, we have = SiL.U; for ease of illustration we consider instead a = ReLU, but notice that similar considerations still hold in this case, in light of the fact that (ir ReLU((z) for dooo: suffices then opportunely scale the inputs.
20

---

Understanding Selectivity in Mamba
1. Ensure that the right key get associated to the right column in the hidden state. This role is covered by the input that only keys in the convolved input pair
matrix (x)) WB ga. By picking WB W [Tx Ovvl] we are used for populating the hidden state. This choice of the input matrix also ensures that information from (v,, vj) pairs does not affect the hidden state.
Propagate information down the sequence without corrupting (i.e., memorization). The state matrix A(x) can take care of this, provided we fix it to the all-one matrix. This can be achieved by prescribing A = 0 in (14a). Notice that with this choice, the SSM layer simply performs cumulative sum h, Σ.= Bg 6 B,.
Retrieve the correct column when needed. This task is taken care of by the output matrix C(x) We As with the input matrix, it suffices have We W = 1 J: also in this case, when an input containing key is encountered, the ((x:) matrix will retrieve information only from the column corresponding to that specifie key.
Output The final output layer simply needs correctly classify us the most likely value from the retrieved vector yc- To this end, it suffices to consider
W.
(59)
I
as vn will return the maximum scalar product with y::
Let us illustrate with a step-by-step example how the above construction yields the required result. Consider a generic input sequence æ for the MQAR task. Under the embedding layer prescribed in (53), the (embeded) input will be in the form
x = ki, mg(1))
kig (5(22)
k; (5(w))
0x
kk,
RdX
(60)
so that the -th key is associated the (im))-th value. The order in which the keys appear at query time also random and denoted by the km subscript. Moreover, keys query time are randomly interwoven with value tokens, denoted as vx- After convolution, this input gets mapped to
Ci, +
Co
C +
C0
CLAbi +
Co +
C1,ck, +
CLAck
à
+
0
((()
C11))((2))
e((ta)
0
0
(61)
where we denote Co.A ogk - b, C1,E cık b, Cr. civ - b to slim notation, and separate the components of à into key-related (top) and value-related (bottom) for illustrative purposes. Notice how, for the initial part of the input, distinct keys (the various c,,, at the top) are only associated with their respective values (the various Cj(i.)) at the bottom), or with 0. Before moving on to the SSM layer, let us remind its action (14f):
H I
y: (A((BOII) o 14-1 + (Ar O xx) B.) C
Σ
AolAmoin) (Ø. B.) (0..B.) C. m=
(62)
0.4...m
y: -h So B C. Σ (a. B,) Ct
Notice how the hidden state naturally admits a matrix structure h, R due to the outer product a, B, whose columns are updated by 2t, where B, determines which columns are affected. If B, x e; (i.c., only nonzero at component i), then o, will contribute to only the i-th column in h. Similarly, taking C, x e; ensures that only the i-th column from the hidden state is retrieved. This is precisely what we achieve with the choice outlined above. In light of this, the hidden state after the initial part of the input (where the key-value pairs are listed) admits the following form time 2x
kee
kez
kex +
h2
+
i = Co,k with i k= cū + c
(63)
(ef(1)
and remains unchanged until the first key kb, is encountered at query time, as B(x) = B( v,)) = 0 for all the tokens in between. For each key k; encountered at query time, C(x:) then takes care of extracting the corresponding i-th column of
21

---

Understanding Selectivity in Mamba
h, which holds vector proportional the embedding of its associated value, e)()- Finally, the output matrix computes scalar products of all the values embedding with the extracted vector, which will then be maximum for ej(1) thus accurately solving the MQAR task.
There is an edge case to this construction, which occurs when, at retrieval time, two keys appear adjacent to each other in a key-key pair (k; k.). In this case, K represents the actual query, while k; acts as noise: the convolution will fact contain information from two keys, implying that C(x) will have two nonzero components at that point. However, with the correct parameter choice, we can have the query information dominate the noise, and still recover the desired solution. We have by (57c)
C'conv (k; ki)) (col: b)e, (c. b)er.
(64)
When tested against the hidden state at that instant, then, the output matrix will return a linear combination of the values associated to the two keys:
(kconv(ks ki)) h4 = (cok- - b) (ke, + +ee(o) + (ezk- -) b) (kee + be5())
(65)
To ensure the key-value pair of ky dominates that of k. we need to impose
cyk: - b < cak - b = 11,
(66)
which is already satisfied by the parameter choice (56).
Step 2: Dimensionality Reduction to d = Ox + log [V). Having proved the construction with d= Ok + [VD, we now apply JI. Lemma to reduce the embedding dimension and suitably adjust the Mamba architecture weights to ensure the desired output. Concretely:
Embedding We use the same key embeddings as above fei...., e,) while reducing the value embedding dimen- /e,) 1V and 1
sionality to satisfy almost-orthogonallyy. By JL Lemma (c.f. Lem. 4), given the value embeddings {e,
(0, 0.5), there exists Me 1R + Ip where p = Olog [V]) and each of its entry is in {- 3/ vp¹ / } such that (M.e,, M I for any J1 # j2. Let M I M, E R(*+p)x(*+IVI) be the direct sum of the identity matrix preserving the one-hot key embeddings and the JL matrix projecting the value embeddings. Let the normalized value embedding be ū, = Me, Let the shifted value embedding be
Bj = Me, + 310 11] = >
EMetns + 8)eita E RR**,
(67)
1=1
where each component of fy is in the range
| + B,
1 vP
+ B = =[umia: Umax] Note that by letting 3 >
vp' we
can ensure all components of bj are nonnegative.
Convolution We use a size-2 convolution as (54), to retain the value from (k;, 0y) pairs and discard the value from (v,, k2) pairs, Ideally we want
conv(k,, 02) (e,k( + cuv b) x (co- b)k; + (aI b)o,
(68a)
conv(v,, k;) ReLU(coU; + ea ky b) 1 span(( ui ........ 1lv1)).
(68b)
Given the shifted value embeddings, we must impose the following constraints to achieve (68):
k > b, Ci Umin > b,
Umaax S b.
(69)
A feasible parameter combination satisfying (69) is given by
A 2 - [Umin, Urnax] C [1, 3] since p 2 1), Consequently, we have the desired convolution outputs
k 10,
co 1,
C == 10, and b= (70)
conv(k,, vy) = (cok- 1 -b)e++ +
(ciMMtnn + +) - b) e = (cok- - -b)e + 01 + (c,8 - b) Σ
es4x
(7la)
i=1
conv(v;,kt) kt)) (cik- - -blej.
(71b)
22

---

Understanding Selectivity in Mamba
Note that the edge case for conv ky) is taken care of since co < c1, while the (v,, v2) pairs will be ignored in the SSM layer, following the same argument in step
Mamba SSM The SSM layer the same in step (58), propagating information through the hidden state (c.f. (63)). At the query time for key k; i = 1...., K. the output y/ ((ogk 6)2 + (c1 - 6)3) e; + (cok b) (c: (C)) + +)
contains the (scaled) normalized value embedding shifted by constant vector
Output The output layer undoes the shift and classifies based the normalized value embeddings. Recall M, RIV is the JL matrix from the value embedding projection. We set the output linear layer with the weight matrix W. [o E R ()+(<>p)) and the bias vector b. = -cok - ss. s. The final output is given by
Woy + b. (cyl b) (c18)() + 8) (col: b)W,s . (cyle b), M.ej(-
(72)
which yields the maximum component j(i) since (Moe)(t) M,e;) S e for any 1 * j(i) by JL Lemma. This completes the proof.
Theorem 3. There exists a 1-layer Mambu-2 model without gating that solves MQAR with k. pairs using embedding size d Ollog k + log V), and state size N = log K.
Proof. The idea is to execute the Mamba solution outlined in the proof of Thm. 2, but in a leaner manner due to the additional degrees of freedom in Mamba-2: particularly, we leverage the fact that the convolutions for the value (14b), key (14d), and query (14e) can be chosen independently (instead of using the same convolution in Mamba). The proof is divided into two steps: We first present the construction using standard basis vectors with d = x + [V]; We then reduce the embedding dimension by applying JL lemma.
Step 1: Construction With d Ok + [V]).
Embedding Same as Mamba The role of the embedding layer is to clearly distinguish values and keys. We achieve this by letting it map to independent directions: let the embedding dimension be d = |V| + K, and e; E Rd denote the standard basis vector. We let
k; - G. = E; and
62 ト U, == Ex+i
(73)
Convolution Differently from Mamba that uses the same convolution kernel to compute the SSM input and parameters B, C, Mamba-2 has the additional flexibility of using three independent convolutions for computing u, B,C C (see details in Tab. 2). We now exploit this flexibility by setting:
a B conv (zt-1, xt.) o(co O Fz-1 C₁ xt b) = 2r-11
ã; = conv (2 = Dt.
(74)
c c
Ze.-
Consequently, the output from conv shifts the input sequence to the right by one position, whereas the outputs from conv,, conve are the same as the input.
Mamba-2 SSM Same as Mamba The role of the SSM layer is to associate key-value pairs, propagate information through the state, and retrieve the correct value given a query(=key). Unlike Thm. 2, the convolved input (74) contains only key or value information but never mixes them. A simple choice is to set B, C' as the identity matrices, but this requires the state size be the same as the embedding size N = d. To further reduce the state size to N K, we use the same choice as Mamba (58) by letting
A = 0. B(x,) = C(..) = [T. W.
(75)
Output Same as Mamba Even in this case, as classifier suffices to pick
v
W
(76)
21 [
23

---

Understanding Selectivity in Mamba
With the definitions above, we can simplify the outcome of the Mamba-2 layer application. This in fact reduces to
y: = c B) c(t.) = Σ (xs (Waos-1) WaD = Σ w. (z) (2,___1 (
(77)
where 2y = 0 is the zero padding vector. Note that the output yr for t 2 2x contains the desired key-value association, Te (WWWW1,s-1 = k for s = 2, ,,......tt 2x.
At query time where x, = kü thanks our construction of the Embedding layer, (k:, for all 8, except when 2x k; in that case have instead (k;; value x, is precisely the one we are seeking, i.e. the value embedding Uj() immediately following the key matching the corresponding query kū Thus, y/ Uz(a)- Now, applying the output matrix amounts to computing scalar products between (y( and all possible value vectors, which will return only at the desired value per the orthogonal embedding construction.
1. Whatever the index at which this occurs, the associated
Step 2: Dimensionality Reduction to = O K + log VI). We now reduce the embedding dimension d = [V| + K to d Ollo [ + log xi). To this end, we apply JL Lemma (c.f. Lem. 4) to construct nearly orthogonal embedding vectors, while keeping the convolution, SSM, and output layers the same as step 1. By JI. lemma, fixed (0, 0.5), there exists a for all i. [d]. We apply JL lemma separately for the key embedding subspace (with JL matrix Ms E Reog and the value embedding subspace (with another JL matrix Me E I [V) ). We collect the final JL matrix via a direct sum, M- = M₁ 1 M.. Let K == Me,, Vy = Mex+> Then for i i j. we have (As,; ky) I C, (v;. U2) S c. On the other hand, (k:, kj) = E M2 s 1. Similarly, we see that (oy, vz) Rs 1. It remains to show such low-dimensional embeddings, followed by convolutions, SSM, and output layer, yields the desired solution. The action of the convolution layer on the low-dimensional embeddings is the same as step 1, shifting the attention-keys by one position to the right while keeping the attention-queries and attention-values unchanged. Also the action of the SSM layer is the same step 1 (77), where the hidden state is a sum of key-value association matrices Dyte) TT except with embedding dimension Olog V + log ks). At query time, upon encountering k; the retrieved output is Ye ()(() + el ((ij)))) B1) for < 1, in light of the fact that the low-dimensional embeddings are nearly orthogonal by construction. Then, applying the output matrix W. yields the maximum component at j(), as desired.
matrix M R for p = Olog d) such that (Me, Mey)l
Remark 7. The above construction with N log works for generic imputs where the values are drawn randomly from the vocabulary with sufficient size. Such construction may fail in the adversarial case where most values are the same and the number of keys k. is large. Concretely, suppose the input is [K1, U1. k2, V. V....,k.. ke,, v), such that all the values at time
l = 4, 6,..,, 2x are the same loken v, and in # v. Then at the retrieval time for the query token ki, the signal value is oy, whereas the noisy values from other keys contribute to E(K- - -Lyo. Then if e ( 1, the model might fail to retrieve the desired value. This can be counteracted by decreasing e; as per JL Lemma, though, this might come at the cost of scaling the embedding dimension - without however impacting its logarithmic behavior (remember p = O d)).
Theorem 4. There exists a 1-layer Mamba model with an S4D mixer that solves MQAR with pairs, using embedding size d = OO log VV), and state size N 1.
Proof. The proof is divided into two steps: first construct a 1-layer Mamba model using S4D SSM layer, that solves MQAR using d = (x(VV)); we then apply JI. Lemma with a shifting trick to complete the proof for d = O log [V)).
Step 1: Construction With d = (xx On a high level, the idea is to organize the hidden state of the SSM in chunks, each collecting a vector representing the value associated to a specific key, similarly to the proof in Thm. 2. However, unlike the Mamba layer, the S4D layer does not have access to input-dependent matrices Bax), C(xx), implying that the N columns of the S4D state are the same up scaling, and hence cannot encapsulate additional information regarding the input. We then decide to work with single-column hidden state, and instead partition it along the embedding dimension d. Ideally, we want the hidden state before retrieval to be long column vector,
h4
2
U22
(78)
This can be achieved by specifying each component of the full Mamba architecture as follows.
Embedding The goal of the embedding layer to organize keys and values a form that is suitable to assemble a
24

---

Understanding Selectivity in Mamba
hidden state in (78). More in detail, let key k; and value ey be mapped to, respectively
k + k; k 10,
0
10,...,01 0 ...
(79)
[
+ 10..... 1
0
10,.... 1 0].
(80)
K|VH+i
e
e 1
(81)
for some fixed parameters A > 0, O. Notice that, in light of this, any combination k; + uy can form a dictionary: it has a maximum value at a component uniquely defined by the pair i.j.
Convolution The (short) convolution layer is responsible for filtering out irrelevant information from the sequence, and retaining only the one pertaining (k,, vý) pairs. To this end, we limit ourselves to convolution with kernel size 2: this acts on any pair (I,, x++1) of the input sequence by mapping it to
(at, 1e++) ) + (Co @ at Ci 1e+1 b)-
(82)
By carefully picking the parameters k, u, Co. Ci2 we can ensure that only (k,, uz) pair "survives" the operation, and everything else gets mapped to the null vector. This for example can be achieved by setting k 10, 1, 10 11 C1 1,b kco- With this choice, see that the convolution maps
(As, v2) + .0.....0 0
10..... 1 Of 10....00 OTT
(83)
[ii++
0
e,
...
0 1T
(84)
(k;, k2) + (0,....00 T
(85)
(w,, ky) - O0.
(86)
(vi, 0j) (0,
(87)
S4D SSM The SSM layer simply needs to accumulate and propagate the combined values down the sequence. We remind that from (5) and (6), the output of S4D is given by
540 h
A - B.
(88)
s=1
We make use of a "trivial" SSM where A = B = 1, resulting in hidden state which, after collecting the initial (k,, U2(1)) pairs, is constant in the form:
h = [0.....
0/0,
1 ,....
J0..... 1 O NIVI+
|V e]_ LT.
(89)
Gate The role of the gating mechanism (14g) to retrieve the part of the hidden state which refer to the requested key. The gate branch acts on linear transformation of the original sequence: by picking this transformation as the identity. we ensure that, when key encountered, only the corresponding value is retrieved from the hidden state, in fact:
U = k ② h4 = ....... ... .0...00......................... 10....,
,/.,,,00,...................,... 0 .....
(90)
Output The final output layer simply needs the retrieved vector yı against all values vyt only the correct one will return scalar product different from 0. suffices to pick
v
e
e
SAD
(91)
V
[
25

---

Understanding Selectivity in Mamba
Having proved the construction with O(VV), now apply JL Lemma to reduce the embedding dimension and suitably adjust the Mamba architecture weights ensure the desired output. Concretely:
Step 2: Dimensionality Reduction O(x 1)).
Embedding We use JL lemma to identify = | almost-orthogonal vectors within a p O(log( + 1))- dimensional space. Namely, there exists a random matrix M TR + Rp such that any two (different) vectors in the set {Me....... Me vI, M are almost-orthogonal, in the sense that (Me,, Me,)) < for some small e E (0, 0.5). Without loss of generality, we ensure that the last of these vectors parallel the all-one vector 1p notice this always possible via opportune rotation, which does not affect the scalar product of the recovered vectors (and hence, their almoss-orthooonnlltt)) We let the embedding layer perform the following map
k; + k; = k. ........... ......,0
1.....1
00 0O.....0.....]
OI....]
E ERP,
1 i= 1...K
(92)
pi-p(i+1)
1 - v2 = [[Me; + +1)
I(Me: + TT E ER",
i i= i=1...[V IV.
(93)
where M _2 + IR, Mi E - { the (rotated) projection matrix recovered with JL Lemma, and that each component of the value embedding is nonnegative, and falls in the range of 1
vF
J
Vp
Convolution We use size-2 convolution in step requiring it to retain only the (k;, wz) pair information while sending other pairs (v,, k2), (v, wj), (ki, k;) zero. To this end, we let
B = 1, 10,
cp 10,
Ci 1,
b kco-
(94)
Note that by the choice of B and p 2 the range of components in is limited to [0, 2]. Consequently, we have
conv(ks, vz) = ReLU(cok; + ci b) 0 conv(u; k,) conv(k, k,) conv(u, 2z)) =
(Me, + 31)
1 O l
(95)
(96)
S4D SSM The SSM layer proceeds similarly as the construction in step 1, yielding the state at > 2x as
hr [Me,, + (MM)), B1)T
[(Mej, + +1
(97)
Gate and Output Also the gating layer and the output layer proceed similarly as the construction in step After gating, when a key k; encountered, have
J = h1 ① K4 (0, ...
(Mej, + +1)
10,... .Q.....
(98)
pip((+1)
And finally, after applying the output matrix, we obtain
c cB
(Mei)
(Mei) : : (Me(v))
WDSAD Ú:
(Mej +131)
1 1+
(99)
(
E eB
in light of the fact that both (Me,)T. Me; c if otherwise (Me)) Me; e per JL construction, and (Me;)T T 1 < per the assumption that Meivi+1 is parallel to 1p. This allows us to recover the correct value, completing the proof.
26

---

Understanding Selectivity in Mamba
C.2. Proofs of Mamba Solving the INDUCTION HEADS Task
The INDUCTION HEADS Task As a reminder, for the INDUCTION HEADS task, the input is a sequence of tokens (1y.... : Fy.] from finite vocabulary V; The output is sequence of tokens ly...... 3/l from the augmented vocabulary [x{x) where y4 equals the imput token right after the latest previous occurrence of the input token 24, ie., 3/4 2
where j(i) = max{j : < i, = ] otherwise yi x. An example input and output drawn from the vocabulary
V {1, 2, 3, 43 with sequence length 8 is shown below:
t 1,
2.
3,
4,
5
6
7,
8]
a 2, 1
3,
2,
4
3
4]. 2,
2,
y/ X, X, X, 1,
X,
4, 3]
Note that the input token 2 appears three times, at instants t == 1,4,7. Thus, at t = 7, the latest previous occurrence of j(7) = 4, which yields the output y/7 = (7)(+)+1 = F4+1 = 4.
Lemma 3. There exists a 1-layer Mamba model with the Mamba-A SSM mixer (15) that solves INDUCTION HEADS with vocabulary V using embedding size d 2 and state size N [V]
Proof. We follow a procedure similar to the MQAR construction, in that we leverage the matrix structure in the hidden state such that its columns are indexed by the key token and store the associated value token. However, differently from the MQAR task, in the INDUCTION HEADS task there no distinction between the key and value set, but rather all tokens are drawn from the same vocabulary V i.c., each token 2; acts as key the (I;, 21+1) pair, and as value in the (2;;1, pair. To resolve this, we use doubling-embedding trick in Mamba, together with suitable choices of convolution. Moreover, the INDUCTION HEADS task requires finding the latest previous occurrence; we will achieve this by leveraging the input-dependent state matrix.
The main idea is to double the embedding size in Mamba, which enables the convolution layer to perform concatenations of the adjacent embedding pairs. Concretely: we let the state size N = [V. and the embedding size 2/ We design the architecture as follows.
Embedding We use V standard basis vectors to embed the vocabulary V = 2,,2, .. IV],,
Cvs Cv:
Rd
RIVI
(100)
Convolution We use 2ize-2 convolution (with left- fett-paddin 0) combining the pair (z;-1, xx) by summing the first |V]-dimensions of -;-1 with the last [V]-dimensions of æ; effectively concatenating (ez,-1 Cx We let
0 1
i conv(x cu x;-1 + C1 O Zt where Cu
0 c1
(101)
Note that describe the proof for linear convolution here to simplify notation, but it holds also for nonlinear convolution (14b), noting that each embedding vector and the convolution weights co, Ci are nonnegative, effectively reducing the nonlinearity = ReLU (or SiLU) to be the identity map. The same reasoning applies to the design of B, C, as we discuss next.
Mamba-AT SSM Since the convolved output के contains x;-1 in its first [V] dimensions and x; in its last [V] dimensions, we choose the state matrix A and the input matrix B to depend on the first V dimensions of & (i.e.. extracting the key), and the output matrix C to depend on the last [V] dimensions (i.e., extracting the query). To this end, we let A- - E N ws 0, and
A(x.) = Soft SoftPlus(Linear(tt)) where Linear( = wAlM |0 E N A(a;)= := e (2.) = - 10A(a)) E IRdxN = R
(102a)
(102b)
((t,) == Linear( (+.) [|v |0 [ E R.V.
(102c)
C(+.) = Linear(ā;) [0μv] [[v]] & E V.V.
(102d)
27

---

Understanding Selectivity in Mamba
Output W.
[0
Mvvl
Rixx
We now show the correctness of such construction. Consider the generic input and output sequences:
t 1,
2.
3.
4.
5,
6,
7,
8,
U2 01
03
22:
14
13
U/
U4.
(103)
y X
X X
01
X
02:
V3:
.....
After the embedding and convolution layers, the SSM input is sequence of d-dimensional vectors for d = 2
Cvi
Cv3
Ce2
Cea
Cvy
Cez
E BRII
(104)
eva
Ev1
Exaa
Er2
Evaa
Evg
eus
where è;; stores the (#1-1, Jii pair.
The action of the SSM layer organizes the hidden state matrix of size d x N = 2| x outer-product with é, followed by retrieving the desired column via the output matrix C(ă;). Now by the choice of SSM parameters, we have
[V] by the input matrix B(t;) taking
( C(a)
Evs
Exi
Ex
Evs
Exi
E R
(105)
evz
Evg
eva
Exa
evg
Cvs
where we stack them together to visualize that ( amounts to shifting C(t)) to the right by one position, due to the design of convB.
We now verify the desired behavior in the SSM layer. Suppose temporarily the state matrix A 1 E R Then the hidden state at time S would be cumulative sum,
h
B(::))
Ez, Ex,
(106)
Thus, the j-th column of the hidden state matrix would store the sum of all è; where the key B(t,) = x; = ey. Yet the INDUCTION HEADS task requires storing the latest associated value only (not all associated values). To this end, we leverage the input-dependence of the state matrix, and particularly of ((x2). Recall the Mamba-A T layer is given by
h. = A0(IS((E)) @ h,-1 +2.B(..
(107)
y. = h. C(t.).
(108)
We design (d+)) E I = g/Vl such that when the input contains the key information, the corresponding key column in the state is erased (while the other columns remain the same). Without loss of generality, suppose B ( ej. By the definition of B (102c), this implies that [ [v] O = e.. By the choice of ws > O and the definition of A (102a), we have
( = Soft [La 0,]2.) wAej.
(109)
Therefore (102b) yields the state matrix as
A, c -10d(e)) (10(w)))))))))mm = & exp [0.
WA OJ "o= 18 (1 RdN B
(110)
which has an all-zeros j-th column and all-ones columns elsewhere. Consequently, the j-th column of the hidden state h,[ts]] is erased by the action e (.. hg-1, and then updated with the current input containing the latest value by the action d., X B ((x,), solving the KEEP n-TH task in Cor. 2, in which we have Mamba approximate a Heaviside by tweaking A A(xz). so to erase information from all tokens before given one. We also see that such selective erasure works consistently well for long sequences when t + oc, since it preserves all other columns (except the j-th one) by setting A(u.)[[ = 0 for I # j. and thereby satisfying the condition in Lem. avoid sensitivity decay.
(..), ), as desired. We remark that such erasure operation is akin to the construction of the S6 layer for
28

---

Understanding Selectivity in Mamba
Finally, the SSM output is given by vs h.C. ez, which retrieves the 2. T-,th thh column of the state that stores the token immediately after the latest previous occurrence of Fs, i.c. y:
Exga)
We then apply the output matrix
W Woys 320++ to obtain the target value, which completes the proof. C200)
29

---

Understanding Selectivity in Mamba
D. Additional Experiment Details
D.I. Training Details
For all experiments, unless otherwise noted, we train with the Adam optimizer (Kingma Ba, 2014), for 600 epochs using an initial learning rate 7 = 0.03 and cosine annealing down 1) 10-0. The training is composed of 105 randomly generated samples with a fixed seed and the batch size 16, which results in up to 3.75 10 gradient updates. We perform early stopping if the validation loss reaches below 10-6 or if six hours have elapsed since the beginning of training. The validation set and test sets have 10x and 10a samples respectively and are generated with the same function but using different seeds. Reported accuracies are always obtained from the test set after the last epoch of training.
D.2. Task KEEP n-TH
In this section, we provide additional details and ablation of the KEEP n-TH task used for Sec. 4, with experimental set-up and partial results reported in Tab. 1.
Model All the results Tab. A.I - Tab. A.4 use 1-layer models. The MAMBA and S4D models (with or without PE) are simplified architectures, which consists of embedding layer, SSM layer, and output linear layer, without convolution nor gating from the Mamba mixer block. The simplification is intended to investigate the role of the SSM layer alone (i.e. S6 versus S4D), without confounding factors from other components in the mixer block.
Experiment Set-up For each input sequence vocabulary with size [V| = 128. The target output is the n-thh token in the input sequence, i.c. = for n < t S T. The predicted output at time = n, . are taken to compute (cross-entropy) loss for training, and accuracy for evaluation.
(ai,...., TT), we draw
Uni randomly from
Discussion When equipped with PE, Mamba manages achieve 100% accuracy on KREP n-TH, regardless sequence length. This is thanks its ability dynamically adjust (via (xx)) for how long the hidden state retains memory of the target token (in position = 5). On the other hand, S4D is lacking such ability, and already fails at the task for T = 20 (Tab. A.1 and A.2). Then again Transformers do not need to retain memory, they can look back to the whole sequence each step, in light of their attention mechanism, and have no issue solving the task for any T. When removing PE from Mamba, however (Tab. A.3 and A.4), the model loses its way to discriminate the specific token that must be retrieved, and performance drops to that of S4D, as expected.
Table A.I: Ablation KEEP n-TH: MAMBA+PE, S4D+PE, TRANSFORMERS with varying sequence length T -0,,00 20, embedding dimension d, and state size N.
T=10 # epch.
T=20
acc.
# prm.
acc.
# epeh.
# prm.
d=8, N=8
1.00 (0.00)
107
2.2k
1.00 (0.00)
37
2.2k
d=8, N=32
1.00 (0.00)
241
2.8k
1.00 (0.00)
39
2.8k
d=8, N=64
(0.00)
161
3.5k
1.00 (0.00)
47
3.5k
d=32, N=8
1.00 (0.00)
9.2k
1.00 (0.00)
15
9.2k
MAMBA (+PF)
d=32, N=32
1.00 (0.00)
19
11.5k
1.00 (0.00)
14
11.5k
d=32, N=64
1.00 (0.00)
88
14.5k
1..0000000
19
14.5k
d=64, N=8
0.99 (0.00)
553
18.7k
0.99 (0.01)
208
18.7k
d=64, N=32
1..0000000
23.3k
(0.00)
23.3k
d=64, N=64
(0.00)
345
29.4k
1.00 (0.00)
12
29.4k
d=8, N=8
0.99 (0.00)
600
2.1k
0.43 (0.04)
600
2.1k
d=8, N=32
0.99 (0.00)
600
2.5k
0.46 (0.05)
600
2.5k
d=8. N=64
0.99 (0.00)
600
3.0k
(0.03)
600
3.0k
d=32, N=8
(0.00)
8.7k
0.72 (0.01)
600
8.7k
S4D
(+PE)
d=32, N=32
0.93 (0.01)
600
10.3k
0.72 (0.01)
600
10.3k
d=32, N=64
0.93 (0.01)
600
12.4k
0.73 (0.01)
600
12.4k
d=64, N=8
0.79 (0.00)
600
17.5k
(0.01)
600
17.5k
d=64, N=32
0.79 (0.00)
20.6k
0.14 (0.00)
600
20.6k
d=64, N=64
0.79 (0.00)
24.8k
(0.00)
600
24.8k
I=, d=16
1,00) (0.00)
16
5.8k
1.00 (0.00)
19
6.0k
TRANS-
1=1, d=32
1.00 (0.00)
12
15.1k
1.00 (0.00)
12
15.4k
FORMER
1=1, d=64
1.00 (0.00)
10
42.3k
1.00 (0.00)
42.9k
30

---

Understanding Selectivity in Mamba
Table A.2: Ablation on KEEP и-TH: MAMBA+PE, S4D+PE, TRANSFORMERS with varying sequence length T 30, 40, 50, embedding dimension d, and state size N.
T=30
T=40
T=50
# epch.
# prm.
epch.
prm.
epch.
prm.
d=8, N
0.94 (0.04)
456
2.2k
0.85 (0.15)
325
2.2k
0.22 (0.15)
600
2.2k
d=8, N=32
1.00 (0.00)
285
2.8k
0.58 (0.37)
416
2.8k
0.28 (0.10)
600
2.8k
d=8, N=64
1.00 (0.00)
3.5k
0.59 (0.31)
420
3.5k
0.98 (0.02)
346
3.5k
d=32, N=8
1.00 (0.00)
20
9.2k
(0.00)
50
9.2k
(0.00)
241
9.2k
МAMBA (+PE)
d=32, N=32
1.00 (0.00)
20
11.5k
1..000
21
11.5k
(..00)
225
11.5k
d=32, N=64
1.00 (0.00)
19
14.5k
(0.00)
23
14.5k
1.00
130
14.5k
d=64,
(0.00)
18.7k
(0.00)
164
18.7k
(0.00)))
600
18.7k
d=64, N=32
1.00 (0.00)
24
23.3k
(0.00)
66
23.3k
(0.00)
600
23.3k
d=64, N=64
1.00 (0.00)
29
29.4k
10,000
226
29.4k
(0.01)
600
29.4k
d=8, N=8
600
2.1k
0.03 (0.00)
600
2.1k
(0.00))
600
2.1k
d=8, N=32
0.46 (0.05)
600
2.5k
(0.01)
600
2.5k
(0.00)
600
2.5k
d=8, N=64
0.43 (0.03)
3.0k
(0.02)
600
3.0k
00.00)))0
600
d=32, N=8
0.72 (0.01)
600
8.7k
0.09 (0.00)
600
8.7k
0.08 (0.00)
600
8.7k
S4D
d=32, N=32
0.72 (0.01)
600
10.3k
0.09
600
10.3k
0.09 (0,00)
600
10.3k
(+PE)
d=32, N=64
0.73 (0.01)
600
12.4K
00.000)000)
600
12.4k
(0.00)
600
12.4k
d=64,
(0.01)
17.5k
(0.00)
17.5k
0.09 (0.00)
600
17.5k
d=64, N=32
0.14 (0.00)
600
20.6k
0.100)
600
20.6k
(0.00)
600
20.6k
d=64, N=64
0.14 (0.00)
600
24.8k
0.09 (0.00)
600
24.8k
0.09 (0.00)
600
24.8k
11, d=16
1.00 (0.00)
16
5.8k
(0.00)
23
6.3k
1.00 (0.00)
17
6.4k
TRANS- FORMER
1=1, d=32
1.00 (0.00)
12
15.1k
1.00 (0.00)
12
16.0k
10.000
13
16.4k
1=1, d=64
(0.00)
42.3k
(0.00)
44.2k
1.00 (0.000
44.9k
Table A.3: Ablation KEEP n-TH: MAMBA and S4D with varying sequence length T 10, 20, embedding dimension d, and state size N.
T=10
T=20 # epch.
# epch.
# prm.
prm.
d=8, N-8
0.20 (0.01)
600
2.3k
0.05 (0.00)
600
2.3k
d=8, N=32
0.18 (0.03)
600
2.9k
0.05 (0.00)
600
2.9k
d=8. N=64
0.17 (0.02)
3.6k
(0.00)
3.6k
d=32, N=8
0.21 (0.01)
600
9.3k
0.100)
600
9.3k
MAMBA
d=32, N=32
0.25 (0.03)
600
11.6k
0.11 (0.00)
600
11.6k
d=32, N=64
0.28 (0.03)
600
14.7k
0.11 (0.00)
600
14.7k
d=64, N=8
0.19 (0.00)
600
18.8k
0.11 (0.00)
600
18.8k
d=64, N=32
0.20 (0.00)
23.4k
0.11 (0.00)
23.4k
d=64, N=64
0.21 (0.00)
600
29.6k
0.11 (0.00)
600
29.6k
d=8, N=8
0.08 (0.00)
600
2.2k
0.04 (0.00)
600
2.2k
d=8, N=32
0.08 (0.00)
600
2.6k
0.04 (0.00)
600
2.6k
d=8, N=64
0.08 (0.00)
600
3.2k
0.04 (0.00)
600
3.2k
d=32, N=8
0.20 (0.00)
600
8.8k
(0.00)
600
8.8k
S4D
d=32, N=32
0.20 (0.00)
600
10.4k
0.100)
600
10.4k
d=32, N=64
0.20 (0.00)
600
12.5k
0.100)
600
12.5k
d=64, N=8
0.20 (0.00)
600
17.7k
0.11 (0.00)
600
17.7k
d=64, N=32
0.20 (0.00)
20.8k
(0.00)
20.8k
d=64, N=64
0.20 (0.00)
600
24.9k
0.11 (0.00)
600
24.9k
31

---

Understanding Selectivity in Mamba
Table A.4: Ablation on KEEP n-TH: MAMBA and S4D with varying sequence length T 30, 40, 50, embedding dimension d, and state size N.
T=30
T=40 # epch.
T=50
# epch.
# prm.
acc.
prm.
epch.
prm.
d=8, N=8
0.04 (0.00)
600
2.3k
0.04 (0.00)
600
2.3k
0.03 (0.00)
600
2.3k
d=8, N=32
0.04 (0.00)
600
2.9k
0.04 (0.00)
600
2.9k
0.03 (0.00)
600
2.9k
d=8, N=64
0.04 (0.00)
3.6k
0.04 (0.00)
3.6k
0.03 (0.00)
600
3.6k
d=32, N=8
0.09 (0.00)
600
9.3k
00.9000000)
600
9.3k
0.08 (0.00)
600
9.3k
MAMBA
d=32, N=32
0.09 (0.00)
600
11.6k
600
11.6k
0.08 (0.00)
600
11.6k
d=32, N=64
0.09 (0.00)
600
14.7k
0.09 (0.00)
600
14.7k
0.08 (0.00)
600
14.7k
d=64,
(0.900.00)
18.8k
0.09
18.8k
(0.00)
600
18.8k
d=64, N=32
0.0900.00)
600
23.4k
0.09 (0.00)
600
23.4k
(0.00)
600
23.4k
d=64, N=64
(0.00))
600
29.6k
0.090))00)
600
29.6k
0.090)
600
29 6k
d=8, N=8
0.03 (0.00)
600
2.2k
0.03
600
2.2k
(.,00)))0)
600
2.2k
d=8, N=32
600
2.6k
0.03
600
2.6k
0.00000000
600
2.6k
d=8, N=64
(0.00)
3.2k
(0.00)
3.2k
0.03 (0.00)
600
3.2k
d=32, N=8
0.08 (0.00)
600
8.8k
600
8.8k
0.09 (0.00)
600
8.8k
S4D
d=32, N=32
0.08 (0.00)
600
10.4k
0.09 (0.00)
600
10.4k
0.09 (0.00)
600
10.4k
d=32, N=64
0.08 (0.00)
600
12.5k
0.09 (0.00)
600
12.5k
0.09
600
12.5k
d=64, N=8
0.09 (0.00)
17.7k
(0.00)
17.7k
0.09 (0.00)
600
17.7k
d=64, N=32
0.09 (0.00)
600
20.8k
0.09 (0.00)
20.8k
0.09 (0.00)
600
20.8k
d=64, N=64
0.09 (0.00)
600
24.9k
0.09 (0.00)
600
24.9k
0.09 (0.00)
600
24.9k
32

---

Understanding Selectivity in Mamba
D.3. Task MQAR
Models All the results in Fig. A.2 use 1-layer models. The Mamba and Mamba-2 models explicitly disable the gating branch. This simplification is intended to verify our constructions without gating in Thm. 2 and Thm. 3. The Mamba-S4D model retains the full Mamba architecture, but only swapping the S6 layer with the S4D layer; not be confused with the original S4D model proposed in (Gu et al., 2022a).
Experiment Set-up We generate the data described App. C.I as follows. For each input sequence of the form
x [k1, UI, Kn,
Kip kigl:
we draw the key token As 1nd Uni ...., x}), and value token vs masked tokens except at the the query chunk where the input query is a key (c.g., at , kiz in the example above). We compute loss during training (and accuracy for evaluation) only at the query positions, informed from the target output sequence.
dnd Unif((L,.... .VI...I.... The target output sequence consists of
Discussion Here we expand on the results from Fig. in the main text. by reporting the accuracy of Mamba, Mamba-2. and Mamba-San trained on MQAR for varying model sizes. In Fig. A.2, we sweep over values of the value vocabulary size [V], to show how our theoretical bounds hold while varying this parameter. The bounds are still reasonably tight, and the observations drawn from Fig. 2 still hold in this case. The extra caveat (which does not however invalidate our claim) is that by increasing [V] we are making the task more difficult solve, and our training procedure for the simplest Mamba-S4DD fails to achieve satisfactory performance for the model sizes considered. Notice also that, for Mamba-2, the simplest task 4 can achieve 100% accuracy even below the thcoretical curve proposed our theorems. This can attributed to the following factors. On the one hand, our bounds rely on JL Lemma. which provides only asymptotical behaviors which might not be verified in practice for x so small. On the other hand, it is perfectly feasible that at this regime the architecture can recover a more efficient solution than the one theorized.
Our bound
100% accuracy
299% accuracy
X
<99% sccuracy
V = 256
MAMRA
MAMBA-2
256-
32-
Model dimension
x
16.
X
X
X
x
x
X
X
X
X
Number of keys
V 512
S4D
MAMBA
MA MBMA-22
XXooo
256.
x XXoos
Model cimension Id)
128-
32-
64-
36-
12.
x
X
x
Number of Kays (x)
Figure A.2. Trained models accuracy on MQAR (best of seeds) across k, and d. For S4D N 4, for Mamba N 2 x Ki and for Mamba2 N= = 8 InR. T = 100 and VV E [256, 512) for all runs.
With Fig. A.3, we further complement our results by sweeping over values of the N state size parameter. We remind that, according to Thm. 2 to 4, our theorized MQAR solutions require a value of at least N = 1, N k and N ~ log k for S4D, Mamba and Mamba-2, respectively. Indeed, in Fig. A.3 we observe that varying N does not have particular impact the final accuracy of S4D. For Mamba, on the other hand, we see that for N < x the training procedure fails to recover an exact solution to MQAR. Similarly, for Mamba-2, 22, no solution recovered for N < 4 log K. These results further validate the tightness our theoretical solutions.
33

---

Understanding Selectivity in Mamba
Our bound
100% securacy
299% accuracy
99%
S4D
N
N
N
N = 8
X X3000X
X Xx
Modkl dinatsson (e))
X
12
x
16
Number of keys
MAMBA
N U.D.
N 2
N 4
X XXX0XX
Mooki dimenesoo
37-
12-
77-
16
x
X
keys
MAMBA-2
K
log
X
Modit dimenskon bl
X
x
X
X
x
x
x
X
X
x
x
X
x
X
x
x
x
X
X
X
X
X
x
X
X
X
X
16
16
16
Number of keys
Figure A.3. Trained models accuracy on MQAR (best of seeds) across N. K, Mamba and Mamba-2 [V| = =52, while for S4D [V| 256 (as it failed to reach satisfactory accuracy for larger [V]. T = 100 for all runs.
34

---

Understanding Selectivity in Mamba
D.4. Task INDUCTION HEADS
Model All the results in Fig. A.2 use 1 1-layer Mamba that explicitly disables the gating branch. This simplification is intended to verify our constructions without gating in Lem. 3.
Experiment Set-up We generate the data described in App. C.2 follows. The data generation requires four scalar parameters: vocabulary size IV, sequence length 7 > 2/V] + 1, hard case probability p [0, 1], and special range ratio y E (0, 0.1]. We first draw X ~ Bernoull(p):: X O, we sample from the standard setting, otherwise the hard setting. The standard setting generates the input
sequence (#) ..., .T) by randomly drawing x; ~ V {1,- ..... [V]] for i 1, ..... T. The hard setting is intended to evaluate the long-range memorization capability (i.e., placing repeated tokens at the beginning and the end of the sequence), which consists of the following steps.
1. Randomly pick special token v" E V
2. Generate the input sequence by randomly drawing æ = (x1....,,27) 2r) where x, '24, V N0 for i I....,
3. Randomly draw position from E Uuif([]......777)
4. Place the special token v" at positions r, T -
the experiments for Fig. 3, we use T = 100, |V| E (5, 10, 20, 40). = 0.75, 0.1. In Fig. A.4, we further ablate the choice of state size N E (1V4.2)V].4411).
Discussion We design the hard setting hetter differentiate the capahilities from Mamba and our proposed Mamba-A T. Specifically, solving the standard setting of the INDUCTION HEADS task requires memorizing the latest previous occurrence. or forgetting the earlier previous occurrences. Note that the input sequence generated from the standard setting consists of many repeated tokens (by the requirement T' > 2| + 1). and the expected time for reappearance of any token is [V]. Thus, for small and medium-size [V), Mamba solve for these cases by using the state matrix with negative cigenvalues discount the remote past pairs, and thercby correctly output the latest previous occurrence. However, solving the hard setting additionally requires the model to memorize long-range information due to the special token (occurring at the beginning part and the end part of the soquence). We see that the Mamba solution with negative eigenvalues is at odds with the long-range memorization, as shown in Lem. 1. On the other hand, Mamba-A can satisfy both selective forgetting and long-range memorization via the input-dependence state matrix that erases outdated information specific to the input key, while retaining other information in the hidden state, illustrated in the proof of Lem. (see details in App. C.2).
Below we expand on the results Sec. 5.2 by reporting a sweep on the hidden state size N for the models used the INDUCTION HEADS task experiments, complementing the findings shown in Fig. 3.
bound
100% accuracy
accuracy
<99% securacy
[V]
2|V]
]|V|
*
dimeniion
8
16
88
88
88
88
88
88
83
Vocabulary size (VD
Figure A.4. Trained models accuracy on INDUCTION HEADS task (best of 5 seeds), varying [V] and d, with N E 21V (left. middle, right). Mamba-A 's performance (outlined) is equal or better than Mamba's (filled) and only hits 100% above the theoretical bound from Lem. 3 (black).
35


Reviews:
[{"review_id": "am6WF4fjU5", "reviewer": "Reviewer_WpKK", "summary": "This paper aims to understand the effect of gating in Mamba models in terms of function approximation power, long-term memory, and associative recall capabilities. Both theoretical derivations and empirical results are provided.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Common Response\n\nWe are grateful for the positive comments from reviewers on our paper, particularly regarding its **clarity**\n(kLHf: *“this paper is of very high quality”*; R4Me: *“quite clearly written”*),\n**soundness** (kLHf: *“very refreshing yet solid”*; R4Me: *“These results are proven rigorously”*),\nand **impact** (kLHf: *“I believe this work is significant. The theory echos closely with practice and provides practically useful insights for practitioners”*; R4Me: *“a useful addition to such literature”*; WpKK: *“this work provides a valuable perspective in understanding a core design element...and might inspire follow-up work”*).\n\nWe answer the remaining questions from the reviewers in the Individual Response section. Additionally, if the reviewers have any other remarks or recommendations that can help further improve the quality of the paper, we remain at their disposal to address them.\n\n## Individual Response\n***Does Mamba2 fit into the function approximation analysis?*** \\\nIndeed it does, and the reviewer is right in that it is worth highlighting more clearly. We do briefly mention (Sec3, after Eq(7)) how the Mamba2 layer is a simplification of Mamba that prescribes a state matrix parameterized by a single scalar $\\boldsymbol{\\Lambda}=\\lambda \\boldsymbol{I}$. Substituting this into (8) would remove the dependency of $\\lambda$ from the hidden-state component $n$, that is we would have $\\lambda_n \\equiv \\lambda$ for all $n=1, \\ldots, N$ in the integral. Nonetheless, for the proof in Sec4.1, it suffices to set $\\lambda_n = -1$ for all $n$ (see Line 605), so the proof still holds even in the Mamba2 framework. A similar reasoning also applies to Sec4.2, where again for Mamba2 $\\lambda_n \\equiv \\lambda$ would be constant over $n=1, \\ldots, N$, but it still does not affect the overall derivation. Following the reviewer's remark, in the main text we have added a note on the validity of both Thm1 and Lem1,2, for Mamba2.\n\n***Does the analysis shed light on the expressivity differences between Mamba1 and Mamba2?*** \\\nThe analysis in Sec4 (see also the response above) highlights that, *from the point of view of function expressivity alone*, if we consider *only the SSM-layer*, then there is no expressivity difference between Mamba and Mamba2 on approximating Haar Wavelets: Mamba2's simplification of setting $\\lambda_n \\equiv \\lambda$ across state dimension does not hinder this ability. \\\nNonetheless, the performance difference between Mamba and Mamba2 has been clearly shown empirically. Motivated by this, we extended our analysis to *the whole architecture* in Sec5. Our results in Thm2,3 highlight that, thanks to its per-SSM-parameter short-convolution layers, Mamba2 can recover more parameter-efficient solutions than Mamba for the MQAR synthetic task, hinting at a possible explanation for its superior performance. We believe analyzing other, more complex tasks to be an interesting next step in shedding more light on the expressivity differences between Mamba and Mamba2. \\\nWe thank the reviewer for raising this point. We took this chance to update the manuscript to highlight more clearly the implications that Thm2 and Thm3 have on the relative parameter-efficiency of Mamba and Mamba2 in solving MQAR.", "claims_and_evidence": "Three major claims as outlined in the paper:\n\n* S6 layer is more expressive than S4D because S6 can represent projections onto Haar wavelets, which presumably reflects the capability of approximating discontinuous functions.\n* S6 (akin to RNN models) suffers from memory decay.\n* There exists an analytical solution to MQAR using Mamba architecture, which reveals the advantage of input selectivity.\n\nAll claims are supported by theoretical derivations, along with some constructed tasks to show the tightness.", "methods_and_evaluation_criteria": "The empirical evaluation relies on benchmarks of MQAR which is widely used to assess associative recall capability.", "theoretical_claims": "I have checked the function approximation analysis and sensitivity analysis in sec 4.1 and 4.2. Below are questions on whether the analysis can be generalized:\n\n* Does Mamba2 fit into the function approximation analysis?\n* Does the analysis shed light on the expressivity differences between Mamba1 and Mamba2?", "experimental_designs_or_analyses": "The MQAR and Induction Head experiments are well-suited for differentiating the expressive power of S6 and S4D.", "relation_to_broader_scientific_literature": "The submission connects to the broad research effort on sub-quadratic/efficient architectures of LLMs. While most existing work focuses on empirical evaluations of pretraining results, this work provides a valuable perspective in understanding a core design element (i.e., gating) of Mamba models, and might inspire follow-up work on better parameterization of gating.", "essential_references_not_discussed": "NO.", "other_strengths_and_weaknesses": "NO.", "other_comments_or_suggestions": "No.", "questions_for_authors": "See the question before.", "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": "saWyk2UvhR", "reviewer": "Reviewer_R4Me", "summary": "This paper analyzes the flexibility of MAMBA showing that the S6 layer can i) it can project the input into Haar wavelet basis, ii) counteract memory decay, iii) solve multi-query associative recall (MQAR) problem proposed by Arora et al. tasks. While this is mostly a theoretical paper, the authors demonstrate applicability of their theory on approximation of discontinuous functions and the counteraction of memory decay via the KEEP n-TH task, requiring memorization the n-th token in a sequence. For the full Mamba model, they confirm empirically that the model sizes prescribed theoretically by their analytical solutions to the MQAR and INDUCTION HEADS are tight in practice.\n\n## Update after rebuttal\n\nI did not see anything to alter my score. I maintain a positive impression of the work.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Common Response\n\nWe refer the reviewer to the Common Response in the Rebuttal to WpKK. \n\n## Individual Response\n\nWe thank the reviewer once again for their positive feedback, and remain open to include any suggestion for improvement, or answer any question they might have.", "claims_and_evidence": "These results are proven rigorously.", "methods_and_evaluation_criteria": "The empirical studies add a little bit as they show that the theoretical bounds are not too loose.", "theoretical_claims": "I have read the proof sketches in the main text but not the detailed proofs in the supplementary material.", "experimental_designs_or_analyses": "I have looked at the experiments but not delved to deeply into it, this being mostly a theoretical paper.", "relation_to_broader_scientific_literature": "For several years now, there has been theoretical work probing of which tasks are possible to solve by transformers. The work contrasting transformers and state space models mostly focused on what transformers can do and state space models cannot (see for example Jelassi etal, 2024). However, as new state space models start to competitive performance, the theoretical interest in state space models capacity increases. This work is a useful addition to such literature.", "essential_references_not_discussed": "I am not in a position to comment on it.", "other_strengths_and_weaknesses": "The paper is quite clearly written.", "other_comments_or_suggestions": "None!", "questions_for_authors": "None!", "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": "DA8eWB2OTo", "reviewer": "Reviewer_kLHf", "summary": "This paper provide theoretical justifications for selective SSM layer (S6) in Mamba architecture. They show\n1. S6 has better expressiveness than S4D layer \n2. S6 suffers from exponential memory decay\n3. 1-layer Mamba (with S6) solves MQAR tasks with SSM mixer. MQAR is an information retrieval tasks that normally requires >2 transformer layers to solve.\n\nThe authors use sufficient numerical results to backup their theoretical findings.\n\n\n## update after rebuttal", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Common Response\n\nWe refer the reviewer to the Common Response in the Rebuttal to WpKK. \n\n## Individual Response\n\n***SD4 is mentioned without definition in abstract (and the 1st page).*** \\\nWe thank the reviewer for pointing out the mistake: this has been corrected in the text. We further thank the reviewer for their favorable review, and remain ready to address any question or suggestion for improvement they may have.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes. I check the lemmas and theorems and their proofs.", "experimental_designs_or_analyses": "yes. I check all verification experiments and their figure and tables.", "relation_to_broader_scientific_literature": "this work connects to both transformer and mamba. the literature are covered in sec 2.", "essential_references_not_discussed": "not that I can think of", "other_strengths_and_weaknesses": "### Strengths\n\n* Clarity: the language and formating of this paper is of very high quality. \n\n* Originality: the results and proposed method are original. There are lots of mamba papers, but this one is very refreshing yet solid.\n\n* Significance: I believe this work is significant. The theory echos closely with practice and provides practically useful  insights  for practitioners.\n\n### Weaknesses\n\n* (minor): SD4 is mentioned without definition in abstract (and the 1st page).\n\n\nOverall, good theory paper well-executed numerical support. I lean toward acceptance.", "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:

Meta-review: The paper received strong reviews. Reviewers emphasize clarity, originality, and significance. Several reviewers checked the proofs and attested to their rigor. The paper is well-written and tackles important open problems in state space modeling.
Decision: Accept (poster)