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

We thank the reviewer for their thorough comments and their appreciation of our contribution. We address the points in the same order. For additional figures (indicated by Roman numerals), please see https://anonymous.4open.science/r/spectral_analysis_transformers-C633/figures_rebuttal.pdf.

1. Thank you for bringing this up. We have addressed this in detail in our response to reviewer uxct, providing extensive illustrations showing that our findings hold in practice, even without these assumptions, which are only needed for our proofs. We will make sure to clarify these theoretical limitations in our revised version.

2. Please see the response regarding $\gamma$ in the answer to reviewer uxct. A precise quantitive characterisation of the spectral bulk without these assumptions is extremely interesting and we will explore it, but this is a multi-year programme of study and we believe it is best that the results we have so far be presented to the machine learning community sooner than later. Regarding Llama3, it is interesting to consider the case of a decoder-like architecture, for which an extra layer of complexity comes in due to the causal mask in the attention, which thus falls outside the scope of this paper. Thus, we kindly refer you to our paragraph GPT-architecture in our answer to uxct.

3. Thank you for highlighting that this was not clearly explained. The reason we do not include input dependence in later layers is that the resulting inputs do not satisfy the orthogonality conditions and, as mentioned above, a rigorous analysis in the non-orthogonality setting is beyond the scope of what can be theoretically justified at this stage. We will make sure to include a clear discussion on this limitation of our work.

---

1. Apologies for these parameters having not been clearly stated. The parameters are now explicitly stated in the revised manuscript, $d=768$ with input length as sentences from our own abstract that are stacked together before being processed with a pre-trained tokeniser that comes with Hugging Face's checkpoints of the corresponding model. All experiments are repeated $5$ times, with the average result shown as a solid line and individual runs displayed as faded lines.

2. We sincerely thank the reviewer for these interesting suggestions which we believe enhanced our manuscript. We have undertaken these experiments and are happy to share the results with you. (a) We show in Fig. V how rank collapse in width is affected when modifying the information propagation in RoBERTa according to our proposed fix. RoBERTa is such that $d=768$. Whilst this small improvement might seem marginal in width (and the discrepancy with our theoretical guarantees can be explained by the extensive complexity of the architecture beyond our theoretical framework), its impact on the rank collapse in depth is remarkable, see Fig. VI. (b) The equivalent of Fig. 5 in our draft with the y-axis being the stable rank for consecutive layers following equation (1) is exactly Fig. 4b in our submission. The label '$\mathbf{A}(\mathbf{X})$' indicates that information flows as in equation (1), whereas '$\mathbf{A}$' (e.g. in Fig. 5) is used when the attention matrix is taken as an i.i.d. Markov matrix, independent of the input. We realise this distinction was only clearly stated in Appendix A.4 and have moved this up to the main body in the revised version. Nevertheless, rank collapse occurs across consecutive layers, as shown in Fig. 4b, and the only modules that address it—among LayerNorm and skip connections—do so by removing the spectral gap. In case the reviewer was interested in seeing the same plot with i.i.d. Markov matrices instead, we upload an additional figure for which similar comments can be made; see Fig. VII. (c) Fig. 4b in our current submission seems to be what the reviewer is asking for. Regarding the gradients in deeper key-query attention layers, we provide an additional plot in Fig. VIII that reaffirms our analysis and give us some empirical insights on the relationship between rank collapse in width, in depth, and gradients explosion.

---

• Thank you for highlighting these manuscripts, which we had not originally cited. We will include citations to them in our revised manuscript.

We have taken into consideration your overall impression that the limitations of our work were not clearly stated enough and have modified our draft accordingly. Should you have any further concerns that might affect your appreciation of our work, we would be happy to address them.

Thanks for your kind support of our theoretical development and for highlighting the "SAMformer" paper "SAMformer" that we looked into and will cite. For extra figures (indicated by Roman numerals), see https://anonymous.4open.science/r/spectral_analysis_transformers-C633/figures_rebuttal.pdf.

• Addressing rank collapse is motivated in multiple ways: (i) Improved quantization, crucial for compressing LLMs. Bounding the largest singular value by solving rank collapse ensures entries stay within a defined range, aiding quantization; (ii) Expressive initial representations. While it's still unclear what defines a good initialization, one might argue for avoiding rank collapse so that token representations remain diverse (as opposed to collapsing into the same vector); (iii) Controlled gradient norms. [2] shows gradients vanish after rank collapse, highlighting the importance of maintaining rank; (iv) Better generalization accuracy. [4] empirically supports our approach. While they don't directly address width, they propose subtracting the leading first-order term $\frac{1}{T} \mathbf{1}_{T\times T}$ from the attention. Their motivation lies in addressing oversmoothing in graph networks, tied to extremal eigenvalues. Their results show this subtraction increasingly benefits performance as datasets grow more complex or deeper, significantly boosting accuracy. (v) Finally, as reviewer qzvJ pointed out, our work helps bridge a theoretical gap in the understanding of attention layers in the literature.

• Regarding orthogonal inputs, we refer you to our response to reviewer uxct.

• We agree that including the impact of additional network modules (as in [1, 2]) is important. We do so empirically in Fig. 4, which shows rank collapse persists even with LayerNorm and skip connections. Fig. 3 further compares our theory to real-world transformers, such as BERT, which include many modules beyond attention.

• Fig. 3 uses a pre-trained tokeniser, so inputs are not isometric (see response to uxct). As our work focuses on theoretical insights into attention at initialisation, we do not analyse training dynamics, which lies beyond our scope. We instead refer to [4], which offers an extensive empirical study on removing the spectral gap. Their work complements ours: they focus on experiments, while we examine the mathematical consequences of centering. Rather than replicate their study, we discuss it in the next paragraph. For signal propagation, we analyse the stable rank of token representations across depth.

• Thank you for highlighting this. We’ve expanded our discussion of how our results explain [3, 4]. In [3], the authors heuristically subtract two softmax matrices, implicitly reducing energy along the dominant eigenvector. This helps mitigate rank collapse, albeit less directly and without theoretical grounding. Nonetheless, it shares the insight from [1, 2]—that rank collapse is central. Our resolution, removing the spectral outlier, has a similar effect but is rigorously understood. In [4], the same proposal of subtracting the dominant $\frac{1}{T} \mathbf{1}_{T\times T}$ term is made to “centre” attention. They show this benefits training. Our contribution is the mathematical rigour: we analyse what remains after centring and, by understanding the spectral gap, indicate what to do for other activations (e.g. ReLU, sigmoid) as shown in Fig. 1. We believe our theory supports practical efforts like [3, 4] to address rank collapse.

---

Q1. Rank decreases with matrix multiplication, which leads to collapse with depth. For a training instability link, see [2] (Theorem 3.2).

Q2. Please refer to the 'Scaling' paragraph in our response to uxct.

Q3. Assuming uniform attention (as in [2]) removes the role of attention by treating all tokens equally. We aim to refine this by analysing the attention once the leading $\frac{1}{T}1_{T\times T}$ term is removed. Think of uniform attention as the constant in a Taylor expansion—our $A^{\perp}$ captures the next-order term.

Q4. Great question. [2] proposes scaling residuals with depth to counteract collapse. But as we study rank collapse in width (even in the first layer), their fix does not apply. As shown in Fig. 4(a), modules from [1] do not significantly prevent collapse in width.

Q5. Please see above.

Q6. We agree with the reviewer and refer to [1] (Fig. 2), which also shows minimal changes in depth. However, we emphasise that even small changes in width can drastically affect depth. For instance, removing the spectral outlier in width causes a small change (Fig. V) but drastically alters depth behaviour (Fig. VI).

---

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

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

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

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

We thank the reviewer for their positive view of our work. We are happy to see that they value our contribution. We would like to address the points they raised as weaknesses:

• The reviewer rightly points out that the direct connection between the pathological rank collapse behaviour at initialisation and the training dynamics is not yet as well established as the analogous edge-of-chaos line of research for fully connected and convolutional networks, e.g. Figs. 5 and 6 in [1]. In the case of transformers in particular, one could argue that as long as the input tokens do not entirely collapse into a single representation during the initial forward pass—i.e. if some information manages to propagate through the network—the early steps of the optimisation algorithm may help recover from the collapsed rank and steer training back on course. Nonetheless, empirical evidence suggests that escaping such a pathological landscape can be more difficult in practice; see Fig. 1 of [2].

• While our analysis is admittedly not comprehensive of the whole transformer block, it provides valuable insights and sheds light on a hitherto overlooked phenomenon (even at the first layer) and its origin, namely rank collapse in width and softmax activation within attention layers. We also believe there is value in our results showing how the same rank collapse is impacted by changing the softmax activation to other options being advanced, such as ReLU and sigmoid. These other choices exhibit a smaller spectral gap than the softmax activation, which might explain why they are being advocated. Moreover, our resolution to subtracting the spectral gap can also be applied to ReLU and other attention activations. We believe this insight will be useful to practitioners studying alternatives. We wish to emphasise that, since rank collapse in width has a totally separate cause from its better-studied counterpart, rank collapse in depth, the ideas proposed to fix the latter do not necessarily affect the former and further distinct research is required. Thank you for the reference.

• Finally, in response to some reviewers' concerns about the assumptions used in our proof of concept, we have conducted extensive ablation studies that the reviewer may find useful in confirming their positive assessment of our work. They can be found here: https://anonymous.4open.science/r/spectral_analysis_transformers-C633/figures_rebuttal.pdf

---

[1] Schoenholz S. et al. (2016). Deep Information Propagation.

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

Thank you for your comments. Those with an editorial nature are applied in the revised draft, so we will address their main concern regarding assumptions. For additional figures (indicated by Roman numerals), please see https://anonymous.4open.science/r/spectral_analysis_transformers-C633/figures_rebuttal.pdf.

Orthonormal input tokens. It is true that input tokens are usually not orthogonal. Nonetheless, we wish to point out that this assumption is only made for the sake of our mathematical analysis which should be seen more as a proof of concept. The emergence of spectral gap and rank collapse in width still occurs in practice without those assumptions. Let us clarify the experimental setup of Fig. 3. Sentences are tokenised with a pre-trained tokeniser, making inputs to our randomly initialised BERT clearly non-isometric. We include $\mathbf{X}_0 \mathbf{X}_0^\top$ to illustrate its deviation from identity. When tokens are instead assigned random embeddings, the covariance matrix is closer to orthogonality. Yet, the spectral gap (Figs. I(d) and II(d)) and rank collapse (Fig. 3 of the draft) persist in both cases. We have clarified this setup in our draft to emphasise this aspect. Another illustration of why orthogonality is not needed per se comes directly from Fig. 6, where the spectral gap persists in deeper layers where representations are definitely correlated. Thanks to your review, we have also added to the appendix additional plots (Fig. IX) of the spectra of a random attention layer with (synthetic) non-isometric input tokens to mirror Fig. 1.

Characterising the spectrum of a generic row-normalised matrix remains an open problem in Random Matrix Theory. Our work advances the understanding of randomly initialised attention layers by assuming orthogonal input but we should acknowledge this is a limitation of our work and will emphasise this aspect in our revised version thanks to your comment. Moreover, we would like to share with the reviewer that we consciously chose to put an assumption on the input data (for our proofs to hold) to fully capture the complexity of the attention mechanism, rather than simplifying the attention mechanism as it is commonly done in the current mathematical treatment of transformers (e.g., in [1], where rank collapse is addressed under the assumption of uniform attention---essentially, no attention).

Ablation on $\gamma$. Note that without $\gamma \leq 1$, the inputs cannot be isometric so this extra assumption was solely made for this reason. To more directly speak to your question we included new plots that explore information propagation in TinyBert, for which $d>T$ hence necessarily the tokens must be non-orthonormal. The spectral gap persists (Fig. III(a)) and rank collapse in width follows (Fig. III(b)).

Large width. Although the theorems are formally stated as $T,d \to \infty$, we have also determined precise rates of convergence, which allows one to derive bounds in the finite case. Moreover, the convergence is sufficiently fast that, for typical practical values of $T$ and $d$ (on the order of $10^2$ to $10^3$), the quantities are already well approximated by their limiting values, see Fig. 1. In fact, it is one of the key messages we want to convey that, given the increasing scale of these architectures, asymptotic analyses (such as Random Matrix Theory) can be appropriate tools for their study.

Scaling. As noted in Remark 3.3, we scale values as $\mathcal{N}(0,1)$ instead of $\mathcal{N}(0,1/d)$ to compensate for the softmax attention that shrinks in all but one direction. This ensures layerwise tokens remain of order one after removing the spectral gap. Alternatively, scaling by $d^{1/2}$ post-softmax would yield the same effect. Since we omit LayerNorm in our analysis, rescaling is crucial to maintain token magnitude. Note that stable rank is independent of scale so we lose no generality here. In contrast, scaling does impact gradients, so vanishing gradients in Noci et al. under traditional scaling become exploding ones in our setting.

GPT architecture. We investigated signal propagation in a GPT2 transformer, as suggested by the reviewer (see Figs. IV(a), IV(b)). The attention matrix is now triangular due to the causal mask and its eigenvalues (real) can be read off the diagonal of this matrix. Although interesting, we feel that because it is a decoder-like architecture, an additional layer of complexity comes with considering causal masks that our theory does not cover, so this setting falls outside of our scope. Yet, we share with the reviewer the results of such an experiment and hope to maybe gain insight from them for potential follow-up work.

We added a proof sketch in our revised work, thanks for the great suggestion. We hope you will consider updating your score based on our rebuttal.

---

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