PaperlayerView exact revision on arXiv

Source-linked AI summary

Force without transmission: a depth-induced rank collapse that no loss on the representation reopens

Martin Hofmann, Patrick Mäder

arXiv:2610.09958v1cs.LG

TL;DR

This paper asks whether an added loss can repair rank collapse after training has driven token representations into one direction. It tests rank-based loss terms and skip-connection restoration on collapsed transformers, finding that repair depends on gradient transmission to the weights that must change, not force size alone.

  • Problem

    Whether a loss term can repair an established rank collapse, where training no longer supplies the gradient that sharpens attention, is unclear.

    Source: §1 · pp. 1–2

  • Method

    The authors weaken skip connections to collapse small transformers, then continue saved copies under loss treatments, skip restoration, or no treatment while measuring blockwise gradients.

    Source: §1 · p. 2

  • Results

    Loss-based treatments failed in weakened-skip rank collapse because their gradients scarcely reached needed weights, while restoring the skip connection reopened the path and rank recovered.

    Source: §9 · p. 7, §11 · p. 8

  • Takeaways & Limitations

    Gradient-path measurement diagnoses whether gradients reach the intended weights, but in registered tests it did not predict recovery better than a simple loss measure.

    Source: §9 · p. 7

    The tested collapse was driven by weakened skip connections in 9–11-million-parameter models trained on one data stream, so the paper does not claim loss-based repair fails for rank collapse generally.

    Source: §10 · pp. 7–8

Abstract

from arXiv · show

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.

1 Introduction

The paper asks whether an added loss can repair an established rank collapse, where training no longer supplies a gradient that sharpens attention. It distinguishes the size of a corrective force from whether that force reaches the weights that need to change, and introduces two rank-based terms to test this distinction.

1 Introduction

  • The paper asks whether an added loss can repair a settled collapse when the architecture is fixed and training alone does not restore attention.The related attention-entropy collapse provides a contrast: a bounded-force term repaired it, motivating a test of whether force behavior also matters for rank collapse.Source: §1 · p. 1
  • A large force at the loss may still fail to act if the backward path does not transmit it to the weights that need to change.The paper calls the gradient reaching a block’s weights its transmission, and the route through the stack its path.Source: §1 · p. 2
  • The study collapses small transformers by weakening the skip connection, then compares loss treatments, untreated copies, and skip restoration while measuring blockwise gradients.It also registers tests of whether path measurements predict recovery on new networks and after a high-learning-rate burst.Source: §1 · p. 2

2 Two kinds of corrective force

  • Stable rank measures how many directions token representations span, with rank one when all tokens point in one direction.The terms target this same quantity but differ in whether their corrective gradient vanishes near collapse.Source: §2 · p. 2
  • The rank penalty’s gradient vanishes as rank approaches one, whereas the hinge retains a constant-size gradient toward spreading tokens.Both are applied across blocks and calibrated to add the same loss when switched on; a mean-squared cosine term is included as a literature comparison.Source: §2 · p. 2

3 How the collapse is produced and treated

The study creates a settled rank collapse by gradually weakening the skip connection in small language models. It then continues copies of the same saved collapsed networks under loss-based treatments, no treatment, or skip restoration.

3 How the collapse is produced and treated

  • The models are pre-LN decoder transformers of width 128, four heads, context 256, and depths 12 or 24, trained on a fixed FineWeb-Edu stream.The two depths have about 9 million and 11 million parameters, respectively.Source: §3 · p. 2
  • The skip scale α stays at 1 for 1,500 steps, then ramps down to 0.02 by step 3,000 to drive collapse.Each block scales both its attention and MLP residual branches by α.Source: §3 · p. 3
  • A network is saved after its last-block rank settles below 1.5; recovery means its 50-step median rank reaches half the healthy value.Each treatment continues the saved network for 4,000 steps, and the loss is reported separately from rank recovery.Source: §3 · p. 3
  • Treatments include penalty and hinge losses at matched added-loss levels, a cosine term, untreated copies, and a reversed ramp restoring α to 1.Two depths and five seeds yield ten main-study collapse events; the treatment design was fixed before the main runs.Source: §3 · p. 3, Table 1 · p. 3

4 The collapse is harmful and does not heal

The weakened-skip collapse is stable, harmful, and not left behind by training alone.

4 The collapse is harmful and does not heal

  • Collapsed losses were 1.73–1.94 nats above healthy networks, and untreated rank stayed below recovery in all 18 study and replication events.Collapsed attention was nearly uniform, with normalized entropy 0.94–0.99 versus 0.65–0.69 in healthy networks.Source: §4 · p. 3

5 No added loss repairs it, although its force is large

No added loss repaired the settled rank collapse, even when its corrective force was large. The task gradient scarcely reached query and key weights, while the hinge gradient faded before the blocks where collapse completes.

5 No added loss repairs it, although its force is large

  • The hinge force was 0.0887 relative to the task gradient, about 28 times the penalty’s 0.0032; at four times strength it was 0.3071.The bounded force remained large, yet it did not repair the collapse.Source: §5 · p. 3
  • Only the reversed ramp recovered; no loss-based treatment recovered rank, including when added gradients were exempted from clipping at up to ten times strength.The loss-based treatments had 0/32 recoveries, and untreated copies also did not recover.Source: §5 · p. 3, Figure 1 · p. 4

6 The force does not reach the weights that must change

  • Query and key weights received less than 10^-6 task-gradient norm per block in 17 of 18 collapsed networks, versus 3·10^-2 to 5·10^-2 in healthy networks.The task gradient was concentrated on value projections; the figure compares per-block gradients in collapsed and healthy networks.Source: §6 · p. 4, Figure 2 · p. 5
  • The hinge gradient was 0.34 in the first block but 6·10^-4 in the last at depth 12, leaving the output-near blocks with little corrective force.At the end of the budget, rank was 8–9 in the first block but only 1.6–2.7 in the middle block, as restored diversity was lost downstream.Source: §6 · p. 4

7 Restoring the path recovers the rank, with hysteresis

Restoring the skip connection reopened the gradient path and recovered rank, but only after the skip scale rose well above the collapse point. The path measurement ranked recovery in one registered test but did not transfer reliably to a learning-rate-induced collapse.

7 Restoring the path recovers the rank, with hysteresis

  • Restoring the skip recovered rank in 8/8 main-study and 10/10 replication networks, reopening the query–key gradient immediately without changing any weights.At the first step after the switch to α = 1, the query–key task gradient rose by about 107-fold even in the block with the smallest increase.Source: §7 · p. 5
  • Rank recovered at α = 0.25–0.97 at depth 12 and α = 1.00 at depth 24, at least 12 times above the collapse scale; Figure 3 compares the down- and up-ramps.With α fixed, 7/8 and 8/8 networks recovered at 0.3 and 0.6, versus 2/8 at 0.1.Source: §7 · pp. 5–6, Figure 3 · p. 6
  • Delaying the switch to α = 0.3 increased recovery time at depth 12, and no delayed depth-24 branch recovered in the main study or replication.The authors did not identify what changes during time spent collapsed.Source: §7 · p. 6
  • All 42 networks that recovered rank after skip restoration still ended 0.21–0.56 nats above the healthy network’s loss, with a median gap of 0.44.The authors note that some of this gap may reflect training time lost while collapsed.Source: §7 · p. 6

8 Does the path predict recovery?

  • The path predictor ranked recovery times in 9 replication networks (Spearman −0.67) but classified only 8/10 outcomes at the weaker switch, matching the registered test one short of its target.The simple loss-gap measure performed as well on the weaker-switch classification.Source: §8 · p. 6
  • Under learning-rate multipliers 16 and 32, Figure 4 compares query–key gradients and rank trajectories: all 8 networks recovered untreated at 16, none within budget at 32, and the path predictor did not transfer.At 16 the predictor did not order recovery times (Spearman −0.10); the loss gap did (+0.79).Source: §8 · p. 6, Figure 4 · p. 7

9 Discussion

The discussion identifies gradient arrival at the weights that must change—not force size—as the binding constraint in the weakened-skip rank collapse. It frames blockwise gradient measurement as a mechanism diagnostic, while cautioning that rank recovery does not establish functional recovery and that the findings have limited scope.

9 Discussion

  • Added loss terms failed because almost none of their gradient reached the weights that needed to change; restoring the skip reopened that path without changing weights.In the learning-rate-burst case, the path stayed open and rank recovered without treatment.Source: §9 · p. 7
  • Blockwise gradient measurement distinguished closed-path weakened-skip collapse from open-path rank loss after a learning-rate burst, but did not predict recovery better than the loss gap.The authors present it as a diagnostic of mechanism rather than a better forecast.Source: §9 · p. 7
  • Rank recovery did not establish functional recovery: every rank-recovered network ended with loss above its control, so the authors recommend continuing training and measuring loss.The reported gaps were 0.21–0.56 nats after skip restoration and 0.86–1.04 after the learning-rate burst.Source: §9 · p. 7

10 Scope and limitations

  • The study’s main collapse was induced by weakening the skip; its networks had 9–11 million parameters and used one data stream, with registered tests on 10 and 8 held-out networks.The authors do not claim that loss-based repair fails for rank collapse in general.Source: §10 · pp. 7–8

11 Conclusion

Restoring the skip connection reopened the gradient path and restored rank, but only well above the collapse scale; the loss remained above that of a healthy network. The paper also reports its reproducibility arrangements and AI-assisted work.

11 Conclusion

  • Restoring the skip connection reopened the gradient path and recovered rank only well above the collapse scale, without restoring a healthy-network loss.The conclusion attributes recovery to reopening the gradient path, not to adding a loss term.Source: §11 · p. 8

Reproducibility statement

  • The code, preregistration, protocols, constants, analysis scripts, and manuscript-generation scripts are available from the authors on request.The statement identifies a frozen commit and SHA-256 hashes for the preregistration and calibration summary.Source: Reproducibility statement · p. 8

Use of AI-assisted tools

  • Claude Code assisted with research code, run orchestration, analysis, literature lookup, and drafting; the authors designed and verified the study and take responsibility for the final text.The statement names Claude Fable 5.1 and Claude Opus 5.5.Source: Use of AI-assisted tools · p. 8

A Design, calibration and detection

The study calibrated a skip-connection driver and operational rules for detecting collapse and recovery, then compared loss-term forces at collapse. Calibration located a narrow collapse threshold, while the force measurements distinguished a fading penalty from a persistent hinge force.

A Design, calibration and detection

  • Collapse required a stable rank below 1.5 on the last block; recovery required the 50-step median to reach half its healthy value.Runs were saved as settled after a measurement at least 100 steps later still had a ten-measurement median below 1.5; two brief depth-12 dips healed untreated.Source: A Design, calibration and detection · p. 10
  • None of 23 calibration checks recovered with the hinge, including tests applying it only to the last block, to every block, or at 10× Λ.The hinge uses the top singular pair, which the passage says is separated at collapse; the loss added on the first step was within a tenth of target in 31/40 loss-treated branches.Source: A Design, calibration and detection · p. 10
  • At depth 12, only αmin ≤0.02 collapsed, and collapse occurred when α reached the floor; αmin=0.03 preserved rank at that depth.The calibration identified αmin=0.02 as the largest grid value at which every seed collapsed and no untreated network recovered.Source: A Design, calibration and detection · p. 10, Figure 5 · p. 11

B Force measurement

  • At δ=0.0001, the penalty gradient was 3.5 · 10−8, while the hinge gradient stayed at 0.022; at a collapsed checkpoint, their force ratios were 0.0013 and 0.074, respectively.The synthetic-matrix penalty gradient fell from 4.9 · 10−4 at δ=1; the cosine ratio was 0.0618 at the collapsed checkpoint.Source: B Force measurement · p. 11, Figure 6 · p. 12

C Post hoc checks

Post hoc checks found no rank recovery from the loss treatments, including when the added-loss gradient bypassed clipping. Recovery after restoring α to 0.3 became slower or failed with longer delays, while exploratory intact-skip drivers did not produce settled collapse.

C Post hoc checks

  • With the added-loss gradient exempt from clipping, neither the penalty nor the hinge recovered rank in any branch, including hinge treatments at 4× and 10×.The post hoc checks report 0/8 recovery for penalty 1×, hinge 1×, hinge 4×, and hinge 10×.Source: C Post hoc checks · pp. 12–13
  • None of the exploratory intact-skip drivers produced a settled collapse, so no treatment branches were made from them.The weight-decay ramp and initialization-gain attempts each detected and settled zero of six runs.Source: C Post hoc checks · p. 13
  • After a further 500, 1,500, or 3,000 steps at αmin, switching to α=0.3 took longer or failed; depth-24 delayed branches did not recover.At depth 12, recovery time rose with delay in every seed; the small event counts make the counts alone weak evidence, and a delayed full restoration was not run.Source: C Post hoc checks · p. 13, Figure 7 · p. 13
  • After 3,500 untreated steps, normalized attention entropy remained 0.94–0.99.The paper reports that the time-related change was not the attention pattern.Source: C Post hoc checks · p. 13

D Replication with new seeds

In the new-seed replication, no loss-based treatment or untreated branch recovered rank within the budget. Every network recovered on the reversed ramp, but only after α rose far above the collapse value.

D Replication with new seeds

  • In the replication, none of 10 collapse events recovered under a loss treatment or untreated within the budget.The same outcome was reproduced for both loss terms, with the untreated network also staying collapsed.Source: D Replication with new seeds · pp. 13–14
  • In the reversed-ramp runs, every network recovered, but only after α rose far past its collapse value; no loss-based treatment or untreated network recovered rank within budget.At the parameters, penalty force was 0.0037 of the task gradient and hinge force was 0.0782, as medians over collapse events.Source: D Replication with new seeds · p. 13

E Registered tests in detail

Registered tests compared the gradient-path predictor across replication networks, delayed restorations, and a learning-rate-burst collapse. The predictor worked on some replication outcomes but did not transfer to rank-recovery timing under the burst; recovered rank also did not imply recovered loss.

Path after the switch

  • After 50 updates following a skip restoration, rank exceeded the recovery threshold in 9 of 18 networks at α = 0.3 and 8 of 18 at α = 1.Before updates, none at α = 0.3 and 4 of 18 at α = 1 exceeded the threshold; the query–key gradient rule was met in 1 of 18 and 15 of 18, respectively.Source: E Registered tests in detail · p. 14

First test: the replication networks

  • The replication predictor was the median across blocks of the first post-switch task gradient on query–key weights, with its α = 0.1 threshold frozen at 3.1 · 10^-5.That threshold was fitted on eight main-study networks, which it separated without error; predictions and evaluation were committed before the replication branches ran.Source: E Registered tests in detail · p. 14
  • At α = 0.1, the replication predictor classified 8 of 10 networks correctly; at α = 0.3, predictor and recovery time had Spearman correlation −0.67.The registered hinge-force prediction and rival-measure prediction were missed; the loss gap classified 8 of 10, matching the path predictor. Figure 8 shows the α = 0.1 classification and α = 0.3 recovery-time comparison.Source: E Registered tests in detail · p. 14, Figure 8 · p. 15
  • Delayed restorations matched the registered prediction: none of 15 depth-24 branches recovered, while depth-12 recovery became less frequent and slower as delay grew.At depth 12, recovery fell from 5 of 5 without delay to 0 after 3,000 steps; within each seed, recovery time rose with delay.Source: E Registered tests in detail · p. 14

Second test: a learning-rate burst

  • At multiplier 16, all 8 learning-rate-burst networks recovered by rank, but at multiplier 32 none recovered within 4,000 steps.The path remained open in both conditions, with query–key gradients above the weakened-skip collapsed networks and below healthy networks; training remained numerically stable.Source: E Registered tests in detail · p. 15
  • Within multiplier 16, the query–key gradient did not rank recovery time (Spearman −0.10), while the loss gap (+0.79) and gradient norm (+0.55) did better.Across both multipliers, every measure ranked outcomes because each separated the two conditions; the frozen threshold classified 3 of 8 at 32, and development thresholds did not beat the multiplier alone.Source: E Registered tests in detail · p. 15
  • After the burst, rank recovery did not restore loss: all 8 networks ending above the rank threshold finished 0.86–1.04 nats above paired controls.Under this driver the path stayed open and rank recovered untreated; the loss gap remained despite rank recovery.Source: E Registered tests in detail · p. 15

F Alternative explanations

The alternative-explanation checks distinguish supported interpretations from unresolved scope: the learning-rate excursion preserved the path, but did not produce a settled collapse with the skip intact. The registered timing order of collapse markers was not observed.

F Alternative explanations

  • A settled collapse with the skip intact was not produced, so whether it would close the gradient path remains untested.The learning-rate excursion with the skip intact left the path open and rank recovered untreated; this does not test a settled collapse under that architecture.Source: F Alternative explanations · p. 16
  • Immediate switches argue against time spent collapsed as the source of hysteresis: most networks recovered at α = 0.3 and 0.6, while fewer than half did at α = 0.1.Those α = 0.1 switches occurred without any delay.Source: F Alternative explanations · p. 16

G Interpretation fixed in advance

  • The registered order of collapse markers—spectral gap, mean squared cosine, then loss gap—occurred in 0 of 10 events, with crossings clustered near α reaching its floor.All crossings fell within a few measurements of the floor, and the loss gap was no later than either spectral marker in 10 of 10 events.Source: G Interpretation fixed in advance · p. 16
Loading 2610.09958v1…