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

TL;DR

This paper introduces lambda-skip connections as an architectural component that guarantees prevention of rank collapse in sequence models, extending the theory from transformers to State Space Models and validating with experiments.

Contribution

It provides the first general theoretical guarantee for preventing rank collapse and extends the analysis to SSMs, highlighting the importance of skip connections.

Findings

Lambda-skip connections prevent rank collapse under certain conditions.

Architectural components like skip connections and gating mechanisms are crucial.

Theoretical analysis applies to both transformers and SSMs.

Abstract

Rank collapse, a phenomenon where embedding vectors in sequence models rapidly converge to a uniform token or equilibrium state, has recently gained attention in the deep learning literature. This phenomenon leads to reduced expressivity and potential training instabilities due to vanishing gradients. Empirical evidence 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 models, 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 framework that captures both architectures. We study how a parametrized version of the classic skip connection component, which we call \emph{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 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 architectural components such as skip connections and gating mechanisms in preventing rank collapse.