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
Venue: ICLR 2025 Poster (2025)
Keywords: State Space Models
Long-context LLMs
Over-smoothing
Primary area: foundation or frontier models, including LLMs
TL;DR: We theoretically and empirically reveal and tackle locality and over-smoothing bottlenecks for state space models.

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.

Published as a conference paper at ICLR 2025
UNDERSTANDING AND MITIGATING BOTTLENECKS OF STATE SPACE MODELS THROUGH THE LENS OF RE- CENCY AND OVER-SMOOTHIN
Peihao Wang', Ruisi Cai', Yuehao Wang', Jiajun Zhu Pragya Srivastava', Zhangyang Wang', Pan Li
'University of Texas at Austin, *Zhejiang University, Google DeepMind, "Georgia Tech (peihaowang, ruisi cai, yuchao, at atwaswang @uteeas.ccoo junniane z ju. edu. en, pragyaBrrivastava com, panli@gatech.ens
ABSTRACT
Structured State Space Models (SSMs) have emerged as alternatives to trans- formers. While SSMs often regarded 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 mod- els' ability to recall distant information and introduces robustness issues. Our scaling experiments then discovered that deeper structures 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-smo,,,ing, e.g., token representations becoming increasingly indistinguish- able. This fundamental dilemma between recency and over-smoothing hinders the scalability of existing SSMs. Inspired by our theoretical findings, we pro- pose 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 consis- tently enhances the associative recall accuracy of Long-range tokens and unlocks SSMs to benefit further from deeper architectures. All source codes are released at htttps://gitai com VVTA
Www.AVATTAT--SSSA- Group SSM Bottllneck
INTRODUCTION
Sequence processing architectures have evolved from RNNs (Hochreiter & Schmidhuber 1997]
Sutskever et al. 2014 Cho et al.) 2014 Cho 2014), to transformers (Vaswani al. [017]] 7 Devlin
al. 2019] Radford al. 2018 2019 Brown et al. 2020), and more recently to State Space Models
(SSMs) (Gu et al. 202 a Gu Dao 2023). particular, SSMs (Gu al. 2021a; Gu & Dao
2023 Dao & Gu 2024) have emerged as a compelling alternative to transformers, enabling more
efficient handling of long sequences. SSMs operate in two modes: convolution and recurrence (Gu al. 2021b). During convolutional mode, SSMs assume visibility of the entire sequence and utilize hardware-optimizzd convolutions to propagate information across all tokens in parallel. This approach avoids the calculation of pairwise correlations inherent in attention mechanisms, thereby accelerating training. Mamba (Gu & Dao 2023) enabled convolution with a parallel scanning algorithm, which yielded more expressive sequence-level mixing without sacrificing efficiency. During recurrent mode, SSMs process one token at time while maintaining a compact recurrent hidden state encoding the sequence history. The outputs are sequentially decoded from this hidden state, avoiding storing all past key-value pairs (Dai et al. 2019) and thus reducing inference memory usage.
Furthermore, SSMs have been meticulously tailored effectively capture long-range dependencies and filter contextual information. These models are grounded in HiPPO theory (Gu et al.) 2020). which demonstrates that a first-order Ordinary Differential Equation (ODE) can encapsulate long-term memory through a designated state matrix known as the HiPPO matrix. Subsequent research (Gu et al. 2021b a Gupta et al.] 2022; Gu et al., 2022a) has simplified this state matrix to diagonal form, significantly improving computational efficiency while retaining the ability model long- range dependencies. Mamba (Gu & Dao 2023) introduced a selection mechanism that selectively aggregates pertinent information from the context into the state. Concurrently, linear attention models

---

Published as a conference paper at ICLR 2025
have been derived from streamlined attention mechanisms (Katharopoulos et al., 2020; Sun et al.) 2023; Peng et al. 2023| Yang et al. 2023). Collectively, these advances can be interpreted through a unified lens more structured SSMs (Dao & Gu 2024).
Despite their initial empirical successes, recent findings indicate that SSMs may not match transform- ers in their ability to recall information from long contexts (Arora et al. [023)] Poli et al. (024) or in handling more complex retrieval patterns (Park et al. (024). Furthermore, it has been observed that Mamba continues underperform compared to transformers on larger scales (Walefte et al. 2024).
These shortcomings have not yet been systematically elucidated.
In this paper, identify two fundamental limitations of SSMs in their ability to model complex long- range dependencies. First, we argue that the long-term memory capabilities of modern SSMs may be misinterpreted. Our analysis reveals that an SSM layer exhibits a strong recency bias, limiting tokens to primarily interact with the nearby context. This bias is intrinsic to SSMs and many linear attention models, regardless of the content-infer inniennniifooiingggg techniques employed, such as the selection mechanism introduced by Mamba (Gu & Dao 2023). We further posit that the loss of long-range capabilities may stem from the oversimplification of HiPPO-induced SSMs, which trades efficiency off performance. To substantiate this claim, we perform long-range retrieval task on an industrial-scale language model (Jiang et al. 2023) based on Mamba. Our test results indicate that Mamba catastrophically forgets distant content once the context length exceeds its memory capacity. Furthermore, raise a novel robustness concern regarding SSMs with recency bias: our empirical outcomes show that Mamba is more susceptible perturbations local tokens, making it vulnerable to adversarial attack, as these local tokens can be easily manipulated to serve as backdoors.
We then conduct series of scaling experiments with varying context lengths during the pre-training of SSMs. Our results indicate that increasing the model depth is crucial for enhancing its ability to utilize long contexts by expanding the receptive field. However, we observe that depth scaling encounters another bottleneck as performance begins to saturate as depth continues to increase. We analyze the feature dynamies across the SSM layers and theoretically revealed that SSMs inherently function as smoothing operators, leading to over-smoothing in deep architectures (NT & Maehara, 2021 Oono & Suzuki 2019; Cai & Wang 2020). As result, token representations become increasingly uniform and indistinguishable with each additional layer.
Recognizing such fundamental dilemma between recency and over-smoothing, we introduce a unified approach called the polarization, which simultaneously addresses recency bias and over-smoothing. Specifically, we reserve two dedicated channels in the state transition matrices of SSMs and reset them to zero and one, respectively. The all-one channel helps preserve historical information and prevents catastrophic forgetting caused by locality, while the zero channel inhibits excessive fusion of information from past tokens, effectively slowing the smoothing rate. Our experiments associative recall (Arora et al., 2023) demonstrate that the polarization technique significantly improves recall accuracy by mitigating locality artifacts and can gain more performance when combined with deeper architectures by effectively alleviating over-smoothing.
2 PRELIMINARIES
In SSMs (and alike), we represent the a discrete-time sequence of T tokens as x [x1 æT] E RT. For vector-valued input sequences, SSMs process each channel inde- pendently. To simplify notations, we focus on scalar-valued sequences without loss of generality. The impact of multi-channe inputs will be addressed in the relevant context. SSMs learn to represent and forecast the next token by integrating past information. Formally, SSMs can be viewed as a soquence-to-sequeme transformation from inputs RR to outputs y E IRT through memory state he E HN,, which is iteratively updated with linear recurrence. A general form can be written as:
h = A + (x:) y: = chh,), hg - 0, Vt E [T].
(I)
where & E [7]] denotes the time step. Intuitively, A, E NxN extracts information from the previous
state h 1 b, : R + N projects every input token to the hidden space, A, E E controls how
much information of the new token will be fused into the hidden memory, and c IR + IR decodes
In most scenarios discussed in this paper, we assume real parameterization by default, as it is the standard approach in the cases of our primary interest, such as language modeling (Gu & Dao 2023
2

---

Published as a conference paper at ICLR 2025
the hidden state at time to the final prediction. In all SSMs considered in this paper, it is necessary to assume (A,, b,, Ct: At) only depends on the inputs at the time t. SSMs are trained from end to end to optimize for the parameters {[A,, b,, Ct, he(r) for which the different SSMs adopt various types of instantialization. Below, we list some representative examples.
S4, DSS, and S4D. The seminal works (u al. 2020 2021b 2022b) demonstrate that discretizing time-invarinn ODE ht( Ah(t) + (z) with some special realization of matrix A can yield an efficient recurrent network for long-sequence modeling. The follow-up works Gu al. (2021a) together with Gupta et al. (2022); Gu et al. (2022a) simplifies A to be diagonal matrix. Applying the zero-order hold rule for discretization, as suggested by Gupta et al. (2022), we can summarize this series of models in the form of Eq.
(S4) A. exp (AA),
bı(1)) b c(hs) c h A
where (A, b, c. A) are learnable parameters. In particular, A is restricted to be diagonal matrix and can be complex valued. However, A must have negative real part (Gu et al. 2022a). A (0, 1] is often interpreted as the time interval for discretization. We call this family of SSMs S4 following the
naming convention in Gu & Dao (0023).
Mamba. A recent breakthrough Mamba (Gu & Dao 2023) introduces the selection mechanism to extend $4. Instead of learning (A, b, c. A) in Eq. as free parameters, Mamba conditions (A, b, c, A) on the inputs, which enables each iterative step in Eq. to filter useful token information during the recurrence. Specifically, Mamba computes (A,, b,, Ct2 A.) as follows:
(Mamba) A. = A)), by (ax)) = (WBa,)a,, ca (Wex:)) his A. a (WaT)) (3)
where WA E R. WB E R-N, Wo E N are learnable weights in addition to A, and (+)) denotes softplus activation. When handling multi-dimensional token embeddings, WA, WB, Wo are ex- tended on the input dimension, different (Ar, b,, C, A,) are assigned each token channel and Eq. is running channel-wiselyy A, is varying across different channels, while b, Ct are shared across token channels. That is, if the token embedding dimension is D, then each token channel will be allocated with a distinct WA E RR,, and shared Wn. E RN and We R In language modeling, A has strictly negative real-valued diagonal, which ensures A, E (0, 1)N Additionally, Mamba is integrated with the H3 architecture (Fu et al. 2022), wherein the selective SSMs is working with a local convolution and sandwiched by two gated connections.
Linear Attention. Concurrent with SSMs, there is another line of work streamlining attention to linear time complexity. With slight abuse of terminology, we name all of them collectively as Linear Attention Models (LAMs). We observe that many of them can be written in the form of Eq, D such that in the remainder of this paper, we extend the definition of SSMs to include LAMs without introducing ambiguity, as LAMs and SSMs are dual each other (Dao & Gu 2024). We leave full
summary to Appendix B
3
CAN SSM EFFECTIVELY REPRESENT LONG-RANGE DEPENDENCIES?
3.1 SSMS ARE LOCALLY BIASED
In this section, we investigate the ability of SSMs to learn long-range dependencies. Recent studies find that SSMs seem more effective than transformers on this task (Gu et al.) 2020 Tay et al. 2020 Li et al. 2022, Gu & Dao 2023). However, in Sec. 3.T we theoretically show a negative result that an SSM layer is inherent to local bias and loses long-term memory exponentially. In Sec. 3.2
empirically justify our claim by showing SSMs struggle to retrieve from distant context. We also
demonstrate that the local bias may lead to robustness issues in Sec. 3.3
To understand how information is propagated and long-range dependencies are modeled in SSMs, we aim uncover the relationship between the output time E [T] and the input token at time s & t. We define the derivatives (8 0xx as the influential score to represent the importance of the s-th
2 Moor rigorously, by zero-order hold, b should be parameterized as b (AA) (AA)- Coxpp-AAA -AA) T)b.
However, the presented form more commonly used in practice as in Gu Dao 2023
3

---

Published as a conference paper at ICLR 2025
imput token to the t-th output token. Note that jay,/Oz, is well-defined for every s, t [T] long as (A,.bs b,, Ct, A.) are all differentiable in terms of æ. Intuitively, if @0y//xz,l is larger, then the s-th input token is more influential on the t-th output token, and vice versa.
Below we present a formal result regarding the influential score.
Theorem 3.1 (Recency of SSMs). Consider an SSM defined in Eq. 7 Twith {(A,, be, ct, bem
Assume that (i) the input space X C. IRT is compact, (ii) {(A,, bt, Ct2 hejm are continuous and have continuous derivatives, and (iti) A, E (0, (0.1)Nx3 are diagonal matrices for all t E [T]. Let Amaz = [[ESEeT]],NEEN]] (At)n.n. Then for arbilrary x E X and every 8, I E (T] such that < I Jdy. /0z = (exx(-=n(t - s)) for some K = ((llo(AAma)))
The proof can be found in Appendix D.I The first two assumptions are standard and always satisfied. The third assumption also holds for most of SSMs discussed in
Sec. 2 Therefore, Theorem 3.1 applies to numerous
SSMs including S4 (Gu et al. 2021a 2022a), Mamba
10
(Gu & Dao 2023), and many LAMs (Sun et al. 2023]
Peng et al. 2023 Yang et al.) 2023 Qin et 2024 De et al. 2024). Theorem 3.1 states that influential scores
log Score
130M pretrained
1.4B pretrained
between two tokens modeled by SSMs are exponentially diminishing with respect to their relative distance. The decay rate is determined by the maximal values among all A,'s elements. The closer Amaz is to zero, the faster the influential scores decay. The practical implication is that SSMs are factually recency-biased models. Tokens farther away are reaching and forgotten rapidly while the information of closer tokens dominates the final output. This can significantly limit their ability of fitting complex long-range relationships.
training
1.48 wio training
0
100
200
300
400
500
Relative token position
Figure 1: Visualization of log influential
scores log (0y:/Ox.. versus distance (t s)-
Empirical Validation. In Fig. [ we plot the logarithmic influence scores against relative distances. Across various model sizes, Mamba consistently exhibits linear decay rate of influence scores, both at initialization and throughout training. This observation suggests that recency arises from an inherent model bias, described in Theorem [.1] rather than solely from capturing data statistics. Moreover, the adopted initialization scheme further amplifies the locality bias.
Is the design of decay necessary or desirable? One key observation from theory is that the parameterization of A, within the interval (0, 1) leads to strictly decaying dependencies among tokens based on their relative distances. This design choice appears to be standard practice, and perhaps
intentional, in several recently proposed SSMs (Gu & Dao 2023
Dao & Gu 2024 Beck et al. 2024
Yang al. 2023 Peng et al. 20244] De et al. 2024]] Ma et 2024 Liu al. 2024 Yang al. 2024).
Interestingly. ly, it has been demonstrated that this " "intentional" decay of long-range dependencies not only avoids degrading the perplexity of language models but also improves generalization for length extrapolation. This observation aligns with the empirical success of soft gating mechanisms adopted in traditional RNNs (Cho 2014 Cho et al. 2014 Gu et al. 20216) and the decaying patterns imposed on transformers (Raffel et al.) 2020) Press et al.) 2021] Sun et al. 2022). We find the constraint NN is likely inherent to SSMs as plays a critical role ensuring the numerical
A, E (0, N stability for length generalization during long-context recurrence (Gu et al. 2022a Yang et al.) 2023). Promoting the importance of local tokens could also lead to nearly correct bias, as natural language generation mostly utilizes recent contexts. However, design comes with a significant drawbacks: it result in substantial loss of long- cong-distance information (Sec. 3.2) and may raise potential security concerns (Sec. 3.3). This observation underscores that the validation perplexity may not fully capture all aspects of a model's capabilities and can be an inadequate metric for assessing long-range dependencies.
will detail in subsequent sections, this
3.2 LOST IN THE DISTANCE: LONG-CONTEXT RETRIEVAL TEST
To assess the ability of large language models (LLMs) to effectively utilize long-conted data, we evaluate open-source SSM using the "Needle in Haystack" benchmark and compare its performance
4

---

Published as a conference paper at ICLR 2025
10240
0.49
0.49
0:47
0.47
0.46
0.49 0.52
0.61
0.74
10240
0.72
0.72
0.72
0.72
9.72
0.72
0.72
0.72
0.72
-08
Full Conteit Length
0.7
14336
0.41
0.39
0:39
0.37
0.47
Contost
14336
0.7
0.7
0.6
18432
0.34
0.33 0.32
0.3 0.29
0.31
0.33 0.39
0.54
18432
+0.72
0.72
0.72
0.72
0.72
0.72
0.72
0.72
0.72
0:29
0:28
0.32
0.37
22528
0.75 0.75 0.75
10g
20%
70%
800%
90%
10%
20%
30%
90%
Neodle Position
Needle Position
Figure 2: Comparison between SSM and Transformer on "Needle in Haystack" benchmark. The left figure shows the retrieval accuracy of the Mamba-Codestral- Codestral-78 7B model, while the right figure presents the retrieval accuracy of the Mistral- 7B model. We present a heatmap where "full context length" refers to the total length of the document, and "needle position"" denotes the relative position of the statement to be retrieved within the context. See
fine-grained visualization Appendix F.2
with that of Transformer. In this benchmark, randomly generated statement is embedded within the middle of long document, and the models are tasked with retrieving the statement. By varying the insertion position of the statement, we measure the retrieval accuracy at each location, which reflects the model's positional bias. To enforce LLMs using the data within context, instead of recalling information memorized by its model weights, we carefully design the statement with factual error. See detailed examples in Appendix E.2
We compare the retrieval accuracy of the Mamba- Baibba-Codestral723 model, a representative SSM capable of handling long-come context inputs of up to 256k tokens, with Mistral-7B (Jiang et al.) 2023), which utilizes a transformer architecture. As illustrated in Figure the retrieval accuracy of the Transformer remains stable regardless of the needle position. In contrast, the SSM achieves higher accuracy when the needle is placed closer to the end of the context (i.e., larger needle position values), while its accuracy drops when the needle is located near the beginning of the document. This indicates positional bias towards local tokens in the SSM.
3.3 POTENTIAL RISK ON MODEL ROBUSTNESS
We conduct quantitative experiments show the recency-biased nature of SSMs will lead potential hazards. The downstream task in this study is image classification on sequences of pixels (Tay et al. 2020), where W H images are flattened to sequences of pixel tokens and fed sequence models for classification. We test a family of SSMs, including H3 (Fu et al. 2022). RWKV (Peng et al. 2023), and Mamba (Gu & Dao (023), and compare them against a transformer baseline (Vaswani et al. (017) on CIFAR- 10 dataset. To adapt SSMs for this task, we append learnable class token after the last token of the input sequence. The output state of this class token is then mapped to logits using a classifier head. Experiment details are given in Appendix BK In the following, two attack patterns on the input data are introduced, which degrade the robustness of SSMs in this task.
Adversarial Attack. To assess the bias of SSMs towards corrupted data, we perturb the leading and trailing tokens input sequences with random noise. unbiased models, perturbations both leading and trailing tokens cause similar performance drops. However, in locally biased models, where the class token is appended after the last input token, the trailing tokens are supposed to have greater impacts on classification outcomes than leading tokens. Table presents experimental results on the CIFAR- dataset under two corruption ratios. For each ratio, the same number of
Corrupted region (seq. length = 1024) [0:32]
Models
(no corrupl)
[992:1024]
1928:
[0:96]
H3
0.654
0.569 (L 13.04%)
0.654 (L 0.03%)
0.477 (L 27.07%)
0.650 0.72%)
Transformer
0.580
0.535 (). 7.81%)
0.447 (, 22.95%)
0.431 (4, 26%7
0.370 (L 36.32%)
RWKV
0.474
0.150 (L 68.35%)
0.466 (} 1.58%)
0.138 (L 70.88%)
0.460 2.91%)
Mamba
0.674
0.126 (L 81.24%
0.658 (1 2.30%)
0.098 (1 85.46 %)
0.647 (L 3.98%)
Table 1: Results of adversarial attack experiments on the CIFAR- 10 dataset, evaluated using classification accuracy, Each input sequence contains 1,024 tokens. Two corruption ratios (32/1024 and 96/1024) applied perturb the leading and trailing tokens, respectively.
5

---

Published as a conference paper at ICLR 2025
0.8
region
0.6
Attack region
[0:256]
[0:480]
Success rate
[768:1024]
Success rale
[544:1024]
0.4
0.4
0.2
0.2
0.0
0.0
H3
Transformer
RWKV
Mamba
H8
Transformer
RWKV
Mamba
(a) Attack ratio 256/1024
(25.00%)
(b) Attack ratio 480/1024
(46.875%)
Figure 3: Results of target attack experiments on CIFAR-10, where "horse" is the target class. (a) and (b) present target attack success rates under two attack ratios. Lower success rates suggest higher robustness in the corresponding attack regions.
leading and trailing tokens are corrupted with Gaussian noise. Among all the SSM family methods compared, the performance drops caused by trailing token corruption are significantly larger than those caused by leading token corruption. Notably, for Mamba, perturbing the last 32 out of 1024 tokens results an 81.24% drop classification accuracy, whereas corrupting the first 32 tokens only reduces accuracy by 2.30%. In contrast, the transformer baseline shows relatively smaller impacts from trailing token corruption. Instead, our experiments indicate that more informative features from transformers tend to sink in the leading tokens, aligning with the observations in Xiao al. (2023)
Target Attack. Beyond degrading the performance of SSMs by attacking trailing tokens, we also demonstrate that local bias creates backdoor for target attacks. In this scenario, a target class is selected, and pixel tokens from that class are used to replace those in images from other classes. The attack succeeds when models mis-classify images from other classes as belonging to the target class. Due to the local bias, trailing tokens are expected to be more effective attack region for SSMs, leading to a significantly higher attack success rate compared to leading tokens. Fig. 3 shows the success rate comparisons across different attack regions and ratios. When trailing regions are replaced with pixels from the target class, SSMs achieve much higher success rates than when leading regions attacked. This phenomenon is observed at both 25% and 47% attack ratios. By comparison, the transformer model possesses greater robustness, maintaining similar success rates between attacks on leading and trailing tokens.
Implications for Language Models. While our adversarial attack experiments are conducted on image datasets, the findings have broader implications for language models. System prompts, which are typically group of confidential tokens prepended to user inputs, play critical role in controlling the behavior of language models and preventing undesirable outputs. However, our targeted attack experiments reveal that SSM- based language models are particularly vulnerable to jailbreak attacks (Perez& & Ríbeiro, 2022 Zou al.) 2023). This 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 human feedback to enforce adherence to system prompts may not resolve this vulnerability, as the recency bias remains inherent to SSM models regardless weight configurations.
UNDERSTANDING SCALABILITY BOTTLENECK OF SSMS
4.1 NECESSITY AND LIMITS OF DEPTH SCALING
In Sec. 3.1 we have seen that the dependencies between tokens are exponentially decaying with their relative distances in an SSM layer. Consequently, SSMs resemble localized kernels, similar to those employed in various neural architectures such Convolutional Neural Networks (CNNs) (LeCun et al., 1998) and Graph Neural Networks (GNNs) (Kipf & Welling 2016). It a reasonable postulation that increasing the number of layers can extend the model's receptive field (Goodfellow ett al.) 2016). We justify this hypothesis via scaling-up experiment with various context lengths and model architectures.
We pretrain Mamba using causal language modeling with two context lengths, (2048, 8192). Besides, fix the number of layers at (16, 24, 32, 48, 64, 72) and vary the hidden dimension. We defer

---

Published as a conference paper at ICLR 2025
Context Length = 2048
Context Length = 8192
3.5
Layer=
Layer=24
3.4
Layer=32
Valid. Loss
3.1
Layer=48
3.3
# Layer-64
Layer= 72
3.2
2.9
3.1
3.0
2.7
200
300
400
500
200
300
400
500
Num. Params (M)
Num. Params (M)
Figure 4: We empirically observe that deeper models become increasingly advantageous the context length grows. However, beyond a certain depth, the performance SSMs begins to plateau and eventually declines.
more experiment details in Appendix E.4 The validation loss versus the number of parameters is plotted in Fig. 4 Under the 2048 context length, models of different configurations exhibit similar performance, consistent with the findings of Kaplan et al. (2020). However, the context length increases, the scaling behavior across depth-width configurations begins to diverge. Notably, deeper models outperform shallower ones, likely because deeper architectures can more effectively utilize the extended context to meet the training objectives. Nevertheless, we observe that the performance gain starts to saturate when we keep increasing the depth (ef. the 32-layer and 48-layer models). When the depth of the model continues to increase, the validation perplexity starts to rise, indicating a decline performance (ef. the 64-layer and 72-layer models). In Mamba with 2048 context length, models with more than 48 layers perform worse than 16-layer models. Longer-context models appear to be more tolerant of increased depth, whereas shorter-context models experience a rapid performance degradation once the depth exceeds certain threshold.
4.2 UNVEILING - IN SSMS
To explain the depth scaling bottleneck revealed in the previous section, we conduct a theoretical and empirical investigation of the feature and state dynamics in SSMs. Our key finding is that token embeddings, after being processed by SSM layers, tend become increasingly similar, which leads to a phenomenon commonly referred to as over-smoothing (NT & Machara 2021 | Cai & Wang 2020 Oono& & Suzuki, 2019). Over-smoothing occurs when token representations become indistinguishable, rendering the state uninformative.
First of all, we warm up by studying continuous-time S4 with constant (A, b, c). Recall that a continuous-tim time S4 layer can be described by group of ODEs: h'(t) AAh(t) + br((t) u(t)
c (t). Our analysis starts with the equivalence between convolution and S4 (Gu al. 2021b). This is, the analytic solution to the time-invariant ODE can be expressed as y ] exp A(t s))br(s)ds. Now we analyze the filtering property of this convolution operator from the Fourier domain perspective. We define convolutional operator as a low-pass filter if it suppresses high- frequency components (see Definition D.3). We summarize the main finding in the following proposition, whose formal version and proof provided in Appendix D.2.1
Proposition 4.1 (Low-pass filtering of continuous S4). Consider a continuous-time S4 with pa- rameters (A. b, c). Assume A is diagonal with all values negative. Then y(() = - exp(A(t- s))br(s (s) ys defines a low-pass filter.
Proposition 4.1 states that $4 is inherently a low-pass filter regardless of how (A, b, e) are trained. Therefore, the high-frquunccy components of input signals are being constantly removed at each layer. Presumably. stacking many S4 layers might cause wer-smoothing when all high-frrqqeney components are suppressed to zero.
Now we consider more general scenario when SSMs work on discrete-time regime and (A,, b., c,, A,) are time-varying or even data-dependent. Formally, we prove the following result showing the sharpness of input signals will be reduced as well:
7

---

Published as a conference paper at ICLR 2025
1.0
Pythia
0.9
Sharpness
Sharpness
Sharpness
0.6
0.9
0.8
0.7
0.4
0.8
0.6
0.2
0.7
D:
0.5
Mamba
0.0
0.6
48
48
48
Depth
Depth
Depth
(a) b. and h
(h) Mixer output.
(c) Block output.
Figure 5: Visualization of feature smoothness across layers in pre-trained Mamba and Pythia. The y-axis represents the average pairwise differences among tokens. Mixer outputs (b) solely consider the Mamba or attention module, while Block outputs (c) include all other components (c.g., MI.P).
Theorem 4.2 (Over-smoothing of SSMs). Consider an SSM specified in Eq. with {(A,, by, At) Assume an input space X C R² such that for every x E E. (i) (At)nn + A S 1 for every E [M] and E (T), (ii) min b, S 0 and max b,(xt)) O for every n E [N]. Let Amin = minggrr),neit
(At)n, (A)
Then for any x E X and the memory
states (h; : [T]} generated by the SSM, we have:
max [
1h - helloo S (1- AT-1) E
max [bb((x)) (o(x)))loo
(4)
Proof can be found in Appendix D.2.2 We first justify our assumptions here. (At)a. A, S 1 generic condition to ensure the recurrence of SSMs is non-expansive, which is crucial to guarantee memory states stay numerically stable. The second assumption requires the data to be well-distributed and centered around the origin, which can be casily satisfied by normalization techniques. We find
that prevalent SSM models such as (Gu & Dao 2023 Peng et al. 2023 De et al. 2024 Qin et al. 2024 can easily achieve these two assumptions.
Moreover, if (A.) + A = 1 is always true letting each recurrent update be conservative Peng remove the second assumption well (see Theorem D.5)
al. 2023 Ma al. 2022h. then
Theorem 4.2 examines the relationship between the pairwise distances of memory states and encoded tokens within the sequence. This result indicates that the pairwise discrepancies among memory states are diminished by a factor less than one, suggesting that the memories undergo smoothing following the application of SSM in Eq. [ We deem that if the memory losing its discriminative capacity, the intermediate hidden feature space will similarly collapse.
Delving deeper, the decay rate intricately linked both the context length and the minimal value within {A,, t E [r]}. As the context length increases, it requires more time to effectively mix all tokens. This can be understood from the pressage-passinggg perspective (Gilmer al. 2017): the message of the first token needs to go through the whole sequence to be mixed with the last token. When Amin approaches one, the decay rate maximized, as the entire SSM essentially performs a uniform pooling over the entire sequence, which smoothens the signal via a box-like filter. It is worth noting that the smoothing nature of SSMs is intuitive; one can conceptualize the recurrent operation of SSMs performing running average of the encoded token signals.
Empirical Validation. We adopt a pairwise distance between tokens to quantify the sharpness of signal: {()) uيーニ Σi ||:; xz, (l)) (E, I| l1)) E(x) being small means the token representations are close to each other and become less discriminative.
plot the feature smoothness of a 1.4B Mamba in Fig. 5 In Fig. 5a b, is above h, among all Mamba blocks. This suggests the sharpness of input signals is consistently higher than the sharpness the memory state output from Mamba, verifying Theorem 4.2 In addition, Fig. 5b and show the sharpness of Mamba mixer and Mamba block output, which tends to decrease rapidly in deeper layers. We also provide a comparison with transformer (of the same size). Although transformers suffer from mmoo-smgg in theory (Dong et al. 2021; (Shi et al. [022 Wang et al. 2022), we observe that transformers have a slower decay of feature sharpness. See theoretical comparison in Appendix
8

---

Published as a conference paper at ICLR 2025
5 DISCUSSIONS
In this section, we provide further discussions based on our theory while deferring the remaining parts to Appendix Clue to the page limit. We also introduce more related work in Appendix A
Revisiting HiPPO theory. HiPPO, introduced in (Gu et al. 2020) and extended by Gu et al. (202hb) 2022b). forms the theoretical basis of SSMs. optimally reconstructs signal x up to time by minimizing ||are - y ( with respect to a measure wit) supported on (-o0, t]. The solution projects >5; onto N basis functions, producing a coefficient vector (t), which synthesizes y! via linear combinations. [Gu et al. (2020) showed htt) evolves (t( (tt)/((t - b(t)x(().. where
((t)) = I[O. fl/t yields closed form ((t = Anippo/ tt. Subsequent works like S4 (Gu et al. 2021a)
and Mamba Gu & Dao 2023) utilize this matrix as initialization but omitted 1/t, resulting in
warped measure (s) x exp(s - -t)L[s < t] (Gu et al. 2022b). This measure emphasizes recent history when approximating F.C.. Hence, findings of Gu et al. (2022b) do not contradict our results but also align with Theorem 3.1 The practical implementations often disconnect from HiPPO theory by neglecting the unitary matrices associated with Anippo Whereas, we focus on discrete-domain SSMs, adhering to practical parameterizations. Additionally, Gu et al. (202 demonstrates that SSMs' expressiveness spans all convolutions and RNNs. We contend that the low-pass filtering property also arises from the parameter simplifications of A, (see Proposition 4.1
The effect of selection mechanism. Traditional S4 architectures operate as linear time-invariant systems. To introduce more non-linearity, Mamba (Gu & Dao 2023) proposes modeling (b,, Ct2 A,) as function of inputs, mechanism known as selection. This is motivated by the selective copying synthetic task, wherein A, and b, need to adapt based on content to filter relevant information for memory updates. Despite this adaptation, Theorem 3.1 still holds in scenarios involving the selection mechanism, meaning selective SSMs like Mamba may continue to suffer from recency bias. In the meantime, Theorem 4.2 also applies to Mamba, suggesting that the selective SSMs do not demonstrate higher expressiveness in filtering signals and perform similarly to linear S4 as low-pass filters (Proposition 4.1. Nevertheless, we note that selection can alleviate these issues by adaptively controlling the values in Aç. According to our theory, the selection mechanism can potentially make the upper bound Amaz closer one and the lower bound Amiin closer to zero. However, parameter A in Eq. iis initialized with negative integers (Gu et al.] (2022a). which exacerbates the bound in Theorem 3.1 by accelerating the decay rate of the influence score.
Complex parameterization. Our analysis in the main text primarily focuses on the case where (A,,b, b,,c)) (z) are both realevalued. While most modern SSMs adopt real parameterizations, complex parameterizations particularly complex-valued A, - have been explored in prior works such as Gupta et al. (022)); Goel et al. 2002):(Com Gu et al. (2022a). Notably, we show that both our locality result (Theorem 3.1 and over-smoothing result (Proposition 4.1p remain valid for complex-vale valued SSMs. We formalize these extensions Theorems D.2 and D.4 These findings collectively demonstrate that complex parameterization does not eliminate the locality over-smoothing bias inherent in SSMs.
Importance of context scaling. Theorem 4.2 also highlights the importance of context-length scaling to mitigate the over-smoothing issue. As the smoothing rate decreases with increased context length, lengthening training sequences not only relieves swoothng but also maximizes the utility hardware efficiency of SSMs. This also evidenced by Fig. 4 where models with longer training contexts have better tolerance of deeper architectures. To enable SSMs to fully utilize the context, increasing the model depth essential. Consequently, a synchronized scaling of both model depth and context length is required. Investigating the scaling laws governing these two dimensions presents an intriguing direction for future research.
6 MITIGATING RECENCY AND OVER-SMOOTHING VIA POLARIZATION
In this section, we propose a simple solution to mitigate recency and over-smoothing simultaneously. First of all, results in Theorem 3.1 and Theorem 4.2 can be interpreted in conjunction. To relieve the smoothening rate, one might aim to minimize the values in A, However, this could inadvertently enhance the locality of SSMs, as a decrease in Amaz may occur. A practical implication of this

---

Published as a conference paper at ICLR 2025
relationship is that the values in A, should be as diverse possible to simultaneously mitigate the artifacts of recency and over-smoothing.
Although Amin 0 and Amaz 2 1 could theo- retically occur simultancously, our empirical find- ings show that these values largely concentrated within a narrow range. To illustrate this, we visu- alize the distribution of [Amax - Amin) across dif- ferent channels in Fig. [6] Each bin represents the proportion of channels whose memory state satisfies (Amas Amii)) being smaller than the corresponding threshold on the x-axis, Notably, over 60% of chan- nels have (Amaz Amin) values smaller than 0.5, indicating that most channels cannot simultaneously achieve Amaz & and Amin s O. Memory repre- sentations will inevitably undergo either exponential diminishing or weotshno.
Cumulative Distribution of (Amer - Amial
1.0
Cumulaatie Percentage
0.8
0.6-
0.4
0.2
0.0
0.1
0.4
0.5
0.6
Amea -m
Figure 6: Cumulative histogram (Amaz Aman). The height of each bin represents the cu- mulative proportion of (Amez - Amonn ) less than or equal to the corresponding value on the x-axis.
To this end, we propose to maintain one component in A, constant 1, another others freely learnable. We term this approach as polarization. Polarization ensures that the dimension in the state memory consistently focuses the current token, counteracting over- smoothing by preventing mixing with previous tokens. Simultancously. another dimension polarized one exclusively retains information from past tokens, avoiding locality issues by preserving the complete history. In our implementation (see Appendix E.5) we polarize the first state channel to one (i.e. (At)1.1 1) and the last state channel zero (i.e., (A1)N.N 0).
constant O, and
polarized
We empirically validate our polariza- tion technique through the associative recall tasks, where Mamba models (Gu Dao 2023) are trained over se- quences of key-value pairs to retrieve the associated value according to the query from the context. Please refer to Arora et al.] (2023) and Appendix f.5 more details. If an SSM re-
Configurations
# Layers
Avg.
64
256
Default A, Default A
2
98.38
81.81
36.00
72.06
4
99.23
82.08
33.52
71.61
(A)). -
2 99.81
94.70
56.39
83.63
(A.)N.N =0
2
98.41
81.35
36.55
72.10
[AI)NN
99.74
92.20
52.21
81.38
(A.)..
(A,
2
99.23
95.54
54.74
(Ai). (A(A)N
4
99.94
98.80
81.56
93.43
call information with higher accuracy from a larger number of key-value pairs, then it has better long-context capability. We the performance of Mamba with neither, one of, or both zero- and one one-polarized channels, across different numbers of layers. The empirical results are reported in Tab. 2 (an extended version in Appendix (E.5). The key observation is that the default parameterization of A, suffers from information retrieval from long context, and deepening the architecture even harms the performance (potentially due to over-smoothing). However, once we introduce channel polarized to one, even shallow Mamba could perform high- gech-accuracy associative recall (row 3). We can also gain the performance by deepening Mamba if the over-smoothing issue can be bypassed by adopting the zero pero-polarized channel (row 5). Furthermore, we could apply one- and zero- pero-polarized channels, and in the meanwhile deepening the architecture, we find it achieves the best performance all settings (row 7).
Table 2: Results of polarization. Rows 1-2 have polarization,
rows 3-5 only polarize channel either zero, and rows 6-7 polarize hoth channels.
CONCLUSION
In this study, we identify two critical limitations of SSMs. First, we reveal that SSMs exhibit a pronounced recency bias, which undermines their ability to model long-range dependencies, recall distant information, and maintain robustness. Second, we find that increasing the depth of SSMs induces over-smoothing, rendering token representations indistinguishable and hindering further performance improvements. To address these challenges, we propose polarization technique that re- serves two channels in the state transition matrices, assigning a value of zero and the other value of one. Both theoretical analysis and empirical evaluations demonstrate that polarization significantly enhances the long-range modeling capabilities of SSMs while alleviating over-smoothing.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
PW thanks Bo Liu and Songlin Yang for the insightful discussions when preparing this manuscript. Work was done while JZ and PS were interning at UT Austin. PL is supported by NSF awards IIS IIS-2239665 and IIS IIS-2427777 for this project.
REFERENCES
Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. arXiv preprint MXXv:22003.01500, 2024.
Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. arXiv preprint (rXih:2006.05205. 2020.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. arXiv preprint arXiv:2312.04927, 2023.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended 2024.
long short-term memory. arXiv preprint
Assaf Ben-Kish, Itamar Zimerman, Shady Abu-Hussein, Nadav Cohen, Amir Globerson, Lior Wolf, and Raja Giryes. Decimamba: Exploring the length extrapolation potential of mamba. arXiv preprint arXiv:2406. 14528, 2024.
Yoshua Bengio, Patrice Simard, and Paolo Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions on neural networks, 5(2):157-166, 1994.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Chen Cai and Yusu Wang. A note over-smoothing for graph neural networks. In International Conference on Machine Learning Workshop (ICMLW), 2020.
Kyunghyun Cho. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint 1rXiv:1406.1078, 2014.
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint arXiv: 1409. 1259, 2014.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 147944 2020.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprini arXiv:1901.028660 2019.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.2 2024.
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murard Albert Gu, Ruba Haroun, Leonard Berrada. Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402. 19427, 2024.
11

---

Published as a conference paper at ICLR 2025
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL, 2019.
Francesco Di Giovanni, Lorenzo Giusti, Federico Barbero, Giulia Luise, Pietro Lio, and Michael M Bronstein. On over-squashing in message passing neural networks: The impact of width, depth, and topology. In International Conference on Machine Learning, pp. 7865-7885. PMLR, 2023.
Yihe Dong, Jean-Baptinn Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning (ICML), 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.11999 2020.
Leo Feng, Frederick Tung, Mohamed Osama Ahmed, Yoshua Bengio, and Hossein Hajimirsadegh. Were rnns all we needed? arXiv preprint arXlv:2410.01201, 2024.
Daniel Y Fu, Tri Dao, Khaled Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212. 1405,, 2022.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspec- tive transformers. arXiv preprint arXiv:2312.10794 2023.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1272. PMLR, 2017.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state-space models. In International Conference on Machine Learning, pp. 7616-7633. PMLR, 2022.
Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint arXiv:2312.007522 2023.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances neural information processing systems, 33: 1474-1487, 2020.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:21 1.00096, 2021a.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems, 34:572-585, 2021b.
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems, 35:35971- 1- 35983, 2022a.
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How to train your hippo: State space models with generalized orthogonal basis projections. arXiv preprint arXiv:2206.12037 2022b.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces effective as structured state spaces. Advances in Neural Information Processing Systems, 35:22982-22994, 2022.
Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-171, 2020.
12

---

Published as a conference paper at ICLR 2025
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. /55,: 2022.
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Trans- formers are better than state space models at copying. arXiv preprint arXiv:2402.01032, 2024.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.0682. 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0840 003311, 2020.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. International conference on machine learning, pp. 5156-5165. PMLR, 2020.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv: wet.::6099000000 2016.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-232) 1998.
Jeffrey Li, Alex Fang, Georgios Smyrnis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, et Datacomp-Im: In search of the next generation of training sets for language models. arXiv preprint arXiv:2406.11794, 2024.
Qimai Li, Zhichao Han, and Xiao-Ming Wu. Deeper insights into graph convolutional networks for semi-supervised learning. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018
Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes convolutional models great long sequence modeling? arXiv preprint arXiv: 10.09299, 2022.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Frez Safahi, Shaked Mcirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. Jamba: A hybrid transformer- mamba language model. arXiv preprint arXiv:2403. 19887, 2024.
Bo Liu, Rui Wang, Lemeng Wu, Yihao Feng, Peter Stone, and Qiang Liu. Longhorn: State space models are amortized online learners. arXiv preprint arXiv:2407. /4207, 2024.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint arXiv:2209.10655 2022.
Xuezhe Ma. Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient I1m pretraining and inference with unlimited context length. arXiv preprint (40.::2440.88800 2024.
William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state state-space models. arXiv preprint 240.0:24000088/9 2024.
Hoang NT and Takanori Maehara. Revisiting graph neural networks: All have is low-pass filters. In International Conference on Pattern Recognition (ICPR), 2021.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint ar 2209: 1895, 2022.
13

---

Published as a conference paper at ICLR 2025
Kenta Oono and Taiji Suzuki. Graph neural networks exponentially lose expressive power for node classification. In International Conference on Learning Representations (ICLR), 2019.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kang- wook Lee, and Dimitris Papailiopoulos. Can mamba learn how learn? a comparative study on in-context learning tasks. arXiv preprint urXiv:2402.00 04248, 2024.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. Rwkv: Reinventing rnss for the transformer era. arXiv preprint arXiv:2305. /0448, 2023.
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemystaw Kazienko, et al. Eagle and finch: Rwkv with matrix-valuued states and dynamic recurrence. arXiv preprint arXiv:2404.00 05892, 2024.
Fábio Perez and Ian Ribeiro. Ignore previous prompt: Attack techniques for language models. arXiv preprint arXiv:2211.099
2022.
Michael Poli, Stefano Massaroli, Eric Nguyen. Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference Machine Learning, pp. 28043-28078, PMLR, 2023.
Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björ Deiseroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, al. Mechanistic design and scaling of hybrid architectures. arXiv preprint arXiv:2403.17844, 2024.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. / /2409, 2021.
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong, Hgrn2: Gated linear rmss with state expansion. arXiv preprint 2rXhv:2404.07904, 2024.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. 2018.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. Mach. Learn. Res., 21(140):1-67, 2020.
Han Shi, Jiahui Gao, Hang Xu, Xiaodan Liang, Zhenguo Li, Lingpeng Kong, Stephen Lee, and James T Kwok. Revisiting over weo-thng in bert from the perspective of graph. arXiv preprint arXiv:2202.0862.. 2022.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputiig, 568:1278000 127063, 2024.
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al. Learning to (learn at test time): Rnns with expressive hidden states. arXiv preprint arXiv:2407.064 04620, 2024.
Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, and Furu Wei. A length-extrapplltable transformer. arXiv preprint arXiv:2212.10554, 2022.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621, 2023.
Ilya Sutskever, Oriol Vinyals, and Quoe Le. Sequence to sequence learning with neural networks. Advances in neural information processing systems, 27, 2014.
14

---

Published as a conference paper at ICLR 2025
Yi Tay, Mostafa Dehghami, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers. arXiv preprint arXiv:201..04000 2020.
Jake Topping, Francesco Di Giovanni, Benjamin Paul Chamberlain, Xiaowen Dong, and Michael M Bronstein. Understanding over-squashing and bottlenecks on graphs via curvature. arXiv preprint 11X1::2111.14522, 2021.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Moreney, and Ruslan Salakhut- dinov. Transformer dissection: a unified understanding transformer's 's attention via the lens of kernel. arXiv preprint arXiv: 1908. / 1775, 2019.
Jos Van Der Westhuizen and Joan Lasenby. The unreasonable effectiveness of the forget gate. arXiv preprint Xrviv: 0004. 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz
Kaiser, and Illia Polosukhin. Attention All You Need. In Proceedings of NeurIPS, 2017.
Aaron Voelker, Ivana Kajié, and Chris Eliasmith. Legendre memory units: Continuous-tima repre- sentation in recurrent neural networks. Advances in neural information processing systems, 32, 2019.
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, et al. An empirical study of mamba- based language models. arXiv preprint arXiv:2406. 07887, 2024.
Peihao Wang, Wenqing Zheng, Tianlong Chen, and Zhangyang Wang. Anti-oversmoothing in deep vision transformers via the fourier domain analysis: From theory to practice. arXiv preprint 20X1v:2203.05962 2022.
Shida Wang and Beichen Xue. State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. Advances in Neural Information Processing Systems, 36, 2024.
Xinyi Wu, Zhengdao Chen, William Wang, and Ali Jadbabaie, A non-asymptotic analysis of oversmoothing in graph neural networks. arXiv preprint arXiv:2212. 10701, 2022.
Xinyi Wu, Amir Ajorlou, Yifei Wang, Stefanie Jegelka, and Ali Jadbabaie. On the role of attention masks and layemorm in transformers. arXiv preprint arXiv:2405. 187811 2024a.
Xinyi Wu, Amir Ajorlou, Zihui Wu, and Ali Jadbabaie. Demystifying oversmoothing in attention- based graph neural networks. Advances in Neural Information Processing Systems, 36, 2024b.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprini arXiv:2309. / 17453, 2023.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efffcient training. arXiv preprint arXiv:2312.06633 2023.
Songlin Yang. Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta nule sequence length. arXiv preprint XXXv:24066064844 2024.
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry. arXiv preprint arXiv:2402.04347, 2024a.
Tao Zhang, Xiangtai Li, Haobo Yuan, Shunping Ji, and Shuicheng Yan. Point could mamba: Point cloud learning via state space model. arXiv preprint arXiv:2403.00762 2024b.
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint mrXiv:2401.09411 2024.
Andy Zou, Zifan Wang, J Zico Kolter, and Matt Fredrikson. Universal and transferable adversarial attacks on aligned language models. arXiv preprint arXiv:2307. /5043, 2023.
15

---

Published as a conference paper at ICLR 2025
A OTHER RELATED WORK
Despite the empirical success of SSMs in various long-range applications (Lieber et al. 2024 Zhu et al. 2024 Zhang et al. 2024b), their theoretical properties and limitations remain underexplored. Arora et al. (2033) leverages associative recall tasks to theoretically analyze the expressiveness of convolutional SSMs, while advocating for input-dependent kernels. Jelassi et al. (2024) separates the representation capacity of transformers and SSMs via coping tasks. Recent work by Merrill et al. (2024) identifies failure modes of SSMs in state-tracking problems through circuit complexity theory. Hahn (2020) characterize the expressivity of SSMs using linear controlled differential equations, and et al. (2024) reveal structural similarities between selective SSMs and attention mechanisms. While these works provide valuable insights consistent with our findings, none directly examine the long-range modeling capability of SSMs. Ben-Kish et al. (024) points out that the product of gating matrices (formalized in our Lemma D.I exhibits locality issues, potentially hindering their ability to model long-range dependencies. However, their analysis lacks a formal justification.
The most relevant prior work to our study recency bias is perhaps Wang & Xue (2024), in which Their Theorem 3.13 reveals exponentially decaying memory for SSMs. Distinct from their findings, our primary contribution lies in analyzing nonlinearity and input-dependent mechanisms widely adopted in modern SSMs (Gu & Dao 2023] Yang et al., 2023; Arora et al. 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 (2024). In particular, [Wang & Xue (2024) only addresses non-linear activations after performing linear SSMs (i.e.. S4), with the state transition matrix A, being independent of inputs. Our analysis, however, considers both A, and b, as input-dependent, aligning with settings in Mamba and other follow-up works. The approach Wang & Xue (2024) does not naturally extend to this case, as they assume the linearity of the sequence mixer (cf. Eq. 9 in Wang & Xue (2024)). Furthermore, we conduct a finer-gran aner-grained analysis, quantitatively relating the decay rate to the specific values within the input-dependent gating matrix A. To our best knowledge, none prior work revealed the smoothening nature of SSM operators.
B UNIFIED FORMULATION OF SSMS
Linear Attention Models (LAMs) represent a vast family of architectures evolving rapidly. We present reformulation of several representative LAMs using the SSM framework described in Eq. [1] This reformulation is expected to extend to models not explicitly covered below, including but not limited
to Megalodon (Ma et al., 2022 2024). Hyena (Poli et al., 2023), HGRN2 (Qin et al. 2024). TTT
(Sun et a 2024), Longhorn [Liu et al. 2024), xL.TM (Beck et al. 2024). and Smin.STMs//min
(Feng et al. 2024).
Linear Attention. Attention notoriously suffers from quadratic computation and memory complex- ities. To address this limitation, Tsai et al. (2019); Katharopoulos et al. (2020) finds that if we express self-attention as a linear dot-product of kernel feature maps, then we can interchange the order of matrix products to achieve linear time complexity. This finding leads to a series of LAMs. Ignoring the denominator therein, can formalize linear attention via the recurrence in Eq. [
(LA) II bı(z)) k(xz), ch(h)) q(xx)" h, A v(zz),
where k. IR + RN are the kernel feature maps and v R + R transforms the input features. Considering multi-channel inputs, k. q. and v become functions processing vectors, and for each channel, the and q are shared, while A, being specified individually. The roles of k. q. are similar to the key, query. and value transformations in standard transformer (Vaswani et al. 2017). While kt, q, e are often chosen as linear mappings, many other options exist. Katharopoulos et al. (2020) adds E+LL after linear mapping, [Chormmanki et al. (2020) leverages orthogonal random Fourier features with hyperbolic cosine as activations, and Zhang et al. (2024a) finds MLP with exponential activation effective.
Retentive Networks (RetNet). RetNet Sun et al. 2023 is another variant of LAMs, proposed as a successor to transformers given its remarkable performance. Each layer of RetNet consists of key,
16

---

Published as a conference paper at ICLR 2025
query, and value transformation, akin to linear attention. In addition, it imposes new decaying term over the past states. We find RetNet can be summarized via our formulation in Eq.
(RetNet)
A -I,
(xx)) = k<(x),, (e( = q(222) he, A = o
where 7 E (0, 1) is (learnable) scalar, k. IR + R. R - R are linear functions, Similar to Mamba (Gu & Dao 2023). RetNet shares b, and e across channels while assigning distinct A, for each channel when handling multi-channel inputs.
Gated Linear Attention (GLA), GLA Yang et al. 2023) introduces gating mechanism, originally from RNNs (Van Der Westhuizen & Lasenby 2018). to LAMs. Its computational mechanism can be encompassed by our formulation in Eq.
(GLA) A, = diag(a(z:)), b((x)) = k(x,), c(hh)) = q(x,) he, A, = v(x)),
where R- + 10, converts input to gating logits, k q IR R RNN R R are linear key, query, value mappings. When the inputs are multi-dimensional we share a across channels while assigning each channel with separate k. q, and extending their input dimension accordingly.
Linear attention (Katharopoulos et al. 2020) can be regarded as GLA with constant A,, while RetNet (Sun et al. 2023) can be formulate aanu as GLA with input-independent A, (Liu et al. 2024).
RWKV. RWKV a series of models that linearize attention computation (Peng et al. 2023 2024). We focus RWKV-4 (Peng al. 2023) and demonstrate it can also be reformulated into the structure of SSMs:
I
A, exp( w) + exp(A((x,))
(x,,) v(xx),
(RWKV)
exp(k(xxt)
c+(h) = q(xt)) h1, A. exp(-w) + cxp(k(x,))
where w E IR., is a learnable coefficient, k: RR + R and q, v R + EN are linear mappings. If the inputs are vector-valued, becomes vector while q, u turns into functions that take vectors as inputs. One important property of RWKV is that (At),, + A. = 1. Later in RWKV-5 and RWKV-6 (Peng et al. 2024), the normalizer is removed for numerical stability.
Griffin.
The recurrent unit in Griffin (De et al. 2024) can be re- re-formulated kind of SSMs:
(Griffin)
A, = diag (a(x;)),
b, (xx)) = diag (i(x,)), (((he) hus, A. = diag (VI [V==)(=)))) a(x))²)
where ((xx)) = sigmoid (W a.; + b.) is an input gate, a is computed in log-space: log (<(x;)) = -E (T)) @ sigmoid (W, + b.), is Hadamard product, & is a constant, and r, W.. b.. W.. b, are learnable parameters. In particular, the dimension of h, in Griffin is equal to the dimension of xt. we consider single-channel Ft, then A,, b,, and A, are all scalar-valued.
C DEFERRED DISCUSSIONS
Extended discussion with HiPPO theory. HiPPO established in (Gu al.] 2020). extended by Gu et al.) 1b 2022b) is the theoretical foundation of SSMs. Consider a signal and its reconstruction y( up to time f. To optimally memorizes the history of x using ( HiPPO minimizes || (a))))) w.r., a measure (t) supported on (-o0, fJ. The solution is to project the history of before time i onto N basis functions (e.g. Legendre polynomials), which yields time-continuous coefficient vector htt),, and (t) can be synthesized by linearly combining the N basis using h().. Gu et al. (002)) shows that the evolution of (t) follows an ODE h²(t) ()) (L) + b(t)x((t.. In
particular, Gu et al. (2020) chooses a uniform measure over the past history (1) 110 tt/t, which places approximation bias over the time horizon in contrast to earlier work (Voelker et al. 2019).
As a result, Att) in Eq. can be written in a closed form: A(t) = - Ansppo /t, where Anappo is a time- independent constant called the HiPPO matrix. Its various forms have been used initialization in subsequent works including S4 (Gu et al. 2021a) and Mamba (Gu & Dao (023). While HiPPO theory seems to guarantee the long-rangeness for SSMs, the actual form of A(t) employed in S4 and
17

---

Published as a conference paper at ICLR 2025
Mamba drops the normalizer 1/t. Gu et al. (2022b) shows that this change causes warp of measure from uniform to (s) x exp(s t)I((Cc) [o,, t]. We note that this warped measure assigns more importance to recent history, and thus, our Th. does not contradict HiPPO theory but also matches the findings in Gu et al. (2022b). We also point out that when adopting the diagonalized form of Anippo (Gu et al. 2021a] Gupta et al. 2022; Gu et al., 2022a), the unitary matrices decomposed from Atippo is sometimes not applied to b, and C₁, which introduces disconnect between HiPPO theory and its practical implementation. Our paper directly studies the discrete-domain SSMs and aligns with the parameterization used in practice. Another less-discussed property of SSMs is their approximation power for a broad family of operators. In Gu et al. (0211)) SSMs are shown to possess expressiveness that encompasses both convolutions and RNNs. It is worth noting that the original HiPPO-based SSM is not necessarily a low-pass filter. Rather, it is the successive simplifications in the parameterization of A, that impart the smoothing characteristics of SSMs (see Proposition 4.1)
Does hungry hungry hippos help? The key innovation of Hungry Hungry Hippos (H3) (Fu et al. 2022) lies in the introduction of self-gating connections and locally shifting convolutions to improve in-context recall for state space models (SSMs). This design has quickly become standard backbone for various SSMs (Gu & Dao 2023 Beck et al. 2024). However, we question its effectiveness in addressing the local rangeness issue in SSMs. The gating mechanism operates at the token level, which impacts the bound in Theorem 3.1 only by a constant factor. Additionally, the introduced convolutions typically use small kernels, which are insufficient to mitigate the exponentially decaying relevance between tokens. As we empirically show in Fig. [ while Mamba with H3 performs adequately in associative recall tasks when the state size is sufficiently large, a locality bias begins to emerge as the number of key-value pairs exceeds the model's memory capacity. This highlights the limitations of the architecture in handling long-range dependencies under constrained memory.
Does gradient vanish in SSMs? Vanishing gradients refer to a challenge in RNNs, where backpropagation-basees learning is impeded due to gradient magnitudes decaying exponentially over time. The diminishing dependencies among distant tokens are a fundamental cause of this expeliiiiti modeling long-range dependencies, as highlighted in (Gu et al. 2020; (0216). 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 33] lies in theoretically challenging this assumption. 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.
issue (Bengio et al., 1994). SSMs were initially proposed address this limitation by explicitly
Connection with over-squashing theory in GNNs. We can compare recency of SSMs to over- squashing in GNNs. The influential score defined in Sec. 3. is also used for over-squashing analysis Graph Neural Networks (GNNs) identify information bottleneck (Topping et al. 2021; Di Giovanni et al.) (023). The sensitivity analysis in Topping et al. (2021) demonstrates similar exponentially decaying dependencies among graph nodes, dependent on the underlying graph topology. We postulate that propagating information from long distances remains challenging for SSMs because the model needs to encapsulate all history information into fixed- dimen-dimnnsoo hidden vector, which is also observed as one major problem with RNNs (Bengio et al. 1994 Alon & Yahav
2020 Sutskever et al. 2014 Cho et al. 2014: Cho 2014).
Connection with over-smoothing theory in GNNs and transformers. Over-smoothing issues were first identified in GNNs (Li et al., 2018, NT & Maehara 2021 Oono & Suzuki] 2019 Cai & >Wanng 2020 Wu et al. 2022 2024b) and later explored in transformers (Dong et al. 2021 I Wang et al.) 2022 Shi et al.) 2022 Wu et al. 2024a, Geshkovski et al. 2023). In both cases, over-smoothing manifests as feature representations becoming increasingly uniform with greater model depth. To the best of our knowledge, our work is the first to uncover this phenomenon in SSMs. In GNNs, the smor-smog effect typically follows a linear decay rate governed by the second-largest eigenvalue of the adjacency matrix: O(2)), ), where I. denotes the number of layers (Oono & Suzuki 2019] Cai & Wang, 2020). More closely related to our setting, Wu et al.] (2024a) analyzes the convergence rate of feature smoothness in causal attention, demonstrating a rate of O((1 - AT: )./T), where Amin the minimal value in attention maps. In comparison, our result in Theorem 4.2 shows rate of O((1 Aprn) indicating that the smoothening speed of SSMs is
18

---

Published as a conference paper at ICLR 2025
faster than that of transformers. Our Fig. supports this claim, the transformer- architecture exhibits mild decrease of sharpness at deeper layers, whereas SSMs show à constantly steeper decay slope. Moreover, due to the inherent locality of SSMs (Theorem B.I. 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. This distinction is partially supported by the observation that modern SSMs often adopt architectures that are roughly twice as deep as transformers. Consequently, depth-scaling limitations are more critical for SSMs than for transformers.
D PROOFS
D.I EXPONENTIALLY DRCAYING DEPENDENCY WITH RELATIVE DISTANCE
In this section, we extend and prove Theorem 3.1 more general case where inputs and parameters are all complex-valued. Below we by default assume that A, E C XN b, C + CN C.. CC + C, and A IR. First of all, we present the following auxiliary lemma, reformulating SSM recurrence in an explicit parallel form.
Lemma D.1 (Parallel form). For any {(Ai, be, Ct. Ail and æ E CT, CT computed via an SSM defined Eq. iis equal to:
h4
Σ s=1
t-1
A, A. r=
II
) A.b,(x)), y: = c((a), e [T].
(5)
Proof. Proof by induction. First let examine the base case when = 1: h1 A ( indicating that Eq, hholls trivially. Now we make inductive hypothesis that Eq. hholds for some t e T 1]. Then for time step t E [T]:
h1+1 Actih + De+ di+ (x++1)
1-1
= A++1
Σ
A,
A,b,(x,)+ + (441b41(0001)
1r=s11
A4t1
II
A
A.b.(x.) (rb1 + Ae-ibe
r=st1
[+1
[+1
Ap
A.D.(T.))
+
II
A, (ab)( + (
yrs+1
r=1+1
[+1
A
A.b.(x.)) +
r=s+1
which also satisfies Eq.
Then we conclude the proof by induction.
remark of Lemma D.I Lemma D.T provides alternative perspective how SSMs compute the outputs. The predicted value for the E-th token is obtained via decoding a weighted aggregation over representations of all past tokens. The encoding and decoding stage element-wise independent of the context. Whereas, the "weight" associated with each past token in the summation reflects the pairwise relationship, playing similar role to attention weights transformers (Dao Gu 2024] Ali et al., (2024). The weight corresponding to one past token is calculated as the cumulative product Ar, where r E [s + 1, traverses from the past token (at time s) to the target token (at time f). Assume A, E 1) NxN, which satisfied by most of SSMs discussed in Sec. 2 we can show that (II Ar)n, < (IT==0'41 Ar)m. for any s < s' < and n E [N]. By this interpretation, SSMs assign strictly higher "attention" to the nearer tokens than the further tokens.
Below we formally state and prove the complex version of Theorem 3.1 Theorem 3.1 is a straightfor- ward corollary of Theorem D.2
19

---

Published as a conference paper at ICLR 2025
Theorem D.2 (Recency of SSMs). Consider SSM defined in Eq. with {(Ac. br, ca, Am)
Assume that:
(i) The imput space X C Cf is compact.
(i) {(At, bi, Ct, SE)TET(T) and {
[AA,
ab. 8an: Bay
02. 8my Dm,
Jeeryy are continuous.
(iii) A₁ is diagonal and ((A))n < 1 for all [
Let Amax maxx (At)mol
Then for arbitrary X and every s, E [T] such that
8 < t
dy: ax,
< C-ep(-x((1 -
for some constant C > O independent of and 8, and log( (A (Amie)
Proof of Theorem 3.7
First of all, we note that by compactness of input space (Assumption ( and
continuity (Assumption (ii),, (Ai, bi, Ct2 As) and aA..
0b,
dc.
0d:
are all bounded by some 0w:
Der
constant for every / [T]-
By Lemma
D.I
Eq.
Tean can expressed the closed form as follows:
h4-
II A,
(x;;(x;) + (x+)
exp
log Ar
A,bi(a;;) + A,br(w.)
1=i+1
exp
log Ar
A,bi(z;)) +
exp Σ
log Ar
(ritt1
where the second equality is by simply rewriting the cumulative product as an exponential of cumulative summation of logarithms. The logarithmic terms are well defined and constantly negative due to Assumption (ii) We decompose h, into three components, where the first term, denoted as ths depends on V, only through (A,) the second term, denoted as vt, only relies on æ, via (A,b: while the remaining part is independent of æ..
Now we tackle each component separately. First, we simplify du./ 0ax as:
Ou 0x
Σ exp
log Ar A(b((zz))
Ox
1
Σ log Ar )) Aib((as))
exp
dx.
|mi++1
Öogg A,
cxp
log Ar
(,b,(z;)
dx
v=++1
r=i+1
aA
exp
log Ar
A-
A,b;(æ;)
dx.
r=
20

---

Published as a conference paper at ICLR 2025
where the last steps follow from the basic chain rule. Then bound its l norm by:
Oue
s-1
0A,)
exp
log A.
A, Ox. Abb;(
(6)
1
mmm
78,22
N
C1
exp
log/ (Ar)n,o
(7)
n=1
1=1
r=i+1
CNN
exp (log Amaz(t i))
(8)
C1N 1 exp (- log Amaz(s 1))
exp (log Amav (t 1))
Amaz
CNN exp (log Amaz (t s)) exp (log Amau (t 1))
Amer
C1N
Amirn - exp (- log Amaz (t (t-s)) s))
(10)
We elaborate on the derivation step by step. Eq. is due to triangular inequality. To obtain Eq. 7 we note that 1..'0A,// A; b, are all wiseeyy-wisely bounded, thus, we can extract their uniform upper bound C1 > 0 out of the summation Eq. can be derived by applying the supremum Amoz Eq. follows from the summation of geometric series. And finally, inequality in Eq. holdd by dropping negative term.
over all {I(At) n, wll.err..eee
Similarly, we rewrite 0v,/0x, as below:
Dv: δ
Σ exp
Σ log Ar (z.) Ami+tl
Fx.
dE
O(A O(A,b,(a;))
cxp
log A,
Ox
i=s
Ar=it1
exp
log Ap
Oxc
by which we can yield the following upper bound on its b, norm:
Ov
O(A b, (x. mm)
exp
log Ar
dx.
dx.
n=1
r=st1
n,m
N
< C2
Σ
exp
Σ \r=++1
logl
(Ar)n,l
(11)
< NC exp (- A (t s))
(12)
where Eq. 11 obtained by applying uniform upper bound C2 over (((A.b.
Eq.
12 is induced by leveraging the supremum Amar over all {I(A,) m. bhetrr.mem
Combining Eqs. T0 and
12
have;
0h:
CiN
S
+ NC2
exp (-log log A (1 s))
(13)
dzs
Amaz
Finally, we conclude the proof based on Eq. 13
Oy:
(h.)
ah
dc, (h:))
Oh+
<
Ox.
Sh.
Ox.
Sh,
Oo.
CNN
O's
+ NC2 exp (- log A A (t 8))
Amiz -
C Ceepp (-i(t s))
21

---

Published as a conference paper at ICLR 2025
where Hölder inequality in the first equation and upper bound
again via
constant C3 > due to continuity (ce (ft) is a composition series continuous functions as in
Assumption This is desired after letting C' > 0 absorb all constants and x log A M
D.2 OVER-SMOOTHING IN SSMS
D.2.1 LOW-PASS FILTERING PROPERTIES
In this section, we formally state and prove Proposition 4.1 To begin with, we define low-pass filters as below:
Definition D.3 (Low-pass Filter). Suppose 2t)) : + C is absolutely integrable. Then the Fourier
transform (w) + C exists and (ww) - x(1) exp(-iw()dl. We say z(1) is an (e, (YyS.l-lossssss filter for some 0 < < and Ω > if [w)] for every |w| 2 Ω. Furthermore, we say z(t) is low-pass filter if for every < E < 1, there exists > such that z(t) is an (e, (e, filter.
Note that definition of low-pass filters focuses its effect on removing high-frequency compo- nents. By Definition D.3] it is equivalent to say system performs low-pass filtering if and only if the responses converge to zero as the frequency increases. Fixing e small enough, we can say (c, (c)-looosssss filter has cut-off frequency at W = Q.
Now we present formal version of Proposition 4.1 as below, which quantifies the cut-off frequency of a filter induced by continuous time S4. Note that Proposition D.4 holds for complex complex-valued valued (A, b, c). Below we assume A E CN
CN c E CN by default.
Proposition D.4 (Formal version of Proposition 4.1). Consider continuous S4 with parameters (A, b, c). Assume A E CN is diagonal and its diagonal values have negative real parts. Then (t) exp(A(t ()) (oo(ssdss ds is an (e, (o.eo(((-)-lo-pppll filter for any sufficiently small.
Proof. First of all, let rewrite the filter induced by S4 as below:
N
2tt)
Σ
b, exp(An,,t)
==1
Then we apply Fourier transform =(t):
7(w)= / =(1)e hot dl
/ Σ cnbn (n, exp((ixil)
Σ bn / exp(An,nt) cxatt)xxpppiic -inttittdd
Σ bn / eep((AAnn wi))tdd
N
n=1
N
C
bn
iw
An,
n=1
where the last integral converges due to R(A,,)) < 0. Then we can upper bound the magnitude of Z (w)) as:
[Z(w]
Σ
N
C
bn
S
N
C₂
b
N
c,
bn
An,
iw
An,m
[
An.a
We consider small enough, thus, is sufficient to consider the scenario when W > maX |m In this regime, we have:
Amaz
1
|||||||ا|||
[Z(w)] <
Σ |anball
[|
Amax
|W
Amaz
22

---

Published as a conference paper at ICLR 2025
where the last two inequalities are due to Hölder's inequality. It is sufficient to set:
M |||2|||||z
Amars
to let [x(w] S for every |س| 2 Ω hold.
D.2.2 GENERALIZED SMOOTHENING PROPERTIES
In this section, we define generalized smoothening operators as those that narrow distances between token features. We provide the following results that formalize this point. Theorem below extends the result of Theorem 4.2 to include the case when (A,)n.. + A. 1 for every n E [M], E [T],
which is more aligned with Ma al. (2022) 2024); Peng et al. (2023) in practice.
Theorem D.5 (Complete version of Theorem 4.2)
Consider SSM specified in Eq. D with
(A,, b., Ct: Lelr) Assume an input space IR. such that for every a E X, either of the following conditions is satisfied:
(i) (At)n, + A forr every Ti E [N] and t [T]. or
(ii) (At)n, + A. & 1 for every n E [N],t t E [T], and minte(T] be & 0 and
maxtelI] by (n 2 0 for every n E [N].
Let Amin minie/TT.nell (At),, Then for any x e X and the memory states (he [T]
generated by the SSM, we have:
max
||; 1allo0 S (1- 10-AT-11 12-1) max
((b((x))) bx (r(x)))llcc (.
[
Proof. To simplify the proof, first only consider the dynamics of channel in the memory state. We denote og (At)n... 2t = and 8 = he. for some n. E [N]. Additionally, we define the following two quantities:
m
min Zt, <E[T]
M max 2t =E[T]
When Assumption iii holds, we know that S 0 and M 2 O. Suppose = pm + (1 P)M for some p [0, 1]. Furthermore, let q = 111ppp. p. Now we consider the following dynamics with 8g =
8; = 04.8-1 1 + Aize, [T).
(14)
Next, we can show the result below (proved later):
Lemma D.6. We have the following inequality for Eq. 14
A g)m + A qM S St S pm (1 A P) Vt [T],
if either of the following conditions holds:
(i) 0: + A. = 1:
(ii) ae + A. and m < 0 M.
Note that the two conditions in Lemma D.6 corresponss to Assumptions i1) and [i) Theorem D.5 respectively. Thus
upper and lower bound the minimum and maximum of memory states
{8t :e:: E [7]} through:
m' A min s; min At 9)m + Ar
[E[T]
LE(2]
= (1 - T-1 4 g)m AT-9M,
as we are moving the convex combination towards the smaller end (i.e., 1 AT-1 > - 9), and similarly.
M' A max < max A 1-1
man pm + (1 A min P)P
[E[T]
[4[T]
< AT pm (1 1T-11) P)
23

---

Published as a conference paper at ICLR 2025
now relaxing the convex combination towards the larger end (i.e., 1 1 -- p Henceforth, we can upper bound:
A
M' m' Σ -1 pm + (1 - Amm - (1 (1-Ammr
18=1 9)m Amii 9M
= (1 ) 10mmmm))) m))
using the fact that p + =-
Now we can apply the above result to all memory channels. Assigning each channel with my min b, (x))) Mp = max be(x))) and Mn m/ accordingly, where n E [N]. We can yield:
max |/, hell20 - max (M" - -m²) [.5E[T]
NE(N]
max (1 mmn) (Mn m.n)
me [N]
A T-1)
max (Mn m.)
nE(N]
A TT-1)
min
NE(N]
max
(E|TT
max (2( min (xx))
tE[T]
T-1) A
min
max
max (b, (xx)))
nE[N
< (1
A T-11
min
max [.
[lba(ax)) - b.(2.))
where we exchange the two maximum in the last step.
Finally, we prove an auxiliary result we used before. This result generalizes Lemma 3 in Wu et al. (2024a) in the scenario when 0₁ + A. < and m S 0 S M.
Proof of Lemma D.6 We first prove the side that is "less than". Note that the desired inequality trivially holds for the base case 81. Then conduct the following inductive steps. Suppose the desired inequality holds for some r E [T - 1}, namely:
8y < A Amoppm + (1 - Am.iP)M
Then we show that for time step + 1 E [T).
S Ar418, Ar412r+1
S Ar41 (A mmmp (1 AmTuP) M) + Ar+1M
(15)
S A (Amrhym + (1 AmoiP) M) + (1 Ar+1)
(16)
S Amin (Amzpmm- + +11 A Amih)) M) + (1 - Amin M
(17)
Am pm + (1 A,
where substitute the upper bounds of 8t (by the inductive hypothesis) and in Eq. 15 Eq. 16
holds because either A Ar+1=
Ar+1 (by Assumption (i) or Ar+1 < 1 - Ar+1 and M > (by
Assumption (ii) Eq. I7 is satisfied we are lowering down the coefficient for the smaller end the
convex combination. This concludes the proof by induction.
The proof for the "greater than" side follows from symmetric argument. We provide a brief proof for completeness. The base case $1 trivially satisfies the desired inequality. Consider an inductive step where s, for some [T 1] also satisfies the desired inequality. Then we have:
1 Ar418, + Ar412r41
> A,41 ((1 A 1 + A A + Ar+1m
2 Ar41 ((1 Am qm A M) + (1 AA41 )
2 Amin ((1 - A qm + + (M) + +(1 -Ami m
= (1 4m + Ammng.
which concludes the proof by induction.
24

---

Published as a conference paper at ICLR 2025
0.40
0.12
Mamba-130M
Mamba-1.4B
0.10
Pythia-160M
Mamba Transformer
Pythia-1.4B
0.08
0.25
Score
150M
0.20
0.06
1.4B
0.15
0.04
0.02
0.00
0.00
400
600
800
1000
40
80
120
Relative token position
Relative token position
(a) Influential scores.
(b) Influential score ratios.
Figure 7: Visualization of influential scores for Mamba and Transformer models under 150M and 11.4B parameters. (a) illustrates the influence score on the last output token by the preceding input tokens. (b) shows Mamba-to-trrasffrmmer ratios of influential scores within the nearest 128 tokens.
E EXTENDED EXPERIMENTS
E. INFLUENTIAL SCORES: TRANSFORMERS VS. MAMBA
We further compare how influence scores vary with the relative distance between Transformers and Mamba. For our study, we select Pythia (Biderman al. 2023), which integrates Rotary Position Encoding (RoPE) (Su et al. 2024) into the attention computation. Fig. 77 illustrates the influence scores of each token contributing the last token. Pythia demonstrates the well-known "lost in the middle" pattern, where tokens at the beginning and end have high influence scores, while those in the middle exhibit lower scores. contrast, Mamba's highly influential tokens are more concentrated in the nearest positions.
We also compare the decay rate of influence scores for these two models at local tokens. Each curve shown in Figure (aaure7/reeeee represeess the ratio of Mamba's influence scores those of Pythia of various model sizes. The dotted curves are fitted with an exponential function. Our findings reveal that Mamba exhibits exponentially faster decay speed compared to Pythia. This suggests that Mamba places greater emphasis on local tokens than transformers.
E.2 NEEDLE IN A HAYSTACK TEST
Our experiments of testing positional bias for Mamba in Sec. 3.2 are based on an open-source project LLMTest_Neede InAHaystack To validate the retrieval capability of the models while preventing them from relying on memorized information stored their model weights, carefully design the inserted statements to contain factual errors. Several examples of such statements are provided in Figure For instance, we insert the statement, "The capital of France is Madrid." and then test the model's retrieval ability by asking the question, "What is the capital of France?" While the correct answer, Paris, likely memorized by the I.LM, if the model "correctly" outputs Madrid based the context provided, it demonstrates that the model is successfully using the contextual information rather than relying on pre-existing knowledge. This approach ensures that the evaluation focuses the model's ability to retrieve and process information from the input context. We also add an instruction "ignore the fact, only use information in the context to answer the questions" to facilitate this behavior. We provide the fine- fine-grained visualization result in Figure
E.3 CIFAR-10 10 IMAGE CLASSIFICATION
Here we present experiment details in Sec. 3.3 where we conduct image classification on the CIFAR-I0 dataset study locality bias in SSMs. Specifically, 32 32 RGB images in the dataset
httts::/githhbe httpe/////tttooooooo//e/ee/eeeee hww.eeeeeeeellaaaeeeaaaa InHHasstac
25

---

Published as a conference paper at ICLR 2025
Inserted Statement: The cupital France Madrid.
Question: What the capital France?
Inserted Statement: The telephone was invented by Elon Musk.
Question: Who invented the telephone?
Inserted Statement: The largest planct our solar system is Earth.
Question: What the largest planet our solar system?
Inserted Statement: The country known
Queetion: What country known for Great Wall?
Figure 8: An illustration of our synthetic data.
10240 0.54
0.54
0.55 0.56 0.56 0.53
0.51 0.51 0.52 0.5 0.48 0.49 0.52 0.54 0.58 0.66 0.74 0.8 0.84
0.8
0.41
0:42
0.46
0.7
14336
0.43
0.43
0.44
0.45
0.44
0.41
0.4 0.4
0 0.42 0.45 0.45 0.43 0.4 0.39 0.42
0.49 0.57
0.66
0.69
Lengin
0.38
0:34
0.6
18432
0.35
0.35
0.37
0.37
0.36
0.34
0.34
0.33
0.32 0.31
0.31
0.32 0.33
0.35 0.36
0.38
0.43
0.52
0.6
0.5
Full
20480
D.27
0.32
0.3
0.29
0.28
0.26
0.28
0.29
0.32 0.35
0.57
0.4
22528
0.27
0.29
0.31
0.31
0.28
0.28
0.3
0.28 0.26 0.26
0.28 0.29
0.29
0.29
0.3
0.32
0.33
0.35
0.37
24576
0.29
0.3
0.31
0.3
0.29
0.28
0.27
0.26 0.26 0.26 0.26 0.28 0.29
0.29
0.3 0.33
0.37
0.39
0.3
20%
Needle Position
10240
0.73
0.73
0.73
0.73
0.73
0.73
0.73 0.73 0.73 0.73 0.73 0.73 0.73
0.73
0.73
0.73
0.73 0.73 0.73
0.7
Length
14336
0.71
0.71
0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71
0.71 0.71
16384
-0.99
0.69
0.69 0.69
0.69 0.69 0.69 0.69 0.69 0.69 0.69 0.69 0.69
0.69
0.6
18432
0.71
0.71
0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71 0.71
Full
0.5
20480
0.72
0.72
0.4
22528
0.76
0.76
0.76
0.76
0.76
0.76 0.76 0.76 0.76 0.76 0.76 0.76 0.76 0.76 0.76 0.76 0.76 0.76 0.76
24576
0.74
0.74
0.74 0.74 0.74 0.74 0.74 0.74 0.74 0.74 0.74 0.74 0.74 0.74 0.74 0.74 0.74 0.74 0.74
0.3
5%
10%
25%
20%
25%
30%
35%
40%
45%
50%
55%
60%
65%
70%
73%
80%
85%
90%
95%
Needle Position
Figure 9: Detailed comparison between SSM and Transformer the "Needle in Haystack" benchmark. The upper figure shows the retrieval accuracy the Cauba-CCdesttra-777 model, while the lower figure presents the retrieval accuracy of the Mistral-7B model. We present heatmap where "full context length" refers to the total length of the document, and "needle position" denotes the relative position of the statement be retrieved within the context.
flattened into sequences with a shape of 1124,3), where 1024 represents the sequence length and corresponds to the RGB channels of the pixel tokens. These pixel tokens are then linearly projected into H -ddmmnssional features, which are then input into SSM or transformer mixers. In addition to pixel tokens, we insert class token at the last position of the input sequence. The output state the class token will be processed by one-layer classifier head to generate the final logits.
Note that while the ViT architecture (Dosovitskiy et al., 2020) places the class token at the first position of the input sequence, this design is incompatible with SSMs, which rely causal sequence modeling. In SSMs, the class token must be positioned last to aggregate features from the entire sequence. We position the class token the last token establish long-range dependencies between
26

---

Published as a conference paper at ICLR 2025
Table Extended results of adversarial attack experiments on the CIFAR-10 dataset. Classification accuracy is used as the metric.
Corrupted region (seq. length = 1024)
Models
(no corrupt)
[1014:1024]
[0:10]
[768: 1024]
[0:256]
[512:544]
[480:576]
H3
0.654
0.629
0,654
0.394
0.639
0.603
0.343
Transformer
0.580
0.571
0.500
0.249
0.263
0.498
0.347
RWKV
0.474
0.194
0,470
0.107
0.448
0.405
0.392
Mamba
0.674
0.348
0.664
0.099
0.597
0,515
0.446
global image features and the leading pixel tokens. Alternative methods for aggregating features across the entire sequence, such as mean pooling (Gu et. 2021a Tay et al. (020) or placing the
class token in the middle of the sequence (Zhu et al. 2024), work more robustly in general but do not
fit the needs for our arguments on locality.
In addition, our image classification setup differs from Tay et al. 20200). where 8-bit pixel intensity lookup table is used as the token embedder. Instead, we employ a linear projection to map RGB pixel colors into -d-dimnsioonll features.
For fair comparison, the same hyperparameters are used across all models: learning rate = 0.001, weight decay 0.1, number of layers feature dimension H 32, and number of states 64. Each model is trained for 100 epochs. The models and training pipelines are built on Arora et al. (2023). No perturbations are imposed on the input sequences in the training stage.
Adversarial Attack. To introduce perturbations to test data for adversarial attack, we first define a corruption length K, which small relative to entire sequence length. We then replace the leading and trailing K tokens with random Gaussian noise. In our experiments, K is set to 32 and 96, corresponding to one row and three rows of pixels, respectively. Table shows more results under other corruption regions.
Target Attack. For the target attack experiments, a target class is first selected. For each image from the other classes, image from the target class randomly sampled, and its leading and trailing pixels are used to replace the corresponding pixels in the original image. We test two attack ratios: 256/1024 and 480/1024. Replacing fewer than 256 pixels generally does not result in considerable success rates based on our trials. In our main text, we show success rates when "horse" is the target class. Similar patterns are also observed across other classes. Fig. 10 shows the average success rates obtained by setting each class as the target.
0.6
Atlack region
Avg. success rate
[0:256]
0.4
[0:320]
[704:1024]
[768:1024]
0.2
0.0
H3
Mamba
RWKV
Transformer
Figure 10: Overall success rate of our target attack experiments on CIFAR-10, calculated by averaging the attack success rates obtained when each class individually the target class.
E.4 SCALING LAW WITH VARYING DEPTH
In this section, we elaborate experiment details for Sec. 4.1 We perform casual language modeling (Radford et al. 2018) with Mamba and report the validation perplexity. For each curve in Fig. 4 fix the depth of the model and vary the number of parameters from 140M to 550M by adjusting the hidden dimension accordingly. The depth is chosen from [16. 24, 32, 48, 64, 72). All other model configurations keep the default values: the memory state size N 16, the intermediate SSM
27

---

Published as a conference paper at ICLR 2025
dimension is two times the hidden dimension, and the intermediate rank for predicting time step A the hidden dimension divided by 16. We adopt aCtmpoLL (Li et al. 2024 as our training corpus for its high data quality. The evaluation set is created by holding out subset of 10M tokens
from the training data. We follow the training recipe provided in Appendix E.2.1 of & Dao (023). which roughly follows the Chinchilla scaling laws (Hoffmann al. 2022) allocate the number of tokens for each model size. We test two block sizes (2048, 8192 During training, we fix the number of tokens one training batch (# sequences sequence length) as 0.5M. The number of training steps is computed by the total number of tokens divided by batch size (# tokens / # tokens in one batch). We use Adam the optimizer with gradient clip value 1.0, weight decay 0.1, and linear learning warmup with cosine decay to le-5, following Gu & Dao (2023). The peak learning rate is x the GPT3 values (Brown et al. 2020) for different model sizes. summarize the training hyperparameters for varying-sized models in Tab. 4
# Params
Training steps
Peak LR
Batch Size (in tokens)
# Tokens
100-250M
4800
3e-3
0.5M
2.5B
250-400M
13500
1.5e-3
0.5M
7B
400-550M
20000
1.25e-3
0.5M
10B
Table 4: Summary of training settings for varying-sized Mamba. The settings are following Chinchilla law (Hoffmann et 2022) and consistent with Gu & Dao (2023).
E.5 SSM POLARIZATION AND ASSOCIATIVE RECALL EXPERIMENTS
In this section, we provide implementation details for our proposed polarization technique. Also, we introduce experiment details for the associative recall tasks we used to validate our techniques.
E.5.1 POLARIZED MAMBA
Recall that we polarize channel in A, with constant zero and another with constant one. Note that in the implementation of Mamba, A, is parameterized by A, exp(A,A). where A a learnable parameters. We change the forward pass of Mamba by prepending zero and appending a number negatively large enough (i.e., -1000). So the pre- and post-exponential A and A, become:
A +
A
A,
exp(A.A)
1000
We also study two variants: O-polarized and 1-polarized Mamba by only pretending or appending one or zero to A₁ to single out the effects of one and zero gating, respectively.
The gradient flow also complies with the original one, meaning polarization does not change the optimization dynamics of Mamba. Below we show that the polarized A, does not influence the
gradient in terms the time interval A. First of all, let G. (a.b.(
(..b then A4.
dE
N-1
d 0(A.).. dE a(A))11 OA,
ae ((A)).. ae
= G.
+
AA.
8 ((A)
0(A.).1 AA. J(A.)NN OA. OA. dt ((Ar).. de exp (sp(A,(A.).. (At)1.1
2=
N-1
= =G +
O ((A)) Am)m
AA. a(A.).1
n=
+ 0( 0(A,) N.N (A)(A(()
20
N-1
d
a ((A,),,
G.
O (())...
aA,
where the term exp( (xp(A,(AA (A N.N) (A.)..N + 0 is achieved by taking (A,) + -00. The component directly contributed by A, is not affected regardless of the polarization.
28

---

Published as a conference paper at ICLR 2025
KV Pairs
Configurations
# Layers
Recency
Over-smoothing
Avg.
64
128
256
Default A,
98.38
81.81
36.00
72.06
Default A,
99.23
82.08
33.52
71.61
(A+) =
99.81
94.70
56.39
83.63
(Ac) N.N =
98.41
81.35
36.55
72.10
(Ae) =0 0
99.74
92.20
52.21
81.38
(At),, = (,(,),, =0 (A)),1 = 1, 1, =
99.23
95.54
54.74
83.17
99.94
98.80
81.56
93,43
Table 5: Extension Tab.
We note the extent locality over-smoothing each configuration. consider 1-polarization mitigates locality most significantly, while deepening architecture only relieves recency mildly but deteriorates over-smoothing. 0-polarization alleviates over-smoothening and unleash the benefits by depth scaling.
E.5.2 ASSOCIATIVE RECALL EXPERIMENTS
Associative Recall (AR) the task of retrieving a specific value tied to a given key. This ability particularly important for in-context learning, where model leverages previously presented information to adapt to new inputs (Olsson et al. 2022 Arora et al.) (023)..
Mamba as the representative SSM, following the setup described in Arora et al. (2023). The input sequences to the model consist of two parts. The first part is a set of key-value pairs, randomly sampled from vocabulary, with keys and values drawn independently. The second part contains random keys, selected from those present in the first part (without repetition), and placed according to a power-law distribution at random. Any remaining slots the sequence are filled with zero token. The target sequence contains the values corresponding to the positions of the respective key occurrences, while all other tokens are excluded from loss supervision and accuracy evaluation.
In experiments,
We employ training strategy with mixed sequence lengths and varying numbers of key-value pairs. Specifically, the total sequence length varies across (64, 128, 256, 512, 1024), with (12.5%, 25%, 50%) ) of the sequence allocated to key-value pairs and the remainder to queried keys. Each configuration includes 20,000 examples, resulting in a total of 300,000 training examples. For evaluation, the sequence length is fixed 1024, while the number of key-value pairs varies among 64, 128, 256. We use a batch size of 128 and a learning rate of le-3 throughout the experiments. We provide extended table for the results in Tab.
29


Reviews:
[{"review_id": "0uWun9b1SU", "reviewer": "Reviewer_2oyp", "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", "questions": "Please see above", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper presents important findings, which open up future directions to resolve the issue.\n2. The section on adversarial attacks is praiseworthy, as it utilizes one of the findings to describe something very practical.\n3. 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.\n2. How do initialization scheme of SSMs interfere with these findings. Can there be some discussion around these?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5Fa23p4lSv", "reviewer": "Reviewer_53Ds", "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.", "questions": "What do you this should the path ahead be for improving SSM models?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 2, "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\n\n1) giving practical examples of failure cases of SSMs.\n\n2) showing high-quality and clear plots that can be used for future research or for reference.\n\nThis 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:\n\n1) 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). \n\n2) 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.\n\n3) 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.\n\n3) The paper does not provide hints into solving the stated problems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qPOc1Dmog4", "reviewer": "Reviewer_LroV", "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.", "questions": "Please see above.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written and easy to follow.\n\nExperiments 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.\n\nThe 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?\n\nThe 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?\n\n\n\nMinor comments: \n- Typo in Eq. 4: The indices $i$ and $j$ are likely meant to be $s$ and $t$.\n- Typos in Eq. 18 and Eq. 19: the inequality signs are likely to be $ \\geq $ and $ \\leq $, respectively.\n\n\nReference\n\n[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.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3WcVE8cZ0b", "reviewer": "Reviewer_veE3", "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.", "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?\n2. 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The limitations of SSMs’ capabilities are theoretically examined from multiple perspectives, including recency bias, robustness to attack, and over-smoothing.\n2. 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).\n\n    [Wang & Xue, 2023] State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. In NeurIPS 2023.\n    \n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)