Demystifying the Token Dynamics of Deep Selective State Space Models

## Abstract

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.

## 1 Introduction

State space models (SSMs)
have undergone significant advancements to mitigate the computational inefficiency associated with the sequence modeling [, , , , , , , ], and they have been successful applied in various domains involving continuous signal data such as audio and vision [, , ].
SSMs can be viewed as a combination of recurrent neural networks (RNNs) and convolutional neural networks (CNNs), drawing on concepts from traditional state space frameworks [].
Unlike transformers, which experience quadratic scaling with respect to input sequence length and exhibit significant computational demands [], 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 [] across various data types, and they have excelled in benchmark tasks such as the Long Range Arena [].

Selective SSMs such as Mamba is a particular SSM which involves the selective state-space layer (S6) as their core component [].
In an S6 layer, parameters are functions of the input, endowing the SSM with the content awareness ability.
Mamba has showcased exceptional performance across a variety of applications, such as language modeling [, ], image processing [, ], video analysis [, ], medical imaging [, ], tabular data analysis [], reinforcement learning [], point-cloud analysis [], graph processing [], and $N$-dimensional sequence modeling [].
It has been argued that Mamba’s significant success across various domains stems from its ability to compute complicated representations for different data types
[, , , ].

Despite their remarkable empirical success, a 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 a selective state-space model is a Mamba block, which is built on top of the $\operatorname{S6}$ layers, along with linear layers, convolutions, and other token-wise operators. The $\operatorname{S6}$ layer, which serves as the core component of the Mamba block, plays a pivotal role in the success of these models.

Formally, an $\operatorname{S6}$ layer is defined as a function that maps a sequence of tokens $\mathbf{x}=(x_{1},\ldots,x_{L})\in\mathbb{R}^{D\times L}$ to another sequence of tokens $\mathbf{y}=(y_{1},\ldots,y_{L})\in\mathbb{R}^{D\times L}$ (with the same number of channels).
In this context, each vector $x_{l}$ (or $y_{l}$) represents a token, while the whole sequence $\mathbf{x}=(x_{1},\ldots,x_{L})$ (or $\mathbf{y}=(y_{1},\ldots,y_{L})$) is called a prompt.
For each $d=1,\ldots,D$, the $d$-th channel output $\hat{y}_{d}=(y_{d1},\ldots,y_{dL})\in\mathbb{R}^{L}$ is determined recurrently from the corresponding $d$-th channel input $\hat{x}_{d}=(x_{d1},\ldots,x_{dL})\in\mathbb{R}^{L}$ via a sequence of hidden states $h_{d1},\ldots,h_{dL}\in\mathbb{R}^{N}$ as follows:

$$ $\left\{\begin{aligned} &h_{dl}=\overline{A}_{dl}\cdot h_{d,l-1}+\overline{B}_{ dl}\cdot x_{dl},\quad h_{d0}=0,\\ &y_{dl}=C_{l}\cdot h_{dl},\end{aligned}\right.$ (1) $$

for $l=1,\ldots,L$.
Here, the input matrices $\overline{A}_{dl}$, output matrices $\overline{B}_{dl}$, and state matrices $C_{l}$ are input-dependent and time-varying, determined by

$$ $\overline{A}_{dl}=e^{\Delta_{d}(x_{l})A_{d}},\quad\overline{B}_{dl}=\Delta_{d} (x_{l})S_{B}\cdot x_{l},\quad C_{l}=(S_{C}\cdot x_{l})^{\top},$ (2) $$

where $\Delta_{d}\,:\,\mathbb{R}^{D}\to\mathbb{R}$ is the step size function at the $d$-th channel and is defined by

$$ $\displaystyle\Delta_{d}(u)=\text{softplus}(S_{\Delta,d}u)=\ln\left(1+e^{S_{ \Delta,d}\cdot u}\right),\qquad u\in\mathbb{R}^{D}.$ (3) $$

The hidden matrices $A_{d}\in\mathbb{R}^{N\times N}$ and the step size vectors $S_{\Delta,d}\in\mathbb{R}^{1\times D}$ for $d=1,\ldots,D$, as well as the input and output matrices $S_{B},S_{C}\in\mathbb{R}^{N\times D}$, are learnable from input data. The matrices $A_{d}$ are diagonal matrices with negative eigenvalues.
In recent developments of selective state space models, the matrices $A_{d}$ are chosen to be of the scalar form $A_{d}=-a_{d}I_{N}$ for some positive number $a_{d}$ (see []), which does not compromise model performance. We will also use this scalar form for $A_{d}$ in our paper.
In our context, the matrix $S_{C}^{\top}S_{B}\in\mathbb{R}^{D\times D}$, 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 to describe the dynamical properties of tokens in a 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 one-dimensional 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.
For the convergent scenario, we empirically verify that this situation negatively impacts the model’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 as it negatively impacts performance, and (ii) we 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 in improving the model’s accuracy and convergence speed compared to the baseline Mamba on the large-scale ImageNet classification task []

##### Organization of the paper

After surveying related works in Section 2, we introduce the dynamical system that governs the continuous-time 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, we propose two refinements for the model in Section 5: excluding unfavorable scenarios and reordering tokens before training. We conclude the paper with a discussion of conclusions and limitations in Section 6.

##### Notations.

We use small bold letters (e.g., $\mathbf{x},\mathbf{y},\mathbf{u}$) to denote sequences of tokens. Tokens can be either vectors or scalars, depending on the context, and are denoted by small normal letters (e.g., $x_{i},y_{j},u_{k}$), where the indices emphasize their position in the sequence. We use 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 perspective. For instance, the authors in [] demonstrated that the S6 layer can be interpreted as a variant of softmax-free attention with linear complexity, subsequently proving that Mamba is more expressive than transformers. Conversely, [] 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, [] utilized circuit complexity theory to establish that both SSM variants and transformers share the same expressive power, as they belong to the complexity class $TC^{0}$, which can be decided by polynomial-sized Boolean circuits. In contrast, the authors in [] provided both theoretical and empirical evidence that transformers surpass state space models in their copying capabilities.
All of these works aimed to provide a theoretical framework to understand the expressivity and generalization of Mamba.

In this paper, we examine the dynamical properties of tokens in a pretrained Mamba model. Following the methodology outlined in [, ], we define an idealized model of Mamba by interpreting the discrete layer indices as continuous time variables. This idealized perspective was used in ResNets
[, ] and neural ODEs [, , , , , ].
Our model focuses on the $\operatorname{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 properties 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 a Deep Neural Network

To ensure a clear presentation of our results, we draw upon the literature concerning the dynamical systems that govern the continuous-time limit of deep neural networks (DNNs).
Generally speaking, data in a DNN is processed sequentially, layer by layer, resulting in a discrete-time dynamical system [].
A notable example is residual neural networks (ResNets) and their continuous-time counterparts known as neural ODEs [, ].
Each layer of ResNet is a residual block, transforms an input vector $x\in\mathbb{R}^{D}$ to an output vector $z\in\mathbb{R}^{D}$ via
a two-layer feed-forward neural network
$x\mapsto y(x,\theta),$ parametrized by $\theta$,
and a skip-connection as:

$$ $z=x+y(x,\theta).$ $$

The continuous-time counterparts of ResNet is conceptualized as a flow map that inputs a vector $x(0)\in\mathbb{R}^{D}$ and outputs another vector $x(T)\in\mathbb{R}^{D}$, which is processed via the corresponding dynamical system

$$ $x^{\prime}(t)=y(x(t),\theta(t)),\quad t\in(0,T).$ $$

There exists a body of work investigating the interpolation, approximation, and controllability properties of these DNN architectures [, , , , , ].

### 3.2 Continuous-time Limit of Mamba

Unlike ResNets and neural ODEs, Mamba represents a function on a sequence of $D$-dimensional tokens rather than solely on an individual input, as discussed in Section 1.1.
In order to derive the dynamical system governing the continuous-time counterpart of Mamba,
we eliminate the hidden states $h_{dl}$ in system (1) to obtain the following form (see []):

$$ $y_{dl}=\sum\limits_{j=1}^{l}P_{dlj}x_{dj},$ (4) $$

or equivalently, $\hat{y}_{d}=P_{d}\hat{x}_{d}$,
where

$$ $\displaystyle P_{d}=\begin{bmatrix}P_{d11}&0&0&\ldots&0\\ P_{d21}&P_{d22}&0&\ldots&0\\ P_{d31}&P_{d32}&P_{d33}&\ldots&0\\ \vdots&\vdots&\vdots&\ddots&\vdots\\ P_{dL1}&P_{dL2}&P_{dL3}&\ldots&P_{dLL}\end{bmatrix},$ (5) $$

and

$$ $P_{dlj}=\left\{\begin{aligned} &x_{l}^{\top}\left(S_{C}^{\top}S_{B}\right)x_{l }\cdot\Delta_{d}(x_{l}),&\text{if }l=j,\\ &x_{l}^{\top}\left(S_{C}^{\top}S_{B}\right)x_{j}\cdot\Delta_{d}(x_{j})\cdot \exp\left(-a_{d}\sum\limits_{k=j+1}^{l}\Delta_{d}(x_{k})\right),&\text{if }l>j .\end{aligned}\right.$ (6) $$

In the above formulation, the upper triangular matrices $P_{d}$ ($d=1,\ldots,D$) represent the hidden attention scores between the $d$-th channel input and output. Following the setting of [], we call $P=[P_{1},\ldots,P_{D}]$ the hidden attention tensor of the $\operatorname{S6}$ layer.

Following the common approaches used in studying the continuous-time counterparts of ResNets [, ] and Transformers [, ], we consider the layer index of a deep selective state space model as a time variable and interpret the selective state space model as the discrete-time version of a parametrized dynamical system of the form

$$ $\displaystyle\frac{d}{dt}x_{dl}(t)=\sum\limits_{j=1}^{l}P_{dlj}(t)x_{dj}(t), \quad t\in[0,+\infty),$ (7) $\displaystyle x_{dl}(0)=x_{dl0}\in\mathbb{R},$ $$

for $l=1,\ldots,L$ and $d=1,\ldots,D$, where

$$ $P_{dlj}(t)=\left\{\begin{aligned} &x_{l}(t)^{\top}\left(S_{C}^{\top}S_{B} \right)x_{l}(t)\cdot\Delta_{d}(x_{l}(t)),&\text{if }l=j,\\ &x_{l}(t)^{\top}\left(S_{C}^{\top}S_{B}\right)x_{j}(t)\cdot\Delta_{d}(x_{j}(t) )\cdot\exp\left(-a_{d}\sum\limits_{k=j+1}^{l}\Delta_{d}(x_{k}(t))\right),& \text{if }l>j,\end{aligned}\right.$ (8) $$

and $\Delta_{d}(u)=\text{softplus}(S_{\Delta,d}u)$.

Similar to [, ], we have focused exclusively on $\operatorname{S6}$ layer, which is the key component of Mamba, and the skip-connection in the above dynamical system.
We have ignored other token-wise operators such as layer normalization, convolution, and linear functions.
In addition, we assume the parameters $A_{d},S_{\Delta,d},S_{B},S_{C}$ are time-independent.
These assumptions are primarily motivated by mathematical convenience.
Nevertheless, such a weight-sharing mechanism is used in practice to reduce the number of trainable parameters, as seen, for example, in ALBERT [], a transformer-based 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, we study the asymptotic behavior of the solution $\mathbf{x}(t)$ of the dynamical system (7), as well as its corresponding attention score $P(t)=(P_{dlj}(t))_{dlj}$ defined in equation (8).
This information encodes the dynamical properties of tokens and hidden attention scores in deep selective state space models.

### 4.1 Convergence and Divergence Scenarios of Mamba’s Dynamics

In the standard form of $\operatorname{S6}$, each $d$-th channel output is determined from the corresponding $d$-th channel input using time-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.
In this case, the input sequence $\mathbf{x}=(x_{1},\ldots,x_{L})\in\mathbb{R}^{L}$ is a list of $L$ real numbers.
We drop the index $d$ from Eq. (7) and rewrite it as

$$ $\displaystyle\frac{d}{dt}x_{l}(t)=\sum\limits_{j=1}^{l}P_{lj}(t)x_{j}(t),\quad t \in[0,+\infty),$ (9) $\displaystyle x_{l}(0)=x_{l0}\in\mathbb{R},$ $$

for $l=1,\ldots,L$, where

$$ $\displaystyle\Delta(u)=\text{softplus}(S_{\Delta}u)=\ln\left(1+e^{S_{\Delta}u} \right),\qquad u\in\mathbb{R},$ (10) $$

and

$$ $P_{lj}(t)=\left\{\begin{aligned} &\left(S_{C}^{\top}S_{B}\right)x_{l}(t)^{2} \cdot\Delta(x_{l}(t)),&\text{if }l=j,\\ &\left(S_{C}^{\top}S_{B}\right)x_{l}(t)x_{j}(t)\cdot\Delta(x_{j}(t))\cdot\exp \left(-a\sum\limits_{k=j+1}^{l}\Delta(x_{k}(t))\right),&\text{if }l>j.\end{ aligned}\right.$ (11) $$

In this case, $S_{B},S_{C}\in\mathbb{R}^{N\times 1}$, thus $S_{C}^{\top}S_{B}\in\mathbb{R}$, and the scalar values $S_{\Delta}$ and $a\in\mathbb{R}$, with $a>0$, are learnable from input data.

To avoid triviality, we will always assume that the initial datum $x_{l}(0)=x_{l0}$ are nonzero for all $l$.
The main goal of this section is to investigate the asymptotic behavior of the solution $\mathbf{x}(t)=(x_{1}(t),\ldots,x_{L}(t))$ of the dynamic (9)
and the corresponding hidden attention matrix

$$ $\displaystyle\mathbf{P}(t)=\begin{bmatrix}P_{11}(t)&0&\ldots&0\\ P_{21}(t)&P_{22}(t)&\ldots&0\\ \ldots&\ldots&\ldots&\ldots\\ P_{L1}(t)&P_{L2}(t)&\ldots&P_{LL}(t)\end{bmatrix},$ $$

when $t$ increases.
The following three scenarios will be considered separately:

- 1.
Convergence scenario: $S_{C}^{\top}S_{B}<0$;
- 2.
Slow-divergence scenario: $S_{C}^{\top}S_{B}>0$ and $S_{\Delta}x_{l0}<0$ for all $l$;
- 3.
Fast-divergence scenario: $S_{C}^{\top}S_{B}>0$ and $S_{\Delta}x_{l0}>0$ for some $l$.

We will see that, in the first scenario, all tokens and hidden attention scores quickly converge to zero as $t\to\infty$, at rates of $O(1/\sqrt{t})$ and $O(1/t)$, respectively (Theorem 4.1).
In the second scenario, all tokens diverge to infinity slowly at a logarithmic rate, while the hidden attention scores still approach zero as $t\to\infty$. We refer to this as the slow-divergence scenario (Theorem 4.3).
In the final scenario, one token diverges to 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 1 summarizes the dynamical properties of tokens in this work.

**Table 1: Summary of the dynamical properties of this work.**
| Parameters | Scenario | Impact on model | Reference |  |
| --- | --- | --- | --- | --- |
| $S_{C}^{\top}S_{B}<0$ | Convergence | Negative | Theorem 4.1 |  |
| $S_{C}^{\top}S_{B}>0$ | $\forall l,\,S_{\Delta}x_{l0}<0$ | Slow-divergence | Positive | Theorem 4.3 |
|  | $\exists l,\,S_{\Delta}x_{l0}>0$ | Fast-divergence |  | Theorem 4.4 |

###### Remark 1 (Convergence vs. divergence scenarios and their impact) .

We collectively refer to the slow-divergence and fast-divergence scenarios as the divergence scenarios. By utilizing model parameters to test the signature of $S_{C}^{\top}S_{B}$, 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 $x_{l,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 convergence 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 a consequence, the contribution of tokens to 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) .

The extension of the above results to higher dimensions is nontrivial. In higher dimensions, $S_{C}^{\top}S_{B}$ is a matrix of size $D\times D$ and may have complex eigenvalues. However, our experiments suggest that the above results still hold if we replace the condition $S_{C}^{\top}S_{B}<0$ (respectively, $>0$) in the one-dimensional case with the condition that $S_{C}^{\top}S_{B}$ is negative (respectively, positive) definite in the higher-dimensional case (see Appendix B). We leave the study of the higher-dimensional case for future work.

In the subsequent subsections, we provide the statements of the main theorems.
The proofs of these theorems can be found in Appendix A.

###### Remark 1 (Convergence vs. divergence scenarios and their impact) .

###### Remark 2 (Unequal contribution of tokens during training) .

###### Remark 3 (Higher dimension) .

### 4.2 Convergence scenario: μ = S C ⊤ ⁢ S B < 0 𝜇 superscript subscript 𝑆 𝐶 top subscript 𝑆 𝐵 0 \mu=S_{C}^{\top}S_{B}<0 italic_μ = italic_S start_POSTSUBSCRIPT italic_C end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_S start_POSTSUBSCRIPT italic_B end_POSTSUBSCRIPT < 0

In the following theorem, we prove that all tokens converge to zero at a rate of $O(1/\sqrt{t})$ as $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 $\mu=S_{C}^{\top}S_{B}<0$.
Let $\mathbf{x}(t)=(x_{1}(t),\ldots,x_{L}(t))$ be the unique solution of the dynamic (9).

- 1.
For each $l=1,\ldots,L$, if $x_{l0}>0$ (respectively, $x_{l0}<0$), then $x_{l}(t)$ is positive and monotonically decreasing (respectively, negative and monotonically increasing) on $[0,+\infty)$.
In addition, $x_{l}(t)=O(\frac{1}{\sqrt{t}})$.
- 2.
$P_{lj}(t)=O(\frac{1}{t})$ for all $l,j=1,\ldots,L$.

The proof of this theorem can be found in Appendix A.1.
To intuitively see why all tokens converge in this scenario, we can rewrite the dynamic (9) as

$$ $\frac{d}{dt}x_{l}(t)=\mu x_{l}(t)^{3}\Delta(x_{l}(t))+g_{l}(t)x_{l}(t),$ (12) $$

where

$$ $g_{l}(t)=\left\{\begin{aligned} &0,&\text{if }l=1,\\ &\sum\limits_{j=1}^{l-1}\mu x_{j}(t)^{2}\Delta(x_{j}(t))\exp\left(-a\sum \limits_{k=j+1}^{l}\Delta(x_{k}(t))\right),&\text{if }l>1.\end{aligned}\right.$ $$

Since $\mu<0$, $g_{l}(t)$ is a nonpositive function.
Let us consider the case when $x_{l0}>0$. Then, $x_{l}(t)$ must remain positive for every $t>0$; otherwise, $x_{l}(t)$ would reach zero at some point, and consequently $x_{l}(t)$ would be identically zero for all $t$, 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 $x_{l}(t)$ is decreasing.
Figure 1 illustrates the graph of tokens and the heat map of the corresponding hidden attention matrices in this scenario.

###### Remark 4 (Negative impact on model’s performance) .

In this scenario, both the tokens and the 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.

Figure: Figure 1: 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 $\mu=-1.58$, $S_{\Delta}=-0.17$, $a=-1.08$, as well as the initial data $\mathbf{x}(0)=(-1.79,-0.34,0.46,-1.25,0.83,-0.83,1.81,-1.16,0.13,-0.19)$. In this case, all tokens and hidden attention scores tend to zero as $t$ approaches infinity.
Refer to caption: x1.png

###### Theorem 4.1 (Convergence scenario) .

###### Remark 4 (Negative impact on model’s performance) .

### 4.3 Slow-divergence Scenario: μ = S C ⊤ ⁢ S B > 0 𝜇 superscript subscript 𝑆 𝐶 top subscript 𝑆 𝐵 0 \mu=S_{C}^{\top}S_{B}>0 italic_μ = italic_S start_POSTSUBSCRIPT italic_C end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_S start_POSTSUBSCRIPT italic_B end_POSTSUBSCRIPT > 0 and S Δ ⁢ x l ⁢ 0 < 0 subscript 𝑆 Δ subscript 𝑥 𝑙 0 0 S_{\Delta}x_{l0}<0 italic_S start_POSTSUBSCRIPT roman_Δ end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_l 0 end_POSTSUBSCRIPT < 0 for all l 𝑙 l italic_l

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 $\mu=S_{C}^{\top}S_{B}>0$ and $S_{\Delta}x_{l0}<0$ for all $l$.
Let $x(t)=(x_{1}(t),\ldots,x_{L}(t))$ be the unique solution of the dynamic (15).
For each $l$, if $x_{l0}>0$ (respectively, $x_{l0}<0$), then $x_{l}(t)$ is positive and monotonically increasing (respectively, negative and monotonically decreasing) on $[0,+\infty)$ with

$$ $\lim_{t\to+\infty}x_{l}(t)=+\infty\quad(\text{respectively, }\lim_{t\to+\infty }x_{l}(t)=-\infty).$ $$

Lemma 4.2 shows that all tokens in the considered case will approach infinity as $t$ approaches infinity. However, it does not help to estimate the rate of divergence or the asymptotic behavior of the hidden attention score $P_{lj}(t)$. In Theorem 4.3 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.

###### Theorem 4.3 (Slow divergence scenario and divergence rate) .

Assume that $\mu=S_{C}^{\top}S_{B}>0$ and

$$ $S_{\Delta}x_{L0}\leq\ldots\leq S_{\Delta}x_{10}\leq-2.12,$ (13) $$

(see Theorem A.7 for a detailed upper bound).
Let $x(t)=(x_{1}(t),\ldots,x_{L}(t))$ be the unique solution of the dynamic (9).

- 1.
For each $l$, if $x_{l0}>0$ (respectively, $x_{l0}<0$), then $x_{l}(t)$ is a positive increasing (respectively, negative decreasing) function on $[0,+\infty)$. In addition, $x_{l}(t)=O((\ln t)^{l})$.
- 2.
$\lim_{t\to+\infty}P_{lj}(t)=0$ for all $l\geq j$.

###### Remark 5 (Unequally tokens’ contribution) .

Theorem 4.3 shows that, in the one-dimensional case within this divergence scenario, tokens with higher absolute values will diverge to 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 in 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 $S_{\Delta}<0$ and thus $x_{l0}>0$. Then $x_{l}(t)$ must always be positive as $t$ increases; otherwise, $x_{l}(t)$ would reach zero at some point, and consequently $x_{l}(t)$ would be identically zero for all $t$, 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 not too large. This implies that $x_{l}(t)$ is decreasing. However, since $S_{\Delta}<0$, the term $\mu x_{l}(t)^{3}\Delta(x_{l}(t))$ converges to zero very quickly, which keeps the right-hand side always not too large. Therefore, $x_{l}(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.

Figure: Figure 2: 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 $\mu=1.79$, $S_{\Delta}=-0.71$, $a=-1.80$, as well as the initial data $\mathbf{x}(0)=(1.55,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 still tend to zero as $t$ approaches infinity. In addition, the hidden attention scores from the second columns tend to zero must fater than those in the first column.
Refer to caption: x3.png

###### Lemma 4.2 (Slow divergence scenario) .

###### Theorem 4.3 (Slow divergence scenario and divergence rate) .

###### Remark 5 (Unequally tokens’ contribution) .

### 4.4 Fast-Divergence Scenario: μ = S C ⊤ ⁢ S B > 0 𝜇 superscript subscript 𝑆 𝐶 top subscript 𝑆 𝐵 0 \mu=S_{C}^{\top}S_{B}>0 italic_μ = italic_S start_POSTSUBSCRIPT italic_C end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_S start_POSTSUBSCRIPT italic_B end_POSTSUBSCRIPT > 0 and S Δ ⁢ x l ⁢ 0 > 0 subscript 𝑆 Δ subscript 𝑥 𝑙 0 0 S_{\Delta}x_{l0}>0 italic_S start_POSTSUBSCRIPT roman_Δ end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_l 0 end_POSTSUBSCRIPT > 0 for some l 𝑙 l italic_l

In the fast divergence scenario, we have the following theorem.

###### Theorem 4.4 (Fast divergence scenario) .

Assume that $\mu=S_{C}^{\top}S_{B}>0$.
Let $x(t)=(x_{1}(t),\ldots,x_{L}(t))$ be the unique solution of the dynamic (15).
If there exists $l=1,\ldots,L$ such that $S_{\Delta}x_{l0}>0$, then $x_{l}(t)$ goes to infinity at finite time.

In addition, assume that $x_{l_{0}}(t)$ tends to infinity first as $t\to T^{-}$ for some $T>0$ while $x_{l}(t)$ is defined on $[0,T]$ for all $l\neq l_{0}$. Then for every $L\geq l\geq j\geq 1$, $\lim_{t\to T^{-}}P_{lj}(t)=+\infty$ if and only if $l=l_{0}$ or $j=l_{0}$.

###### Remark 6 .

The proof of this theorem can be found in Appendix A.3. Intuitively speaking, in this scenario, if $1\leq l\leq L$ is the index of the first token such that $S_{\Delta}x_{l0}>0$, then all tokens with indices $j\geq l$ will diverge to infinity in finite time. Figure 3 illustrates the dynamical properties of tokens and the hidden attention scores in this scenario.

Figure: 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 $\mu=0.76$, $S_{\Delta}=0.59$, $a=1.66$, as well as the initial data $\mathbf{x}(0)=(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 to infinity very quickly at finite time.
Refer to caption: x5.png

###### Theorem 4.4 (Fast divergence scenario) .

###### Remark 6 .

## 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 the input-output matrix $S_{C}^{\top}S_{B}$ on the performance of auto-regressive language modeling tasks and (ii) propose a reordering strategy and illustrate its effectiveness in vision tasks. We provide the details on datasets, models, and training procedures in Appendix C.

### 5.1 Effects of negative eigenvalues of input-output matrix

We use Mamba [] as the baseline and conduct evaluations on the WikiText103 language modeling task []. 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 case. The scenario where $S_{C}^{\top}S_{B}$ encompasses complex eigenvalues is deferred to future work. We report the perplexity (PPL) for each scenario. Lower PPL values indicate better model performance.

Table: Table 2: Test perplexity on WikiText103. A model with fewer negative eigenvalues in $S_{C}^{\top}S_{B}$ achieves lower Perplexity.

| Scenario | Perplexity ($\downarrow$) |
| --- | --- |
| Negative | 17.26 |
| Mix | 16.84 |
| Positive | 16.71 |

Figure: Figure 4: Test perplexity on WikiText103 during training procedure. The positive case consistently demonstrates superior performance compared to the other two scenarios.

Table 5.1 and Figure 4 illustrate the impact of different eigenvalue configurations on the model’s performance. Specifically, our findings demonstrate a correlation between the prevalence of positive eigenvalues in the input-output matrix and model performance. Conversely, the presence of negative eigenvalues is linked to a decrease in performance, suggesting that removing negative eigenvalues could further optimize the model’s efficiency.

### 5.2 Re-ordering input tokens

We now emphasize the significance of token ordering in vision tasks, where tokens naturally lack an inherent structure due to each token representing a 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 $S_{\Delta}x_{l0}$, which is a vector of dimension $D$ in a scenario with $D$ channels. To determine the optimal order of tokens $x_{l0}$, we define an importance score as follows:

$$ $\displaystyle s_{l}=\langle K,S_{\Delta}x_{l0}\rangle$ (14) $$

where $K\in\mathbb{R}^{D}$ is a learnable parameter. Prior to passing the tokens through the $\operatorname{S6}$ layer, they are reordered based on their importance scores in ascending order. To ensure that the reordering process is differentiable, we utilize $\operatorname{SoftSort}$ [], which allows gradients to propagate through the sorting operation during backpropagation. Specifically, $\operatorname{SoftSort}$ is employed to generate a permutation matrix, which is then applied to sort the $s_{l}$ array and rearrange the tokens $x_{l0}$ accordingly.

To assess the effectiveness of our proposed reordering technique, we benchmark our method on the image classification task using the ImageNet-1K dataset []. We employ MambaVision [] as the baseline and compare the top-1 and top-5 accuracy metrics.

**Table 3: Top-1 and Top-5 accuracy ($\uparrow$) on the large-scale ImageNet-1K image classification task. Our token reordering method leads to a slight improvement in MambaVision’s performance.**
| Model | Top-1 (%) | Top-5 (%) |
| --- | --- | --- |
| MambaVision-T (baseline) | 81.90 | 95.86 |
| MambaVision-T + Token reordering | 82.02 | 95.87 |

Figure: Figure 5: Top-1 (left) and Top-5 (right) accuracy ($\uparrow$) on ImageNet-1K during the training process. Our token reordering method boosts the accuracy of the MambaVision baseline and achieves faster convergence compared to the baseline
Refer to caption: x8.png

The results presented in Figure 5 and Table 3 demonstrate that our reordering method
speeds up the convergence of the training loss and improves the performance of the MambaVision baseline. These findings confirm the effectiveness of our technique in vision tasks, highlighting its potential for broader tasks.

## 6 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 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 suggest that these refinements do help improve model performance. Regarding limitations, our theoretical analysis of the dynamical properties of tokens is restricted to tokens with one feature channel. The extension of the results to higher dimensions is nontrivial. We leave the study of higher-dimensional cases for future work.

#### Ethics Statement

Considering the scope of our research, we do not anticipate any negative societal or ethical consequences emerging from this work.

#### Reproducibility Statement

Our implementation details are provided in Section  5 and Appendix  C. Source code is provided in the Supplementary material.

## Appendix A Dynamical Properties Tokens in Mamba

We will consider $\operatorname{S6}$ dynamic in the single channel dimension case, i.e. $D=1$, in our setting to avoid overly complicated technicalities.
In this case, the input sequence $\mathbf{x}=(x_{1},\ldots,x_{L})\in\mathbb{R}^{L}$ is a list of $L$ real numbers.
The dynamic of $\mathbf{x}$ is then governed by the dynamical system:

$$ $\displaystyle\frac{d}{dt}x_{l}(t)=\sum\limits_{j=1}^{l}P_{lj}(t)x_{j}(t),\quad t \in[0,+\infty),$ (15) $\displaystyle x_{l}(0)=x_{l0}\in\mathbb{R},$ $$

for $l=1,\ldots,L$, where

$$ $\displaystyle\Delta(u)=\text{softplus}(S_{\Delta}u)=\ln\left(1+e^{S_{\Delta}u} \right),\qquad u\in\mathbb{R},$ (16) $$

and

$$ $P_{lj}(t)=\left\{\begin{aligned} &\left(S_{C}^{\top}S_{B}\right)x_{l}(t)^{2} \cdot\Delta(x_{l}(t)),&\text{if }l=j,\\ &\left(S_{C}^{\top}S_{B}\right)x_{l}(t)x_{j}(t)\cdot\Delta(x_{j}(t))\cdot\exp \left(-a\sum\limits_{k=j+1}^{l}\Delta(x_{k}(t))\right),&\text{if }l>j.\end{ aligned}\right.$ (17) $$

In this case, $S_{B},S_{C}\in\mathbb{R}^{N\times 1}$, thus $S_{C}^{\top}S_{B}\in\mathbb{R}$, and the scalar values $S_{\Delta}$ and $a\in\mathbb{R}$, with $a>0$, are learnable from input data.

To avoid triviality, we will always assume that the initial datum $x_{l}(0)=x_{l0}$ are nonzero for all $l$.
The main goal of this section is to investigate the asymptotic behavior of the solution $\mathbf{x}(t)=(x_{1}(t),\ldots,x_{L}(t))$ of the dynamic (15)
and the corresponding hidden attention matrix

$$ $\displaystyle\mathbf{P}(t)=\begin{bmatrix}P_{11}(t)&0&\ldots&0\\ P_{21}(t)&P_{22}(t)&\ldots&0\\ \ldots&\ldots&\ldots&\ldots\\ P_{L1}(t)&P_{L2}(t)&\ldots&P_{LL}(t)\end{bmatrix},$ $$

when $t$ increases.
The following three scenarios will be considered separately in the subsequent subsections:

- 1.
Convergence scenario: $\mu=S_{C}^{\top}S_{B}<0$;
- 2.
Slow-divergence scenario: $\mu=S_{C}^{\top}S_{B}>0$ and $S_{\Delta}x_{l0}<0$ for all $l$; and
- 3.
Fast-divergence scenario: $\mu=S_{C}^{\top}S_{B}>0$ and $S_{\Delta}x_{l0}>0$ for some $l$.

### A.1 Convergence scenario: μ = S C ⊤ ⁢ S B < 0 𝜇 superscript subscript 𝑆 𝐶 top subscript 𝑆 𝐵 0 \mu=S_{C}^{\top}S_{B}<0 italic_μ = italic_S start_POSTSUBSCRIPT italic_C end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_S start_POSTSUBSCRIPT italic_B end_POSTSUBSCRIPT < 0

We will start with the following auxiliary lemmas.

###### Lemma A.1 .

Let $t_{0}\in\mathbb{R}$ and let $f(u,t)$ be a differentiable continuous function on $\mathbb{R}\times[t_{0},+\infty)$.
Assume that $u(t)$ is the unique solution of the initial value problem

$$ $\displaystyle\frac{d}{dt}u(t)=f(u(t),t)\,u(t),\quad t\in[t_{0},+\infty),$ $\displaystyle u(t_{0})=u_{0}\in\mathbb{R}.$ $$

If $u_{0}>0$ (respectively, $u_{0}<0$), then $u(t)$ is a positive (respectively, negative) function on its maximal interval of the existence.
In addition,

- 1.
if $f$ is a positive function, then $u(t)$ is monotonically increasing (respectively, monotonically decreasing),
- 2.
if $f$ is a negative function, then $u(t)$ is monotonically decreasing (respectively, monotonically increasing).

###### Proof.

We assume that $u_{0}>0$. The case when $u_{0}<0$ is obtained by replacing $u(t)$ by $-u(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 to the fact that $u(0)=u_{0}>0$.
Therefore, $u(t)$ is positive on its maximal interval of the existence.

Next, assume that $f$ is positive on $\mathbb{R}\times[t_{0},+\infty)$.
Then we always have $u^{\prime}(t)=f(u(t),t)u(t)>0$ for all $t\in[t_{0},+\infty)$.
This shows that $u(t)$ is increasing on its maximal interval of the existence.
The case when $f$ is negative is similar.
∎

###### Lemma A.2 .

Let $t_{0}\in\mathbb{R}$. If a function $u(t)$ is bounded on $[t_{0},+\infty)$, then there exists positive numbers $c_{1},c_{2}$ such that

$$ $c_{1}\leq\Delta(u(t))\leq c_{2},$ $$

for all $t\in[t_{0},+\infty)$.

###### Proof.

Recall that $\Delta(u(t))=\ln(1+e^{S_{\Delta}u(t)})$.
Since $u(t)$ is bounded, there are $a_{1},a_{2}\in\mathbb{R}$ such that $a_{1}\leq S_{\Delta}u(t)\leq a_{2}$ for all $t\in[t_{0},+\infty)$.
Then we can choose $c_{i}=\ln(1+e^{a_{i}})$ for $i=1,2$.
∎

###### Theorem A.3 (Convergence scenario) .

Assume that $\mu=S_{C}^{\top}S_{B}<0$.
Let $\mathbf{x}(t)=(x_{1}(t),\ldots,x_{L}(t))$ be the unique solution of the dynamic (15).

- 1.
For each $l=1,\ldots,L$, if $x_{l0}>0$ (respectively, $x_{l0}<0$), then $x_{l}(t)$ is positive and monotonically decreasing (respectively, negative and monotonically increasing) on $[0,+\infty)$.
In addition, $x_{l}(t)=O(\frac{1}{\sqrt{t}})$.
- 2.
$P_{lj}(t)=O(\frac{1}{t})$ for all $l,j=1,\ldots,L$.

###### Proof.

Assume without loss of generality that $x_{l0}>0$.
The case when $x_{l0}<0$ is proved similarly by replacing $x_{l}(t)$ by $-x_{l}(t)$.
We rewrite the dynamic (15) as

$$ $\frac{d}{dt}x_{l}(t)=\mu x_{l}(t)^{3}\Delta(x_{l}(t))+g_{l}(t)x_{l}(t),$ (18) $$

where

$$ $g_{l}(t)=\left\{\begin{aligned} &0,&\text{if }l=1,\\ &\sum\limits_{j=1}^{l-1}\mu x_{j}(t)^{2}\Delta(x_{j}(t))\exp\left(-a\sum \limits_{k=j+1}^{l}\Delta(x_{k}(t))\right),&\text{if }l>1.\end{aligned}\right.$ $$

Since $\mu<0$, $g_{l}(t)$ is a nonpositive function.
According to Lemma A.1, $x_{l}(t)$ is positive and monotonically decreasing on its right maximal interval of the existence $[0,T)$ with $T\in(0,+\infty]$.
Thus, $x_{l}(t)$ is bounded and $T=+\infty$.
As a consequence, it follows from Lemma A.2 that there is a positive constant $c>0$ such that $\Delta(x_{l}(t))\geq c$ for all $t\in[0,+\infty)$.
Therefore, from Eq. (18), we have

$$ $\frac{d}{dt}x_{l}(t)\leq\mu cx_{l}(t)^{3},$ $$

which imply that $x_{l}(t)\leq(-2\mu ct+x_{l0}^{-2})^{\frac{1}{2}}$ for all $t\in[0,+\infty)$.
Hence, $x_{l}(t)=O(\frac{1}{\sqrt{t}})$ as expected.

For the hidden attention score $P_{lj}(t)$ with $L\geq l\geq j\geq 1$, we have

$$ $\displaystyle|P_{lj}(t)|=\mu|x_{l}(t)x_{j}(t)|\Delta(x_{j}(t))\exp\left(-a\sum \limits_{k=j+1}^{l}\Delta(x_{k}(t))\right)\leq\mu c|x_{l}(t)x_{j}(t)|=O\left( \frac{1}{t}\right),$ $$

as expected.
∎

###### Lemma A.1 .

###### Proof.

###### Lemma A.2 .

###### Proof.

###### Theorem A.3 (Convergence scenario) .

###### Proof.

### A.2 Slow-divergence scenario: μ = S C ⊤ ⁢ S B > 0 𝜇 superscript subscript 𝑆 𝐶 top subscript 𝑆 𝐵 0 \mu=S_{C}^{\top}S_{B}>0 italic_μ = italic_S start_POSTSUBSCRIPT italic_C end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_S start_POSTSUBSCRIPT italic_B end_POSTSUBSCRIPT > 0 and S Δ ⁢ x l ⁢ 0 < 0 subscript 𝑆 Δ subscript 𝑥 𝑙 0 0 S_{\Delta}x_{l0}<0 italic_S start_POSTSUBSCRIPT roman_Δ end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_l 0 end_POSTSUBSCRIPT < 0 for all l 𝑙 l italic_l

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 $\mu=S_{C}^{\top}S_{B}>0$ and $S_{\Delta}x_{l0}<0$ for all $l$.
Let $x(t)=(x_{1}(t),\ldots,x_{L}(t))$ be the unique solution of the dynamic (15).
For each $l$, if $x_{l0}>0$ (respectively, $x_{l0}<0$), then $x_{l}(t)$ is positive and monotonically increasing (respectively, negative and monotonically decreasing) on $[0,+\infty)$ with

$$ $\lim_{t\to+\infty}x_{l}(t)=+\infty\quad(\text{respectively, }\lim_{t\to+\infty }x_{l}(t)=-\infty).$ $$

###### Proof.

Without loss of generality, we assume that $S_{\Delta}<0$. The case when $S_{\Delta}>0$ is then obtained by replacing $x(t)$ by $-x(t)$.
Similar to the proof of Theorem A.3, we rewrite dynamic (15) as

$$ $\frac{d}{dt}x_{l}(t)=\mu x_{l}(t)^{3}\Delta(x_{l}(t))+g_{l}(t)x_{l}(t),$ (19) $$

where

$$ $g_{l}(t)=\left\{\begin{aligned} &0,&\text{if }l=1,\\ &\sum\limits_{j=1}^{l-1}\mu x_{j}(t)^{2}\Delta(x_{j}(t))\exp\left(-a\sum \limits_{k=j+1}^{l}\Delta(x_{k}(t))\right),&\text{if }l>1.\end{aligned}\right.$ (20) $$

Since $\mu>0$, $g_{l}(t)$ is a nonnegative function on $[0,+\infty)$.

According to Lemma A.1, $x_{l}(t)$ must be a positive increasing function on its maximal interval of the existence $[0,T)$ for some $T\in(0,+\infty]$.

First, we claim that $x_{l}(t)$ is unbounded on $[0,T)$.
Indeed, if this is not the case, then $x_{l}(t)$ is bounded and $T=+\infty$.
According to Lemma A.2, there exists a positive constant $c$ such that $\Delta(x_{l}(t))\geq c$ for all $t\in[0,+\infty)$.
Therefore, Eq. (19) yields

$$ $\frac{d}{dt}x_{l}(t)\geq\mu cx_{l}(t)^{3},\quad t\in[0,+\infty),$ $$

which implies that $x_{l}(t)\geq(-2\mu ct+x_{l0}^{-2})^{-\frac{1}{2}}$, which is unbounded at $t=\frac{1}{c\mu x_{l0}^{2}}>0$, a contradiction. The claim is then proved, and $\lim_{t\to T^{-}}x_{l}(t)=+\infty$.

It remains to prove that $T=+\infty$.
Observe that, since $S_{\Delta}<0$, we have

$$ $\lim_{r\to+\infty}r^{3}\Delta(r)=\lim_{r\to+\infty}\frac{r^{3}}{e^{-S_{\Delta} r}}\frac{\ln(1+e^{S_{\Delta}r})}{e^{S_{\Delta}r}}=0.$ $$

Furthermore, since $x_{l}(t)$ is positive and monotonically increasing to infinity at infinity, we also have

$$ $\lim_{t\to+\infty}x_{l}(t)^{3}\Delta(x_{l}(t))=0.$ $$

In particular, in case $l=1$, we have $\frac{d}{dt}x_{1}(t)\leq 1$, which implies that $x_{1}(t)\leq t$, for all $t$ large enough.
Therefore, $T=+\infty$ in this case.

Assume that $l>1$ and we already have $T=+\infty$ for all cases $j<l$.
Then it folows from Eq (20) that

$$ $\displaystyle g_{l}(t)\leq\sum\limits_{j=1}^{l-1}\mu x_{j}(t)^{2}\Delta(x_{j}( t))\leq\sum\limits_{j=1}^{l-1}\frac{\frac{d}{dt}x_{j}(t)}{x_{j}(t)}=\frac{d}{ dt}\sum\limits_{j=1}^{l-1}\ln(x_{j}(t)).$ (21) $$

Therefore, we can bound $\frac{d}{dt}x_{l}(t)$ from Eq. (19) as

$$ $\frac{d}{dt}x_{l}(t)\leq 1+\left(\frac{d}{dt}\sum\limits_{j=1}^{l-1}\ln(x_{j}( t))\right)x_{l}(t),$ $$

and this inequality hold for all $t\geq t_{1}$ for some $t_{1}$ large enough.
As a consequence, we have

$$ $\displaystyle x_{l}(t)$ $\displaystyle\leq\exp\left(\sum\limits_{j=1}^{l-1}\ln(x_{j}(t))-\ln(x_{j}(t_{1 }))\right)\left(\int_{t_{1}}^{t}\exp\left(-\sum\limits_{j=1}^{l-1}\ln(x_{j}(s) )-\ln(x_{j}(t_{1}))\right)ds+x_{l}(t_{1})\right).$ $$

The right hand side is a function defined for all $t\geq t_{1}$.
Thus $T=+\infty$.
∎

###### Remark 7 .

Lemma A.4 shows that all tokens in the considered case will approach infinity as $t$ approaches infinity. However, it does not help to estimate the rate of divergence or the asymptotic behavior of the hidden attention score $P_{lj}(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.5 .

Assume that $\mu=S_{C}^{\top}S_{B}>0$ and $S_{\Delta}x_{10}<0$.
Let $x_{1}(t)$ be the unique solution of the dynamic (15).
Then $x_{1}(t)$ is defined on $[0,+\infty)$ and $x_{1}(t)=O(\ln(t))$.

###### Proof.

Without loss of generality, let us assume that $S_{\Delta}<0$, thus $x_{10}>0$.
The case when $S_{\Delta}>0$ can be obtained by replacing $x_{1}(t)$ by $-x_{1}(t)$.
Then it follows from Lemma A.4 that $x_{1}(t)$ is a positive and monotonically increasing function on $[0,+\infty)$.
Let us fix a small enough $\epsilon>0$ such that $S_{\Delta}+\epsilon<0$.
Then, there exist $c_{1},t_{1}>0$ such that:

$$ $\displaystyle\mu x_{1}(t)^{3}\Delta(x_{1}(t))=e^{(S_{\Delta}+\epsilon)x_{1}(t) }\cdot\frac{\mu x_{1}(t)^{3}}{e^{\epsilon x_{1}(t)}}\cdot\frac{\ln(1+e^{S_{ \Delta}x_{1}(t)})}{e^{S_{\Delta}x_{1}(t)}}\leq c_{1}e^{(S_{\Delta}+\epsilon)x_ {1}(t)},$ (22) $$

for all $t\in[t_{1},+\infty)$.
Therefore, we have

$$ $\frac{d}{dt}x_{1}(t)\leq c_{1}e^{(S_{\Delta}+\epsilon)x_{1}(t)},\quad t\geq t_ {1},$ $$

which yields

$$ $x_{1}(t)\leq-\frac{2}{S_{\Delta}}\ln\left(-\frac{1}{2}S_{\Delta}c_{1}(t-t_{1}) +e^{-\frac{1}{2}S_{\Delta}x_{1}(t_{1})}\right),\quad\text{ for all }t\geq t_{1}.$ $$

Since $x_{1}(t)$ is positive, we finally obtain $x_{1}(t)=O(\ln(t))$ as expected.
∎

To determine the divergence rate of the other tokens, we will need the following lemma.

###### Lemma A.6 .

Let $r_{0}(\approx 2.1160)$ be the unique positive root of the function $h(r)=2\ln(1+e^{-r})-\frac{re^{-r}}{1+e^{-r}}$.
Then for $\lambda>0$, the function $f(r)=r^{2}\ln\left(1+e^{-\lambda r}\right)$ is positive and monotonically decreasing on $[\frac{r_{0}}{\lambda},+\infty)$.

###### Proof.

The lemma follows from the negativity of the function $h(r)$ on $[r_{0},+\infty)$ and the equality $\frac{d}{dt}f(r)=rh(\lambda r)$.
∎

###### Theorem A.7 (Slow divergence scenario and divergent rate) .

Assume that $\mu=S_{C}^{\top}S_{B}>0$ and

$$ $S_{\Delta}x_{L0}\leq\ldots\leq S_{\Delta}x_{10}\leq-r_{0},$ (23) $$

where $r_{0}(\approx 2.1160)$ is defined in Lemma A.6.
Let $x(t)=(x_{1}(t),\ldots,x_{L}(t))$ be the unique solution of the dynamic (15).

- 1.
For each $l$, if $x_{l0}>0$ (respectively, $x_{l0}<0$), then $x_{l}(t)$ is a positive increasing (respectively, negative decreasing) function on $[0,+\infty)$. In addition, $x_{l}(t)=O((\ln t)^{l})$.
- 2.
$\lim_{t\to+\infty}P_{lj}(t)=0$ for all $l\geq j$.

###### Proof.

Without loss of generality, we will assume that $S_{\Delta}=-\lambda$ for some $\lambda>0$.
The case when $S_{\Delta}>0$ is then obtained by changing all $x_{l}$ to $-x_{l}$.
Then the condition (23) becomes

$$ $x_{L0}\geq\ldots\geq x_{10}\geq\frac{r_{0}}{\lambda}.$ (24) $$

According to Proposition A.4, $x_{l}(t)$ is a positive and monotonically increasing function on $[0,+\infty)$ with $\lim_{t\to+\infty}x_{l}(t)=+\infty$.

Before estimating the asymptotic behavior of $x_{l}(t)$ and $P_{lj}(t)$, we first claim that that $x_{l}(t)\geq x_{l-1}(t)$ for all $t\in[0,+\infty)$ and $l\geq 2$.
We will prove this claim by induction on $l$.
Indeed, in case $l=2$, we observe that

$$ $\displaystyle\frac{d}{dt}x_{2}(t)=\mu x_{2}(t)^{3}\Delta(x_{2}(t))+g_{2}(t)x_{ 2}(t)\geq\mu x_{2}(t)^{3}\Delta(x_{2}(t)).$ $$

Since $x_{20}\geq x_{10}$, it follows that $x_{2}(t)\geq x_{1}(t)$ for all $t\in[0,+\infty)$.
Let us assume that $l>2$ and $x_{l-1}(t)\geq\ldots\geq x_{1}(t)$ for all $t\in[0,+\infty)$.
It is noted that $x_{1}(t)\geq x_{10}\geq\frac{r}{\lambda}$.
According to Lemma A.6, we have

$$ $\displaystyle x_{j}(t)^{2}\Delta(x_{j}(t))\geq x_{j-1}(t)^{2}\Delta(x_{j-1}(t) ),\quad j=2,\ldots,l-1.$ (25) $$

Recall that

$$ $\displaystyle\frac{d}{dt}x_{l}(t)$ $\displaystyle=\mu x_{l}(t)^{3}\Delta(x_{l}(t))+g_{l}(t)x_{l}(t)\noindent$ (26) $\displaystyle=\mu x_{l}(t)^{3}\Delta(x_{l}(t))+\mu x_{l}(t)\sum\limits_{j=1}^{ l-1}x_{j}(t)^{2}\Delta(x_{j}(t))\exp\left(-a\sum\limits_{k=j+1}^{l-1}\Delta(x_ {k}(t))\right).$ (27) $$

By removing the term with the index $j=1$ inside the sum and using inequality (25), we obtain

$$ $\displaystyle\frac{d}{dt}x_{l}(t)$ $\displaystyle\geq\mu x_{l}(t)^{3}\Delta(x_{l}(t))+\mu x_{l}(t)\sum\limits_{j=2 }^{l-1}x_{j-1}(t)^{2}\Delta(x_{j-1}(t))\exp\left(-a\sum\limits_{k=j+1}^{l-1} \Delta(x_{k}(t))\right)$ $\displaystyle=\mu x_{l}(t)^{3}\Delta(x_{l}(t))+\mu x_{l}(t)\sum\limits_{j=1}^{ l-2}x_{j}(t)^{2}\Delta(x_{j}(t))\exp\left(-a\sum\limits_{k=j+2}^{l-1}\Delta(x_ {k}(t))\right).$ $$

Since $\Delta(x_{k}(t))\leq\Delta(x_{k-1}(t))$ for all $k=2,\ldots,l-1$, we can proceed the last expression as

$$ $\displaystyle\frac{d}{dt}x_{l}(t)$ $\displaystyle\geq\mu x_{l}(t)^{3}\Delta(x_{l}(t))+\mu x_{l}(t)\sum\limits_{j=1 }^{l-2}x_{j}(t)^{2}\Delta(x_{j}(t))\exp\left(-a\sum\limits_{k=j+1}^{l-2}\Delta (x_{k}(t))-a\Delta(x_{l}(t))\right).$ $$

Therefore, if we set

$$ $\displaystyle f(u)=\mu u^{3}\Delta(u)+\mu u\sum\limits_{j=1}^{l-2}x_{j}(t)^{2} \Delta(x_{j}(t))\exp\left(-a\sum\limits_{k=j+1}^{l-2}\Delta(x_{k}(t))-a\Delta( u)\right),$ $$

then $x_{l-1}(t)$ is a solution of the differential equation $\frac{d}{dt}u(t)=f(u(t))$, while $\frac{d}{dt}x_{l}(t)\geq f(x_{l}(t)$ for all $t\in[0,+\infty)$.
It follows from the initial conditions $x_{l0}\geq x_{l-1,0}$ that $x_{l}(t)\geq x_{l-1}(t)$ for all $t\in[0,+\infty)$.
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 $l=1$, it is already known from Lemma A.5 that $x_{1}(t)=O(\ln(t))$.
In case $l\geq 2$, according to equation (27), we have
$\displaystyle\frac{\frac{d}{dt}x_{l}(t)}{x_{l}(t)}=\mu x_{l}(t)^{2}\Delta(x_{l
}(t))+\sum\limits_{j=1}^{l-1}\mu x_{j}(t)^{2}\Delta(x_{j}(t))\exp\left(-a\sum
\limits_{k=j+1}^{l-1}\Delta(x_{k}(t))\right)$
(28)
For each $j=1,\ldots,l$, since $x_{j}(t)\geq x_{1}(t)$ and the function $r\mapsto r^{2}\Delta(r)$ is monotonically decreasing over $[\frac{r_{0}}{\lambda},+\infty)$, we have $x_{j}(t)^{2}\Delta(x_{j}(t))\leq x_{1}(t)^{2}\Delta(x_{1}(t))$ for all $t>0$ large enough.
In addition, we also have and $\exp\left(-a\sum\limits_{k=j+1}^{l-1}\Delta(x_{k}(t))\right)\leq 1$ for all $t\in[0,+\infty)$.
Therefore, we can upper bound the right hand side of equation (28) as
$\displaystyle\frac{\frac{d}{dt}x_{l}(t)}{x_{l}(t)}\leq l\mu x_{1}(t)^{2}\Delta
(x_{1}(t))=l\frac{\frac{d}{dt}x_{1}(t)}{x_{1}(t)}.$
(29)
Thus, $\ln x_{l}(t)\leq l\ln x_{1}(t)$, which implies that
$x_{l}(t)\leq x_{1}(t)^{l}=O((\ln t)^{l}).$
- 2.
For the hidden attention score, it follows from its definition in equation (17) that
$P_{lj}(t)\leq\mu x_{l}(t)x_{j}(t)\Delta(x_{j}(t))$
for all $t\geq 0$.
Since $S_{\Delta}<0$, the function $r\mapsto r\Delta(r)=r\ln(1+e^{S_{\Delta}r})$ is monotonically decreasing from a large enough $t$.
Therefore, it follows from the claim above that
$P_{lj}(t)\leq\mu x_{l}(t)^{2}\Delta(x_{l}(t))$
for all $t$ large enough.
Hence, $\lim_{t\to+\infty}P_{lj}(t)=\lim_{r\to+\infty}\mu r^{2}\Delta(r)=0$.

The theorem is then proved.
∎

###### Lemma A.4 (Slow-divergence scenario) .

###### Proof.

###### Remark 7 .

###### Lemma A.5 .

###### Proof.

###### Lemma A.6 .

###### Proof.

###### Theorem A.7 (Slow divergence scenario and divergent rate) .

###### Proof.

### A.3 Fast-divergence scenario: μ = S C ⊤ ⁢ S B > 0 𝜇 superscript subscript 𝑆 𝐶 top subscript 𝑆 𝐵 0 \mu=S_{C}^{\top}S_{B}>0 italic_μ = italic_S start_POSTSUBSCRIPT italic_C end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_S start_POSTSUBSCRIPT italic_B end_POSTSUBSCRIPT > 0 and S Δ ⁢ x l ⁢ 0 > 0 subscript 𝑆 Δ subscript 𝑥 𝑙 0 0 S_{\Delta}x_{l0}>0 italic_S start_POSTSUBSCRIPT roman_Δ end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_l 0 end_POSTSUBSCRIPT > 0 for some l 𝑙 l italic_l

###### Theorem A.8 (Fast-divergence scenario) .

Assume that $\mu=S_{C}^{\top}S_{B}>0$.
Let $x(t)=(x_{1}(t),\ldots,x_{L}(t))$ be the unique solution of the dynamic (15).
If there exists $l=1,\ldots,L$ such that $S_{\Delta}x_{l0}>0$, then $x_{l}(t)$ goes to infinity at finite time.

In addition, assume that $x_{l_{0}}(t)$ tends to infinity first as $t\to T^{-}$ for some $T>0$ while $x_{l}(t)$ is defined on $[0,T]$ for all $l\neq l_{0}$. Then for every $L\geq l\geq j\geq 1$, $\lim_{t\to T^{-}}P_{lj}(t)=+\infty$ if and only if $l=l_{0}$ or $j=l_{0}$.

###### Proof.

Without loss of generality, we assume that $S_{\Delta}>0$. The case when $S_{\Delta}<0$ is then obtained by replacing $x(t)$ by $-x(t)$.
Then $x_{l0}>0$.
Recall that

$$ $\frac{d}{dt}x_{l}(t)=\mu x_{l}(t)^{3}\Delta(x_{l}(t))+g_{l}(t)x_{l}(t),$ (30) $$

where $g_{l}(t)$ defined in equation (20).
Since $\mu>0$, $g_{l}(t)$ is a nonnegative function on $[0,+\infty)$.
It follows from Lemma A.1 that $x_{l}(t)$ is a positive increasing function on its right maximal interval of the existence $[0,T)$ for some $T\in(0,+\infty]$.
By equation (30), we have

$$ $\frac{d}{dt}x_{l}(t)\geq\mu x_{l}(t)^{3}\Delta(x_{l}(t)).$ $$

Furthermore, since $S_{\Delta}>0$, the function $\Delta(r)=\ln(1+e^{S_{\Delta}r})$ is monotonically increasing on $[0,+\infty)$.
Therefore, $\Delta(x_{l}(t))\geq\Delta(x_{l}(0))=\ln(1+e^{S_{\Delta}x_{l0}})$ for all $t\in[0,+\infty)$.
As a consequence, we can proceed the above inequality further as

$$ $\frac{d}{dt}x_{l}(t)\geq\mu\ln(1+e^{S_{\Delta}x_{l0}})x_{l}(t)^{3}.$ $$

Thus,

$$ $x_{l}(t)\geq\left(2\mu\ln(1+e^{S_{\Delta}x_{l0}})t-x_{l0}^{-2}\right)^{-\frac{ 1}{2}},$ $$

which goes to $+\infty$ at finite time.

Next, let us assume that $x_{l0}(t)$ goes to infinity at finite time faster than the other tokens, i.e, there exists $T>0$ such that $\lim_{t\to T^{-}}x_{l_{0}}(t)=+\infty$, while $x_{l}(t)$ defined over $[0,T]$ for all $l\neq l_{0}$.
Then $\lim_{t\to T^{-}}x_{l_{0}}^{2}(t)\Delta(x_{l}(t))=+\infty$.
Then for every $L\geq l\geq j\geq 1$, $\lim_{t\to T^{-}}P_{lj}(t)=\infty$ if and only if $l=l_{0}$ or $j=l_{0}$.
∎

###### Theorem A.8 (Fast-divergence scenario) .

###### Proof.

## Appendix B Higher Dimension

Figure 6 shows the tokens’ trajectory (top) and the heat maps of the hidden attention matrices (middle and bottom) when the input-output matrix $\mu=S_{C}^{\top}S_{B}$ is negative definite (left) or positive definite (right).
Below is the parameters and initial values used on each case:

- Case 1 (left).
–
the input-output matrix
$\mu=\begin{bmatrix}-0.0152&-0.0181\\
-0.0181&-0.0232\end{bmatrix},$
–
the step sizes
$S_{\Delta}=\begin{bmatrix}0.0006&-0.0012\\
-0.0003&-0.0005\end{bmatrix},$
–
the hidden matrix
$A_{1}=A_{2}=1.0906I_{2}$,
–
the initial values
$x_{1}=(-0.0446,0.4832)$,
$x_{2}=(-0.4253,-0.2427)$,
$x_{3}=(0.5347,0.0489)$,
$x_{4}=(-0.7332,0.1256)$,
$x_{5}=(0.3604,0.3871)$,
$x_{6}=(0.5285,0.4282)$,
$x_{7}=(0.7644,0.3652)$,
$x_{8}=(-0.3043,1.1064)$,
$x_{9}=(-0.3831,-0.7250)$,
$x_{10}=(-1.0115,-0.4852)$.
- Case 2 (right).
–
the input-output matrix
$\mu=\begin{bmatrix}0.0509&-0.0036\\
-0.0036&0.0062\end{bmatrix},$
–
the step sizes
$S_{\Delta}=\begin{bmatrix}-0.9034&-0.9288\\
-1.1058&-1.6070\end{bmatrix},$
–
the hidden matrix
$A_{1}=A_{2}=1.49165I_{2}$,
–
the initial values
$x_{1}=(2.7727,2.8381)$,
$x_{2}=(2.3896,2.8116)$,
$x_{3}=(0.6050,0.8420)$,
$x_{4}=(0.9214,0.9635)$,
$x_{5}=(0.8279,0.7399)$,
$x_{6}=(2.0382,0.7402)$,
$x_{7}=(2.6117,1.3210)$,
$x_{8}=(0.6065,2.8700)$,
$x_{9}=(1.7087,5.1197)$,
$x_{10}=(1.7190,3.2396)$.

Figure 6 suggests that our results in the one -dimensional case may be extended to higher-dimensional cases with definite input-output matrices.
However, this extension is nontrivial.
We leave the study of the higher-dimensional case for future work.

Figure: Figure 6: The trajectories of tokens (top) and the heat map of the hidden attention matrices (middle and bottom) for Case 1 (left) and Case 2 (right).
Refer to caption: x10.png

## Appendix C Experiment Details

In this section, we outline the experimental setup in detail. All experiments are conducted using a server with four A100 GPUs.

### C.1 Auto-regressive language modeling

##### Dataset

We utilize the WikiText-103 [], created from Wikipedia articles. It includes a training set of approximately 28,000 articles, amounting to 103 million words in total. Each article is split into sections of about 3,600 words. The validation and test sets contain 60 articles each, with word counts of 218,000 and 246,000 respectively, combining for a total of roughly 268,000 words.

##### Parameterization for three scenarios:

We outline the parameterization used to ensure the $\operatorname{S6}$ layer operates within the three scenarios described in Section 5.1. For the first two scenarios, we leverage the $LDL^{T}$ decomposition [] to ensure the positive or negative definite property of input-output matrix. Specifically, we parameterize $S_{C}=L^{T}$ and $S_{B}=DL^{T}$, where $D$ is a diagonal matrix and $L$ is a lower triangular matrix with diagonal values fixed at 1. In the first scenario, the diagonal values of $D$ are constrained to be positive and are parameterized as $D=\operatorname{diag}(\operatorname{softplus}(d^{\prime}))$. In the second scenario, this constraint is reversed by parameterizing $D=-\operatorname{diag}(\operatorname{softplus}(d^{\prime}))$, ensuring that the diagonal values are negative.

In the third scenario, both positive and negative eigenvalues are required, meaning that only complex eigenvalues are excluded. To achieve this, we set $S_{C}=S_{B}$, which guarantees that $S_{C}^{\top}S_{B}$ is a symmetric matrix, thus ensuring that the eigenvalues are real.

##### Training details

We train a 130M parameters version of Mamba model [] on WikiText103 for 25 epochs with GPT-2 tokenizer [] and AdamW optimizer []. Other hyperparameters are set identically across three scenarios and listed in Table 4.

**Table 4: Hyperparameters configuration for Language Modeling task on WikiText-103.**
| Sequence Length | 1024 |
| --- | --- |
| Peak Learning Rate | 0.0015 |
| Momentum | $\beta_{1}=0.9;\beta_{2}=0.999$ |
| Weight Decay | 0.25 |
| Batch size | 128 |
| Learning rate warmup | Linear |
| Learning rate scheduler | Cosine decay |
| Dropout | 0.25 |

### C.2 Image classification

##### Dataset

The ImageNet-1K dataset [] 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 a broad variety of objects, animals, scenes, and more, allowing the model to learn a rich and diverse range of visual concepts.

##### Training details

We train the tiny version (31.8M parameters) of MambaVision model [] with and without our token reordering technique. We follow the training procedure in [, , ] and set the hyperparameters for two cases identical, which is listed in Table 5.

**Table 5: Hyperparameters configuration for Image Classification task on ImageNet-1K.**
| Optimizer | AdamW |
| --- | --- |
| Peak Learning Rate | 0.0008 |
| Momentum | $\beta_{1}=0.9;\beta_{2}=0.999$ |
| Weight Decay | 0.05 |
| Batch size | 512 |
| Learning rate warmup | Linear |
| Learning rate scheduler | Cosine decay |
| Dropout | 0.0 |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Could you extend the theoretical analysis to higher-dimensional cases?
2. 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?
3. Do your findings and proposed refinements generalize to other selective state space models beyond Mamba?
soundness: 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.
- Theoretical Rigor: Provides rigorous mathematical proofs for the asymptotic behavior of tokens, offering clear criteria based on model parameters.
- Practical Implications: The findings lead to actionable refinements that improve model performance, validated through experiments on language modeling and image classification tasks.
- Clarity: The paper is well-structured, with logical progression and clear explanations that make complex mathematical concepts accessible.
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.
Based on these findings, the authors propose two refinements to improve model performance:
* Excluding the Convergent Scenario: By ensuring certain conditions on model parameters, they prevent the tokens from converging to zero.
* Reordering Tokens: They reorder tokens based on their importance scores to account for the unequal contribution during training.

Experimental results on language modeling and image classification tasks validate these refinements, showing improved accuracy and convergence speed.
weaknesses: - Simplifying Assumptions: Using time-independent parameters and excluding other layers for mathematical convenience may oversimplify Mamba's actual behavior in practical settings.
- Experimental Scope: Though supportive, experimental validation could be more extensive. Including additional datasets and comparisons with other models would strengthen the claims.
- 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.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: **Summary of Actionable Improvements**:
- Expand the theoretical analysis to higher dimensions or at least provide empirical evidence for higher-dimensional token dynamics.
- Broaden the experimental validation to include diverse tasks beyond vision (e.g., language modeling and reinforcement learning).
- Discuss the computational costs and trade-offs of the proposed refinements in real-world applications.
- Increase the novelty by exploring more complex dynamical behaviors or offering a deeper comparison with existing models like transformers.
- Provide a more comprehensive experimental comparison to other recent SSMs and transformers, including computational efficiency.
soundness: 3
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.
- **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.
- **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.
- **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.
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.
The 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.
This study offers new theoretical insights into token dynamics and suggests practical adjustments to improve the Mamba model's performance.
weaknesses: - **Theoretical Scope is Too Narrow**: 
While 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.
- **Insufficient Diversity of Experiments**: 
The 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.
- **Shallow Investigation of Practical Impact**: 
The 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.
- **Limited Novelty of Some Contributions**: 
While 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.
- **Lack of Comprehensive Comparison to Prior Work**: 
The 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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: My main concerns I would like addressed are listed in the weaknesses section. Here are some additional comments or questions:

- 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".

- 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?

- Line 075: The input, output, and state matrices are mislabeled


Update: I have increased my score to reflect my updated view of the paper after the rebuttal period.
soundness: 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
- 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
- The paper is well-structured and Figures 1-3 are useful.
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.
weaknesses: - The effects of the assumptions made for the analysis are not thoroughly explored
    - 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
    - The analysis assumes weights are shared across layers
    - The analysis assumes the one-dimensional case and excludes the possibility of complex eigenvalues in the input/output projection matrix 
    - 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.

- 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. 

- 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).

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: See above.
soundness: 4
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.

(Score increased from 5 to 8 after the discussion period)
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.
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:
- 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.
- 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)?
- 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.

While 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$.

Finally, 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.

Meta-review: This paper presents a theoretical analysis of token dynamics in deep selective ssms, particularly focusing on mamba. The key findings are: 

- In the one-dimensional case, tokens either converges to zero or diverge to infinity, with specific criteria determining which scenario occurs; 
- In the divergent case, tokens diverge at different rates based on their position, leading to unequal contributions during training. 

authors propose two improvements: excluding scenarios where tokens converge to zero and reordering tokens based on importance scores. The authors validate these improvements through experiments on language modeling and image classification tasks.

Strengths:

- good theoretical analysis for ssms
- good empirical validation across multiple tasks and architectures

Weaknesses:

- Initial analysis relied on simplifying assumptions (1D case, time-independent parameters)
- Limited discussion of computational complexity and practical implementation challenges
- Initial experimental scope was narrow, focusing mainly on vision tasks

Authors thoroughly addressed reviewer concerns during rebuttal with additional experiments. I vote for acceptance.

Decision: Accept (Spotlight)