Block-Biased Mamba for Long-Range Sequence Processing
Authors: Annan Yu
N. Benjamin Erichson
Venue: NeurIPS 2025 poster ()
Keywords: Sequence Model
State Space Model
Mamba
Long Range Arena
Primary area: deep_learning
TL;DR: We analyze Mamba's shortcomings in long-range sequence processing through its expressiveness, inductive bias, and training stability, and design a remedy called Block-Biased S6 to achieve state-of-the-art performance on long-range tasks.

Abstract: Mamba extends earlier state space models (SSMs) by introducing input-dependent dynamics, and has demonstrated strong empirical performance across a range of domains, including language modeling, computer vision, and foundation models. However, a surprising weakness remains: despite being built on architectures designed for long-range dependencies, Mamba performs poorly on long-range sequential tasks. Understanding and addressing this gap is important for improving Mamba's universality and versatility. In this work, we analyze Mamba’s limitations through three perspectives: expressiveness, inductive bias, and training stability. Our theoretical results show how Mamba falls short in each of these aspects compared to earlier SSMs such as S4D. To address these issues, we propose $\text{B}\_{2}\text{S}\_{6}$, a simple extension of Mamba's S6 unit that combines block-wise selective dynamics with a channel-specific bias. We prove that these changes equip the model with a better-suited inductive bias and improve its expressiveness and stability. Empirically, $\text{B}\_{2}\text{S}\_{6}$ outperforms S4 and S4D on Long-Range Arena (LRA) tasks while maintaining Mamba's performance on language modeling benchmarks.

Block-Biased Mamba for Long-Range Sequence Processing
Annan Yu Center for Applied Mathematics Corelll University Ithaca, NY 14853 ay262@corne11.cm edu
N. Benjammi Erichson
Lawrence Berkeley National Laboratory International Computer Science Institute
Berkeley, CA 94720 erichson@icsi berkeley..yey..... edu
Abstract
Mamba extends earlier state space models (SSMs) by introducing input-dependent dynamics, and has demonstrated strong empirical performance across range of domains, including language modeling, computer vision, and foundation models. However, surprising weakness remains: despite being built on architectures designed for long-range dependencies, Mamba performs poorly on long- range sequential tasks. Understanding and addressing this gap important for improving Mamba's 's universality and versatility. In this work, we analyze Mamba's 's limitations through three perspectives: expressiveness, inductive bias, and training stability. Our theoretical results show how Mamba falls short in each of these aspects compared earlier SSMs such as S4D. To address these issues, we propose B2Sg, a simple extension of Mamba's S6 unit that combines block-wise selective dynamics with a channel-specific bias. We prove that these changes equip the model with better-suited inductive bias and improve its expressiveness and stability. Empirically, B2S, outperforms and S4D on Long-Range Arena (LRA) tasks while maintaining Mamba's performance on language modeling benchmarks.
Introduction
Mamba 25 has recently emerged as exciting alternative to Transformers, demonstrating strong empirical performance not only in language tasks [84] 45 but also in a variety of domains, including vision [07 [03]. audio 43, time series forecasting
43 and operator learning [33]. Based on the
state-space model (SSM) framework S4 and S4D, Mamba replaces attention with a selective state space mechanism. Unlike attention, which computes pairwise interactions across the sequence, this approach uses data-dependent weights a recurrent structure, aligning with the temporal nature of sequential data. Despite its success consistently underperforms on benchmarks such as the Long-Range Arena (LRA) [3]. This limitation is surprising, given that Mamba is derived from architectures specifically designed for modeling long-range dependencies. This behavior raises an important question:
general-purpose sequence modeling tool [5I], Mamba
Why does Mamba struggle with rangrannge sequence modeling?
To address this question, we conduct a theoretical analysis of the S6 state-space unit that underlies the Mamba architecture. Our goal is twofold: first, to explain why Mamba underperforms long-range sequence tasks, and second, to provide a principled remedy. Through a rigorous mathematical analysis, we identify three key factors that limit Mamba's 's ability to model long-range dependencies:
Expressiveness. Unlike S4D, which allows each internal channel to learn independent recurrent unit, Mamba shares parameters across all channels. Therefore, one should not interpret Mamba as
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

S4D
S6 (Mamba)
B2S6
B
1C CCuupe
80
C
B
CO
XX99
"Rangn
B
C
Figure 1: Comparison of S4D, S6, and B2S; units. S4D uses independent linear SSM units for each channel, giving it high capacity (or width) but no input-dependent selectivity. S6 introduces selective mechanism that modulates its internal dynamics based on the input, but shares parameters across channels, limiting its effective width and expressiveness. Our proposed BySg unit partitions the input into smaller blocks, enabling selective behavior across multiple subspaces (see also (20]). Additionally, includes channel-specific, inpu-indeepndeett bias term that further increases model capacity. These design choices improve the performance on long-range sequence tasks.
a straightforward extension of S4D that gains time-variant benefits at no cost. Rather, Theorem shows that while Mamba introduces input-dependent dynamics, it also sacrifices certain degrees of freedom. Most notably, it has weaker ability to learn independent behavior in each channel.
Inductive Bias. Unlike S4D, which uses a position-bas based bias that helps preserve long-range memory, Mamba adaptively decides what information to remember forget based on the input. This input-dependent mechanism too extreme and can cause the model to discard useful long-term information too quickly, leading to poor retention in long-range tasks (see Theorem 2)
Training Stability. Unlike S4D, which tends to train reliably on long sequences, Mamba empiri- cally exhibits unstable training behavior on the LRA benchmark tasks. We theoretically show that this instability comes from the input- dependee selection mechanism, which becomes increasingly difficult to optimize as the length of the sequence increases (see Theorem
Motivated by these findings, we propose the Block-Biased-S6 unit (B2Sg), a simple yet principled extension of the S6 architecture that addresses its limitations in modeling long-range dependencies (see Figure [). Like S6. B2S; retains the selective mechanism that modulates internal dynamics based on the imput. However, instead of computing the selection weights from the full input ug. we divide the model into smaller blocks, with each block operating on subset of Ug This multihead structure was already suggested in 220, but in addition to that, we introduce an input-independen, -hannel-specific bias term to the recurrent unit (see Table DT. We prove that these modifications enhance B2Sg's expressiveness and assign it a more appropriate inductive bias for handling range dependencies. We also apply reduced learning rate the parameters governing the input-dependent sampling intervals, which improves training stability. Together, these modifications enable B2S, to model long-range dependencies more effectively. On the LRA benchmark, B2SS resurrects Mamba from failure and even outperforms $4 and S4D. Furthermore, we show that ByS.g achieves comparable perplexity Mamba when trained on language tasks, demonstrating its versatility across domains.
Contributions. Our main contributions are as follows:
1. We analyze the limitations of Mamba in modeling long-range dependencies and provide the- oretical results that characterize the failure modes from three key aspects: (i) expressiveness (see Theorem (ii) inductive bias (see Theorem 2). and (iii) training stability (see Theorem B).
2. We introduce the B2S; unit, a principled extension of the S6 architecture. We show that B2Ss has universal approximation capabilities (see Theorem 4 and is thus more expressive, while also equipped with a more suitable inductive bias for long-range sequence processing (see Theorem 5)

---

Table 1: comparison of different classes of selective or non-selective state space models.
S4D
S5
Mamba
Mamba2
B2S6
A, Dependency
input- ieput-independent input-independnt
inpu--nndeenndent input-independent
input-dependent input-dependent
depu-ddeenndent input-dependent >
input-dependent both 1
B/C Dependency
Number of Heads
> complex
> complex diagonal
Parameters
real
real scalar
complex diagonal
Parameterization
diagonal
diagonal
We evaluate B2S; on the Long-Range Arena benchmark, where it achieves state-of-the-an perfor- mance. In addition, preliminary examination on the SlimPajama dataset [78) shows that B2Sg matches Mamba's 's performance on language modeling tasks, demonstrating its versatility.
2 The S4D and S6 recurrent units
Introduced in 28 and 27], the S4 and S4D recurrent units process sequences of d-dimensional inputs
u= (u.....,uu uy to produce corresponding d-dimensional outputs I ssss4D(()) = y = =y---.... .yI.....yL) subscripts to index positions in a sequence and superscripts to index channels. For example, ( denotes the ith entry of the kth input vector un. Each channel of S4/S4D model is computed independently with linear system. For the ith channel, the update rule is
Throughout this paper,
( _
( uk
Yk 1
(
(i)
1 1sisd,
(1)
k-1
where X E Cax1 is the hidden state with initial condition x O, and u E TR. The matrices ( Coxn, ( C"x1 and c( Clxn are derived from continuous-time system. While there are many different discretization rules, we focus on the Zero-Order Hold (ZOH) discretization:
=xxp(A (1) A):
A (i) exp(b()))
B --11-
I)B(,,
C =(0) (( (2)
Here, A Caxn B Cax1 (+i) Cixn, and bto) E are trainable parameters, specific to each channel sis a S4 and S4D particularly well-suited for modeling sequences with long-range dependencies and achieve state-of-the-art results on the Long-Range Arena benchmark. This strength is largely due the reparameterized (6) values, which, when chosen small, give the model long-term memory by slowing the system dynamics. S4 and S4D differ only in how the matrix is represented; we henceforth focus on S4D, which uses diagonal A consistent with Mamba.
One limitation of the S4D units is that their dynamics are linear and time-invariant. The S6 units (25) used in the Mamba models improve upon this by making the dynamics input-dependent. More specifically, given d-dimensional input sequence u (ui...., 1,), an S6 unit computes the ith channel of the output s'so (u) = y (y1...., VL) as follows:
(
A A xb-
(8) (i)
(1) B u
yk
(6)
C
isisd,
where x E RR and u .yk E R. The matrices A k
(8)
b and Cg are computed at each step as
exp(( A).
(A))
softplus(w "u ((0),
B (8) (0- - AA 1)Bug, C C. (4)
The trainable parameters of the model are A E Baxn B E Raxd C Rdxn, E RRd, and ((0) E IR for all 1 < < dF The key innovation the S6 architecture lies the input-dependent dynamics. Both B, and Cs change with each input step, introducing nonlincarity into the system. In addition, the sampling interval (1) also input-dependent, allowing the model to vary the rate at which memorizes or forgets information. Comparing eq. (4) to eq. (2), we note that Bu, and ug C in an S6 are shared across all channels, while B and (T) in an S4D are channel-specific. In the next section, we show how this distinction limits the "effective width" and expressiveness of an S6 unit.
Note that we tie the matrix A across all channels, which is consistent with the I. I.A training setup. This choice can be loasened without affecting the essence of the paper.
"In many implementations of Mamba, the selection of A is more complex, and BN may be computed using forward Euler method instead. We follow the formulation in Algorithm of the original Mamba paper 25
3

---

3 A single-layer Mamba is not a universal approximator
universal approximation theorem (UAT) characterizes the expressive power of neural networks by showing that, under certain conditions, they can approximate any target function to arbitrary accuracy, provided they are sufficiently wide or deep [60] 39]. While UAT' do not directly address training dynamics, they serve as useful sanity check for a model's 's theoretical capacity. For deep S4D models, some universal approximation results have been established in [87. In this section, we prove a new UAT of S4D and a new non-UAT of S6 that highlight how the shared state matrices in S6 significantly constrain its expressiveness. We investigate the task of learning a univariate sequential task defined by a ground-truth function G. using single-layer model that takes in an input (#1,-... Ut), linearly embeds the input into a d-dimensional space, passes it through an S4D or S6 unit, and then decodes the output. More precisely, we consider the class of models (see Figure defined by
G(u)= G((ui, ut)) No (T(MMM) .......... (Mu.)...,M.,))), + O),
where M Rdx1 is an encoder, N Rdx1 is decoder, A Rdx1 is a bias term, and R IR is a nonlinear activation function applied entrywise to the last output of an S4D or S6 unit There is only one remaining issue: in the definition of S6, we have C2 = 2n C. That means I ( Muz))ı is zero as long as ur = O, which makes G with I I'so clearly not universal approximators on the domain [0, 1] To avoid "cheating" in this way, we assume that un = 1 and only focus on approximating functions on [0, _ x {1}, which makes our non-UAT stronger.
Theorem 1. The single-layer S4D models are universal approximators of continuous functions, but the single-layer S6 models are not. More precisely, fixing constant for A (0) for all 1 S i S d
in eq. ( and (0) for all i S d and S S L in eq. 4). the following two statements hold:
1. Let o - be any Lipschitz continuous. non non-polynomial activation function. Given any continuous function G: [0, 1]4 - IR and any > 0. There exist some d 2 1, n 2 1, and choice of parameters M. N, 0, A, B6ú, and C(O), where S i S d, such that the map G in eq. (5) with T= ssn in eq. 1b satisfies that (G(u) G(u)] < for any E [0, 1)
Let a : - be any function and let L 2 3 be given. There exists a continuous function G: 1|-1 x {1} + IR and some >> 0 such that for any d 2 1, any n 2 1, and any choice of parameters M, N, 0, A, B, and C, the map G in eq. with r = S in eq. satisfies that (G(u) G(u) > for some u E 10.1] x (1).
You should not interpret Theorem T as a pessimistic claim that Mamba models are incapable of solving complex tasks. In practice, Mamba architectures typically use multiple stacked layers and do not fix A both of which enhance the expressiveness. Rather, it highlights that an S6 unit is not simple extension of S4D that inherits all of its structural principles, and the sharing of the B and C matrices across all channels can reduce its capacity, making it potentially less effective than S4D for certain tasks. Beyond its impact on approximation capacity, this also leads to issues in larger models,
such as training instability [23 16 67]. reduced generalization [5 92
100, 40). and challenges in
interpretability [104]] 94].
While Theorem assumes that A is fixed, we use a synthetic experiment to demonstrate that even when A is trainable, the channel -ineeeedddddeeddentt design of B and C in S6 limits its ability to solve certain problems than S4D. Consider the input function u(t: g) E Pi cos(p:t). where .1..... P10 are the 10 smallest prime numbers and g E 1 is a coefficient vector. The learning task is recover from the function a(f:g), i.e., to learn the mapping G u(t:g) + g. Intuitively, if the model is sufficiently wide, each channel can specialize in capturing one of the cosine components, making the problem
easy. We train both a single-layer S6 model and single-layer S4D model to approximate this mapping, also allowing the A-related parameters to be learned in both cases. As shown in Figure 2 the performance of the S6 model remains nearly constant as the width d increases, reflecting the fact that its effective width does not scale with d. In contrast, the S4D model benefits significantly from increased width.

---

S4D
S6
B2Sc
10
100
glla
gl12
10
10
10
2 1p- JaH
-6 (p-2 (p-222o04)
Soud b.0l
10"
Epoch 8
10
10
10
Epoch
Epoch
Figure 2: The mean loss lا - g/1 between the true coefficient g and the model prediction g. Every model has a single layer and is trained for 10 epochs. Here, d is the number of channels in model. For the B2Sg model, the number of blocks and p is the number of channels each block.
4 Mambas exhibit strong inductive bias
necessary condition for a model obtain the capability of handling long-range dependency is that it possesses long-term memory. In S4D models, this is achieved by learning small values of (1) which slow down the evolution the continuous- time LTI systems and allow information to persist over longer time horizons (74). In contrast, S6 models introduce input-dependent sampling interval ( where certain inputs lead to larger values of A causing faster memory decay, while others result in smaller values and slower decay.
This input-adaptin memory control serves as useful inductive bias in language modeling, where only a limited number of words in a sentence typically carry semantic significance. However, this same mechanism is less suitable for many LRA tasks involving non-linguistic sequences, such as Image, Pathf inder, and PathX, where inputs are derived from flattened pixel arrays. In these settings, while some positional bias may exist (e-g.. central pixels may more informative than peripheral ones), it is difficult to justify that pixels of certain colors in a CIFAR image should receive dramatically more attention than others. As a result, the input-dependent memory dynamics of S6 may act as overly aggressive or misaligned inductive bias in such contexts.
To formalize the notion of input-dependent bias. let I denote either an S4D S6 unit as defined in eq. [I) and eq. (3. respectively. Given a d-dimensional input sequence u = (u....., ut), we ask: how does the output depend on each individual input vector uy?? To study this, we examine the relative gradient of the output with respect to each ug. Specifically, we define:
S₁ JR
a Duk
JN
ut)) E RRd
ug)
(6)
J
where || ||r denotes the Frobenius norm, and Jur is the Jacobian of the final output vector with respect to the k'tth input vector. Intuitively, the quantity Sş. measures the relative influence of un on the final output. attributed to the k-th input. For linear system such as S4D, the Jacobians J are independent of the input u, so the relative gradients Sş remain constant across different inputs. This indicates that the model treats all inputs uniformly and does not introduce bias toward particular input u changes. In contrast, we now show that S6 units introduce a bias when input vectors vary in magnitude.
other words, it tells us how much of the output's sensitivity to perturbations is
Theorem 2. An S6 unit imposes an exponentially large bias as the input magnitude increases. That is, fix a ko S and let A E Rx be a diagonal matrix with negative diagonal entries and w / For almost every (a.e.) input sequence u (u,,... .u,) E R2 and a.e. parameters B E Bxd C e Rdxn, and b E Idd, let Ss6. be defined in eq. (6). where I = I's is defined in eq. As 0o, the following statements hold for any p 0:
1. w uk. > 0, we have
h < ko:
Ss6, & (u,j.... Ukg-1, Cuko: ukot1 ...., uz))
O(c-l)
ku,
O(1)
kn-
5

---

S4D
S6
B2Sc
10
1o7 ( (a) 0
10
tp/ [m) (m)
1of ( (1)
10
10
10
10"
10
10"
10 10"
100
10"
10"
Figure 3: The mean relative loss G(() (u)) /oi for different choices of a and Ø2. The S6 model cannot make useful predictions when 01 is small and 02 is large; B2Sg fixes this. In all experiments, we fix d = 32. For B2Ss, set h 8 and 4.
2. If w uko O, we have
(1)
o ko: Ku-
Ss6, & ((u),.... -ko-1: Cuko Uko+1. uz)) =
Oc (
Moreover, let S be defined in eq. (6 where I SaD is defined in eq. (I We have that SSaD, constant and does not depend on the input u.
The theorem shows that as the magnitude of single input vector Uso grows, one of two extreme behaviors can emerge: either the current input (Case or all previous inputs (Case 1) exert an exponentially diminishing influence on the final output. Neither outcome is desirable when modeling sequences with long-range dependencies. In practice, many implementations of Mamba mitigate this issue through normalization layers that constrain the magnitudes of input vectors. Additionally, unlike the original formulation (25), more sophisticated parameterizations are often used to compute
the sampling intervals A These modifications can be interpreted as efforts to counteract the exponential input bias introduced by maggemmae vectors.
From Theorem 2 we see another perspective of the limitation of S6: in the computation of A softplus(w un + b(t) the vector w defines two half-spaces in R Inputs uy that lie in the positive half-space produce larger values of A, leading slower dynamics and better memorization, whereas those in the negative half-space yield smaller A, resulting in faster forgetting. This setup restricts the model's ability assign distinct memorization biases different nonlinear regions of the input space. In section 6
we will see how B2S6 improves upon this.
To illustrate the inductive bias of S6, consider synthetic task with inputs sampled from
u = (u), O, O, ur.),
u = [u1..... un]" E Rdd ~ ~1 01), uy ~ N' o2la),
where 01, 02 > 0 are fixed. The task learn the mapping G(u) = |ur]. Since the target size scales with Ø1, we report the relative error [G(u) ((u)//o to make comparisons fair across different settings. F
the relative errors of different models under different values of 01 and 02- While the S4D model remains robust across all cases, the S6 model fails significantly when d2 is small and 02 is large. This corroborates Theorem large input uy in the half-space {u | w Fuu 0} triggers fast forgetting, erasing the memory of u and making the prediction unreliable.
Mambas are not stable to train
Our discussion far has focused the expressiveness and generalization of S6 models. In general, a model lacking expressiveness cannot achieve low training eиror on complex tasks, while poor generalization refers to models that fit the training data well but perform poorly have trained S6 models on LRA tasks, then you might observe another issue: the training loss curve does not decrease smoothly. Occasionally, single optimization step causes the model to collapse, turning a high-accuracy model into one that performs no better than random guessing. This reflects a training stability issue rather than of expressiveness generalization.
test data. If you
This phenomenon has been empirically studied in the context of general neural networks 21 [9].
where instability is often linked to sharp curvature in the loss landscape. In this section, we provide

---

a theoretical explanation for the training instability observed in S6 models. Since we are not tied to specific task or loss function, we instead study how the output [(;; )) depends on the model parameters O While curvature involves second-order derivatives, we focus on first-order gradients. If (a/a0)) (u: )) is large in magnitude, large curvatures in the loss landscape might appear when we compose (uu; )) with other functions to form loss function.
We find that instability in S6 models is closely tied to how the sampling interval A is computed. To formalize this, we prove theorem comparing the gradients of S6 and S4D models with respect to their A-related parameters. For simplicity, we restrict our analysis to single-inpu setting, i.e., d 1, and drop all superscripts. Hence, the (1) for S4D in eq. (2) only depends
parameter b = ((2) E R and the ()) for S6 in eq. (4) only depends on two parameters w E IR and b = ((2) E R. Since the specific input distribution is unknown, we make generic assumption that the inputs random variables with mild regularity conditions, ensuring that the analysis is broadly applicable.
Theorem 3. Assume that d and n 2 Let Tss (u; θ) and (u: b) be given eq. !I
and 33. respectively, where O is the collection of all model parameters. Given a diagonal matrix A E I maxn whose diagonal entries are negative, and E R=xi, the following statements hold:
1. An 56 model is less stable train than S4D model the input magnitude increases. That is,
fix some L 2 1, b E IR, and w = 0. Let u = (wi,... ur.) ) be sampled from a distribution D. For every C E R1xx if Eu-D, (0/0w)) ss. Eu-D, ((8/ (0000/00000 [ss Eu-Dy (0/āb)[ssan # O, then
Eured (0/00)) T S6 (u)zl R)c²),
Bu (0/ T (6())
Enwen, , (0/0b) (0 (u))
E En 1 [b] su)
2. An S6 model is less stable to train than an S4D model as the input length increases. That is, fix w
0. There exists sequence b(L) + -00 L 0o, such that given any sequence of distributions
Dz on [0, 001]=-1 x {1} with Eu-DD, aT [u)z] = 0, c1 S Vara<D, I 02 for all 1 1,
where C1, 02 > O are universal constants, and Σ LigiejeL Cov(A u4 A レ-) ug) = OLL),
where A = e we have for a.e. C E R Ixn that
EurD, (0)Ob()))(&s (())) Eu-Dy (0/0b(L)) (4)(
VV > /VVVVZ
lim L-+00
The first part of our result shows that an S6 model becomes increasingly unstable relative to an S4D model as the input size grows. This is expected, as the matrices B and C2 in eq. (3 depend linearly the input u,. However, this effect is typically mitigated in practice by input normalization and is therefore of less practical concern. The more significant insight lies in the second part of the theorem: the sequence length I. increases, the S6 model also becomes less stable train. This is particularly relevant in the context of LRA tasks, where sequences can be extremely long; for example, the PathX task involves sequences of length 16384. In such cases, the training stability of S6 and S4D are dramatically different. Note that we have made some technical assumptions. First, to preserve long-range dependencies as L + oo, the sampling interval A must shrink. This justifies why we assumed b b(L) + 0o so that A + 0+ as L + co. Second, since the pairwise covariances
between A 1-- "4 and A T-jj u, can be both positive and negative, cancellation occurs and assuming their total sum scales as O(L) is mild and reflects the natural behavior of aggregated dependencies in long sequences. We also fixed w = O, leaving the case when w A 0 for future work. More numerical and empirical experiments that illustrate Theorem 3 can be found in Appendix E
6 B2S8: : an expressive, gently biased, and stable selective model
We introduce BySg Block-BBiased-S to address the shortcomings of S6. Given an input sequence (u : )r.), where Hg Rdx1 for < k S L, we assume that d h p for two hyperparameters hpp 21 1. We then partition each input vector into h sub-vectors of size p: U₁ = [u] A "u]]= Note that we use a left-superscript for the block index. Then, the output of B2S. unit is sequence TE,S. (u) y (y1,.... ..), where each output vector can also be partitioned into h sub-vectors of size p. i.e., yk = I
"yI] defined by
i (1) xx
AA
1xp-1 + B
(6)
j uu
(i)
1y(0) (1) = Cr xx
1sisp.
isish, h, (7)
7

---

where 2x E R", &uf E RR,'y E R, and the matrices / A (4) CC E 12 are input-dependent and computed at each step as
Raax, yBx E R and
= exp / A); $ a - softplus(2 w T Jus + $b(0)) B i A ("A 1) B weight -u + +'B bins
ISkSL.
(8)
C₁= u i,
The trainable parameters are A E R=x", B weight E R"xp, 'B (6)
bias E Rax!, &C E R Jw E RPx1, and 1b(4) R. for every 1 < S p and S S h. The pseudocode for BySg is given in Algorithm 2 which compares to Algorithm for S6 found in the Mamba paper [25], where 8A (+) Broadcast, Broadcast(Linnerr, (-)) and i ( Broadcast, (aaadcast,,(Lieeaa (Linear (-)) for every S j S h.
Algorithm S6 Forward Pass
Algorithm 2 B2Sc Forward Pass
Input: [B L
Input: (B L
Output: y : (B d)
Output: y y(B dB)
I: parfor j j= =1 ho
b independent for every block
2: -I (ip 1ip-P. 11
block index
1: Bs (B L n), Bs.(), (.))
3: JB: (B La), /B(i..)) + Ba (.
4:
JB: Lp Brondcss,, ( (BB) Bromdcasty, (Bas))
2: Css (B L n), Cull.d))) - T
5: iC: (B Ln), (C(L)))) x(i. ( Tic
3: A (B L + softplus(aa (x) b)
6: JA : p) softplus( e (x((: ()) +b)
4: A,B B: (B L. d) + discretize(A, A. Bs))
7: JK,/B: (B Lp discreizee A, AB)
5: y + SSM(A, B. (x))()
8: ((. - SSM A, B, (C)x(,
9: end parfor
Comparing the B2Sa unit in eq. 7 to the S6 unit in eq. 3 we highlight two key differences (see Figure []. First, B2Sa partitions the d-dimensional input into h blocks of p-dimensional sub- vectors and applies an independent recurrent unit to each block. Second, we introduce a bias term (B( in the computation of the input matrix BBg. While this term is input-independent, it varies across channels, increasing the model's effective width. We do not claim that we invented the block unit design, which was previously suggested in the Mamba2 paper [20]. but in this work, we rigorously demonstrate how the combination of block structure and channel-speen -specific bias enhances the expressiveness and generalization of the model on long-range sequence modeling tasks.
First, we show that unlike an S6 model, a B2Sg model regains the universal approximation property studied in Theorem by introducing either the block unit or the bias unit.
Theorem 4. Fix constant for A for all i, and k in eq. The following two statements hold:
1. The block unit alone makes B2Sg a universal approximator. More precisely, set B (1) all i and j. and let : - R be any Lipschitz continuous, non-polynomial activation function. Given any continuous function G [0, 1 1 {1} + R and any e <>> 0. There exist some h and p with p = 2 1, > 1, and choice of parameters M, N, 0, A, BBeeihht and C, where
bias 0 for
1 < j < h, such that the map G in eq. (5) with I'= L'B,S, in eq. satisfies that
(ū G(u)) S, for any [0, 1) x1).
2. The bias unit alone makes S2S, universal approximator. More precisely, set h 1 and p d, and let a : R + R be any Lipschitz continuous, non-polynomial activation function. Given any continuous function G j0, 1)1-11 x {1} + E and any > 0. There exist some d 2 1, 2 1, and choice parameters M, N, 0, A, B Beeight (B()) and C, where < < d, such that the map G in eq. (5) with I = 1n2,,, in eq.
bias
7 satisfies that
(G(u) G(u) S,
for any 10, x {1).
The block design and channel-specifie bias in the B2Sg model significantly enhance its expressiveness, enabling to handle more complex sequential tasks that benefit from wider neural networks. To illustrate this, revisit the experiment from section 3 this time using only the block structure

---

(without the bias term), noting that the addition of the bias would only further improve performance. As shown in Figure Z unlike an S6 model, the performance of the B2S, model improves substantially as d increases, holding the ratio h/p constant.
In contrast the S6 model, the B2Sg architecture introduces gentler inductive bias that is better suited for long-range tasks. Instead of Theorem 2
we now establish the following result.
Theorem 5. Fix a ko < L and let A E IR=x" be a diagonal matrix with negative diagonal entries. For a.e. input sequence u (u.,...,, uz) R² such that there exist j1 and J2 with 1 ''w" ii uko
bins ', and 18(0), where < j S h and 1 S i S p. let S be defined eq. (6). where I' B2S. is defined in eq. (7). We have
and JawT ja uk, < O, and a.c. parameters ' 'Beeight. (B(i)
(1dc- k ko
Snoso, Uko 1. cuggs Uko ur.))
o(dd-t) O(1) k > ko.
h ko:
c tO0.
If each ww is randomly initialized for S S h, then the probabilities that a given input satisfies "w" 3 uko < 0 and /w 'u > 0 are equal. As a result, the probability that the assumptions in The- oremBJare are violated vanishes exponentially with the number of blocks h. Compared to Theorem 2 this means that ByS, exhibits significantly milder inductive bias in the presence of large inputs, avoiding exponentially decaying effects as the imput magnitude increases. This behavior is confirmed by the experiment in section 4 where BySa maintains strong performance even when Ø1 is very small and d2 is very large. Lastly, while the block and bias components improve expressiveness and inductive bias, they do not directly resolve the training instability that arises with long sequences. Therefore, for LRA tasks, we reduce the learning rates of /w and d 18()) to improve training stability
Experiments
Ablation. In this paper, we show that the B2S, model achieves strong performance on the LRA benchmark. Compared to the standard Mamba model, our ByS, variant incorporates multihead structure and a bias term Buiss, and uses complex-valued parameters for diag(A) (A). Beiight and B (see [64] 97 47)) Before presenting full LRA results, we demonstrate the utility of each modification in Table 2 The second column shows model accuracy without bias terms, the third without complex parameterization, and all rows below the first the multihead structure. As seen in Table all three components contribute positively to the performance on the SCIFAR-10 task. An interesting observation is that, when d = h x p is fixed, model accuracy increases monotonically with h only in the absence of Baias This should not be surprising: if Buias = O, the model's s effective width is h, so larger h improves expressiveness and generalization. However, once Boias is introduced, the model already attains high effective (though non-selective) width even when h = because ' (t) with i. Thus, increasing h trades off between more selective blocks, which improves the quantity of selectivity, and reduced receptive field Un per block, which impairs the quality of each selection.
bias varies
that our B2So model outperforms many selective and non-selective models on LRA. In particular, is the first selective SSM we are aware of that achieves Mate-of-theean performance this benchmark. A more comprehensive table is found in Appendix F
Long-Range Arena. In Table
Language Tasks. While our primary goal is not to train large language model (LLM), we conduct a preliminary evaluation of B2Sg's 's language modeling capability using sampled version of the SlimPajama-627B dataset 778.. For each model family, we train a model with approximately 30 layers, corresponding to approximately 250M parameters. The precise number of layers varies slightly with models to make the number of parameters roughly the same. We train four models for one epoch and report the perplexity statistics. We only vary the recurrent unit without changing the rest the model architecture. That is, all models are based on the Mamba architecture proposed
in [25 Fig. 3]:
y [((((Con(Lineaut) (x))) o (Linear(x)).
"In general, freezing some of the other parameters indeed helps stabilize the training. but it also puts us at risk of converging to bad minima. For example, we observed that it never works to freeze B and C. and that freezing the matrix A and D usually does not have big impact, given that they initialized properly (see 97).

---

Table 2: Ablation study of our B2S; model. We train a model to learn the SCIFAR-10 task, where we fix 128 and change the number of blocks h and the size of each block p (see eq. 7 For each pair of h and p. we train full B2Sg model with complex parameters, B2S; model with complex parameters but bias term, and full B2Sg model with real parameters. more greenish color indicates better accuracy, while more reddish color labels a worse accuracy.
With Buin (Complex)
Without Bhine (Complex)
With Bii (Real)
h
Accuracy + Std
Accuracy + Std
Accuracy Std
128
81.56 1.22
71.77 + 1.88
47.82 5.42
2
64
84.83 + 0.73
79.44 + 1.10
54,40 + 0.35
32
83.86 1.38
80,93 + 0.94
54.77 1: 0.54
87.19 = 0.26
81.54 + 1.14
53.81 4: 0.87
16
8
87.04 = 0.33
83.78 + 0.79
51.79 +: 3.61
32
4
85.83 = 0.80
84.11 + 0.39
56.18 + 0.18
86.30 0.60
84.27 0.76
52.03 2.04
128
85.32 0.55
84.79 + 0.43
56.45 + 0.71
Table 3: Test accuracies in the Long-Range Arena of different variants of SSMs. A bold number indicates the best accuracy on task while an underlined number corresponds
second best.
Model
ListOps
Text
Retrieval
Image
Pathfinder
Path-X
Avg.
(28
59.60
90.90
88.65
94.20
96.35
86.09
S4D
60.47
89.46
88.19
93.06
91.95
84.89
62.15
89.31
91.40
88,00
95.33
98.58
87.46
(Mamba)
26
38.02
72.14
69.82
69.26
67.32
66.59
S7
63.77
87.22
91.80
61.14
65.62
61.50
71.82
Mamba2 20
41.45
79.23
71.96
75.45
69.07
70.54
B2Sg (ours)
63.85
88.32
91.44
88.81
95.93
97.90
87.71
20:45
20.50
120.36
220.2
20038
20.17
where o stands for the Hadamard product and we only change I in three models to be P S6- I Mambaz, and S respectively. We find that introducing the bias term Boias only makes the training of a B2Sg model 2.7% slower than an S6 model comparable size. Yet, the perplexity of ByS, closely matches that of S6, showing the versatility of model.
Figure 4: Perplexity training steps for S6, Mamba2, B2Sg, and $4D model. The dataset is SlimPajama-6B, a sampled version of the Slim JlimPajam--6277 dataset.
Steps
S6
Mamba2
B2Sg
S4D
B,S
2620
61.5441
59.6233
62.6133
64.1203
SAD
5241
42.7838
39.8305
44.3846
46.5118
Perplexity
7862
34.1740
31.9994
35.0440
39.3496
50
10483
29.7074
27.0987
30.3787
35,3028
27.0171
25.2930
27.3891
32.5946
40
15725
25.2265
24.0179
25.2668
30.6639
18346
23.9789
23.1224
23.9370
29.2052
20967
23.0044
22.3695
22.8394
28.0314
30
23588
22.2368
21.7509
22.1322
27.0897
26209
21.6068
21.2480
21.5206
26.3221
20
20k
28830
21.0785
20.8218
20.9773
25.6661
Steps
8 Conclusion
In this work, provided a theoretical analysis of Mamba's limitations in modeling long-range sequences, focusing on expressiveness, inductive bias, and training stability. We proposed new model, B2Sg. which introduces block structure and channel-specific bias to improve upon each of these dimensions. Empirical results demonstrate that B2S, significantly enhances performance on long-range sequence tasks while maintaining versatility across domains. Future work includes a more detailed analysis of training stability across all Mamba parameters, exploring architectural refinements to further mitigate instability, studying the training dynamics of Mamba and BySş to better characterize their inductive biases, and scaling up B2Sg as a language or foundation model.
10

---

Acknowledgments
AY was partially supported by the Office of Naval Research under Grant Number N00014-23-1-2722 and NSF DMS-231962) NBE would like to acknowledge the U.S. Department of Energy, under Contract Number DE-ACO2-05CHI 1231 and DE-ACO2-05CHI1233 for providing partial support of this work.
References
[1] Naman Agarwal, Daniel Suo, Xinyi Chen, and Elad Hazan. Spectral state space models. arXiv preprint arXiv:2312.06837, 2023.
[2] Zeyuan Allen-Zhu Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparam- eterized neural networks, going beyond two layers. Advances in neural information processing systems, 32, 2019.
[3] Carmen Amo Alonso, Jerome Sieber, and Melanie N Zeilinger. State space models as foundation models: A control theoretic overview. arXiv preprini arXiv:2403./6899 2024.
[4] Marco Ancona, Enea Ceolini, Cengiz Öztireli, and Markus Gross. Towards better under- standing of muudintt-based attribution methods for deep neural networks. arXiv preprint arXiv:1711.06104 2017.
[5] Alessio Ansuini, Alessandro Laio, Jakob H Macke, and Davide Zoccolan. Intrinsic dimension of data representations in deep neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[6] Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks.
In International Conference on Machine Learning, pages 1120-1128, PMLR, 2016.
[7] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International conference on machine learning, pages 322-332. PMLR, 2019.
[8] David Bachrens, Timon Schroeter, Stefan Harmeling, Motoaki Kawanabe, Katja Hansen, and Klaus-Robert Müller. How to explain individual classification decisions. The Journal of Machine Learning Research, 11 11:1803-1831, 2010.
[9] Shaojie Bai, J Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convo- lutional and recurrent networks for sequence modeling. arXiv preprint arXiv:1803.01277 2018.
[10] Peter L Bartlett, Andrea Montanari, and Alexander Rakhlin. Deep learning: statistical viewpoint. Acta numerica, 30:87-201, 2021.
[11] Ronen Basri, David Jacobs, Yoni Kasten, and Shira Kritchman. The convergence rate neural networks for learned functions of different frequencies. Adv. Neur. Info. Proc. Syst., 32, 2019.
[12] Gregory Beylkin. Fast convolution algorithm for state space models. arXiv preprint arXiv:241..17729, 2024.
[13] Bo Chang. Minmin Chen, Eldad Haber, and Ed H Chi. Antisymmetrierm:: A dynamical system view on recurrent neural networks. International Conference on Machine Learning, 2019.
[14] Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Chris- tian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):124018, 2019.
[15] Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IFEE Transactions on Neural Networks, 6(4):9119917, 1995.

---

[16] Anna Choromanska, Mikael Henaff, Michael Mathieu, Gérard Ben Arous, and Yann LeCun. The loss surfaces of multilayer networks. In Artificiil intelligence and statisties, pages 192-204. PMLR, 2015.
[17] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Machine Learning, 2020.
[18] George Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of control, signals and systems, 2(4))303- -344, 1989.
[19] Alex Damian, Eshaan Nichani, and Jason D Lee. Self-stabilization: The implicit bias of gradient descent the edge of stability. International Conference on Learning Representations, 2023.
[20] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint wrXiv:2405.21066 2024.
[21] Yann N Dauphin, Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, Surya Ganguli, and Yoshua Bengio. Identifying and attacking the saddle point problem in high-dimensiona, non-convex optimization. Advances in neural information processing systems, 27, 2014.
[22] N Benjamin Erichson, Omri Azencot, Alejandro Queiruga, Liam Hodgkinson, and Michael W Mahoney, Lipschitz recurrent neural networks. In International Conference on Learning Representations, 2021.
[23] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pages 249-256. JMLR Workshop and Conference Proceedings, 2010.
[24] Riccardo Grazzi, Julien Siems, Arber Zela, Jörg KH Franke, Frank Hutter, and Massimiliano Pontil. Unlocking state-tracking in linear rmns through negative eigenvalues. International Conference on Learning Representations, 2025.
[25] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
[26] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33:1474 202.
[27] Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing 00135971-35983,, -35983, 2022.
Systems, 35:35971
[28] Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022.
[29] Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. Advances in Neural Information Processing Systems, 2022.
[30] Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. International Conference on Learning Representations, 2023.
[31] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[32] Arya Honarpisheh, Mustafa Bozdag, Mario Sznaier, and Octavia Camps. Generalization error analysis for selective state-space models through the lens of attention. arXiv preprint arXiv:2502.0144 2025.
12

---

[33] Zheyuan Hu, Nazanin Ahmadi Daryakenari, Qianli Shen, Kenji Kawaguchi, and George Em Karniadakis. State-space models are accurate and efficient neural operators for dynamical systems. arXiv preprint urXiv:2409.03231, 2024.
[34] Sukjun Hwang, Aakash Sunil Lahoti, Ratish Puduppully. Tri Dao, and Albert Gu. Hydra: Bidirectional state space models through generalized matrix mixers. Advances in Neural Information Processing Systems, 37:1 1777:11086-1111 10876-110908, 2024.
[35] A. Jacot, F. Gabriel, and Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Adv. Neur. Info. Proc. Syst. 31, 2018.
[36] Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. J. Basic Eng. 82(1):35-45, 1960.
[37] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mas: Fast autoregressive transformers with linear attention. In International conference machine learning, pages 5156-5165. PMLR, 2020.
[38] Nitish Shirish Keskar, Dheevatsa Mudigere. Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. International Conference on Learning Representations, 2017.
[39] Patrick Kidger and Terry J. Lyons. Universal approximation with deep narrow networks. CoRR, 985399 2019.
[40] Sungyoon Kim, Aaron Mishkin, and Mert Pilanci. Exploring the loss landscape of regularized neural networks via convex duality. International Conference on Learning Representations, 2025.
[41] Nikita Kitaev, łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Machine Learning, 2020.
[42] Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018.
[43] Kai Li, Guo Chen, Runxuan Yang, and Xiaolin Hu. Spmamba: State-space model is all you need in speech separation. arXiv preprint arXiv:2404.0220 002663, 2024.
[44] Kunchang Li, Xinhao Li, Yi Wang. Yinan He, Yali Wang, Limin Wang, and Yu Qiao. Video- mamba: State space model for efficient video understanding. In European Conference Computer Vision, pages 237-255. Springer, 2024.
[45] Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, al. Jamba: A hybrid transformer-mambe language model. International Conference on Learning Representations, 2025.
[46] Fusheng Liu and Qianxiao Li. From generalization analysis to optimization designs for state space models. arXiv preprint arXivv2405.02677 2024.
[47] Fusheng Liu and Qianxiao Li. Autocorrelation matters: Understanding the role of initialization schemes for state space models. International Conference on Learning Representations, 2025.
[48] Ziwei Liu, Qidong Liu, Yejing Wang, Wanyu Wang, Pengyue Jia, Maolin Wang, Zitao Liu, Yi Chang, and Xiangyu Zhao. Sigma: Selective gated mamba for sequential recommendation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 12264- 12272, 2025.
[49] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. International Conference on Learning Representations, 2017.
[50] Lu Lu, Pengzhan Jin, Guofei Pang. Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):218-223) 2021.
13

---

[51] Haoyu Ma, Yushu Chen, Wenlai Zhao, Jinzhe Yang, Yingsheng Ji, Xinghua Xu, Xiaozhu Liu, Hao Jing, Shengzhuo Liu, and Guangwen Yang. mamba foundation model for time series forecasting. arXiv preprint arXiv:241 (tfi:24/1.020000 1.02941, 2024.
[52] Song Mei, Andrea Montamari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665- E7671, 2018.
[53] Nicola Muca Cirone, Antonio Orvieto, Benjamin Walker, Cristopher Salvi, and Terry Lyons. Theoretical foundations of deep selective state-space models. Advances in Neural Information Processing Systems, 37:127226- 127277, 2024.
[54] Yuqi Nie, Nam H Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam. A time series worth 64 words: Long-term forecasting with transformers. In The Eleventh International Conference on Learning Representations, 2023.
[55] Antonio Orvieto, Samuel Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. arXiv preprint arXiv:2303 2023.
[56] Rom N Parnichkun, Stefano Massaroli, Alessandro Moro, Jimmy TH Smith, Ramin Hasani, Mathias Lechner, Qi An, Christopher Ré, Hajime Asama, and Stefano Ermon. State-free inference of state-space models: The transfer function approach. International Conference on Machine Learning, 2024.
[57] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pages 1310-1318. Pmlr, 2013.
[58] Badri Narayana Patro and Vijay Srinivas Agneeswaran. Mamba-360, Survey of state space models transformer alternative for long sequence modelling: Methods, applications, and challenges. arXiv preprint arXiv:2404. 16111,, 2024.
[59] Yan Ru Pei. Let SSMs be ConvNets: State-space modeling with optimal tensor contractions. International Conference on Learning Representations, 2025.
[60] Allan Pinkus. Approximation theory of the mlp model in neural networks. Acta numerica, 8:143-195, 1999.
[61] Biqing Qi, Junqi Gao, Dong Li, Kaiyan Zhang, Jianxing Liu, Ligang Wu, and Bowen Zhou. S4++: Elevating long sequence modeling with state memory reply. 2024.
[62] Yanyuan Qiao, Zheng Yu, Longteng Guo, Sihan Chen, Zijia Zhao, Mingzhen Sun, Qi Wu, and Jing Liu. Vl-mamba: Exploring state space models for multimodal learning. arXiv preprint 2.Xi::2403.13600 2024.
[63] Gokul Raju Govinda Raju, Nikola Zubić, Marco Cannici, and Davide Scaramuzza. Per- turbed state space feature encoders for optical flow with event cameras. arXiv preprint arXiv:2504. 10669, 2025.
[64] Yuval Ran-Milo, Eden Lumbroso, Edo Cohen- KohennKarlik,, Raja Giryes, Amir Globerson, and Nadav Cohen. Provable benefits of complex parameterizations for structured state space models. Advances in Neural Information Processing Systems, 37:115906-115999, 2024.
[65] Stefano Rando, Luca Romani, Matteo Migliarini, Luca Franco, Denis Gudovskiy, and Fabio Galasso. Serpent: Selective resampling for expressive state space models. arXiv preprint 1rXiv:2501.11729, 2025.
[66] David W Romero, Anna Kuzina, Erik J Bekkers, Jakub M Tomezak, and Mark Hoogendoom. Ckconv: Continuous kernel convolution for sequential data. In International Conference on Machine Learning, 2022.
[67] Grant Rotskoff and Eric Vanden-Fijnden Trainability and accuracy of artificial neural net- works: An interacting particle system approach. Communications on Pure and Applied Mathematics, 1889-198899 2022.
14

---

[68] T Konstantin Rusch and Siddhartha Mishra. Unicornn: A recurrent model for learning very long time dependencies. In International Conference on Machine Learning, pages 9168-9178. PMLR, 2021.
[69] T Konstantin Rusch and Daniela Rus. Oscillatory state-space models. International Conference on Learning Representations, 2025.
[70] Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of werrparametrried neural networks. arXiv preprint arXiv: 17706.04544, 2017.
[71] Anton Maximilian Schäfer and Hans Georg Zimmermann. Recurrent neural networks are universal approximators. In Artificial Neural Networks-ICANN 2006: 16th International Conference, Athens, Greece, September 10-14, 2006, Proceedings, Part 16, pages 632-640. Springer, 2006.
[72] Shuaijie Shen, Chao Wang, Renzhuo Huang, Yan Zhong, Qinghai Guo, Zhichao Lu, Jianguo Zhang, and Luziwei Leng. SpikingSSMs: Learning long sequences with sparse and parallel spiking state space models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 20380-20388, 2025.
[73] Avanti Shrikumar, Peyton Greenside, and Anshul Kundaje. Learning important features through propagating activation differences. In International conference on machine learning, pages 3155-3153, PMIR, 2017.
[74] Jerome Sieher, Carmen Amo Alonso, Alexandre Didier, Melanie Zeilinger, and Antonio Orvieto. Understanding the differences in foundation models: Attention, state space mod- els, and recurrent neural networks. Advances in Neural Information Processing Systems, 37:1 37:134534-13456, 2024.
[75] Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps. International Conference Learning Representations, 2014.
[76] Jakub Smékal, Jimmy TH Smith, Michael Kleinman, Dan Biderman, and Scott W Linder- man. Towards theory of learning dynamics in deep state space models. arXiv preprint vrXiv:2407.07279 2024.
[77] Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023.
[78] Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, 2023.
[79] Taylan Soydan, Nikola Zubić, Nico Messikommer, Siddhartha Mishra, and Davide Scara- muzza. S7: Selective and simplified state space layers for sequence modeling. arXiv preprint arXiv:2410.03464, 2024.
[80] Namjoon Suh and Guang Cheng- survey on statistical theory of deep learning: Approxima- tion, training dynamics, and generative models. Annual Review of Statistics and Its Application, 12, 2024.
[81] Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In International conference on machine learning, pages 3319-3328. PMLR, 2017.
[82] Matus Telgarsky. Benefits of depth in neural networks. In Conference on learning theory, pages 1517-1539. PMLR, 2016.
[83] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[84] Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, et al. An empirical study of mamba-based language models. arXiv preprint arXiv:2406.0 078877 2024.
15

---

[85] Maolin Wang, Sheng Zhang, Ruocheng Guo, Wanyu Wang, Xuetao Wei. Zitao Liu, Hongzhi Yin, Yi Chang, and Xiangyu Zhao. Star-rec: Making peace with length variance and pattern diversity in sequential recommendation. arXiv preprint arXiv:2505.034x. 20344004400 2025.
[86] Shida Wang and Qianxiao Li. StableSSM: Alleviating the curse of memory in state-space models through stable reparameterization. arXiv preprini urXiv:2311. 2023.
[87] Shida Wang and Beichen Xue. State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. Advances in Neural Information Processing Systems, 36, 2024.
[88] Shida Wang and Beichen Xue. State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. Advances in Neural Information Processing Systems, 36, 2024.
[89] Xiao Wang. Shiao Wang. Yuhe Ding, Yuchang Li, Wentao Wu, Yao Rong, Weizhe Kong, Ju Huang. Shihao Li, Haoxiang Yang, et al. State space model for new-generation network alternative to transformers: A survey, arXiv preprint (:Xxv:2444.095111 2024.
[90] Yihan Wang, Lujun Zhang, Annan Yu, N Benjamin Erichson, and Tiantian Yang. A deep state space model for minfall-runoff simulations. arXiv preprint 25.V::2501.198000 2025.
[91] Zihan Wang, Fanheng Kong, Shi Feng, Ming Wang. Xiaocui Yang. Han Zhao, Daling Wang. and Yifei Zhang. Is mamba effective for time series forecasting? Neurocomputing, 119:129178, 2025.
[92] Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and generalization in deep neural networks. In International Conference on Machine Learning, pages 10472. PMLR, 2020.
[93] Tiankai Xie, Caleb Geniesse, Jiaqing Chen, Yaoqing Yang, Dmitriy Morozov, Michael W Mahoney, Ross Maciejewski, and Gunther H Weber. Evaluating loss landscapes from topology perspective. arXiv preprint arXiv:241 1.09807, 2024.
[94] Greg Yang and Edward Hu. Feature learning in infinite-width neural networks, arXiv preprint arXiv:2011.1452,, 2020.
[95] Zhifan Ye, Kejing Xia, Yonggan Fu, Xin Dong, Jihoon Hong, Xiangchi Yuan, Shizhe Diao, Jan Kautz, Pavlo Molchanov, and Yingyan Celine Lin. Longmamba: Enhancing mamba's a's long-context capabilities via training-free receptive field enlargement. In The Thirteenth International Conference on Learning Representations, 2025.
[96] Annan Yu, Chloé Becquey, Diana Halikias, Matthew Esmaili Mallory, and Alex Townsend. Arbitrary-depth universal approximation theorems for operator neural networks. arXiv preprint arXiv:2109. 1354, 2021.
[97] Annan Yu, Dongwei Lyu, Soon Hoe Lim, Michael W Mahoney, and N Benjamin Erich- son. Tuning frequency bias of state space models. International Conference on Learning Representations, 2025.
[98] Annan Yu, Michael W Mahoney, and N Benjamin Erichson. HOPE for robust parameteriza- tion of long-memory state space models. Internation Conference on Learning Representations, 2025.
[99] Annan Yu, Arnur Nigmetov, Dmitriy Morozov, Michael W. Mahoney, and N. Benjamin Erichson. Robustifying state-space models for long sequences via approximate diagonalization. In The Twelfth International Conference on Learning Representations, 2024.
[100] Annan Yu, Yunan Yang, and Alex Townsend. Tuning frequency bias in neural network training with nonuniform data. International Conference on Learning Representations, 2023.
[101] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequenee functions? Internation Conference on Learning Representations, 2020.
16

---

[102] Fengrui Zhang, Jiaoyi Hou, Dayong Ning, Cheng Zhou, Gangda Liang, and Zhilei Liu. Srs4: A stacked residual deep neural network for heave motion continuous prediction of salvage barge. Available at SSRN 4938844, 2024.
[103] Hanwei Zhang, Ying Zhu, Dan Wang, Lijun Zhang, Tianxiang Chen, Ziyang Wang, and Zi Ye. A survey on visual mamba. Applied Sciences, 14(13):5683, 2024.
[104] Guoqiang Zhong, Li-Na Wang. Xiao Ling, and Junyu Dong. An overview on data representa- tion learning: From traditional feature learning to recent deep learning. The Journal of Finance and Data Science, 2(4):265-278, 2016.
[105] Ding-Xuan Zhou. Universality of deep convolutional neural networks. Applied and computa- tional harmonic analysis, 48(2):787-794, 2020.
[106] Tian Zhou, Ziqing Ma, Qingsong Wen, Xue Wang, Liang Sun, and Rong Jin. Fedformer: Frequency enhanced decomposed transformer for long-term series forecasting. In International Conference on Machine Learning, pages 27268 27286. PMLR, 2022.
[107] Lianghui Zhu, Liao Bencheng, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model. International Conference on Machine Learning, 2024.
17

---

A Related works
Sequence Models. Sequential data appear across a wide range of domains, including natural language processing, computer vision, generative modeling, and scientific machine learning. To model such
data, primitive deep learning approaches primarily relied on recurrent neural networks (RNNs) and
their variants [6] [3 22 68 [55], as well as convolutional neural networks (CNNs)
[ 66]. The last
decade has witnessed the dominance of the transformer architecture [37 17 41 106 54), following
the seminal work [83]. Recently, new class of models known as state-space models (SSMs) has
emerged as a promising competitor [28 27 30 77 25 20 These models represent sequences
through underlying continuous-time dynamical system and offer key advantage: they can handle sequences of varying lengths with a fixed number of parameters [85] and be trained with time and
memory complexity that scales linearly with sequence longth [2 77 25).
State Space Models. The term "state space models" (SSMs) originates from control theory and dates back to [36. The first widely adopted SSM in machine learning was the S4 model proposed by 228 Both the original S4 and its successor, Liquid-S4, S4 300, use diagonal- sesoonal-plu-rankee plus-rank-one (DPRO) structure in the state matrix to significantly reduce computational cost compared to traditional RNNs. Later, [27] showed that a purely diagonal state matrix could achieve comparable performance, leading to the simplified S4D model. Since then, most SSMs have adopted diagonal state matrices, including S5 1771, Regularized SSM [46], Stable-SSM [6). S4-PTD, and S5-PTD 99 631. While a diagonal parameterization explores the structural simplicity of SSMs, [59] analyzes the computational accelcration of LTI systems and 12 investigates its numerical stability. Other works explore implicit sequence models that do not rely directly on an LTI formulation rely modified LTI formulation, 69. An extensive list of different SSM architectures has been given in the survey The expressiveness of SSMs has been studied in 98 47, their training dynamics in [76, 97.. and generalization properties in (97 46). Several studies have also investigated the
such as 98 1 56
articles 89 58
representation stability of SSMs, including 86 98
Applications of SSMs in different scientific
fields have been studied in several papers [102 90 63
Mamba Models. Mamba (25 extends SSMs by introducing input-dependent dynamics via selective mechanism, resulting in a highly efficient sequential architecture. Unlike earlier SSMs that use fixed, imput-independent recurrence, Mamba adapts its dynamics at each step, enabling it to rival transformer-bas based models in language tasks. The theoretical benefits of this input-dependent recurrence was studied in [53] through the lens of controlled differential equations (CDEs). Mamba has since inspired a number of extensions and studies. Mamba2 [20] incorporates a multihead structure with softmax gating. Mamba has also been applied beyond language modeling (58),
including computer vision [107 103
44 time-series forecasting (91]. multimodal learning (62), and
audio processing 43). In addition, 15 proposed using Mamba as a foundation model backbone across modalities. Despite its versatility, follow-up studies such as 3 have highlighted the limitations of Mamba on longerange sequence benchmarks such as LRA. The S7 model 77] has been proposed as another selective model with slightly improved performances on LRA tasks; though, the accuracies are still substantially worse than models like S4D and S5, Other efforts to enhance Mamba's 's long memory retention and selectivity mechanism are found in (95] and 655. respectively. The spectral properties of the state matrix are considered in 24 and related to Mamba's sate-iracking capabilities. In addition, the recent work (32) proposes a theoretical generalization upper bound based on the Rademacher complexity of Mamba models. Notably, (74] conducted a thorough comparison of these different recurrent units, including S4, Mamba, and traditional RNNs. Many works have shown the advantages of a bidirectional structure in SSMs and Mambas 34 48
which we also adopt training
the LRA benchmark tasks.
Universal A Approximation Theorems. In this paper, we use the universal approximation theorem (UAT) as a tool to assess a model's expressiveness. UATs ask whether a target function can be approximated arbitrary accuracy by a sufficiently large neural network. This line of work dates back to [18) who proved that two-layer neural networks with sigmoid activation are universal approximators. This result was later extended networks with any continuous, non-polyno pol-polylo activation function. While many classical results focus on shallow but wide networks, recent work by 39 established that deep, narrow networks can also be universal approximators. Universal approximation properties have been studied across various architectures. Notable examples include shallow and wide operator neural networks (50 15, as well as their deep and narrow counterparts (%]. For sequence models, UATs are known for RNNs
CNNs [055, Transformers [101), and
18

---

SSMs [6]. There are limited works on the universal approximation properties of Mamba models. One notable exception [53], where the uniform closure of Mamba (i.e., functions that can be uniformly approximated by Mamba) is analyzed in continuous setting (i.e., the inputs are continuous time series instead of discrete sequences) through linear controlled differential equations. The analysis there highlights the important role of an input-dependent sampling interval A in the expressiveness of Mamba, suggesting that Theorem ! may not hold for a trainable (see Appendix (G).
Width and Depth of Neural Networks. While UATs focus on the density of neural networks in a given topology, more practically relevant questions involve the rate of approximation, e.g., how deep or wide network must be to approximate target function within a given accuracy. One of the most celebrated results is by [2], which shows that depth improves expressiveness more efficiently than width. However, in practice, deep but narrow networks often face optimization challenges
during training (23 57 16 67 contrast, networks with greater effective width tend to enjoy better theoretical guarantees generalization. This advantage can be understood through both the
neural tangent kernel (NTK) perspective [35 7 2
100] and mean-field theory [22 80 10, both
of which highlight the benefits of wide models in training dynamics and generalization performance.
Sensitivity Analysis. Neural networks often operate on high-dimensional inputs, and a large body of work has studied how to quantify the sensitivity of the output to individual input components - a line of research typically referred to attribution. While our paper focuses on relative gradients, many other gradient-basee attribution methods have been developed, including DeepLIFT [73], integrated gradients (T. sensitivity-n 4 and others g [5]. Most of these methods compare given input to a baseline or reference input, whereas our analysis focuses specifically on the effect of large-magnitude inputs in isolation without reference to a benchmark input.
Training Stability and the Loss Landscape. Training stability in deep neural networks has been extensively studied through the lens of loss landscape geometry. Sharp minima and high-curvature
regions are often associated with poor generalization and unstable optimization [38] [4 42). Gradient explosion or vanishing can lead to optimization failure [23 57 for recurrent neural networks. More
recent works have linked curvature and gradient dynamics to the trainability of neural networks [16] [7]] These studies emphasize the importance of smooth, well-conditioned landscapes for stable training, motivating architectural choices [1 93 and learning rate schedulers [49] that mitigate instability.
B Proofs of UAT and non-UAT results
In this section, we prove all technical results related to UATs presented in section and Recall that the univariate neural network architecture that we study in this paper given by
(u)) Č((u)...., wr)) No (o MMr))) +0.
pictorial illustration of this architecture is given in Figure
B.1 Proof of Theorem 1
To prove that such a & is not a universal approximator when r = Ts.. we will see that since an S6 model ties the matrices B, and Cş for all channels, the system becomes a quadratic encoder of the input sequence. We need the following technical lemma to show that a quadratic encoder is not invertible, and therefore, information will be lost when we apply this encoder to an input sequence.
Lemma 1. Given any L 2 3, there exists a continuous function G [0, 1] + R such that given any quadratic polynomial of L variables P [0, 1 x (1) + IR there exist two points x,y E [0, 1|4-1 {1} such that
P(x) = P(y), [G(x) - G(y) > 1.
Proof. Since the restriction of quadratic polynomial of L variables 3 variables is still a quadratic polynomial, without loss of generality, we assume that L = 3. Moreover, instead of assuming that and are functions on [0, 1|-11 x {1), we can assume that they are functions on [0. 1J2. is well-known that there is a number N 2 1 such that every bivariate quadratic polynomial has at most N strict local minima or maxima. We construct by selecting N + arbitrary distinct
points X1, X2, .... XN+1 in (0, 2)2. Around each point Xj. we make small disk D2 Dxxx).
19

---

Input/Outpuu Space
(#1 m1)
M
Embedding Space
I SAD II 96
Figure 5: The architecture of the neural network eq. ( that we study in this paper. In this picture, horizontal operator is applied channel-wise to every sequence in a channel, and a vertical operation is applied dlement-wise every position sequence. green color indicates linear operator while an orange color indicates a nonlinear one.
By taking P small enough, can also ensure that Dj C [0, 1J2 for all 1 S N 1 and that D1, D2,.... DN+1 are mutually disjoint. Let G be continuous function such that G(x,) ) = 0 for every 1 S j s N + and G equals 2 on 0D,, the boundary of Dj. for every 1 S j N 1. Such a G clearly cxists by Urysohn's 's lemma. We claim that G satisfics the condition in the our lemma. To see this, given any arbitrary quadratic polynomial P [0, 1J2 + E. we let S, be the connected component of the set [x| = P(x;)} that contains Xj. There are two possibilities: either Sj intersects 0D, not. If Sj does not intersect 0Dj. then that means there is strict local minimum or a strict local maximum of P in D,. but Dı...., DN+1 are mutually disjoint and there are at most N local minima or maxima of P. Hence, there is least one such that Sj intersects 0Dj. Let E Sj na Then, we have that
P(x,) = P(y),
IG(x,) - GGy) = 10 2 >.
The proof is complete.
We now ready prove Theorem The proof of the first part is built upon a result from (60 for UAT of two-layer wide neural network to approximate any continuous function and result from 88
for UAT of LTI systems approximate any convolutional kernel.
Proof of Theorem 7
We prove the two statements separately.
Proof of Part T. Without loss of generality, assume that is 1-Lipschitz. Let continuous function
G: J0, 1] IR and an error tolerance 0 be given. By [60], there exists function G: [0, 1)
of the form
(u)) = Σ ( (a(w(()) Tu (( )()), i=1
where w(() & and b, R for every 1 < i S d, such that
(E(u) G(u))
for every 1). For any 1 i d, by (87. there exist n, 1, A R=:x*, B()) R
and (s) Rixn, such that the matrices A and C (0) from the discretized LTT system (see eq. satisfy that
2V +1)
*Otherwise, if is E-piscciizz, we can divide by to make I-Lipschitz and multiply by so that value of & does not change.
20

---

for all 1 S I L, where f = [f(l) f..(TT Hence, we have that
1
(((w())) 11 + 0(0))
A (1) ) B u + 0(3)
j=1
E (A (A (A
Σ
(w
aj
S (ww())
- (A (0) 3L "
)² ||u||2 S 2/dL([| +
2vd(1 l2 + 1)
where the first inequality follows from the Lipschitz continuity and the second inequality follows from the Hölder's inequality. Set
Σi= and define block-diagonal matrix A so that
(1)
A (2)
A
A (d)
For every 1 Si < d, define ( E [ to be sparse column vector so that B B i-1
(=== nhy +1):
(Ek= my)) ( and is zero elsewhere. Similarly, define C E R to sparse column
vector so that ((E =1 + 1): (Er=1 n²)) = C and is zero elsewhere. Now, it is easy to
see that given these definitions, the S4D defined in eq. M satisfies that
vz.)) (t)
c A (1) )
4j
for every S i < d. Hence, let GSaD be such that
GSaD [0, 111t R, (ui,...,uz) + (T sap ((Mu)..... Mun)))+ +0),
where M - I|T. For any u E IR4, we have that
A ) ug )) 2
(Gs4D(u)= G(u)) s [[f]|2
Σ
a(w(())Tu)
Σ 0 (A
S [| 2vd([ 1)
<
Now, for any u E R. we have that
(u) G(u)) < (G)
s c Se.
The statement proved.
Proof of Part II. Given any sequence E R-,, any encoder M = [m(1) m (i)TT and any
Mamba parameters A E Ra** B E Rnxd and C E Rdxn we first show how the Mamba system eq. (3) can be simplified. For any < d, we have that
(6)
Ax (6)
(Pux) ( uk.
x41
(9)
(yx) (3) (u: x (i)
where = P- - IJMM, and Q = MTC C are fixed matrices depending only on A, B, C, and M. Hence, the ith entry of the final output is given by
I (Mu1,.... Ml ( ur
LOA' Pug
F((n)..)))))
21

---

Importantly, note that since uz = 1, F is quadratic function in u1. u...... ... i. The output of Gse can then be expressed as
that does not depend on
Gso(u) En (m F()) +
where = =In(1) ... (.))). Let G be the function defined in Lemma] 1 Then, given any choice of M. N, 0, A, B, and C, we know that there is pair of inputs and v such that
F(u) F(v),
(G(u) G(v)l > 1.
Since F(u) F(v), we also have that Gss u) (vs(v). By the triangle inequality, we have
(u) (ss(u)]+ (G(v) - Gss(v)) 2 [G(uu G(v)) 1.
Setting c = 1/2, we are done.
B.2 Proof of Theorem
The proof of Theorem can be easily established upon our proof that S4D models are universal approximators.
Proof of Theorem Note that the construction of a universal approximator in Theorem can be i < (see [88] 97)). The second statement
achieved by using the same (++) = [1 1] for all 1 <
follows immediately from the first statement of Theorem by setting Beeeih O, B be the same ( used saD,, and IC be the rank-1 matrix whose entries are all d-1 so that ur M 1C
[1 ... IJ. For the first statement, consider the function H (u) given by // (u) G(vu), where the square-root is applied elementwise. Clearly, is continuous since G is continuous. Then, we
know, by the second statement, that there exist M [1
N, 0, A, B weight 0, B (i)
bius
and C = =[//di., such that the BySa system I Bas,, defined in eq. with h and p and the
map
HB,SO : [0, 1] + R, (ui,..., uL) (o (a (Mun,..... Muz)) + +
satisfy that
B.SS (u) H(u) 6, for any E [0, 10.11--1x (1).
Now, define h = d and p 1. Let M, N, A, = Á be the same matrices, and let ' Bueight B bias ( = O, and CC = [1
1] for all and j. Then, it is clear that the system In,s, defined by
these matrices satisfies that
(TB,s, (u))z (TB,,, (oass(un)) (u²)),, for any [0, 1]4-1 {1), where u2 is the sequence obtained by squaring every entry of u. Hence, the map
Gn,S, [0, 111-11111- (1) IR, (uI,..., ut) No B2So Pestasssuuua (MMu,..,MMu))),
Mur))
satisfies that
Gass, (u) G(u) = Ha.s, (u²) H(u²) S
for any u 10. {1}.
The proof is complete.
C Proofs of inductive bias results
In this section, prove Theorem and The proof of these results relies very simple argument that helps us avoid the cancellation terms, which we state below.
Lemma 2. Let E IR be given. Assume : R R and g IR - I are two functions such that f()) = ax. o(x²) and g(x) fix² + o(x²) ) as >> pc, for some constants a, B * 0. Then, we have that for a.e. choice of C1, 02 E IR that e₁ of f(x) c2g(z) ya² + o(x²) (z²)) as + for some 7 + O.
Proof. The proof is straightforward: long as we have that ocι + Boz * O. the conclusion is satisfied. Since a and are nonzero, the equation oc + Bez = is satisfied only on a Lebesgue null set of and cz.
22

---

C.1 Proof of Theorem 2
We now prove the main theorems by explicitly calculating the Jacobians.
Proof of T Given an S6 system in eq. (3) and an input sequence u (ui..... Uz), let d, we have that
YL (u) be the last oulput of the system. For any
LL
VL Cx u u
II A(u.) B(uj)u))
1=1 1=++1
where
(u:) exp (A( exp (softplus(w Au)
B(u,) = A- (exp (softplus(w T (,) - I) Bu,
Note that we have changed the notations so that (u,) is exactly the matrix A, and (u,,) exactly the matrix B; in eq. (3.. This helps us better keep track of the dependency of every term on the input u. In addition, is easy to see that b66) does not play role in the asymptotic behaviors when - too0. Hence, we set them to zero. For any we have
(a) u
L II (=j+1
OyL
a
(J,),
(
B(uj)u,
Ou:
Ouk
where
A(uz)
(II A(u:) ((
, <r,
d
dfrra)
[ A(u))
(B(u.) +ut (0) o B(u.)
j r,
ou
0
j >
For any 8 # 1, we have
0yg, (t)
a
(Jr)s,
un
C
II A (=j+1
B(u;)u,
Our (8)
Oui
where
a A(u.) (T A(u.) B(uy)u
, <r,,
Dal!)
p
plriet)
[T A(u))
(1) a
B(u,)
j r,
Bu.""
We now break the proof into two cases.
Case I: + We discuss the cases when < k. == k. and > separately.
When
is easy to check that ||A decays exponentially when c + co. Hence,
we have [ and ||e, [ decay exponentially as c - oo for every s, t. and j. That is, the Jacobian norm || decays exponentially as c + co.
When r = kt the norm of the term
II A(I.) i-r+1
& B(u,) u Fr.
20 for all A and a.e. choice of and B, where * 0 is constant. That is, d (k,s) and p(k....) grow like a linear factor plus o(c) term as c >> co.
grows like ac o(c)
Hence, by Lemma 2 for a.e. choice of C, we have that || @(c) as +
23

---

When > k, it is straightforward to check that | = a(s)c + o(c²) for some
constants ass) * 0 and a.e. choice of and [| [ = t(s,t)c² + o(c²) for some constants ((s, 1) + 0 and a.e. choice of B. Hence, by Lemma 2 have that || = (c²²) as 0c.
for a.e. choice of C,
The first part of the theorem is proved by combining the three statements above.
Case II: c + -0. We discuss the cases when < k. and r > It separately.
When < k; the only thing that changes in the expression of J, when e + -00 is the matrix A (ux), in which case it converges to I. Hence, we have that ||, || (1).
When r = k we have that
i A ux)
o(d
Dug
for any > 0 and s. Moreover, we also have that
[B = O(dd"))
B ((ux)
2uky
for any p > O and That means we have ||d (k.>) and |1)] 1 decay exponentially for any choice of s, L, and j. Hence, have that J ([c)) for any p 0.
When k. note that while | grows linearly, the norm of the vector [B(u))] decays exponentially. Therefore, we have that [B(ue)u decays exponentially for every This shows that for a.e. B, the norms | and || converge to constants for all s, t, and j. By Lemma 2 we have || F ((1) for a.e. choice of C.
The second part of the theorem is proved by combining the three statements above. The claim about an I'sD system is obvious since the system is linear.
Interestingly, from the proof of Theorem we see that when c - oo, the reason why Ss6.6 is large for l > Ko is that u, plays an important role determining the matrix Bx,, that affects a large input Ukg- This is analogous the function ((x, y) xry. When is large and y is small, the gradient (0/Or)f is still small but (0/0g)f is huge.
C.2 Proof of Theorem
Once Theorem 2 is proved, the proof of Theorem 5 can be maintained easily from it.
Proof of Theorem For each 1 < r S L, the Jacobian J, a block diagonal matrix J, diag( J Jy..... J, where every JJ is p x p matrix. Since the cases when is negative and when e is positive are symmetrie, without loss of generality, we assume that c - co. We study the Jacobian norms when < k, k. and A separately. The following statements hold for a.e. model parameters.
When r < k, from the proof of Theorem we know that || decays exponentially when Jw T > and |JJ,|| ((1) when w in < 0. Hence, by our assumption, we know that ||J, (1))
When k. from the proof of Theorem ww Ju > 0. Hence, we have || = (cc).
2 we know that ||J +cc) as long as
When > k from the proof of Theorem 2 we know that |J = +ecc) as long as w Ju > 0. Hence, we have ||J || = O
Combining the three cases, proved the result.
24

---

D Proof of the stability result
We now bring in the last piece of technical details by proving the stability result in Theorem 3 The proof again relies on a tedious expansion of the gradients, which then gives us expressions that are intellectually interesting to analyze with basic probability theory.
Proof of Theorem Since we assumed that d = 1, we drop all superscripts for channel indices. Given an input u (#1,.... ut) be an input and denote by ys4D saD (u), and yss ss (u). the outputs of an S4D system and an S6 system, respectively. Then, we have
dYSAD Øb
a
C
_ ==1
A
Buj
ab
1=/+1
where
A = exp(exp(b)A)),
B A AA - I)B, I)BB
and
Tyse Ow
1.
a
1.
un C
Σ
II i=j+1
A (us)
B(uz)u
dw
j=1
Tyse
ur Σ
1
à
II (wx) (=1+1
L
B(uj)us
ab
Ob
f=1
where
A(u;)= Á e
Bu,;) A AA'(((a)) - I)Bu,.
Note that the matrix (u,) does not really depend on u, because we assumed that but keep the notation consistent with the proof of Theorem In particular, we have that
ab AAex((b),
A Aeepp()) ABexp(().
ab
and
exp exp (spfffopllusb) (b) Au
a
B(u) = exp (softplus(t)))) Au 1
A (u) =
A -BB
Ow
1 b
dw
A(u) exp (softtlluu b)A A 0b
0
B(u) exp (sofiplus(b) A) A
1 A ABB
l
Øb
Using the product rule, compute rj, Sj. and ty:
A h-j 0
3z=up Uj
B
Σ
A
A
B
ab
0b
i=ト+1
A
1-j-1
Uj
A
B +
Σ
A
B
bb
bb
i=j+1
-A
- Ug exxxp(()
A + (L
A - 1)) B.
25

---

and
ÁL-j a
a A(w))(--BB(s))
8, Uj
B(uz)
+
Σ
dw
dw
im++1
A h-j-1
Uj
B B(u +
Σ i=j+1
A(ws) A(s)) B(ag)
dw
Ew
A 44---1 exp (s))t)lus(b)A)
À B + uz (exp (sottplus(()) I)
Σ Ui
1 +
i=j+1
0
1, = Uj
AL-
B(u;)+
Σ
A
A(,)) A A BB(,)
0b
ob
i=+1
U A 0 B(u,) + uu
AL-111
Σ
( B((x))
0b
0b
i=j+1
Á 1-j-1 exp (softplus(b)A) A Buy + uT j) (exp (b)))plu((b)A) 1) 1)BB
e-b
Increasing input magnitudes. From the formulas of sy. t,. and ry, it is straightforward that they homogeneous in u with degrees 3, 2, and 1, respectively. Fixing an L, as long as EuwD, (0/ (0/0ul y Eu-my (0/ htte/Ob))sssk ysol, and Su..a,((0/00)) (b)Db) * O, we have that
Eu-eD_ (1/0w)) jusol Oce),
EuweD (0/0b) E Eu (0/0b) (bsAA)
O
e- oo.
Ba (0/ ab) ysaD
This proves the first part of the theorem.
Increasing sequence length. The proof of the theorem when L + oo is more involved. For clarity, we break into three parts.
The gradient of YSADD We first consider the case when n = 1. We will show that E (0/00) (xp(((L))O(()) as L - xc. Since is diagonal, when n 2 1, 3S4D can be calculated by summing up the outputs from independent LTI systems, all with 1-dimensional state space. This obviously shows that | (100bb))ssn exp(b(Z))O(VT)) when 1. To prove the case when n 1, we use Jensen's ' inequality and obtain that
L
L A
E
A
Buj
E
Bug
2
E
A
Buj
+ Var
Σ A
Buj
Buj
0
+
Var(u;) +
Σ =
1=1
ISKISL
where Uy A - Buy. BByz This shows that
EnnDe
Ob USsD
FunDy
C
rj
<exp(b(L)) EuNDi
C
A
1-++1
Bu,
+EunDy
C
C
1- (L-))A- -I)uu,
j=1
26

---

Note that we just proved that the first expectation is O(VL) as L - oc. By letting ((L) + -00 fast enough, A vanishes exponentially, have the second term is also (V)L). This proves that
En-D, ((/0b) USaD = exp(K(D))O(VT)
L - oc.
(10)
The gradient of Yse when 1. When n 1, the matrices A, B, and C are all scalars. Without loss of generality, assume that B, C > Similar to the previous case, if b(L) - 0o as - oc, the matrix (xp(softpluss((()))4
vanishes. That is, since us 1, have
Eu-Dz
U /56
Ob
= EurDE
exp(b(L))Eu-D))
c E -
Buy
1=l
1
exxp(b(L)) Σ CAL-IB exp(b(L)) Σ CÁL-IB Var exp(b(L))((L))
provided that ([.) decays fast enough so that A (1)) as - oc.
The gradient of Yse when >> 1. Now. suppose >> 1. The output Use can be written as the sum of terms: y/s = 1/1 +... Un. where yj the output of an L.TI system with n 977. By the previous argument, we know that for every 1 S S n,
0 y Ob C
Eu-Dz
texp(((())O(L))
where the expectation can be either positive or negative, depending the sign of B,C,. That we know that
yj @6 C2
1/(csp// ((cxp(b(LL))
EunDr
bounded between two positive numbers or between two negative numbers as L is sufficiently large. We still want to apply Lemma[to control cancellation effects, but the main issue is that we do not have that this term converges number. However, since the metric space IR is complete, i.e., boundedness implies subsequent convergence, know that there exists subsequence L1, L2.... such that
yi Tb C,
E EuDDLE
as + 0o,
1 1 S n,
where (i is nonzero constant for all 1 S <
Hence, by Lemma
2 we have for a.e. C that
Tyso ab
وy Øb
Bu
EunDing U/
Σ
=x))((((L)))
a + 0.
1=1
Hence, we have
0yso
E Eu~Diy
= exx)((((())))(LL)
(11)
bb
Combining eq. 100 and
we prove the theorem.
E
Supplementary experiments on the stability
In this section, we present two experiments to corroborate our discussion training stability sec- tion 5 The first experiment numerically verifies Theorem as the input magnitude and the sequence length grow. The second experiment empirically shows that a reduced learning rate on A-related parameters help stabilize the training of an S6 model.
27

---

E.1 A numerical experiment that verifies Theorem
In this experiment, we set d and randomly sample diag(A) from the left half of the complex plane, and randomly sample B and C. We then use these matrices to construct an S4D system and an S6 system. We sample our length- L input sequences from i.i.d. Gaussian distributions with mean zero and standard deviation e We fix the ratio L. expxpb((()) (ap(b(.))=1 to take into account the fact that as L increases, we need a longer memory window to capture the long-range dependency. We then compute the quantities in Theorem[3 For the first two experiments, we fix L 100 and let increase. Then, we fix over 30 different trials.
and let L increase. The results are shown in Figure All results are averaged
31%
800
8 Bys4D
00
Bysa
yys 8b
C
Bys6
L
+ OC
Drogs Droea
12f
0/
%
Figure 6: Numerical experiments to verify Theorem 3 where compute the ratio between the gradients with respect to the S6 parameters and S4D parameters. For the first two figures, we fix L 100; for the last figure, we fix c The gradients are computed using closed algebraic formulas. The red reference lines in the three log-log plots have slopes of 3, 2, and 1/2, respectively.
In these experiments, we see that the three ratios studied in Theorem 3 follow exactly the pattern of a cubic, quadratic, and square root growth, respectively. This is what expect from Theorem 3
E.2 An empirical experiment on the training stability
We also show a real training example, where we want to learn a ground-truth function that takes in univariate sequential input and returns fixed linear combination of them: G(ui..... ..,)..
Σ/=1 Ojuy. where 0,,.... O, are fixed parameters. We show the loss curves in Figure 7 with different models (i.e., S4D) versus S6) and different learning rate assignments (i.e., whether not is learned). We see that the S4D model very stable during training, with its loss going down smoothly. This is not the case for the S6 model, where the loss goes up and down and even restarts.
We find that by freezing the At parameters w and b, the model trains much more robustly.
S4D
S6
S6 (fixed A)
15
loss
loss
5-
0
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
2000
4000
8000
8000
10000
Step
Stop
Step
Figure 7: The poo--mean-squaree loss [[G(()) - G(u)l between the true output Guu) and the model prediction G()). The first two models are trained with a learning rate of 0.001 on the A parameters, whereas the last model trained with no training of the A parameters.
'Note, in particular, that this sequence has no long-range dependency. We made this choice because it is a natural one without any prior information. It does not affect the training stability lot. One can try different input functions with long-range dependencies, e.g., have tried sinusoidal waves, and the results are similar.
28

---

Details of experiments
In this section, we show the details of the experiments in section 7 Before we provide the detailed configurations, we provide an extended table for LRA results of more variants of state-space models (see Table 4. While many of these models focus on improving certain aspects of the recurrent unit, we do not meorporate and test them in our ByS; model.
Table 4: An extended list of SSM accuracies on the Long-Range Arena benchmark. An entry is left blank no result is found. All but the last three models are input-selective.
Model
ListOps
Text
Retrieval
Image
Pathf inder
Path-X
Avg.
S4
[28
59.60
86.82
90.90
88.65
94.20
96.35
86.09
S4D
127
60.47
86.18
89.46
88.19
93.06
91.95
84.89
DSS
29
57.60
76.60
87.60
85.80
84.10
85.00
79.45
LRU [55
89.00
60.20
89.40
89.90
95.10
94.20
86.30
S4++
57.30
86.28
84.82
82.91
80.24
HOPE-SSM [98]
62.60
89.83
91.80
88.68
95.73
98.45
87.85
S4D-FT 997
62.75
89.76
92.45
90.89
95.89
97.84
88.26
Reg. S4D
46
61.48
88.19
91.25
88.12
94.93
95.63
86.60
Liquid $4
330
62.75
89.02
91.20
89.50
94.80
96.66
87.32
RTF SSM [56
61.59
89.72
92.04
90.51
96.11
96.32
87.71
Spectral SSM
60.33
89.60
90.00
95.60
90.10
Spiking SSM
72
60.23
80.41
88.77
88.21
93.51
94.82
84.33
S5 [77
62.15
89.31
91.40
88.00
95.33
98.58
87.46
S6 (Mamba) 22
38.02
82.98
72.14
69.82
69.26
67.32
66.59
S7 [79
63.77
87.22
91.80
61.14
65.62
61.50
71.82
B2S, (ours)
63.85
88.32
91.44
88.81
95.93
97.90
87.71
In this paper, all models are trained with one more NVIDIA L40 GPUs with 48GB of memory. For the ablation study, we use a 4-layer model and we increase the number of layers for the full experiments on the LRA benchmark. We provide the details of the model and training hyperparameters used for training each LRA task in T For experiments, we set h = 8 so that p #Fcatures 8 Notably, this hyperparameter is not carefully fine-tuned but rather picked randomly. Note that model for training the Path-X tasks is smaller than the corresponding S4D model.
Table 5: Configurations of our B2S, model on the LRA benchmark, where LR, BS, and WD stand for learning rate, batch size, and weight decay, respectively.
Task
Depth
#Features
Norm
Prenorm
LR
BS
Epochs
WD
ListOps
128
BN
False
0.008
32
120
0.03
Text
6
256
BN
True
0.01
16
40
0.05
Retrieval
6
128
BN
True
0.004
60
200
0.03
Image
6
512
LN
False
0.01
48
500
0.05
Pathfinder
256
BN
True
0.004
48
250
0.03
Path-X
6
128
BN
True
0.001
24
120
0.03
G Limitations and Future Work
We acknowledge (and justify) several limitations of this work, which also suggest promising directions for future research:
The universal approximation theorems (UATs) in this paper are derived under a few simplifying assumptions - most notably, the assumption that the sampling interval A is fixed. This setup diverges from practical implementations, where A is input-dependent. However, our goal is not to prove UAT's for all realistic settings, but to expose fundamental expressiveness gaps between S6 and S4D under clean conditions. Extending these results to settings with dynamic A remains interesting theoretical direction. If one show that single-layer Mamba is not universal

---

approximator even when A is not fixed. then that would further stress the benefit of a multi-head design in selective SSMs. Conversely, if a single-layer Mamba with dynamic is a universal approximator, then an important follow-up question is if (and how) trainable channel-specific A can compensate for the channel-independent B and matrices single-head Mamba.
Our experiments language modeling are not designed compete with large-scale models datasets. This choice reflects a deliberate focus: our primary aim is to study and improve the behavior of Mamba on long-range sequence tasks. Demonstrating that B2Sg retains Mamba's performance in language modeling - even at small scale supports its versatility. Scaling to larger models and corpora is natural next step, pending access greater computational resources.
Our implementation of B2S; is based entirely PyTorch. While this makes the method easily accessible and reproducible, further speed improvements would require low-level optimization in CUDA (see (25)). We leave efficient implementation and integration into high-performance inference/training libraries to future work.
30

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and the introduction clearly state the main claims and contributions the paper. It is made clear that the paper focuses on investigating Mamba's capability of modeling long-range sequences, and that the paper analyzes the previous limitations from three aspects (expressiveness, inductive bias, and training stability) and proposes a model that better captures long-range dependencies. All claims supported by both theoretical statements and experiments.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitations of this paper in Appendix G These limitations are also discussed in the main paper, though not in a centralized way.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model vell-specifiiattio, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
31

---

3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: All assumptions are clearly stated, with notations previously defined in the paper clearly cross-referenced. All proof' are provided in the appendices.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide details of all experiments in Appendix F While it is infeasible to state all details of the training and model configurations, we provide anonymous code to enhance the reproducibility of the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
32

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide anony code with our submission. They come along with clear instructions, allowing for faithful reproducibility. The datasets are external and instructions are given to download them.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ /nips.ccc/pullc//uuui
Cooe/Sbmeeeeeeeee wwwiibbmmsssiinni ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if' applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide the table of hyperparameters in Appendix While we do not specify all training and test details, including the data splits and optimizers, most of the experiments come with standard choices, and they can be found in our anonymous code repository.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
33

---

We repeat our ablation study and Long- -angee Arena experiments with multiple random seeds and report the standard deviations in section
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Details of the compute resources are given in Appendix F We do not measure the time of execution of the Long-Range Arena experiments as they not standardly reported in the state-space models community.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https:
Answer: [Yes]
Justification:
The work complies with the Code of Ethies.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
34

---

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work is theoretical and methodological It does not involve large-scale implementation and experimentation and has very limited societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
The paper releases no data or models that have a high risk of misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The existing assets are both acknowledged Appendix F and the anonymous code repository.
Guidelines:
35

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paaassttssppersseaaee paperswi thcode. com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [Yes]
Justification: The anonymous code repository we provide is well-documented and contains a README file that explains the organization of the repository and steps needed to reproduce our experimental results.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
36

---

Answer:
[NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs are only used for word choices and limited polishing of the writing. This research does involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not non involve LLMs as any important, original, non- ntn-standard components.
Please refer to our LLM policy for what should or should not be described.
37


Reviews:
[{"review_id": "8yFuYbKnuu", "reviewer": "Reviewer_rPYn", "summary": "The paper presents a study of current state space models (SSMs) like S4D and Mamba. As a result a number of weka points are discovered: training stability, poor long-range performance, inductive bias and expressiveness. To compensate for these shortcomings, a new architecture is proposed named as B2S6. The proposed architecture combines block-wise selective dynamics (as in Mamba1) with a channel-specific bias. Most evaluations are done on Long-Range Arena.", "questions": "- How does it perform in the hybrid settings (combined with Attention)? Do we need a new head design, or attention can compensate for it? \n- Propose to create a table with a summary of differences between S4D, Mamba, B2S6 and Mamba2.\n- For Figure 2, can the failure of Mamba be explained by not optimal learning rate, or should we increase number of heads even more?\n- What is the impact on speed of using B2S6 vs Mamba1/2?", "limitations": "Limitations are mentioned in the Supplementary Materials.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "- The paper contains a comprehensive analysis of the current state of Mamba1.\n- Key limitations of Mamba's ability to perform well on long-range dependencies are summarized as a lack of expressiveness, inductive bias, and training stability.\n- The background and recap of S4D and S6 are well covered in Section 2.\n- The analysis provided in Section 3 is interesting and highlights the insights mentioned above. Including Table 1 and Figure 2 is great.\n- Section 4 is well organized.\n- Section 5 analyzes the training stability of Mamba. From personal experience, we have not encountered issues when training models with Mamba1. How robust are the conclusions in this section to applying weight decay or using pseudo second-order optimization frameworks like Adam?\n- B2S6, introduced in Section 6, appears novel and well-grounded. However, more discussion is needed regarding comparisons to Mamba2. Lines 253–255 suggest that the proposed architecture is very close to Mamba2, though the authors claim to provide additional insights. Is this accurate?\n \n\nWeaknesses:\n1) The experimental analysis is weak—it takes less than one page in the main paper. It is important to assess whether B2S6 scales to LLMs and whether observations from the small-scale model will generalize to larger models. Training instabilities might emerge.\n2) The intuition behind incorporating per-channel separation in B2S6 is related to the Mamba2 extension proposed in [20]. However, the current evaluation lacks comparison to [20].\n3) Overall, the paper lacks a deeper analysis of the proposed model and B2S6. For example, Mamba2 should be included in Figure 2 and Table 1.\nIdeally, B2S6 should be evaluated in an LLM setting (at least a 1–2B parameter model trained on 200B tokens) and compared to other works mentioned in the paper, such as Mamba1. Current evaluations are insufficient. I recommend the authors use lm-eval-harness to obtain more comprehensive evaluation benchmarks.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the careful review and insightful comments. Below we address the questions and comments raised in this review.\n\n* ### Multi-head Design in a Hybrid Model (Q1)\n\n  This is a very interesting question. In a hybrid model, such as Jamba, Mamba layers and attention layers induce different biases and are used to capture different dependencies in the sequence of tokens (e.g., attention is good at handling short-term, high-salience interactions while Mambas are good for long-range dependencies). That said, from a universal approximation perspective, the attention layers alone are capable of learning most smooth ground truth functions; in practice, however, the multi-head design in attention layers cannot totally compensate for the lack of heads in a Mamba layer in order to better learn all aspects of the sequence.\n\n  Another interesting thing to point out is that, even though the word \"multi-head\" is used for both attention and Mamba, the precise mechanisms are slightly different. For a multi-head attention, the input into each \"head\" is the entire sequence; hence, increasing the number of heads mainly introduces a trade-off between the variance and the computational time. For Mamba, however, the input into each \"head\" is only a sub-block of the sequence, consisting of a subset of channels, and this induces a trade-off between the variance (in terms of the number of different time-variant systems we have) and bias (in terms of how much we can achieve with each system). Hence, we advocate an independent multi-head design for Mamba layers and attention layers in a hybrid model.\n\n* ### Conceptual Comparison with Mamba2 (Q2 + S7)\n\n  Thank you for your comment. We agree that a table will improve the clarity of our method. We plan to include the following table to summarize the differences between each state-space model.\n\n  | | S4D | S5 | Mamba | Mamba2 | B2S6 |\n  | --- | :---: | :---: | :---: | :---: | :---: |\n  | $\\Delta_t$ Dependency | input-independent | input-independent | input-dependent  | input-dependent | input-dependent |\n  | $\\mathbf{B}$ / $\\mathbf{C}$ Dependency | input-independent | input-independent | input-dependent | input-dependent | input-dependent + input-independent |\n  | Number of Heads | $\\geq 1$ | $\\geq 1$ | $= 1$ | $\\geq 1$ | $\\geq 1$ | \n  | Parameters | complex | complex | real | real | complex\n  | Parameterization of $\\mathbf{A}$ | diagonal | diagonal | diagonal | scalar | diagonal|\n\n  Also, we confirm that our discussion on Line 253–255 suggests that the multi-head design is not new to SSMs, but in this paper we provide further theoretical analysis to motivate it.\n\n* ### Additional Experimental Comparison with Mamba2 (W2 + W3)\n\n  We are happy to run some experiments to incorporate Mamba2 in our comparison. The following table compares Mamba2 with Mamba1 and B2S6 on the LRA benchmark:\n\n  | | ListOps | Text | Retrieval | Image | Pathfinder | Path-X | Avg. |\n  | --- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n  | Mamba |38.02| 82.98| 72.14| 69.82| 69.26| 67.32| 66.59|\n  | Mamba2 |41.45| 86.09| 79.23| 71.96 | 75.45| 69.07 | 70.54 |\n  | B2S6 |63.85| 88.32| 91.44| 88.81| 95.93| 97.90| 87.71|\n\n  Moreover, we also pretrained the Mamba2 model following the design in Appendix G. The table below shows the perplexity results in addition to Figure 7. As shown in the table, B2S6 converges slightly slower initially than both Mamba and Mamba2, and catches up in the end. In this small example, B2S6 shows a matching performance.\n\n  |Steps | Mamba | Mamba2 | B2S6 |\n  | --- | :---: | :---: | :---: |\n  |2620 | 61.5441 | 59.6233 | 62.6133 |\n  |5241 | 42.7838 | 39.8305 | 44.3846 |\n  |7862 | 34.1740 | 31.9994 | 35.0440 |\n  |10483 | 29.7074 | 27.0987 | 30.3787 |\n  |13104 | 27.0171 | 25.2930 | 27.3891 |\n  |15725 | 25.2265 | 24.0179 | 25.2668 |\n  |18346 | 23.9789 | 23.1224 | 23.9370 |\n  |20967 | 23.0044 | 22.3695 | 22.8394 |\n  |23588 | 22.2368 | 21.7509 | 22.1322 |\n  |26209 | 21.6068 | 21.2480 | 21.5206 |\n  |28830 | 21.0785 | 20.8218 | 20.9773 |\n\n  We also reran the experiment in Figure 2 with Mamba2. While the results are hard to deliver in a table, we confirm that Mamba2, unlike Mamba1, does benefit from an increasing hidden dimension $d$. We provide more explanation about this below, and if our paper is accepted, we will include all these results in the camera-ready submission.\n\n* ### Training a Large-scale Model (W1)\n\n  This is a very good point. While we spent most of the rebuttal period running additional experiments with Mamba2, we will start training some larger models on a large-scale dataset `SlimPajama-627B` and evaluate them using `lm-eval-harness` and be sure to incorporate these results in the camera-ready submission if the paper is accepted.\n\n* ### Explaining Figure 2 (Q3)\n\n  This is an interesting point, and we have done some further hyperparameter tuning over a grid of 10 learning rates, ranging from $\\texttt{1e-4}$ to $\\texttt{1e-1}$. We observed that none of these learning rates are able to save S6's performance on this task. Hence, we conclude that this is mainly an issue with a lack of the number of heads. Note that the only difference between models in the middle panel and right panel is that the B2S6 model has more heads than Mamba, which is a single-head model. (On Line 271-272, we clarify this by noting that \"we revisit the experiment from section 3, this time using only the block structure (without the bias term).\")\n\n  Another interesting point to make here is that the S4D model performs better than the B2S6 model on this task, and this is true even if we set the number of heads in B2S6 to be equivalent to the hidden dimension $d$. The reason is that the LTI systems in an S4D model are intrinsically more capable of learning a single Fourier mode (see [96] in the bibliography). Hence, the primary lesson from Figure 2 is not that S4D is better than S6 because it has more \"heads,\" but that increasing the number of heads from 1 will enhance Mamba's expressiveness. We will make this clear in our paper.\n\n* ### Efficiency Comparision (Q4)\n\n  We report the training throughputs (tokens/sec) of Mamba and B2S6 models. We remark that the code we used for training the Mamba and B2S6 models is purely PyTorch-based and does not involve any CUDA optimization. We mentioned this in the limitation section, and once the CUDA optimization is incorporated into the B2S6 implementation, its efficiency will be improved.\n\n  | Model Size | Mamba | B2S6 |\n  | :--- | :---: | :---: |\n  | 250M | 14,667 | 14,218 |\n  | 370M | 9,914 | 9,676 |\n  | 500M | 7,218 | 7,010 |\n\n* ### Applicability of the Stability Analysis (S6)\n\n  This is another interesting thing to comment on. In training a Mamba-based LLM, we indeed have not observed instability induced by $\\Delta$. However, when training Mamba models on LRA benchmark tasks, without carefully reducing the learning rate of $\\Delta$, the training is very volatile. One reason is that the models we train on LRA are much smaller than LLMs (and hence, the loss landscapes are less smooth than those of a larger model), and the sequences in LRA tasks can also be very long, e.g., in Path-X, a sequence has a length of 16,384. This is very large with respect to the size of the model we trained.\n\n  You are absolutely right to think of Adam and weight decay as ways to speed up the convergence and stabilize the training. Our analysis in Theorem 3, on the other hand, is independent of the optimizer or the learning rate scheduler. Since we aim to analyze the stability without having a specific loss function, dataset, optimizer, or scheduler in mind, we go for the most general strategy and investigate the gradient norm. In general, a larger gradient norm leads to less stable training, and this holds for any optimizer and scheduler. In that sense, the Adam optimizer, weight decay, and potentially gradient norm clipping and learning rate reduction can all be viewed as efforts to counterbalance the stability issue raised in Theorem 3.\n\nThank you again for your careful review and constructive comments. We are committed to improving this manuscript, and please let us know if you have any other questions!"}, {"review_id": "9xZLd395ji", "reviewer": "Reviewer_8YXd", "summary": "Driven from the concrete motivation -- why Mamba performs poorly on long range tasks -- the paper points out a simple solution: use of bias. The authors theoretically investigate critical limitations (supported by proofs) of S6 such as less expressivity with a single layer than S4D, strong inductive bias which can be particularly harmful for long sequences,  and instability. The paper empirically validates that simply adding group-wise (similar to Mamba2) input-independent bias, S6 becomes much stronger in long range modeling while not hurting the language modeling performance.", "questions": "Does the same effects arise in Mamba2 SSD kernel too? What's the main reason of choosing S6 to compare against instead of SSD?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The paper is very well written, and I really appreciated its strong motivation and simple solution. In high-level, the addition of bias harmonizes the strengths of both S4D and S6, making Mamba become much more powerful. The claims are well supported in both theoretical and empirical validations, making the paper more credible. A weakness (which is a nitpick) is lack of evaluation on language modeling benchmarks (and the training scale is too small: 6B tokens), and it would be interesting to see if this simple change can lead to gains in real-world language modeling capacities such as retrieval.", "was_revised": true, "final_justification": "Most of my concerns are addressed during the rebuttal, thus I maintain my score.", "rebuttal": "We thank the reviewer for the careful review and insightful comments.\n\nWe chose Mamba1 as a reference model mainly because it is the first selective state-space model proposed. Nonetheless, we are happy to include additional comparison of our B2S6 model with Mamba2. First, conceptually, as suggested in section 6 of the paper, the multi-head structure is implemented in Mamba2, so our main contributions are theoretical analysis of the multi-head structure, together with a new design of a bias term $\\mathbf{B}_{\\text{bias}}$. We will add the following table to our manuscript to further compare the design choices of different state-space models in a clear way:\n\n| | S4D | S5 | Mamba | Mamba2 | B2S6 |\n| --- | :---: | :---: | :---: | :---: | :---: |\n| $\\Delta_t$ Dependency | input-independent | input-independent | input-dependent  | input-dependent | input-dependent |\n| $\\mathbf{B}$ / $\\mathbf{C}$ Dependency | input-independent | input-independent | input-dependent | input-dependent | input-dependent + input-independent |\n| Number of Heads | $\\geq 1$ | $\\geq 1$ | $= 1$ | $\\geq 1$ | $\\geq 1$ | \n| Parameters | complex | complex | real | real | complex\n| Parameterization of $\\mathbf{A}$ | diagonal | diagonal | diagonal | scalar | diagonal|\n\nOn the experimental side, we have trained Mamba2 models on the LRA benchmark during the rebuttal period, and the following table summarizes the results:\n\n| | ListOps | Text | Retrieval | Image | Pathfinder | Path-X | Avg. |\n| --- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| Mamba |38.02| 82.98| 72.14| 69.82| 69.26| 67.32| 66.59|\n| Mamba2 |41.45| 86.09| 79.23| 71.96 | 75.45| 69.07 | 70.54 |\n| B2S6 |63.85| 88.32| 91.44| 88.81| 95.93| 97.90| 87.71|\n\nMoreover, we also pretrained the Mamba2 model following the design in Appendix G. The table below shows the perplexity results in addition to Figure 7. As shown in the table, B2S6 converges slightly slower initially than both Mamba and Mamba2, and caught up in the end. In this small example, B2S6 shows a matching performance.\n\n|Steps | Mamba | Mamba2 | B2S6 |\n| --- | :---: | :---: | :---: |\n|2620 | 61.5441 | 59.6233 | 62.6133 |\n|5241 | 42.7838 | 39.8305 | 44.3846 |\n|7862 | 34.1740 | 31.9994 | 35.0440 |\n|10483 | 29.7074 | 27.0987 | 30.3787 |\n|13104 | 27.0171 | 25.2930 | 27.3891 |\n|15725 | 25.2265 | 24.0179 | 25.2668 |\n|18346 | 23.9789 | 23.1224 | 23.9370 |\n|20967 | 23.0044 | 22.3695 | 22.8394 |\n|23588 | 22.2368 | 21.7509 | 22.1322 |\n|26209 | 21.6068 | 21.2480 | 21.5206 |\n|28830 | 21.0785 | 20.8218 | 20.9773 |\n\nWe hope these experiments enhance the comprehensiveness of the manuscript. We also agree with the reviewer that the paper would benefit more from larger-scale experiments on language datasets. To this end, we will start training some larger B2S6-based LLMs on a large-scale dataset `SlimPajama-627B` and be sure to incorporate these results in the camera-ready submission if the paper is accepted.\n\nThank you again for your careful review and constructive comments. We are committed to improving this manuscript, and please let us know if you have any other questions!"}, {"review_id": "PoBSmcxmyj", "reviewer": "Reviewer_y9vW", "summary": "This paper identifies and addresses a critical limitation of Mamba models, their underperformance on long-range sequence tasks despite architectural suitability. The authors attribute this to three factors: limited expressiveness due to channel-wise parameter sharing (Theorem 1), extreme input-dependent inductive bias causing premature forgetting of long-range information (Theorem 2), and training instability on long sequences (Theorem 3). To resolve these, they propose Block-Biased S6 (B2S6), which partitions inputs into blocks for multi-head selective dynamics and adds channel-specific bias terms. Theoretically, B2S6 achieves universal approximation (Theorem 4) and balanced inductive bias (Theorem 5). Empirically, it outperforms S4/S4D on Long-Range Arena (LRA) (Table 2) while matching Mamba’s language modeling performance on SlimPajama.", "questions": "1. Mamba incorporates parallel computing for scaling to larger language model training. Could B2S6's design hinder parallel computation?\n\n2. Could B2S6 match or surpass Mamba in large-scale language modeling (e.g., 1B+ parameters, 100B+ tokens)? \n\n3. Would the universal approximation of B2S6 (Theorem 4) hold with input-dependent Δ (as in standard Mamba)? \n \n4. Section 5 attributes instability to $\\Delta$ optimization. Does freezing other parameters also help? \n\n5. Does Mamba2 suffer from the same Expressiveness, Inductive Bias, and Training Stability issues? Could the same toy model experiments (like those in Figures 2 and 3) be used to test Mamba2?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**:\n1. The paper provides a comprehensive theoretical diagnosis of Mamba’s limitations, supported by formal theorems on expressiveness, bias, and stability. The non-universal approximation proof for S6 (Theorem 1) and the bias characterization (Theorem 2) are particularly insightful.\n2. B2S6 elegantly combines block-wise processing with channel-specific bias, addressing all three limitations. Theorems 4–5 formally validate its expressiveness and bias correction.\n3. B2S6 achieves competitive performance on LRA (e.g., +18.5% avg. gain over Mamba in Table 2) while preserving language modeling efficacy (Fig. 7). Ablations (Table 1) and synthetic tasks (Figs. 2–3) robustly justify design choices.\n4. Detailed proofs (Appendices B–D), hyperparameters (Table 4), and open code enhance reproducibility.\n5. The toy model experiments and theory validate each other very well.\n\n**Weakness:**:\n1. Experiments use a small SlimPajama subset (6B tokens) and a 250M-parameter model. Validation on larger model scales is needed to confirm generality.\n\n2. Theorems assume fixed $\\Delta$, while practical Mamba uses input-dependent $\\Delta$. While simplifying proofs, this may weaken the real-world relevance of expressiveness claims.\n\n3. The computational cost of block partitioning and bias terms is unquantified. Comparisons against efficient S4/S4D and Mamba in latency/throughput would strengthen practicality.\n\n4. minor typos: Lines 157 & 159: Standardized \"k'th\" and \"kth\" → \"k-th\".", "was_revised": true, "final_justification": "The author addressed my concern. I recommend accepting the paper.", "rebuttal": "We thank the reviewer for the careful review and insightful comments. Below we address the questions and comments raised in this review.\n\n* ### Efficiency of B2S6 (Q1 + W3)\n\n  We report the throughputs (tokens/sec) of Mamba and B2S6 models below. The block design in B2S6 does not hinder parallel computing, as shown in the Mamba2 paper [20]. The bias term $\\mathbf{B}\\_{\\text{bias}}$ is benign either, because the Mamba algorithm computes the discrete matrices ${^j}\\overline{\\mathbf{A}}$ and $^j\\overline{\\mathbf{B}}$ of shape $(B \\\\; L \\\\; p \\\\; n)$ (see Algorithm 2) and one can simply broadcast $^j\\mathbf{B}_{\\text{bias}}$ to $^j\\mathbf{B}$ to achieve this. In the original Mamba implementation, however, this step is done in the CUDA code; our implementation right now is purely PyTorch and does not involve any CUDA optimization. The throughputs reported below are based on comparing pure PyTorch implementations of the models. We mentioned this in the limitation section.\n\n  | Model Size | Mamba | B2S6 |\n  | :--- | :---: | :---: |\n  | 250M | 14,667 | 14,218 |\n  | 370M | 9,914 | 9,676 |\n  | 500M | 7,218 | 7,010 |\n\n* ### Relaxing the Fixed $\\Delta$ Assumption (Q3 + W2)\n\n  It is an interesting question whether Mamba with a trainable $\\Delta$ is a universal approximator. While relaxing $\\Delta$ breaks our non-UAT proof, it is hard either to directly leverage some conventional wisdom to prove a UAT result for Mamba with trainable $\\Delta$, as we did for S4D models. Nonetheless, as we mentioned in the manuscript, whether or not a single-layer Mamba (with or without a fixed $\\Delta$) is a universal approximator is not a practically important question anyway; instead, the moral of the Theorem 1 is that it highlights that a single-head Mamba, while extends the S4D model by making the system time-variant, also loses the channel-dependent $\\mathbf{B}$ and $\\mathbf{C}$; therefore, they should not be thought of as a direct extension of S4Ds. This causes some convergence, generalization, and training stability issues, as pointed out on Line 122-127. For the theoretical community, it is indeed important to stress that the non-UAT result we proved holds only in a restricted setting -- we thank the reviewer for bringing this up and will be sure to include a more justified discussion of the theorem (see our response to Reviewer PVYx).\n\n* ### Stability of B2S6 (Q4)\n\n  This is another very interesting question. In general, freezing some of the other parameters indeed helps stabilize the training, but it also puts us at risk of converging to bad minima. We have done some empirical analysis and would be happy to provide some general observations we have made:\n\n  1. It never works to freeze $\\mathbf{B}$ and $\\mathbf{C}$. These two matrices are in some sense the most crucial ones that define a system and play the same role as the weight matrix $\\mathbf{W}$ in an MLP, and freezing these two parameters will lead to a significant performance drop.\n\n  2. For LRA tasks, freezing the matrix $\\mathbf{A}$ usually does not impact the performance by a lot. For instance, if the matrix $\\mathbf{A}$ is properly initialized using HiPPO-LegS [26] or its scaled variants [96], then the performance is dropped by less than $2\\\\%$ when $\\mathbf{A}$ is frozen. In the meantime, as pointed out in [26], freezing $\\mathbf{A}$ or reducing its learning rate indeed stabilizes the training. For pretraining LLMs, training the matrix $\\mathbf{A}$ rarely causes a stability issue, so there is no need to freeze the parameter.\n\n  3. Freezing $\\mathbf{D}$ usually does not have a big impact, on either the training stability or the performance. This parameter introduces a skip connection and causes fewer issues (while also being less important) than other parameters.\n\n  We hope this summary is useful in answering the reviewer's question. We will include this discussion in our camera-ready submission if the paper is accepted.\n\n* ### Additional Comparison with Mamba2 (Q5)\n\n  We thank the reviewer for raising a discussion of Mamba2 as we believe this is an interesting model to compare with. To better compare B2S6 across different models, we will include the following table in our camera-ready submission if the paper is accepted.\n\n  | | S4D | S5 | Mamba | Mamba2 | B2S6 |\n  | --- | :---: | :---: | :---: | :---: | :---: |\n  | $\\Delta_t$ Dependency | input-independent | input-independent | input-dependent  | input-dependent | input-dependent |\n  | $\\mathbf{B}$ / $\\mathbf{C}$ Dependency | input-independent | input-independent | input-dependent | input-dependent | input-dependent + input-independent |\n  | Number of Heads | $\\geq 1$ | $\\geq 1$ | $= 1$ | $\\geq 1$ | $\\geq 1$ | \n  | Parameters | complex | complex | real | real | complex\n  | Parameterization of $\\mathbf{A}$ | diagonal | diagonal | diagonal | scalar | diagonal|\n\n\n  As we indicated in the discussion after Theorem 4 and 5, the multi-head design will enhance the expressiveness of the model and equip it with a gentler inductive bias. Given that, Mamba2 is expected to be more expressive than Mamba1. We have rerun the experiments in Figure 2 and 3 with Mamba2. While the results are hard to present without being able to upload figures, we confirm that Mamba2 performs quite as good as B2S6 in the experiment in Figure 2. This is not surprising because in this experiment, we remove the bias term from B2S6 to avoid \"cheating,\" so Mamba2 is very close to B2S6. Mamba2's performance is in between that of Mamba1 and B2S6 in the experiment in Figure 3, which shows the advantage of using the bias term $\\mathbf{B}_{\\text{bias}}$ in B2S6.\n\n  Mamba2 still suffers from the stability issue, so when trained on LRA, one needs to reduce the learning rate, just as B2S6. Moreover, without complex parameterization, Mamba2 falls short significantly on the LRA benchmark. The following results were obtained during the rebuttal period.\n\n  | | ListOps | Text | Retrieval | Image | Pathfinder | Path-X | Avg. |\n  | --- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n  | Mamba |38.02| 82.98| 72.14| 69.82| 69.26| 67.32| 66.59|\n  | Mamba2 |41.45| 86.09| 79.23| 71.96 | 75.45| 69.07 | 70.54 |\n  | B2S6 |63.85| 88.32| 91.44| 88.81| 95.93| 97.90| 87.71|\n\n  Moreover, we also pretrained the Mamba2 model following the design in Appendix G. The table below shows the perplexity results in addition to Figure 7. As shown in the table, B2S6 converges slightly slower initially than both Mamba and Mamba2, and caught up in the end. In this small example, B2S6 shows a matching performance.\n\n  |Steps | Mamba | Mamba2 | B2S6 |\n  | --- | :---: | :---: | :---: |\n  |2620 | 61.5441 | 59.6233 | 62.6133 |\n  |5241 | 42.7838 | 39.8305 | 44.3846 |\n  |7862 | 34.1740 | 31.9994 | 35.0440 |\n  |10483 | 29.7074 | 27.0987 | 30.3787 |\n  |13104 | 27.0171 | 25.2930 | 27.3891 |\n  |15725 | 25.2265 | 24.0179 | 25.2668 |\n  |18346 | 23.9789 | 23.1224 | 23.9370 |\n  |20967 | 23.0044 | 22.3695 | 22.8394 |\n  |23588 | 22.2368 | 21.7509 | 22.1322 |\n  |26209 | 21.6068 | 21.2480 | 21.5206 |\n  |28830 | 21.0785 | 20.8218 | 20.9773 |\n\n  We will incorporate these new results and discussion into our paper. We hope they enhance the comprehensiveness of our discussion.\n\n* ### Training a Large-Scale Model (Q2 + W1)\n\n  This is a very good point. While we spent most of the rebuttal period running additional experiments with Mamba2, we will start training some larger models on a large-scale dataset `SlimPajama-627B` and test them using `lm-eval-harness`, as Reviewer rPYn suggests, and be sure to incorporate these results into the camera-ready submission if the paper is accepted.\n\n* ### Minor Typo (W4)\n\n  Thank you also for pointing out the inconsistency on Line 157 and 159. We will change both into \"$k$-th.\"\n\nThank you again for your careful review and constructive comments. We are committed to improving this manuscript, and please let us know if you have any other questions!"}, {"review_id": "SJ88C92C83", "reviewer": "Reviewer_PVYx", "summary": "This paper introduces Block-Biased-S6 (B2S6), a modification of the S6 recurrent layer from Mamba, that adds a channel specific bias term to the linear map acting on the input sequence, $\\mathbf{B}_w\\mathbf{u}_k$$\\rightarrow$$\\mathbf{B}_w\\mathbf{u}_k+\\mathbf{B}_b$. Additionally, the model adopts the block unit design suggested in the Mamba2 paper. \n\nThe paper motivates these changes through a number of theoretical results and toy experiments, which explore the expressivity, inductive bias, and training stability of Mamba. They also demonstrate that B2S6 achieves superior performance on the long range arena benchmark compared to other selective state-space models.", "questions": "## Questions\n\n1) Can a result similar to Theorem 2 be proved when you assume that a single term in the input dominates, whilst the norm of the input remains fixed?\n\n## Suggestions\n\n1) The discussion of the theoretical results should be modified to align better with the contents of the Theorems. Most critically, I do not believe your theoretical results are sufficient to claim that \"Mamba does not have universal approximation properties\" (line 35) or that \"S6 has structurally limited width compared to S4D\" even if $\\Delta^{(i)}_k$ is dynamic (line 949). Given this adjustment, I would be happy to raise my score.", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n- The paper is well written, I found no errors in the proofs of the theoretical results, and the theoretical results are explored empirically using suitable toy experiments. \n- The ablation study is a good demonstration of the benefits of including the bias term and the trade-off between the number of independent blocks and the dimension of the input to each block.\n- The empirical results on the long range arena benchmark demonstrate the benefits of using a block unit design and including the bias term $\\mathbf{B}_b$. \n\n## Weaknesses\n\n### Theory\n\nThe theoretical results of the paper are interesting, but the resulting claims about Mamba are overstated. In particular:\n- \"A single-layer Mamba is not a universal approximator\" does not follow from Theorem 1. In Theorem 1, the channel specific dynamics of Mamba have been removed by fixing the $\\Delta^{(i)}_k$, reducing the model to being identical copies in each channel. As the author's note on line 124, this obviously hurts expressivity. Removing the channel specific dynamics from Mamba also weakens Theorem 4 and the motivation for the channel specific bias in B2S6, as you are adding in an alternative type of channel specific dynamics, motivated by removing the channel specific dynamics that were already present. From the empirical results in Figure 2 and Table 2, the channel specific dynamics from $\\Delta^{(i)}_k$ and $\\mathbf{B}_b$ clearly have different effects, but this is not explored theoretically.\n- Similarly, \"Mambas exhibit strong inductive bias\" does not follow from Theorem 2, which relies on one diverging sample $c\\mathbf{u}_{k_0}$, a case which does not arise in practice due to normalisation. Theorem 2 would be much stronger if you could instead assume that one term dominates the others whilst the norm of the input remains fixed. \n\nAdditionally, on line 672 you state that \"To the best of our knowledge, however, very little has been done to establish universal approximation properties for Mamba models\". However, a paper you cite, [52], established the explicit expressivity of Mamba in the continuous setting without restricting the model to fixed $\\Delta^{(i)}_k$. There should be a discussion of your expressivity results in relation to the results presented in [52].\n\n### Experiments\n\n- The bold and underlined results in Table 2 are slightly disingenuous, given that the three models which achieve a higher average test accuracy are relegated to the Appendix. However, this does not weaken the claim of the paper, as B2S6 is still the best performing selective state-space model.", "was_revised": true, "final_justification": "The authors have resolved my key concern, which was that their claims about Mamba were not sufficiently backed up by their theoretical results. Given the interesting theoretical results, strong empirical results, and quality of the manuscript, I am now happy to recommend acceptance.", "rebuttal": "We thank the reviewer for the careful review and insightful comments. Below we address the questions and comments raised in this review.\n\n* ### Overclaims in Theorem Discussions (W1 + Suggestion)\n\n  We totally agree that the fixed $\\Delta$ in Theorem 1 is a strong assumption, and that a few overclaims were made surrounding it. Since we are unable to submit a rebuttal revision, we outline our planned modifications below based on your suggestion. We hope these adjustments resolve your concern.\n\n  1. Line 32-35 will be changed into\n\n  > **Expressiveness.** Unlike S4D, which allows each internal channel to learn an independent recurrent unit, Mamba shares parameters across all channels. Therefore, one should not interpret Mamba as a straightforward extension of S4D that gains time-variant benefits at no cost. Rather, Theorem 1 shows that while Mamba introduces input-dependent dynamics, it also sacrifices certain degrees of freedom. Most notably, it has a weaker ability to learn independent behavior in each channel.\n\n  2. Line 948-950 will be changed into\n\n  > Extending these results to settings with dynamic $\\Delta$ remains an interesting theoretical direction. If one can show that a single-layer Mamba is not a universal approximator even when $\\Delta$ is not fixed, then that would further stress the benefit of a multi-head design in selective SSMs. Conversely, if a single-layer Mamba with dynamic $\\Delta$ is a universal approximator, then an important follow-up question is how a trainable channel-specific $\\Delta$ can compensate for the channel-independent $\\mathbf{B}$ and $\\mathbf{C}$ matrices in a single-head Mamba.\n\n  In addition to the two cases pointed out by the reviewer, we plan for the following two changes as well:\n\n  3. Line 124-125 will be changed into\n\n  > Rather, Theorem 1 highlights that an S6 unit is not a simple extension of S4D that inherits all of its structural principles, and the sharing of the $\\mathbf{B}$ and $\\mathbf{C}$ matrices across all channels can reduce its capacity, making it potentially less effective than S4D for certain tasks.\n\n  4. Before Line 128, we will add the following sentence:\n\n  > While Theorem 1 assumes that $\\Delta$ is fixed, we use a synthetic experiment to demonstrate that even when $\\Delta$ is trainable, the channel-independent design of $\\mathbf{B}$ and $\\mathbf{C}$ in S6 limits its ability to solve certain problems than S4D.\n\n  We also appreciate the reviewer for pointing out [52] as a related work to our expressiveness discussion. We will change Line 672-673 into\n\n  > There are limited works on the universal approximation properties of Mamba models. One notable exception is [52], where the uniform closure of Mamba (i.e., functions that can be uniformly approximated by Mamba) is analyzed in a continuous setting (i.e., the inputs are continuous time series instead of discrete sequences) through linear controlled differential equations. The analysis there highlights the important role of an input-dependent sampling interval $\\Delta$ in the expressiveness of Mamba, suggesting that Theorem 1 may not hold for a trainable $\\Delta$ (see Appendix H).\n\n* ### Comparison with Selective SSMs (W3)\n\n  We thank the reviewer for pointing out this rendering issue in Table 2. We agree that the bold and underlined numbers do not take other SSMs into account. We will add the following sentence to the caption:\n\n  > While we compare B2S6 only to a few SSM prototypes in this paper, there are many variants of SSMs that can achieve slightly better performance (see Table 3).\n\n  In addition, we have also evaluated Mamba2 on the LRA benchmark during the rebuttal period. We will update the manuscript with the following sub-table:\n\n  | | ListOps | Text | Retrieval | Image | Pathfinder | Path-X | Avg. |\n  | --- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n  | Mamba |38.02| 82.98| 72.14| 69.82| 69.26| 67.32| 66.59|\n  | Mamba2 |41.45| 86.09| 79.23| 71.96 | 75.45| 69.07 | 70.54 |\n  | B2S6 |63.85| 88.32| 91.44| 88.81| 95.93| 97.90| 87.71|\n\n  We hope this comparison makes our evaluation more comprehensive.\n\n* ### Variants of Theorem 2 (W2 + Q1)\n\n  This is another great point. You may have noticed that the distinction between a time-variant system in Mamba and a time-invariant system in S4D in the context of relative gradients comes from two sources:\n\n  1. The system in S4D is linear while the system in S6 is quadratic (ignoring the trainable $\\Delta$ factor). The former linearity makes the gradient norm independent of the input $\\mathbf{u}_k$ while the latter quadratic dependency makes the gradient norm dependent on $\\\\|\\mathbf{u}\\_k\\\\|_2$.\n\n  2. In S6, the input-dependent sampling interval $\\Delta(c\\mathbf{u}_k)$ increases exponentially when $c$ grows to one end of $\\pm \\infty$ and decays exponentially when $c$ grows to the other end.\n\n  The reason why Theorem 2 involves an exponential gap is the second source above. That said, if we fix the norm of the entire input sequence, increasing the ratio between $\\\\|\\mathbf{u}\\_{k_0}\\\\|\\_2$ and $\\\\|\\mathbf{u}\\_{k'}\\\\|\\_2$ by decreasing the magnitude of $\\mathbf{u}\\_{k'}$, then the sampling interval $\\Delta(\\mathbf{u}_k)$ will be in a bounded range for all $1 \\leq k \\leq L$. So asymptotically, we will not see an exponential gap under this setting. The first source above is still relevant under this setting, meaning that as we increase the ratio $c = \\\\|\\mathbf{u}\\_{k_0}\\\\|\\_2 / \\\\|\\mathbf{u}\\_{k'}\\\\|\\_2$, the ratio between the gradient norms $S\\_{\\text{S6},k_0} / S\\_{\\text{S6},k'}$ will be asymptotically proportional to $c$.\n\n  While normalization weakens the theoretical merits of Theorem 2, we highlight an additional limitation in the original S6 block proposed in [25]: in the computation of $\\Delta_k^{(i)} = \\text{softplus}(\\mathbf{w}^\\top \\mathbf{u}_k + b^{(i)})$, the vector $\\mathbf{w}$ defines two half-spaces in $\\mathbb{R}^d$. Inputs $\\mathbf{u}_k$ that lie in the positive half-space produce larger values of $\\Delta$, leading to slower dynamics and better memorization, whereas those in the negative half-space yield smaller $\\Delta$, resulting in faster forgetting. This setup restricts the model’s ability to assign distinct memorization biases to different nonlinear or even linear regions of the input space. As discussed in the manuscript, modern implementations of Mamba often replace the vector $\\mathbf{w}$ with a low-rank projection, which allows for a finer partitioning of $\\mathbb{R}^d$. From this perspective, the multi-head design in Mamba can be seen as another architectural strategy to allocate different memorization biases to different subspaces. We will incorporate a fair and clear discussion of these points into the manuscript.\n\nThank you again for your careful review and constructive comments. We are committed to improving this manuscript, and please let us know if you have any other questions!"}]

Author rebuttal:

Meta-review: The paper at hand proposes a few modifications to the Mamba architecture targeted to improving expressivity and optimization properties of the block. Going beyond the original Mamba design, keeping efficiency while improving long-range capabilities is a highly relevant, timely topic. I particularly enjoyed the careful notation and precision of the paper, as well as the thorough ablations.


All reviewers agree that the paper deserves publication, and the authors provided additional experiments supporting their claims during the rebuttal period. For this reason, I recommend accepting this work.
Decision: Accept (poster)