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

We thank the reviewer for their careful comments and appreciate their recognition of our work's "sound'' and "interesting'' approach. We want to add a few points regarding the listed weaknesses:

1. Our Markov matrix model is used because of its known limiting spectral distribution; see Theorem 2. Although it captures the attention matrix in the first layer only, the idea of studying the spectrum of attention matrices at initialisation can certainly be extended to deeper layers as well. In fact, we are currently working on generalising our analysis by allowing more complicated Markov matrix models to model attention in deeper layers. The main issue would be to relax the assumption of independence of pre-normalisation attention entries and still obtain an analogue of Theorem 2. Such an analysis is complicated by the need to develop new results in random matrix theory about this broader class of random Markov matrices.
2. Thanks to your comment, we have now relegated the previous Figure 6 to the appendix as we are convinced it was misleading. To clarify, we do not make any claim about the possible benefits of our fix to the "training accuracy'' once training does take place. It is the very "trainability'', i.e. the success rate of the training process in converging to any reasonable local min, which is at risk in the presence of rank collapse and exploding gradients at initialisation. We have now added Figure 12 which shows how our proposed fix enables training in a simple case (whilst without it no training occurs), leaving a large-scale verification of its effectiveness in real-world transformers for future investigations. We, however, included in our related work section pointers to empirical works exploring this idea, such as [1], [2], and [3].
3. We appreciate your meticulous reading and raising of this interesting point. First, we wish to emphasise that the soundness (or lack thereof) of a certain scaling depends on the scaling of the input signal. It is only under the assumption of $\mathbf{X}_0 \mathbf{X}_0^\top = \mathbf{I}$ that we show that Xavier's initialisation of key and query matrices leads to degenerate attention. Therefore, wherever Xavier is used in practice, the input covariance matrix is scaled differently, e.g. $\mathbf{X}_0 \mathbf{X}_0^\top \approx T\mathbf{I}$. Now, back to your question, if we stick to our assumption of isometric input, the same degenerate behaviour is also observed in deeper layers. That is because Xavier initialisation generates uniform attention at the first layer $\mathbf{A}_1 = \frac{1}{T} \mathbf{1}_T$ which forces $\mathbf{X}_1 = \mathbf{A}_1 \mathbf{X}_0 \mathbf{W}^{V}_1 = \frac{1}{T}\mathbf{1}_T \mathbf{X}_0 \mathbf{W}^{V}_1$ to be rank-one. Considering the key-query products at the second layer $\frac{\mathbf{X}_1 \mathbf{W}^{Q}_2 {\mathbf{W}^{K}_2}^\top \mathbf{X}_1^\top}{\sqrt{d_q}},$ one observes that the problem is all the more exacerbated. Not only $\mathbf{X}_1$ is not scaled up compared to $\mathbf{X}_0$, but also it is projected down to one dimension.

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

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

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

We thank the reviewer for their careful comments and support of both the novelty and clarity of the analysis. The points they raised as weaknesses are addressed at some length in our general response. To summarise:

1. Thanks to your comment, we have now added a plot (Figure 1c) showing the absence of any spectral gap in linear transformers. Furthermore, although our mathematical statements are proven for a simple model, we have provided empirical evidence that (i) rank collapse (in depth and width) occurs in other variants (Figure 2) and even real-world models (Figure 3), and (ii) the insight drawn from our analysis is indeed relevant for practical purposes. In the related works section, we have included citations to recent empirical studies that suggest a similar "mean subtraction'' method as evidence of our work's relevance and applicability.

2. We appreciate the reviewer's recognition of the theoretical significance of our work. We have now included plots (Figure 3) demonstrating the occurrence of rank collapse (in depth and width) in real-world architectures. We believe it is a great addition to our manuscript and thank the reviewer for raising this point.

We hope the reviewer is content with our edits and explanations and will consider raising their score.

We thank the reviewer for their helpful comments and support of our rigorous analysis. Regarding the listed weaknesses:

• We admit that our model does not capture the full complexity of state-of-the-art transformer architectures used in practice. We would like, however, to highlight that analysing even this simplified model posed significant technical challenges. Additionally, the model is sufficiently meaningful to provide relevant insights about the real case. Finally, this work is by no means a conclusive analysis of transformers at initialisation and further extensions are currently being developed, based on the framework set by this manuscript. That said, the harmful consequences of the spectral gap are empirically shown to persist even in the presence of other mitigating architecture designs (see Figure 4); removing this spectral gap has a clear promising benefit.
• Thank you for raising this point. As we discussed in a remark at the beginning of Section 2, we have chosen (non-scaled) Gaussian initialisation for value matrices based on the fact that softmax attention is essentially a downscaling. In fact, if one initialises the value matrices using Xavier or He scheme, one would get an even faster collapse of the stable rank. In practice, this is mitigated by the use of other modules such as LayerNorm and skip connections as we have discussed at the end of Section 3; see Figure 8b.
• Our paper has a clear focus on the theory. While we acknowledge the need for further large-scale implementations, we believe such a task falls out of the scope of the current manuscript. Instead, we point the reader to recent results by [1], [2], and [3], where similar adaptations to the attention mechanism are explored with a more experimental focus. [1], in particular, consider the same removal of the leading spectral component and give a substantial experimental section exploring its efficacy. Our theoretical contribution complements their more experimental results. We believe that both the theory and large-scale simulations are substantial enough to warrant separate investigations and acceptance at this venue.
• As pointed out in our general response, our paper contributes to a body of work concerned with neural networks at initialisation and does not address the training dynamics altogether. Our study, however, has consequences for the success of the training process, i.e. trainability, and can be used to decrease the risk of failed training. This is highlighted in the initial excessive stability of the network as a consequence of rank collapse which implies the network is approximately projecting all inputs onto a one-dimensional subspace. Moreover, the linear growth of the gradient norm with token length $T$ in even a single layer suggests that exploding or vanishing gradients can be an expected consequence of the spectral gap. Both of these unfortunate consequences of the softmax-based attention can be resolved by removing the spectral outliers. A more comprehensive study of the spectral properties during or after training can be further explored in follow-up works.

Regarding the question:

• We have added a plot to the main body (Figure 3) to show that famous (untrained) transformers suffer from rank collapse in both width and depth. We believe it to be a great addition to our manuscript and thank the reviewer for raising this point. However, as we discussed in our general response, (pre-)trained networks are not the subject of our study.

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

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

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

We appreciate the reviewer’s feedback and would like to highlight the strengths of this paper to ensure they are not overlooked:

• We introduce a mathematical framework to analyse transformers based on free probability and random matrix theory, that can open avenues for further theoretical developments.
• This framework proved effective in our case as we could (i) recover previously known phenomena like vanishing gradients and rank collapse in depth (Propositions 3 and 4), (ii) unveil for the first time the phenomenon of rank collapse in width (Proposition 3), (iii) quantify exactly how the latter exacerbates the former (Proposition 3), (iv) propose a simple fix that comes across as natural through this new perspective (Section 2.2), and (v) show its positive impacts on the information flow in our transformer model at initialisation (Propositions 5 and 6). Proposition 1 demonstrates how an attention layer fits exactly within our theoretical framework with no further modification, highlighting its practical relevance.
• All statements are rigorously proven and empirically verified, making a strong case for the set of techniques used in our proofs to be useful and to inspire other researchers in the field.
• Our experiments also explore the potential universality of our findings when more modules like layer normalisation or residual connections are integrated into the transformer block (Figure 4). At the same time, we transparently address the limitations of our study, offering insights into why generalising our analysis to deeper layers would require non-trivial efforts that we hope this foundational work will facilitate.

Regarding individual mentioned weaknesses:

1. We appreciate your suggestion, however, reviewers Ua2p and Htgk have found the paper logically organised and well-written.
2. The phenomenon of rank collapse is indeed observed in real-world transformers at initialisation, as shown in the original paper that identifies rank collapse (in depth); see Figure 2 in [1]. Thanks to your comment, we have added a similar plot (Figure 3) to the paper showing rank collapse in depth and width for a few standard real-world transformers, e.g., BERT and its variants.
3. As pointed out in our general response, this paper has a theoretical essence and therefore does not claim nor attempt to yield comprehensive empirical validation across tasks. We believe that the foundational theory developed in this manuscript is sufficient on its own and fits well within the aim and scope of ICLR and would encourage the community to take it further and give it an empirical dimension. Considering all the comments, we decided to move the training accuracy plots to the appendix (Figure 11 in the current version) as they were distracting the reviewers/readers from the primary focus of our work: initialisation. We have also added plots (Figure 3) showing rank collapse, both in width and depth, indeed occurs in real-world architectures and updated the section title to "Experiments and Further Insights". Regarding the inference speed, the reviewer may find it interesting that by controlling the spectra of attention matrices via our fix, the signal entries are also regulated in magnitude. This, in turn, enhances the effectiveness of the quantisation of the signal using fewer bits, as discussed in Section 3.7 of [2].
4. Thank you for raising this question. We have added a panel to Figure 1 to showcase the absence of outliers in the spectrum of a linear attention matrix. Figure 9 also exhibits the spectra of other attention variants, e.g. sigmoid- and ReLU-based, described in [3]. Interestingly, these last two exhibit a spectral gap that our theoretical framework could now fully determine, explain and fix.

We hope these clarifications will help the reviewer appreciate our work and revise their score.

[1] Dong, Y. et al. (2021). Attention is not all you need: Pure attention loses rank doubly exponentially with depth.

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

[3] Wortsman, M. et al. (2023). Replacing softmax with relu in vision transformers.