Lambda-Skip Connections: the architectural component that prevents Rank Collapse
Authors: Federico Arangath Joseph
Jerome Sieber
Melanie Zeilinger
Carmen Amo Alonso
Venue: ICLR 2025 Poster (2025)
Keywords: Rank Collapse
Skip Connections
Sequence Modeling Architectures
Primary area: learning theory

Abstract: Rank collapse, a phenomenon where embedding vectors in sequence models
rapidly converge to a uniform token or equilibrium state, has recently gained at-
tention in the deep learning literature. This phenomenon leads to reduced expres-
sivity and potential training instabilities due to vanishing gradients. Empirical ev-
idence suggests that architectural components like skip connections, LayerNorm,
and MultiLayer Perceptrons (MLPs) play critical roles in mitigating rank collapse.
While this issue is well-documented for transformers, alternative sequence mod-
els, such as State Space Models (SSMs), which have recently gained prominence,
have not been thoroughly examined for similar vulnerabilities. This paper extends
the theory of rank collapse from transformers to SSMs using a unifying frame-
work that captures both architectures. We introduce a modification in the skip
connection component, termed lambda-skip connections, that provides guaran-
tees for rank collapse prevention. We present, via analytical results, a sufficient
condition to achieve the guarantee for all of the aforementioned architectures. We
also study the necessity of this condition via ablation studies and analytical exam-
ples. To our knowledge, this is the first study that provides a general guarantee to
prevent rank collapse, and that investigates rank collapse in the context of SSMs,
offering valuable understanding for both theoreticians and practitioners. Finally,
we validate our findings with experiments demonstrating the crucial role of archi-
tectural components in preventing rank collapse.

Published as a conference paper at ICLR 2025
LAMBDA-SKIP CONNECTIONS: THE ARCHITECTURAL COMPONENT THAT PREVENTS RANK COLLAPSE
Federico Arangath Joseph ETH Zurich Farangath@ethz.co
Jerome Sieber
ETH Zurich js jsieber@ethz.ch
Melanie N. Zeilinger
Carmen Amo Alonso Stanford University camoalon@ Lunoalon@@taafood.ccc
ETH Zurich
mzeilingerOethz.chh
ABSTRACT
Rank collapse, a phenomenon where embedding vectors in sequence models rapidly converge to a uniform token or equilibrium state, has recently gained at- tention in the deep learning literature. This phenomenon leads to reduced expres- sivity and potential training instabilities due vanishing gradients. Empirical ev- idence suggests that architectural components like skip connections, LayerNorm, and MultiLayer Perceptrons (MLPs) play critical roles in mitigating rank collapse. While this issue is well-documented for transformers, alternative sequence mod- els, such as State Space Models (SSMs), which have recently gained prominence, have not been thoroughly examined for similar vulnerabilities. This paper extends the theory of rank collapse from transformers to SSMs using unifying frame- work that captures both architectures. We study how a parametrized version of the classic skip connection component, which we call lambda-skip connections, provides guarantees for rank collapse prevention. Through analytical results, we present a sufficient condition to guarantee prevention of rank collapse across all the aforementioned architectures. We also study the necessity of this condition via ablation studies and analytical examples. To our knowledge, this is the first study that provides general guarantee to prevent rank collapse, and that inves- tigates rank collapse in the context of SSMs, offering valuable understanding for both theoreticians and practitioners. Finally, we validate our findings with exper- iments demonstrating the crucial role of architectural components such as skip connections and gating mechanisms preventing rank collapse.
INTRODUCTION
The phenomenon rank collapse has been recently reported in the deep learning literature Dong et al. (2023) y Noci al. (2022); Shi et (202)) [ee al. (2023): Geshkovski Geshoooo al. 2024) Wu al.
(2024a] b). Rank collapse is a phenomenon by which embedding vectors exhibit a fast convergence rate into uniform token, or equilibrium embedding vector. Various explanations why this phe- nomenon arises have been proposed: Dong et al. (2023) proposed a path decomposition argument and suggests that of the reasons on why rank collapse happens is that the attention matrix row stochastic. (Geshkovski et al. (2024) instead proposes an analysis of rank collapse from statistical physics perspective by treating transformers as mean-fiel interacting particle systems and observes that the resulting system converges practice, this convergence
metastable state corresponding to Dirac delta measure. In uniform token reduces model expressivity [Daneshmand et al. (2020); Dong al. (2023); Noci et al. (2022): Wu et al. (2024b) since low rank in the output matrix means that the model is unable to capture complex relationships between tokens, also potentially causing vanishing gradients during training (Noci al. 2022). Therefore, gaining deeper understanding of this phenomenon its origins, conditions for occurrence, and potential prevention strategies is essential for designing models that are more robust, stable, and expressive.
Rank collapse has primarily been studied in the transformer architecture Vaswani et al. (202)). In particular, it has been reported that the convergence rate the equilibrium embedding vector of self-

---

Published as a conference paper at ICLR 2025
attention Vaswani et al. 2023) (corresponding to rank-1 matrix) is doubly exponential (Dong al. (2033). ). Both empirical observations and theoretical results support that the presence of skip con- nections (He et al. (2015) as well as the LayerNorm component (Ba al. (0166) in the architecture mitigate the issue of rank collapse Dong et al. ) Wu et al. (2024a). In particular, Dong et al.
(2023) shows that by adding skip connections, there exist an infinite amount of parametrizations of
the network for which rank collapse does not happen. [Wu et al. (2024a) instead argues that adding LayerNorm slows down rank collapse by showing that in this setting the upper bound on the rank collapse measure is higher than for networks without these components. This implies that, in order to have convergence to 0, stricter conditions on input sequence and attention weights are needed.
Sequence models alternative to transformers have recently gained widespread attention. In particu- lar, State Space Models (SSMs) such as S4 Gu et al. 2022b). S5 Smith et al. (2023). H3 [uu et al.
2023) and Mamba Dao and Gu 2024); Gu and Dao (2024) have been very prominent, where at-
tention was replaced by SSM blocks, which operate combination with other components such as Gating Mechanisms (Hochreiter and Schmidhuber (997); Chung et al. (2014)), MLPs (Popescu)
ctt al. (2009) and skip connections (He et (2015))) Furthermore, comparisons between these ar- chitectures have been discussed in (Sieber et al.) [024] Hè and Kabic 2023] Tiezzi et al. (024), and a larger framework where both the transformer and SSMs are seen as different realizations of the same general model has recently been proposed (Ali et al. (2024), [Do and Gu (2024)). However, whether they suffer from similar weaknesses, and whether rank collapse happens more broadly in these architectures, is still an open question.
In this work, we focus on the specific issue of rank collapse and explore how slight modifica- tions in the architecture, referred to as lambda-skip connections, can help to prevent it. Although parametrized versions of skip connections have been explored in the past (see e.g. He et al. 2016 Srivastava et al. (215)), ), their impact in providing guarantees to prevent rank collapse has not been studied. Moreover, here we develop general theory of rank collapse using the unifying frame- work for sequence models presented in Ali et al. (2024) and (Dao and Gu 2024), comprising both transformers and SSM architectures. To the best of our knowledge, the rank collapse phenomenon in the SSM setting has not yet been studied and we are the first to provide a general lower bound for rank collapse which holds for both Transformers and SSMs. Additionally, we offer mecha- nistic solution to prevent rank collapse by introducing a skip strength parameter, which comes with minimal computational overhead while enhancing the model's 's expressivity and stability. Since rank collapse can be a very serious issue for practitioners, investigating broadly when rank collapse oc- curs and designing robust architectures that prevent rank collapse is of central interest not only to theorists but also to experimentalists. Our contributions can be outlined as follows:
* We extend the theory on rank collapse originally developed for transfomers to SSMs by using the recently proposed framework unifying these two architecture types (Ali et al.) 2024 Dao and Gu 2024).
We investigate how lambda-skip connections influence rank collapse. In particular, give sufficient conditions that guarantee that rank collapse does not occur in the finite layers setting. Our conditions apply to both transformers and SSMs. We also study the necessity of this condition via ablation analysis and analytical examples.
We provide experimental details in support of our theory, by showing how adding a parame- ter controlling the skip connection strength could be beneficial in preventing rank collapse.
We also empirically show the role of gating mechanisms in preventing rank collapse. To the best of our knowledge, we are the first to make this connection between gating mecha- nisms, which were originally constructed for improving memory capabilities of the models (Hochreiter and Schmidubeer 1997), and rank collapse.
2 RELATED WORKS
Rank Collapse. Rank collapse is the phenomenon in which the rank of the output layer of the network collapses to (Dong et al. (2023) ) To mathematically capture this, measuring the rank is not ideal since for every finite amount of layers, the output matrix will always be full rank with probability 1. This led to the development of an alternative metric to measure rank collapse Dong et al. 2023). which intuitively evaluates how "ill-conditioned" is the matrix of the output layer, i.e.,
2

---

Published as a conference paper at ICLR 2025
how close it is being rank 1. In Dong et al. (2023), this metrie is defined by computing how far the columns of the matrix are from the mean of all columns. It was shown in e.g. Dong et al. (2023);
Wu et al. 2024a) that this metric quickly decays to O the depth of the model increases. Hence, rank collapse is an intrinsic issue of deep-learning models. Rank collapse has mainly been studied for transformers (Vaswani et al. (2023)), Graph Neural Networks (Zhou al. (202T) and ReLUU networks (Agarapl (2019), but not SSMs. The work in Daneshmand et al. (2020) studied rank collapse in randomly initialized linear and ReLU networks and proved that Batch Normalization can be helpful to mitigate the phenomenon. In Dong et al. (2023) the authors show that in self- attention-onle -onllooooonyyyyoo transformers, rank collapse occurs doubly exponentially, i.e. with rate O(a²)) for some b E R. In the same paper, the authors propose for the first time that skip connections and MLPs (Popescu et al. (2099) ) are helpful in preventing rank collapse from happening. Moreover, the work in Noci et al. 2022) shows that rank collapse is not only an issue during inference but also hinders training due to vanishing gradient problems at initialization. In Wu et al. (2024b) the authors show that rank collapse also happens in graph attention networks (in this setting the phenomenon is also known as oversmoothing). also causing the model to lose expressive power. Finally, the work in Wu al. (2024a) studies the effect of LayerNorm on rank collapse in transformers, by proving that this component can help preventing or mitigating rank collapse.
Skip Connections. Skip connections were first introduced in ResNets He et al. (2015) with the purpose of easing training and optimization in deep neural networks (He et (2016); Veit et al.) (2016;; Balduzzi al.) (2018). More precisely, the introduction of skip connections addresses the vanishing gradient problem by decomposing baok-propagation of gradients into multiple paths and thus allowing them to skip and bypass layers. As an alternative explanation, was shown in Li et al. (2018) that adding skip connections also results in loss function which is much smoother and hence it is harder for gradients to get stuck in local minima, facilitating the training procedure. In transformers and SSMs, the introduction of the skip connection followed the same motivation as for ResNets and other deep neural networks, i.c. to improve and facilitate optimization. The work in Dong et al. (2023) is the first to also connect the benefit of skip connections to avoiding rank collapse by studying the paths decomposition in transformers. Regarding modifications to the skip connection. [Bachlechner al. (2020) also considers a control parameter in residual networks. However, their approach differs in that it is used as a gating with the main network layer again with the purpose of stabilizing optimization, which allows for training of much deeper models. The work in Srivastava et al.) (2015) proposes adding an extra feedforward layer to each residual layer. The output of each residual layer is then expressed as convex combination of two terms: one gated by the residual connection and the other by the output of the primary network layer. The effect of these modifications on skip connection have not been analyzed. Lastly, HHe et al. (2016) studies the lambda-skip connection proposed here in terms of training efficiency and concludes that the standard choice of À = 1 is usually beneficial.
Normm LayerNorm (Ba et al. (2016) has become the standard choice of normalization layer for sequence models such as transformers ( /Vaswani et al. (2023) and Mamba (Dao and Gu (2044) )- In contrast to batch normalization (Loffe and Szegedy y 2015)), LayerNorm performs the shifting and the normalization token-wise and not sequence-wise wise. This is particularly beneficial as the embed- ding dimension is usually much smaller than the sequence dimension, resulting in much cheaper computational overhead than BatchNorm. It has also been shown that LayerNorm is helpful to stabi- lize training (Xiong et al. (2020)) and increase the expressivity of the attention layer (Brody et al. 2023 Related instead to rank collapse, the work in Dong et al.] (2023) was the first to analyze a potential connection between rank collapse and LayerNorm. However, they conjecture that Layer- Norm has no impact on rank collapse. Later on, the work in Wu et al.) (2024a) refutes this claim, showing both empirically and theoretically that LayerNorm helps slow down rank collapse.
PROBLEM SETUP
3.1 SEQUENCE MODEL DEEP-LEARNING ARCHITECTURE
Rank collapse is defined as the convergence in the limit of infinite depth of the output matrix of the network's layers to a Rank-1 matrix. In this sense, rank collapse is intrinsic to deep-learning architectures. In what follows,
outline the different components on single-layer of the archi- tecture considered in this paper. We note that the presented architecture is general enough to capture both transformers and SSM-like architectures.
3

---

Published as a conference paper at ICLR 2025
We start by introducing some notation. We index each layer of the architecture with & {1,..., K). We denote X E RNxd and Y (k) E Rxxx to be the input and output of the k-th layer, respectively, where N is the sequence length. In this setting, (()) is the input sequence to the model. Each layer has the following components.
Main mechanism. The main mechanism differs between attention in transformers, and recurrent block in SSMs. However, as shown in Ali ct al. (2024) and Dao and Gu 2024), both models can be seen part of general framework
(k) = (())
In the case of attention,
w (X(())T Q (
(() w()
(b)
ERNxd and ( softmax
RNN,
Vdox
with w E Rdxd w respectively.
Rdxd and w Raxd are the query, key and value matrices
K
In the case of the recurrent block, (()) (A), and (k) a lower triangular matrix with
(
M (k) c,Ah)
II
Bit)
[=D
where A(t) iddl xdH B(()) E gd c(R) E RR U,(,) E glx1 (usually it is indicated with
D(4) but we decide to use different notion to avoid confusion with LayerNorm). The first versions
of the SSM blocks consisted of Linear Time Invariant (LTT) representations, i.e., A (k).
B(a) = B(a), c,A)) c()) and = (k) for all time steps (Gu et al. 2022b a Smith et al. 2023 Recently, in order to improve the expressivity and the performance such models, selective State Space Models have been proposed (Gu and Dao 2024 Dao and Gu 2024). In these models, matrices A,, B,, and C₁ are input-dependent. We refer to Appendix A. for a detailed explanation of attention and the recurrent block.
Skip Connection. The skip connection sums the unmodified input to the output of the self- attention layer:
y-()) (+) (A))
Although it was originally proposed to stabilize training by avoiding vanishing gradients (He et al. (2015), was recently shown that it also has a role modulating rank collapse Dong et al. (2023). In SSM architectures, skip connections are often preceded by gating mechanisms, i.e., ( in equa-
tion 2 is replaced by (kk) (kk) o (W()) x (()) However, we ignore these in the theoretical part of this paper for simplicity. In this work, we modify the skip connection by allowing it to have an additional parameter (k) E I, that controls the strength of the skip connection, i.c..
Definition 3.1. Given a standard skip connection as per equation and parameter ( E, lambda-skip connection is defined as
(N) (()) (k))
(3)
In the theoretical section, we will consider (8) to be fixed for all the layers for simplicity. However, in the experimental section, we will also consider models with learnable (4k), i.c. it might vary across layers.
Layer Norm. The LayerNorm (Ba et al. 2016) computation shifts and normalizes the output (*) after the skip connection. Similar the skip connection, the purpose the LayerNorm also to stabilize training. It was recently shown to also help to mitigate rank collapse (Wu et al. (022a))
Here, we consider a slightly simplified version LayerNorm, similar to Wu et al. (202aa) where we only apply normalization and not shifting, namely:
((k) = (()) (4)
(4)
where (*) = diag(d (R)
(k)
E NXN and d
E
1
(-) is non-linearity, often SILU, (kk is weight matrix and @ the Hadamard product.
4

---

Published as a conference paper at ICLR 2025
Other components. After the above-mentioned computations, the input Y passes through MLP (Popescu et al. 2009). which can intuitively be seen as a non-linear feature extractor or a kemnel MLPs might play role in mitigating rank collapse since they increase the Lipschitz constant of the network (Dong et al. 2023). In particular, the Lipschitz constant of the MI.P layer appears in the double exponential decay term of the upper bounds of the rank collapse measure. This is beneficial since it allows for a wider range of choices for the parametrization of other components while maintaining high-enough upper bound. Since the role of the MLP layer on rank collapse is well-understood, in this work we focus on the effect that LayerNorm and skip connections have on rank collapse. We anticipate that the addition of MLP will further improve the bound, and hence the rank collapse avoidance conditions.
3.2 PROBLEM STATEMENT
In mathematical terms, the phenomenon of rank collapse means that the columns of the final output matrix are proportional to each other. The representational capacity of the model is severely hindered by this phenomenon: since the final representations of each token proportional to each other, there is limited diversity of tokens to be chosen from. Given that the rank of the output layer will be full rank with probability one (since rank-deficient matrices are zero measure subspace of the space of all matrices), a different metric was introduced in Dong et al. (2023) : Wu et al. 2024a to assess rank collapse:
(5) where ((k) E RNxd is the output matrix resulting from the k-th layer, Ty = N 1 (x), and Il I| is the Frobenius norm.
((Y("))= = ||Y -1
In order to study the impact of the architecture on the phenomenon of rank collapse, we make use of the equations describing the architecture: equation I equation 3 and equation 4 Moreover, we
note that the output of layer k - 1 is the input to layer k. i.e. ( = (k)), Combining these equations obtain
(+) ((*) (M () DC (k) (k)y(()-)))
where c/ w for the attention block and c for the recurrent block.
The goal of this paper is to show how the introduction of (k) in the lambda-skip connection influ- ences rank collapse. For this purpose. the main quantity of interest is (k) i.c., the value of the saken-similarity metric on the output of the last layer. We provide a mathematical analysis of this metric in terms of the different architectural components in the next section.
4 LAMBDA-SKIN CONNECTION: NECESSARY AND SUFFICIENT TO PREVENT
RANK COLLAPSE
Here, we present the analysis resulting from studying rank collapse with lambda-skip connections. In Section 4.J. we analyze the scenario where we replace the regular skip connection by a lambda- skip connection and show that, together with LayerNorm, an appropriate choice of in equation is sufficient to prevent rank collapse in the finite layers setting. For simplicity, here we consider fixed value of À for all the layers, i.e. (k) = A Vk. Specifically, we provide the value of À E R in the lambda-skip connection that guarantees the absence of rank collapse. This result holds for any arbitrary amount of layers and all the three architectures studied in the paper (transformers, LTT SSMs and selective SSMs). In Section 4.2 we study the necessity of this condition. First we show that in the ablated architecture without skip connection, rank collapse occurs exponentially, i.e., (K)()) converges to O exponentially. If, additionally, LayerNorm is also ablated, rank collapse occurs doubly exponentially. Second, we highlight via examples that the strength of the lambda- skip connection (parametrized through À E R) plays crucial role in determining the effect of skip connections on rank collapse.
4.1 LAMBDA-SKIP CONNECTION: SUFFICIENT TO PREVENT RANK COLLAPSE
Here we show the role that the parameter À E Ik plays in guaranteeing, under appropriate conditions, the prevention of rank collapse. To do this, provide X-dependent lower bound for the rank
5

---

Published as a conference paper at ICLR 2025
collapse metric ((((()) in equation 5
We start by defining the following quantity:
1 2ANdSCM
= k X2 a(SCM + +
with Cm := sup, |I sup 11 Il. We also define the collapse rate as:
Definition 4.1. The collapse rate, O, is the rate at which the lower bound for the rank collapse
metric u(Y(k))² decays with the number of layers, i. u(Yik))a a u(Y(0))2
Theorem 4.1 (Lower Bound on Rank Collapse). Let the input sequence ( be such that (0) 2 b. If the skip connection strength is chosen to satisfy
X² aSSSss + 11)) > O,
then we can lower bound (KY(K)) by (( 2 a ((Y10)) for all K KEN E.N.
Proof. The key part of the proof is show the following relationship between the rank collapse measure two consecutive layers: 1)) 2 [SCM+|/l]Jz (A² 2ANdCMr). which is done by algebraic manipulation and upper bounds on || The claim follows by imposing the right-hand side above greater of equal to a))((()))2 rearranging the terms and imposing the condition on À to guarantee the feasibility of this. The full proof can be found in Appendix A.4
We note that this result is very general, since it holds for any model that can be expressed in the form of equation [ independent of its specific implementation or parametrization. Moreover, the presence of LayerNorm implies that the bound provided in equation doos not depend on the input even when the matrix M is input-dependent. This is because the input is normalized before being fed through M. Specifically, Cλ is only dependent on the weights and on the sequence length
Theorem also provides an important intuition for the choice of X. In particular, we note that to guarantec prevention of rank collapse, the ideal choice is a = 1. However, this choice will be mediated by the values of C'M and S in the different architectures. In particular, it is easy to that the only way to guarantee a solution to is by having a > 0, for which needs to satisfy [사| (a+va)SCMM
Remark 4.1. In the case of Mamba, Cy and c c= c= = This implies that the only possible choice for the collapse rate is a < 1. Although this does not guarantee ((Y(K))2 2 ()), in practice can be chosen be very close 1. For instance, one can set a 0.9999. Given K = 64 (standard number of layers in Mamba), a so 0.993. Hence, in practice, if we choose X appropriately, this choice still prevents rank collapse (see e.g. Figure 2
Finally, we conclude this section by briefly discussing how Theorem can be generalized to the case where we let À vary across different layers. The following small modification will be enough to make the result still hold this setting: order to still satisfy the condition (((())) 2 ap(YY (k-1)),
we can simply adapt equation to be X - a(S.CM, + |Ax )² > O. where S₃ = ||c| || and Cm. || Additionally, the flexibility to choose A allows for different values of a to be used across layers. For instance, instead of setting a single a value for all layers, one can vary A to meet different layer-specific conditions. In one layer, Y might satisfy the condition for smaller value of a, while in a subsequent layer, A satisfies the condition for a larger value of a.
4.2 LAMBDA-SKIP CONNECTION: NECESSARY TO PREVENT RANK COLLAPSE?
Next, we explore the role that the lambda-skip connection plays in terms of being necessary to pre- vent rank collapse. Although we do not provide formal necessary condition, explore this idea in two ways. First, we provide theoretical results on how rank collapse occurs when the (standard) skip connection is ablated. Then, give two examples of how rank collapse may occur when the lambda-skip connection does not attain the bound in Theorem 4.1
In transformers Cx VN, in structured LTT SSMs Cs 2 sup₁ ||4001| [| [|C())] and in selective SSMs Cи sup; ||W|"1| [|W>[*)] [|r (since in Selective SSMs it is common practice choose A with eigenvalues smaller than onc).

---

Published as a conference paper at ICLR 2025
4.2.1 RANK COLLAPSE OCCURS WITHOUT SKIP CONNECTIONS
Here we show how, for different architectures, rank collapse occurs in the absence of skip connec- tions. Here we focus on transformers and selective SSMs. Additional results on LTI SSMs are presented in Appendix A.8
Transformers. We start by reporting for completeness the result in Wu et al. (2024a) for trans- formers architectures with Self-Attentio and LayerNorm only.
Theorem 4.2 (Corollary 1, Wu al. (2224) (Informal)). Under certain assumptions the self-
attention weights and the input sequence, there exist C and
> 0, > 0 and r > 0 such that Ne <
((Y(K)) S e
AK 2 O.
meaning that tokens converge to common point on s exponentially.
We refer to Appendix A.5 for rigorous presentation of Theorem 4.2 The main takeaway from this result is that, under appropriate conditions on the architecture and the input sequence, removing skip connections and only using attention layers and LayerNorm can cause exponentially fast convergence to rank-1 matrix. This highlights the importance of skip connections to mitigate rank collapse. It is also interesting to compare the result in Theorem 4.2 to the setting of self-attention only transformers (without both LayerNorm and skip connections), reported in Dong et al. (2023). Under certain assumptions the attention matrix and the input sequence, rank collapse occurs at a doubly exponential rate. We report the full result in Theorem [A.3] and Corollary |A.3.1 (where we adapted the Theorem the rank collapse measure provided in the problem statement).
Selective SSMs. According to Dao and Gu 2024), for selective SSMs where A, = =al (e.g.,
Mamba-2) 2), the matrix M can be compactly written as (a) (s) @
( (r(k (k)"
where (a) is a lower-triangular 1-semiseparable matrix. Due to space limitations, we formally define 1SS(av in Equation 12 in Appendix A.3] In what follows, we introduce the following as-
sumption on selectivity:
Wn and C, Y (R) We are input-dependent, while A; A = al VL (with a S 1) is independent of the input.
Assumption 4.1. Only the matrices Y
(k)
Y!!
This choice is done for ease of exposition. Although it may seem restrictive in theory, show experimentally that eliminating skip connections can cause rank collapse for selective SSMs where A is input dependent as well, such as Mamba (see e.g. the plot corresponding to À 0 in Figure D. In the following derivation, we neglect gating mechanisms for the sake of the theoretical analysis-we will thoroughly analyze their effect on rank collapse simulation (see e.g. Figure 3).
Theorem 4.3 (Rank Collapse for selective SSMs without skip connection). Let ((k))
min (Y-(A) y((k), where (,,0) indicates the inner product. Under Assumption 4.1 and if some e some
(00)
O R, Amin 0, E, Mij 2 Vi, then holds that:
(KY VN 22
2 2N)
AK 2 0,
where Amin and Amax the minimum and maximum eigenvalues of WoW W.
Proof. The proof adapts the derivation in Corollary 1 in Wu et al. (2024a) to selective SSMs. In particular, the key is to find a recursive relation between の and Φ (++1) The full proof can be found in Appendix A.7
In particular, when Amin 2N the model suffers from exponential rank collapse. Note that the above result holds for a model with layers having the same parametrization.
In Appendix A.9 Theorem A.10
we show that if we instead have selective SSM with ablated the skip connections and LayerNorms, then rank collapse can happen doubly exponentially under certain conditions, similarly to what we have in transformers architecture. In contrast with the LTI SSM case (Theorems A.7 and A.8), which exhibits an exponential decay, the matrix (k) depends quadratically on the input for selective SSMs: || S / [[ l|r. where WBC WCWT This dependency is what causes the doubly exponential behavior.

---

Published as a conference paper at ICLR 2025
4.2.2 RANK COLLAPSE CAN STILL OCCUR WITH SKIP CONNECTIONS UNDER SPECIFIC CHOICES OF A: EXAMPLES
In the following, we show an analytical example of how, for Selective SSMs, ank collapse can still oceur in the presence of lambda-skil skip connections for specific choice of X. A similar example for Structured LTI SSMs can be found in Appendix A..In.
Selective SSMs. Consider the system, which we will denote as [Sys-2], A, = A al (for simplicity, we consider not be input dependent here), B YWB B and C YWc input dependent (where is the f-th row of Y) with LayerNorm applied to the output of each layer. Once
( a
9) By
again, we let N = 2 and d = 2. We have M = Wa @ (FWoWEX") where WA =
choosing a = and WB = We = I, we get that M = ( Vi y2 9) since the input to the layer is normalized by LayerNorm.
Proposition 4.3.1. Given [Sys-2], if we then choose
M as above and y
0
for any oo. Po with og > O, then we have that ((Y + On the other hand, if we choose X < < and M and Y as above, then rank collapse is avoided.
VaB+08
Vag+28
Proof. Again, the result follows by calculating (h) inductively and by showing that the correspond-
ing p (Y (K) decays to 0 for A > while this does not happen X < 3 We refer to Appendix
2'
A.1I for the full proof.
4.2.3 TIGHTNESS OF LOWER BOUND IN THEOREM
We conclude this section with discussion of the tightness of the lower bound provided in Theorem
4.1 In particular, we explore whether there exist a system such that (x) = O (((((()0))))
for a choice of that satisfies Equation 7 We show in Proposition 4.3.2 that such a system exists. Hence, without additional assumptions on the architecture or on the input sequence, it is not possible to provide a better guarantee for the lower bound of the rank collapse metric.
Proposition 4.3.2. There exist an architecture of the form of Equation 6 such that, when choosing A to satisfy Equation 7 i.e. |A| = S2( ( it holds that u(Y = O (a (oo=(YY (0))).
Proof. We again consider [Sys-2] from Proposition 4.3.1 with og 30 2. From the proof of
Proposition 4.3. we calculate the values of (k) and μ (((())) for all layers Then, we calculate the value of a by simply considering the ratio of the rank collapse measure at two consecutive layers.
Then, the result simply follows by upper bounding the value and comparing it with the desired bound. The full proof can be found in Appendix A.12
5 EXPERIMENTS
In this section, we aim to both explore the link between gating mechanisms and rank collapse and to empirically validate our theoretical findings. In Section 5.1 we empirically validate Theorem [. demonstrating the importance of selecting the appropriate skip connection strength to mitigate rank collapse. In Section 5.2. we show that for the Mamba-2 architecture (Dao and Gu (2024)). gating mechanisms indeed play crucial role in preventing rank collapse. Finally, we provide additional experiments on the S4 architecture (Gu et al. 20226) in Appendix A.13.I
5.1 COMPARISON OF RANK COLLAPSE FOR MAMBA ARCHITECTURE WITH DIFFERENT VALUES OF À
We start by validating our main finding, i.e. that by controlling the value of À in the skip connection, we can prevent rank collapse from happening. In order to do so, we consider a billion parameters and pre-trained Mamba-2 model available on HuggingFace and remove the gating mechanisms. Instead, we add an additive skip connection with varying A. For simplicity, we keep the value of X constant every layer. Similar to Dong et al. (2023): Wu al. (202aa), sample 32 text excerpts from
8

---

Published as a conference paper at ICLR 2025
Wikipedia using the Wikipedia API and tokenize them in sequences of at least 128 tokens using the BERT tokenizer Devlin et al. (201)). We then forward pass these inputs through the pre- pre-trained model and evaluate the rank collapse measure j (normalized by the norm of the layer output) for each of the 64 layers for different values of X, as shown in Figure We observe that when [^] is small, the model suffers from rank collapse. This happens even when À = 1, which is usually the default design choice. After sufficiently increasing or decreasing X, we note that we avoid rank collapse, consistent with our theoretical results, particularly Theorem 4.1 This suggests that rather than fixing the skip connection strength at À = 1 (as is standard), it may be beneficial to treat À as a learnable parameter or hyperparameter improve model stability. In Figure 2
we analyze
Normalized Rankc Collanse Measure vi Desth Nombo
Rank Collapse in last layer function A
10-103
1.0
0.6
0.6
0.4
0.4
1:10
Albart
1:20
Kamba
2=102
0.0
-100
-75
-50
-25
0
25
50
75
100
30
Skip Strength
Layer
Figure I: (Normalized) Rank Collapse measure plotted for different values of the skip connec- tion's 's strength X. Again, shaded areas represent one standard deviation from the mean calculated over the 32 examples.
Figure (Normalized) Rank Collapse measure at the last layer plotted for different values of X for both Albert and Mamba. Shaded areas repre- sent standard deviation from the mean cal- culated over the 32 examples.
how the rank collapse measure at the last layer varies with different skip connection strengths for both the Albert's architecture (Lan et al.] (2020)) and Mamba-2, keeping the full architecture (also LayerNorm). In both cases, note the importance of tuning and choosing the correct value of X to avoid rank collapse. We see how a value of À very close to O, as we expected, is detrimental and causes a very low value of μr, even in the presence of LayerNorm, which is consistent with our findings in Theorem 4.3, Additionally, we see that a negative À leads to higher μ.. Intuitively, the reason for this could be that a negative À can be interpreted as negative feedback in the system, hence helping stabilizing it. Note that even if in both cases our condition on À in Theorem 4.1 is too conservative, practice much lower values of X are good enough to prevent rank collapse. In Figure 2 we note that as X increases, there is little variation in μ.. This can be explained using the lower bound proposed in Theorem 4.1 Specifically, as A + oc, the parameter a + 1. This implics that the rank collapse metric decreases much more slowly with the number of layers. In the limit, it does not decrease at all, remaining at or above the rank collapse measure at the input.
We investigate whether using À as a trainable parameter affects the learning performance. To this end, we train two transformer- based architectures and two SSM- architectures on the image task of the long range arena (LRA) benchmark (Tay al. 2020) and on multi-query associative recall (MQAR) task (Arora et al. 2023). The accuracy of all four architectures on both tasks are reported in Table and the details on the experimental setup are provided in Appendix AA.14 From our observations, we conclude that learning \ does not affect the performance and even outperforms the models with fixed À in some cases.
Model
Task [%]
Transformer
Lin. Transformer
Mamba
Mamba-2
A A=
var. À
X =
var.
A = 1
var. À
A = 1
var. A
Image (LRA)
32.64
32.85
34.10
32.80
63.04 81.5
62.92
42.28
38. 92
MQAR (L 512)
99.6
98.9
1.6
1.55
85.3
97.3
99.1
Table 1: Model performance in terms of test accuracy on the Image task of the LRA benchmark and the MQAR task (T. = 512, KV-pairs 64).

---

Published as a conference paper at ICLR 2025
5.2 COMPARISON OF RANK COLLAPSE FOR MAMBA ARCHITECTURE WITH AND WITHOUT GATING MECHANISMS AND LAYERNORM
We also analyze the effect of gating mechanisms and LayerNorms in the Mamba architec- ture. Because gating mechanisms can be viewed as a multiplicative form of skip connec- tions, we aim to empirically explore whether they, like skip connections, contribute to sta- bility and prevent rank collapse. We repeat the same procedure as above, i.c., forward passing 32 excerpts from Wikipedia to the pre-trained model and calculating the rank col- lapse measure across all layers. We replicate this for all considered cases, namely gat- ing plus LayerNorm, only gating, only LayerNorm, and neither gating nor LayerNorm.
As shown in Figure 3 when removing the gat- ing mechanism, we either observe the rank col- lapse measure approaching zero especially for later layers if the LayerNorm is present, or observe initial drop in the rank collapse measure followed by a dramatic increase if we remove the LayerNorm. This highlights the importance of having LayerNorm to keep the model stable, as was also observed in Dao and Gu (0024). Indeed, even when the model con- tains a gating mechanism, although we prevent rank collapse, removing LayerNorm leads to sudden increase in the rank collapse measure, indicating instability. Interestingly, we find that gating mechanisms, originally designed to en- hance memory capabilities, also play crucial role in preventing rank collapse. To the best of our knowledge, this is the first time this connec- tion has been made.
Norma lzze Rank Collagse Measure vs Layers Depth Mamba
2.0
0.4
0:2
50
Layer
Gating-True, UN-
GatingFFlice, LN-True Gatin-False, LN=False
Gating-True, LN-False
Figure 3: (Normalized) Rank Collapse measure of the Mamba-2 model plotted depth. The shaded areas represent one standard devi- ation from the mean calculated over the 32 examples. Gating= True/False indicates we use a Mamba2 archi- tecture with/without gating whereas LN=True/False in- dicate use the architecture with/without LayerNorm
function of layer
CONCLUSION
This paper studies the phenomenon of rank collapse across general sequence model architectures, using the unifying framework from Ali et al. (2024) Dno and Gu (2024). Our main contribution the introduction of lambda-skip connection, which adds a regulatory term to the standard skip connection component to control its strength. We provide sufficient conditions for this parameter under which rank collapse does not occur. We do this via a lower bound on the rank collapse mea- sure, which holds for transformers, LTI SSMs, selective SSMs and more generally, all architectures that can be written in the framework from Ali et al. (004) Dao and Gu (2024). We further study the necessity of the lambda-skip skip connection by showing that when this component is ablated, selec- tive SSMs also experience rank collapse even in the presence of LayerNorm. Moreover, we show that similarly Transformers, selective SSMs can also be affected by doubly exponential rank col- lapse when both skip connections and LayerNorm are ablated. We suggest that the primary cause of this doubly exponential rank collapse is the input dependence of the M)) matrix. We validate our theoretical analysis with simulation experiments and further analyze other components such as the Gating Mechanism, which we have found also play a role in rank collapse.
Limitations and Future Work. A limitation of this work is the exclusion of gating mechanisms from the theoretical analysis. Since gating mechanisms help to mitigate the effect of rank collapse, we expect their inclusion in the derivations would result in tighter upper bounds on the rank collapse measure for Mamba. Furthermore, one other line worth exploring is to consider the effect on lower bounds of MLPs addition to skip connections and LayerNorm, Deriving tighter lower bound that matches the results in our experiments is also a very promising avenue of research. Additionally, Remark 4. offers intriguing insights from a control-theoretic perspective, as it suggests that the parameter can be interpreted as control gain. We believe that further theoretical analysis of the stability and performance associated with A-skip connections could lead to more principled design of learning architectures. Finally, exploring alternative measures of rank collapse, like effective rank, and their correlation with our metric is a promising direction to expand on how our results generalize across metrics.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
A. F. Agarap. Deep learning using rectified linear units (relu), 2019. URL https://arxiv.
org/ hte//ab//aa abs/aaa3...... 08375
A. Ali, Zimerman, and L. Wolf. The hidden attention of mamba models, 2024. URL https:
//axxiv.org/abs/22 http://aabb/203.05555/
S. Arora, S. Eyuboglu, A. Timalsina, I. Johnson, M. Poli, J. Zou, A. Rudra, and C. Ré. Zoology: Measuring and improving recall in efficient language models, 2023. URL bttps: org/ abs/ 2312 04927
J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization, abs/1607 06450
T. Bachlechner, B. P. Majumder, H. H. Mao, G. W. Cottrell, and McAuley. Rezero is all you need: Fast convergence at large depth, 2020. URL https: // wtwpv.axiivv..... org at////b///aa 2000./2003.044 04887
D. Balduzzi, M. Frean, L. Leary, J. Lewis, K. W.-D. Ma, and B. McWilliams. The shattered gradients problem: If resnets the answer, then what is the question?, 2018. URL bRttps:/://axx
org htt//bb/17702 08591
S. Brody, U. Alon, and E. Yahav. On the expressivity role of layernrm in transformers' attention. arXiv preprint 2335 230..23200000 2023.
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert- Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language Models are H. Larochelle, M. Ranzato, Hadsell, M. Balcan, and H. Lin, edi- tors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Cur- ran Associates, Inc., 2020. URL http://////////oooooee/e.e.re...
Few-Shot Learners.
files/
J. Chung. C. Gulcehre, K. Cho, and Y. Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling, 2014. URL/attps.//arxiv.orr/////
3555
H. Daneshmand, Kohler, Bach, T. Hofmann, and A. Lucchi. Batch normalization provably avoids rank collapse for randomly initialised deep networks, 2020. URL https://arxiv: org/abs/ 200..011000312222222
Dao and A. Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024.
J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional trans- formers for language understanding. 2019. URL/https: httptpps:://axi...e/
http/rixvivv///ee www.ariww...g..////aaaaaa 1180://411100005 04805
Y. Dong. J.-B. Cordonnier, and A. Loukas. Attention not all you need: Pure attention loses rank doubly exponentially with depth, 2023. URL http///axvvvvvv.///vv//ee./...av www./rrgg..............e.......e /atbb///213.. 2t1//2113...0 03404
Y. Fang, K. Loparo, and X. Feng. Incqualities for the trace of matrix product. IEEE Transactions on Automatic Control, 39(12): (12):24489 2499, 1994. doi: 10.1 00.10099.36284...
D. Y. Fu, T. Dao, K. K. Saab, A. W. Thomas, A. Rudra, and C. Ré. Hungry hungry hippos: Towards language modeling with state space models, 2023, URL 14052
www. aarg/bbb/2212.
B. Geshkovski, C. Letrouit, Y. Polyanskiy, and P. Rigollet. A mathematical perspective on trans- formers, 2024. httpss/ts://axii..vo org. abs 2tt./2312.1100 10794
A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024. URL fictps://arxiv wttpav/.rrr//a///21...ooo
A. Gu, T. Dao, S. Ermon. A. Rudra, and C. Re. Hippo: Recurrent memory with optimal polynomial projections, 2020. URL Utl.hhttss///axiiit www.:orrggssso.o 22000 0008 07669

---

Published as a conference paper at ICLR 2025
A. Gu, I. Johnson, K. Goel, K. Saab, T. Dao, A. Rudra, and C. Ré. Combining recurrent, volutional, and continuous-time models with linear state-space layers, 2021. URL https:
A. Gu, K. Goel, A. Gupta, and C. Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems, 35:35971-35983, 2022a.
A. Gu, K. Goel, and C. Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference Learning Reprexentations (ICLR), 2022b.
B. He, J. Martens, G. Zhang, A. Botev, A. Brock, S. L. Smith, and Y. W. Teh. Deep transformers without shorteuts: Modifying self-attention for faithful signal propagation, 2023. URL https:
7/arxiv. wwg/abs.aaa.ooo.o... 302 10322
K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition, 2015, URL pettps//aaxsen org/ / abs /ass/1122 1512. 03385
K. He, X. Zhang, S. Ren, and Sun. Identity mappings in deep residual networks, 2016. URL
org.orgg.ee als// 603 05027
Hochreiter and Schmidhuber. Long short-term memory. Neural Comput., 98):1733-1780,, doi: 10. 1tt.////cc...99........
Nov. 1997. ISSN 0899-7667.
URL httptttppps.... //d.///d..//// .org/
10.1162/neee neco. 1997. 9.8.1....1 1735
H. Hè and M. Kabic. A unified view of long-sequence models towards modeling million-scale dependencies, 2023. /tt.S/////ttt////////rrrrra/a/// 001167
S. loffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift, 2015. URL httttttt/aaaaaat wwwpxiiwwoo///b///52222222
A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are RNNs: fast autoregres- sive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, (CM JMLR.org, 2020.
Z. Lan, M. Chen, S. Goodman, K. Gimpel, P. Sharma, and R. Soricut. Albert: A lite bert for self- supervised learning of language representations, 2020. URL httpps://ersvv.orr//a...
1909. 11942
H. Li, Z. Xu, G. Taylor, C. Studer, and T. Goldstein. Visualizing the loss landscape of neural nets, 2018. BRLhttps:
I. Loshchilov and F. Hutter. Decoupled Weight Decay Regularization. In International Confer- ence on Learning Representations, 2019. URL http:://openrvviww http://enreeee..eeffoorrmomm met/forummii 3k<6RiCqYY
M.-T. Luong, H. Pham, and C. D. Manning, Effective approaches to translation, 2015. httttttt////riiiiie
neural machine
htsr//bs//5555 04025
L. Noci, S. Anagnostidis, L. Biggio, A. Orvieto, S. P. Singh, and A. Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse, 2022. URL https:
// //arxiv. http:/iwwrggabb////222 22w.bb//2226..e 03126
A. Orvieto, S. L. Smith, A. Gu, A. Fernando, C. Gulcehre, R. Pascanu, and De. Resurrecting re- current neural networks for long sequences, 2023. URL 06349
M.-C. Popescu, V. E. Balas, L. Perescu-Popescu and N. Mastorakis. Multilayer perceptron and neural networks. WSEAS Trans. Cir. and Sys., 8(7):599-588, July 2009. ISSN 1109-2734.
H. Shi, J. Gao, H. Xu, X. Liang, Z. Li, L. Kong, S. M. S. Lee, and J. T. Kwok. Revisiting over- smoothing in bert from the perspective of graph, 2022. URL https: /ttp///a/vvv////////////aa//awa/ org/ abs/ 2202. 08625
12

---

Published as a conference paper at ICLR 2025
J. Sieber, C. Alonso, A. Didier, M. N. Zeilinger, and A. Orvieto. Understanding the differences in foundation models: Attention, state space models, and recurrent neural networks, 2024. URL
hntpss:/aarxe
pttgabs///aeeooooeo.ggggeeee org/ abs 2405. 157311
J. T. H. Smith, A. Warrington, and S. W. Linderman. Simplified state space layers for sequence modeling, 2023. BRRh:sts:///rat
org/ abs/abas....... 22t.//2288.44933
S. Song, C. Lan, J. Xing, W. Zeng, and J. Liu. An end-to-end spatio-temporal attention model for human action recognition from skeleton data, 2016. URL
1611.0605
R. K. Srivastava, K. Greff, and J. Schmidhuber. Highway networks, 2015. URL //aaxii.org/abs//155 00387
https:
M. Stollenga, J. Masci, F. Gomez, and J. Schmidhuber. Deep networks with internal selective atten- tion through feedback connections, 2014. URL (tttttt:///axxttii org/ abs / 1407. 3068
Y. Tay, M. Dehghani, Abnar, Y. Shen, D. Bahri, Pham, Rao, L. Yang. S. Ruder, and D. Metzler. Long range arena: A benchmark for efficient transformers, 2020. URL https: //arxiv.org/abs/ /201 1.000
M. Tiezzi, M. Casoni, A. Betti, M. Gori, and S. Melacci. State-space modeling in long sequence processing: A survey recurrence in the transformer era, 2024.
abs/2408 09062]
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polo- sukhin, Attention is all you need, 2023. URL htttts://aaiiiii wwtp:/iwwwgggabb////ao/ 03762.
A. Veit, M. Wilber, and Belongie, Residual networks behave like ensembles relatively shallow networks, 2016. WtLhhtps::
06431
A. Voelker, I. Kajié, and C. Eliasmith. Legendre memory units:
Continuous-time representation in recurrent neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, d' d'Alché-Buee E. Fox, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 32. Cur- ran Associates, Inc., 2019. URL /ttp://///ttss///prreeeeee.o.oo.
cttppcpeeepeeeeeeee////c/ files/ paper/2019/Fi1 / 952285522 ataaa/aaaaaa Paw.:00---ppppppp....
X. Wu, A. Ajorlou, Y. Wang, S. Jegelka, and Jadbabaie. On the role of attention masks and layemorm in transformers, 2024a. URL SoRlhtttps://axlv.coom abs/g//bb/24555 2tt.:////2288811111118888
X. Wu, A. Ajorlou, Z. Wu, and A. Jadbabaie. Demystiffyingg in attention- based graph neural networks, 2024b. UK//Lppp..//axrrr.r///////2.....
R. Xiong, Y. Yang, D. He, K. Zheng, S. Zheng, C. Xing, H. Zhang, Y. Lan, L. Wang, and T. Liu. On layer normalization in the transformer architecture. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Pro- ceedings of Machine Learning Research, pages 10524-10533, PMLR, 13-18 Jul 2020. URL mlr https: http::///ooceeddeeee www.:mmmrr..reeeemmmm www.seressiiiii 1ttp111//xxiilbbb/ html
J. Zhou, G. Cui, S. Hu, Z. Zhang, C. Yang, 7. Liu, L. Wang, C. Li, and M. Sun. Graph neural networks: A review of methods and applications, 2021. URL https: / //ax/vivvvv..////////aa///aaaael org/abs/ 1812. 08434
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 PRELIMINARIES
We present a brief description of the sequence models considered in this work, namely attention (Stollenga al. (2014); Luong et al. (2015): Song al. (2016); Vaswani al. (2023) and recurrent
blocks in State Space Models (Voelker et al. (2019); p Gu et al. (2020) 2021
Fu et al. (2023)
and their corresponding architectural components (from Transformers Vaswani et al. (2023 and
Mamba Gu and Dao 2024
A.2 ATTENTION
Attention is considered to be the backbone of the Transformer architecture. This is thanks to its ability to capture correlations and dependencies between tokens in the sequence. In particular, the attention block can be expressed follows:
XWOWTXT
0 softmax
XW. E Rxxd
(8)
vdox
where the softmax operation is applied row-wise and dQQ IR acts as of temperature parameter modulate the strength of the interactions between tokens. Wq Iddd WK E R and Wv [dxd are the query, key and value matrices respectively, and X E Nxd is the input sequence, where d is the embedding dimension and N is the length of the sequence. In preparation for later sections, we also provide an alternative form for the attention block, as proposed by Dao and Gu 2024):
O = MV
with V XW, Nxd and M softmax XW Vdox
N for attention.
A.3 RECURRENT BLOCK
SSMs lie their foundations in linear (often input-dependent) dynamical systems, which be as linear RNNs (Orvieto et al. (2023)) ). particular, the classic recurrent block for a generic SSM
takes the following form:
hi = Agha-1 1 + B&CK Cyhua DRTE
(10)
where Ia E Rd is the input at time k. where &o 21, hy E Rd is the hidden state and on E Rd is the output the SSM layer at time We note the connection with the Transformer's notation Fe == XR.> and similarly denote O E R.vx to be the output matrix of the SSM block, i.e. Ok = OR.- Here, given a matrix A, we indicate Ak.: to be the k-th row.
One fundamental difference between attention and the recurrent block is that the attention mecha- nism does not have the capacity for memory, so the whole sequence of previous inputs X has to be fed every time step. In contrast, the recurrent block in equation 10 is able capture previous inputs' information, so only the last token has to be fed in. Yet, this expression also allows for a convolutional representation, where equation i0 unrolled in time as:
O
Σ ==
C
II 1=0
Ax-1
BJd, + DR
(11)
where A E H B2 E B C2 E Rd Dr E Rddd The first versions of the SSM blocks
consisted of Linear Time Invariant (LTI) representations, i.e., A = A, B B, C = C and D₁ D for all time steps h. Recently, in order to improve the expressivity and the performance of such models, selective State Space Models have been proposed /u and Dao (2024); : Dao and Gu 2024). In these models, matrices Ag. Be, and C input-dependent.
14

---

Published as a conference paper at ICLR 2025
As shown in Dao and Gu (2024). the SSM update rule in equation 11 can be equivalently expressed
as in equation 9 where M is lower triangular matrix with My = Cj (II Aj-i) B; and = X. In the reminder of the paper, this formulation will be used as covers both the attention and the recurrent block.
According to Dao and Gu (2024), for selective SSMs where A, = ag! (e.g., Mamba-2), the matrix M can be compactly written as M (a) @ WWWwWwyr where ISS(a) is lower-
triangular& I-semiseparable matrix, which can be expressed as:
1
01 02201
1
ISS(a) =
02
(12)
aN-1 0g
0N-1
02
ON-1
1
A.4 PROOF OF LOWER BOUND RANK COLLAPSE MEASURE
Lemma A.1. Let A RT and B RuXP matrices. Then it holds that ||AB|| 2
min! [B)] where mmin minn(B) is the smallest singular value of B.
Proof. From Fang et al. (1994). it holds that ||CD]| 2 "min [C] By the properties of the Frobenius norm, we have that ||AB|E | Hence, have
||B T A || 2 mia B ||[A" ||p. The result then simply follows from the fact that m B
minn (B).
Theorem 4.1 (Lower Bound on Rank Collapse). Let the input sequence ( be such that
(0) 2 b. If the skip connection strength À is chosen to satisfy
X² -a(SCM 111) 0.
then we can lower bound (K) by ((Y(())22 2 a (( for all K E N.
Proof. In the following, we denote v((A) = ((t) y,(,) = ((k) and (k) = p(Y(t)) for brevity.
|
11
p(++1
((() (1
11T
(M))) + ) 1
N
N
=12((r(k)
y(k)
11T
(()
2
M(())
M(
(k)
N
20
v(())
((())
Σ 21
M (8)
11 (1()
y-(A) 11
N
N
2
X2
((k)
11
y(k)
2
+ 21
(1- (1- 117)
(--N)
11
y.(A)
F
N
N
x2((r(0) r(
11
(4)
(()) (K)T
11
y(())
21
N
N
1=1
Continuing from here, we get:
((++1)² 2 22
12((1(00 y(k)
11
((k) 1
2A Σ
d
|
(1~
n
1
1117)
|
N
1=1
15

---

Published as a conference paper at ICLR 2025
=>²/ =
(k)
11T
N y (1k)
1r- 2 2 (1)
2
d d ΣΣ 1=
c y(x)' y,
=
1²
(K)
11T
N (())
1|= - (+9)
Σ
d
d Cv. y (k)7 ((1))
y:!
F
1=1
2 12
11 y (8)
1,2 -2
[1)
max
Σ
Y
(k) y;
N
F
1=
11T
N
x²
N y (8)
21
|1]
mnax/Crs
Σ
(k)
(k)
1j
2
x²
11 T
y
1
2 (4)
max
(R)
N
2 12
11
y (1)
2
2
(1)
max CC
Σ
d
(1)
N
1.j
j=1
2 x2/(1r00) - 11 y(b)
1|- F
-2
22
max (CV, .dd
N
2 22p00000 21 |[Cv|leND
where in (i) we used Stuchy-Schwarte inequality and the fact that using LayerNorm and hence every row of (A) has norm 1 and in the last step we used the definition of the rank collapse measure.
We now need to analyze the relationship between ( and p((+t1)
/[K+1)² = ((
11T D(+t))((11)
=1(11=
11
|
=
N
F
N
F
11T
((++))
|
31
12
11T
((+1))
N
| 112
I max
2
1
[1-
11T
y((+11)
(k+1)
N
st
[N
(++1)
[MX
Combining the previous two results and recalling that CM = supx || ||r and = supp |[Cvll,, we then get:
(k+1))
(22 2ANdSCM)
max
||
16

---

Published as a conference paper at ICLR 2025
We will now proceed in upper bounding maxie[N] ||
Σ
M
V + +
max [E[N]
max [E[N]
max SEIN
S max <E[N]
Σ
MM
V(02²
+1g + 2
Σ M
(b)
(k)
+1²+ + 22 Σ
(K) Σ
max (E[N]
1j
S92
max KE[N]
S
M (R)2
+22 + 22 Σ
M
(k)
Vim
my
(A)
y()))(()))
21
(S M + X2 + 220M8)
max iE[M]
S (CMS 2)2
By putting everything together, we get:
(k+1))
(SCM + 12/)3 (AP)LL)P - 2ANdSC
To guarantee that y (k+1)² 2 ap(03 for some a < 1, we can simply lower bound the right hand side. Hence, we want that (SC (22rult)² 2ANdSCM) aa((t)
Hence, we need:
aSS-m + (21)2) (k) 2 2ANdSCM
To do so, we need to guarantee that X2 a(SOM + [A|)2 > 0. Hence, if satisfies the condition
above and we have that (x)) >
X X4 + then (*++1)² 2 (1) Note that in order to
2ANdSCy
2ANdSC satisfy the condition on 1 (k)² up until time K. we need to choose (0)² ak X a(3Csr
which concludes the proof.
In the following, we provide a brief discussion on the key variables of interest in the bound above, their typical values in practice, their relationship and implications on the bound, in particular fo- cusing describing how it is possible to choose suitable values for X. the presented bound, the key variables of interest are X, and (k)). Specifically, we aim for ((Y(()) be as far from as possible, as values close to 0 indicate rank collapse. Furthermore, as we mentioned in the main text, the ideal value of a would be (since this would guarantee that the rank collapse metric is non decreasing over layers), although in order to satisfy Equation 7 this value cannot be chosen. In practice, the typical value of is 1. The key relationship between a and À is the following: in order to guarantee values of closer and closer to I (and hence to ensure higher values of the rank collapse metric at the final layer) we must choose larger values for jt]. Additionally, N represents the input sequence length, which varies based on the task. For example, N might be on the order of tens for simple question answering tasks, but it could scale to hundreds or thousands when sum- marizing long document. d instead represents the embedding dimension, typically in the order of
17

---

Published as a conference paper at ICLR 2025
tens hundreds. Regarding the selection of X, propose possible approaches: treating it as hyperparameter or making it learnable. In the first approach, À would be chosen through a standard hyperparameter optimization procedure, testing different values and evaluating their impact both performance and the rank collapse measure. While effective, this method requires multiple training to identify the best value. In contrast, the second sarroach-makinng À learnable- -oferr signif- icant advantages. It automates the process of finding an optimal A, eliminating the need for manual hyperparameter tuning and requiring only single training run. This both more efficient and prac- tical. For these reasons, we adopted the learnable approach in our experiments, as illustrated in Table
A.5 THEOREM RANK COLLAPSE SELF-ATTENTION ONLY TRANSFORMERS WITH LAYERNORM
We first define some notation that is used in the Theorem. First, we denote Φ(+) where (.....} is the inner product. Then, G represents the graph induced by the attention mask. In particular, there is directed edge from j to i, this means that token attends token j. A Quasi-Strongly Connected graph is graph G in which there exists a node from which every other node in the directed graph G is reachable (i.e. there exist a directed path connecting the two nodes). Finally, the authors consider the following assumptions:
(0)
Assumption A.1. G contains self-loops, i.e. every token in the sequence attends to itself
Assumption A.2. There exists a constant IR such that maxteN
[IIW ll2. I| K SC
Theorem A.2 (Corollary 1. Wu et al. 2024a ). Consider the architecture defined in 6 by choosing
A O, i.e. Self-Attention Network with LaverNorm and no skip connection. Let & be quasi- strongly connected graph. Under A.1 and A.2 if w is orthogonal for all and 6 z 0, there exist C >> 0 and > 0 such that Ne < and
u(Y(K)) S C(1 - )
vK
where is the diameter g meaning that tokens converge common point on gd-1 exponentially.
A.6 THEOREM ON RANK COLLAPSE IN SELF-ATTENTIO ONLY TRANSFORMERS WITHOUT LAYERNORM
Let rYs(Y)== = 11TY Note that with this notation, we can redefine p(Y) = [|re((Y]] Then, it holds that:
Theorem A.3 (Rank Collapse Transformers Dong al. 2023)
For any single-head Self-Attention
network consisting of K layers with w (k) with) (k)
w
S 3 and for term that depends
1,00
on the attention entries, we have that
3
4
res (Y S
res( Y (0))
V /dgk
which amounts doubly exponential convergence rank-1 matrix. Corollary A.3.1. Under the same conditions of Theorem A.3 it holds that
( (4r/min( NN, (iy/minN,,d.
(ees(y(K)))la
(1))
dug
Proof. Follows straightforward from the relations between Frobenius norm and the 1- and co- norms, i.c. given a matrix we have A E R ||l|| 2 ||A]], ||A]|r 2 ||Alloes |||]|r S
min(NNd)]AAl |||]|r min( Alpo-
18

---

Published as a conference paper at ICLR 2025
A.7 PROOF OF RANK COLLAPSE FOR SELECTIVE SSMS WITH LAYERNORM BUT WITHOUT SKIP CONNECTIONS
Lemma A.4. Ifa S 1, We have that M 2 Amin o sori
Proof. If 2 j. we have that:
(R) (k) = a (R) WoWBY Y 2 a Amin Amin (x xxf. a Amin o(k) M
Amio
Note that the worst case is when i - j bounding a--J 2 N-1 2 aN since I 1.
1. which leads to the desired lower bound by lower
Lemma A.5. We have that D N Anar
Vk
Proof. From the definition of LayerNorm, we have ((k)
i,
||
Furthermore, we have that
|| 1. 1l2 - ش M ( Σدك M
(k-1)
(R- 1))
1l2 N Amax.
Theorem 4.3 (Rank Collapse for selective SSMs without skip connection). Let g(k) (k) ((k)), where (:.)) indicutes the inner product. Under Assumption 4.1 and if e S (10) 1ffor some e > O R, Amin > 0, E, Mij 2 1 Vi, then holds that:
min (Y, 3:
u(Y(K)) SVN 2 a AK 2 0,
where Amin and Amax the minimum and maximum eigenvalues of W WW
Proof. For the fist part of the proof, we will follow the proof of Corollary in Wu et al. (022a
(((+1)= min (x: (k++), ((++))) = 12
(((--)) (--)) jit n..)
(JE(N) YY
(((4))())))
y-(1)
(D(C)
((()) se
kn.
10
N
(k)
(k)
(R)
N Ammax
131
A
where the last inequality follows from Lemma A.5 From here, using Lemma A.4 and the assumption that Amax S NV we have that:
(b(k+1)
M
2 Min (h y!!!))
(k)
k1=1d
N
N
= Σ
(hk
Mik
(R)
(T( (((k)
Σ
M() Rettl
(K)
(x((M) Ki.s y/())
(()1) 31
mi:!
m.!
m=1
2 2000)
2 (A))
2N +(1 2 a ((0))) ((())
2 a (1 (1-c2 A2 2N (A)
where the last two inequalities follow since E, M,, 2 Vi and ((()) 2 e respectively. Note also that by iterating the above recurrence starting from 1, it easy show by induction that since (5(0) 2 c, then (b) 2 c vkk By rearranging the terms, have that:
19

---

Published as a conference paper at ICLR 2025
1 - ((+1) $ 5(1- Nai (1- o(k))
By recursively unrolling the relationship, we get:
1 1- b(0k)) < (1 min (1-600)
I (1 0² 1min a K
since 0 < c g(0)) <1
Since by definition, 1 - (b()) 2 1 (X(())
(A) (k)
(1) | y 12/2 it follows that:
(K)
N
(( || - 11 (k))
/
|M (K)
1 TYYYY//NN
N
N
(x( 1,
( 11
2N
V (1 c. 2inn a2 2N1 K
A.8 RANK COLLAPSE FOR LTI SSMs WITH NO SKIP CONNECTIONS AND LAYERNORM Lemma A.6.
|| S ||
Proof. We have that ||Y[(++) || = || < [[]]]||(()|]) By simply solving the recur- rence we get the desired result.
We now analyze rank collapse for Structured LTT SSMs. For simplicity of exposition, here we
neglect the gating mechanism and focus the recurrent block. In particular, will analyze network consisting of K layers of stacked Structured LTI SSM blocks. We recall that for structured LTI SSM blocks, which satisfy with A₁ = a,I, we have that the matrix M takes the form (
1
a1 (2011
(( o (WW = W(k), where WA =
02
GN-1 01
aN-1 02
ON-1 1
We first consider the case where all the layers have the same parameters, i.e. it holds that ((k) W Vk for some matrix W. We then have the following result:
Theorem A.7 (Rank Collapse Structured LTI SSMs). For K layer network of only Struc- tured LTI SSM blocks where all layers have the same parametrization, it holds that ((y( s
|| ||
Proof. We have that:
(k) 17441 = (1 (1- 11
())
N
From here, by using Lemma A.6 we get:
20

---

Published as a conference paper at ICLR 2025
u(Y(k++))
||M 1mstille
/ 1 ||
11
1r S |||(kt)|||
N
Hence, whenever |[W]|r < 1, we have rank collapse for the case of structured L.T-SSMs. Note that however here the overall dynamics and behavior is different: the convergence is exponential rather than doubly exponential like in the Attention case. This due to the fact that the matrix M is not input dependent, which is what causes the double exponential convergence. We will see that for selective SSMs, which also have a similar input dependence to attention (apart from the softmax) for the matrix M, the convergence will again be doubly exponential. Furthermore, note that here the initial input sequence does not have any influence on rank collapse, whereas in the attention case this was not the case (since there was also a doubly exponential dependence on ||res(Y (Y))))))))) Again, will see similar behavior in the selective SSM setting.
For the case where we have different parametrizations for different layers, we then have: Theorem A.8 (Rank Collapse Structured LTI SSMs General). For K layer network of only Struc- tured LTI SSM blocks it holds that u(Y (K)) s 11 [| ||Y|()]]=>
we have that ||Y |F = || S [| [| Applying this recursively gives us [|Y(K]]| S II || || Then, similarly to the proof of TheoremA.7 we have:
Proof. From the proof of Lemma A.6
((Y(K+1)) = ||ar(K))(()) 17x+ille
11
| S VN /
w
N
K+1
[w())] |||1
In this case, simple sufficient condition to have exponential convergence of the rank collapse measure is that ||W())) Vk.
A.9 PROOF OF UPPER BOUND OF RANK COLLAPSE MEASURE FOR SELECTIVE SSMs WITHOUT LAYERNORM AND SKIP CONNECTIONS
Lemma A.9. For selective SSM, it holds that:
||Y(()|| (VNIWWBlle) 3h-1
||
where WBC = WCWT
Proof. We have that:
||y(++1|||= = (1 (x) < 1| ||r
( (y r(())) ee [|
= -||SSS(a) @
< [ (A) le ||
By simply solving the recurrence we get the desired result.
21

---

Published as a conference paper at ICLR 2025
Theorem A.10 (Rank Collapse for Selective SSMs). Civen selective SSM with same parametriza- tion for each layer with no skip connection, it holds that:
m(Y(1)) = (VNIIWBOlEe)
| 1
Proof. We start by observing that:
((Y(h)) = 11 I1T Mik M S ( || N
S [V]]WBclll Ball 1|
By using Lemma A.9 we get that:
16-1
u(Y S (VNIIWBolle)
||
proving the result.
A.10 PROOF OF COUNTEREXAMPLE FOR LTI SSMs
Structured LTI SSMs. Consider the system A, = A = al, B, = = and C. = C I
and let N 2 and d 2. By using the relation Mji = C, Aj f..... A4 whenever j 2 Dao
(1 (1+) 1+ X
0
and Gu 2024), we get that M + Al = LayerNorm, [Sys-1].
Let us denote this system, equipped with
a 1+
1 1 2
-2 -2 and ( I, -9.22-2
Proposition A.10.1. Given [Sys-1], if we choose 2, i.. M
then we have that (k) - O. Conversely, given the choices above with A < -2 (with A + -1, then rank collapse does not happen, i.e. p(Y(()) A O.
Proof. We will prove the theorem by induction. In partieular, we will prove that (k) will take the following form:
1
y())
(13)
av-1
Cpp.
VaLa
where ag = 1 and og satisfy the recurrence: ax+1 = axk (1
(1+1)²
+
1+1 v Vax - -I -Ivak- For
4
-1, we simply have that ax+1 2 ax (1 [1-1]= implying that ay 2 ap (1 (I+1Y=
which causes ag to diverge. Similarly, if -2 < X < 11, we have that ax+1 2 az [ + r+x i.e.
3+
24 which also diverges. This implies that (k) ->
1 0) which rank one
0% 2 0g
[+x
i 0
matrix, hence also implying that μu (+)) +
Let's consider now the case where À -2. We will prove that the sequence aj remains bounded from above. Note that this automatically proves that rank collapse does not happen, since necessary Let define
condition for this is that ag by induction that ag < Byk. Since 3 > 1. the statement clearly holds for k 0. Suppose
1+A and 3 1-,- > 1 for X < -2. We prove
now that the statement holds for k. We prove that the statement holds for + 1. We have that
akt1 3(1 + 4r2) + 4 T. We want the quantity the right hand side to be smaller than
22

---

Published as a conference paper at ICLR 2025
3. This is equivalent to saying that r2B < -rvByB -1 By squaring both sides and solving the inequality we can see that --p- satisfies this inequality.
1.
We now go to prove the above statement in Equation about the structure of y using induc- tion. Clearly, the statement true for k = 0. Now, let's s assume the statement holds for k. We will prove that then the statement holds for 1. We have that:
((++)) = (M (M+
1 1 A
O
K
1+1
a--1
(L+))Vau-++22a
1+1
Vau
VaE
Hence, after we apply LayerNorm y(k+l), by letting ? Vavav/ - we get that:
(4 + (1)) ax 4(1 +
y(++2)
[1+Alaa-- T+2y/g
1+1
VT
Choosing ax+1 (1+
1 (1+2}² at + 1#x Va I Vax clearly results being of the form
((+1))
PELL-1
Vou++
which proves the induction hypothesis.
A.11 PROOF OF COUNTEREXAMPH FOR SELECTIVE SSMS
Proposition 4.3.1.
Given [Sys-2], if we then choose
3 M as above and (
0
Ba
for any as So with og O. then we have that ((Y O. On the
Vot+B0
Vap+op
other hand, if we choose < - and M and (0) as above, then rank collapse is avoided.
Proof. Again, we proceed by induction. In particular, we will show that under the assumptions of the theorem, (k) takes the following form:
(*)
(2+ ((2+))=*
(1+ de (2+x)2* (1+
Note that if this result holds, as long as ag and + A| (which happens iff À
3)) 0). which a rank one matrix, hence also implying that (x( 0.
we have that (A) + (1 0)
On the other hand, < we have that (k), as increases, tends to oscillate between the matrices () and () -1) Although this case the limit of (k) does not exist, we have that both the matrices above have full rank and hence rank collapse does happen.
We now proceed to prove the statement above. Clearly, the statement true for = 0. Let's suppose that the statement holds for k. We will prove that it holds for k + 1. We have that under the choice of M and X for the theorem, we have that:
1+λ (2+1)* dg ((2+))
0
(l++1)
1 1+
+
(2+1) ap
(1+A)* 24
(1+A)
(22-X)
24
23

---

Published as a conference paper at ICLR 2025
1+1 (2+ 00
(1+ 8g
22+ (1+
Hence, after applying LayerNorm to ((++1,, get that:
y(k)
(2+1)-
(1+1)+1
(2+X)
(2+x)
which proves the claim and finishes the proof.
A.12 PROOF OF THE TIGHTNESS OF THE LOWER BOUND THEOREM 4.1
Proposition 4.3.2. There exist architecture of the form of Equation 6 such that, when choosing to satisfy Equation 7 i.e. [^] = (2( (2()-)). it holds that u(Y = O ( (()())))).
Proof. We again consider [Sys-2] from Proposition 4.3.1 with do 30 2. First, calcu- late the value of the rank collapse metric at every layer. Following the proof of Proposition 4.3.1
0
where we saw that (
(2+A)
(1+A)* ((
a simple calcula-
((2-))
tion leads to (( = 1+(17x) 100 We are now interested the value of a, which can be simply
calculated by a m(Y((+))) 1+(722)
++(2=xja443 It is easy to see from the previous expression that
((Y(())
[XI = =m,inn(, (++2))) = ((min(1, T+[T]) In particular, we notice that in the region where the model suffers from rank collapse the limit of infinite depth (i.e. > - )) we have that
-O( Note that this corresponds to choosing [^] 22 =a)=)) hence concluding the
proof.
A.13
ADDITIONAL EXPERIMENTS
A.13.1 COMPARISON RANK COLLAPSE FOR THE S4 ARCHITECTURE WITH AND WITHOUT SKIP CONNECTION AND LAYERNORM
In the following, we consider a different SSM model, i.e. the S4 architecture Gu et al. 2002b), to analyze the effect of skip connections and LayerNorm S4 variant with diagonal A, i.e. S4D et al. (2022a). We train the S4D architecture with 32 layers and 1.6 million parameters the Cifar10 dataset. We then sample 32 images from this dataset and forward pass them through the trained model. Again, we calculate the normalized rank collapse measure μ each layer. We repeat this for all considered cases, namely skip connections plus LayerNorm, only skip connections, only LayerNorm, and neither skip connections or LayerNorm.
different model. In particular,
the
Figure 4 shows that for S4D, even when removing the skip connection and the LayerNorm from the model, rank collapse is not observed. This an be explained by the fact that the weight matrices at different layers have a large Frobenius norm. From Theorem A.8 in order have rank collapse, we must have that the weights matrices have Frobenius norm smaller than 1. Hence, the fact that we do observe rank collapse does not violate our theory as in the case of weight matrices with large Frobenius norm, the upper bound presented in Theorem A.8 becomes vacuous. Additionally, note that while for Transformers and for Mamba-2 the observation of rank collapse depended (at least theoretically) on both the values of the weight matrices and the inputs values, for standard SSMs (e.g. S4) rank collapse only depends on the weight matrix. This renders the SSM model more "robust" to rank collapse, in the sense that to observe rank collapse of the model and cannot simply find an "adversarial" input sequence causing rank collapse (which instead could be done for every parametrization of Transformers and selective SSMs).
needs to find precise parametrization
24

---

Published as a conference paper at ICLR 2025
Normalized Rank Collapse Measure vs Layers Depth S4D
1.02
1.01
11
1.00
0.99
Skip= Norm= True
Skip=True, Norm=False
0.98
Skip=false, Norm=True
Skip=False, Norm=false
0
5
10
20
25
Layer
Figure 4: (Normalized) Rank Collapse measure function of layer depth for the S4D model. The shaded areas represent one standard deviation from the mean calculated over the 32 examples. Skip=True/False indicates Norm=True/False indicate we use S4D architecture with/without LayerNorm
S4D architecture with/without skip connection whereas Layer-
A.13.2 COMPARISON OF PERFORMANCE ON IMAGE LRA FOR DIFFERENT VALUES OF SKIP STRENGTH
In the following, we test and report the performance for different values of À for Transformers for the Image LRA task. In particular, we report table with the performance of Transformers (with layers) trained from scratch on the Image LRA dataset for different values of X, namely -2, -1, O, 1, From the results, we can make the following conclusions: we can clearly see that rank collapse is happening for the case where À O, proving the importance of having skip connections in the architecture. Indeed, the model for this value of performs random guessing, meaning that not able to build useful representations of the inputs due to rank collapse. For the other cases, the performance comparable across different values of X.
Model Transformers
À = -2 A
x=
À
x=
38.61
35.89
10.00
40.22
38.90
Figure Performance metrics of Transformer and Mamba-2 for different values of
1. EXPERIMENTAL DETAILS
For both the LRA image task and the MQAR task we use the standard code bases provided online We use the standard transformer architecture, the linear attention model proposed in Katharopoulos
al. (202)), the Mamba architecture (Gu and Dao 2024). and the Mamba-2 architecture (Dao and For the MQAR experiments we use the following training protocol:
Gu 2024
Optimizer and schedule:
Weight decay of 0.1, linear warmup with duration of 10%, AdamW optimizer (Loshchilov and Hutter] (001)). For each run, we sweep the learning rates in np. oospace(-44- 4) and train for 64 epochs. This is the same setup as in (Arora et al. 2023).
Initialization: For all models we use their standard initialization and initialize - in each layer.
range- arena
25

---

Published as a conference paper at ICLR 2025
Training duration: We use global batch size of 64.
For all runs, we use two layers (each with a sequence model and MLP, interleaved with layer normalization). The model dimensions d = 128, state dimen- sion 64, sequence length L 512, and number of KV pairs = 64 are kept constant for all four architectures.
Width and depth:
Position information:
Positional embeddings (Brown al. 2020) are used for the tention architectures, but not for the SSM architecture classes. This is the same setup as in (Arora et al. 2023).
Data: Each model is trained on 100,000 datapoints and evaluated on 3,000 datapoints. The data and its order are constant for all runs. This is the same setup as in (Arora et al. 2023
For the LRA image experiments we use the following training protocol:
Optimizer and schedule: Linear warmup with duration of 10%, AdamW opti- mizer (Loshchilov and Hutter 2019). For attention-based models we use weight decay 0.0 and learning rate 5e 4 and for SSM-based models we use weight decay 0.01 and learning rate 2e 4.
Initialization: For all models we use their standard initialization and initialize = -1 in each layer.
Training duration: We use global batch size of 64.
Width and depth: For all runs, we use four layers (each with sequence model and MLP, interleaved with layer normalization). The model dimensions d = 256 and state dimension n 64 are kept constant for all four architectures.
Data: Each model is trained on 35,200 datapoints and evaluated on 7,850 datapoints. The data and its order are constant for all runs.
Since keep the model sizes constant among all architectures and do not optimize the hyperpa- rameters, the accuracies reported in Table | are generally lower than in the literature. However, the main point of these experiments investigate if learning affects performance.
26


Reviews:
[{"review_id": "JuON1TuK6P", "reviewer": "Reviewer_gCrW", "summary": "Dao and Gu [https://arxiv.org/pdf/2405.21060] established a form of equivalence between transformers and continuous-time state-space models.  In a different development, Dong et al. [https://arxiv.org/abs/2103.03404] showed that self-attention layers without skip connections or MLPs suffer from \"rank collapse\" ― with increasing layers, the output matrix tends to rank-1, i.e., all token positions tend to the same representation.\n\nThe present submissions puts these together to show that rank collapse is a problem also for state-space models.  It shows that the skip connection provides vital protection against rank collapse, but that a weighted addition (with weight $\\lambda$ which may be regarded as a hyperparameter, or perhaps trainable) with the skip connection is more flexible.", "questions": "I have a reasonable estimate of creativity and technical depth, but it is difficult for me to assess impact.  I am not familiar with the area and my assessment has limited confidence. I would not, for example, know if rank collapse is widely appreciated within even the transformer \"community\" (if there is any such thing).  I have not seen MAMBA become that visible or widely used compared to standard transformer-based LLMs, but cannot speculate if rank collapse played a role.  $\\lambda$-tuning for robustness seems quite useful, but again I do not know the area well enough to know if, in practice, $\\lambda=1$ is frequently dangerous.  If the authors point to specific places in the paper where the above issues are discussed, or add some more motivating support, that would be helpful.\n\nA few writing style and notation nits:\n\nL156-L159 set up $X^{(k)}$ as layer input and $Y^{(k)}$ as layer output.  However, equation (1) introduces $O^{(k)}$ without explaining it will provide skipping ability in equation (2).\n\nL174-L178 There seem to be some inconsistent subscripts and superscripts. On one side we see $A^{(k)}_t, B^{(k)}_t$ etc. But just after the displayed equation, for LTI systems we see the superscript $(k)$ disappear, without an explanation if this is because the LTI system is assumed to have one layer.\n\nL888-L903 While setting up expressions and bounds with so many variables, it helps to afterward highlight the most import 1-2 variables, and give qualitative connections between their typical values in practice and the implications on the bounds.  E.g., how easy or difficult would be to choose an acceptable $\\lambda$ in a typical LLM?  Also, some of the definitions like $S$ are very far from the proofs in the appendix.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Reasons to accept:\n* Identifies rank collapse problem in state-space models like MAMBA, similar to earlier discovery of this problem in transformer-type networks.\n* Identifies skip strength parameter $\\lambda$ as an important knob to limit the damage of rank collapse.", "weaknesses": "Reasons to reject:\n* Given the two papers on which this paper builds, it might be argued that the present work is relatively incremental.  (That being said, I appreciate the candor while setting up the contributions of this paper, and I learnt something from it.)", "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": "COGe0FpNNN", "reviewer": "Reviewer_Av2c", "summary": "This paper analyzes the rank collapse of SSM due to identical $\\lambda$ skip connections. The authors provide a rigorous convergence rate for the rank collapse and offer sufficient guarantees to prevent it. Experimental results demonstrate the effectiveness of their analysis.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The lower boundary of the rank collapse of $\\lambda$ skip connections is analytically derived. The results agree well with empirical analysis.\n2. The paper presents the convergence rate in the absence of skip connections, contributing valuable insights.", "weaknesses": "1. The authors analyze the $\\lambda$ skip connections. However, the skip strength $\\lambda_k$ may vary on different layers.The paper should discuss how the findings hold up under these varying conditions. Additionally, many models implement skip connections selectively across layers rather than uniformly. A discussion on the generalizability of the results would enhance the paper.\n2. Theorem 4.1 paves the way to choose suitable $\\lambda$. However, in Figure 2, it appears that when $\\lambda$ is sufficiently large, the rank collapse index shows little variation. Clarification on how to determine the optimal value of $\\lambda$ would be beneficial.\n3. Based on theorem 4.1, could the authors explore adding constraints to the parameters to optimize $C_M$, $S$ and $c$ for improved neural network performance?", "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": "nisUTXTL8v", "reviewer": "Reviewer_jNQJ", "summary": "This paper examines the phenomenon of rank collapse in general sequence model architectures, including transformers and state space models. To mitigate this issue, the paper proposes a parameterized version of the skip connection that multiplies the residual stream by a constant factor. Theoretical analysis identifies the conditions on the parameter sufficient to prevent rank collapse, and an analytical example demonstrates that neither the absence of skip connections nor the standard implementation prevents rank collapse. Finally, empirical evaluations support the findings of theoretical analysis.", "questions": "The rank collapse metric is not normalized in the definition. Would it be enough to lower bound the rank collapse metric, when the norm itself evolves across layers?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper addresses the significant issue of rank collapse in sequence model architectures. It offers both theoretical analysis and empirical evaluation to support the proposed architectural component aimed at resolving this problem. I like the remark that provides the parameters corresponding to the practical architectural settings.\n\nAdditionally, the theoretical development and overall presentation of the paper are commendably clear and well-structured.", "weaknesses": "The theory investigates the sufficient conditions for preventing rank collapse in the worst-case scenario. This could imply that the required conditions are overly stringent.", "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": "kkuegMqqHc", "reviewer": "Reviewer_LT9i", "summary": "This paper addresses rank collapse, a phenomenon where embedding vectors in deep learning models converge to a uniform state. Building on previous studies that focused on transformers, this paper extends the analysis to State Space Models (SSMs). The study employs theoretical and empirical analysis to demonstrate how lambda-skip connections, LayerNorm, and gating mechanisms contribute to both the stability and expressivity of transformers and SSMs.", "questions": "Q1. Between Eq. (3) and Eq. (6), there is ambiguity regarding the residual term, specifically whether X or V serves as the residual component. This inconsistency could impact the theoretical derivations that follow. Could the authors clarify this definition? Additionally, using the same symbol D for both SSM and LayerNorm contexts creates potential confusion. Distinct notations would enhance clarity.\nQ2. The theoretical conditions for λ appear to be conservative compared to empirical findings. Could the authors explain this discrepancy? Furthermore, the appendix notes cases of rank stability without skip connections, which might challenge the theory. An analysis of these cases would be valuable.\nQ3. Could the authors provide additional experiments showing the model’s downstream performance as a function of layer depth and skip strength? Also, would the inclusion of alternative metrics, such as effective rank, offer a more comprehensive assessment of rank collapse?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "S1. The paper tackles the problem of rank collapse, extending its analysis from transformers to SSMs.\nS2. Through theoretical proofs, the paper demonstrates that lambda-skip connections prevent rank collapse, preserving model expressivity in both transformers and SSMs.\nS3. Experimental results show that lambda-skip connections and other components enhance expressivity and stability across different model architectures.", "weaknesses": "W1. The definition of the residual term between Eq.(3) and Eq.(6) is inconsistent, with ambiguity around whether X or V serves as the residual term. This inconsistency impacts the theoretical derivations that follow and should be clarified to ensure precise interpretations. Additionally, certain symbols, such as D, are used in both the SSM and LayerNorm contexts but represent different meanings. Distinct notation would improve readability and reduce potential confusion.\nW2. While the experiments generally align with the theoretical predictions, some disparities remain unaddressed. For example, the theoretical threshold for λ appears more conservative than the empirical results suggest, and additional clarification would help. Further, the appendix notes rank stability even without skip connections, which might challenge the presented theory.\nW3. The paper primarily focuses on rank collapse within the model’s architecture but does not connect this phenomenon to downstream task performance. Adding experimental results that measure downstream task performance in relation to model depth and skip connection strength could provide a more comprehensive assessment.", "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)