Demystifying the Token Dynamics of Deep Selective State Space Models

[Reply 2/2]

[W3]. 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.

Answer: We would like to clarify that our hyperparameter settings follow those established in previous works, specifically Mamba [Gu2023] for language modeling task on WikiText103 and MambaVision [Hatamizadeh2024] for image classification on ImageNet. Additionally, we have provided our code in the supplementary materials, ensuring that our results can be fully reproduced.

References [Gu2023] Gu, A., & Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. [Hatamizadeh2024] Ali Hatamizadeh and Jan Kautz. Mambavision: A hybrid mamba-transformer vision backbone. arXiv preprint arXiv:2407.08083, 2024.

---

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

[Reply 1/2]

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

---

[W1]-[Q1]. 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.

Answer: As outlined in [Geshkovski2024] (page 6, paragraph 2), tokens in a purely self-attention setting typically exhibit exponential divergence to infinity. However, our theoretical analysis demonstrates that the token dynamics in 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 to infinity within a finite time or diverge more gradually to infinity over an infinite time horizon, with the divergence following a logarithmic rate. We have added this discussion to Remark 4 in the revised version of the paper.

References [Geshkovski2024] Geshkovski, B., Letrouit, C., Polyanskiy, Y., & Rigollet, P. (2024). The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36.

[Q2]. 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)?

Answer: Thanks for your suggestion. We plan to extend our analysis to the bidirectional case in the next few days.

[Q3]. 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.

Answer: Theorem 1 in [Gu2023] establishes that the classical gating mechanism of RNNs represents a specific instance of the selection mechanism in Mamba. Consequently, our theoretical analysis of Mamba encompasses those RNNs as special cases. A full analysis for a more general gated RNNs such as LSTMs or GRUs would require nontrivial works. We leave this interesting works for future study.

References [Gu2023] Gu, A., & Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752.

[W2]. 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$.

Answer: Regarding conclusions in high-dimentional case: We conjecture that our theorems still hold in the high-dimensional case. In the high-dimensional case, the input-output matrix $S_C^{\top}S_B$ is a square matrix of size $D \times D$ and it may have complex eigenvalues. In this case, we consider the symmetric matrix $\mu:=\frac{1}{2}(S_C^{\top}S_B)^{\top} + \frac{1}{2}S_C^{\top}S_B$, which is the symmetric part of the input-output projection matrix $S_C^{\top}S_B$. This term is motivated by the observation that: the term $x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l$ in the considered dynamical system is a scalar and it can be rewritten as $x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l = \frac{1}{2}[x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l]^{\top} + \frac{1}{2}[x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l] = x_l^{\top} \cdot \mu \cdot x_l.$ Since $\mu$ is a symmetric matrix, it always has only real eigenvalues. We particularly conjecture that, in the high-dimensional case, all tokens will converge to zero if $\mu$ has only negative eigenvalues, while all tokens will diverge to infinity if $\mu$ has at least one positive eigenvalue.

To illustrate this conjecture, we run simulations in two-dimensional cases and draw the trajectories of tokens in Figures 6 in Appendix B of the revised version. We observe the following:

• In Figure 6A and 6B, the matrix $\mu$ has at least one positive eigenvalue, and the divergent behavior of the tokens is also observed to persist.
• In Figure 6C, the matrix $\mu$ has only negative eigenvalues, and the convergent behavior of the tokens is observed to persist.

Regarding simulations with Mamba setting involving practical operators (such as layer normalization, linear, or convolution): We plan to provide a few simulations in the next few days to observe the dynamics of tokens.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback. Beside of our responses above, we would like to provide missing answer on the possibility of extending the analysis to the bidirectional case below.

[Q2] Regarding the possibility of extending the analysis to the bidirectional case: Our theoretical results on the dynamic behaviors of tokens may no longer hold for bidirectional SSMs (see [Zhu2024] for a definition of bidirectional SSMs). In a bidirectional SSM, input tokens are reordered in two distinct sequences: one in the forward order and the other in the backward order. These two sequences, containing the same tokens but in different arrangements, are processed separately by the SSM before being combined into a single sequence using a summation operator. According to our theoretical results, tokens in earlier positions tend to diverge at a slower rate under the divergence scenario. Consequently, the summation operator mixes token dynamics that diverge at different rates, leading to more complex overall behaviors. We believe that a comprehensive investigation of token dynamics in this bidirectional setting is an interesting direction for future research.

[Zhu2024] Zhu, Lianghui, et al. "Vision mamba: Efficient visual representation learning with bidirectional state space model." arXiv preprint arXiv:2401.09417 (2024).

We would be happy to do any follow-up discussion or address any additional comments.

[Reply 1/2]

[W1]. The effects of the assumptions made for the analysis are not thoroughly explored.

Answer: Our primary objective was to examine the most basic form of Mamba that allows for thorough mathematical analysis. We carried out numerical simulations to further validate our tokens dynamic theorems in high-dimensional cases with additional practicle settings in Appendix B of our revised manuscript (with updates highlighted in blue). Interestingly, we found that in general our theorems on tokens dynamic behaviors remained consistent in the high-dimensional case.

[W2-W6-Q1]. 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. [...] 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?

Answer: Thank you for your suggestion. We plan to extend our experiment in the coming days by including an ablation study on the importance of each layer/component.

[W3]. The analysis assumes weights are shared across layers.

Answer: In case when the weights are time-dependent, we believe that if the entries of the input-output projection matrix $S_C(t)^\top S_B(t)$, the step size matrix $S_\Delta(t)$, and the scalar $a(t)$ are bounded both above and below, the conclusions of Theorems 4.1, 4.3, and 4.4 remain valid. However, this would require modifications to the proofs to account for these assumptions.

[W4]. The analysis assumes the one-dimensional case and excludes the possibility of complex eigenvalues in the input/output projection matrix.

Answer: We conjecture that our theorems still hold in the high-dimensional case. In the high-dimensional case, the input-output matrix $S_C^{\top}S_B$ is a square matrix of size $D \times D$ and it may have complex eigenvalues. In this case, we consider the symmetric matrix $\mu:=\frac{1}{2}(S_C^{\top}S_B)^{\top} + \frac{1}{2}S_C^{\top}S_B$, which is the symmetric part of the input-output projection matrix $S_C^{\top}S_B$. This term is motivated by the observation that: the term $x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l$ in the considered dynamical system is a scalar and it can be rewritten as $x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l = \frac{1}{2}[x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l]^{\top} + \frac{1}{2}[x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l] = x_l^{\top} \cdot \mu \cdot x_l.$ Since $\mu$ is a symmetric matrix, it always has only real eigenvalues. We particularly conjecture that, in the high-dimensional case, all tokens will converge to zero if $\mu$ has only negative eigenvalues, while all tokens will diverge to infinity if $\mu$ has at least one positive eigenvalue.

To illustrate this conjecture, we run simulations in two-dimensional cases and draw the trajectories of tokens in Figures 6 in Appendix B of the revised version. We observe the following:

• In Figure 6A and 6B, the matrix $\mu$ has at least one positive eigenvalue, and the divergent behavior of the tokens is also observed to persist.
• In Figure 6C, the matrix $\mu$ has only negative eigenvalues, and the convergent behavior of the tokens is observed to persist.

[W5]. All of these assumptions [one-dim, time-independent weights, no additional operators beside S6 layer, no complex eigenvalues] 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.

Answer: We have discussed all of these assumptions in the answers of [W1-W4] above, except an ablation study on the effects of different components/layers which we plan to provide in the next few days. We have also added explorations of these assumptions and additional empirical results in Appendix B (for the high-dimensional case) and Appendix D (for additional experiments) of the revised version of the paper.

In addition to these discussions, we would like to add one more comment on possibly more complex behaviors in the interaction between tokens. After conducting several simulations with random choices for the inputs in the two- and three-dimensional cases, we observe that the dynamic of tokens in purely Mamba setting typically has no periodic, chaotic, nor bifurcation behavior, i.e., if the tokens converge to zero (diverge to infinity), they will converge to zero (diverge to infinity) directly after a few time steps. In particular, the dynamic behavior is quite stable, even when a little noise is added to the initial conditions. We conjecture that the system is stable under a generic assumption of the parameters. A full proof for this conjecture will require adaptations. We leave it for future study.

[Reply 2/2]

[W7]. More care should be taken to distinguish between the depth direction of the network and the sequence direction of the network. [...]

Answer: Similar to other continuous-time limit of a deep neural network, the time direction in the dynamical system in our paper is exactly the depth direction. Following reviewer's suggestion, we have added an explanation about this time direction below the considered dynamical system in page 4 of the revised version.

[Q2]. 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?

Answer: To assess the robustness of our findings across three scenarios of input-output matrices on the WikiText103 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 in Table 1 suggest that our findings are robust to different random seeds. We also report these results and findings in Table 9 in Appendix D.3 of our revision.

Table 1: Test perplexity on WikiText103 ($\downarrow$). Mean and standard deviation are computed over three runs with different random seeds.

[Q3]. Line 075: The input, output, and state matrices are mislabeled

Answer: Thank you very much for pointing out this typo. We have updated the sentence. Indeed, the input and output matrices are $S_B$ and $S_C$, respectively.

---

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback. We would like to provide missing answers on an ablation study to analyze the effect of practical components on the reordering refinement below.

[W2-W6-Q1] Regarding the effects of practical components of Mamba in the token reordering technique: We conducted an ablation study on WikiText103 and examined the following configurations, both with and without the reordering technique:

(i) the full Mamba model,

(ii) the Mamba with weight sharing across layers,

(iii) the Mamba without layer normalization,

(iv) the Mamba without the short convolution layer.

To ensure fairness, all configurations are standardized to approximately 129M parameters. The results which are shown in Table 4 below highlighting the contribution of each component to the model's performance under our reordering technique.

Table 4: Ablation study on language modeling task with WikiText103.

We have added these results to Appendix D.5 of the revised paper.

We would be happy to do any follow-up discussion or address any additional comments.

[Reply 3/3]

[W4&Q4]. 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.

Answer: Regarding more complex behaviors in the interaction between tokens: After conducting several simulations with random choices for the inputs in the two- and three-dimensional cases, we observe that the dynamic of tokens in purely Mamba setting typically has no periodic, chaotic, nor bifurcation behavior, i.e., if the tokens converge to zero (diverge to infinity), they will converge to zero (diverge to infinity) directly after a few time steps. In addition, the dynamic behavior is quite stable, even when a little noise is added to the initial conditions. We conjecture that the system is stable under a generic assumption of the parameters. A full proof for this conjecture will require adaptations. We leave it for future study.

Regarding a comparison of token dynamics with transformer: As outlined in [Geshkovski2024] (page 6, paragraph 2), tokens in a purely self-attention setting typically exhibit exponential divergence to infinity. However, our theoretical analysis demonstrates that the token dynamics in 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 to infinity within a finite time or diverge more gradually to infinity over an infinite time horizon, with the divergence following a logarithmic rate.

References [Geshkovski2024] Geshkovski, B., Letrouit, C., Polyanskiy, Y., & Rigollet, P. (2024). The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36.

[W5&Q5]. 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. </font>

Answer: Thank you for your suggestion. We plan to extend our experiment in the coming days by including a more detailed comparison with other SSMs and transformers.

---

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

[Reply 2/3]

[W3&Q3]. 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.

Answer: Regarding the analysis of computational overhead, we first present the pseudo-code for calculating token importance scores and forwarding through a single SSM layer in Algorithm 1 in Appendix D.2 of our revision and also include this pseudo-code below.

Algorithm 1: Forwarding through a SSM layer with token reordering

Input: Input sequence $\mathbf{x} \in \mathbb{R}^{D \times L}$, hidden matrix $A \in \mathbb{R}^{N \times N}$, input matrix $S_B \in \mathbb{R}^{N \times D}$, output matrix $S_C \in \mathbb{R}^{N \times D}$, step size matrix $S_\Delta \in \mathbb{R}^{D \times D}$, learnable vector $K \in \mathbb{R}^{D \times 1}$, hyperparameters $\tau$ and $p$ of $\operatorname{Softsort}$

Output:

• Output sequence $\mathbf{y} \in \mathbb{R}^{D \times L}$

---

Steps:

1. Calculate importance score for every token: Compute token score vector $\mathbf{s} = (S_\Delta \mathbf{x})^\top K \in \mathbb{R}^{L \times 1}$.

2. Sort the score vector in decreasing order: $\mathbf{s}_{sorted}=\operatorname{Sort}(\mathbf{s}) \in \mathbb{R}^{L\times 1}$

3. Calculate the reordering matrix: $P=\operatorname{Softmax}(-\frac{(\mathbf{s}_{sorted} \mathbb{1}_L^\top-\mathbb{1}_L\mathbf{s}^\top)^p}{\tau}) \in \mathbb{R}^{L \times L}$, where the power function is applied element-wise and $\operatorname{Softmax}$ is applied row-wise.

4. Reorder the input: $\mathbf{x}_{\text{reordered}} = \mathbf{x}P^\top \in \mathbb{R}^{D \times L}$.

5. Calculate the output sequence with selective scan: $\mathbf{y} = \operatorname{SSM}(A, S_B \mathbf{x}\_{\text{reordered}}, S_C \mathbf{x}\_{\text{reordered}})(\mathbf{x}_{\text{reordered}})$

6. Return $\mathbf{y}$.

The computational overhead of our technique lies in steps 1, 2, 3, and 4, which require $O(D^2 \times L + D\times L + L\times L + D\times L^2)$ operations per input sequence.

Additionally, we present the empirical computational costs of our experiment in Section 5.2, benchmarked on the ImageNet classification task using MambaVision-T [Hatamizadeh2024] as the baseline. Table 4 summarizes key metrics, including memory usage, training time per epoch, parameter count, and FLOPs. Table 4: Comparision of computational cost on model training with MambaVision-T as Baseline.

The results show that our reordering technique introduces minimal overhead, with only slight increases across all metrics.

[Hatamizadeh2024] Ali Hatamizadeh and Jan Kautz. Mambavision: A hybrid mamba-transformer vision backbone. arXiv preprint arXiv:2407.08083, 2024.

[Reply 1/3]

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

---

[W1&Q1]. A more comprehensive exploration of higher-dimensional token dynamics should be included. [...] Additionally, comparisons with more sophisticated state-space approaches (such as diagonal vs. non-diagonal state space models) could show how the proposed ideas generalize.

Answer: Regarding the high-dimensional case: We conjecture that our theorems still hold in the high-dimensional case. In the high-dimensional case, the input-output matrix $S_C^{\top}S_B$ is a square matrix of size $D \times D$ and it may have complex eigenvalues. In this case, we consider the symmetric matrix $\mu:=\frac{1}{2}(S_C^{\top}S_B)^{\top} + \frac{1}{2}S_C^{\top}S_B$, which is the symmetric part of the input-output projection matrix $S_C^{\top}S_B$. This term is motivated by the observation that: the term $x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l$ in the considered dynamical system is a scalar and it can be rewritten as $x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l = \frac{1}{2}[x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l]^{\top} + \frac{1}{2}[x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l] = x_l^{\top} \cdot \mu \cdot x_l.$ Since $\mu$ is a symmetric matrix, it always has only real eigenvalues. We particularly conjecture that, in the high-dimensional case, all tokens will converge to zero if $\mu$ has only negative eigenvalues, while all tokens will diverge to infinity if $\mu$ has at least one positive eigenvalue.

To illustrate this conjecture, we run simulations in two-dimensional cases and draw the trajectories of tokens in Figures 6 in Appendix B of the revised version. We observe the following:

• In Figure 6A and 6B, the matrix $\mu$ has at least one positive eigenvalue, and the divergent behavior of the tokens is also observed to persist.
• In Figure 6C, the matrix $\mu$ has only negative eigenvalues, and the convergent behavior of the tokens is observed to persist.

Regarding comparisons with more sophisticated state-space approaches: As described in Algorithms 1 and 2 on page 6 of [Gu2023], S4 and its variants (diagonal, diagonal plus low-rank, HIPPO) can be regarded as special cases of Mamba when the input matrix $B$, output matrix $C$, and step sizes $\Delta$ are input-independent. Consequently, our theoretical analysis of Mamba encompasses those for S4 and its variants as special cases.

References

[Gu2023] Gu, A., & Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752.

[W2&Q2]. 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.

Answer: Thanks for your suggestion. We have expanded our experiment to provide evidence that our proposed method can benefits in a more wide range of tasks. In particular, we have benchmarked on two additional tasks, language modeling and text classification.

For language modeling task, we benchmarked on WikiText103 dataset and achieved perplexity reduction of 1.49, as shown in Table 2.

Table 2: Test Perplexity ($\downarrow$) on Wikitext-103

For text classification, we benchmark on two datasets, IMDB [I] and AGNews [Ag]. The results in Table 3 indicate that our re-ordering technique further improves the baseline model's performance on these tasks.

Table 3: Test accuracy ($\uparrow$) on IMDB and AGNews datasets

These results suggest that our refinement has the potential to generalize effectively beyond vision-related tasks.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback. Beside of our responses above, we would like to provide missing answer on the effectiveness of our reordering technique across a broad range of architectures below.

[W5&Q5]: The authors should include a more detailed experimental comparison against state-of-the-art models in SSMs and transformers [...]

Answer: We conducted an additional experiment to evaluate the effectiveness of our reordering technique across a broad range of architectures. We use the language modeling task on the WikiText-103 benchmark to test three models: Mamba, H3 [1], and Transformer [2]. For the latter two models, no direct equivalent of the step size matrix $S_\Delta$ exists for calculating the token importance score, as defined in Equation (14) of our paper. To adapt the reordering technique to these models, we made the following adjustments:

• H3 Model: The token importance score is computed as $s = \langle x, K \rangle$, where $K$ is a learnable vector.

• Transformer Model: We converted the pretrained GPT-2 small model into the Mamba format and fine-tuned it, following the method outlined in [3]. On the converted model, we calculated the token importance score as described in Equation (14).

The results, presented in Table 5, demonstrate that our reordering method consistently improves perplexity (PPL) across all tested models. These findings highlight the potential of our approach to enhance the performance of a wide range of architectures.

Table 5: Test perplexity on different models

We have added these results to Appendix D.6 of the revised paper.

We would be happy to do any follow-up discussion or address any additional comments.

References

[1] Fu, Daniel Y., et al. "Hungry hungry hippos: Towards language modeling with state space models." arXiv preprint arXiv:2212.14052 (2022).

[2] Vaswani, A. "Attention is all you need." Advances in Neural Information Processing Systems (2017).

[3] Wang, Junxiong, et al. "The mamba in the llama: Distilling and accelerating hybrid models." arXiv preprint arXiv:2408.15237 (2024).

[Reply 3/3]

[Q3]. Do your findings and proposed refinements generalize to other selective state space models beyond Mamba?

Answer: Yes, our findings and proposed refinements generalize to other selective state space models beyond Mamba. Indeed, as described in Algorithms 1 and 2 on page 6 of [Gu2023], S4 and its variants (diagonal, diagonal plus low-rank, HIPPO) can be regarded as special cases of Mamba when the input matrix $B$, output matrix $C$, and step sizes $\Delta$ are input-independent and are chosen carefully. Consequently, our theoretical analysis and the proposed refinements of Mamba encompasses those for S4 and its variants as special cases.

References [Gu2023] Gu, A., & Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752.

---

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

[Reply 2/3]

[Q1]. Could you extend the theoretical analysis to higher-dimensional cases?

Answer: We conjecture that our theorems still hold in the high-dimensional case. A full proof for this conjecture will require a nontrivial work. In the high-dimensional case, the input-output matrix $S_C^{\top}S_B$ is a square matrix of size $D \times D$, and it may have complex eigenvalues. In this case, we consider the symmetric matrix $\mu:=\frac{1}{2}(S_C^{\top}S_B)^{\top} + \frac{1}{2}S_C^{\top}S_B$, which is the symmetric part of the input-output projection matrix $S_C^{\top}S_B$. This term is motivated by the observation that the term $x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l$ in the considered dynamical system is a scalar and it can be rewritten as $x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l = \frac{1}{2}[x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l]^{\top} + \frac{1}{2}[x_l^{\top} \cdot (S_C^{\top}S_B) \cdot x_l] = x_l^{\top} \cdot \mu \cdot x_l.$ Since $\mu$ is a symmetric matrix, it always has only real eigenvalues. We particularly conjecture that, in the high-dimensional case, all tokens will converge to zero if $\mu$ has only negative eigenvalues, while all tokens will diverge to infinity if $\mu$ has at least one positive eigenvalue.

To illustrate this conjecture, we run simulations in two-dimensional cases and draw the trajectories of tokens in Figures 6 in Appendix B of the revised manuscript. We observe the following:

• In Figure 6A and 6B, the matrix $\mu$ has at least one positive eigenvalue, and the divergent behavior of the tokens is also observed to persist.
• In Figure 6C, the matrix $\mu$ has only negative eigenvalues, and the convergent behavior of the tokens is observed to persist.

We have added these discussions with further details in Appendix B.

[Q2]. 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?

Answer. To provide a clearer understanding of our re-ordering technique, we present the pseudo-code for calculating token importance scores and forwarding through a single SSM layer in Algorithm 1 in Appendix D.2 of our revision and also include this pseudo-code below.

Algorithm 1: Forwarding through a SSM layer with token reordering

Input: Input sequence $\mathbf{x} \in \mathbb{R}^{D \times L}$, hidden matrix $A \in \mathbb{R}^{N \times N}$, input matrix $S_B \in \mathbb{R}^{N \times D}$, output matrix $S_C \in \mathbb{R}^{N \times D}$, step size matrix $S_\Delta \in \mathbb{R}^{D \times D}$, learnable vector $K \in \mathbb{R}^{D \times 1}$, hyperparameters $\tau$ and $p$ of $\operatorname{Softsort}$

Output:

• Output sequence $\mathbf{y} \in \mathbb{R}^{D \times L}$

---

Steps:

1. Calculate importance score for every token: Compute token score vector $\mathbf{s} = (S_\Delta \mathbf{x})^\top K \in \mathbb{R}^{L \times 1}$.

2. Sort the score vector in decreasing order: $\mathbf{s}_{sorted}=\operatorname{Sort}(\mathbf{s}) \in \mathbb{R}^{L\times 1}$

3. Calculate the reordering matrix: $P=\operatorname{Softmax}(-\frac{(\mathbf{s}_{sorted} \mathbb{1}_L^\top-\mathbb{1}_L\mathbf{s}^\top)^p}{\tau}) \in \mathbb{R}^{L \times L}$, where the power function is applied element-wise and $\operatorname{Softmax}$ is applied row-wise.

4. Reorder the input: $\mathbf{x}_{\text{reordered}} = \mathbf{x}P^\top \in \mathbb{R}^{D \times L}$.

5. Calculate the output sequence with selective scan: $\mathbf{y} = \operatorname{SSM}(A, S_B \mathbf{x}\_{\text{reordered}}, S_C \mathbf{x}\_{\text{reordered}})(\mathbf{x}\_{\text{reordered}})$

6. Return $\mathbf{y}$.

The computational overhead of our technique lies in steps 1, 2, 3, and 4, which require $O(D^2 \times L + D\times L + L\times L + D\times L^2)$ operations per input sequence.

Additionally, we present the empirical computational costs of our experiment in Section 5.2, benchmarked on the ImageNet classification task using MambaVision-T [Hatamizadeh2024] as the baseline. Table 4 summarizes key metrics, including memory usage, training time per epoch, parameter count, and FLOPs. Table 4: Comparision of computational cost on model training with MambaVision-T as Baseline.

The results show that our reordering technique introduces minimal overhead, with only slight increases across all metrics.

[Hatamizadeh2024] Ali Hatamizadeh and Jan Kautz. Mambavision: A hybrid mamba-transformer vision backbone. arXiv preprint arXiv:2407.08083, 2024.

[Reply 1/3]

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

---

[W1]. Simplifying Assumptions: Using time-independent parameters and excluding other layers for mathematical convenience may oversimplify Mamba's actual behavior in practical settings.

Answer: Our primary objective was to examine the most basic form of Mamba that allows for thorough mathematical analysis. We carried out a few numerical simulations to further validate our tokens dynamic theorems in high-dimensional cases in the revised version (with updates highlighted in blue). Interestingly, we found that in general our theorems on tokens dynamic behaviors remained consistent in the high-dimensional case.

Regarding the time-indepedent weights assumption, we believe that if the entries of the input/output projection matrix $S_C(t)^\top S_B(t)$, the step size matrix $S_\Delta(t)$, and the scalar $a(t)$ are bounded both above and below, the conclusions of Theorems 4.1, 4.3, and 4.4 remain valid. However, this would require modifications to the proofs to account for these assumptions.

Regarding a more complicated setting with additional practical layers, we will provide a few simulation results in the next few days to illustrate the possibility of generalizing our theorems in these cases.

[W2]. Experimental Scope: Though supportive, experimental validation could be more extensive. Including additional datasets and comparisons with other models would strengthen the claims.

Answer: Following your suggestion, we have conducted additional experiments to better assess the performance of our proposed method. Specifically, we benchmarked on two additional tasks: language modeling and text classification. We also plan to conduct an experiment to with other models (SSMs and Transformers) in the next few days.

For the language modeling task, we evaluated our method on the WikiText-103 dataset and achieved a perplexity reduction of 1.49, as shown in Table 2.

Table 2: Test Perplexity ($\downarrow$) on WikiText103

For text classification, we benchmark on two datasets, IMDB [Maas2011] and AGNews [Zhang2015]. The results in Table 3 indicate that our re-ordering technique further improves the baseline model's performance on these tasks.

Table 3: Test accuracy ($\uparrow$) on IMDB and AGNews datasets

Reference [Maas2011] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, 2011. URL https://aclanthology.org/P11-1015.

[Zhang2015] Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. Advances in Neural Information Processing Systems, 2015. https://proceedings.neurips.cc/paper_files/paper/2015/file/250cf8b51c773f3f8dc8b4be867a9a02-Paper.pdf.

[W3]. 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

Answer: Following your suggestion, we have added discussions on the limitations of the proposed refinements and potential challenges in practical implementation to the limitation part of the conlusion section in the revised version. In particular, our token reordering refinement, based on importance 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 a learnable affine line, as used in practical implementation, is optimal. We leave a systematic study of token reordering refinements for future work.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback. Beside of our responses above, we would like to provide missing answers on the actual behavior of tokens' dynamic in practical setting and effectiveness of our reordering technique across a broad range of architectures below.

[W1]. [...] Mamba's actual behavior in practical settings: Following reviewers' suggestion, we provide additional experiments to analyze the token dynamics in the standard Mamba architecture in a practical setting. We evaluate three cases:

(i) There is no constraint on the input-output projection matrix.

(ii) The input-output projection matrix is symmetric and has at least one positive eigenvalue.

(iii) The input-output projection matrix is symmetric and has only negative eigenvalues.

We first conduct simulations using a pretrained standard Mamba architecture without LayerNorm on WikiText103. In Figure 7 in Appendix D.4 in the revised version, we visualize the evolution of token norms across the blocks. Figure 7 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 it to always belong to the class of positive definite or negative definite matrices on every layer, the tokens' dynamic behaviors become quite complicated. Figure 7 suggests that our theoretical analysis of the tokens' dynamic behaviors is still valid in the standard Mamba model without LayerNorm.

In addition, we also conduct simulations using a pretrained standard Mamba architecture with LayerNorm in a practical setting on WikiText103. In Figure 8 in Appendix D.4 in the revised version, we visualize the evolution of token norms across the blocks. Figure 8 shows that the tokens' dynamics in a full Mamba setting can be quite complicated. When the input-output projection matrix has only negative eigenvalues, the tokens do not necessarily tend toward zero. This situation explains the effect of LayerNorm on the tokens' dynamic. Hence, with LayerNorm in the Mamba architecture, our theoretical results on the tokens' dynamics might no longer hold.

[W2]: [...] Including [...] comparisons with other models would strengthen the claims.

Answer: We have conducted an additional experiment to evaluate the effectiveness of our reordering technique across a broad range of architectures. We use the language modeling task on the WikiText-103 benchmark to test three models: Mamba, H3 [1], and Transformer [2]. For the latter two models, no direct equivalent of the step size matrix $S_\Delta$ exists for calculating the token importance score, as defined in Equation (14) of our paper. To adapt the reordering technique to these models, we made the following adjustments:

• H3 Model: The token importance score is computed as $s = \langle x, K \rangle$, where $K$ is a learnable vector.

• Transformer Model: We converted the pretrained GPT-2 small model into the Mamba format and fine-tuned it, following the method outlined in [3]. On the converted model, we calculated the token importance score as described in Equation (14).

The results, presented in Table 5, demonstrate that our reordering method consistently improves perplexity (PPL) across all tested models. These findings highlight the potential of our approach to enhance the performance of a wide range of architectures.

Table 5: Test perplexity on different models

We have added these results to Appendix D.6 of the revised paper.

We would be happy to do any follow-up discussion or address any additional comments.

References

[1] Fu, Daniel Y., et al. "Hungry hungry hippos: Towards language modeling with state space models." arXiv preprint arXiv:2212.14052 (2022).

[2] Vaswani, A. "Attention is all you need." Advances in Neural Information Processing Systems (2017).

[3] Wang, Junxiong, et al. "The mamba in the llama: Distilling and accelerating hybrid models." arXiv preprint arXiv:2408.15237 (2024).