Force without transmission: a depth-induced rank collapse that no loss on the representation reopens
Authors: Martin Hofmann, Patrick Mäder
Organizations: Data-intensive Systems and Visualization Group (dAI.SY), Technische Universität Ilmenau, Max-Planck-Ring 14, 98693 Ilmenau, Thuringia, Germany · German Centre for Integrative Biodiversity Research (iDiv) Halle–Jena–Leipzig, Deutscher Platz 5e, 04103 Leipzig, Saxony, Germany · Faculty of Biological Sciences, Friedrich Schiller University, Fürstengraben 1, 07745 Jena, Thuringia, Germany
Training can drive a transformer into a rank collapse: all token representations point in one direction, and learning stops. In a related collapse of attention, a loss term with a bounded corrective force repairs the network during the run. We ask whether such a term repairs rank collapse. We collapse small transformers by weakening their skip connection and treat copies of the collapsed network. No added loss term repaired the collapse, although the stronger kind pushed with about a tenth of the task gradient. The reason was the path, not the strength. The task gradient no longer reached the query and key weights, which decide where attention looks, and the added term's gradient faded before the blocks where the collapse forms. Restoring the skip connection, which changes no weight, reopened this path at once. The rank then recovered, but only far above the scale of collapse. After a burst of high learning rate the path stayed open and the rank recovered untreated. Registered predictions from the path ranked recovery times but did not transfer to this cause. In every case the loss stayed above that of a healthy network after the rank recovered. Whether a collapsed network can be repaired depends on whether the gradient still reaches the weights that must change, not on how strongly a loss term pushes.
Figures & tables
Treatment
Added loss
Force class
Added loss when switched on
penalty 1× / 4×
β(1−Q2/(d−1))
vanishing
Λ / 4Λ
hinge 1× / 4×
βmax(0,s0−Q)
bounded
Λ / 4Λ
cosine 1×
βU
neither
Λ
untreated
–
control
–
reversed ramp
–
driver reversed
–
Table 1: Treatments started from each saved collapsed network. s0=srankhealthy−1 .
Figure 1: Stable rank of the last block after the treatment is switched on, all treatments and collapse events of the main study (50-step median, as in the recovery rule; dotted line: recovery threshold). Only the reversed ramp recovers.
Figure 2: Gradient path in the saved collapsed networks of the main study and the replication (lines: median over networks; bands: range). Top: per-block gradient norm at the collapsed state for the task loss on the attention input weights (solid) and on the query and key part alone (dotted), for the per-block hinge at matched added loss (dashed), and for the task loss in the healthy network of the same seed (black). Bottom: task gradient on the query and key weights after the skip scale is set to 0.3 or 1, at the first step before any update (solid) and after 50 updates of the branch (dashed). Values below 10−13 are numerically zero in single precision and are drawn at the lower edge.
Figure 3: Stable rank of the last block against the skip scale α while it is ramped down (drivers) and, from the saved collapsed state, while it is ramped back up at the same rate (main study). Dotted line: recovery threshold.
Figure 4: Registered test under a learning-rate excursion with the skip intact. Left: query–key gradient at the saved state (median over blocks) for the depth-24 networks collapsed by a weakened skip ( section 6 ), the held-out excursion networks at multipliers 16 and 32, and healthy depth-24 networks; filled: recovered by rank, open: not, triangles: development networks ( appendix E ); dashed: threshold frozen in the first registered test; dotted: development threshold. Right: stable rank of the last block after the saved state, excursion networks against the untreated depth-24 networks of the main study (50-step median; dotted: recovery threshold).
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Left: depth-12 drivers over the αmin grid; at this depth only αmin≤0.02 collapses, and it does so when α reaches the floor. Right: the healthy network ( α=1 ) dips to the detection threshold early in training and heals by itself.
Figure 6: Force of each loss term. Left: gradient norm with respect to the synthetic token matrix as its distance δ from rank one shrinks; the penalty fades linearly, the hinge stays constant, the cosine term is non-monotone because its row normalisation diverges for tokens near the mean. Middle: ratio ρ of added-loss to task gradient at the parameters of a collapsed calibration checkpoint and at points along the straight line towards the healthy network of the same seed and, as a control, the same distance along a random direction. Right: ρ per treatment over the eight treated collapse events.
Treatment
events
ρ median
range over events
clip fraction
penalty 1 ×
8
0.0032
0.0027–0.0043
0.22
hinge 1 ×
8
0.0887
0.0550–0.1630
0.44
penalty 4 ×
8
0.0130
0.0111–0.0170
0.21
hinge 4 ×
8
0.3071
0.1942–0.6194
0.60
cosine 1 ×
8
0.0618
0.0380–0.0881
0.40
Appendix
Table 2: Ratio ρ of added-loss to task gradient norm per treatment (median over measurements per collapse event; median and range over the eight events) and the fraction of steps on which gradient clipping was active.
Figure 7: Filled: main study, post hoc. Open: the same switches on the replication networks, registered before they ran. Steps to recovery after the switch to α=0.3 against the steps spent at αmin before the switch (0: undelayed). Crosses at the top: no recovery within the budget.
Figure 8: Registered prediction on the ten replication networks. Left: outcome of the switch to α=0.1 against the predictor; dashed line: threshold fitted on the main study (grey squares) and frozen before the branches ran; crosses: wrong predictions. Right: steps to recovery after the switch to α=0.3 against the predictor at that scale.
A widely cited result by Dong et al. (2021) showed that Transformers built from self-attention alone, without skip connections or feed-forward layers, suffer from rapid rank collapse: all token representations converge to a single direction. The proposed remedy was the MLP. We show that this picture, while correct in the regime studied by Dong, is incomplete in ways that matter for architectural understanding. Three results are established. First, layer normalisation is precisely affine-rank-neutral: it preserves the affine rank of the token representation set exactly. The widespread claim that LN "plays no role" is imprecise; the correct statement is sharper. Second, residual connections generically obstruct rank collapse in real Transformers such as BERT-base, in a measure-theoretic sense, without contribution from the MLP. The MLP's irreplaceable function is different: generating feature directions outside the linear span of the original token embeddings, which no stack of attention layers can produce. Third, a phenomenon distinct from rank collapse is identified: head-channel non-identifiability. After multi-head attention sums per-head outputs through the output projection, individual contributions cannot be canonically attributed to a specific head; n(H-1)d_k degrees of freedom per layer remain ambiguous when recovering a single head from the mixed signal. The MLP cannot remedy this because it acts on the post-summation signal. A constructive partial remedy is proposed: a position-gated output projection (PG-OP) at parameter overhead below 1.6% of the standard output projection. The four collapse phenomena identified in the literature -- rank collapse in depth, in width, head-channel non-identifiability, and entropy collapse -- are unified under a symmetry-breaking framework, each corresponding to a distinct symmetry of the Transformer's forward pass.
Giansalvo Cirrincione
aLaboratoire LTI, Université de Picardie Jules Verne, Chemin du Thil, 80025 Amiens, France
Deep decoder-only Transformers often replace the original Post-Norm architecture with Pre-Norm variants because Post-Norm training is highly sensitive to warmup and learning rate under conventional initialization schemes. Although prior work has identified rank collapse and gradient vanishing as related symptoms, it remains poorly understood how causal attention creates high-similarity representations and why training dynamics fail to repair them. We give a two-stage analysis of Post-Norm rank collapse using token similarity as a scalar state variable. First, at initialization, causal attention acts approximately as a prefix-averaging operator that increases token similarity across depth, while the SwiGLU branch contributes only a smaller damping effect. Second, once training enters a high-similarity regime, growth of pre-normalization residual norms makes the RMSNorm backward factor contractive; under mild conditions, gradients to earlier layers decay geometrically. As a complementary result, we characterize the properties of a collapsed network: its best predictor is frequency distribution with relatively high loss floor, and gradients in collapsed layers vanish at frequency distribution. Experiments on 48-layer decoder-only Transformers trained on C4 dataset match the predicted initialization-time similarity growth and collapse-time gradient contraction, and show that collapsed runs stay near the predicted frequency loss. Together, these results distinguish the forward similarity amplification and backward repair incapacity in Post-Norm collapse, while also characterizing the behavior of collapsed networks.
Xingjian Wang, Qingyu Han, Xiaodong Luo +1
The Chinese University of Hong Kong, Shenzhen · Shenzhen Research Institute of Big Data
We investigate how each component of the Transformer feedforward block architecture design determines how much rank survives across depth at initialization. We reinterpret skip connections and normalization, long understood as controlling magnitude, as mechanisms for preserving gradient rank across depth, since the very matrix multiplications and nonlinear activations that make the network expressive also reduce the rank. We show that skip connections trade off rank collapse against ensemble-like behavior, controlled by the relative scales of the branch and the skip: skip connections route the gradient around the residual branch, where rank is lost, rather than along the long gradient paths that encourage the layers to compose. The placement of the normalization layer controls this same tradeoff by setting the branch-to-skip ratio across depth, unifying much of the normalization placement and depth scaling literature, in particular why rank collapses for Post-Norm but plateaus for Pre-Norm. Other aspects of the architecture, like the two-matrix structure that expands and contracts the width, use additional parameters to preserve the representation or branch Jacobian rank. The second matrix decorrelates a coherent mean spike that would grow across blocks with a single matrix and uncentered activation, preventing the residual representation from collapsing. The width expansion between the two matrices keeps the branch Jacobian full rank: applying the rank-reducing activation in this expanded space leaves enough directions to span the original, at a width that follows a Marchenko--Pastur law. The initialization rank of the input--output Jacobian predicts which networks train on CIFAR-10. Taken together, we recast architecture design for deep networks as navigating an intrinsic tradeoff among rank collapse, ensemble-like behavior, and parameter count.