Two failure modes of deep transformers and how to avoid them: a unified theory of signal propagation at initialisation
Authors: Alessio Giorlandino
Sebastian Goldt
Venue: ICLR 2026 Poster (2026)
Keywords: Transformer
Signal Propagation
Theory of Neural Networks
Physics for Machine Learning
Primary area: learning theory
TL;DR: We provide a complete analysis of signal propagation across a Transformer block with self-attention, residual connections, LayerNorm, and a ReLU MLP, making it possible to predict the trainability of deep Transformers.

Abstract: Finding the right initialisation for neural networks is crucial to ensure smooth
training and good performance. In transformers, the wrong initialisation can
lead to one of two failure modes of self-attention layers: rank collapse, where
all tokens collapse into similar representations, and entropy collapse, where
highly concentrated attention scores lead to training instability. While previous work has studied different scaling regimes for transformers, an asymptotically exact, down-to-the constant prescription for how to initialise transformers has so
far been lacking.  Here, we provide an analytical theory of signal propagation
through deep transformers with self-attention, layer
normalisation, skip connections and MLP. Our theory yields a simple algorithm to compute trainability diagrams
that identify the correct choice of initialisation hyper-parameters for a given
architecture. We overcome the key challenge, an exact treatment of the self-attention layer, by establishing a formal parallel with the Random Energy Model from statistical
physics. 
We also analyse gradients in the backward path and determine the regime where gradients vanish at initialisation. We demonstrate the versatility of our framework through three case studies. Our theoretical framework gives a unified perspective on the
two failure modes of self-attention and gives quantitative predictions on the
scale of both weights and residual connections that guarantee smooth training.

Published as a conference paper at ICLR 2026
Two FAILURE MODES OF DEEP TRANSFORMERS AND HOW TO AVOID THEM: A UNIFIED THEORY OF SIGNAL PROPAGATION AT INITIALISATION
Alessio Giorlandino & Sehastian Goldt International School of Advanced Studies (SISSA) Trieste, Italy
(agiorlan, sgoldt 18sissa.
ABSTRACT
Finding the right initialisation for neural networks is crucial to ensure smooth training and good performance. In transformers, the wrong initialisation can lead to one of two failure modes of self-attention layers: rank collapse, where all tokens collapse into similar representations, and entropy collapse, where highly concen- trated attention scores lead to training instability. While previous work has studied different scaling regimes for transformers, an asymptotically exact, down-to- the constant prescription for how initialise transformers has far been lacking. Here, we provide an analytical theory of signal propagation through deep transform- with self-attention, layer normalisation, skip connections and MLP. Our theory yields a simple algorithm to compute trainability diagrams that identify the correct choice of initialisation hyper-paranee parameessssssrss for given architecture. We overcome the key challenge, an exact treatment of the self-attention layer, by establishing a formal parallel with the Random Energy Model from statistical physics. We also analyse gradients in the backward path and determine the regime where gradients vanish at initialisation. We demonstrate the versatility of our framework through three case studies. Our theoretical framework gives unified perspective on the two failure modes self-attention and gives quantitative predictions on the scale of both weights and residual connections that guarantee smooth training.
INTRODUCTION
Finding the right initialisation for neural network is a key challenge facing every practitioner before training. Choosing the right scale for the initial weights is critical for maintaining information flow during the forward and backward passes, which in turn ensures smooth training and good final performance. In fully connected networks, a series of works established a clear picture of signal propagation and trainability initialisation by computing how the similarity of two inputs evolves as they propagate through the same network at initialisation (Poole et al. 2016 SSchoenholz et al. 2016)) Xiao al. 2018 Hanin and Rolnick 2018 Hanin 22018) Doshi et al. 2023). ultimately establishing
the optimal initialisation at the "edge of chaos" " (Poole et al. 2016).
In Transformers (Vaswani al. 2017), where fully-connected layers alternate with self-attention layers (Bahdanau et 2015), the key quantity for measuring information flow through the network is the similarity between tokens sequence faces additional challenges due to two key failure modes of self self-attention layers. The first is rank collapse, where self-attention maps all input tokens to identical representations, producing an output matrix of rank one. This phenomenon manifests as the attention pattern shown in fig. a). Dong et al. (202I) demonstrated that networks composed solely of self-attention layers will inevitably collapse any input sequence into uniform token representations at rate that is double exponential in the number of layers. This rank collapse fundamentally destroys input sequence information and impedes effective training by inducing vanishing gradients (Noci et al.) 2022). Entropy collapse represents the second failure mode, where queries attend to only a small, frozen subset of tokens irrespective of the input, leading to low Shannon entropy of the attention distribution (hence the name) and, more
propagates through the network. Signal propagation

---

Published as a conference paper at ICLR 2026
[a]
(b)
strongty)
The
The
The
The
quick
quick
quick
quick brown
brown
brown brown
fox
fox
fox
fox
jump
jump
jump
Collapse
#As ##s #ts
##fs
(selt att
Rank Collapse
0.5
the
the the
fence fance fence
0.5
A:
(Q/K weights scaling)
1.0
7.0
6.0
0.8
Cosine Similarity
5.0
Test Loss
4.0-
1.5
0.4
3.0
1.0
Bank Collapse
1.0
2.0
1.0
0
10
20
30
40
50
60
1x10*
2x10*
3x10*
Layer
Steps
Figure 1: Two failure modes of Transformers at initialisation, and how to avoid them. (a) Rank collapse occurs when the self-attention layer attends uniformly all tokens, mapping input tokens into the same output token. (b) Entropy collapse is regime of highly saturated attention matrices which attend to random, semantically meaningless patterns, leading to training instability (Zhai et al., 2023). (c) Trainability diagram for 60-layer BERT Transformer, obtained from analytical theory of signal propagation, see algorithm [ Depending on the strength of the self-attention residual connections OSA (eq. (5) and the scale of initial key and query weights / (eq. (9). delineate the three regimes of rank collapse, entropy collapse, and the regime where the Transformer is trainable (blue). (d) Average cosine similarity between token embeddings of sequence taken from the StoySto dataset as it propagates through the layers of a vanilla BERT model for different self-attention residual strengths; empirical measurements (dots) closely follow theoretical predictions (solid lines). Sufficiently large residual connections asa are key to preventing the similarity between tokens from becoming unity, which would indicate rank collapse. (e) Test loss of 60-layer BERT model on TinyStories for two initialisations from each regime. Models suffering from rank or entropy collapse at initialisation fail to train, as predicted by theory. Full experimental details in section A.
importantly, unstable training (Zhai al. 2023 An attention matrix exhibiting this pathology is shown in (bb).
1.1 RELATED WORK
Previous work has highlighted the importance of skip connections in mitigating rank collapse (Dong et al. 2021; Noci et al. 2022 Wang et al., 2024), and Zhai et al. (2023) suggested modification of the self-attention layer that helps avoiding entropy collapse. However, quantitative, unified description of how these two distinct phenomena emerge has been lacking. Our result] [ fills this gap, providing clear guidance on how to jointly avoid both issues by appropriately choosing the strength of residual connections and the scale of initial weights.
The main challenge in studying information propagation in Transformers arises from the self-attention layers. Previous works either assume uniform attention (Noci et al., 2022) approximate the average behaviour of a self-attention layer by taking expectations separately over the numerator and denominator of the softmax (Cowsik al. 2024) strong simplification that fails the behaviour of large initial query and key weights, which is responsible for entropy collapse. In contrast, we analyse the standard softmax Transformer in the complementary limit of infinitely long sequences
2

---

Published as a conference paper at ICLR 2026
by leveraging tools and concepts from statistical physics. Our theoretical framework provides unified explanation for - and a practical solution to the emergence of both failure modes observed in practice: rank collapse and entropy collapse.
Further related works. Geshkovski et al. 22023a b) developed an elegant mathematical framework that models Transformers interacting particle systems, revealing emergent clustering phenomena in token representations. BBncoraglio et al. (2025) determines Bayes-optimal generalisation error for attention-indexed models at finite sequence length T,, while we consider the complementary limit of + oc. [ooddelon et al. (2024) analyses mean-field training dynamics in various limits: infinite query/key dimensions, infinite heads, and infinite depth, leveraging dynamical mean field theory to derive feature and gradient kernels. Noci et al. (2023) provide a description of the limiting distribution of the covariance matrix of token embeddings the proportional limit of infinite width and depth for
modified Transformer architecture while we explore the complementary limit of infinite sequence
length. We revisit their result in section 5
1.2 CONTRIBUTIONS
Our main contribution is an analytical theory of signal propagation in deep Transformers with self-attention layers, skip connections, layer normalisation, and MLPs. Our theoretical framework yields simple algorithm to compute the evolution of the typical overlap between token embeddings as sequences propagate through deep, off-the-shelf Transformers at initialisation. This enables us to construct trainability diagrams such as the one shown in fig. [c) for a 60-layer BERT-style Transformer. By varying the scale of initial key and query weights B, eq. (9). and the strength of the skip connections of the self-attention layer aSA, eq. (5). we identify three regimes for signal propagation: entropy collapse dominance (red), rank collapse dominance (yellow), and trainable regime characterised by small initial weights and strong skip connections (blue). The critical threshold for query/key variance, Bc, emerges global property, while the critical threshold a, for residual strength depends on network depth: for any given model depth, our theory predicts the minimum residual strength required guarantee signal propagation and ensure trainability.
We validate our theory in two ways. In fig. dd), we show that our theory (solid lines) accurately predicts the average cosine similarity between tokens (averaged over all token pairs in a sequence) when propagating sequences from the TinyStories dataset (Eldan and Li 2023) through a vanilla BERT model (Devlin et al. 2019) at initialisation (dots). In fig. [fc). we show the test loss of a 60-layer BERT model on TinyStories for two initialisations from each regime; models suffering rank or entropy collapse at initialisation indeed fail to train, as the trainability diagram predicts. We provide further applications of our theory in section 5
We also make the following additional contributions:
We identify sharp phase transition in the behaviour of self-attention in the limit of infinite sequence length, using tools from statistical physics (Derrida 1981). Below critical value 8, of the initial scale 3 of query/key weights, eq. (9). attention matrices are approx- imately uniform, while for 3 >> 3.,. they are highly saturated, causing rank and entropy collapse, respectively, see result
Our results are exact in the limit of very large sequences. We calculate finite-size corrections that are essential for obtaining precise, practical predictions, see section 3.2
We analyse the backward pass by deriving an exact expression for the norm of the gradient of query and key weights initialisation, and find that gradients vanish for 3 < Bo, see section 3.3
We integrate our results on self-attention and normalisation layers with existing results on MLPs to derive simple algorithms like algorithm T that yield trainability diagrams like the one shown in fig. Te. see section 4
We demonstrate the versatility of our framework by considering three case studies: signal propagation in standard BERT architecture; comparing different placements of LayerNorm; and comparing variations of the self-attention mechanism itself, see section
In summary, we provide a full theory of signal propagation in deep Transformers that unifies two phenomena previously discussed separately, namely entropy and rank collapse. Tuning initialisation
3

---

Published as a conference paper at ICLR 2026
hyperparameters in accordance with theoretical analysis enables accurate prediction of the trainability very deep models, offering practitioners a principled approach to designing and initialising Transformer architectures for reliable training at scale.
SETUP AND NOTATION
We consider a vanilla Transformer encoder (Vaswani et al, 2017) that processes se- quences (X1):=1...., of T tokens embedded in d-dimensional space. We analyse signal propaga- tion through a complete Transformer block comprising self-attention, residual connections, layer normalisation, and a reed-forward MLP.
Layer Norm. We consider layer normalisation (Ba et al. 2016) which centers each token embedding and rescales it by its standard deviation; for simplicity, omit the affine transformation:
X
LN(X.)
Var X
XXII
where |e is the Euclidean norm. We consider both pre-norm and post-norm variants.
Self-Attention. Given the query, key and value weight matrices We, WK, Wv E IRdxd, the attention score between tokens X, and X2 is
1
(W X
ag
These scores are normalised via a softmax operation to obtain the attention weights, which are then used compute the weighted sum of the value projections:
T
Au
S(X), = Σ
Av
ΣJ
P=1
MI.P. Each embedding X, is processed independently by shared two-layer feed-forward network with hidden dimension and non-linearity Ø:
MLP(X,) W2 / WWiX, + b1) + b2,
The weight matrices W,. W2 and the biases b,, by are initialised with i..d. entries of zero mean and variance ad/a and o7 respectively.
Residual Connections. We write the residual connections for the self-attention and MLP blocks as RESsA = S(X) + O,AX, RES (X) MLP(X) + C.MLPX,
(5)
where a 0 are the strengths of the skip connections.
Sequence Structure. Since token embeddings are sampled i.i.d. at initialisation, standard concen- tration results (Vershynin) 2018, sec. 3.1) in high dimension d imply that their norms concentrate, while pairwise overlaps scale (d2-1/)) Under self-attention, these tokens mix and their overlaps evolve as the sequence propagates in depth. Two key quantities to describe the evolution of token similarity the overlap and cosine similarity matrices
Gts
X, X,
Pts
V4uqss
which measure the degree of alignment between two tokens in a sequence. Our theory provides framework analytically track the evolution of these quantities through transformer blocks.
Notation. We denote the average over tokens or token pairs with brackets (+). whilst E[|] denotes the average with respect the initialisation of the parameters Wq. WK. Wv-
4

---

Published as a conference paper at ICLR 2026
3 A THEORY FOR SIGNAL PROPAGATION IN TRANSFORMERS
The main challenge in analysing information flow through a Transformer is handling the self-attention mechanism with its strong non-linearity and its normalisation step. In this section, we derive a theory for signal propagation by leveraging a formal similarity between self-attention and the Random energy model of Derrida (I98T), see section 3.1] Using tools and concepts from statistical physics, we give an asymptotically precise characterisation of signal propagation in self-attention layers in the forward pass in section 3.2 We analyse the backward pass in section 3.3 To complete the picture, we finally add the remaining elements of the Transformer architecture in section which yields algorithm T to compute the trainability diagram of Transformers.
3.1 SELF-ATTENTION AND THE RANDOM ENERGY MODEL
The starting point of our analysis is exact mapping of attention at initialisation model of disordered systems from statistical physics, called the Random Energy Model (Derrida 1981). This mapping was previously identified by Lucibello and Mézard (2024) the context of associative memories; our novelty lies in leveraging it to provide an exact description of signal propagation.
The Random Energy Model (REM). The REM is simple model of a disordered system with N spins that has (N)) possible configurations. Each configuration has an energy E(s) drawn at random from a Gaussian distribution with variance O(N) hence the name Random Energy Model. The probability of each state is given by the Boltzmann distribution with inverse temperature 3 > 0:
-8)(ss)
p(s)
Tr. (-DB(s)
The probability of a state p(s) and the attention Aue between a pair of tokens, eq. 3). thus share the same mathematical structure: normalised exponential with a random argument. While the energies in the REM are random by construction, the self-attention scores dur are random at initialisation due to their dependence on the random initial weights. We therefore interpret a row of the attention matrix as REM, with the full attention matrix representing T copies of the system.
Mapping of Self-Attentii to the REM. We initialise the query/key projection parameters i.i.d. with variance of o a./d so that the attention scores, defined in eq. (2), are normally distributed with mean and variance o2 While both REM energies and self-attention scores ag are Gaussian, the key difference between them is that attention scores are correlated. Averaging over the realisation of Wo. WK. we get the following correlation structure (see section B.T.:
E [orsuro] = o (X, Xx) Xx) = T.
(8)
We extend the REM framework to account for these correlations; consequently, the critical inverse temperature derived in result depends on these underlying geometric quantities.
The only free parameter is o² and thus the initialisation of query/key weights, which controls the scale of attention scores and, as we will see later, the sharpness of the self-attention mechanism. The remaining step is to match the scaling of attention scores with sequence length to the random energies.
Scaling of Query/Key Weights. In MLPs, the natural scaling for weight variance ((1/d) He al. 2015), ensuring that neuron pre-activations remain (()). For self-attention the REM analogy suggests that the natural scaling for query and key weight variance should make the attention score variance Olog T), where T is the sequence length, match the energy of a REM with N degrees of freedom, whose energy typically scales (LV) with fluctuations of O while the partition sum runs over Oee") configurations. Self-attention only sums over 7' tokens in the sequence, suggesting that attention scores should scale like O O'Vlog T). This motivates us control the variance of attention scores via a constant inverse "temperature" parameter 3 E R as:
of = o = =a
0a BVIoT.
This scaling is, to our knowledge, novel and complements the existing literature various reparametrisations of the infinite-width limit (Yang and Hu 2020) by accounting for the infinite sequence length limit.
5

---

Published as a conference paper at ICLR 2026
Rc
Inpul Cosine Similarity
0.8
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.2
2
5
Phase diagram for single layer of self-attention. We use result to plot the average cosine similarity between pairs of tokens after one layer of self-attention as function the query/key variance parameter 3 and the input average cosine similarity p. (Left): Theoretical phase diagram obtained from result] (with all input tokens map single output direction and the cosine similarity saturates at 1. For 3 > 8.. token diversity is preserved, but entropy collapse emerges. (Right): Simulations with embedding dimension d = 512 and sequence length T 1024 qualitatively reproduce the theoretical transition, with deviations attributed to finite-size effects, as discussed in section 3.2
Figure 2:
and p = p). For f < Se. we observe rank collapse phase, where
3.2 FORWARD SIGNAL PROPAGATION THROUGH A SELF LAYER
By tracking the average squared norm E(qn) := and the average overlap E(gss) := p. we describe the evolution of the average cosine similarity o and of the inverse participation ratio (IPR), which serves as an indicator of entropy collapse. The IPR of an attention row is defined for all queries E T and
T
(
Σ A.
(10)
For r = 2, the IPR estimates the effective number of tokens receiving attention: a small IPR, scaling (1),, corresponds to delocalised attention, while non-vanishing IPR indicates that only Or keys matter, signalling sparse self-attention and entropy collapse. Larger values sharpen the distinction between localised and delocalised attention vectors.
We state the result below and provide its derivation in section B.2
Result (Average Cosine Similarity Update under
Let WQ and WK be initialised with i...d. entries with variance o2 = o B / T /d, and Wy with variance oî o / For a sequence with average squared token norm q and average pairwise overlap define the critical initialisation scale 3el..)) p) = pp In the limil of infinite width d + oc and infinite sequence length T oo, we then have that:
1. The evolution of the average cosine similarity takes the form:
< Bclg.))
$ = (1 (1= 1-3 22(11p))
(11)
B 3clq.p).
2. The average inverse participation ratio (IPR) y2?)) satisfies Vt E [T]:
0,
8< Bclg. p)
lim Ey T
(2)
(12)
B.(g.))
B (c(
Result aallows draw phase diagram that describes the behaviour of relf-attention for different values and the input cosine similarity P. see fig. [ Before we discuss the phase diagram in detail, comment the definition of the critical B, and finite-size effects.

---

Published as a conference paper at ICLR 2026
(8)
(6)
1c-4
Attention Entropy
1.2
3-
Test Loss
1.0-
2
0.8
0.6
0.4
10
10 10 Steps
1.5
10*
10
10* Steps
1.5
x0
0.2
13
0.0
04
14
14
Figure 3: (a, b) A phase transition in the impact of query / key initialisation training dynamics. Average Shannon entropy of attention's row and the test loss of Transformer with single layer of self-attention trained on masked language modelling on TinyStories we vary the scale of the initialisation from small to large initial weights (blue to red). Small initial weights (blue) permit attention diversify over time, supporting effective learning, while large-variance initialisation (red) collapses the attention to only few tokens, visible in an entropy that quickly goes near zero. Here 3e(p 0) = v2. (c) Norm of the query gradient. Frobenius norm of the gradient of the loss with respect to query weights for various combinations of sequence length T = 2048, 4096, 8192 and embedding dimension d = 256, 512, 1024. As predicted by result 2 gradients collapse for different T and d, and vanishing gradients afflict the low-B regime.
Definition of B. and Finite-Size Effects. A key element of our analysis the convergence behaviour of the softmax normalisation 7(3) = Σ e0t, whicc depends on 33 through the scores atr (see eq. 2 We define 8e as the value of at which the law of large numbers for Z breaks down, i.e. the value for which Z ((3) no longer concentrates around its mean as T + co. However, as noted by Ben Arous et al. (2005), the central limit theorem for Z-which guarantees Gaussian fluctuations around EZ- El-already fails lower threshold B. Be/2. Since T is large but finite in practice, we therefore expect our asymptotic predictions to remain accurate for / < Be. where fluctuations are still approximately Gaussian. For 3 3,,. these finite-size fluctuations become non non-Gaussian, signalling gradual crossover toward the broken-LLN regime rather than sharp transition exactly at 3.. This mechanism accounts for the discrepancies between theory and simulation observed near the critical inverse temperature in fig, 2
Rank Collapse. The phase diagram shown in fig. exhibits a phase transition in the cosine similarity between tokens after self-altention. For small initialisation < Be. self-attention operates in 'spread attention' phase where the attention layer effectively outputs the average of all tokens, making the pair-wise cosine similarity between tokens saturate one, resulting in rank-one representation matrix. In the absence of skip connections, this behaviour leads to rank collapse, which makes the Transformer untrainable. This behaviour is also reminiscent of the "clustering" property analysed by Geshkovski et al. (003bb); Bruno et al. (2025); Chen et al. (2025).
Entropy Collapse.For 3 Ba, the self-attention layer preserves diversity, and hence information, among input tokens, suggesting this regime as a viable initialisation. However, the non-vanishing value of the IPR for Be, eq. (12), reveals that self-attention is in localised phase in this regime: it only attends to a few tokens which are determined by initialisation, rather than learnt, leading to the training instabilities observed by Zhai et al. (003)). In other words, for B > Be self-attention suffers from entropy collapse, which cannot be avoided using skip connections. This behaviour is a central novelty of our analysis: Noci et al. (2022) assumes the model operates entirely in the rank-collapse phase, thereby overlooking this distinct failure mode, while the "annealed" approximation of Cowsikk ctt al. 2024) does not capture the large-deviation behaviour underlying entropy collapse.
Take-Home Message. The main implication of result I is that the only viable initialisation for self-attention lies in the amall-vriiane regime 3 < 3cc supplemented by skip connections to maintain information flow. To illustrate this, we train one-layer Transformer task at varying 3. As shown in fig. 3 the two phases - spread vs. localised exhibit distinct training behaviours. In the low-variance regime, skip connections help mitigate rank collapse,
masked language modelling
7

---

Published as a conference paper at ICLR 2026
enabling successful training. In contrast, in the high-variance regime, the attention vectors' entropy collapses, and the model fails to train effectively, as escaping the pathological initialisation and learning meaningful patterns would require prohibitively long training. Overall, this result unifies two previously observed phenomena, namely rank and entropy collapse, within a single theoretical framework: the onset of either phenomenon is separated by sharp phase transition, governed by the variance of the query and key weight initialisation, as parametrised by B.
Sequence Geometry in Depth. We assumed i.i.d. token embeddings, which implies that before any processing the sequence forms a simple random geometry: by concentration of measure, the tokens lie near the vertices of high- digensimnniooal simplex, with qu ~ 1 and qts ~ 0. As the sequence is propagated through Transformer at initialisation, this simplex structure is preserved in the long-sequence limit (i.e. qu = q and Its = p), as shown in section B.2.3 although the values of q and p themselves evolve with depth. Thus, result] provides way to track how the geometry changes across layers: given the geometrie pair (g. p) at layer I, we plug it into the update map to obtain the new pair-and and hence the updated geometry at layer ++1.
Before introducing the remaining Transformer components- reesdual connections, MLPs, and layer normalisation-we we analyse the backward pass through the self-attention layer.
3.3 THE BACKWARD PASS
To complete our theory of signal propagation, derive the following result on the norm of the gradients of query and key weights at initialisation (see section f.4 the derivation):
Result 2. (Query/Key Gradient Analysis) In the limit T - xc, under the same hypothesis of resull] ! the expected squared Frobenius norm of the query gradient, and analogously for the key, is given by
T
C33o% 4(a (apeaaaappppe -D)( (30-v(0--((0-00000000 -2 + p((12)))))
E
8C DW
(13)
d &-oog
where C a constant independent of and d. The same result holds for WK- K-
Vanishing Gradients. Result 2 shows that gradients can vanish under two conditions. First, we see that if q = i.e. if attention is uniform and maps all input tokens into the same output token, gradients vanish. In this case, we recover the well-known result Noci et al. (2022) that showed that if the inputs to the self-attention layer are already collapsed, and hence all the tokens are the same, gradients vanish. Result 2 goes beyond their result to show that even if input tokens are diverse, i.e. p o q, gradients vanish if 3 < Be because in that regime, the inverse participation ratios (2) and (3) tend to zero as the sequence length T + For long but finite sequences, finite-size effects reduce the threshold down B < Be Be/2 (see section 3.2). We numerically verify result 2 in fig. BJ. where we show that (1) the curves of gradient norm versus initialisation strength collapse with the scaling suggested by result 2 and that (2) gradients do indeed tend to zero for 3 < 3c.
Our analysis of the backward pass raises a paradox: initialise with 8 9ec and gradients vanish; initialise with 3 > Se. and self-attention is stuck with entropy collapse. We discussed in the previous section that entropy collapse can only be avoided by initialising with small weights 3 < 3, and adding skip connections. So where do the gradients come from? Result relies on the simplex structure of the embeddings. While this structure holds exactly at initialisation, the embeddings receive non-zero gradient through the skip connections already in the first backward pass, independently of the value of B. These updates break the simplex structure on which our derivations depend, and thus also invalidate the vanishing-an -rrddient result, which applies only at initialisation.
FULL TRANSFORMER BLOCK ANALYSIS
We now describe how to track signal propagation through skip connections, LayerNorm, and the MLP, and then show how these components combine into simple iterative algorithms. Such algorithms allow practitioners to predict the evolution of the average cosine similarity across Transformer layers and to generate trainability diagrams that reveal viable pperrpaaaeettr regimes for given Transformer architecture. An example for BERT-style post-norm Transformers is given in algorithm|] [ additional variants, including pre-norm, reported in section B.5.2
8

---

Published as a conference paper at ICLR 2026
Self-attention layer. The order parameters characterizing sequence geometry evolve in the self- attention layer as derived in section B.2.3
Is = (p+ (g p) (8, Ps o.p.
(14)
where we recall that [ Wvl = of.ld. Taking the ratio of these expressions recovers the cosine- similarity update of result
Adding residual connections. Because Wy is independent of the residual stream, adding residual connection of strength ass contributes additively (see section B.5.I):
q - qs + qAAA.
p ps POSA-
(15)
An analogous relation holds for the MLP residual path.
Adding MLPs. How the signal propagates through the MLP layers follows from the results of Poole et al. (2016); Schoenholz et al. (2016). For a two-layer ReLU MLP, the evolution of squared norms
qd)) and pairwise inner products p:)) is:
/ D (VQ(-V()
?
og. 1= 1=2....,4
n o² / Da (VQ 21) Φ (Q(l-)) 20) + t
where ø is the activation function and T2) is pair of standard Gaussian variables with covariance ((t-1) p(l-1)/0 /q(1-1) The initial conditions after the first linear layer are q(d)
q goth + of and p(2) po + +oz. For two-layer ReLU MLP, the second-layer outputs simplify to:
(2)
2 (2) of
2
2 zhf (2) (() + off
(16)
where f(p) captures correlations after the ReLU nonlinearity Cho and Saul (2009):
Algorithm 1 Post-norm Block Update
f(p) = (VI- 0² p(T arccos(p)))
(
1: Inputs: 3, q. p. aSs, CMLP, d oī o?
2: Attention layer + residual
After including the MLP residual connection, the updated quantities become:
3: 8c wb
4: ( (() + max - Bc/3)
g (2g))+++ g²q(2) a
p p (-pp2)+pah (17)
5: + of (p+ (q-p) ( (8, +g-
6: t p. (o) + aza)
Taking their ratio gives the full block-level update of the average cosine similarity.
7: P Post-norm LN
8: - p/9: 2+1
LayerNorm. The remaining component is layer
9: resduull
normalisation, which simply enforces
10: 91 + 0 oj:
Pi t p
11:
pa + = f4)p +o
p+6 2
1.
12: g f + a...r
p - P2 + aMaP
This update should be applied at the position dic- tated by the specific Transformer variant under consideration.
13: > Post-norm LN
14: p <- p/q: +-
15: return (q.
We conclude the paper by outlining several applications of the resulting algorithm.
5 APPLICATIONS
To showcase the versatility of our approach, we consider three case studies: signal propagation in standard BERT architecture; comparing different placements of LayerNorm; and comparing variations of the self-attention mechanism itself.
Signal propagation in vanilla Transformer. We first used algorithm in -tRT-style Transformer (Devlin et al., (219).. Since BERT uses the post-norm convention for Norm, state the algorithm for the post-norm architecture; algorithm 2 for the pre-norm
To to analyse signal propagation

---

Published as a conference paper at ICLR 2026
version. The algorithm states the update for the average norm q and average overlap p. the average cosine similarity can be read off after each block simply as p p/q. Iterating the algorithm for different values of 3 and ass finally yields the trainability diagram shown in fig. [d).
Placement of LayerNorm. Xiong et al. (2020) showed that placing LayerNorm before the self- attention layer and before the MLP greatly stabilises the training of deep Transformers. Comparing signal propagation using our algorithm, we show in fig. 4 that rank collapse, corresponding to an average token similarity of (p) confirming pre-LN as the more stable choice.
does indeed occur much later for pre-LN than for post-LN,
Avoiding All Collapses: gain- gain-controlled attention. A recent line of work has sought to alleviate rank col- lapse by directly modifying the self-attention layer it- self. Noci et al. (2023): Naderi et al. (2024) proposed to enforce the attention layer to perturbation of the identity. Noci et al. (2023) derive an SDE descrip- tion the limiting distribution the overlap (or neu- ral covariance) matrix in the proportional limit where both width and depth go to infinity. To obtain this limit, they introduce width-dependent temperature parameter and remove layer normalisation. However, layer normalisation is crucial for preventing entropy collapse: the critical value (e(g. p) identified in re- sult depends inversely on the squared average norm Using our framework, we can show that simply removing the mean of the values along the sequence from the output of the standard self-attention layer, also explored in (Naderi et al. 2024) and reminiscent of gain control in neuroscience, can be combined to great effect with either post-LN or pre-LN. In fig. 4 show the theoretical prediction for the evolution of the average cosine similarity, illustrating how this modification alleviates rank collapse. Our prelim-
1.0
post-LN
0.8
pre-LN
GC post-LN
GC pre-LN
0.4
0.2
0
60
100
Layer
Figure 4: Theoretical prediction of the evo- lution with depth of the average cosine sim- ilarity for the standard Transformer and the Gain-controlled Transformer under both LN strategies. Rank collapse is avoided simply by removing the mean value in the self-attention layer. Here, we set aSA = OMLP = 1.
inary experiments with a thirty-layer BERT-style Transformer trained on TinyStories show that gain-controlled Transformers succeed in regimes where vanilla attention fails, see fig. 6 encouraging further experiments at seale which are however out of the scope for the present paper. As a final note, order propagate the signal infinitely deep. one should simply use gain-controlled attention and initialise the MLPs at the edge of chaos.
Autoregressive Models. Our theory relies on the observation that, when starting from embeddings with a simplex structure, non-causal transformers preserve this structure in depth in the limit of infinite sequence length. In this regime, the representation dynamics can be fully described by only two parameters: the average squared norm g and the average overlap p. A key ingredient in this result the infinite-sequence limit itself. For autoregressive models, however, this limit introduces additional subtleties. Early tokens do not satisfy the simplex geometry because they have access to only a small fraction of the sequence context, and this regime the geometry more complex and two parameters are no longer sufficient. As one moves to later tokens, the geometry is progressively restored. Consequently, in causal models rank collapse does not occur at the level of the full sequence; instead, for sufficiently long sequences, it can emerge in the representations of tokens appearing later in the sequence. Our theory approximates late-position tokens, tracking how their pairwise similarity changes with depth and predicting when their representations collapse (fig. T2).
CONCLUSIONS
We developed a theory for signal propagation in Transformers that unifies the understanding of the two main failure modes in Transformer training: rank collapse and entropy collapse. Building on these insights, we find new evidence for the viability of simple architectural modification of self-attention, the gain- controlndd self-attention, that avoids both failure modes, which would be interesting explore at scale in future work.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
We thank Antoine Maillard for helpful discussions. AG gratefully acknowledges funding from Next Generation EU, in the context of the National Recovery and Resilience Plan, Missione Componente 1, Investimento 4.1 "Estensione del numero di dottorati di ricerca dottorati innovativi per la Pubblica Amministrazione e patrimonio culturale" (CUP G93C23000620003) and also gratefully acknowledges financial support from Fondazione Zegna. SG gratefully acknowledges funding from the European Research Council (ERC) for the project "beyond2", ID 101 166056; from the European Union-Nee((GenneatiinnE,,, in the framework of the PRIN Project 2022E3WYTY CUP G531D23000780001, and from Next Generation EU, in the context of the National Recovery and Resilience Plan, Investment PEI Project FAIR "Future Artificial Intelligence Research" " (CUP G53C22000440006).
(code
REFERENCES
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in neural information processing systems, 29, 2016.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. arXiv preprint 161i 0/222,, 2016.
Lechao Xiao, Yasaman Bahri, Jaseha Dick-Diccst,,n, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and mean field theory of cmns: How train 10,000-layer vanilla convo- lutional neural networks. In International Conference on Machine Learning, pages 5393-5402. PMLR, 2018.
Boris Hanin and David Rolnick. How to start training: The effect of initialization and architecture. Advances in neural information processing systems, 31, 2018.
Boris Hanin. Which neural net architectures give rise to exploding and vanishing gradients? Advances in neural information processing systems, 31, 2018.
Darshil Doshi, Tianyu He, and Andrey Gromov. Critical initialization of wide and deep neural networks using partial jacobians: General theory and applications. Advances in Neural Information Processing Systems, 36:40054-40095. 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In International Conference of Learning Representations, 2015.
Yihe Dong, Jean-Boa Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages
PMLR, 2021.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation transformers: Theoretical perspectives and the role of rank collapse, 2022.
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Josh Susskind. Stabilizing transformer training by preventing attention entropy collapse, 2023. URL httpttts:///riiv..rrr//sss///ees
06296
Hongyu Wang, Shuming Ma, Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers to 1,000 layers. IEEE Transactions Pattern Analysis and Machine Intelligence, 2024.
Aditya Cowsik, Tamra Nebabu, Xiao-Liang Qi, and Surya Ganguli. Geometric dynamics of signal propagation predict trainability of transformers. arXiv preprint arXiv:2403.02579, 2024.
11

---

Published as a conference paper at ICLR 2026
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspec- tive on transformers. arXiv preprint arXiv:2312.1079, 2023a.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. Advances Neural Information Processing Systems, 36:57026-57037, 2023b.
Fabrizio Boncoraglio, Emanuele Troiani, Vittorio Erba, and Lenka Zdeborová. Bayes optimal learning of attention-indee models. arXiv preprint arXiv:2506.00 582, 2025.
Blake Bordelon, Hamza Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head transformer dynamics. Advances in Neural Information Processing Systems, 37:35824-35878, 2024.
Lorenzo Noci, Chuning Li, Mufan Li, Bobby He, Thomas Hofmann, Chris J Maddison, and Dan Roy. The shaped transformer: Attention models in the infinite depth-and-widhe limit. Advances in Neural Information Processing Systems, 36:54250-54281. 2023.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprint arXiv:2305.07759, 2023.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume (long and short papers), pages 4171-4186, 2019.
Bernard Derrida. Random Random-energy energy model: An exactly solvable model of disordered systems. httttps://apl.semmaatt
Physical Review B, 24:261332622 1981.
org/
Corpus ID: 122288449
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450. 2016.
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47, Cambridge university press, 2018.
Carlo Lucibello and Mare Mézard. Exponential capacity of dense associative memorics. Physical Review Letters, 132(7):077301, 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-1034, 2015.
Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint 1rXiv:2011.14522. 2020.
Gérard Ben Arous, Leonid Bogachev, and Stanislav A Molchanov. Limit theorems for sums of random exponentials. Probability theory and related fields, 132(4))57--612 2005.
Giuseppe Bruno, Federico Pasqualotto, and Andrea Agazzi. Emergence of setle clustering in mean-field transformer models, 2025.
Shi Chen, Zhengjiang Lin, Yury Polyanskiy, and Philippe Rigollet. Quantitative clustering in mean-field transformer models. arXiv preprint arXiv:2504.14697, 2025.
Youngmin Cho and Lawrence Saul. Kernel methods for deep learning, Advances in neural information processing systems, 22, 2009.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International conference on machine learning, pages 10524-10533. PMLR, 2020.
Alireza Naderi, Thiziri Nait Saada, and Jared Tanner. Mind the gap: a spectral analysis of rank collapse and signal propagation in transformers. arXiv preprint 11X.722410.07799 2024.
12

---

Published as a conference paper at ICLR 2026
Jesse Vig. A multiscale visualization of attention in the transformer model. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics: System Demonstrations, pages 37-42, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/
v1/P19-3007. URL httpps:/.. www. ac web. httbbborggaathlooy//eegee 9-007
M. Mézard and Montanari. Information, Physics, and Computation. Oxford Graduate Texts.
OUP Oxford, 2009. ISBN 9780198570837. URL http::////oooossoooeeeeeeeee.
id= id= (b=j)hc77O@UUSS
Charles Bordenave, Pictro Caputo, and Djalil Chafai. Circular law theorem for random markov matrices. Probability Theory and Related Fields, 152(3):751-779, 2012.
A
FIGURE DETAILS
All experiments have been conducted using a single NVIDIA A100-PCIE-40GH Training times vary from approximately one hour (one layer models) to one day (60 layer models).
A.1 FIGURE
(a, b) Attention visualizations obtained using BertViz (Vig 2019) on a single-layer, single-head Transformer at initialization. The attention maps illustrate the effect of varying 3 (directly related to variance of queries/keys): in yellow, a model initialized with 3 = 0.1 (low-variance regime, resulting in approximately uniform attention distributions), and in red, model initialized with B 1.8 (high-variance regime, leading to sharp attention).
(c) Trainability Diagram for a 60-Layer BERT Transformer. The diagram based on the critical value 3. = 22, which marks the threshold at which entropy collapse occurs in the first self-attention layer under the assumption of i.i.d. and normalised token embeddings. In post- norm architectures (such as BERT), we always have g 1 at the self-attention layer, so 3 22/(1 p). Since p increases with depth, the smallest critical value the one at the first layer (where = 0), and we therefore this value single critical threshold for the entire network (i.e. all layers initialised with the same variance). In other Transformer variants, such as pre-norm, the critical threshold also depends on the average squared norm, so one must determine the appropriate critical value for each layer individually.
The residual-stee residual-strength threshold is defined the smallest value of the residual scaling factor oSA below which rank collapse (loss of representation diversity) occurs across all 60 layers.
(d) Evolution of cosine similarity between token embeddings across layers in a 60-layer Transformer initialized in the low-variance regime (using standard HuggingFace initialization for queries and keys. which corresponds to small J8. Lines denote theoretical predictions. while dots indicate empirical averages over 10 random initializations and 10 input sequences. (Error bars are the standard deviation.) High similarity indicates representational collapse.
(e) Same as (d), but for 12-layer Transformer initialized in the high-variance regime (B = 1.8). Again, lines show theoretical predictions and dots indicate empirical means. Remark: we intentionally use shallower model in this regime to ensure that any training failure observed in panel (f) is attributable to entropy collapse rather than rank collapse, as signal propagation across 12 layers is guaranteed.
(f) Pre-training results for BERT-style encoder models using masked language modeling task (masking probability 0.15) on the TinyStories dataset. We compare models with 60 layers (blue and yellow: small 3 ~ 0.02 < 3.) and 12 layers (red: 8 = 1.8 > 3.). In particular learning curves in the three phases are: trainable (3 ~ 0.02, USA 1.5, 2), rank collapse (8 ~ 0.02, aSa = 1.0), entropy collapse (8 1.8, OSA = 1.0, 1.5, 2.0). Remark: 3 0.02 corresponds to the standard initialization from Hugginface given the set of hyper--araminn paramters we are using. All models use ReLU activations, 6 attention heads, embedding dimension 600, and absolute positional embeddings.
13

---

Published as a conference paper at ICLR 2026
Initialization: di 0.2, of 0.0004, of o standard HuggingFace initialization for queries/keys; no biases or affine transformations LayerNorm.
Residual scaling: aSA shown in the figure, OMLP 1.0.
Optimizer.AdamW with learning_rat=-led rate=le-4, num wum_train_epochs-l.
size=_,
max_grad_norm-ll norm=1.0, Ir_ sesssccheeeerrree setillerr_typeelineeett weight deggh__deca=0000, and warmup warmup_ ratio=0.005.
Note: 3. (p = 0) 2 is the critical threshold for the first layer of self-attention, which at initialization takes inputs that are approximately orthogonal.
A.2 FIGURE
(Left)
Theoretical prediction.
(Right) Cosine similarity averaged over all pairs of tokens when passing a sequence of length
1024 with embedding dimension d 600 through single self-attention head.
A.3 FIGURE 3
Pre-training 1-layer Transformer with 1 head using masked language modeling (masking prob- ability = 0.15) on TinyStories. Embedding dim d 768, standard residual strengths, ReLU activation, absolute position embedding. Custom init of queries/keys with As as in figure; standard init for other weights, no biases for queries/keys, no affine in LayerNorm. Opti- mizer: AdamW, Bearning_rate=5e-4 num_true num_train_eppee wrain_epochs=1, batch size__s max lar___rrrdnnmmmlll scheduler type= type="linear", "linear", weight decay=0.01,, warmup_ ratii00.50
B THEORY APPENDIX
B. ATTENTION SCORES ARE CORRELATED GAUSSIAN VARIABLES.
Consider the attention scores defined in eq. (2). By the Central Limit Theorem, they converge in distribution to Gaussian random variables with zero mean and variance o where is the embedding dimension (or head dimension in the multi-head case), and o determined by the initialisation of Wq and WK as described in section 2 Although the scores ag are individually Gaussian in the infinite width limit (d oc), they not independent; in fact, they correlated. To quantify these correlations, we compute:
Cov(a Aro) d
1
Σ
XI B[(Wq),: ([((K))WK (Wk))))))WW)nn mmn).
J.m.M.M.....l
Since the query and key weights are independently initialised with variances of = 0% = co/d, this simplifies
[arsaro] = o600 (Xe Xx) Xx) 0.4t44aa.
(18)
B.2 DERIVATION OF RESULT
B.2.1
COMPUTATION OF (8))
Consider the partecipation ratio of the t-th row of a self-attention matrix average over the initialisation We: WK.
20c2
Ey.2) E C eo
We also define
(r(3,h) h) = E log (1,
(19)
where
T
h3) = Σ
hasa
h MER R
==1
14

---

Published as a conference paper at ICLR 2026
and we recall that 3 enters in the definition of the covariances between the a's, they are all proportional to o = 3 log
We want to exploit Stein's s lemma, which yields:
Σ aese hoess
0 (, = E
Σ
IF| (aatsass
E
(20)
choos
8,8'=1
with the correlation between attention scores given by eq. (88). We assume qu q for all due to concentration of measure. For the pairwise overlaps, we proceed as follows: at the first layer, all tokens are approximately orthogonal, making it safe to assume 9ts ~ 0 Vt + s [T]. One then derives the update for this layer and observes that, in the limit of infinite sequence length (as we will briefly show), the update becomes independent of the indices to leading order. By repeating this argument across layers, it is therefore justified to treat the gts as having a common mean p together with sub-leading Gaussian fluctuations, which we neglect since our analysis focuses on the average overlap. Applying this argument, we get:
hha
2hot En charm) )
chargghaer (2x chaen)) )
On (,,,h ho Σ E 1 'En
Σ
phosu
= =hoâq (q -p) p) (1-BY,
Finally, this leads to:
lim EY(22) =1 024 (q p) A-1 11T T T-+00
lim lim 0, (,,,
(21)
To proceed, are left with computing the expectation (8, E [log E, cnres
The remaining computation amounts to evaluating a variant of the Random Energy Model (REM), but with correlated energy levels due the structure of the ats variables. This problem can be
tackled using the Replica method or, alternatively, via a micro-canonical argument, as discussed in Mézard and Montanari (2009). Here, we proceed with the Replica method.
We compute the replicated partition function as:
T
E, 20.2((30 (3, h) E,
hors
Σ
En
exp
h
Orea
a=
12022 n
Σ
exp
E atss asmol
a,b=1
h²o2
Σ
4 Σ
I1(aa 88) + qp ΣΣ I1s so)I(s' sb)
exp
$1
a.b
(22)
We now introduce the empirical overlap matrix:
Qab = I[ssa = 8t)
and perform a change of variables from replica indices to overlap structures Q, giving:
1²022
E, (8,
ΣΣN
8(Qab, 1(5. = =b)) exp
a
(g
p) Σ
Qab + 0(13)
2
15c) a.b
a.b
1:'22 Σ exp)
a
p) Σ Qai 0(13)
(23)
2
15

---

Published as a conference paper at ICLR 2026
p=0
0.8
p=0.25
p=0.5
p=0.75
0.6
Yg(8)
0.4
0.2
0.0
B
Figure 5: Theory and experiments (T" 10%) comparison of the computation of (3)) finite size
effects are visible around the phase transition.
Now we take the I-RSB ansatz for Q: the n replicas are diveded into groups of æ elements which in the energy configurations. Moreover, we need to consider exponentially long sequences, i.c. we take T = e and control N. This implies:
S(Q) = ee
Σ Qut = ne
ab
So exploiting the replica trick and recalling the definition of o we get:
8²h²
(, max s<1 9 p) x
(24)
which leads the existence of a critical temperature Be q,2) 1o2 // a(a-p) phase transition takes place. In particular we have:
where a condensation
1+
hhh 2 ( p) 8< 3ch.
s ( 1)/N
Bh 2g (1 24
8> 3h,
Due to the fact that we took the expectation over all the a's there is no dependence on v, rather only on the similarity matrix, and can drop the subscript from 4. Finally and we get:
put all together,
0 3< Bc(ap)
lim (Y((8)= = 1- 1 lim lim O, (, = (2) T-200
03 (g p) R-1
(egg.) 8> (c(..))
where Be(g.p) p) = 4e(h = 1, = V Vi
5 the result of the computation and 2 We plot in fig.
some numerical simulations.
B.2.2 COMPUTATION OF (8)
We need to compute:
E, tan
(g( = Σ. Cn Ee clusd
Consider the partition function with auxiliary fields h (Meer (t )6,0=11
Zo,u (3, = Σ
(eu
26"
16

---

Published as a conference paper at ICLR 2026
Let's observe that:
Cavo + ana )emooelorrr aanrrann
Σ Dn.. B log Zo, (3, E
Σ Canstianntaa
E(avo + auo (ara aua)
Σ
Chh tand)
E
(25)
Σ Chhoolavotass
(ourtaa))
Assuming the first term self-averaging, we approximate:
Σ (aea + auo ) (ae
Σ choolauro
Σ On. E log Zu, (3. = E
E
Σ choo (Ave taya)
cheerloos
Under this approximation:
lim Σ Oner E log
Σ cauttn)))
(3)
(26)
lim E,g
hla
The computation the free entropy in the denominator is straightforward. It closely resembles the derivation for the free entropy appearing in the calculation of (B).
One gets:
1+
1
8² 8
h2h
h
22
E Elog
bauvstuus)
(27)
N
2,3h
1
which gives
2,82 1 -
B <
lim 0
E log
Σ Mavoraar)
(28)
h-51
2
23
p
B
The calculation of the free entropy the numerator bit more involved, but we can observe the following:
(aarttatt E, (ays + )e
lim h-
Σ On.o E log Σ
En, edaatauel
5,5
La + aue' De twu
EE,s''i +
En,er gavataue
We define:
Σ Cus
(avz)
Σ enu
Σ (ave + Que' )aur)ttutta tays
8,²==1
((ave + aus a)s
sts'
Σ
=1
17

---

Published as a conference paper at ICLR 2026
Thus, obtain the approximation:
lim h-11
Σ Tnao
E log
Σ 5,8
chastaanttnnl (ava)s + (aus) {(ays + aus' Msps' 2
There is also an intuitive way to see this. Consider the two limiting cases:
As B -
The attention weights become uniform, i.e., Aes -> 7- Then,
BORV (S(X), S(X) ExQ Σ Av Sea + T E Se0 P O
1
q
meaning both qIt,, pall + q'
As B + oo: The attention becomes fully peaked:
Aus (v)) with 8 (v) arg max aps-
In this limit, the dot product becomes:
T S. 1, ydt (g-(w))o- P
since (²) * o² (w) with probability
L/T.
Hence, as we vary 3 from 0 to oc, the quantity Sg(g) interpolates between and 1, while SA(p) remains nearly constant.
B.2.3 UPDATE OF AVERAGE AVERAGE COSINE SIMILARITY.
We begin our analysis by averaging over the value projection matrix Wy. Since Wy is independent from We, QK at initialisation, the scalar product between self-attention outputs then becomes:
E S(X), S(X) = E Σ A.sA B Wy] X 8,0
d
Choosing, without loss generality, the variance of the value projection weights of l/d, obtain [, Wv] Lg, the expression simplifies to:
E
Σ
AceAvo
(29)
With the limit T + oo in mind, we neglect sub-leading fluctuations and focus on the leading terms:
Gsa = for 8 + a, Isa q.
This approximation holds due concentration of measure immediately after the embedding layer. Moreover, the result we derive in the unfiite-sequencc-lenggie limit for the first layer is independent of the specific token indices, meaning that the simplex geometry is preserved. Consequently, the same reasoning applies recursively at every layer, implying that the dynamics are fully captured by the two scalar quantities and p: they the only relevant geometric degrees of freedom precisely because the simplex structure persists throughout the network.
E
9
AleAvs
+ pp > AleAvo
(30)
sifo
Since each row of the attention matrix sums to one, we can simplify this expression further:
q Σ
AlsAps +p
AlgAeo AloArs
(q Σ ArsAps +p.
18

---

Published as a conference paper at ICLR 2026
Now, upon averaging and taking the limit T oc, two distinct quantities emerge depending whether t = t' or t t':
y(2 = lim E T-000
Σ
A
Yx()) = lim
Σ AcsAUs for ot,
(31)
T-+00
which quantify the self- and cross-overlap attention distributions. The expression for the first one derived in section B.2. and in section B.2.2 We show that the second term is sub-leading.
Substituting the result, get the update for the average squared norm:
p B < P)
f p (q-
q
(g p) (1 B.g..))
8> Bc(q.p)
(32)
B
where Be(g.p)
2
V ple-pi
On the other hand, the scalar product p is not updated, Y (3) is sub-leading. Taking the ratio between the updates yields, at leading order in d, the following update for the average cosine similarity:
s
P
1+ 1+ (1-1)))(((BB)
which gives rise the phase diagram in fig.
B.3 FINITE-SIZE EFFECTS
Here we give a non-rigonn non-rigorous argument on the finite size effects that afflict our asymptotic theory. In the low-3 regime, the attention is spread approximately uniformly over number T S(B.p) of keys, given by an entropic quantity S(3, p) = (B, p) $0 p) (where the free entropy was defined in eq. (19)). A derivation of the entropy for the REM, very much related our problem, is explained in depth by Mézard and Montanari (209)). For 3 < (p). this turns out to be:
S(B, (1
g²
(33)
(p(p)
Since the IPR is the inverse number of the expected number of state that matter:
(3, = ~ -(8.p)
(34)
In the limit N O(log T), this non-rigorous argument suggests that the corrections to eq. 12 scale
a² Beto)
As long as 3 3. = B./2, these fluctuations can be neglected since they remain Gaussian. For B./2 < 8., however, the fluctuations around the asymptotic solution become non-Gaussian, as the central limit theorem breaks down at Be/2.
are
B.4 DERIVATION OF RESULT 2
Consider square matrices A, a. Q T where A = softmax(a) is the standard attention matrix
computed from logits and Q ((ts) is the overlap matrix. Let JT denote the T identity matrix. Also, consider matrices X E R Txd Wv, Wv,Wo, Wx E Rddxd
Define the attention operation
S(X) AXW..
(35)
We are interested in computing the squared Frobenius norm of the gradient of S(X) with respect to the query matrix Wo:
OS(X) tr
[SS(X)
[SS())
)
(36)
aW2
aWe
@Wee
FF
19

---

Published as a conference paper at ICLR 2026
The first part of the derivation parallels the proof of related result in Noci et al. (2022) However, unlike their approach, we do not assume uniform attention. Instead, we retain the attention explicitly, which allows the previously derived participation ratio naturally emerge.
Chain rule decomposition:
as(x) (Ir 8
wTXT)
0A da
1 Vd
X @ xwx).
aWq
Consequently, the squared Frobenius norm equals the trace of
0A
(Ir & WxxxT)
da Vd
x Vd x ② w
(oa)
(AA) T
(Ir 8 XWv).
0a
Simplification of the middle terms:
( XWxxx (
Io (XWxWxXT)
Assuming Wx W concentrates as
WKW = da
we obtain
d xx @ (X(d))))x)) of (XXx xxx).
Taking the trace and using its cyclic property, we define
0A
0A
G
(XXT XXXT)
da
da
and write
(S(X)
= o tr (G XWv) (L
aW
Assuming Wv WT concentrates as
WvWW se do?
we simplify further as
aS(X)
do tr (G (G(r
aW4
Recall the definition of the overlap matrix
xxT - E / d
we get the compact expression
0S(X) 2
d o tr (G(Iz @ Q))
aW F
20

---

Published as a conference paper at ICLR 2026
where
00A da
0A da
G =
(Q
2 Q)
This expression reveals how the gradient norm depends on the structure of Q. the Jacobian of the attention, and the variance parameters associated with the key and value projections.
Let's compute the trace term:
0A ( Q)
0A da
(1r (Tr @Q)
24 is the Jacobian matrix of size 7 x T2. Let's write the Jacobian in components. Using the where Du fact that A softmax(a), so the Jacobiam components are:
aAy dandrAAs AunA
(37)
Daw
Now, the trace can written as
DA
(Q
Q)
0A T
(Ir s Q)
kudgu-
da
da
(t))
LJ,,,=1
Expanding indices leads to
T
tr
Σ
(400049 (mm)) (emn (m
Simplifying the deltas:
T Σ
tr
(kd) (Gs))
t.
Now we can substitute the Jacobian components given by eq. Assume Einstein's 's notation.
37
Gi,i Aij Ainiinnn Aniiiiinnn Ainqnin
To leading order in we can substitute with its expectation value,
at + (g p)dze.
Expanding each contribution and taking the limit:
(1) AtjAmnj
p² (2p(g p) (q- p)9) ((2)
(2) AisissAinin
p² + 2p(a - (2) + (g - y2
(3) An
p² + 2p(g 22) + bg 2)2
(4) A.g
p² + 2pla p)r + plg p) (2(2))2
21

---

Published as a conference paper at ICLR 2026
putting all together:
tr = 9l9 -r=lo-ppl--p((() () + +
Finally:
[S(X) 2 DW d' o E
Aig p)2 + Σ A2y(1 - p)² o(a - -p)² Σ
AlyA2.
(38)
ij
fjm
Recall that we took x = 3 (T) /d and of = 0²/ so:
x B
= 3 /oo((T)oo ala p) (v 22(31) + +
(39)
d²
F
Now consider instead the gradient of the loss:
3 E
1
DC BW
I
B(X)) aaa p) () - 2 (
(40)
d2
where B(X) bounded quantity of X (see the proof of theorem 3.2 Noci al. 2022
We assume that this quantity scales like Or(T-1) and check it numerically fig. c). So putting all together:
T
E
8C 8Wa
x 3o q1a p)((0- - + ((V(2))3)
-
(41)
d log T
F
Since the (r) + 0 in the small /3 regime (Mézard and Montanari 2009 let's check our predictions for the case where = 1, = O. In this case 3g = v2, but in light of the discussion on finite size effects in section 3.2 we actually expect our prediction be sharp up Be/2 se 0.7 and then
a crossover between 3./2 2 and Be to the other solution. This behaviour is correctly predicted, see fig. c)..
B.5
SIGNAL PROPAGATION IN THE FULL TRANSFORMER BLOCK
B.5.1 ACTION OF RESIDUAL CONNECTIONS
Recall the action of the skip connection in self-attention:
RESsa(X) = S(X) + aSA = AXWv + a.A..
Consider the quantity
Eq.ky |RESSA (X); RESSA (X).).
The expectation over the value matrix vanishes in the mixed terms, leading to
Eq.K.v (S(X), + az X₂ X..
Overall, considering
E(g(s)
0. + uZAA
=
E(qu)
dz (2)) aZa9
where we used [ Wy] = oāLa a and the updates for and described section B.2.3
B.5.2 FULL BLOCK COSINE SIMILARITY UPDATE ALGORITHMS
22

---

Published as a conference paper at ICLR 2026
Algorithm 2 Pre-norm Block Update
1: Inputs: 3, q. OSA, OMLP, o of, of
2: > Pre-norm LN before attention
3: PLN f p/q; QLN
4: > Attention layer (normed input) residual
5: 8. T(1 PLN)
6: y (3) t max(O, Bc/3)
7: t of PIN + (gEN PIN) (3)( asSA
8: p + J.PLN + asAP
9: > Pre-norm LN before MLP
10: PLN p/q/ QLN 1
11: MLP layer (normed input) residual
q1 +- ou aß:
pi < PLN + ob c
13: 02 + " + of;
Pa t 2 f(pi/u)) + og
14: - 92 + amary:
p - P2 + aaar
15: return (g,P)
Algorithm 3 Gain-controlled Transformer Block Update with Post-norm
1: Inputs: 3. q. P. OSA. OMLP. 2 of o?
2: > Attention layer residual (centered-value updae))
3: A. V o(a-p)
4: (z) max(0, 1 Be/3)
g < ot (q (B) + aáA
6: p asAp
> Post-norm LN
p- + ppp/: q g+1 1
9: > MI.+ residual
10: Q1 -o-q + di:
P1 +o?
11: @ "" of:
P2 f(r//1)) + of
12: Q2 + o
pa ahar
13: Post-norm LN
14: + p/q: +
15: return (g.p)
23

---

Published as a conference paper at ICLR 2026
C SUPPLEMENTARY FIGURES
C.I TRAINING GAIN-CONTROLLED TRANSFORMERS ON TINYSTORIES.
We train gain-controon controlled transformers with post-LN with skip connections strength asa = 4MLP = 1 of one and twenty layers. The latter case would fail for standard transformer in the same setting due to rank collapse. The training dynamics reported in fig. 6
30 layers
8
Vanilla, post-LN
Guin-contralled, Post-LN
6
Eval loss
4
0.0
0.5
1.0
1.5
Training step
1e4
Figure 6: Training 30 layers of vanilla and Gain-controlled Transformer on TinyStories.
Details of training: 30-layer, single-head BERT-style model with embedding size 480 and ReLU activation, using masked language modeling with 15% masking probability, a learning rate of 5e-4, batch size 64, warmup ratio 0.05, weight decay 0.01, for 0.5 epochs.
C.2 VISUALIZING THE THE QUERY/KEY VARIANCE EFFECT ON THE SPECTRUM OF THE SELF-ATTENTION MATRIX
As shown by Bordenave et al. (202)) the spectral bulk of a random stochastic matrix-sun as an attention matrix initialization radius (T2-1/2). This result applies under standard initialization schemes, where the variance of the attention scores remains fixed and independent of sequence length.
In contrast, under our proposed rescaling of the query and key weight matrices, the variance of the attention scores scales o2 log which depends on the dimension of the attention matrix. This places the system in different spectral regime and, crucially, avoids the spectral bulk to have (TT 1/2) radius, key driver of rank collapse in conventional transformer models, as visualized in fig.
8 = 0.016 (BERT)
β = 0.75
β = 1.25
β =
2
Imaginary Part
0.001
-0.001
-1
-1 -1
-1 -1
Real Part
Real Part
Real Part
Real Part
Figure 7: Spectrum a 512 512 self-attention matrix for various values the query/key variance parameter
24

---

Published as a conference paper at ICLR 2026
C.3
ENTROPY COLLAPSE CAN MITIGATED BY LOW LEARNING RATE
Figure is obtained with the set-up fig. but with smaller learning rate
Attention Entropy
4
Test Loss
3
2
103
104
1.5
104
5 x 10a
104
1.5
10
Steps
Steps
0.2
0.4
0.6
0.8
1.0
1.2
1.4
1.6
1.8
2.0
Figure 8: Entropy collapse can be partially mitigated by smaller learning rates.
C.4 PHASE TRANSITION IN INFINITELY DEEP TRANSFORMERS
natural question is whether signal propagation remain stable at infinite depth. While this possible with ReLU activations. it becomes a genuine phenomenon when using tanh. This behaviour was first observed by Poole et (2166) in MIPs, and more recently extended to Transformers by Cowsik et al. (2024). In fig. 9 we confirm the phase transition in forward signal propagation, but emphasise that achieving this behaviour requires placing the MLP in the chaotic phase, which is associated with exploding gradients and unstable training.
Output Cosine Similarity
1.00
1.00
Cosine Similarity
0.75
0.75
0.50
0.50
Collapse
0.25
0 - 1.0
0.25
of, 1.0 of
of =.5
0.00
0.00
0.25
0.50
0.75
1.00
50
150
200
Input Cosine Similarity
Layer
Figure 9: Phase Transition to Infinitely Deep Signal Propagation. (Left) Cosine-similanne update map of full transformer block with tanh activations in the MLPs. By tuning the MLP variance to enter the chaotic regime, the collapsing effect of self-attention can be counterbalanced resulting in non-trivial fixed point in the similarity dynamics. (Right) Iterating the update map reveals the evolution of cosine similarity with depth-predictions align closely with experimental observations.
(Left) Theoretical predictions for full Transformer block initialized in the low query/key variance regime, with MLP weights initialized as d.e 1.0, 2.5, oī = 0.1. Residual connection scaling
factors OSA 6.0 and OMLP 1.0.
25

---

Published as a conference paper at ICLR 2026
(Right) Evolution of the overlap by iteratively applying the map from the left panel over 200 layers, starting from an initial overlap of O(d-1/2). Solid lines denote theoretical predictions, while dots represent experimental results averaged over 25 random initializations, each evaluated on 10 input sequences (error bars are the standard deviation).
C.5 TRAINING DYNAMICS FOR THE EXPERIMENTS THE MAIN TEXT
10.0
Train Loss
6.0
4.0
1x10*
2x10*
3x10
Steps
Figure 10: Training loss corresponding to the experiment shown in fig. De.
8
Train Loss
6
4
2
5
103
10a Steps
1.5 x 104
Figure 11: Training loss corresponding to the experiment shown in fig. 3p.
C.6 AUTOREGRESSIVE MODELS
All experiments performed 50-layer autoregressive transformer with post-normalization. We used hidden size 720, intermediate (MLP) size 720, and ReLU activations. The weight and bias variances were o of 0.2 and of x 10 respectively. Both the self-attention and MLP blocks employed scaling coefficients OSA = 1.0 and QMLP 1.0. Results were averaged over 10 independent model initializations and 10 sequences of length approximately 200 tokens.
26

---

Published as a conference paper at ICLR 2026
First 50 Tokens
Full Saquence
Last 50 Takens
Costre Simieriy
0.8
p
Cosine Cosine Simierfa
0.8
5=2
Cosine Simiariy
0.8
0.6
0.4
0.4
044
0.0
0.0
10
20
30
40
10
20
40
50
20
30
40
Layer
Layer
Figure 12: Propagation in Autoregressive Models. Cosine similarity between token representations function of layer depth for three values of 3. The solid lines show theoretical predictions, while the markers with error bars show empirical measurements obtained by averaging over all token pairs, across 10 sequences of length 200 and 10 independent model initialisations. Left: For the first 50 tokens, the simplex geometry is not yet established, and the empirical evolution departs from the two-paraan two-parameter theory. Middle: Averaging over the full sequence partially restores the structure but still mixes early and late tokens. Right: For the last 50 tokens, the geometry is well approximated cosine similarity, and the
by our theory: error bars are small, indicating strong concentration empirical curves follow the theoretical predictions. This illustrates that in causal models the simplex structure-and hence rank-collapse dynamics-emerg. progressively and are captured accurately only for sufficiently late tokens.
27


Reviews:
[{"review_id": "snvzKZROZS", "reviewer": "Reviewer_zigt", "summary": "This paper presents a unified analytical theory to explain and predict two failure modes of deep Transformer networks due to initialization: rank collapse and entropy collapse. The authors' central innovation is to establish a formal parallel between the self-attention mechanism and the Random Energy Model (REM) from statistical physics. This mapping allows them to analyze the behavior of self-attention as a phase transition. The authors present a parameter, $\\beta$ (related to the variance of query/key weights), that functions like \"inverse temperature\" in the physics model. The theory finds that operating below a critical threshold ($\\beta < \\beta_c$) leads to rank collapse, while operating above it ($\\beta > \\beta_c$) leads to entropy collapse. The paper's practical solution is to operate in the rank-collapse regime (to avoid the more fatal entropy collapse) and use strong residual connections ($\\alpha_{SA}$) to preserve signal diversity. The framework yields a simple algorithm to compute \"trainability diagrams,\" which provide a principled and quantitative guide for selecting initialization hyperparameters (weight scale and residual strength) to ensure stable training. The authors' theoretical predictions are shown to closely match empirical measurements in BERT-style models.", "questions": "The authors map self-attention to the Random Energy Model, but REM assumes independent energies while attention scores are correlated. How do these correlations affect the validity of the phase transition predictions, especially for moderate sequence lengths where finite-size effects matter?", "limitations": "", "rating": 4, "confidence": 5, "soundness": 4, "presentation": 2, "contribution": 2, "strengths": "The paper's strengths lie in its rigorous theoretical approach to a critical practical problem.\n\nStrong Theoretical Analysis: The paper provides a comprehensive and asymptotically exact analytical framework. It goes beyond simple heuristics to derive precise update equations for token similarity (Result 1) and an exact expression for the gradient norm at initialization (Result 2). This allows it to analyze both forward signal propagation and backward gradient flow, providing a complete picture of the network's dynamics at initialization.\nNovel Parallel with the Random Energy Model (REM): The analogy to the REM is the paper's key insight. By formally mapping the self-attention $\\text{softmax}$ (Eq. 3) to the Boltzmann distribution (Eq. 6), the authors import a powerful analytical toolbox from statistical physics. This allows them to model the complex, random behavior of the attention layer and identify the \"inverse temperature\" parameter $\\beta$ (Eq. 8) as the crucial knob controlling the system's behavior.\nEffective Visualization of the Phase Transition: The REM analogy naturally leads to the discovery of a sharp phase transition at a critical threshold $\\beta_c$. This core finding is visualized clearly and effectively. Figure 2 presents the theoretical phase diagram for a single layer, while Figure 1c translates this into a practical \"trainability diagram\". This diagram provides a clear map for clearly delineating the trainable zone from the Rank Collapse and Entropy Collapse zones.", "weaknesses": "My main concern is the lack of convincing experiments supporting the framework. While it is totally fine that the paper's focus is theoretical, it's a shame that there is no better evidence to support what is ultimately a very important practical problem - specifically when strong evidence would not be so difficult to provide. The core results are derived in the \"limit of infinite sequence length.\" The authors acknowledge this creates \"finite-size effects\" and a \"discrepancy between theory and simulations\" (Fig. 2) in real-world. Also, the gradient analysis (Result 2) assumes token embeddings are i.i.d. (independently and identically distributed). The authors note this is only true at the moment of initialization ($t=0$) and is immediately broken by the very skip connections their solution relies on. Both these theoretical assumptions alone justify better experimental evidence for the hypothesis. \n\nFurthermore, better support of how should the community benefit from more informed initialization - which is the whole point of the paper given the abstract - would be welcome. \n\nFinally, I find the paper needs some restructuring and rewriting. Here are some pointers: Section 4 for instance is standing alone, and a bit odd with the rest of the paper. A standard \"Related Works\" section would be good, as opposed to simply a \"Further Related Works\" paragraph after stating the main paper contributions. A discussion section as a 3.2.3, followed by 3.3 The Backward Pass also reads rather strangely. You mention training a 20 layer BERT (line 472) in the Main manuscript, but in the supplementary material, it becomes a 30 layer BERT model (line 1214) - which is true? In addition, the paper is filled with typos...", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ru2nVvBaVc", "reviewer": "Reviewer_X3PM", "summary": "The paper conducts a unifying analysis of two failure modes in signal propagation in Transformer models, namely the rank and entropy collapse, under the assumption of infinite width and sequence length. The authors use a known mapping of the self attention matrix to a Random Energy Model (ERM) [1], which allows them to conduct analysis that, contrary to the previous work, neither assumes uniform attention scores, nor taking expectations in the numerator and denominator of the self-attention separately.\nFollowing that they propose a new sequence-length dependent initialization scheme for the query and key weights, and derive the formula for the evolution of the average cosine token similarity as a function of the sequence-length independent initialization variance and the input embedding correlation. They also extend our understanding of the vanishing gradient problem beyond the uniform attention assumption, and provide an analysis of a full Transformer block, which lets them construct cosine-similarity evolution diagrams also for full Transformers. They verify their theoretical results with experiments on a BERT-style Transformer models.\n\n[1] Carlo Lucibello and Marc Mézard. Exponential capacity of dense associative memories. Physical Review Letters, 132(7):077301, 2024.", "questions": "1. Can you explain (briefly in the main text of the paper and maybe more in detail in the appendix) how one arrives at equations 7-15 from algorithm 1? As I understand it, it comes from previous work, but it would be useful to walk the reader through how one arrives at these formulas.\n2. How does $\\beta_c$ depend on the model depth, assuming models with and without layer normalization? How does $\\beta_c$ depend on the layer number assuming a model with a fixed depth (with and without layer normalization)?\n3. Can you produce a trainability diagram in the deep transformer case like the one in Figure 1d, just for varying values of $\\beta$?\n4. Is the entropy collapse regime also associated with a vanishing gradient problem? From fig. 1e it seems that models in this regime quickly stop learning. If the whole norm of the gradient in eq. 13 does not disappear, maybe the gradients wrt to just some subset of query and key parameter do? Do the attention scores stay roughly the same (with the same few large scores) during training or do they change between many low-entropy settings?\n5. How does masking influence the results? Would a causal mask like in GPT-style attention change Result 1 and eq. 14? How would the phase diagram from fig. 2 look under attention with a causal mask?\n6. In the paper $q_{tt}$ is referred to as the token norm (for example in line 295). Are you sure it is not a squared token norm?\n\nSuggestions on writing:\n* The sentences around eq 14. could better suggest that the equation describes propagation through a self-attention layer with skip-connection. From the section title I was expecting a single equation describing the full Transformer block, and the sentence leading to the equation just mentions skip-connection.\n* Lack of consistency with capitalising Transformer/transformer (lines 14 and 42 for example)\n* Broken sentence in line 113 (the strength the scale)\n* Line 131: „highly highly”\n* While I appreciate the explanation in appendix B.1, the writing there could be improved (take a look at the second sentence in line 652).\n* $q_{ts}$ is used in line 233 before its defined way later in the paper", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Unifying view of two failure modes of attention is novel and valuable for the signal propagation community. I think that working under less strict assumptions than previous work (uniform attention, separate treatment of nominator and denominator in attention) is valuable.\n2. Using ERM in the context of signal propagation analysis is novel and original.\n3. The part of the paper that focuses on a single self-attention layer is good quality and well presented. It’s commendable that the authors extend their analysis to a full Transformer block, and provide experiments for deep Transformer models.", "weaknesses": "I’m open to increasing my score if the authors address the following weaknesses:\n\n1. The analysis of the average cosine similarity while passing through the full Transformer block (Section 4) is rushed. I believe some other parts could be shortened (like the setup in Section 2) in order to make space for a more thorough analysis of the full Transformer block. The result from eq. 14 is not interpreted and authors do not clearly walk the reader through how they arrive at Algorithm 1 from it. I believe this type of analysis is valuable but should be presented better in the main text of the paper.\n2. The analysis of the propagation in terms of $\\beta$ and $\\beta_c$ seems to be done only for a single layer Transformer. I would be curious to see some discussion and some experiments similar to what’s presented in Figure 3d on how it behaves with increased depth. Especially, as the paper already discusses deeper models, but focuses only on the residual scaling $\\alpha_{SA}$.\n3. The experiments concern only BERT-style architecture. I believe the authors should provide some empirical results and comment on how their theoretical result and proposed algorithms for obtaining trainability diagrams behave in architectures with causal masking, like GPT-style autoregressive Transformer. These architecture is arguably more relevant for the DL community these days.\n4. There is no discussion about the gradients in the entropy-collapse regime. Section 3.3 just discusses the gradient formula for $\\beta < \\beta_c$. I would be interested in seeing if there is a way to explain with gradients how it happens that for large $\\beta > \\beta_c$ we get stuck in the entropy collapse regime.\n5. Regarding figures: 1) Figure 3c does not have a legend. 2) I would appreciate seeing also the training loss curves (in the appendix) for the experiments for which the evaluation/validation loss is presented.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UT23z1rIjA", "reviewer": "Reviewer_1xRi", "summary": "The authors propose to use the REM to provide an asymptotic characterization of signal propagation in Transformer (self-attention layer, skip connections, Layer Norm, MLP) and explain how different initialization scalings cause rank and entropy collapse.", "questions": "see weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper provides a precise asymptotic analysis on the evolution of the average cosine similarity and the average IRP with respect to the initialization scaling.", "weaknesses": "- There are many unclear arguments.\n    - Why is the variance of the attention scores $\\sigma_a^2$? (Doesn't it depend on the value of the variance of $X_t$?)\n    - Sec B.1 L666 says $\\sigma_Q^2=\\sigma_K^2=\\sigma_a^{\\color{red}2}/d$ which is different from eq (8). Also the conclusion in (15), $Cov(a_{ts}a_{\\tau\\sigma})=\\sigma_a^2q_{ts}q_{s\\sigma}$ is $d^2$ times larger than eq (7).\n    - Why does the attention score variance should $O(\\log T)$?\n\n- The asymptotic prediction (Result 1) does not perfectly match the empirical results. The authors said that this is because $T$ is finite. How about the nonasymptotic results with concentration bounds? What happens if we increase $T>1024$?\n\n- Why $\\tilde \\beta_c=\\beta_c/2$ does not act as a critical value in Fig 3?\n\n- Violation of the iid assumption means that (13) may not hold, but it does not answer why gradient does not vanish with the skip connection. Why do we have nonvanishing gradients when $\\beta<\\beta_c$ with skip connection? Also it seems like the iid assumption is violated during training, even without skip connection, so we may not conclude that $\\beta<\\beta_c$ implies vanishing gradients.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CUPW2bTYod", "reviewer": "Reviewer_5iSN", "summary": "The paper analyses the signal propagation, i.e. the evolution of similarities between token embeddings in the self-attention layer and other layers of the transformer architecture. The paper focuses on the infinite context length limit (while also discussing the finite number approximations), or regimes where the context length is very large. They compute the correct asymptotic context length-dependent initialization scaling of the attention parameters that makes the limit non-trivial and also compute the concrete multiplication-constant threshold under which the attention matrix converges to uniform attention, causing rank collapse and above which the attention matrix converges to entropy-collapsed one where tokens only focus on finite number of other tokens despite asymptotic context length. Both these regimes constitute trainability issues. The authors further analyze the size of the gradient and show that it is vanishing in the small initialization regime. The authors then analyze the strength of the residual connection, which can alleviate the rank collapse in the small initialization regime. The authors obtain concrete algorithm to determine the trainability for given hyperparameters, which is then used to construct trainability diagrams. Finally, authors experimentally confirm their theory and evaluate some methods to boost trainability.", "questions": "-\tQ1: To what extent are your proofs inspired by the either of the cited works [1,2,3] or some other works? \n-\tQ2: You haven’t disclosed LLM usage. I assume this means you haven’t used LLMs beyond some text adjustments and grammar checks. Is that right? \n-\tQ3: In Figure 1e, would the red and orange figures eventually decrease all the way down to blue curves or do you think they are stuck in local minima? \n-\tQ4: The normal distribution assumed in line 229 of the attention scores is also enforced in practice? If not, can it interfere with your theory because you are potentially creating heavier-tailed distributions? \n-\tQ5: In Figure 3(a), why does the entropy for the red curves actually go down? I understand they start lower than the blue curves, which is predicted by your theory, but why go down? As you show in the Appendix, smaller learning rate can alleviate this. Is it only an optimization issue then? Could LR warm-up help? \n-\tQ6: In result 2, you normalize with $T/log(T)$. This basically means that per-weight gradient converges to zero at an inverse rate. Does this mean we get vanishing gradients even in the optimal case? \n-\tQ7: What can be said in result 1 when $\\beta=\\beta_c$?\n-\tQ8: Can you elaborate on point (1) in line 396? How can we see this statement? \n-\tQ9: Can you elaborate on your discussion in lines 401-405? How can we see that this actually explains the phenomenon? \n-\tQ10: How do your results adjust in the case of multi-head attention? \n-\tQ11: More importantly, what can we learn about causal attention? It seems to me that causal attention cannot be solved by initialization and your results sort-of show it. Is that right? What would you recommend practitioners instead? \n-\tQ12: In line 463 you say that layer normalization is crucial to avoid entropy collapse. Where is this shown in your work? \n\n[1] Carlo Lucibello and Marc Mézard. Exponential capacity of dense associative memories. Physical\nReview Letters, 132(7):077301, 2024.\n\n[2] Alireza Naderi, Thiziri Nait Saada, and Jared Tanner. Mind the gap: a spectral analysis of rank\ncollapse and signal propagation in transformers. arXiv preprint arXiv:2410.07799, 2024.\n\n\n[3] Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien\nLucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse,\n2022. URL https://arxiv.org/abs/2206.03126.\n\n**Summary:**\nI consider the paper to have strong and novel contribution. While there are some aspects that can be improved and some questions that deserve answering, I think the paper should at least be accepted. Based on the discussion period and other reviewers’ points, I will then try to refine my judgement so as whether to recommend also spotlighting.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "-\tS1: The theoretical contribution of the paper is quite strong and, to the best of my knowledge, novel. Regarding novelty: while there are some pieces of the puzzle scattered around the related works, none provides such quantitative characterization, especially in the infinite context length limit or large contexts. None of the works unifies rank and entropy collapse regimes with precise threshold that divides them. The paper also convincingly discusses the non-asymptotic deviations, provides precise gradient norm analysis (novel) and connects it with the rest of the transformer (partially novel). Regarding the strength, I think the quantitative nature of the results, reasonable scaling regime and the potential for practical use make the paper useful for both theoreticians and practitioners. \n-\tS2: Besides being mostly novel, the authors seem to also accurately contextualize their novelty respective to the related work. I manually checked the related work (and also used AI to double-check) and it seem the authors correctly cited the literature in places where their results built on or aligned with some of the previous results. \n-\tS3: The paper is mostly quite clear and fairly easy to read.", "weaknesses": "-\tW1: Despite general clarity, at times the paper would benefit from more detailed discussions. For instance, it should be discussed in the main body, why do the results require IID embeddings. If I guessed it correctly from the Appendix it seems that only to guarantee some geometric properties in the first layer. But this should be discussed because if this is not the case, then it is not clear how the IID property applies in the intermediate layers. Another example is that terms used in Equation 7 are only defined in Equation 9. Some statements are a bit vague and not justified-enough. In particular the discussion in lines 401-405 seems to reflect authors’ intuition, but I don’t see why this should be the correct explanation. Furthermore, Figure 1(e) would benefit a legend. I raise a couple more unclear points in the questions section. \n-\tW2: One peculiar decision is to scale up the residual connections instead of the output of the modules (for instance value matrix or MLP). As far as I am concerned, this doesn’t provide any benefit compared to the other option and, to the contrary, in the case of pre-norm LN causes exponential growth of embeddings, which is a big issue because the contributions of attention layers and MLPs will diminish at initialization as we go deeper through the network (plus, numerical instability might be an issue at large depths). \n-\tW3: The paper could go a bit more into depth about how to use Algorithm 1 to effectively determine the correct per-layer $\\beta$ values in practice. \n\nTypos: \n\n-\tl125 “applicaitons”\n\n-\tl141 fig 1c\n\n-\tl228 “initialise”\n\n-\tl319 the $\\tau$ should be indexed", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)