Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Transformers

Authors: Alireza Naderi, Thiziri Nait Saada, Jared Tanner
Primary area: learning theory
Keywords: transformers, attention mechanism, spectral analysis, initialisation, random matrix theory, signal propagation

Abstract:
Attention layers are the core component of transformers, the current state-of-the-art neural network architecture. However, \softmaxx-based attention puts transformers' trainability at risk. Even \textit{at initialisation}, the propagation of signals and gradients through the random network can be pathological, resulting in known issues such as (i) vanishing/exploding gradients and (ii) \textit{rank collapse}, i.e. when all tokens converge to a single representation \textit{with depth}. This paper examines signal propagation in \textit{attention-only} transformers from a random matrix perspective, illuminating the origin of such issues, as well as unveiling a new phenomenon---(iii) rank collapse \textit{in width}. Modelling \softmaxx-based attention at initialisation with Random Markov matrices, our theoretical analysis reveals that a \textit{spectral gap} between the two largest singular values of the attention matrix causes (iii), which, in turn, exacerbates (i) and (ii). Building on this insight, we propose a novel, yet simple, practical solution to resolve rank collapse in width by removing the spectral gap. Moreover, we validate our findings and discuss the training benefits of the proposed fix through experiments that also motivate a revision of some of the default parameter scaling. Our attention model accurately describes the standard key-query attention in a single-layer transformer, making this work a significant first step towards a better understanding of the initialisation dynamics in the multi-layer case.

Decision: reject
Reviewer scores: 3, 6, 5, 5

Reviews and discussion:
<thread>
creation date:
1730453811277

thread_type:
official_review

<thread_object>
creator: reviewer_smfh

<content>
summary:
The manuscript studies the self-attention mechanism at initialization through free probability and random matrix theory. Free probability deals with non-commuting random variables such as random matrices, making it a good fit for the analysis of attention matrices in transformers. The attention matrices are usually large due to numerous tokens, thus the corresponding spectral properties become predictable. In particular, the eigenvalues of softmax-attention matrix at the first layer are similar to the ones of the random Markov matrix - there is an outlier eigenvalue that equals one while the remaining eigenvalues form a bulk centered at zero with radius $T^{-0.5}$. The manuscript then links this discrepancy in eigenvalues, referred to as spectral gap, with rank collapse and exploding gradients. Fixing these issues could improve the optimization. Thus, the manuscript proposes a solution that adapts the softmax-attention matrix in order to remove the spectral gap. Namely, 1/T is subtracted from each element of the square attention matrix, where T is the matrix dimension. This solution is validated in toy experiments with multi-layer transformers. The experiments indicate that the proposed solution is effective in slowing rank collapse when the spectrum contains a single outlier.

strengths:
S1. Viewing softmax-attention matrices as random Markov matrices and analysing them via free probability seems like a sound approach.

S2. Explaining rank collapse and exploding gradient with the spectral gap (i.e. discrepancy in eigenvalues) is an interesting takeaway.

S3. The proposed solution for mitigating the spectral gap is convenient and easy to implement.

S4. Experiments indicate that the proposed solution slows the rank collapse.

weaknesses:
W1. The proposed analysis holds only for the first layer attention matrix. This limits the usefulness of the proposed analysis. Could the proposed analysis incorporate the attention matrices in deeper layers? What are the key issues?

W2. Removing the spectral gap does not benefit optimization in transformers with two layers, as indicated by the bottom row of Figure 6. What happens in the case of even deeper transformers? Does the removal of the spectral gap have any practical relevance?

W3. Figure 8 and the corresponding text raise questions about the validity of Xavier's initialization. However, this analysis is conducted only for a single-layer transformer (line 488). Can we observe the same issues with initialization in deeper transformers?

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

</thread>
<thread>
creation date:
1730624895522

thread_type:
official_review

<thread_object>
creator: reviewer_htgk

<content>
summary:
- The study reveals that softmax-based attention mechanisms during the initialization of Transformers can cause pathological signal and gradient propagation through the network, manifesting as vanishing/exploding gradients and rank collapse—where all tokens converge to a single representation at depth.
- This paper identifies and analyze rank collapse in width, which is caused by the spectral gap between the two largest singular values in the attention matrix.
- A solution is proposed to address the rank collapse and exploding gradients issues.

strengths:
- The paper is clearly written and well organized
- This paper provides a novel concept named rank collapse in width, advancing understanding of initialization issues. I.e, the gap between the largest eigen value and other eigen value bulk will lead to rank collapse, preventing the processing of input tokens.
- Rigorous analysis is provided about the reasons of gap, how gap leads to randk collapse and exploding gradients.
- A reasonable solution is proposed to deal with the rank collapse, in which, rank-one perturbation is removed from every attention matrix to close the gap. The effectiveness of this solution is analyzed in Proposition 5, 6.
- extensive experiments are provided to support the analysis.

weaknesses:
- The analysis is based on simple model structure with attention-only single-head transformer. Analysis on more complex models with multiple layers, skip connections and layernorm is preferred.
- The paper assumes Gaussian-distributed initialization for parameters, which is not always optimal in practice. Transformers often use Xavier or He initialization, will the results still hold?
- lack empirical validation. Most experiments are conducted to evaluate the rank collapse, gradient exploding and loss curve, there are limited experiments on real datasets and across different tasks.
- The paper analyzes the initialization stage of transformers without considering parameter updating during training. Spectral properties change over training, and over multi-layer transformers, which may significantly affect rank collapse and gradient stability analysis in this paper.

questions:
Does the rank collapse and gradient exploding also exist in the existing pre-trained large language models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730643047046

thread_type:
official_review

<thread_object>
creator: reviewer_ua2p

<content>
summary:
In this paper, by analyzing the spectral gap in the eigenvalues of the attention matrix, we reveal the rank collapse in width that exists in the classical transformer structure, and point out that this phenomenon leads to the disappearance of the training gradient as well as the rank collapse, which in turn causes the training instability of the transformer, and subsequently The corresponding solutions are given in this paper.

strengths:
1. The paper is logically organized, with a clear theoretical narrative and a high degree of readability
2. This paper approaches the issue of Transformer performance from a relatively rare angle, pointing out that rank collpase occurs not only as the number of layers (“depth”) of the transformer increases, but also as the “width” of the It is pointed out that rank collpase occurs not only with the increase of the number of layers (“depth”) of the transformer, but also with the increase of the “width”, and also leads to the emergence of training instability, and further pointed out the causes and solutions, which is of high theoretical significance.

weaknesses:
1. What is discussed in this paper has strong theoretical implications, but is studied in the context of an extremely simplified Transformer model, and it is not possible to determine whether the phenomenon also exists for various types of Transformer model variants, such as linear transformer. It is also difficult to determine whether the problem also occurs in transformer models that add many tricks to mitigate rank collpase in depth.
2. Taking up the weakness 1, although this paper is of high theoretical significance, it is hoped that the authors' experiments will reflect whether the various Transformers in practical use have similar problems, or, under what circumstances similar problems may affect the results.

questions:
see weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730798006051

thread_type:
official_review

<thread_object>
creator: reviewer_alfq

<content>
summary:
The paper addresses the issue of rank collapse in softmax attention. While the well-known rank collapse in depth is discussed, the authors also identify a new variant: rank collapse in width. By modeling softmax-based attention at initialization using Random Markov matrices, the paper pinpoints the root cause of the problem as the spectral gap between the two largest singular values of the attention matrices. By eliminating this spectral gap, the issue of rank collapse has been effectively resolved.

strengths:
The finding is interesting and the solution is simple.

weaknesses:
1. I propose that the current introduction section be divided into two distinct parts: an introduction section and a preliminary section. The introduction should encompass the existing situation, motivation, a succinct overview of the proposed solution, and the results obtained. The preliminary section should offer the necessary background knowledge essential for understanding the paper.

2. The phenomenon of width collapse is frequently observed in current transformer training, or it may necessitate specific data conditions to manifest. While transformers have been extensively utilized in domains such as computer vision (CV) and natural language processing (NLP) for several years, reports of width collapse remain absent. Given that the proposed method solely addresses this width collapse, I am interested in its actual efficacy in real-world applications.

3. The proposed solution involves the direct substitution of A with its variant within the attention layer. However, the performance implications of this adjustment have not been comprehensively validated. The author has provided only a comparison of training loss, indicating that this modification may influence the training loss curve. I would appreciate the inclusion of ablation studies conducted on actual transformer models in either language modeling (LM) or CV tasks. For example, in the case of LM, it is essential to supply benchmark results within controlled scenarios. Furthermore, it would be beneficial to understand how this modification may affect both training and inference speed.

4. I understand this paper only investigates the rank collapse in softmax attention transformers. However, I am curious whether this problem exists in other attention variants, such as linear attention.

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

</thread>
<thread>
creation date:
1731964256441

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for their taking the time to review the manuscript. A few general themes emerged which we think beneficial to discuss before a point-by-point response.

- The clarity of the theoretical contributions and the simplicity of the model are listed as strengths and weaknesses of the manuscript, respectively. We have chosen our model to maximise the clarity of the theoretical analysis and its main takeaway: the softmax attention mechanism has a spectral gap, which is a key contributor to rank collapse — a topic under discussion in the machine learning community. We develop a rigorous theory and show that it is consistent with the key-query model in the first attention layer. Subsequent layers exhibit greater complexity as the details of their associated Markov matrices are unknown. Achieving analogous precise results for these layers would require developing some new random matrix theory, which lies beyond the scope of this initial investigation. Experimental results, now detailed in the updated Section 3, demonstrate that more complex architectures — including famous transformer models such as BERT — also suffer from exacerbated rank collapse due to the use of softmax attention. Let us emphasize that employing simplified models is a common approach in papers investigating the theory of deep learning, as it facilitates a rigorous analysis of the issue under study; see, for instance, [1], [2], and others. Our contribution is, in fact, at the cutting edge of the analysis of transformers at initialisation. We ask the reviewers to reconsider their scores by acknowledging the complexity and elegance of the significant theoretical contributions reflected in the appendices.
    
- The lack of extensive empirical validation is highlighted as a weakness by reviewers. While we agree that such an addition would give the manuscript an extra dimension, we believe that it should be addressed separately. Notably, three recent manuscripts with a more experimental focus — [3], [4], and [5] — propose solutions similar to ours. For instance, [3] empirically explores "centred attention'', which corresponds to the removal of the spectral gap, formalised in our work; see their Section 4 for more extensive experiments that validate our proposal. We have included citations to these manuscripts to further substantiate the relevance of our work and to guide readers interested in the practical implications of our theory. While we agree that our work opens exciting avenues for future empirical studies, our findings are robust, stand-alone and empirically supported within our theoretical framework. ICLR lists "theoretical issues in deep learning" as a key topic, and our work clearly aligns with this, meeting ICLR’s standards without necessarily requiring comprehensive validation across real-world tasks. We respectfully ask the reviewers to evaluate our work solely based on its scope and content, irrespective of potential personal preferences.

- Reviewers identified as a weakness that the manuscript does not consider training dynamics. While we agree that this would be a key addition to the field, it is beyond the scope of the current work. That said, we believe the theoretical contributions and the novelty of our results will appeal to a significant number of attendees at ICLR, making the manuscript deserving of acceptance. This work addresses signal propagation in transformers at initialisation, which is the subject of a sizeable and growing body of literature; see [6], [7], [1], and [2]. Please also note that our paper introduces a set of tools — random matrix theory and free probability — that have not previously been applied to the analysis of transformers. We believe these tools have significant potential to inspire and drive future theoretical developments, as they have demonstrated in the present work.


Having clarified the above common points raised by reviewers, we now discuss individual comments accordingly and invite the reviewers to reconsider their evaluation of our work considering this note.

----

[1] He, B. et al. (2023). Deep transformers without shortcuts: Modifying self-attention for faithful signal propagation.

[2] Noci, L. et al. (2022). Signal propagation in transformers: Theoretical perspectives and the role of rank collapse.

[3] Ali, A. et al. (2023). Centered self-attention layers.

[4] Noci, L. et al. (2024). The shaped transformer: Attention models in the infinite depth-and-width limit.

[5] Ye, T. et al. (2024). Differential transformer.

[6] Hron, J. et al. (2020). Infinite attention: NNGP and NTK for deep attention networks.

[7] Dong, Y. et al. (2021). Attention is not all you need: Pure attention loses rank doubly exponentially with depth.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732478746267

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Willing to discuss

comment:
We thank the reviewers for their feedback and for taking the time to evaluate our paper. We have carefully addressed all the concerns and questions raised in the reviews and provided detailed responses. We hope that our clarifications have resolved the mentioned issues.

If there are any additional questions or suggestions, we would be happy to address them promptly. We would also greatly appreciate it if the reviewers could revisit our responses and consider updating their evaluations if the clarifications have adequately addressed their concerns.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732557264044

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Gentle Reminder

comment:
We kindly remind the reviewers to please submit their feedback and update their scores by 26 November to help us refine and improve our work to the fullest extent. Thank you!
</content>
</thread_object>

</thread>
<thread>
creation date:
1732902979904

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Summary

comment:
With no further responses received despite the extended discussion period, we will summarise our rebuttal once again. Our work presents a cutting-edge analysis of transformers' initialisation, utilising advanced results from random matrix theory and opening up numerous exciting avenues for future research. In most cases, the reviewers have mistaken these promising directions—such as the extension of our results to deeper layers, training dynamics, and extensive empirical validation—for weaknesses and grounds for rejecting the paper, despite not having identified any specific flaws in our presented analysis.

We have made every effort to explain why these emerging open questions not only highlight the significance of our work but also emphasise its potential for large-scale impact. They can be viewed as additional reasons to accept the paper, rather than excuses for rejection. **Will our contribution benefit the field by inspiring further advancements if published? We firmly answer in the affirmative, and we trust the Area Chairs will recognise this and make their final decision accordingly.**
</content>
</thread_object>

</thread>


Metareview:
metareview:
a) The paper considers the issue of rank collapse in softmax attention, which can cause training instabilities. The authors identify a new variant of rank collapse: in width, which is the focus of this paper. By modeling softmax-based attention at initialization using Random Markov matrices, the paper found the root cause of the problem as the spectral gap between the two largest singular values of the attention matrices. By eliminating this spectral gap, the issue of rank collapse is resolved.

b) The paper is well organized, with a clear theoretical narrative and easy to read. The proposed analysis is simple, the found issue is interesting and the proposed solution is also simple and easy to implement.

c) The main weakness of the paper is the lack of empirical validation. Author provided a comparison of training loss, however, it is important to validate the theoretical finding on actual transformer models in either language or vision tasks and measure what is the impact of the issue and the found solution on the final tasks. In addition the proposed analysis is simplified, based on an attention-only single-head transformer. What would happen on more realistic configurations is not analyzed, nor theoretically (which can be hard) neither empirically. 

d) The main reason for rejection are the weaknesses exposed above. The paper is well written and easy to follow, however, without an empirical validation of the proposed solution to contrast rank collapse in width, it is difficult to estimate the real impact that this contribution can have on the research community. Therefore, the theoretical contribution, without empirical evidence is not enough for acceptance especially considering that the proposed analysis has many simplifications which make it far from real setting.