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

Authors: Federico Arangath Joseph, Jerome Sieber, Melanie Zeilinger, Carmen Amo Alonso
Primary area: learning theory
Keywords: Rank Collapse, Skip Connections, Sequence Modeling Architectures

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.

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

Reviews and discussion:
<thread>
creation date:
1730530783206

thread_type:
official_review

<thread_object>
creator: reviewer_lt9i

<content>
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.

strengths:
S1. The paper tackles the problem of rank collapse, extending its analysis from transformers to SSMs.
S2. Through theoretical proofs, the paper demonstrates that lambda-skip connections prevent rank collapse, preserving model expressivity in both transformers and SSMs.
S3. 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.
W2. 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.
W3. 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.

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.
Q2. 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.
Q3. 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730720839328

thread_type:
official_review

<thread_object>
creator: reviewer_jnqj

<content>
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.

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.

Additionally, 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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730725296579

thread_type:
official_review

<thread_object>
creator: reviewer_av2c

<content>
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.

strengths:
1. The lower boundary of the rank collapse of $\lambda$ skip connections is analytically derived. The results agree well with empirical analysis.
2. 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.
2. 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.
3. 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?

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730741216632

thread_type:
official_review

<thread_object>
creator: reviewer_gcrw

<content>
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.

The 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.

strengths:
Reasons to accept:
* Identifies rank collapse problem in state-space models like MAMBA, similar to earlier discovery of this problem in transformer-type networks.
* Identifies skip strength parameter $\lambda$ as an important knob to limit the damage of rank collapse.

weaknesses:
Reasons to reject:
* 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.)

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.

A few writing style and notation nits:

L156-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).

L174-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.

L888-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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732055151018

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We would like to thank all the reviewers for their time and effort evaluating our paper. We believe the insightful reviews helped us to greatly improve the paper. The main contents of the rebuttal are clarifications on the novelty and the impact of our work, a discussion on the tightness of our bound in Theorem 4.1. and how to choose $\lambda$ in practice and extended evaluations on the Image LRA benchmark for different values of $\lambda$. We have addressed all the concerns that have been raised by the reviewers and we hope they find our answers, modifications and explanations satisfactory. Furthermore, we have introduced the suggested changes in the updated version of the paper. The reviewers can find the newly introduced or modified parts highlighted in blue.

**Novelty and Impact**: In response to the reviews, we clarified our theoretical contribution, its novelty and the practical importance and impact of rank collapse. To the best of our knowledge, Theorem 4.1. represents the first instance of a general lower bound on rank collapse being proposed. Additionally, we offer for the first time a mechanistic solution to prevent rank collapse by introducing a skip strength parameter, which can be seen as a new architectural component. Models suffering from rank collapse have both limited inference and training capabilities, caused by the reduced representational capacity. Hence, since State Space Models and Mamba in particular are gaining significant traction and rank collapse has not yet been studied for these models, we think it is of great importance to shed light on when this phenomenon arises and how it can be prevented. Finally, the introduction of a learnable parameter $\lambda$, while it does not influence the cost of training, improves models’ expressivity and stability.

**Tightness of lower bound**: Reviewers expressed some concerns relatively to the fact that there is a discrepancy between the lower bound we presented in Theorem 4.1. and experiments. To address this, we added a new paragraph (section 4.2.3) in the updated version of the paper, where we discuss and show our lower bound is tight. The discrepancy with experiments can then be explained by the fact that the lower bound is a worst-case guarantee and hence it accounts for all possible input sequences. On the other hand, it is likely the case that for most input sequences much lower values of $\lambda$ are enough to prevent rank collapse from happening, as we showed in the experimental section.

**Additional experiments**: In response to the reviews, we performed additional experiments exploring the relationship between different values of $\lambda$ and the performance of models on the Image LRA task. In particular, we train Transformers for 5 different values of $\lambda$ from scratch on the task. We then report the final performance for the different values of $\lambda$. We note that for $\lambda=0$ the model performs random guessing, hinting that in this case the model suffers from rank collapse and hence showing the importance of including skip connections in the architecture. The models with the remaining values of $\lambda$ instead achieve similar performance. 

**Finding the optimal $\lambda$**: In response to the reviewers, we clarified both how we can select acceptable values of $\lambda$ (i.e. satisfying the bound in Equation 7) and optimal values of $\lambda$. Although a concrete and precise notion of optimality would need to be defined to discuss this rigorously, in both cases there are two ways in which we can address these two tasks: the first is by letting $\lambda$ to be a hyperparameter whereas the second is to make $\lambda$ learnable. We think that the second option is more efficient and practical, since it automates the search for an optimal $\lambda$, eliminating the need for manual hyperparameter tuning and requiring only a single training run. 

**Normalization and choice of rank collapse metric**: In response to reviews, we address the questions related to our choice of rank collapse metric and in particular the possibility of using a normalized version of the metric or other metrics (such as effective rank) to assess rank collapse. First, we clarify that our choice of metric was driven by the fact that this metric was used in previous papers addressing rank collapse. Hence, employing this metric enabled us to compare our findings with previous results. Furthermore, we think that the use of a normalized metric would be interesting but this metric would still be equal to ours in settings where LayerNorm is deployed.

**Typos and notation**: In response to the reviews, we addressed their concerns about typos and notation individually. We would like to bring to the attention of all reviewers the changes introduced in Equation 6, which lead to slight modifications of some constants in Theorem 4.1. and Equation 7.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732520320169

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Any additional questions from the Reviewers on our Rebuttal?

comment:
Dear Reviewers,

We would like to thank you again for your thoughtful reviews and valuable feedback. As the rebuttal period is coming to an end, we would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any additional questions about our rebuttal.
We would be happy to do any follow-up discussion or address any additional comments.

Best regards,   
The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper discusses the phenomenon of rank collapse in LLMs, extending the study done on transformers to the SSM architecture. This issue is highly motivated given the popularity of SSM-based LLMs. The reviews note that the paper provides valuable insights to the problem coupled with a theoretically backed solution to it. This solution is mentioned to be practical and shown to be effective with convincing empirical evidence. 

The reviews did mention some vague parts in the paper requiring clarifications and suggested ways in which the analysis can be made more thorough. From the discussion, these issues are either resolved by the authors explanations or appear to be minor. The changes required to integrate the needed parts of the discussion into the paper are quite minor and can be made towards a camera-ready version. 

I believe the paper’s strengths outweigh the relatively minor weaknesses pointed out in the reviews, leading me to recommend accepting it to ICLR.