Research map / Architectures and efficiency / Attention and transformers
Lambda-Skip Connections: the architectural component that prevents Rank Collapse
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Extends rank-collapse theory from transformers to SSMs via a unified framework and proves that a learnable lambda-skip connection (scaled identity skip) with LayerNorm can prevent rank collapse under a sufficient condition on lambda.
Problem
Rank collapse (token representations converging to a uniform/rank-1 state) is well studied for transformers, but unexamined for state space models; it reduces expressivity and can cause vanishing gradients.
Approach
Parametrize the skip connection as lambda*X + O, and derive a lower bound on the rank-collapse measure showing that when lambda^2 - a(SC_M+|lambda|)^2 > 0 the measure decays at most at rate a per layer, unifying transformers, LTI SSMs and selective SSMs; complemented by examples and ablations showing gating and LayerNorm also help.
Key results
- Rank collapse measure of pretrained Mamba-2 (2B) with lambda=1 still collapses in early layers; small |lambda| leads to collapse while sufficiently large or negative lambda avoids it.
- Removing gating and LayerNorm from Mamba-2 causes instability (sudden increase of mu), and removing gating with LayerNorm leads to mu approaching zero in later layers.
- Learning lambda as a trainable parameter does not hurt accuracy on LRA image and MQAR and sometimes improves it.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Tune skip strength away from small values, including the default λ = 1 | Mamba-2 without gating · Wikipedia excerpts | normalized rank collapse measure | Helped |
| Use λ values very close to 0 instead of a tuned skip strength | Albert and Mamba-2 with LayerNorm · Wikipedia excerpts | normalized rank collapse measure | Hurt |
| Use negative rather than positive skip strength | Albert and Mamba-2 with LayerNorm · Wikipedia excerpts | normalized rank collapse measure | Helped |
| Make λ learnable instead of fixed | Two transformer-based and two SSM-based architectures · LRA image task and MQAR | accuracy | No clear effect |
| Remove skip connections (ablation) | Transformer with LayerNorm · Theoretical rank-collapse analysis | rank collapse measure | Hurt |
| Remove skip connections (ablation) | Selective SSM with LayerNorm · Theoretical rank-collapse analysis | rank collapse measure | Hurt |
| Remove both skip connections and LayerNorm (ablation) | Transformer · Theoretical rank-collapse analysis | rank collapse measure | Hurt |
| Remove both skip connections and LayerNorm (ablation) | Selective SSM · Theoretical rank-collapse analysis | rank collapse measure | Hurt |
| Remove gating while retaining LayerNorm (ablation) | Mamba-2 · Wikipedia excerpts | normalized rank collapse measure | Hurt |
| Remove LayerNorm while retaining gating (ablation) | Mamba-2 · Wikipedia excerpts | normalized rank collapse measure | Mixed |
| Remove both gating and LayerNorm (ablation) | Mamba-2 · Wikipedia excerpts | normalized rank collapse measure | Mixed |
| Use a sufficiently strong positive skip strength instead of λ > -3/2 | Selective SSM [Sys-2] · Analytical example | rank collapse measure | Hurt |
Surprising findings
- Negative lambda values keep the rank collapse measure higher than positive ones.
- Gating mechanisms, designed for memory, also play a crucial role in preventing rank collapse.
Limitations
- Gating mechanisms and MLPs are excluded from the theoretical analysis, so bounds are not tight for Mamba.
- The sufficient condition on lambda is conservative; in practice much smaller |lambda| suffices.
- Only the rank-collapse measure mu is studied; no formal necessary condition is proven.
Experimental setup
Datasets: Wikipedia excerpts (32 samples, >=128 tokens) Long Range Arena image task MQAR
Benchmarks: Long Range Arena (image) MQAR
Metrics: normalized rank collapse measure mu accuracy
Baselines: standard skip connection (lambda=1) lambda=0 (no skip) ALBERT
Builds on: rank collapse analysis of transformers (Dong et al. 2021) Transformers are SSMs / structured state space duality (Dao & Gu 2024) Mamba ResNet skip connections
Closest papers by meaning
- On the Role of Attention Masks and LayerNorm in Transformers (NeurIPS 2024 · similarity 0.85)
- The Persistence of Neural Collapse Despite Low-Rank Bias (NeurIPS 2025 · similarity 0.85)
- Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Attention Layers (ICML 2025 · similarity 0.85)
- Maestro: Uncovering Low-Rank Structures via Trainable Decomposition (ICML 2024 · similarity 0.85)
- Implicit bias of SGD in $L_2$-regularized linear DNNs: One-way jumps from high to low rank (ICLR 2024 · similarity 0.85)
- Neural Collapse is Globally Optimal in Deep Regularized ResNets and Transformers (NeurIPS 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.