Understanding and Mitigating Bottlenecks of State Space Models through the Lens of Recency and Over-smoothing

4. Difference between Mamba and Transformer when the context length is around 1000.

We use different pre-trained models in Fig. 1 and Fig. 2. For plotting the influential score, we adopt author released Mamba-130M/1.4B, while the Mamba-Codestral-7B pre-trained by MistralAI is used for the long-context retrieval test. Consequently, these two figures are not directly comparable to each. Our observations show that Mamba-Codestral-7B performs similarly to transformers when the context length is around 1000, likely because it is well-trained for this range of context lengths. However, assuming that the transformer baseline Mistral-7B is pre-trained under similar conditions (a reasonable assumption since both models are trained by the same team), Mistral-7B demonstrates a clear advantage over Mamba-Codestral-7B when the context length is extended.

We conducted a more controlled experiment by training both Mamba and a transformer on the associative recall task, following Zoology [1]. We selected models of size closer to those in Fig. 1 (few millions) and fixed the context length to 1024. The results, presented in the table below, show that Mamba continues to exhibit significant context loss and lower retrieval accuracy compared to transformers, particularly for tokens at longer distances.

[1] Arora et al., Zoology: Measuring and improving recall in efficient language models

5. Does over-smoothing occur to transformers.

As highlighted in numerous prior works, over-smoothing is also a known issue in transformers [1,2,3,4]. While the theory of over-smoothing in transformers has been extensively studied, our work is the first to theoretically establish the over-smoothing effect in SSMs. Compared to the smoothening rate of causal attention in [4], we show that the feature sharpness (measured by Dirichlet energy, as defined in Sec. 4.2) in SSMs decays at a rate of $O((1-A_{min}^T)^L)$, which is faster than transformers' rate of $O((1-A_{min}^T)^{L/T})$, where we denote $L$ as the number of layers. Empirically, we validate this theoretical result by visualizing token smoothness in our revised Fig. 5, which shows that the pairwise token differences in SSMs decays more rapidly, highlighting a more pronounced over-smoothing issue in SSMs compared to transformers.

Moreover, due to the inherent locality of SSMs (Theorem 3.1 in revision), achieving effective long-range interactions necessitates deeper architectures. In contrast, transformers, which allow arbitrary long-range interactions among tokens, do not have the same requirement. Consequently, depth-scaling limitations are more critical for SSMs than for transformers. We have included an extended discussion into our revised Appendix B.

[1] Dong et al., Attention is not all you need: Pure attention loses rank doubly exponentially with depth

[2] Shi et al., Revisiting over-smoothing in bert from the perspective of graph

[3] Wang et al., Anti-oversmoothing in deep vision transformers via the fourier domain analysis: From theory to practice

[4] Wu et al., On the Role of Attention Masks and LayerNorm in Transformers

We sincerely thank Reviewer veE3 for reviewing our work and providing insightful suggestions on connecting our findings with the existing literature and conducting additional experiments to further strengthen our claims. We address your concerns as follows.

1. Comparison with results in [Wang & Xue, 2023].

We thank Reviewer veE3 for pointing out the recent work by [Wang & Xue, 2023]. Distinct from their findings, our primary contribution lies in analyzing nonlinearity and input-dependent mechanisms, which serves as the key to the success of modern SSMs [1,2,3]. The approach in [Wang & Xue, 2023] does not naturally extend to this case, as they assume the linearity of the sequence mixer (cf. Eq. 9 in [Wang & Xue, 2023]). To the best of our knowledge, our work presents the first counterargument showing that even with input-dependent SSMs, effective context filtering may not be achieved. Instead, these mechanisms impose a strict recency bias, inheriting the limitations of linear SSMs as previously highlighted by [Wang & Xue, 2023]. Our analysis considers both $A_t$ and $b_t$ as input-dependent, aligning with settings in Mamba and other follow-up works [2,3]. Furthermore, we conduct a finer-grained analysis, quantitatively relating the decay rate to the specific values within the input-dependent gating matrix $A_t$. We have summarized this discussion into our related work (Sec. 6) and include a full version in Appendix B.

[1] Gu & Dao, Mamba: Linear-time sequence modeling with selective state spaces

[2] Yang et al. Gated linear attention transformers with hardware-efficient training

[3] Arora et al. Zoology: Measuring and improving recall in efficient language models

2. Experiments in Fig. 2 are limited to a few pre-trained models. Whether increasing hidden dimension and the number of heads can mitigate recency?

We thank Reviewer veE3 for the suggestion to reproduce our empirical observations across a variety of models. However, we would like to point out that the model used in Fig. 2 is Mamba-Codestral-7B, which is currently the state-of-the-art SSM model known to us. At present, we are unable to find an open-source pretrained Mamba model that is larger and more advanced. Nevertheless, we still observe that such industry-level model aligns well with our theoretical findings. In contrast, transformer-based model Mistral-7B developed by the same team does not exhibit similar limitations. If the reviewer has any recommendations for models that are larger and more capable of long-context tasks than this baseline, we would greatly appreciate it if you could share them with us.

3. Training dynamics of $A_{max}$ and $A_{min}$.

We agree with Reviewer veE3 that $A_{max} \approx 1$ and $A_{min} \approx 0$ can happen at the same time. However, our empirical findings show that these values are largely concentrated within a narrow range. To illustrate this, we visualize the distribution of $(A_{max} - A_{min})$ across different channels in Fig. 6 of our revised manuscript. Each bin represents the proportion of channels whose memory state satisfies $(A_{max} - A_{min})$ being smaller than the corresponding threshold on the x-axis. Notably, over 60% of channels exhibit $(A_{max} - A_{min})$ values smaller than 0.5, indicating that most channels cannot simultaneously achieve $A_{max} \approx 1$ and $A_{max} \approx 0$. Memory representations will inevitably undergo either exponential diminishing or over-smoothing.

Thank you for your comments and for the reminder. I am also grateful for your additional experimental results. Since many of my concerns were resolved by the authors' reply, I raised my score by one.

We are grateful for Reviewer LroV's efforts and time on reviewing our work. We have taken your valuable suggestions and carefully revise our manuscript to incoporate some initial solutions for mitigating the stated problems. Please see our responses to your questions:

1. Discussion on possible methods for mitigation.

We thank Reviewer LroV for this thoughtful suggestion. We added the following dicussion on the potential solutions based on our theoretical results.

As our theory suggests, the maximal and minimal values in $A_t$ directly governs the locality and over-smoothing issues in SSMs, respectively. Encouraging the diversity of values in $A_t$ is thus crucial. While $A_{max} \approx 1$ and $A_{min} \approx 0$ could could theoretically occur simultaneously, our new experiments in Fig. 6 reveal that this does not happen in practice, as the values in $A_t$ are often concentrated within a narrow range. A simple yet promising solution to address this issue is to fix one component of $A_t$ as a constant 1, another as a constant 0, and others freely learnable. This approach ensures that one dimension in the state memory consistently focuses on the current token, counteracting over-smoothing by preventing mixing with previous tokens. Simultaneously, another dimension exclusively retains information from past tokens, avoiding locality issues by preserving the complete history. We are currently in the process of empirically validating this simple solution and will notify the reviewer and update the draft as soon as we have the results.

2. Discussion comparing recency and vanishing gradients.

We appreciate Reviewer LroV's comments on connecting our theory with vanishing gradients. Vanishing gradients refer to a challenge in RNNs, where backpropagation-based learning is impeded due to gradient magnitudes decaying exponentially over time. The diminishing dependencies among distant tokens are a fundamental cause of this issue [1]. SSMs were initially proposed to address this limitation by explicitly modeling long-range dependencies, as highlighted in [2, 3]. Subsequent work, such as Mamba, extends this approach by adopting their initialization alongside newly proposed selection mechanisms, which are widely believed to enhance these capabilities. Our Theorem 3.1 (formerly Theorem 3.2) lies in theoretically challenging this claim. We demonstrate that modern SSMs still suffer from the recency bias, which not only undermines their ability to capture long-term dependencies but also potentially exacerbates the vanishing gradient problem. We have included this discussion to Appendix B of our revision.

[1] Bengio et al., Learning long-term dependencies with gradient descent is difficult

[2] Gu et al., Hippo: Recurrent memory with optimal polynomial projections

[3] Gu et al., Combining recurrent, convolutional, and continuous-time models with linear state space layers.

3. The experimental details for Section 4.1 are lacking.

Our scaling experiments are aligned with the setting adopted in Mamba (see Appendix E.2.1 in [1]). Specifically, in each curve of Fig. 4, we vary the size of models from 140M to 550M by changing the hidden dimension. Then we train these varying-sized Mamba with a certain amount of tokens determined by Chinchilla laws. We have included a detailed version in our Appendix D.3 of the revised draft.

[1] Gu & Dao, Mamba: Linear-time sequence modeling with selective state spaces

4. Performance downgradation caused by oversmoothing does not appear in Section 4.1.

Over-smoothing does not always result in performance degradation. In many prior studies [1,2,3], over-smoothing has been observed to manifest as performance saturation, where the model's performance plateaus as its depth increases.

In our new experiments, we further increased the depth of Mamba, following the setup described in Sec. 4.1. We observed a noticeable performance drop when the depth exceeded 48 layers. As shown in Fig. 4 (of the revision), the validation perplexity begins to rise for both models with sequence lengths of 2048 and 8192, indicating a decline in performance (e.g., in the 64-layer and 72-layer models).

[1] Zhou et al., Deepvit: Towards deeper vision transformer

[2] Wang et al., Anti-oversmoothing in deep vision transformers via the fourier domain analysis: From theory to practice

[3] Chen et al., Bag of tricks for training deeper graph neural networks: A comprehensive benchmark study

5. Typos and writing errors.

We thank Reviewer LroV for pointing out writing errors. We have resolved them in our revision. We have moved Lemma 3.1 (currently Lemma C.1) to Appendix C.1 to save some space for other more meaningful discussions.

3. CIFAR robustness examples are weak and needs reframe.

The image classification with CIFAR is used as a showcase to demonstrate the robustness issues of Mamba caused by recency. We note that testing sequential models with image datasets is a standard practice in sequence modeling, as used in well-known works [1, 2]. Image dataset is also a part of Long-Range Arena (LRA) [3], a popular benchmark for long-sequence modeling.

While our adversarial attack experiments focus on image datasets, the findings extend to language models. In particular, it is reasonable to generalize our conclusions on targeted attack experiments to LLMs that: SSM-based LLMs can be particularly vulnerable to jailbreak attacks [4, 5]. This is because SSMs prioritize recent information over past tokens, making it easier to bypass system prompts by appending jailbreak instructions at the end of the input. Moreover, our theoretical analysis suggests that fine-tuning LLMs with instructional datasets or human feedback to enforce adherence to system prompts may not resolve this vulnerability, as the recency bias remains inherent to SSM models regardless of weight configurations. We have reframed our Sec. 3.3 to include these discussions.

[1] Gu et al., Combining Recurrent, Convolutional, and Continuous-time Models with Linear State Space Layers

[2] Gu et al., Efficiently modeling long sequences with structured state spaces

[3] Tay et al., Long range arena: A benchmark for efficient transformers

[4] Perez & Ribeiro, Ignore previous prompt: Attack techniques for language models

[5] Zou et al., Universal and transferable adversarial attacks on aligned language models

4. No hint to solve the state problems.

We appreciate Reviewer 53Ds for suggesting the inclusion of hints to address the stated problems. While this aspect was initially beyond the scope of our work, our theoretical findings do point to some intriguing potential solutions worth exploring.

Our theory highlights that the maximum and minimum values in $A_t$ play a pivotal role in determining the extent of locality and over-smoothing in SSMs, respectively. Therefore, promoting diversity in the values of $A_t$ is a potential approach of mitigating locality and over-smoothing simultaneously. Although it is theoretically possible for $A_{max} \approx 1$ and $A_{min} \approx 0$ to happen at the same time, our new experiments in Fig. 6 reveal that, in practice, the values in $A_t$ tend to cluster within a narrow range. To mitigate this issue, a straightforward yet effective strategy involves fixing one component of $A_t$ at a constant 1, another at a constant 0, while allowing the remaining components to be freely learnable. This design ensures that one dimension of the state memory (with all-pass gating value) focuses exclusively on retaining past information, addressing locality concerns by preserving the full history. The dimension with zero gating value remains dedicated to the current token, thereby mitigating over-smoothing by preventing unnecessary mixing with previous tokens. We are actively validating this approach empirically and will update the draft with results and inform the reviewer as soon as we have the results.

We appreciate Reviewer 53Ds's for the positive initial assessment for our work. Regarding your concerns on novelty and experiment settings, please see our responses to your questions:

1. The "exponential memory" results are not novel or surprising.

While the concept of "exponential memory" may seem intuitive to many researchers in this field, we emphasize that our work is the first to theoretically demonstrate this property in the context of input-dependent mechanisms such as selection [1], contextualized gating [2], and input-dependent convolution [3], which are integral to modern SSMs. Existing works [3, 4, 5, 6] provide valuable insights into SSMs from various perspectives. However, the results in [4] are limited to linear SSMs (e.g., S4), considering nonlinearity only after applying the linear SSMs. The method proposed in [4] does not naturally extend to this scenario, as it assumes linearity in the sequence mixer (cf. Eq. 9 in [4]). Please refer to Appendix B of our revision for a detailed comparison. [3] and [5] primarily analyze the limitations of SSMs using specific synthetic tasks, such as copying and associative recall, which do not directly expose the long-range dependencies among tokens. [6] establishes a strong benchmark for modern SSMs. Our theoretical results further elucidate these limitations. Additionally, our long-context retrieval experiments extend these findings to scaled-up industrial-level models, demonstrating that even with scaling-up, the long-range limitations of SSMs persist. Furthermore, we identify a novel security vulnerability in SSMs, which can lead to easier jailbreaks in LLMs (see additional discussions in Sec. 3.3 of our revision). To cover these points, we have added a new related work section in our revision.

[1] Gu & Dao, Mamba: Linear-time sequence modeling with selective state spaces

[2] Yang et al. Gated linear attention transformers with hardware-efficient training

[3] Arora et al. Zoology: Measuring and improving recall in efficient language models

[4] Wang & Xue, State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory

[5] Jelassi et al., Repeat After Me: Transformers are Better than State Space Models at Copying

[6] Waleffe et al., An Empirical Study of Mamba-based Language Models

2. Exponential memory does not imply loss in perplexity and length generalization.

We agree with Reviewer 53Ds that exponential memory does not necessarily lead to a loss in perplexity and can even aid in length generalization, as demonstrated by xLSTM [1] and Alibi [2]. As discussed in Sec. 3.1, recency serves as a near-optimal bias for language modeling, meaning its imposition does not negatively impact perplexity.

However, we wish to emphasize that a key contribution of our work is showing that, despite not affecting perplexity, this decaying design has notable impacts on many downstream tasks. These issues include, but are not limited to, substantial loss of long-distance context and more severe security vulnerabilities (refer to our responses to (3)). This observation also underscores that validation perplexity may not fully reflect all aspects of a model's performance and is often an inadequate metric for evaluating long-range dependencies.

[1] Pöppel et. al., xLSTM: Extended long short-term memory

[2] Press et. al., Train short, test long: Attention with linear biases enables input length extrapolation

Dear Reviewer 53Ds:

Thank you again for your valuable feedback and thoughtful suggestions during the review process. We have made the following updates based on your feedback:

1. Added a dedicated related work section (Sec. 7) to discuss relevant prior work and included detailed comparisons with [Wang & Xue, 2023] in Appendix B.
2. Expanded Sec. 3.3 to generalize the implications of our robustness experiments to LLM jailbreaking.
3. Proposed a unified technique in Sec. 6 to simultaneously mitigate locality and over-smoothing, supported by validations aligning with our theory (Tab. 2). Also see our previous replies.

As the discussion period is nearing its end, have you gotten a chance to read our responses and finalize your assessment? We sincerely wish our hard work, revealing the recency issue and over-smoothing nature of modern SSMs for the first time, could receive a more positive feedback.

Best,

Authors of Paper 1195

We appreciate Reviewer 2oyp's for the positive initial assessment of our work. We have included more experiments to strengthen our over-smoothing part. Please see our detailed responses below:

1. Empirical results for over-smoothing are not extensive.

To strengthen our experiments on over-smoothing, we have added more results in Fig. 4 and Fig. 5. Specifically, we extended the depth-scaling experiments in Sec. 4 and further increased the depth of Mamba, following the setup in Sec. 4.1. We observed a significant performance drop when the depth exceeded 48 layers. As shown in Fig. 4 (of the revision), validation perplexity rises for both models with sequence lengths of 2048 and 8192, indicating a decline in performance (e.g., in the 64-layer and 72-layer models). This finding suggests that over-smoothing caused by increased depth not only saturates performance gains but also degrades performance, akin to previous observations in GNNs [1].

We have additionally measured the smoothness of trasformers' hidden states in Fig. 5. We demonstrate that SSMs are more prone to over-smoothing.

We would greatly appreciate Reviewer 2oyp's detailed advice on how we might further enhance our experimental section.

[1] Li et al., Deeper Insights into Graph Convolutional Networks for Semi-Supervised Learning

2. How do initialization scheme interfere with these findings?

We thank Reviewer 2oyp for initiating this insightful discussion. Our theoretical results confirm that initialization plays a crucial role, as it impacts the distribution of values in $A_t$, where the maximum and minimum values determine the extent of locality and over-smoothing, respectively. To address both issues, our theory suggests that $A_t$ should span a wider range of values, which could be achieved through carefully designed initialization schemes. However, we note that the current initialization scheme used in Mamba performs poorly in promoting diversity in $A_t$. As shown in our newly added Fig. 6, the majority of values are closely clustered, leading to either short-range dependencies or over-smoothing of token representations.

Moreover, we point out why the initialization scheme derived from HiPPO theory fails to work effectively for Mamba in Sec. 5 (formerly Sec. 3.1). The primary reason is that Mamba simplifies the parameterization to a diagonal form, which undermines the merits of HiPPO matrices. This limitation has also been noted in previous work (Proposition 3 in [1]).

[1] Gu et al., How to train your hippo: State space models with generalized orthogonal basis projections