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

Authors: Peihao Wang, Ruisi Cai, Yuehao Wang, Jiajun Zhu, Pragya Srivastava, Zhangyang Wang, Pan Li
Primary area: foundation or frontier models, including LLMs
Keywords: State Space Models, Long-context LLMs, Over-smoothing

Abstract:
Structured State Space Models (SSMs) have emerged as alternatives to transformers.
While SSMs are often regarded as effective in capturing long-sequence dependencies, we rigorously demonstrate that they are inherently limited by strong recency bias.
Our empirical studies also reveal that this bias impairs the models' ability to recall distant information and introduces robustness issues. Our scaling experiments then discovered that deeper structures in SSMs can facilitate the learning of long contexts.
However, subsequent theoretical analysis reveals that as SSMs increase in depth, they exhibit another inevitable tendency toward over-smoothing, e.g., token representations becoming increasingly indistinguishable.
This *fundamental dilemma* between recency and over-smoothing hinders the scalability of existing SSMs. 
Inspired by our theoretical findings, we propose to *polarize* two channels of the state transition matrices in SSMs, setting them to zero and one, respectively, simultaneously addressing recency bias and over-smoothing.
Experiments demonstrate that our polarization technique consistently enhances the associative recall accuracy of long-range tokens and unlocks SSMs to benefit further from deeper architectures.
All source codes are released at https://github.com/VITA-Group/SSM-Bottleneck.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730468557433

thread_type:
official_review

<thread_object>
creator: reviewer_vee3

<content>
summary:
This paper investigates why the neural networks based on state space models (SSMs) are inferior to Transformers for sequence data with long-range dependencies. First, the authors theoretically and experimentally demonstrate that the SSM output’s dependency on input diminishes exponentially with the distance from the output position. They call this phenomenon recency bias. They also reveal through experiments that SSMs are vulnerable to attacks targeting the trailing part of the sequence due to recency bias. Furthermore, they explore the effects of increasing the layers of SSMs and find that while this reduces recency bias, it causes over-smoothing.

strengths:
1. The limitations of SSMs’ capabilities are theoretically examined from multiple perspectives, including recency bias, robustness to attack, and over-smoothing.
2. Since multiple experiments closely linked to the theory are conducted, the paper’s claims are easy to understand.

weaknesses:
1. One of the contributions of the paper is that the authors demonstrate the exponential decay of the input-dependency on the SSMs’ output. However, similar findings have already been established in studies such as [Wang & Xue, 2023] (Theorem 3.13).

    [Wang & Xue, 2023] State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. In NeurIPS 2023.
    
2. The experiments on recency bias, such as those in Figure 2, are limited to a few pre-trained models. As with the experiments in Figure 4, it would be preferable to observe the change in performance across the number of parameter counts. In particular, it should be checked whether recency bias can be reduced when the dimensions of the hidden state and the number of heads are moderately increased.
3. Line 463 suggests a trade-off between A_max (related to recency bias) and A_min (related to over-smoothing). However, it is not trivial whether this is actually a challenge when training SSMs since $A_{\max}\approx 1$ and $A_{\min}\approx 0$ can hold at the same time. To clarify this point, additional experiments would be beneficial, such as observing the dynamics of A_min and A_max during training.

questions:
1. In the experiment in Figure 1, the context length is around 1000, whereas in Figure 2, it exceeds 10,000. In the latter case, would there be any difference between SSMs and Transformers with a context length around 1000?
2. Are there any experimental/theoretical results showing that over-smoothing does not occur in Transformers? Figures 4 and 5 contain experimental results demonstrating over-smoothing in Mamba; is it possible to conduct similar verification for Transformers?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730559010050

thread_type:
official_review

<thread_object>
creator: reviewer_lrov

<content>
summary:
This paper identifies two issues in State Space Models: the recency bias in long-context modeling and the over-smoothing effect that arises as the model scales. The authors provide a theoretical framework to analyze these challenges and substantiate their findings with empirical evidence.

strengths:
The paper is well-written and easy to follow.

Experiments are extensive and effectively support the theoretical results.

weaknesses:
While the authors have comprehensively identified and examined these two issues, they did not provide any discussion on possible methods for mitigation. Including a brief exploration of potential strategies could further enrich the paper's contributions.

The problem of recency in SSMs appears somewhat similar to the vanishing gradient problem in RNNs mentioned in, e.g., [1]. Could the authors provide a discussion comparing these two phenomena and highlight the novelty of this work?

The experimental details for Section 4.1 are currently lacking, which makes it challenging to assess the validity of the results fully. Additionally, if I understand correctly, the over-smoothing problem is characterized by a performance degradation when additional layers are stacked beyond a certain threshold (as observed with GNNs, see [2]). However, this effect does not appear to manifest in the experiments presented in Section 4.1. Could the authors clarify whether this phenomenon is inherently absent in SSMs or if it might be a limitation of the current experimental setup?



Minor comments: 
- Typo in Eq. 4: The indices $i$ and $j$ are likely meant to be $s$ and $t$.
- Typos in Eq. 18 and Eq. 19: the inequality signs are likely to be $ \geq $ and $ \leq $, respectively.


Reference

[1] Y. Bengio, P. Simard and P. Frasconi, "Learning long-term dependencies with gradient descent is difficult," in IEEE Transactions on Neural Networks, vol. 5, no. 2, pp. 157-166, March 1994.

[2] L. Zhao, Y. Song, C. Zhang, Z. Liu, Y. Wang, and H. Lin, "PairNorm: Tackling Oversmoothing in GNNs," in Proceedings of the International Conference on Learning Representations (ICLR), 2020.

questions:
Please see above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730645085638

thread_type:
official_review

<thread_object>
creator: reviewer_53ds

<content>
summary:
This paper lists inspects and provides a theoretical investigation of some of the pitfalls of SSMs. Such pitfalls are caused by a "recency bias," which alters the "token importance," making it flow exponentially down to zero. This fact presents an issue for robustness, as well as for needle-in-a-haystack problems.

strengths:
What I really like about this paper is that it is easy to understand, it is well written and interesting. The contributions are, in my opinion

1) giving practical examples of failure cases of SSMs.

2) showing high-quality and clear plots that can be used for future research or for reference.

This paper indeed seems a bit like a "first step". I like it, though: definitely SSMs have such problems (as intuitively the case and also as pointed out in recent literature) -- yet is **very good to have a reference** - will be interesting for future research!

weaknesses:
I think the following are weaknesses in the current status:

1) The derivation of "exponential memory" results are not novel nor surprising: it is very well known that SSMs have exponential memory (e.g., https://openreview.net/forum?id=i0OmcF14Kf). It is also known that they might struggle with long memory (https://arxiv.org/abs/2402.01032) and to retrieve past information (https://arxiv.org/abs/2312.04927) -- yet scale helps (https://arxiv.org/pdf/2406.07887). 

2) Note that exponential memory per see does not imply loss in perplexity and length generalization; this was shown in recent papers such as https://arxiv.org/pdf/2410.11135v1 (or the xLSTM paper https://arxiv.org/abs/2405.04517). Note also that Alibi positional embedding induce a similar exponential decay.

3) The CIFAR robustness example I find a bit weak because here you are implicitly constructing a 1D model for 2D data: in 1D data, recency bias is actually helpful! In 2D there is no notion of recency. I think the example is interesting for improving or researching Visual state-space models such as VMamba. Please frame this a bit more carefully.

3) The paper does not provide hints into solving the stated problems.

questions:
What do you this should the path ahead be for improving SSM models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730669612529

thread_type:
official_review

<thread_object>
creator: reviewer_2oyp

<content>
summary:
The paper presents 2 limitations of current SSM. The first points out to the fact that SSMs suffer from a strong recency bias. This in turn limits the SSMs ability to recall distant information and introduces robustness issues. Secondly, it is revealed that as SSMs increase in depth, they exhibit a tendency toward over-smoothing, resulting in token representations becoming increasingly indistinguishable

strengths:
1. The paper presents important findings, which open up future directions to resolve the issue.
2. The section on adversarial attacks is praiseworthy, as it utilizes one of the findings to describe something very practical.
3. Both findings are substantiated with some theoretical understanding.

weaknesses:
1. While, empirical results are satisfactory for the recency of SSMs, I don't feel these are extensive for the claim of over-smoothing in SSMs.
2. How do initialization scheme of SSMs interfere with these findings. Can there be some discussion around these?

questions:
Please see above
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper presents a comprehensive analysis of the limitations of State Space Models (SSMs) in capturing long-range dependencies in sequential data. The authors theoretically and empirically demonstrate that SSMs suffer from a recency bias, where the model's output dependency on input diminishes exponentially with distance from the output position. Additionally, they show that increasing the number of layers in SSMs can lead to over-smoothing, which degrades the model's performance.

The reviewers have provided thoughtful feedback, and the authors have addressed their concerns through revisions and additional experiments. The authors have clarified the differences between their work and existing studies and have provided more detailed comparisons with Transformers.

Overall, the paper has been significantly improved through the revision process, and the authors have demonstrated a good understanding of the limitations and potential solutions for SSMs. The paper's contributions, including the theoretical analysis of recency bias and over-smoothing in SSMs, are valuable and relevant to the field.

Hence, after discussion with the authors and among themselves, the reviewers find the paper much improved, with all reviewers leaning towards acceptance. We are therefore happy to accept the paper. We would still like to encourage the authors to address all the reviewers' comments in the camera-ready version.