Demystifying the Token Dynamics of Deep Selective State Space Models
Authors: Thieu Vo
Duy-Tung Pham
Xin T. Tong
Tan Minh Nguyen
Venue: ICLR 2025 Spotlight (2025)
Keywords: Selective state-space model
continuous-time limit
dynamical system
asymptotic behavior
token reordering
Primary area: foundation or frontier models, including LLMs
TL;DR: We describe the dynamical properties of tokens in a deep selective state-space model, and based on that, we suggest improvements to the model by excluding negatively impactful scenarios and reordering tokens using importance scores.

Abstract: Selective state space models (SSM), such as Mamba, have gained prominence for their effectiveness in modeling sequential data. Despite their outstanding empirical performance, a comprehensive theoretical understanding of deep selective SSM remains elusive, hindering their further development and adoption for applications that need high fidelity. In this paper, we investigate the dynamical properties of tokens in a pre-trained Mamba model. In particular, we derive the dynamical system governing the continuous-time limit of the Mamba model and characterize the asymptotic behavior of its solutions. In the one-dimensional case, we prove that only one of the following two scenarios happens: either all tokens converge to zero, or all tokens diverge to infinity.  We provide criteria based on model parameters to determine when each scenario occurs. For the convergent scenario, we empirically verify that this scenario negatively impacts the model's performance.  For the divergent scenario, we prove that different tokens will diverge to infinity at different rates, thereby contributing unequally to the updates during model training.  Based on these investigations, we propose two refinements for the model: excluding the convergent scenario and reordering tokens based on their importance scores, both aimed at improving practical performance.  Our experimental results validate these refinements, offering insights into enhancing Mamba's effectiveness in real-world applications.

Published as a conference paper at ICLR 2025
DEMYSTIFYING THE TOKEN DYNAMICS OF DEEP SELECTIVE STATE SPACE MODELS
Thieu N. Vo
Duy-Tung Pham
Department of Mathematics National University of Singapore
FPT Software AI Center cungpa108110 com
thieuvo@nus.co edw.euuuuuuuuuuuuuueeu..eu. sg
Xin T. Tong"
Tan M. Nguyen"
Department of Mathematics National University of Singapore mattxin@nus edu. sg
Department of Mathematics
National University of Singapore Canmn@nus edu. sg
ABSTRACT
Selective state space models (SSM), such as Mamba, have gained prominence for their effectiveness in modeling sequential data. Despite their outstanding em- pirical performance, a comprehensive theoretical understanding deep selective SSM remains elusive, hindering their further development and adoption for ap- plications that need high fidelity. In this paper, we investigate the dynamical properties of tokens in pre-trained Mamba model. In particular, we derive the dynamical system governing the continuous-time limit of the Mamba model and characterize the asymptotic behavior of its solutions. In the one-dimensional case, we prove that only one of the following two scenarios happens: either all tokens converge to zero, or all tokens diverge to infinity. We provide criteria based on model parameters to determine when each scenario occurs. For the convergent scenario, we empirically verify that this scenario negatively impacts the model's s performance. For the divergent scenario, we prove that different tokens will di- verge to infinity at different rates, thereby contributing unequally to the updates during model training. Based on these investigations, we propose two refinements for the model: excluding the convergent scenario and reordering tokens based their importance scores, both aimed at improving practical performance. Our experimental results validate these refinements, offering insights into enhancing Mamba's 's effectiveness in real-worle applications. The code is publicity available at https::/ / teuii.hu.b... com/F Ffom/Fsottt.. -oof-AiCCCCC...... CMMcC///aaaaaaaaa token ddnamiic.
INTRODUCTION
State space models (SSMs) have undergone sigmificant advancements to mitigate the computational inefficiency associated with the sequence modeling (Gu al., 2022b; 2021; 2022a; Gupta et al., 2022; Li et al., 2023; Kosma et al., 2023; Orvieto et al., 2023; Smith et al., 2023), and they have been successfully applied in various domains involving continuous signal data such as audio and vision (Goel et al.. 2022; Nguyen et al., 2022a; Saon et al., 2023). SSMs can be viewed as a combi- nation of recurrent neural networks (RNNs) and convolutional neural networks (CNNs), drawing on concepts from traditional state space frameworks (Kalman, 1960). Unlike transformers, which expe- rience quadratic scaling with respect input sequence length and exhibit significant computational demands (Vaswani et al., 2017), SSMs are designed for efficient computation through recurrence or convolution, achieving linear or near-linear scaling as the sequence length increases. Furthermore, SSMs incorporate robust mechanisms for capturing long- range dependencies (Gu et al., 2020) across various data types, and they have excelled in benchmark tasks such as the Long Range Arena (Tay al., 2021).
"Co-last authors. Please correspond thieuvo nus.edu.sg and tanmn @nus.edu.sg

---

Published as a conference paper at ICLR 2025
Selective SSMs such as Mamba is a particular SSM which involves the selective state-space layer (S6) as their core component (Gu & Dao, 2024). In an S6 layer. parameters are functions of the input, endowing the SSM with the content awareness ability. Mamba has showcased exceptional performance across variety of applications, such as language modeling (Pióro et al., 2024; Lenz et al., 2025), image processing (Liu & Li, 2024; Zhu et al., 2024), video analysis (Yang al., 2024; Li et al., 2024a), medical imaging (Ma & Wang, 2024; Wang et al., 2024b), tabular data analysis (Ahamed & Cheng, 2024), reinforcement learning (Ota, 2024), point-cloud analysis (Liang et al., 2025), graph processing (Wang et al., 2024a), and N-dimensional sequence modeling (Li et al., 2024b). It has been argued that Mamba's ss significant success across various domains stems from its ability to compute complicated representations for different data types (Rezaei Jafari et al., 2025; Patro & Agneeswaran, 2024; Xu et al., 2024; Zhu et al., 2024).
Despite their remarkable empirical success, thorough theoretical understanding of deep selective SSM is still lacking, which poses challenges for their advancement and application in high-fidelity tasks. As these models are scaled at a remarkable rate, understanding their internal mechanisms has become a major problem. In this paper, we take an initial step toward addressing this gap by analyzing the dynamical properties of tokens in a pretrained Mamba.
1.1 BACKGROUND: MAMBA
The building block of selective state-space model is Mamba block, which is built top of the S6 layers, along with linear layers, convolutions, and other token-wise operators. The S6 layer, which serves as the core component of the Mamba block, plays a pivotal role in the success of these models.
Formally, an S6 layer defined as a function that maps sequence of tokens x (Zi...., 21) DDL to another sequence of tokens y (y1,..., E DxL (with the same number of channels). In this context, each vector Ji (or y1) represents a token, while the whole sequence = (#1...., ,L) (or y (V1.... ................................ ....)) is called a prompt. For each d = I...., D, the d-th chan- nel output y4 = (Mal.-... yar.) E B. is determined recurrently from the corresponding d-th channel input Z2= = (I.I) ................................ .ar.) E IR L via sequence of hidden states hai, hai...., har. E IR- as follows:
haa Aa ha,i- Bay Tdl,
hao O,
yau = C2 -ha,
for L, I...., L. Here, the matrices Aat, Bai and Oj are input-dependent and time-varying, deter- mined by
Aa Au(rr)As Ba Aa(1)) m. C (Sc 21)T
where Ao DRD - R is the step size function at the d-th channel and is defined by
Ad(u) (oftplus(Ss,au) In C)
R.
(3)
The hidden matrices Aa E NNN and the step size vectors SA.d E RIID for d = 1, D, as well as the input and output matrices SB. Sc E NxD, are learnable from input data. The matrices Ag are diagonal matrices with negative eigenvalues. In recent developments, the matrices Aa chosen to be of the scalar form A = a n for some positive number aa (see (Dao & Gu, 2024)), which does not compromise model performance. We will also use this scalar form for Aa in our paper. In our context, the matrix ST.S E DxD which we will refer to as the input-output matrix, plays an essential role in characterizing the dynamical properties of tokens.
1.2 CONTRIBUTION
This paper aims describe the dynamical properties of tokens in pre-trained Mamba model. To achieve this, we consider the dynamical system governing the continuous-time limit of Mamba. The dynamical properties of tokens are described through the asymptotic behavior of the solutions of the dynamical system. Additionally, we empirically investigate the relationship between the dynamical properties of the tokens and model performance. In summary, our main contributions are three-fold:
1. In the mieedimensional case, we prove that only one of the following two scenarios occurs: either all tokens converge to zero, or all tokens diverge to infinity. We provide criteria based on model parameters to determine when each scenario occurs. Our experiments suggest that these observations generally hold in high-dimensional cases.
2

---

Published as a conference paper at ICLR 2025
2. For the convergent scenario, empirically verify that this situation negatively impacts the model's 's performance. In contrast, for the divergent scenario, we prove that different tokens will diverge to infinity at different rates, thereby contributing unequally to the updates during model training.
3. Based on these investigations, we propose two refinements for the model: (i) we exclude the convergent scenario before training negatively impacts performance, and (ii) reorder the tokens according to their ascending importance scores, as different tokens are prioritized unequally during training.
We empirically demonstrate the benefits of our token rendering method improving the model's accuracy and convergence speed compared to the baseline Mamba the large-scale ImageNet classification task (Deng et al., 2009)
Organization of the paper After surveying related works in Section 2, we introduce the dynam- ical system that governs the continuous-ting limit of Mamba. In Section 4, we characterize the dynamical properties of tokens based on model parameters and determine the divergence rate in the divergence scenario. Based on the findings in Section 4, propose two refinements for the model in Section 5: excluding unfavorable scenarios and reordering tokens before training. We conclude the paper with discussion of conclusions and limitations in Section 6.
Notations. We use small bold letters (e.g., x, y. u) denote sequences of tokens. Tokens can be either vectors or scalars, depending on the context, and are denoted by small normal letters (e.g., Te. Mi uk), where the indices emphasize their position in the sequence. We capital letters (e.g.. X, Y, P, A, B, C) to represent matrices.
2 RELATED WORK
Several studies have analyzed the expressivity and generalization of Mamba from a theoretical per- spective. For instance, the authors in (Ali et al., 2024) demonstrated that the S6 layer can be in- terpreted as a variant of softmax-free attention with linear complexity, subsequently proving that Mamba is more expressive than transformers. Conversely, (Muca Cirone et al., 2025) employed tools from Rough Path Theory to show that diagonal selective state space models (SSMs), such as Mamba, possess less expressive power than their non-diagonal counterparts. Furthermore, (Merrill et al., 2024) utilized circuit complexity theory to establish that both SSM variants and transform- share the decided by polynomial-sized Boolean circuits. In contrast, the authors in (Jelassi et al., 2024) pro- vided both theoretical and empirical evidence that transformers surpass state space models in their copying capabilities. All of these works aimed to provide theoretical framework to understand the expressivity and generalization Mamba.
expressive power, they belong to the complexity class TC", which can be
In this paper, we examine the dynamical properties of tokens in a pretrained Mamba model. Fol- lowing the methodology outlined in (Geshkovski et al., 2023; 2024), we define an idealized model of Mamba by interpreting the discrete layer indices as continuous time variables. This idealized perspective was used in ResNets (Chen al., 2018; 8 Haber & Ruthotto, 2017: Abdullaev Nguyen, 2025) and neural ODEs (Lin & Jegelka, 2018; Zhang et al., 2020; Xia et al., 2021; Li et al., 2022; Tabuada & Gharesifard, 2022; Nguyen et al., 2022b; Ruiz-Baler & Zuazua, 2023; Cheng et al., 2025). Our model focuses the S6 layer, which is central to Mamba and contributes to its success across various domains.
3 SELECTIVE STATE SPACE DYNAMICS
In this section, we introduce the dynamical system governing the continuous-time counterpart of Mamba. This dynamical system will provide an effective tool to characterize the dynamical proper- ties of tokens in Mamba as well as the impact of these properties on model performance in the next section.
3.1 BACKGROUND: CONTINUOUS TIME LIMIT OF DEEP NEURAL NETWORK
To ensure clear presentation of our results, draw upon the literature concerning the dynamical systems that govern the continuous time limit of deep neural networks (DNNs). Generally speaking, data in DNN is processed sequentially, layer by layer, resulting in discrete-time dynamical system
3

---

Published as a conference paper at ICLR 2025
(LeCun et al., 2015). A notable example is residual neural networks (ResNets) and their continuous- time counterparts known as neural ODEs (Chen et al., 2018; Haber & Ruthotto, 2017; Nguyen et al., 2020; Wang et al., 2022). Each layer of ResNet is residual block, transforms input vector x & RDD to an output vector 2 E TDD via a two-layer feed-forward neural network æ - y(r, 0), parametrized by 0, and skip-connection as:
== + (z, 0).
The continuous-time counterparts of ResNet is conceptualized as a flow map that inputs vector z<() E RD and outputs another vector z(T) E RR", which processed via the corresponding dynamical system
=
E (O,T).
There exists body of work investigating the interpolation, approximation, and controllability prop- erties these DNN architectures (Lin Jegelka, 2018: Zhang et al., 2020; et al., 2022; Tabuada & Gharesifard, 2022; Ruiz-Balet & Zuazua, 2023; Cheng et al., 2025; Nguyen et al., 2024).
3.2 CONTINUOUS-TIMA LIMIT OF MAMBA
Unlike ResNets and neural ODEs, Mamba represents a function on sequence of D-dimensional tokens rather than solely on an individual input, as discussed in Section 1.1. In order derive the dynamical system governing the continuous- time counterpart of Mamba, we eliminate the hidden states hat system (1) obtain the following form (see (Ali al., 2024)):
Va
Pay
(4)
or equivalently, da- Patd, where
Pan
0
0
Paz
Pazz
0
Pa
(5)
PALa
Pata
Patt
and
xo (S 31 Aa(a))
i = =j
Paj
(6)
T (5) 2j Ad(T;) - exp
- -Od Σ Aa(TR)
if j.
k=j+1
In the above formulation, the upper triangular matrices Pa (d 1,..., D) represent the hidden attention scores between the d-th channel input and output. Following the setting of (Ali et al., 2024), call P .P.....,P the hidden attention tensor the S6 layer.
Following the common approaches used in studying the continuous-ting counterparts of ResNets (Chen al., 2018; Haber & Ruthotto, 2017) and Transformers (Geshkovski al., 2023; 2024), we consider the layer index of deep selective state space model as a time variable and interpret the selective state space model as the discrete-time time version of parametrized dynamical system of the form
/i =a
Σ (f) E +00), 1=1
(0) Fat R.
for = and 11,.,.... 1,...,D. D. where
( (S) (1) Ad(ai()),
ifl =
Paz(()
I TT(1) SB) 2z(1) Aa()((t) exp -ad Σ sa(ax(()) k=j+1
(8)
if
and (²( = (oftplus(SA.ai) is important to mention that the time in this dynamical system represents the depth direction of Mamba model.
4

---

Published as a conference paper at ICLR 2025
Similar to (Geshkovski et al., 2023; 2024). have focused exclusively S6 layer, which is the key component of Mamba, and the skip-connection in the above dynamical system. We have ig- nored other token-wise operators such as layer normalization, convolution, and linear functions. In addition, we assume the parameters Ad, SB,Sc are time-independent. These assumptions are primarily motivated by mathematical convenience. Nevertheless, such weight-sharing mechanism is used in practice to reduce the number of trainable parameters, as seen, for example, in ALBERT (Lan al., 2020), a transformer-bast large language model. The dynamical properties of tokens in deep selective state space models with these additional layers will be discussed in future works.
4 DYNAMICAL PROPERTIES OF TOKENS IN MAMBA
In this section, study the asymptotic behavior of the solution x(t) of the dynamical system (7), as well as its corresponding attention score P(t) = (a))(t))a defined in equation (8). This information encodes the dynamical properties of tokens and hidden attention scores deep selective state space models.
4.1 CONVERGENCE AND DIVERGENCE SCENARIOS OF MAMBA'S DYNAMICS
In the standard form of S6, each d-th channel output is determined from the corresponding d'th channel input using vari-variant input-dependent parameters. Therefore, we will consider the single channel dimension case, i.e. D 1, in our setting to avoid overly complicated technicalities. this
case, the input sequence x (I1,- (#1,..., IL) R_ is list of L real numbers. We drop the index d from Eq. (7) and rewrite as
dt mt Σ (t)))) te 10,+0000) j=1
(9)
(0) E I,
for = 1, I...,, L, where
(wu) softplus(Sauu In (1 + eSau), MER
(10)
and
(ST 21) ((x((t)),
if
Ptt(t)
(ST.S)) 2i (t) (tz;(t) exp -a Σ A(Tx(()) if j. k=jt1
(11)
In this case, SB, Sc E Nx1, thus STSB E IR, and the scalar values SA and E R, with > 0, are learnable from imput data.
To avoid triviality, we will always assume that the initial datum (0) = are nonzero for all The main goal of this section to investigate the asymptotic behavior of the solution x(t) (ai(t)..... . (t)) of the dynamic (9) and the corresponding hidden attention matrix
(n)()) 0
Pai(1)
P2z(1)
0
P(t)=
PLi(f) Pta(t)
(L)
when t increases. The following three scenarios will be considered separately:
1. Convergence scenario: ST.SB 0;
2. Slow-divergence scenario: ST and SAZIO 0 for all I:
3. Fast-divergence scenario: ST.B > and SATI > 0 for some
We will see that, in the first scenario, all tokens and hidden attention scores quickly converge to zero as - 20, at rates (LIVVF) and O(1/t), respectively (Theorem 4.1). In the second scenario, all tokens diverge to infinity slowly at logarithmic rate, while the hidden attention scores still approach zero as co. We refer to this as the siow-divergence scenario (Theorem 4.3). In the final scenario, one token diverges infinity very quickly in finite time, causing a blow-up in the corresponding hidden attention scores. We refer to this as the fast-divergence scenario (Theorem 4.4). Table summarizes the dynamical properties of tokens in this work.
5

---

Published as a conference paper at ICLR 2025
Table 1: Summary of the dynamical properties of this work.
Parameters
Scenario
Impact on model
Reference
Convergence
Negative
Theorem 4.1
SB > 0 Vl, SA310 < 0
Slow-divergence
Theorem 4.3
Positive
3l, Sa > 0
Fast-divergence
Theorem 4.4
Remark (Convergence vs. divergence scenarios and their impact). We collectively refer to the dlvedivergennce and fast-divergence scenarios as the divergence scenarios. By utilizing model pa- rameters to test the signature of St. we can distinguish between the convergence scenario and the divergence scenarios. However, to distinguish between the slow-divergence and fast-divergence scenarios, additional information from the input data, specifically F1,0- is required. In addition, our experiments suggest that the convergence scenario negatively impacts model performance in practice, while the divergence scenario positively impacts model performance. Therefore, the con- vergence scenario should be excluded before training (see Subsection 5.1).
Remark 2 (Unequal contribution of tokens during training). In the slow-divergence scenario, when tokens diverge to infinity at infinity, we further prove that tokens with different initial positions will have different rates of divergence. As consequence, the contribution of tokens the updates during model training is unequal. Intuitively speaking, the tokens should be rearranged in such a way that those carrying more important information diverge faster than those carrying less important information. In Subsection 5.2, we suggest a method for reordering the tokens before training by projecting them into a learnable line. Our experimental results show that this ordering method leads to better model performance.
Remark 3 (Higher dimension). When the dimension of the input tokens is D 2 1, the input-output projection matrix St is square matrix of size D D, which may have complex eigenvalues. We conjecture that the dynamic behavior of the tokens in system (7) can be determined by analyzing the matrix
1 SS + SST
/ 2
which represents the symmetric part of the input-output projection matrix Sg.S.. It is worth noting that in case D 1, μ SC expected. The matrix μ has only real eigenvalues. We particularly conjecture that all tokens will converge to zero if μ has only negative eigenvalues, whereas the tokens will diverge to infinity if fu has least positive eigenvalue (see Appendix B for further detail). We leave the study of the higher- dighes-dimennional case for future work.
Remark 4 (Compare with tokens' dynamic in Transformer). As outlined in (Geshkovski et al., 2024) (page 6. paragraph 2), tokens in purely self-attention setting typically exhibit exponential divergence to infinity. However, our theoretical analysis demonstrates that the token dynamics the purely Mamba setting are much more complicated. Specifically, tokens in the Mamba framework may either converge to zero or diverge to infinity. In cases of divergence, the behavior can vary: tokens may diverge rapidly infinity within finite time diverge more gradually to infinity over an infinite time horizon, with the divergence following a logarithmic rate.
In the subsequent subsections, we provide the statements of the main theorems. The proof's of these theorems can be found in Appendix A.
4.2 CONVERGENCE SCENARIO: == ST <0
In the following theorem, we prove that all tokens converge to zero rate of O(1/Vt) t approaches infinity, regardless of the input data. Consequently, the hidden attention scores will also tend toward zero.
Theorem 4.1 (Convergence scenario). Assume that μ STSB < 0. Let = be the unique solution of the dynamic (9).
1. For each I = l,...,, if xno > 0 (respectively, 2ia < 0), then Iz(t) is positive and mono- tonically decreasing (respectively, negative and monotonically increasing) on [0, +00). In addition, 2t( = 0 7.)
2. (t)(t) (()) for all 1.j 1,...,I. L.

---

Published as a conference paper at ICLR 2025
The proof of this theorem can be found in Appendix A.1. To intuitively why all tokens converge in this scenario, we can rewrite the dynamic (9) as
dt =1) = (()())))))))))))))((((((())))))) + (()))())))
(12)
where
0,
if
((
pay (t)) exp Σ Sloult)) k=j+1
if
Since μ < O, gt) is nonpositive function. Let us consider the case when 210 > 0. Then, (t(t) must remain positive for every t > O; otherwise, x1( would reach zero at some point, and consequently zt) would be identically zero for all according to the uniqueness of the given initial value problem, which is impossible. As a result, the right-hand side of equation (12) is always negative. This implies that s7(() decreasing. Figure illustrates the graph of tokens and the heat map of the corresponding hidden attention matrices in this scenario.
Remark (Negative impact on model's performance). In this scenario, both the tokens and hidden attention scores collapse to zero, leaving no room for diversity or randomness, which may be harmful to the model's performance. Our experiment in Section 5.1 verifies these effects in high dimensions.
X1
X4
2
X6
Xg
-1.
Figure 1: The graphs of tokens (left) and the heat map of the hidden attention matrices (right) for the
convergence scenario with tokens and parameters 14 11558, SA = -0.17, a -1.0 as well as initial data x0) (-1.79, 00.34, 0.46, - 1122, 0.83, -0.83, 1.81, -1.16, 0.13, 0.199). In this case, all tokens and hidden attention scores tend to zero as approaches infinity.
4.3 SLOW-DIVERGENCE SCENARIO: f STSa > 0 AND SA3io < 0 FOR ALL
In this case, we first observe that all of the tokens will tend to infinity as we will see in the following lemma.
Lemma 4.2 (Slow divergence scenario). Assume that μa = S > and SANI < O for all I. Let sx() (x) (t)) be the unique solution of the dynamic (15). For each I, if ano > (respectively, Σin < 0), then si(t) is positive and monotonically increasing (respectively, negative and monotonically decreasing) on [0, +oo) with
lim (1) = +00 (respectively, lim si(() = -oo) +
t0+
Lemma 4.2 shows that all tokens in the considered case approach infinity as approaches infinity. However, it does not help to estimate the rate of divergence or the asymptotic behavior of the hidden attention score (t). In Theorem 4.3 below, we provide the rate of divergence for the tokens and the asymptotic behavior of the hidden attention scores, but under certain additional assumption.
Theorem 4.3 (Slow divergence scenario and divergence rate). Assume that μ ST > 0 and
SATIO S ... S SAAT1O < -2.12
(13)
(see Theorem A.7 for a detailed upper bound). Let x(t))))= = (t) ...))....,........... (t)) be the unique solution of the dynamic (9).
7

---

Published as a conference paper at ICLR 2025
1. For each I. if 3/t 0 (respectively, vi < 0), then zi(t) is positive increasing (respec- tively, negative decreasing function on [0, +.0). In addition, 21(t) = ((Int))).
2. lim (t) = 0 for all i 2 j. Remark (Unequally tokens' contribution). Theorem 4.3 shows that, in the mne-dimensional case within this divergence scenario, tokens with higher absolute values will diverge infinity faster than those with smaller absolute values (while all hidden attention scores still converge to zero). Intuitively speaking, this means that different tokens will contribute differently to the updates during training. Based on this observation, we conjecture that reordering tokens based on their importance scores before training might be beneficial for model performance. Our experiment Subsection 5.2 empirically confirms this conjecture.
The proof of this theorem can be found in Appendix A.2. We can intuitively see why the tokens tend to infinity slowly by looking at the values of the right-hand side of equation (12). Indeed, let us assume that SA (t) would reach zero at some point, and consequently (t) would be identically zero for all t, according to the uniqueness of the given initial value problem. which impossible. As result, the right-hand side of equation (12) is always not too large. This implies that 2t(+) is decreasing. However, since SA < O, the term par (t))(() A (z)) converges to zero very quickly, which keeps the right-hand side always not too large. Therefore, (t) goes to infinity slowly. Figure 2 illustrates the graph of tokens and the heat map of the corresponding hidden attention matrices in this scenario.
and thus Ti0 > O. Then (t) must always be positive t increases; otherwise,
X1
200
Xs
150
X6
X7
100
Xg
100 X10
Figure 2: The graphs of tokens (left) and the heat map the hidden attention matrices (right) for the convergence scenario with L = 10 tokens and the model parameters fr -.779, S3 = 00771, a -1880, as well the initial data x(0) (1.35, 2.84, 3.81, 4.57, 5.99, 6.94, 7.71, 8.96, 9.59, 10. 75). In this case, tokens tend to infinity at the log-rate, and the one which larger initial value diverges faster, while the hidden attention scores tend to zero as f approaches infinity. In addition, the hidden attention scores from the second columns tend to zero must faster than those in the first column.
4.4 DAVE-DIVERG SCENARIO: μ = ST.B > 0 AND SAZIO >> O FOR SOME In the fast divergence scenario, we have the following theorem. Theorem 4.4 (Fast divergence scenario). Assume that I = ST > O. Let x(t) = (t), ..... (t)) be the unique solution of the dynamic (15). If there exists that SAT > O, then st( goes to infinity at finite time.
1,..., L such
In addition, assume that Fi, (t) tends to infinity first as t T for some T > 0 while z (1) is defined on [0, T), for all + lg. Then for every 2212j 2 1, lim, (t)(t) = too if and only if 1= lo or j lo-
Remark 7. The proof of this theorem can be found in Appendix A.3. Intuitively speaking, in this scenario, if 1 < < L is the index of the first token such that SAZIo > 0. then all tokens with indices j 2 will diverge to infinity in finite time. Figure illustrates the dynamical properties of tokens and the hidden attention scores in this scenario.
5 EXPERIMENTS
In this section, we empirically validate two key observations derived from the theoretical analysis presented in Section 4. We aim to (i) demonstrate the detrimental impact of negative eigenvalues in
8

---

Published as a conference paper at ICLR 2025
X1
1.10
X5
3.10
X6
2
Xy
X6
10
0.3 0.6 0.5 X10
0.1
Figure 3: The graphs of tokens (left) and the heat map of the hidden attention matrices (right) for the convergence scenario with L. 10 tokens and the model parameters μ 0.76, SA 0.59, a 1.66, as well as the initial data x0) = (0.83, 0.91, 0.64, 0.78, 0.66, 0.99, 0.68, 0.72, 0.61, 0.90). In this case, tokens and the hidden attention scores tend infinity very quickly at finite time.
the input-output matrix ST B on the performance of auto-regressive language modeling tasks and (ii) propose a reordering strategy and illustrate its effectiveness in vision tasks. Details on datasets, models, and training procedures are provided in Appendix C.
5.1 EFFECTS OF NEGATIVE EIGENVALUES OF INPUT-OUTPUT MATRIX
We use Mamba (Gu & Dao, 2024) as the baseline and conduct evaluations on the WIKITEXTION 03 language modeling task (Merity et al., 2017). We consider three distinct scenarios, including cases where the input-output matrix contains (i) only positive eigenvalues, (ii) only negative eigenvalues, as discussed in Section 4 and Appendix B, and (iii) both positive and negative eigenvalues as an intermediate (mixed) case. The scenario where Sg encompasses complex eigenvalues is deferred future work. We report the perplexity (PPL) for each scenario. Lower PPL values indicate better model performance.
30.0
Nogative elgenvalues
27.5
Mixed
Table Test perplexity on WIKI- TEXT103. A model where Sy has a smaller proportion of negative eigen- values achieves lower perplexity.
Perplexity
eigenvalues
25.0
22.5
20.0
Scenario
Perplexity (1)
17.5
Negative 17.28
0
10
Epoch
Mixed 16.82
Positive 16.66
Figure 4: Test perplexity on WIKITEXTI0S 103 during train- ing procedure. The positive case consistently demonstrates superior performance compared to the other two scenarios.
Table and Figure 4 illustrate the impact of different eigenvalue configurations on the model's performance. Specifically, our findings demonstrate correlation between the prevalence of positive eigenvalues in the input- matrix and model performance. Conversely, the presence of negative eigenvalues linked to decrease in performance, suggesting that removing negative eigenvalues could further optimize the model's 's efficiency.
5.2 RE-ORDERING INPUT TOKENS
We now emphasize the significance token ordering in vision tasks, where tokens naturally lack an inherent structure due to each token representing patch from a 2D or 3D image. As indicated in Theorem 4.3, the contribution of each token to model updates is primarily influenced by the term SATM, which is a vector of dimension D in scenario with D channels. To determine the optimal order of tokens Tin we define an importance score as follows:
(K, SANI)
(14)

---

Published as a conference paper at ICLR 2025
where K RP is learnable parameter. Prior passing the tokens through the S6 layer, they reordered based on their importance scores in ascending order. We use SoftSort (Prillo & Eisen- schlos, 2020) enable gradient propagation through the sorting operation. Specifically, SoftSort is employed to generate a permutation matrix, which then applied to sort the er array and rearrange the tokens Zin accordingly. Details of the procedure and analysis of the computational cost provided in Appendix D.2.
To assess the effectiveness of our proposed reordering technique, we benchmark on the image classification task using the ImageNet-IK dataset (Deng et al., 2009). We employ Mamba Vision (Hatamizadeh Kautz, 2024) the baseline and compare the top-1 and top-5 accuracy metrics.
Table 3: Top-1 and Top-5 accuracy (1) on the large-scale ImageNet- image classification task. Our token reordering method leads to slight improvement in Mamba Vision's ' performance.
Model
Top-1 1 (%)
Top-5 (%)
MambaVision- (baseline)
81.90
95.86
Mamba VisionVissonTTT Token reordering
82.02
95.87
Top 1 Accuracy (%)
Top 5 Accuracy
NambaVision (baselin))
N Nambavision (baselin) Nambavision tokeen
Nambaiision reordaring
160
240
240
Epoch
Epoch
Figure 5: Top-1 (left) and Top-5 (right) accuracy (1) on ImageNet-IK during the training process. Our token reordering method boosts the accuracy of the Mamba Vision baseline and achieves faster convergence compared to the baseline.
The results presented in Figure and Table 3 demonstrate that our reordering method speeds up the convergence of the training loss and improves the performance of the Mamba Vision baseline. These findings confirm the effectiveness of our technique in vision tasks, highlighting its potential for broader tasks. Experiments language modeling and text classification are in Appendix D.
CONCLUDING REMARKS
We characterize the dynamical properties of tokens in a pre-trained Mamba model. In the one- dimensional case, we show that either all tokens converge to zero or all tokens diverge to infinity. In addition, for the convergent scenario, we empirically verify that this scenario negatively impacts the model's performance; therefore, it should be excluded from model training. For the divergent scenario, we prove that tokens at different locations will diverge to infinity at different rates, thereby contributing unequally the updates during model training. Based on these investigations, propose two refinements for the model: excluding the convergent scenario and reordering tokens based on their importance scores, both aimed at improving practical performance. Our experimental results suggest that these refinements do help improve model performance.
Regarding limitations, to tokens with one feature channel. In addition, our token reordering refinement, based on impor- tance scores, introduces additional computational overhead during training, although this increase is relatively small. Furthermore, we are uncertain whether the method of reordering tokens based on their orthogonal projection onto learnable affine line, as used in practical implementation, is optimal. We leave a systematic study of token reordering refinements as well as generalization of our theoretical analysis dimees-dimensional cases for future work.
theoretical analysis of the dynamical properties of tokens is restricted
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This research project is supported by the National Research Foundation Singapore under the Al Singapore Programme (AISG Award No: ALSG2-TC-20223012-SSIII supported by the Ministry of Education, Singapore, under the Academic Research Fund Tier This research / project is supported by the NUS
This research project is
(FY2023) (A-8002040-00-000, A-8002039-00-000 Presidential Young Professorship Award (A-0009807-01-000 and the NUS Artificial Intelligence Pustitute-Seed Funding (A 1--8003062-00-00), This research project is also supported by Singapore National Academy of Science under the SASEA Fellowship Programme (Award No: NRF-MP- 2025-0001).
Reproducibility Statement. and Appendix C. Source code is provided at https:// com/Fsoft-AIC/
Our implementation details are provided in Section
Lamba-tokke-dyyel dynamic..
Ethics Statement. Considering the scope of our research, we do not anticipate any negative societal or ethical consequences emerging from this work.
REFERENCES
Laziz Abdullaev and Tan Minh Nguyen. Transformer meets twicing: Hamessing unattended residual information. In The Thirleenth International Conference on Learning Representations, 2025. n forum? 1 6kG5aNleS.
URL wwt.hhtp////peeeee http:///pppeeiieeeeeeee wwt.iwwweeeeee...........m..eeee
Md Atik Ahamed and Qiang Cheng. Mambatab: A plug-and model for learning tabular data. In 2024 IEEE 7th International Conference on Multimedia Information Processing and Retrieval (MIPR), pp. 369-375, 2024. doi: 10.11000.10.0.. 11119/MIPR22022222200000000
Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. arXiv preprint arXiv:2403.01590 2024.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Jingpu Cheng, Qianxiao Li, Ting Lin, and Zuowei Shen. Interpolation, approximation, and con- trollability of deep neural networks. SIAM Journal Control and Optimization, 63(1):625-649) 2025.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning, pp. 10041- 10071. PMLR, 2024.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei. Imagenet: A large-scale hier- archical image database. 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp.
2009. doi: 10. 109/CVPR.2009.5200622.
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry Hungry Hippos: Towards language modeling with state space models. In International Confer- ence on Learning Representations, 2023.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. mathematical per- spective on transformers. arXiv preprint arXiv:2312.1079, 2023.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clus-
ters in self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state- space models. In International Conference on Machine Learning, pp. 7616-7633, PMLR, 2022.
Gene H. Golub and Charles Van Loan. Matrix Computations. The Johns Hopkins University Press, Baltimore, Maryland, 4th edition, 2013. ISBN 978-1-4414-009444.
11

---

Published as a conference paper at ICLR 2025
Albert Gu and Tri Dao. Mamba: Linear-tine sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024. URL https:://///... httttp::///pereviviie net forum?? id= tEYskwIVY2.
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-1487, 2020.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Com- bining recurrent, convolutional, and continuous-time models with linear state space layers, Ad- vances in neural information processing systems, 34:572-585. 58022,5552500 2021.
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems, 35:35971- 35983, 2022a.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022b. URL https:
//openreview. wwwpppeeewww..eeeeeeeeeeeeeeeeee wewtt//ooommddduuuLoooooo
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces effective structured state spaces. Advances in Neural Information Processing Systems, 35:22982-22999 2022.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse problems, 34( 1):014004, 2017.
Ali Hatamizadeh and Jan Kautz. Mambavision: hybrid mamba-transformer vision backbone. arXiv preprint urXiv:2407008083, 2024.
Ali Hatamizadeh, Hongxu Yin, Greg Heinrich, Jan Kautz, and Pavlo Molchanov. Global context vi- sion transformers. In International Conference on Machine Learning, pp. 12633-12644 PMLR, 2023.
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Trans- formers are better than state space models at copying. In International Conference on Machine Learning. pp. 21502-21521. PMLR, 2024.
RE Kalman, A new approach to linear filtering and prediction problems. Journal of Basic Engineer- ing. 82(1):35-45, 1960.
Chrysoula Kosma, Giannis Nikolentzos, and Michalis Vazirgiannis. "imme-paraeeterized convo- lutional neural networks for irregularly sampled time series. arXiv preprint arXiv:2308.032210 2023.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Machine Learning. PMLR, 2020.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. Nature, 52177553):436-44 2015.
Barak Lenz, Opher Lieber, Alan Arazi, Amir Bergman, Avshalom Manevich, Barak Peleg, Ben Aviram, Chen Almagor, Clara Fridman, Dan Padnos, Daniel Gissin, Daniel Jannai, Dor Muhl- gay, Dor Zimberg, Edden M. Gerber, Elad Dolev, Eran Krakovsky, Erez Safahi. Erez Schwartz, Gal Cohen, Gal Shachaf. Haim Rozenblum, Hofit Bata, Ido Blass, Inbal Magar, Itay Dalmedi- gos, J.onathan Osin, Julie Fadlon, Maria Rozman, Matan Danos, Michael Gokhman, Mor Zus- man, Naama Gidron, Nir Ratner, Noam Gat, Noam Rozen, Oded Fried, Ohad Leshno. Omer Antverg, Omri Abend, Or Dagan, Orit Cohavi, Raz Alon, Ro'i Belson, Roi Cohen, Rom Gilad, Roman Glozman, Shahar Lev, Shai Shalev-Shwartz, Shaked Haim Meirom, Tal Delbari, Tal Ness, Tomer Asida, Tom Ben Gal, Tom Braude, Uriya Pumerantz, Josh Cohen, Yonatan Belinkov, Yuval Globerson, Yuval Peleg Levy, and Yoav Shoham. Jamba: Hybrid wansformer-mandi language models. In The Thirteenth International Conference on Learning Representations, 2025. URL https:/ / /ttp//peeeeeeeeeeeeeee/eeeeepeee net forum?
12

---

Published as a conference paper at ICLR 2025
Kunchang Li, Xinhao Li, Yi Wang, Yinan He, Yali Wang, Limin Wang, and Yu Qiao. Videomamba: State space model for efficient video understanding. In European Conference on Computer Vision, pp. 237-255. Springer, 2024a.
Qianxiao Li, Ting Lin, and Zuowei Shen. Deep learning via dynamical systems: An approximation perspective. Journal of the European Mathematical Society, 25(5):1671-1799, 2022.
Shufan Li, Harkanwar Singh, and Aditya Grover. Mamba-nd: Selective state space modeling for multi-dimensional data. In European Conference Computer Vision, pp. 75-92. Springer, 2024b.
Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes convolutional models great long sequence modeling? In The Eleventh International Conference on Learning Representations, 2023. URL https: /tapttps://ppeeeeeeeee /tpp//ppeereeeee//eeeeee ew net/ forum G=GGSSPbRp.TX-
Dingkang Liang, Xin Zhou, Wei Xu, Xingkui Zhu, Zhikang Zou, Xiaoqing Ye, Xiao Tan, and Xiang Bai. Pointmamba: A simple state space model for point cloud analysis. Advances in neural information processing systems, 37:32653-32677, 2025.
Hongzhou Lin and Stefanie Jegelka. Resnet with one-neuron hidden layers universal approxi- mator. Advances in neural information processing systems, 31, 2018.
Fusheng Liu and Qianxiao Li. From generalization analysis to optimization designs for state space models. arXiv preprint AXXv:2445.026700 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint
17X00:171..05101, 2017.
Chao Ma and Ziyang Wang. Semi-mamba-unet: Pixel-level contrastive and cross-supervised visual mamba-based unet for semi-supervised medical image segmentation. Knowledge-Basee Systems, 300: 1 12203, 2024.
Andrew Maas, Raymond E Daly, Peter Pham, Dan Huang, Andrew Y Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings the 49th annual meeting of the association for computational linguistics: Human language technologies, pp. 142-150, 2011.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mix- ture models. In International Conference Learning Representations, 2017. URL https: / //openzeview. net forum? Bid
William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. In International Conference on Machine Learning, pp. 35492-35506. PMLR, 2024.
Nicola Muca Cirone, Antonio Orvieto, Benjamin Walker, Cristopher Salvi, and Terry Lyons. The- oretical foundations of deep selective state-space models. Advances Neural Information Pro- cessing Systems, 37: 37:127226-122 127272, 2025.
Eric Nguyen, Karan Goel, Albert Gu. Gordon Downs. Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4nd: Modeling images and videos multidimensional signals with state spaces. Advances in neural information processing systems, 35:2846- 2861, 2022a.
Ho Huu Nghia Nguyen, Tan Nguyen, Huyen Vo, Stanley Osher, and Thieu Vo. Improving neural ordinary differential equations with nesterov's 's accelerated gradient method. Advances Neural Information Processing Systems, 35:7712-7726, 2022b.
Tam Minh Nguyen, César A Uribe, Tan Minh Nguyen, and Richard Baraniuk. Pidformer: Trans- former meets control theory. In Forty-first International Conference on Machine Learning, 2024.
Tan Nguyen, Richard Baraniuk, Andrea Bertozzi, Stanley Osher, and Bao Wang. Momentummnn: Integrating momentum into recurrent neural networks. Advances in neural information processing systems, 33:1924-1933 1936, 2020.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference Machine Learning, pp. 26698. PML.R, 2023.
13

---

Published as a conference paper at ICLR 2025
Toshihiro Ota. Decision mamba: Reinforcement learning via sequence modeling with selective state spaces. arXiv preprint arXiv:2403.19922, 2024.
Badri Narayama Patro and Vijay Srinivas Agneeswaran. Mamba-360: Survey of state space models as transformer alternative for long sequence modelling: Methods, applications, and challenges. arXiv preprint arXiv:2404.16112 2024.
Maciej Pióro, Kamil Ciebiera, Krystian Kröl, Jan Ludziejewski, and Sebastian Jaszezur. Moe- mamba: Efficient selective state space models with mixture of experts. arXiv preprint M4Kiv:2401.0408/ 2024.
Sebastian Prillo and Julian Eisenschlos. Softsort: A continuous relaxation for the argsort operator. In International Conference on Machine Learning, pp. 7793-7802. PMLR, 2020.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188)99 2019.
Farnoush Rezaei Jafari, Grégoire Montavon, Klaus-Robert Müller, and Oliver Eberle. Mambalrp: Explaining selective state space sequence models. Advances in Neural Information Processing Systems, 18:40-050-1 2025.
2025.
Domenec Ruiz-Balet and Enrique Zuazua. Neural ode control for classification, approximation, and transport. SIAM Review, 65(3):735-773, 2023.
George Saon, Ankit Gupta, and Xiaodong Cui. Diagonal state space augmented transformers for speech recognition. ICASSP 2023-2023 IEEE International Conference on Acousties, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023.
Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for quence modeling. In The Eleventh International Conference on Learning Representations, 2023.
URL https://p////..
/oppnreviee net/ / forum? id=Ai BHxkkss
Paulo Tabuada and Bahman Gharesifard. Universal approximation power of deep residual neural networks through the lens of control. IEEE Transactions on Automatic Control, 68(5):2715- 2728, 2022.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena A benchmark for efficient transformers. In International Conference on Learning Representations, 2021. URL https: //openreview. wwt. wtwp://wwwwweeeetttt/eeieeeeee.e www.iiidiivyeeegevveeeee
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Bao Wang, Hedi Xia, Tan Nguyen, and Stanley Osher. How does momentum benefit deep neural networks architecture design? few case studies. Research the Mathematical Sciences, 9(3): 57, 2022.
Chloe Wang, Oleksii Tsepa, Jun Ma, and Bo Wang. Graph-mamba: Towards long-range graph sequence modeling with selective state spaces. arXiv preprint 47Xhv:2402.00789 2024a.
Junxiong Wang, Daniele Paliotta, Avner May, Alexander Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models. Advances in Neural Information Processing Systems, 17762432-62457. 2025.
Ziyang Wang, Jian-Qing Zheng, Yichi Zhang, Ge Cui, and Lei Li. Mamba-unet: Unet-like pure visual mamba for medical image segmentation. arXiv preprint arXiv:2402.05072 2024b.
Hedi Xia, Vai Suliafu, Hangjie Ji, Tan Nguyen, Andrea Bertozzi, Stanley Osher, and Bao Wang. Heavy ball neural ordinary differential equations. Advances in Neural Information Processing Systems, 18x18666-1888 18659, 2021.
14

---

Published as a conference paper at ICLR 2025
Rui Xu, Shu Yang, Yihui Wang. Bo Du, and Hao Chen. A survey vision mamba: Models, applications and challenges. arXiv preprint arXiv:2404.18861 2024.
Jianwei Yang, Chunyuan Li, Pengchuan Zhang, Xiyang Dai, Bin Xiao, Lu Yuan, and Jianfeng Gao. Focal attention for long-range interactions in vision transformers. Advances in Neural Information Processing Systems, 34:30008-30022, 2021.
Yijun Yang, Zhaohu Xing, and Lei Zhu. Vivim: video vision mamba for medical video object segmentation. arXiv preprint arXiv:2401.14168 2024.
Han Zhang, Xi Gao, Jacob Unterman, and Tom Arodz. Approximation capabilities of neural odes
and invertible residual networks. In International Conference Machine Learning, pp. 11086- 11095. PMLR, 2020.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text clas- sification. Advances in neural information processing systems, 28, 2015.
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vi- sion mamba: Efficient visual representation learning with bidirectional state space model. In International Conference on Machine Learning, pp. 22429-62444 PMLR, 2024.
15

---

Published as a conference paper at ICLR 2025
Supplement to "Demystifying the Token Dynamics of Deep Selective State Space Models"
A DYNAMICAL PROPERTIES TOKENS IN MAMBA
We will consider S6 dynamic in the single channel dimension case, i.e. = 1, in our setting to avoid overly complicated technicalities. In this case, the imput sequence x (#1...., .L) E R is a list of L real numbers. The dynamic of x is then governed by the dynamical system:
d (t) Σ (t). te 10,+
(15)
dt
(0) 310 E R,
for ..... L, where
(u) softplus(SAu) In (1+e uER,
(16)
and
(5) 21) ((x((t)),
if
(t)()=
(17)
(((S))=((taa) (xxs(t) exp -a Σ s(Ta(t)) kmj+1
if if
In this case, SB, Sc E Rx1, thus ST. e and the scalar values Sa and R, with 0, are learnable from imput data.
To avoid triviality, we will always assume that the initial datum 27) 3ıo are nonzero for all
1. The main goal of this section is to investigate the asymptotic behavior of the solution x(t) (t)) of the dynamic (15) and the corresponding hidden attention matrix
(t)
Pai(t)
P22(t)
P(t)
[Pi(())
Pra(()
(t)
when & increases. The following three scenarios will be considered separately in the subsequent subsections:
1. Convergence scenario: μ ST
2. Slow-divergence scenario: μ = ST. SB > 0 and SAZIo < for all I; and
3. Fast-divergence scenario: S 0 and SATI > 0 for some l.
A.1 CONVERGENCE SCENARIO: μ = ST <
We will start with the following auxiliary lemmas.
Lemma A.1. Let lo E R and let [u, !) be differentiable continuous function on RR x [lo +0x). Assume that u(t) is the unique solution of the initial value problem
d u(t) f(u(t), u(t). E [tto. +00),
dt u(ta) uo R. If ug (respectively, uo 0), then u(t) is positive (respectively, negative) function its maximal interval of the existence. In addition,
1. if f is a positive function, then u(t) is monotonically increasing (respectively, monotoni- cally decreasing).
16

---

Published as a conference paper at ICLR 2025
2. if f is negative function, then u(t) is monotonically decreasing (respectively, monotoni- cally increasing).
Proof. We assume that ug The case when uo is obtained by replacing u(t) by (t).
The function u(t) cannot be equal to zero on its maximal interval, since otherwise, u(t) 0 is the unique solution of the given differential equaion, which contradicts the fact that u(D) uo > 0. Therefore, u(t) is positive its maximal interval of the existence.
Next, assume that I is positive on R x |/o +o0). Then we always have (1) = ((((),, ()((l) for all E [4o, +oo). This shows that u(() is increasing on its maximal interval of the existence. The case when negative is similar.
Lemma A.2. Let to E R. If function u(t) is bounded on [to, +00), then there exists positive numbers c1 C2 such that
S Au((t) C2.
for all t [to. +00).
Proof. Recall that (uu(t) In(1 Saa((t)), Since u(t) bounded, there are 01 02 E R such that I Stu(t) & 02 for all E [fo, too). Then we can choose = =In(1 Eat) for i = 1,2. L
Theorem A.3 (Convergence scenario). Assume that STSB < O. Leix(t) = (t be the unique solution of the dynamic (15).
1. For each I 1......; L, æıo (respectively, Zí0 0), then rz(() is positive and mono- tonically decreasing (respectively, negalive and monotonically increasing) on [0, +x)). In addition, x1) = 0 (T)
2. (t) for all L,j L....,L. L.
Assume without loss of generality that æ1o > The case when z < is proved similarly by replacing (t) by (t). We rewrite the dynamic (15)
Proof.
dt =1( = (m) (i()))((t),
(18)
where
0,
ifl
(t)
1-1 H (1)2 exp
Σ A(or(()) k=
ifl > 1.
Since A < 0, (1) is nonpositive function. According Lemma A. 27((1) is positive and monotonically decreasing on right maximal interval of the existence [0, with T E (0, +oo]- Thus, (t) is bounded and T +oo. As consequence, follows from Lemma A.2 that there is a positive constant e >> 0 such that (tx)(t) 2 c for all E [0, +oc). Therefore, from Eq. (18), we have
d
dt (t) S μcz; ) (t)
which imply that (4) < (-2pet + xT²²) for all I E [0, +o0). Hence, .T((1) ((2)) as expected.
For the hidden attention score Pış(!) with L > 2 j 2 1, we have
= = (() -a Σ A(Ea(t)) (c(e(()),,())= =
k=j+1
expected.
17

---

Published as a conference paper at ICLR 2025
A.2 SLOW-DIVERGENCE SCENARIO: μ SGSp > 0 AND SaZ < 0 FOR ALL
In this case, we first observe that all of the tokens will tend to infinity as we will see in the following lemma.
Lemma A.4 (Slow-divergence scenario). Assume that μt = ST > 0 and Sa < 0 for all Let zxt) (z) (t),... ....,........................... (t)))) be the unique solution of the dynamic (15). For each I, if a1o > (respectively, Tio < 0), then (t) is positive and monotonically increasing (respectively, negative and monotonically decreasing) on [0, +00) with
lim (t) = +o0 (respectively, lim xz(t) -o0).
+0000
1-+00
Proof. Without loss of generality, we assume that SA < O. The case when SA > is then obtained by replacing zr() by -x((). Similar to the proof of Theorem A.3, rewrite dynamic (15)
dt =1((t) =
+ (()))((t))
(19)
where
0,
if
(((1) =
Σ S(Ca(t))) k=j+1
(20)
pac , (tz)))z((t) exp
exp
ifl > 1.
Since μ 0, (t) is nonnegative function on [0, +00).
According to Lemma A.I, mt(() must be positive increasing function on its maximal interval of the existence [0, T) for some 7' (0, +oo].
First, we claim that æ4 (1) is unbounded on [o, T). Indeed, if this is not the case, then (2) is bounded and T +00. According to Lemma A.2, there exists a positive constant such that (z)( 2 c for all t E [0, +00) Therefore, Eq. (19) yields
d (t) 2 (t) E (00)
which implies that (t) ( (-2pct+ ) ) 1, which is unbounded at The claim is then proved, and lim,-7- 21 +o0.
1 > 0, contradiction. cu2fo
It remains prove that T +00. Observe that, since SA < O, we have
3 In(1 + e
lim r Ar) = lim Sar ear
0.
--- e
Furthermore, since (t) is positive and monotonically increasing to infinity at infinity, we also have
lim ( tu+oc
O.
In particular, case 1, we have $ (t) < 1, which implies that (1) S t, for all large enough. Therefore, T +oo in this case.
Assume that / and we already have T too for cases j Then it folows from Eq (20) that
(t,(t) d
(t(
Σ pay (1) (1))
In(2, (t))
(21)
zt(t) dt
1=1
Therefore, we can bound
d dt (f) from Eq. (19)
d
(t) 1
Σ
(t) z((t).
dt
dt
1=1
and this inequality hold for all / for some large enough. As a consequence, have (S'ex)
H
(1) S exp Σ
- (n(e(())))
Σ
(n(2(())) - (n(2,((1)) ds + s2((1)
J=1
The right hand side is function defined for all h. Thus T +00.
18

---

Published as a conference paper at ICLR 2025
Remark 8. Lemma A.4 shows that all tokens in the considered case will approach infinity t approaches infinity. However, it does not help to estimate the rate of divergence or the asymptotic behavior of the hidden attention score Pty (t). In Theorem A.7 below, we will provide both the rate of divergence for the tokens and the asymptotic behavior of the hidden attention scores, but under a certain additional assumption.
We provide the divergence rate of the first token in the following lemma.
Lemma A.S. Assume that μ Sn 0 and Saxio 10 0. Let on (1) be the unique solution of the dynamic (15). Then 1 (1) is defined on [0, +x)) and xı (1) = (In(()).
Proof. Without loss of generality, let us assume that Sa < 0, thus 210 > 0. The case when SA can be obtained by replacing xi (t) by (t). Then it follows from Lemma A.4 that 2 (1) is positive and monotonically increasing function on +oc). Let us a small enough such that SA + e < 0. Then, there exist C1,
such that:
pmz (t)( (()))(())))))) (x)(()) = (S (t) Lad (t) In(1 + (85
I (Satet)))))))
(22)
(Sni(())
for all t E [t1, too). Therefore, we have
d dt (1) (S+t t ti:
which yields
2 In ( - t1) +e - (e-ssmu(h)) for all t 2 f1. SA
1 (t) -
Since æi (l) is positive, we finally obtain (1) Oln(()) as expected.
To determine the divergence rate of the other tokens, we will need the following lemma.
Lemma A.6. Let ro(== 2.1160) be the unique positive root of the function h(r) 2 2 + + re_ Then for > 0, the function f(r) = r² In (1 e e>rr is positive and monotonically de- creasing on [ +00).
I-e-r-
Proof. The lemma follows from the negativity of the function h(r) on [ro. +oo) and the equality d f(r rh(Ar).
Theorem A.7 (Slow divergence scenario and divergent rate). Assume that j = ST B > 0 and
SATIO _... SATI0 S -Tn,
(23)
where ro( 2.1160) is defined in Lemma A.6. Let xxt) = (t) (t)) be the unique solution of the dynamic (15).
1. For each I, if 3ıo > 0 (respectively, < 0). then x1(t) is a positive increasing (respec- tively, negative decreasing) function on [0, +o0). In addition, 2;(t) O(I
2. lim (t) for all 2 j.
Proof. Without loss of generality, we will assume that SA - for some À > 0. The case when SA > 0 then obtained by changing all i to -t, Then the condition (23) becomes
To
L 2 2 x10 Z
(24)
According to Proposition A.4, .1) is a positive and monotonically increasing function on [0, +00) with lim; (1) +00
Before estimating the asymptotic behavior of (t) and (t)(t), first claim that that 2x(() 2 it)) for all / E +oo) and / 2 2. We will prove this claim by induction on l. Indeed, in case I 1= 2, we observe that
d di =2( = ((=(()) (2())a(()) 2 " ((xx(())
((z(())
19

---

Published as a conference paper at ICLR 2025
Since T20 2 Tiu. it follows that (2()) 2 (t) for all t [0, +oo). Let assume that / >
and (t))((t 2 ... 2 (t) for all t E +0)) It is noted that x1(() 2 x10 2 x. According to
Lemma A.6, have
2y (x())(((2((1)) 2 (x) j= 1.2...,1-1
(25)
Recall that
d
dt z1(() (urr(t) ((xr(t) + ((t)aa(()
(26)
1-1
A(me(1)
(z) + (t) Σ ==)()))(z)((0))m) exp
Σ
(27)
k=
By removing the term with the index = inside the sum and using inequality (25), we obtain
d
L-1
1-1
dt x1(t) 2 (t)(t) + (t)
exp
a
Σ
A(xx(t))
j=2
k=j+1
1-2
- (()(()) + (m)(1)
Σ
xz (1; exp
Σ
(axx(1))
j=1
k=j+2
Since (xxx(1)) - (( (1)) for k= 2, ...., 1, we can proceed the last expression as
1-2
((rx(1) - (asca(1)))
di ==((1= 2 p., (I=)))(((((())) + (1)
Σ
j=1
k=j+1
Therefore, if set
L-2
((xx(t)) as(m)
f(u) = pu + pau Σ
exp
Σ
k=j+1
then (t) is solution of the differential equation 4. u(t) (u(()), while di (t) 2 (z) for all t E [0, +oo). It follows from the initial conditions 2 21-1.9 that 2 21 2 21-1(t) for all i E [0, +00). The claim is then proved.
Next, let us go back to the analysis of the asymptotic behaviour of tokens and hidden attention scores.
1. For 1 = 1, it is already known from Lemma A.5 that æ1 (1) Oln(()). In case 1 > 2. according to equation (27), we have
(1)(()) x1(()
1-1
= parz
+ Σ (1),()))(( exp
-a
Σ ((Ta(0)) k=
(28)
For each j = 1,...,,, I, since x, (t) 2 21 (t) and the function r (r)(r) is monotonically decreasing over 1 +oo), we have x, (t)² (t) (zz
S x (t) for all > 0 large
プ (t)x(t)) < 1 for all t E +oo).
enough. In addition, we also have and exp kmj+1
Therefore, we can upper bound the right hand side of equation (28) as
d #1((1)) < [uoo (t))((t))((z(())))) = _ (1)
(29)
(t)
(t)
Thus, (n) S In (lnzz)()))) (t), which implies that
m2(1) x (1)² = O((n))).
20

---

Published as a conference paper at ICLR 2025
2. For the hidden attention score, it follows from its definition in equation (17) that
Ptş(t) pm (1)(
for all 2 O. Since Sa < 0, the function r ~ rr(r) = r In(1 + e is monotonically decreasing from large enough t. Therefore, it follows from the claim above that
(t)(t) S (x(()))(((((t))
for all / large enough. Hence, lim Ptş(() = lim, pr Ar)) = 0.
The theorem then proved.
A.3 FAST-DIVERGENCE SCENARIO: μ # == ST > 0 AND SAZI0 > 0 FOR SOME
Theorem A.8 (Fast-divergence scenario). Assume that μ St > 0. Let x(t) (ri(t)....,, (t) be the unique solution of the dynamic (15). If there exists = 1,..., such that SATI 0, then (t) goes to infinity at finite time.
In addition, assume that Fio (t) tends to infinity first as t T for some T > while z (t) is defined on T) all 1 + lo. Then for every I. 1.2 12 j 1, li Piş(() +000 if and only if 1 1= lo or lo-
Proof. Without loss of generality, we assume that Sa 0. The case when Sa is then obtained by replacing x(f) by -x()). Then T1o > 0. Recall that
d =t)= (m)
+ (t)))())))
(30)
where (t( defined in equation (20). Since μ > 0, (t) a nonnegative function on [0, +oo). follows from Lemma A. that (t) is positive increasing function on its right maximal interval of the existence [0, 7) for some T E (0, +oo]. By equation (30), we have
d m1( 2 pay (
di
dt
Furthermore, since SA > 0, the function A = In(1 + (ar) is monotonically increasing on [0, too). Therefore, ((xx(t)) 2 (())()) Im for all e [0, +oo). As consequence, we can proceed the above inequality further as
d
(t(t) 2 pIn(1 + Sa )z) (t)²
dt
Thus,
m1(1) (2 + eSaio -)
which goes to +00 at finite time.
Next, let us assume that (t) goes to infinity at finite time faster than the other tokens, i.e, there exists > such that lim,-rr 31, (1) = +oc while s1(() defined over [0, T] for all + lo- Then lim,-T 2 (())(a((() +o. Then for every L 12120 j 2 1, lim--- Ptş(() = x if and only if to or J lo-
B HIGHER DIMENSION WITH/WITHOUT ADDITIONAL COMPONENTS
When the dimension of the input tokens is D > 1, the input-output projection matrix ST a is a square matrix of size Dx x D, which may have complex eigenvalues. We conjecture that the dynamic behavior of the tokens in system (7) can be determined by analyzing the matrix
(S.Sp + SB)
# = 2
21

---

Published as a conference paper at ICLR 2025
which represents the symmetric part of the input-output projection matrix Sg... This conjecture is motivated by the observation that the term x SB) in the considered dynamical system (7) a scalar, and it can expressed as:
कf SS I - [ x (S( x1 + (z) (ST 2:) =1) x
21.
The matrix μ has only real eigenvalues. We particularly conjecture that tokens will converge to zero if has only negative eigenvalues, whereas the tokens will diverge to infinity if μ has at least one positive eigenvalue.
Figures illustrate this conjecture by visualizing the trajectories of four tokens in the two- dimensional case under different scenarios. In these figures, the parameters and initial values used on each case given follows:
Figure 6A (μ has two negative cigenvalues):
0.552679
-0.843293
STSB =
thus the eigenvalues of μ are 00.967455 and
0.869146
-0.967042
-0.552776,
0,287585
0.99662
SA -0.201208
-0.964587
-A 0.
initial values #1
(-1.47982, 0..27882 32 (-0.406453, 1.24415), 23
(1.8491,-00 -0..,-0.62500 0..655355.. 24 (1.26989, -1.91216).
Figure 6B (μ has one positive eigenvalue and one negative eigenvalue):
-0.155283 0.542694
STSB =
thus the eigenvalues of are 0.846723 and
0.989821
0.260748
-0.741258,
0.430263
0.555071
SA
-0.654555
0.422737
-A -0.573698T.
-initial values 21 = (1.39902, 1.60688). 22 (-0.342366, 0.845203). Is
(0.616744, 1.63846), .14 = (-0.185335, -..34666).
Figure 6C (μ has two positive eigenvalues):
0.981721 0.342524
-0.802299
STSB =
thus the eigenvalues of are 1.39646 and
0.605171
0.190429,
0.653732
-0.228578
SA 00.960714
-0.495344
A 0. .99940872.
initial values F1 (1.21287, -0.816296). IA = (-0.462258, 1.44549).
(1.1932. -0.702409) 22 (-1.49159, 00.733305,, 23
We observe from Figures 6C that all tokens converge to zero when y has only negative eigenvalues. While all tokens diverge to infinity when y has at least one positive eigenvalue.
C EXPERIMENT DETAILS
In this section, outline the experimental setup in detail. All experiments are conducted using server with four A100 GPUs.
C.I AUTO-REGRESSIVE LANGUAGE MODELING
Dataset We utilize the WIKITEXT103 (Merity et al., 2017), created from Wikipedia articles. It includes a training set of approximately 28,000 articles, amounting to 103 million words total. Each article is split into sections of about 3,600 words. The validation and test sets contain 60
22

---

Published as a conference paper at ICLR 2025
(A)
(C)
2.0
15
1.0
0.5
-1.0
(B)
00
X3
-00
02
-0.05
-30 10
-0.0
0.50
0.05
0.10
0.15
Figure 6: Tokens dynamic behaviors in two-dimensional case where the input/output projection matrix Sp has two positive eigenvalues (A), or one positive eigenvalue (B), or two negative eigenvalues (C). We observe that, in cases (A) and (B), all tokens will diverge infinity, and in case (C), all tokens will converge to the origin.
articles each, with word counts 218,000 and 246,000 respectively, combining for total roughly 268,000 words.
Parameterization for three scenarios: We outline the parameterization used to ensure the S6 layer operates within the three scenarios described in Section 5.1. For the first two scenarios, lever- age the LDL" decomposition (Golub & Loan, 2013) to ensure the positive or negative definite property of input-output matrix. Specifically, lower triangular matrix with diagonal values fixed at 1. In the
parameterize Sc I and Sa DT. where
D is a diagonal matrix and first scenario, the diagonal values of D) constrained to be positive and are parameterized as D = (sof(sssspuu(dd)) In the second scenario, this constraint is reversed by parameterizing D diag (oog(sfftlus(d")) ensuring that the diagonal values are negative. In the third scenario, where both positive and negative eigenvalues are required, we modify the diagonal matrix D by changing the sign of half of its diagonal elements. Specifically, we use D diag(s @ softplus( (d')) ),) 1] is vector with half of its elements set to and the other half set
where = [1,..., 1, -1
1.
Training details We train 130M parameters version of Mamba model (Gu & Dao, 2024) on WIKITEXTIO3 for 25 epochs with GPT-2 tokenizer (Radford et al., 2019) and AdamW optimizer (Loshchilov & Hutter, 2017). Other hyperparameters set identically across three scenarios and listed Table
C.2 IMAGE CLASSIFICATION
Dataset The ImageNet-IK dataset (Deng et al., 2009) is a widely recognized benchmark in the field of computer vision, commonly used for training and evaluating models on the task of large-scale image classification. It contains a collection of 1.28 million labeled training images and 50,000 validation images. Each image is labeled with one of 1,000 distinct classes, representing broad variety objects, animals, scenes, and more, allowing the model to learn a rich and diverse range of visual concepts.
23

---

Published as a conference paper at ICLR 2025
Table 4: Hyperparameters configuration for Language Modeling task on WIKITEXTIO3.
Sequence Length
1024
Peak Learning Rate
0,0015
Momentum
81
0.9; (32 0.999
Weight Decay
0.25
Batch size
128
Learning rate warmup
Linear Cosine decay
Learning rate scheduler Dropout
0.25
Table 5: Hyperparameters configuration for Image Classification task ImageNet-IK.
Optimizer
AdamW
Peak Learning Rate
0.0008
Momentum
81
0.9; (32 0.999 0.05
Weight Decay
Batch size
512
Learning rate warmup
Linear Cosine decay 0.0
Learning rate scheduler Dropout
Training details We train the tiny version (31.8M parameters) of Mamba Vision model (Hatamizadeh & Kautz, 2024) with and without our token reordering technique. We follow the training procedure in (Yang et al., 2021; Hatamizadeh et al., 2023; Hatamizadeh & Kautz, 2024) and set the hyperparameters for two cases identical, which is listed in Table 5.
D ADDITIONAL EXPERIMENTAL RESULTS
D.1 RE-ORDERING INPUT TOKENS
In this section, we conduct an additional experiment to assess the performance of our re-ordering technique across tasks beyond the vision task reported in Section 5.2.
D.1.1 LANGUAGE MODELING
We first present experiment to evaluate the effectiveness of our reordering technique across broad range of architectures. We use the language modeling task on the WikiText- 103 benchmark to test three models: Mamba (Gu & Dao (2024)), HB (Fu et al. (2023)), and Transformer (Vaswani al. (2017)). For the latter two models, our importance score Equation 14 cannot be directly calculated as these models do not have the step size matrix Sa. To address this, we made the following adjustments:
H3 Model: We calculate the token importance score as (z, K), where K a learnable vector.
We converted a pretrained model (GPT-2 small (Radford et al. (2019)) in experiment) into the Mamba format and fine-tuned it, following the method outlined in (Wang et al. (2025)). On the converted model, we calculated the token importance score described in Equation 14.
Transformer Model:
The results, presented in Table 6, demonstrate that the reordering method consistently improves perplexity across all tested models. These findings highlight the potential of the reordering approach to enhance the performance of a wide range of architectures.
D.1.2 TEXT CLASSIFICATION
We now evaluate our method on text classification tasks using the IMDB (Maas et al. (2011)) and AGNews (Zhang al. (2015)) datasets, using Mamba (Gu & Dao (2024)) baseline model. The
24

---

Published as a conference paper at ICLR 2025
Table 6: Test perplexity on different models
Model
With reordering
Without reordering
Improvement with reordering
Mamba
14.97
16.46
1.49
H3 Transformer
20.72
21.61
0.89
13.44
16.98
3.54
results in Table 7 indicate that our re-ordering technique further improves the baseline model's performance on these tasks.
Table 7: Test accuracy (1)
IMDB and AGNews datasets
Model
IMDB
AGNews
Mamba (baseline)
88. 88.46
92.08
Mamba + Token reordering
88.66
92.37
D.2 IMPLEMENTATION DETAIL AND COMPUTATIONAL. COST
To provide clearer understanding of our re-ordering technique, we include the pseudo-code for calculating token importance scores and forwarding through a single SSM layer in Algorithm The computational overhead of our technique lies in steps 1, 2, 3, and 4, which require O(D² L+ D L²) operations per input sequence.
D L +
Algorithm Forwarding through a SSM layer with reordering
Require:
Input sequence x E IR DxL, hidden matrix A E RNXN, input matrix SB E RNXD output matrix Sc E R NxD, step size matrix Sa E RDDD learnable vector E R hyperparameters I and p of Softsort.
Ensure: Output sequence y E R
1: Calculate importance
for each token in sequencce, combined into vector
(SAx) K E _
2: Sort the score vector decreasing order: Ssorted Sort(s) E 4x1
3: Calculate the reordering matrix P Softmax (Scoread -1us "") {The power function is
applied element-wise and Softmax is applied row-wise)
4: Reorder the input Xycordeedd x/T 13 D
5: Calculate the output sequence with selective scan
y SSM(A, Sgrreordeed,, Scxreordered) (Xreordered) (
6: return
Additionally, we present the empirical computational costs of our experiment in Section 5.2, bench- marked on the ImageNet classification task using MambaVisions (Hatamizadeh & Kautz (2024)) as the baseline. Table 8 summarizes key metrics, including memory usage, training time per epoch, parameter count, and FLOPS. The results show that our reordering technique introduces minimal overhead, with only slight increases across all metrics.
D.3 ROBUSTNESS ON WIKITEXTI03 EXPERIMENT
To assess the robustness of our findings across three scenarios of input-output matrices on the WIKI- TEXTI03 (Merity et al. (2017)) language modeling task, we conduct the same experiment described in Section 5.1 with three different random seeds for each scenario and report the mean and standard deviation. The results Table 9 suggest that our findings are robust to different random seeds.
D.4 DYNAMIC OF STANDARD MAMBA ARCHITECTURE
In this section, we present an experiment to analyze the token dynamics in the standard Mamba architecture in practical setting. We evaluate three cases:
25

---

Published as a conference paper at ICLR 2025
Table Comparision of computational cost on model training with Mamba Vision-T (Hatamizadch & Kautz (2024)) as Baseline.
Model
#params
Training time / epoch
Memory / /GP
FLOPS
Baseline
31.79M
308s
7232MB
8.93GFLOPs
reordering
31.79M
322s
7334MB
8.94GFLOPs
Table 9: Test perplexity on WIKITEXTIO3 (L). Mean and standard deviation are computed over three runs with different random seeds. Models with fewer negative eigenvalues in Sg consis- tently achieve lower perplexity.
Scenario
Perplexity ())
Negative eigenvalues
17.28 0.02
Mixed eigenvalues
16.82 + 0.02
Positive eigenvalues
16.66 + 0.05
12.
20-
10-
13.
0.8
0.6
02-
15
20
25
10
15
20
25
10
15
Dapth
Dopth
(a) Unconstrained eigenvalues
(b) At least
positive eigenvalue
(c) Negative eigenvalues
Figure 7: Token norm evolution in standard Mamba architecture without LayerNorm.
12
12.5
15
7.5-
2.5
02.
0.0
10
15
10
15
20
25
0
10
15
Depth
Dopth
(a) Unconstrained eigenvalues
(b) At least positive eigenvalue
(c) Negative eigenvalues
Figure 8: Token norm evolution in standard Mamba architecture with LayerNorm.
(a) There is no constraint the input-output projection matrix.
(b) The input-output projection matrix is symmetric and has at least one positive eigenvalue.
(c) The input-output projection matrix is symmetric and has only negative eigenvalues.
We first conduct simulations using pretrained standard Mamba architecture without LayerNorm on WIKITEXTI03. In Figure shows that tokens converge toward zero when the input-output projection matrix has only negative eigenvalues. while tokens diverge when the input-output projection matrix has at least one positive eigenvalue. Moreover, when there is no constraint on the input-output projection matrix, i.e., when we do not force to always belong to the class of positive definite or negative definite matrices, the tokens' dynamic behaviors become quite complicated. This figure suggests that our theoretical analysis of the tokens' dynamic behaviors in the standard Mamba model without LayerNorm might still be valid.
visualize the evolution of token norms across the blocks. Figure
In addition, we also conduct simulations using a pretrained standard Mamba architecture with LayerNorm in practical setting WIKITEXT103. In Figure visualize the evolution of token norms across the blocks. Figure shows that the tokens' dynamics in a full Mamba setting can be quite complicated. When the input-output projection matrix has only negative eigenvalues,
26

---

Published as a conference paper at ICLR 2025
the tokens do not necessarily tend toward zero. In particular, with LayerNorm in the Mamba archi- tecture, our theoretical results on the tokens' dynamics might no longer hold.
D.5
ABLATION STUDY
We perform ablation study to analyze the effects of various components of Mamba in the token ordering technique. We examine the following configurations, both with and without the reordering technique:
The full Mamba setting,
(2) Mamba setting with weight sharing across layers,
(3) Mamba setting without layer normalization,
(4) Mamba setting without the short convolution layer.
To ensure fairness, all configurations are standardized to approximately 129M parameters. The re- sults, shown Table 10, highlight the significant contributions of each component and the reorder- ing technique to the model's performance, emphasizing their collective importance in achieving optimal results. Our reordering method consistently yields at least 1.49 PPL improvement in all settings, which is significant for language modeling task with WikiText1033
Table 10: Ablation study language modeling task with WikiText103.
Model
With reordering
Without reordering
A
Mamba
14.97
16.46
1.49
Mamba weight sharing
17.77
20.51
2.74
Mamba without LayerNorm
17.82
19.31
1.49
Mamba without Convolution
21.81
23.38
1.57
27


Reviews:
[{"review_id": "I6vBIt5mfc", "reviewer": "Reviewer_FpJk", "summary": "This paper investigates the dynamical properties of tokens in a pre-trained Mamba model, a type of selective state space model (SSM). Despite the empirical success of Mamba, a comprehensive theoretical understanding of its internal mechanisms is lacking in various areas. The authors address this gap in one area by deriving the dynamical system governing the continuous-time limit of the Mamba model and analyzing its asymptotic behavior. In the one-dimensional case, they prove that the system exhibits one of two behaviors: either all tokens converge to zero or diverge to infinity. They provide criteria based on model parameters to determine which scenario occurs. Empirically, they verify that the convergent scenario negatively impacts model performance, while in the divergent scenario, tokens diverge at different rates, contributing unequally during training.\nBased on these findings, the authors propose two refinements to improve model performance:\n* Excluding the Convergent Scenario: By ensuring certain conditions on model parameters, they prevent the tokens from converging to zero.\n* Reordering Tokens: They reorder tokens based on their importance scores to account for the unequal contribution during training.\n\nExperimental results on language modeling and image classification tasks validate these refinements, showing improved accuracy and convergence speed.", "questions": "1. Could you extend the theoretical analysis to higher-dimensional cases?\n2. Can you provide more details on how the token importance scores are computed and how the reordering is implemented? What is the computational overhead of this process?\n3. Do your findings and proposed refinements generalize to other selective state space models beyond Mamba?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Originality: The paper addresses a gap in the theoretical understanding of selective state space models, specifically Mamba, by analyzing the continuous-time dynamics of tokens.\n- Theoretical Rigor: Provides rigorous mathematical proofs for the asymptotic behavior of tokens, offering clear criteria based on model parameters.\n- Practical Implications: The findings lead to actionable refinements that improve model performance, validated through experiments on language modeling and image classification tasks.\n- Clarity: The paper is well-structured, with logical progression and clear explanations that make complex mathematical concepts accessible.", "weaknesses": "- Simplifying Assumptions: Using time-independent parameters and excluding other layers for mathematical convenience may oversimplify Mamba's actual behavior in practical settings.\n- Experimental Scope: Though supportive, experimental validation could be more extensive. Including additional datasets and comparisons with other models would strengthen the claims.\n- Limited Discussion on Limitations: The paper could benefit from a more thorough discussion of the limitations of the proposed refinements and potential challenges in practical implementation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JrrOsu3fFK", "reviewer": "Reviewer_3W4C", "summary": "The paper \"Demystifying the Token Dynamics of Deep Selective State Space Models\" explores the theoretical behavior of tokens in the Mamba model, a type of Selective State Space Model (SSM) used in sequence tasks like language modeling and image processing. Despite Mamba's empirical success, its internal dynamics are not well understood.\nThe authors show that in a one-dimensional case, tokens either converge to zero or diverge to infinity, depending on model parameters. Convergence harms performance, while divergence leads to unequal contributions from tokens during training. To address this, the paper proposes two improvements: eliminating the convergent scenario and reordering tokens by importance. Experiments confirm these refinements improve accuracy and speed, particularly in image classification tasks.\nThis study offers new theoretical insights into token dynamics and suggests practical adjustments to improve the Mamba model's performance.", "questions": "**Summary of Actionable Improvements**:\n- Expand the theoretical analysis to higher dimensions or at least provide empirical evidence for higher-dimensional token dynamics.\n- Broaden the experimental validation to include diverse tasks beyond vision (e.g., language modeling and reinforcement learning).\n- Discuss the computational costs and trade-offs of the proposed refinements in real-world applications.\n- Increase the novelty by exploring more complex dynamical behaviors or offering a deeper comparison with existing models like transformers.\n- Provide a more comprehensive experimental comparison to other recent SSMs and transformers, including computational efficiency.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- **Originality**: The paper makes a meaningful contribution by exploring the often overlooked internal dynamics of tokens in Selective State Space Models (SSMs), specifically the Mamba model. While SSMs are well-known for their efficiency in sequence modeling, this work tackles a fresh problem—understanding how token behavior impacts model performance. By analyzing the continuous-time limit of token dynamics, the authors open up a new line of inquiry that could lead to more informed model designs and optimizations. The refinements proposed—eliminating the convergent token scenario and reordering tokens by their importance—are practical ideas that could lead to significant performance gains.\n- **Quality**: The paper is technically solid, with rigorous theoretical analysis that provides clear criteria for determining whether tokens in the Mamba model will converge to zero or diverge to infinity. This kind of formal, mathematical foundation is a strong aspect of the paper, ensuring that its claims are well-supported. The empirical results on large-scale tasks such as ImageNet classification offer clear evidence that the proposed refinements are effective in improving model accuracy and convergence speed. The experiments directly tie into the theoretical findings, showcasing a strong link between the analysis and practical outcomes.\n- **Clarity**: The paper is clearly structured, with a logical flow from the problem statement, through the theoretical analysis, to the empirical validation. The technical content, while dense in some areas, is well-explained, and the use of diagrams and tables helps to clarify key findings. The authors do a good job of contextualizing their work within existing research, making it clear how their contributions build on prior work in SSMs and neural ODEs. Overall, the clarity of the presentation allows readers to follow the complex theoretical insights without becoming lost in technical details.\n- **Significance**: The significance of the work lies in its potential to impact both the theoretical understanding and practical application of SSMs. By uncovering how token dynamics influence model performance, the paper provides a foundation for improving not only Mamba but other models that rely on similar state-space mechanisms. The proposed refinements—removing harmful convergent behaviors and reordering tokens based on importance—offer actionable insights that can directly improve model training efficiency and accuracy, as shown by the experimental results. Given the growing use of SSMs in fields such as language modeling, computer vision, and reinforcement learning, these contributions could lead to broader advancements in how such models are designed and optimized.", "weaknesses": "- **Theoretical Scope is Too Narrow**: \nWhile the paper makes a valuable contribution by exploring token dynamics in the one-dimensional case, the focus is too limited. The analysis is restricted to this specific setting, which weakens the generalizability of the conclusions. The authors mention extending their findings to higher-dimensional cases, but without any concrete exploration of these cases, the scope of the theoretical contribution feels incomplete. Actionable Insight: A more comprehensive exploration of higher-dimensional token dynamics should be included. If full proofs are too complex for this paper, at least preliminary results or conjectures, backed by empirical evidence, would make the work stronger. Additionally, comparisons with more sophisticated state-space approaches (such as diagonal vs. non-diagonal state space models) could show how the proposed ideas generalize.\n- **Insufficient Diversity of Experiments**: \nThe empirical validation is focused on tasks like ImageNet classification, which is a strong benchmark, but does not provide sufficient diversity to substantiate the claims made in the paper. The experiments demonstrate the refinements' effectiveness in vision tasks, but the proposed improvements—such as reordering tokens by importance—are presented as potentially broadly applicable to sequence models across various domains. The paper could have benefited from testing the refinements on different types of data and tasks, such as language modeling or reinforcement learning, where SSMs are also commonly applied. Actionable Insight: Expanding the experiments to include tasks like language modeling(e.g., WIKITEXT or other sequence-based datasets) or sequence-based tasks in reinforcement learning would provide stronger evidence that the proposed refinements generalize beyond vision tasks. This would substantiate the claim that these refinements are broadly beneficial across domains.\n- **Shallow Investigation of Practical Impact**: \nThe paper claims that eliminating the convergence scenario and reordering tokens significantly improves the performance of Mamba. While the empirical results show improvements in accuracy and convergence speed, there is little discussion on the practical implications of these refinements in real-world applications. For instance, the computational cost of implementing token reordering or of checking for convergence scenarios during training is not discussed. Additionally, the analysis of why these changes work, especially in larger-scale, practical deployments, is somewhat superficial. Actionable Insight: Including an analysis of the computational overhead or potential trade-offs of implementing these refinements in large-scale models is crucial for practical relevance. Furthermore, deeper explanations of why these adjustments improve performance, with specific references to token dynamics, would enhance the practical significance of the work. For instance, analyzing the impact of token reordering on different architectures or showing how it interacts with other model components would provide more actionable insights for model developers.\n- **Limited Novelty of Some Contributions**: \nWhile the paper presents the analysis of token convergence and divergence as a central contribution, this analysis—though useful—builds heavily on well-known concepts from continuous-time neural networks (e.g., neural ODEs) and does not introduce significantly new ideas. The observation that token dynamics can either converge to zero or diverge to infinity, depending on model parameters, is mathematically rigorous but not particularly surprising in the context of dynamical systems. This reduces the novelty of the contribution. Actionable Insight: To strengthen the originality of the work, the authors could extend the analysis beyond what is currently expected in dynamical systems and neural ODEs. For instance, they could explore more complex behaviors in the interaction between tokens, such as the emergence of periodic behavior, chaotic dynamics, or the impact of noise on the stability of the token trajectories. Alternatively, a deeper comparison with transformer-based models (such as interpreting token dynamics in transformers vs. Mamba) would increase the novelty of the work.\n- **Lack of Comprehensive Comparison to Prior Work**: \nThe paper briefly acknowledges prior work on SSMs and transformers but lacks a comprehensive comparison to alternative models. Given the ongoing research in state space models and their applications, a more thorough evaluation against recent models like GSS (Generalized State Spaces) or HiPPO-based models would provide a stronger benchmark for the proposed refinements. Additionally, since the paper suggests that Mamba can be a more efficient alternative to transformers, it should compare the computational efficiency and performance against contemporary transformer models on long-range sequence tasks. Actionable Insight: The authors should include a more detailed experimental comparison against state-of-the-art models in SSMs and transformers. They should also explicitly compare their proposed refinements with other model optimizations, such as low-rank approximations or efficient attention mechanisms in transformers. This would give the community a clearer picture of where Mamba, with its refinements, stands in the current landscape of sequence modeling.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fzQPjGDDHy", "reviewer": "Reviewer_biXN", "summary": "The paper provides a theoretical analysis of the dynamical properties of tokens in input-varying SSMs, particularly focusing on the Mamba architecture. The authors derive and analyze the continuous-time limit of the Mamba model, characterizing how tokens evolve through the network. Their key finding is that under several assumptions, in the one-dimensional case, tokens either all converge to zero or all diverge to infinity, with specific criteria determining which scenario occurs. For the divergent case, they prove that different tokens diverge at different rates, suggesting unequal contributions during training. Based on these theoretical insights, they propose two practical improvements: 1. excluding scenarios where tokens converge to zero, as this negatively impacts performance, and 2. reordering tokens based on importance scores. They validate these improvements experimentally on language modeling and image classification tasks.", "questions": "My main concerns I would like addressed are listed in the weaknesses section. Here are some additional comments or questions:\n\n- Perhaps you could start with a model pretty close to that presented in the theory: no layernorm, convs or extra MLPs; weight sharing, 1 dimensional, and then ablate these different choices. How much do each of these choices affect the results? Would also be interesting to see how the results vary with a varying number of layers, e.g. increasing the number of layers to better approximate the \"continuous-time limit\".\n\n- Wikitext 103 is very much a toy language task that is known to be very sensitive to hyperparameters. How robust are these results to different random seeds? Or different hyperparameters?\n\n- Line 075: The input, output, and state matrices are mislabeled\n\n\nUpdate: I have increased my score to reflect my updated view of the paper after the rebuttal period.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- I believe this is the first paper to present theoretical analysis related to how the token dynamics evolve through the networks across the depth of deep SSMs. This can be valuable to help provide insights into how these models process information\n- The mathematical analysis appears to be sound and the empirical results provide some evidence that the theory can lead to insights that can improve performance\n- The paper is well-structured and Figures 1-3 are useful.", "weaknesses": "- The effects of the assumptions made for the analysis are not thoroughly explored\n    - The analysis ignores the effects of important practical additions used in these networks such as layer normalization, short convolutions and the dense linear operations or MLPs often used between layers\n    - The analysis assumes weights are shared across layers\n    - The analysis assumes the one-dimensional case and excludes the possibility of complex eigenvalues in the input/output projection matrix \n    - All of these assumptions are fine for the sake of getting started with a simple case of analysis, but I would have liked to have seen a more thorough discussion of these assumptions and exploration in the empirical results.\n\n- The empirical results seem promising but modest. In addition, it is unclear how much the proposed improvements actually help. In particular, the scenarios investigated empirically seem to have strayed far from the assumptions in the theory (e.g. I believe no weight sharing, higher dimensions, and layer norm, short conv and MLPs included). I would have liked to have seen a series of ablations that build up from the simplified case considered in the theory to the full models that are used in practice. This would give a better sense of how important the assumptions are, and how much the proposed improvements are contributing to the performance boost. \n\n- More care should be taken to distinguish between the depth direction of the network and the sequence direction of the network. SSMs such as S6 come from the prior S4 and S5 methods which are based on continuous-time parameterized ODEs that evolve along the sequence length (often thought of as the \"time\" direction) and are then discretized. In this work, the continuous-time formulation is taken to be across the depth (i.e. layers) of the network and the \"time\" is referred to in this depth direction. However, this could lead to much confusion and I would recommend being much more explicit about what the \"time\" direction is more often throughout the paper to prevent confusion. E.g. in line 212, it is very unclear that time-independent parameters here is meant to mean the parameters are shared across layers and not shared across the sequence length (i.e. a LTI SSM sytem).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1U1sOyC6t6", "reviewer": "Reviewer_jxmT", "summary": "Deep selective state-space models have achieved promising results on sequence modelling tasks. Yet, theoretical understanding of these models remain limited. This paper proposes to study how the embedding of different tokens evolve as they are processed through Mamba layers. In particular, they show that they either converge to 0 or diverge to infinity, based on some criteria that they provide. Importantly, this divergence speed depends on the position of the token. Motivated on their findings, the authors propose some modifications to the standard Mamba architecture that improve its performance.", "questions": "See above.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "Given the very little theoretical understanding of these models, such a theoretical contribution is valuable to the community. Despite its theoretical focus, the paper is easy to follow and one can grasp the main insights without having to delve into the math. The empirical improvements the theory motivates seem to be significant.\n\n(Score increased from 5 to 8 after the discussion period)", "weaknesses": "My main concern with the paper is that its conclusions remain limited compared to what the analysis could, in my opinion, deliver. Here are a few questions that would help, once answered, to improve this point:\n- How does the token dynamics in Mamba compare to the ones of Transformers? Are there some qualitative differences? Does it lead to some insights on the different inductive biases of these architectures.\n- Is it possible to extend the analysis to the bidirectional case, i.e. without the causal masking (the implicit attention matrix would not be lower triangular anymore)?\n- Similar question than the previous point but this time to extend the analysis to gated RNNs such as LSTMs or GRUs, whose simplified versions have some similarity with Mamba.\n\nWhile it is totally fine to restrain the theoretical analysis to the 1D, simulations assessing whether the conclusions of the analysis on the 1D case also hold in higher dimension are missing. One example of such simulations would be to initialize a Mamba layer in the same way as it is usually done in practice, give it random inputs, and observe how tokens evolve. It might be interesting to play with how correlated the random inputs are and with the magnitude of $a$.\n\nFinally, I remain suspicious of the empirical results given that no learning rate tuning was reported. Without such tuning I find the results hard to trust.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)