Lambda-Skip Connections: the architectural component that prevents Rank Collapse

Thank you very much for the helpful feedback and pointers! We have reported the suggested changes in the updated version of the paper. In what follows we provide a further discussion on the points that were raised in the review:

Question 1. We very much thank the reviewer for pointing this out. The correct definition is the one in Equation 3, whereas the correct formula for Equation 6 is $Y^{(k+1)} = D^{(k)}(M^{(k)}V^{(k)}+\lambda Y^{(k)})$. We have corrected Theorem 4.1 and Equation 6 according to this change, the final result still holds with only minor corrections in the constants in the final bound and assumption. We reported the corrected version of the Theorem in the updated version of the paper. The corrected formulas are marked in blue. In particular, the following changes have been made: in the definition of the variable $b$ in Section 4.1., the denominator is $2 \lambda N d S C_M$ instead of $2 \lambda N S^2 C_M$ and the denominator becomes $\lambda^2-a(SC_M+|\lambda|)^2$ instead of $\lambda^2c^2-aS^2(C_M+|\lambda|^2)$. Equation 7 follows the change in the denominator of $b$. We apologize for the confusion arising from using the variable $D$ for both SSM and LayerNorm and we agree with the reviewer that this was not the best choice of variables. We have changed this in the updated version of the paper, introducing different variables for the SSM and LayerNorm to improve clarity.

Question 2. We thank the reviewer for raising this point. The main reason for the discrepancy between our theoretical predictions in the lower bound of Theorem 4.1 and the experiments is that the guarantee we propose is a worst case bound. This means that it takes into account all the possible input sequences that satisfy the assumptions stated in the Theorem. To clarify more the tightness and stringency of our bound, we have added in the updated version of the paper (in Section 4.2.3) an analysis of the tightness of the lower bound. In particular, we find an architecture and an input sequence for which, up to constants (in particular we study the interplay between $a$ and $\lambda$, considering $C_M$ and $S$ as constants), the lower bound is tight. One could see this as an adversarial sequence for which the smallest value of $\lambda$ to avoid rank collapse in the finite layer setting is the one proposed in our Theorem. However, it is likely that for “common”/arbitrary sequences (i.e. not chosen adversarially), the value of $\lambda$ to avoid rank collapse might be much lower, i.e. on the “average case” the lower bound is less tight, thus explaining the observed discrepancy. On the other hand, without any further assumptions on the input sequences or additional properties of the architecture, it is not possible to go beyond this bound to obtain a guarantee on rank collapse not occurring. Nevertheless, we think that extending this theory to study “average case” scenarios instead of worst cases is an interesting line of research and we plan to explore this in future work.

The fact that in Figure 4 in the appendix we do not observe rank collapse even in the absence of skip connection can be explained by the fact that the weight matrices at different layers have a large Frobenius norm. From Theorem A.8, in order to have rank collapse, we must have that the weights matrices have Frobenius norm smaller than 1. Hence, the fact that we do not 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 one needs to find a precise parametrization 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). We have expanded our discussion of Figure 4 in the appendix by including this remark and a better explanation of why we observe this.

Thank you for the insightful comments! We address in the following the points that were raised:

Skip connection strength. This is a really good question and we thank the reviewer for having raised this point. In the case where we let all the layers have different values of the skip strength $\lambda_k$ (which for instance would be the case if we make the parameter $\lambda$ learnable), we can generalize the result in Theorem 4.1. by making the following small modification: in order to still satisfy the condition $\mu(Y^{(k)}) \geq a \mu(Y^{(k-1)})$, we can simply adapt Equation 7 to be $\lambda_k^2-a(S_kC_{M_k}+|\lambda_k|)^2>0$, where $S_k = ||C^{(k)}_V||_F and C_{M_k} = ||M^{(k)}||_F. Furthermore, the fact that we can choose$\lambda_k$, means that we can also choose different values of$a$in different layers. Consider the following example: suppose we aim to ensure that$\mu(Y^{(2)}) \geq 0.25 \mu(Y^{(0)})$. According to the setting of Theorem 4.1, the only way to achieve this is by selecting$a=0.5$and adjusting$\lambda$to satisfy the condition in Equation 7 for this specific value of$a$. However, if we allow$\lambda_k$to vary across layers, we gain additional flexibility. For instance, in the first layer, we could choose a$\lambda_k$that does not satisfy the condition in Equation 7 for$a = 0.5$(and thus cannot guarantee$\mu(Y^{(1)}) \geq 0.5 \mu(Y^{(0)})$), but instead satisfies the condition for a smaller value, for instance$a = 0.4$. In the second layer, we could then select a$\lambda_k$that satisfies the condition in Equation 7 for$a = 0.625$. This setup still achieves the desired result:$\mu(Y^{(2)}) \geq 0.625 \times 0.4 \mu(Y^{(0)}) = 0.25 \mu(Y^{(0)})$.This demonstrates the added flexibility and robustness provided by allowing$\lambda_k$ to vary across layers. However, when skip connections are not applied uniformly, greater caution is required to generalize the results, as specific lower bounds would need to be developed for cases without skip connections—an aspect we do not address in this work. That said, to the best of our knowledge, the vast majority of Transformer-based models, SSMs, and Mamba (the architectures considered in this paper) consistently apply skip connections uniformly across layers. Thus, our theoretical results remain applicable to these architectures. Nonetheless, we agree with the reviewer that extending our results to scenarios where skip connections are applied non-uniformly across layers would be an interesting direction for future research. We have addressed these important points at the end of Section 4.1. in the updated version of the paper.

Optimal value for $\lambda$. This is a great question. Currently, there is not a definition of what it means for $\lambda$ to be optimal. Does the reviewer refer to something like the smallest $\lambda$ such that the rank collapse metric at the output layer is greater than some threshold? If this is the case, the optimal $\lambda$ would be threshold dependent. Intuitively, the lower the threshold we want to achieve, the lower the optimal $\lambda$ will be (and vice-versa). One way to find this optimal $\lambda$ would be to use the lower bound proposed in Theorem 4.1: by tuning the value of $a$ we can control the threshold we want the metric to satisfy and then we can choose $\lambda$ accordingly to satisfy Equation 7 with the chosen value of $a$. However, as noted by other reviewers, the bound in some situations is too conservative and lower values of $\lambda$ are enough to guarantee that the rank collapse metric is above the desired threshold. Although in our work we do nor provide an explicit way for finding an optimal $\lambda$ (since we do not introduce any notion of optimality), another approach for finding good values of $\lambda$ is by making $\lambda$ a learnable parameter and provide good initialization for it (for example, for the experiments in Table 1, we found that $\lambda$= -1 was a good choice for initialization). Moreover, the small variation observed in Figure 2 for larger values of $\lambda$ can be explained using the lower bound proposed in Theorem 4.1. Specifically, as $|\lambda| \to \infty$, the parameter $a \to 1$. This implies 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 have included a brief discussion on this point in the updated version of the manuscript (see Section 5.1).

We thank the reviewer for the positive feedback and comments! We really appreciate that they liked the paper and found the issue or rank collapse which we target significant. In the following, we address the mentioned weaknesses and the question raised:

Overly stringent condition: The lower bound we consider in Theorem 4.1. is indeed a worst case bound and takes into account all the possible input sequences that satisfy the assumptions stated in the Theorem. To clarify more the tightness and stringency of our bound, we have added in the updated version of the paper (in Section 4.2.3) an analysis of the tightness of the lower bound. In particular, we find an architecture and an input sequence for which, up to constants (in particular we study the interplay between $a$ and $\lambda$, considering $C_M$ and $S$ as constants), the lower bound is tight. One could see this as an adversarial sequence for which the smallest value of $\lambda$ to avoid rank collapse in the finite layer setting is the one proposed in our Theorem. However, it is likely, as the reviewer suggests, that for “common”/arbitrary sequences (i.e. not chosen adversarially), the value of $\lambda$ to avoid rank collapse might be much lower, i.e. on the “average case” the lower bound is less tight, thus explaining the observed discrepancy. On the other hand, without any further assumptions on the input sequences or additional properties of the architecture, it is not possible to go beyond this bound to obtain a guarantee on rank collapse not occurring. Nevertheless, we agree with the reviewer that extending this theory to study “average case” scenarios instead of worst cases is an interesting line of research and we plan to explore this in future work.

Normalization of rank collapse metric: We thank the reviewer for the interesting question. By normalized metric, we assume the reviewer refers to something like $\tilde{\mu}^{(k)} = \frac{\mu^{(k)}}{||Y^{(k)}||_F}$ (please let us know if this is not the case and what other type of normalization they had in mind, we will be happy to discuss this further). First of all, we observe that our result in Theorem 4.1. considers an architecture with LayerNorm. This means that we have that $||Y^{(k)}||_F=N$ for all $k$, i.e. the Frobenius norm of the output of the k-th layer is always $N$. Hence, considering the normalized version of the metric leads to exactly the same result of the un-normalized one we use in this case (only scaled by a constant $N$). Hence, studying these two versions of the metric would lead to the exact same conclusions for the setting of Theorem 4.1. However, it would be interesting to study the normalized version of the metric in cases where no LayerNorm is present in the architecture, as it is likely that under these circumstances the two metrics could present different behavior and dynamics. Finally, we conclude by saying that our choice of using the un-normalized metric was driven by the fact that previous works in the literature ([1], [2]) use this version and hence studying this metric made comparing our and previous findings easier. Please let us know if this addresses the question raised, or if we have misunderstood the question.

References:

[1] Y. Dong, J.-B. Cordonnier, and A. Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth, 2023. URL https://arxiv.org/abs/2103.03404.

[2] X. Wu, A. Ajorlou, Y. Wang, S. Jegelka, and A. Jadbabaie. On the role of attention masks and layernorm in transformers, 2024a. URL https://arxiv.org/abs/2405.18781

Thank you for the insightful comments! We address in the following the points that were raised:

Novelty: We thank the reviewer for raising this point. We will make sure to clarify the novelty of our contribution in the Introduction and Conclusion sections. While it is true that some of our findings—particularly those concerning the upper bounds of rank collapse for Mamba—share similarities with results on rank collapse in Transformers presented in the two foundational papers we reference, our primary theoretical contribution, namely the lower bound on the rank collapse measure in Theorem 4.1, is entirely original. To the best of our knowledge, this represents the first instance (even within the context of Transformers) of a general lower bound on rank collapse of this kind being proposed. Additionally, we offer for the first time a mechanistic solution to prevent rank collapse by introducing a skip strength parameter. This parameter, which can be viewed as a novel architectural feature, deviates from the standard practice of defaulting the skip strength to 1. We also highlight that the simplicity of this proposed novel mechanism makes our solution straightforward to implement.

Impact: We appreciate the points brought up by the reviewer. We will add clarifications on these points throughout the paper. We have broken down the impact concern into two parts: (i) why is rank collapse important?, (ii) why is rank collapse relevant in State Space Model architectures (such as Mamba)?, and (iii) what is the usefulness of $\lambda$-skip connections?

(i) Rank collapse poses a significant challenge during both inference and training. When a model is susceptible to rank collapse, it results in tokens being mapped to nearly identical representations, severely limiting the model's expressivity and representational capacity ([1], [2], [3], [4]). Furthermore, previous studies [3] have demonstrated that rank collapse during training can lead to instabilities, primarily caused by vanishing gradient issues. Therefore, gaining a 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.

(ii) To the best of our knowledge, the Rank Collapse phenomenon is well-recognized within the Transformer community. For example, the original paper on rank collapse [5], which first highlighted the importance of skip connections in mitigating this issue, has garnered nearly 400 citations, and several subsequent studies have built upon its findings, as outlined in our Related Work section. While it is true that Mamba has not yet reached the level of prominence of Transformers—likely due to its relatively recent introduction less than a year ago [6]—it is beginning to gain traction across a variety of fields, such as biology [7]. Given this, we believe it is crucial to raise awareness within the SSM/Mamba community that rank collapse is not a phenomenon exclusive to Transformer architectures, but could also negatively impact the performance of recurrent-like architectures. In doing so, we hope our work will provide valuable insights for developing more robust SSM-based architectures, thereby expanding their potential applications.

(iii) While the choice of $\lambda = 1$ might, in some circumstances, not result in rank collapse (often when combined with LayerNorm), allowing for a learnable $\lambda$ or treating it as a hyperparameter offers several advantages. First, it enhances the model's expressivity, leading to comparable and in some cases improved performance (as evidenced by our experiments in Table 1). We remark that this is possible with minimal computational overhead: adding just one additional learnable parameter per layer is negligible compared to the vast number of learnable parameters in Transformers or SSM layers, typically ranging from hundreds of thousands to millions in the smallest cases. Second, it improves the model's stability. Specifically, initializing $\lambda$ with a negative value (e.g., $\lambda$=−1, as in our experiments in Table 1) can be viewed through a control-theoretic lens as implementing a negative feedback loop. Negative feedback systems are generally more stable compared to positive feedback systems, which is effectively the case when $\lambda$=1, as is commonly used. Therefore, we argue that introducing a learnable $\lambda$ offers significant potential benefits—enhancing stability and expressivity—at a very low computational cost. This makes it a promising addition to traditional architectures. In this work we simply lay the foundation, and we believe that additional analysis on the strength of the skip connection could be done to prevent and control the behavior of foundation models from a control-theoretic standpoint.

Dear Reviewer gCrW,

Thank you once again for your insightful feedback.

With the Discussion Phase extended until December 2, we would greatly appreciate it if you could let us know if there are any additional aspects we could address to further improve our paper. If our responses so far have not fully resolved any remaining concerns, please don’t hesitate to let us know, and we would be happy to work on addressing them during the extended timeline.

Thank you again for your time and thoughtful comments!

Best regards,
Authors