Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
Authors: Shikai Qiu
Lechao Xiao
Andrew Gordon Wilson
Jeffrey Pennington
Atish Agarwala
Venue: ICML 2025 oral (2025)
Keywords: Scaling Laws
Optimization
Primary area: deep_learning
TL;DR: Loss curves from compute-optimally trained models collapse onto a universal shape, from which we can derive both theoretical insights and practical diagnostics for scaling.

Abstract: What scaling limits govern neural network training dynamics when model size and training time grow in tandem? We show that despite the complex interactions between architecture, training algorithms, and data, compute-optimally trained models exhibit a remarkably precise universality. Specifically, loss curves from models of varying sizes collapse onto a single universal curve when training compute and loss are normalized to unity at the end of training. With learning rate decay, the collapse becomes so tight that differences in the normalized curves across models fall below the noise floor of individual loss curves across random seeds, a phenomenon we term supercollapse. We observe supercollapse across learning rate schedules, datasets, and architectures, including transformers trained on next-token prediction, and find it breaks down when hyperparameters are scaled suboptimally, providing a precise and practical indicator of good scaling. We explain these phenomena by connecting collapse to the power-law structure in typical neural scaling laws, and analyzing a simple yet surprisingly effective model of SGD noise dynamics that accurately predicts loss curves across various learning rate schedules and quantitatively explains the origin of supercollapse.

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
Shikai Qiu !f
Lechao Xiao
Andrew Gordon Wilson'
Jeffrey Pennington
Atish Agarwala
Abstract
2023; Achiam et al., 2023; Xiao, 2024). While the complex- ity of modern architectures, optimizers, and datasets often renders exact, first-principes analyses intractable for any individual model, recent work shows that some key aspects of training are predictable when we focus on their seal- ing behavior across family of models. Examples include empirical power-law relations linking optimal final loss, model size, dataset size, and compute budget under compute- optimal training, known as neural scaling laws (Hestness et al., 2017; Kaplan et al., 2020; Sharma & Kaplan, 2022; Hoffmann et al., 2022), as well as hyperparameter transfer from small to large models based on infinite-width or depth limits of training dynamics (Yang et al., 2021; Bordelon et al., 2023; Everett et al., 2024; Bordelon et al., 2024c).
What scaling limits govern neural network train- ing dynamics when model size and training time grow in tandem? We show that despite the com- plex interactions between architecture. training algorithms, and data, compute-optimally trained models exhibit remarkably precise universality. Specifically, loss curves from models of varying sizes collapse onto single universal curve when training compute and loss are normalized to unity at the end of training. With learning rate decay, the collapse becomes so tight that differences in the normalized curves across models fall below the noise floor of individual loss curves across random seeds, a phenomenon we term supercol- lapse, We observe supercollapse across learning rate schedules, datasets, and architectures, includ- ing transformers trained on next-token prediction, and find it breaks down when hyperparameters are scaled suboptimally, providing precise and practical indicator of good scaling. We explain these phenomena by connecting collapse to the power-law structure in typical neural scaling laws, and analyzing a simple yet surprisingly effective model of SGD noise dynamics that accurately predicts loss curves across various learning rate schedules and quantitatively explains the origin of supercollapse.
In this work, show the entire training process follows highly predictable scaling, beyond final losses and optimal hyperparameters. We find that the entire loss curves of compute-optimalle trained models exhibit a precise scaling symmetry, collapsing onto a single universal curve across models after a simple normalization. Learning rate decay amplifies this effect dramatically, producing what we call supercollapse: collapse so tight that cross-scale differences fall below the noise floor of individual loss curves due to random seeds. Figure (a-d) summarizes these results.
These findings advance our understanding in two key ways. First, while Kaplan et al. (2020, Figure 11) found the loss curves roughly follow a sum of power laws, we identify that loss curves follow universal shape with far greater precision. For typical learning rate schedules, this shape deviates from simple power laws and may not admit any obvious functional form. Second, our work provides com- pelling empirical evidence for a well-defined joint scaling limit where model size and training time grow together un- der compute-optime allocation. This limit contrasts with traditional infinite- wiith or depth limits that fix training duration (Yang & Hu, 2021; Bordelon & Pehlevan, 2022). While these theories predict initial dynamical consistency, accumulating finite-size effects lead to gradual divergence as training progresses, as demonstrated by Vyas et al. (2023). In contrast, the collapse we observe reveals a joint scaling limit that preserves consistency throughout training, pre- cisely the regime relevant for practical large-scale training.
1. Introduction
As machine learning systems grow in scale, accurate predic- tive models of their training dynamics become increasingly valuable, both for interpreting costly experiments and for de- signing robust, efficient training pipelines (Wortsman et al.,
Work done partly during an internship Google Deep- Mind 'New York University "Google DeepMind. Correspon- dence to: Shikai Qiu <sq2129@nyu.edu Atish Agarwala <thetish google..ooo.............
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Scaling Collapse Reveals Universal Dynamics in Computc-Optimally Trained Neural Networks
10-1
3.19
6 ae
Widts
= 0.154
Norma lized Loss
1.8
768
3.18
4. m 3.132
0.191
1.6
Rel. Variation
10-2
1152
1280
3.17
1.4
10-1
3.16
1.2
2040
Noise floor
Collapse deviation
4
1.0
10-4
10f Compute (PetaFLOPs)
10
0.0
0.2
0.4
0.5
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Normalized Compute
Normalized Compute
(a) Computc-Opting Scaling Law
(b) Supercollapse
(c) Collapse Below Noise Floor
1-
e.52 + onjnx))
1- x (20%)
10-
(20))
3.20
Collapse Dev. ///
Normalized Loss
1.5
on ca(exs) 011 (ott-o)
Oiseved
3.19
Prediction
10-2
1- (9S) 050 eegttz))) (205) mx(ferl)
3.18
10-3
3.17
vI
10-4
10
10-2
10t2 f
10a
Normalized Compute
Tokens (Billion)
(d) Collapse in Different Schedules
(c) Predicting Loss Curve per Schedule
(1) Predicting Collapse Quality
Figure 1: Scaling collapse of compute-optimal transformer loss curves and its explanation through model of SGD noise dynamics.
(a) compute-optimas loss curves and fitted scaling law on CIFAR-5M using a linear learingg rate decay schedule. (b) Normalized reducible loss curves collapse onto a single universal curve independent of model size, with both final compute and reducible loss normalized to unity. (c) Collapse deviation A (cross-model variation of normalized loss) falls below per-model noise floor a (variation of reducible loss across random seeds) for much of training, phenomenon decay phase of various learning rate schedules, each producing its own universal curve. To explain these phenomena, we show that simple model of SGD noise dynamics (e) accurately predicts loss curves for different schedules across model scales (Section 3.2) and (f) quantifies how learning rate decay improves the collapse due to the predicted scaling A vn(1 +), where 1j is the instantaneous learning rate and normalized gradient flow time (Section 3.3). We observe supercollapse in other arhitectures and datasets (Figure 4).
term supercollapse. (d) Supercollapse occurs during the
We provide elementary theoretical analysis that reveals the key mechanisms behind this precise collapse. We first show that for loss curves following typical neural scaling laws, collapse occurs precisely when models are trained for constant multiples of their compute-optimal horizons (Sec- tion 3.1). We then analyze a simple theoretical model of the SGD noise dynamics that predicts loss curves under a vari- ety of learning rate schedules remarkably well (Section 3.2), and explains two key observations: why normalized curves retain universal form despite losing their power-law struc- ture, and how learning rate decay suppresses variance to produce supercollapse (Section 3.3).
2.1. Experiment Setup
In each task, we train sequence of models with increasing compute, scaling hyperparameters such as data, initializa- tion, and learning rate with the model. We refer to se- quence of training configurations as a scaling ladder. We provide further experimental details in Appendix A. We fo- cus on width scaling, where hyperparameter transfer most well-studied, but find scaling transformer depth leads to similar results in Appendix B, suggesting our observations may generalize to more general scaling ladders where width, depth, batch size, weight decay, etc. can be co-scaled.
Transformers Next-Token Prediction. We consider two next-token prediction tasks: 1) CIFAR-5M (Nakkiran et al., 2020), a dataset of 6M generated CIFAR-like images, and 2) Lichess, collection of chess games recorded in alge- braic chess notation where the goal is to predict the next move in the game. Our scaling ladder includes models with about 10M to 80M parameters, approximately log- uniformly spaced, by scaling the width (embedding dimensions) from 768 to 2048 and fixing the number of blocks to 3. All models use jP (Yang & Hu, 2021; Yang et al., 2021) for initialization and learning rates, and are trained with Adam.
Beyond theoretical interest, supercollapse provides prac- tical scaling diagnostic, as we find that deviations from collapse can signal misconfigured scaling choices, such as suboptimal scaling of learning rate and data (Figure 4). Overall, our results suggest supercollapse provides a novel, powerful tool to study scaling. Our code can be found here.
2. Empirical Observations
We demonstrate our main empirical findings in this section, independently multiple tasks and architectures which can be studied even in academic settings.
ML.Ps on Power-Law Fourier Features. To investigate
2

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
other architectures and training objectives, we train 7-layer MLPs with varying widths from 384 to 2048 on a synthetic regression task. The target function has power-law Fourier spectrum, designed to elicit the power-law scaling laws observed in natural data. We count each example as token.
22n1
Nomaltzd Los
Rel. Viriation
1024
1.4
1223
Nuis floor
Collapse deviation
10
66 60
Normalized Compute
Normalized Compute
2.2. Estimating Compute-Optimit Scaling Laws
Figure 3: Collapse with constant LR schedule. (Left) Esti- mated mean and 90% confidence interval (shaded) of the normal- ized loss curves. (Right) A is comparable to a without LR decay.
Let [(t, p, w) be the test loss after tokens (proportional to steps) for a model with parameters trained with random seed w. We estimate the compute-optimal training hori- zon in tokens for a p-parameter model as t" (p) = (p/po)", where y is the data exponent. by extracting the Pareto fron- tier of expected loss (estimated using 5 seeds) vs. compute under a constant learning rate schedule, following pro- cedure similar to Approach 1 in Hoffmann et al. (2022), with compute estimated as c = 6èp FLOPs (Kaplan et al., 2020). We reuse the same t* (p) as the training horizon for other learning rate schedules, which prior work suggests is optimal up to a constant factor (Pearce& & Song, 2024). For each task and schedule,
for some offset P. We refer Note the denominator uses the stochastic final loss value specific to the random seed. We set ī, = Lo to subtract the estimated irreducible loss that bottlenecks the asymptotic performance, leading { p, w) C(xt" (p).p.)) [P)
the normalized compute.
Remarkably, we observe that the family of normalized loss curves is nearly identical across p. revealing equal rates of relative progress (Figure lb). We say these curves collapse, as the phenomenon resembles the ubiquitous scaling col- lapse found in statistical physics, theoretical biology, and other sciences, where observables from systems of different sizes collapse onto a single curve after appropriate rescaling (see Appendix D for further discussion). We found setting L. Log achieves the best collapse (Figure 2).
fit the resulting compute-optinal scaling law using the form Lo + ac-b (Figure 1a), for con- stants Lo, a, b 2 0. Following Sharma & Kaplan (2022) and Hoffmann et al. (2022), we refer to Lo as the estimated irre- ducible loss. Using the best-fit Lo, we define the reducible loss curve [(I, p) = [(Ipp w) - Lo- We detail the pro- cedure for fitting the compute-optimal training horizon in Appendix C. "Compute-optimal" here primarily refers to the choice of training horizon, not of all hyperparameters.
2.4. Quantifying the Collapse Quality
We quantify the quality of collapse using the collapse devia- tion A, defined as:
2.3. Scaling Collapse of Compute-Optimal Loss Curves
The loss curves for different model sizes cover varying ranges of compute and loss values, but appear to follow consistent shape, which motivates us to affinely rescale them to the normalized loss curve given by
V (p,ll(x, Epp.o [ p.
A(x)
(2)
where Ep and Vn denote the expectation and variance over the random seed and the empirical distribution of model size p in the scaling ladder (approximately log-uniormmly distributed). The collapse deviation measures the relative variation of the normalized curves across For perspective, we compare it to the per-model (relative) noise floor:
(x).p.w) = L(xt" ((tt(p). w)) - L ((t - LL L
EE xE
(1)
0
[ 512
Nomallteed Love
1.000
Manralized Losk
Vu[C(at (p) (.][((((())
1.004
(x,
(3)
1.00
36
36 Normafzed 3.14
which measures the relative fluctuation in the reducible loss curve for each model size p across random seeds.
Normaliadd
[ Lg 3132
By the definition of e, A(1) always. As seen in Fig- ure 3, for constant learning rate, A(z) quickly rises to level comparable to (x,,), and remains at that level for most x < 1. This observation shows that variations in the normalized curves arise primarily from seed-to-seed fluctua- tions rather than model-to model differences, quantitatively demonstrating that the observed collapse is non-trivial.
Nornaltedd
65
144
2.
3.1
06
Normaliee Compane
Foemalooed Compate
Figure 2: Subtracting irreducible loss leads to the best collapse.
Setting L to values far from Lo breaks collapse on CIFAR-5M.
3

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
0.048
Const IR
Vieth
Widn
0.040
AP
Nornalized Loes
1.6
512
384
Normalized Loss
2.5
312
641
645
Lo65
0.032
1.4
612
1254
2.0
312
10N
1230
1200
0.024
1/2
1.5
1425
0.016
10
1.0
10
10a
10a
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Compute (PetaFLOPs)
Normalized Compute
Normalized Compute
(a) MLP μP vs Constant LR
(b) μPP (supercollapse)
(c) Constant LR (no collapse)
0.750
Over-trained Optimal
Normalized Loss
1.4
Virth
755
Normalized [8
1.3
With
0.725
Loss
0.700
1152
1.2
1152
1.2
1335
0.675
1.1
2048
1732
1.1
2048
1700
0.650
1.0
1.0
10'
10 Compute (PetsFLOPs)
10'
0.0
0.8
1.0
0.4
0.5
Normalized Compute
Normalized Compute
(d) Transformer Optimal vs Over-trained
(c) Optimal (supercollapse)
(f) Over-trained (no collapse)
Figure 4: Collapse provides a practical indicator good scaling, as suboptimally scaling key hyperparameters breaks the collapse. With the default setup, we observe supercollapse in MLP regression (b) and transformer trained on chess (c), but even changes that only lead to minor worsening in the scaling law can manifest as significant disruption to the collapse. (Top) Replacing JPP with a constant learning rate cross models for MLPs. (Bottom) Increasing the data exponent from estimated compute-optina value 1.02 to 1.2 for Transformers trained on chess. We perform a separate power-law fit to determine the value Lo for each scaling ladder.
2.5. Supercollapse: Consistency Below the Noise Floor
Model Parameterization. Carefully parameterizing the model, i.c., scaling the initialization, learning rate, and pos- sibly other hyperparameters as model size increases, is cru- cial for achieving stable and efficient training at scale (Yang et al., 2021; Bordelon et al., 2023; 2024c; Everett et al., 2024). When models are trained in the wrong parameteri- zation, we expect the loss curves not to collapse due to a lack of consistent training dynamics across scales. Using the MLP setup, we show that replacing jP with a constant learning rate across widths breaks the collapse (Figure 4, top row). Remarkably, the normalized loss curves expose inconsistent dynamics even at small scales where the fi- nal losses are virtually identical between constant and jP scaling, demonstrating that the collapse is more sensitive probe of sealing behavior than final performance alone.
Remarkably, with learning rate decay. find that the col- lapse deviation is less than the noise floor for a significant
fraction of training; that is, (xx) < or (x.,)) p) for x > 1 5 for some moderate 8 as large as 0.5 (Figure lc). We refer to this stronger form of collapse as supercollapse (in contrast to the collapse in Figure 3). Supercollapse appears in the decay phase of all tested learning rate schedules that decay to zero (Figure Id). All schedules are defined in terms of relative training fraction, i.e., the learning rate is fixed function of the normalized compute x across model sizes.
Under supercollapse, seff-normalized loss curves from dif- ferent models collapse better than our ability to predict any individual model's loss. Normalizing by the final loss of the particular realization of the stochastic loss curve is key to supercollapse, which reduces variance by exploiting correla- tions at different times along single optimization trajectory. We explain this mechanism in detail in Section 3.3.
Compute-Opting Data Exponent. For language models, Kaplan et al. (2020) showed that compute-optimal training corresponds to training each model to a fixed multiple of its converged loss. If this principle generalizes to our set- ting, the data exponent y should match the compute-optimal value for collapse to occur. For example, when y exceeds the optimal value, larger models will make more rapid rel- ative initial progress but decelerate later as a function of normalized compute, causing their normalized curves to shift downward. We indeed find this shift in Figure 4 (bot- tom row). This sensitivity suggests a novel application: rather than fitting power laws to sparse points on the Pareto frontier, one could tune 7 to maximize collapse quality, leveraging the full statistical power of entire loss curves.
2.6. Suboptimal Scaling Breaks Supercollapse
Supercollapse provides practical method for comparing inherently noisy training loss curves across model scales with precision that exceeds naive noise floor estimates, with- out the need for expensive multi-seed experiments typically required to obtain equally clean signals. This comparison can provide valuable diagnostic information about scaling where the ability distinguish small signal from noise is often crucial (Xiao, 2024), which we now demonstrate.
4

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
Scaling
0.9h"*
T--7"
7 1,1""
50
mormellied Lows
2000
9.00
3.21
( 0.999)
3134 c 38
RM
2.0
1.75
175
Compate-Opcinal
10'
133
3.20
10F
1.55
123
3.19
1.0
10
20
1.0
0.00
1.00
6.00
0.25
3.18
205
2rr1 (enti yote
Norma Loss
3.17
3.16
1.25
10 Compute (PetaFLOPs)
10a
12a
107 10
6.0
023
0.50
0/75
0.00
0.5
0.50
0.73
1000
640
0.25
0,30
2.15
L02
Compute (6ip)
Normalized Compute
Normalized Compute
Normalized Compute
(a) CIFAR-5M Fit
(b) Exact Power Laws
Figure 5: Scaling collapse from sum power-law curves. (a) CIFAR-SM expected loss curves (averaged over seed) without learning rate decay agree well with the sum-ol-power-laws fit L(t,p) p) = Lo + -- +pp (constant multipliers not shown), a form commonly observed in natural data. We omit the first 1B tokens to avoid fitting the early time transients. (b) Simulated exact sum-of-power-bee loss curves show scaling collapse precisely when the data exponent y is the theoretical compute-optine value y². Small variations of around 7 lead to nearly negligible worsening in the resulting scaling law but dramatically disrupt the collapse.
3. Explaining Loss Curve Scaling Collapse
for future theoretical work.
In this section, we investigate theoretical explanations for the scaling collapse of compute-optimal loss curves and supercollapse. Our analysis starts with a simple observa- tion: the numerator of the collapse deviation ((x) can be decomposed as:
3.1. Scaling Collapse from Power-Law Scaling
In this section, we consider deterministic models of the loss curves and assume all randomness has been averaged out.
Power-Law Pareto Frontier is Necessary. For family of differentiable loss curves (t,,)) the compute-optimal loss frontier after subtracting L must follow power law for our affine transformation to induce scaling collapse (proof in Appendix E). The key insight is that collapse requires the transformed loss curves to be related by multiplicative scaling, equivalently translation in log-log space, where the frontier must have constant log-log slope since it remains tangent to shifted versions of the same curve. This motivates choosing f, = Lo, which by definition yields the best power- law Pareto frontier. However, a sufficient condition for scaling collapse requires an explicit form of L(t,p).
"p../((z,,,uu)) = = V,Ewl((x, + E.V &b w] (4) The first term corresponds to the variation between different scales p after averaging over all sources of randomness. We will first show how this term can be small:
In Section 3.1, we prove that for family of power- law neural scaling laws, compute-optimal loss curves indeed collapse after normalization. We show loss curves in our experiments fall into this family when using a constant learning rate schedule.
In Section 3.2, we develop simple theoretical model that successfully predicts the empirical loss curves un- der various learning rate schedules and explains why they collapse despite deviating from power laws. Given its effectiveness, we believe this model has value for understanding learning rate schedules more broadly.
Neural Scaling Laws. Motivated by empirical neural scal- ing laws in natural data (Hestness al., 2017; Kaplan et al., 2020; Hoffmann et al., 2022), we consider expected loss curves following a ofm-of-power-et
scaling of the form
L (t Lo - +
(5)
for constants Lop 2 14 V > O, with potential constant multipliers absorbed via an appropriate choice units. In Figure 5a, we show the CIFAR-5M loss curves are well-fit by Equation (5) if trained under a constant learning rate schedule (averaged across 5 seeds). We also find decent fits in other datasets in Figure 11.
We then analyze the second term, which captures the loss variance due to random seeds, averaged across model sizes:
In Section 3.3, we show the same noise model enables to reason about the noise in the loss curves, and quantitatively predict the variance reduction effect in supercollapse.
Equivalence by Balance of Power Laws. As before, let t" (p) denote the training horizon. We will examine con- ditions under which t* (p) (a) is compute-opiimal, and (b) results in scaling collapse. We assume deterministic loss
Together these findings provide an initial theoretical expla- nation for supercollapse, and uncover promising directions
5

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimall Trained Neural Networks
curves for now and omit the argument w. To find compute- optimal t* (p), we fix e so that t(p) = c/ (6p) and minimize t(p) +p- with respect to by set-
exact theoretical model seems out of reach for the realis- tic training setup, we show that simple model based on quadratic loss analysis proves surprisingly effective. Un- der this model, we demonstrate that although learning rate schedules deform the loss curves in a schedule-dependent way, the deformation is approximately independent of p. We consider stochastic effects that depend on the random seed w, but omit w as an explicit argument for brevity and use bar to denote expectation over w.
the loss [(t(p).p) ting 0 = di 3C dt
= --t-t- 1- (-t/p) - vp --1
dp oc dp
Op
pt vp
(6)
which yields 6* (0) = p /e, with r = w/10 Under this scaling, the normalized loss curves are:
((r,p) = + rp
ræ " +
1
rx
3.2.1. A SIMPLE MODEL FOR LR SCHEDULES
r r+1
Let w=t) and L(w(t)) denote the parameters and loss at step I, we can model the dynamics of full-batch gradi- ent descent under a small learning rate (7()) with a gra- dient flow dw = - (y(t) ((w))) To model stochastic dr
All p dependencies cancel, leaving the final expression inde- pendent of and giving us an exact collapse. Moreover, it is clear that this is the unique choice for 1" (p) up to a constant multiplier that leads to such cancellation. This agreement is not an accident: compute-optinn optimal scaling requires balancing the derivatives of two power laws, while collapse requires balancing the power laws themselves. For power laws, these two conditions coincide, up to a multiplicative constant.
effects, a noise term is added to the gradient, leading to the SDE die dF (t) (V((t)(V(())) (V[w)) (w))(w)((()) (Li et al., 2017; Malladi et al., 2022), where the mini-batch gradient noise Σ' ('((()(((t) satisfies [[g(t){(t)) = ō(t - I')I, and we allow its covariance (which depends on batch size) E(w) to be a function of the parameters. Prior works have used the SDE model or discrete variants to study learning rate schedules in analytically tractable problems (Zhang et al., 2019; 9; d' Ascoli al., 2022; Wen al., 2024), but we will show it can make surprisingly accurate predictions in real models. We work in gradient flow time T(t) = So nss)ds, where
In Figure 5b, we numerically verify the agreement between collapse and compute-optimal scaling. When the data expo- nent y deviates from the optimal value s/p, we observe a suboptimal scaling law and no collapse. Note that the ab- sence of irreducible term in E is also necessary. Had we set i. = Lo + E for some F * 0 in Equation (1), would instead have f(xr,) = (att) "+p "+E
dw (VI((w) ==1/2(1)((T)
[e-)---p-' " +EE where no l" (p) can leave the numerator and denominator homogeneous in p.
dr
and ([E())((7') = 5(t +4) nT))6( L')L. We over- load the notation and use 77()) w(().. and L(r) denote the evolution of these quantities in gradient flow time.
In Appendix F, we study the more general form
m a,t Hip
L(l,p) Lo
Quadratic Loss. For the moment, let us suppose the loss function is quadratic L(w) = 1 ww" Hw, where we assume the minimum is at the origin without loss of generality. Then V ((w)= = H w and standard calculation shows
which naturally arises in theoretical models of neural scaling laws (Paquette et al., 2024b; Bordelon et al., 2024a;b), and implies scaling collapse by
show that compute-optimat balancing the two dominant terms, though with m >> the collapse is only exact asymptotically.
L
w(T) e Hr (0) (0)
ds e H(r H((--) 1/2 (w
((s))
(10)
Together with the close empirical fit in Figure 5a, this anal- ysis explains scaling collapse in the constant learning rate setting; however Equation (5) fails to fit the empirical loss curves with most learning rate schedules, as varying the learning rate can modulate the loss curve in quite arbitrary ways, clearly shown in Figure ld. Why, then, does the collapse transfer to other schedules?
Letting E(s) = E (w)(s)) the expected loss is then
I(T) E E|le Hr w (0)
5(rr
1 I ds dsn(s) Tr
(ie He 211(7 (())-))((()))))))))))))))))))(
+ 2
(1I)
3.2. Universality of I Learning Rate Schedules
8(T)
To understand why scaling collapse is robust across learning rate schedules, we develop a quantitative model for how learning rate schedules affect the loss curves. While an
The first term $ (7) is the forcing function, equal to the
expected loss curve in the noiseless limit n/> - 0 and
6

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimaly Trained Neural Networks
p
20M. T 10B
=11MM 31M, 79M}, T = 10B
p p= 20M, T T= (2B, 20B)
3.20
uag + conftrs))
1-x (20%)
3.20
coeet
1-x
3.19
0.501 oo(cx])
saaBaa)) (20%)
1-x (22%)
(27%)
3.19
3.50 (n( (20%)
0.50 codBoo)) (20)
Loss
08
3:18
3.18
3.18
3.17
3:17
3.17
3.16
2
10
0
4
8
10
10
15
20
Tokens (Billion)
Tokens (Billion)
Tokens (Billion)
(a) Vary Schedules
(b) Scale Model Size
(c) Scale Training Horizon
Figure 6: A simple model predicts Transformer loss curves trained across learning rate schedules, model sizes p, and training horizons T on CIFAR-SM. Dashed curves show the predicted loss according to Equation (18), with a = 0.21, which closely match with the true curves in solid. Fach curve is smoothed with an exponential moving average with half-life equal to 1% of total steps.
independent of the learning rate schedule. The second term E(T) is the excess loss due to SGD noise, which is a sum of exponential moving averages (up to normalization) of the gradient variance scaled by the learning rate over each eigenmode. Substituting in the specific forms for Σ recovers the convolutional Volterra equation for linear regression ana- lyzed in Paquette et al. (2021; 2024a), or the noisy quadratic model in Zhang et al. (2019) for small learning rates.
3.2.2. PREDICTING Loss CURVES ACROSS SCHEDULES
We apply this simple model to predict empirical loss curves in the CIFAR-5M experiments. We measure the trace of the preconditioned gradient covariance on a fixed set of 2M tokens (see Appendix A for experiment details).
Let L, 1), E be a given reference trajectory and L' L + 8L, 17 + 8y, = Σ dE be the target trajectory, Equation (14) allows us to predict the target loss via
If E varies slowly compared to the timescale of the expo- nential moving average, we can make the approximation (s)) (TT)((T) inside the integrand, giving us:
sL(T) 1 T=((((+))(6)))
(15)
I mग(T) TH (S(+) L' dse (2H((=)))))
We use (12) a constant learning rate for the reference trajectories and various schedules sharing the same peak learning rate for the (13) target. Decomposing 6() = 6σ + +706E, we find the first term is typically 3 to 10 times larger than the second as the learning rate decays, which can be attributed to how learning rate interacts with curvature (Figure 13). In Figure 6, we only keep the first term, and predict the target loss
where 8(n(r)E(()) = (
E(T) ~
1 11 - + =
For large the expected loss is then approximately
4 (T))(( (T)())
(14)
T'(T) R l(r) + dn(()) T+(E'(T)),
(16)
Given access to (T)())), we can derive prediction for how the loss changes we change the learning rate schedule without knowing F
where a is a shared hyperparameter. We find a single a 0.21 fits the target loss curves surprisingly well across schedules, model sizes, and training horizons. In Appendix I, we show even better fits for MLPs, though puzzlingly, including the second term can lead to worse fits.
General Case. In Appendix G. we discuss how this analysis can be generalized to more realistic setups. For general loss functions, we show via perturbation theory that, to first order T one can make similar approximations to derive Equation (14) given an additional assumption that the Hessian is slowly varying, and with the forcing function F (r) no longer admitting a quadratic form. We also show in Appendix G that Σ should be the preconditioned gradient covariance when using adaptive optimizers. We absorb the layerwise, width- detth-dependent learning rates in jP into the preconditioner, similar to Noci et al. (2024), so 17(t) E [0, 1] reflects only the schedule.
Recent works proposed more complex functional forms for how learning rate schedules affect loss curves, derived pri- marily from empirical observations (Tissue et al., 2024; Luo et al., 2025). The accuracy of our simple model suggests it captures the essential dynamics, and erucially, the correct scaling of the excess loss through Tr(E') so that a single a is predictive across model sizes, schedules, and training horizons. Notably, Luo et al. (2025) experimented with a similar form to Equation (16) but with a constant Σ', which likely explains the reduced effectiveness they observed.
7

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
1 x
0.5(1 cos(3mxx)
x (20%)
normalized loss curves under two schedules S and S Let v(*: map the normalized compute under S' to the normal- ized compute under S at matching gradient flow time, where y is independent of for schedules defined in terms of the normalized compute. Let 0fj(x)) = 6n) (p).p) p) be the dif- ference between the two schedules. Assuming small relative fluctuations (Bic(x)/ R ([[(()) we have:
Tr(E)/C
1.0
3.5
2001
70M
25
50
25 Takens (B)
50
25
90
Tokens (B)
Takens
1.5
Tr(EE(C
('(xt(pp)PP)
2.1
(x, 2
(p)t
(19)
((y(x))(()),1))) (o0[]-=(())(())))
35 Normalized Compute
0.5 Normalized Compute
0.5
(20)
Normalized Compute
(ly(L)+<p)) ( ah(1)on(11)
(1(
Figure 7: Universality of gradient noise CIFAR-5M. Fixing a learning rate schedule, the ratio TC is approximately function of normalized compute alone, independent of model size. We show similar results with MI. regression in Figure 14.
= =(((()), ah(z)oń(z)
(21)
independentt of
which shows that, in expectation, collapse under one sched- ule (e.g. constant) implies collapse under any other schedule, provided we take Equation (18) be exact. Since collapse under a constant learning rate can be attributed to the sum- of- (f-powerrlawwe scaling law, this result helps explain why we also observe collapse in other schedules.
3.2.3. UNIVERSAL SCALING OF GRADIENT NOISE
For typical loss functions, the gradient covariance can be related to the loss itself. For example, in noiseless high- dimensional linear regression with Gaussian features drawn from.N(O, K), we have Tr(E) x Tr(K) (Paquette et al., 2021), an intuitive result since the gradient scales with both the prediction error and the input. For non-linear regression, Newton
This analysis also suggests that collapse can serve as a filter for identifying interventions that yield scalable improve- ments: those that multiplicatively shift the reducible loss curve by the same factor across all model sizes.
K should be taken to be the time- varyin matrix for a first approximation. In this case, Tr(K) is known to depend strongly with the learning rate (Agarwala & Pennington, 2024), but we expect weak dependence on model size given our models are trained with μP (see Noci et al. (2024) for evidence that curvature statistics depend weakly on model size in jPP). Since the schedule is function of the normalized compute t/t" alone, we hypothesize there exists schedule-dependden function h(z)) such that
3.3. Supercollapse as Variance Reduction
Lastly, we turn to understanding the "super" in supercol- lapse: why does learning rate decay significantly improve the collapse, to the extent that the collapse deviation A(x) drops below the per-model noise floor o p) for a substan- tial fraction of training? Again, the simple quadratic model provides quantitative insights into this phenomenon.
Tr(E(zt" (p))) (()))((()(())))))))))))))))))()) x h(z),
(17)
V (e and the decomposition:
Recall A(x) =
which we verify in the regression (Figure 14) and next-token prediction experiments (Figure 7).
By [[(=,p.))]
Vp.w [1. p)] E. V. [((x.))) V V,E,
(22)
Combining Equation (16) and Equation (17) and making p-dependence explicit:
The first term measures variance due to the seed alone, aver- aged over model sizes, while the second term measures vari- ance due to varying the model size, having averaged over the seeds first. Since we observed that variations in the normal- ized curves primarily arise from seed-tt-seed fluctuations rather than model-to-model differences (Section 2.4) under a constant schedule, and switching to other schedules does not significantly increase the model-to-mode differences (Section 3.2), we will assume the first term V., p)] dominates, which implies Ax = (x, where (r = V. (x, ()).,)))/// (x, is the squared modll collapse deviation.
('(T,p) x C(T, ah(2))n((T.)))
(18)
where x is the normalized compute at gradient flow time r. We leave to future work an explanation of why this relation appears to hold for cross-entropy loss despite the presence of non-negligible irreducible loss, as this setting is analogous to regression with label noise, where the gradient covariance should scale with the total loss rather than just the reducible component, i.e., Tr(E) R 2T. Tr(K).
Scaling Collapse Across Schedules. Combining our in- sights so far, we can now understand why collapse happens across schedules. Let E(x, p) and 8 ((xx,) p) be the expected
To simplify notation, we temporarily omit p-dependence and write in terms of instead of x. Letting L(t)
8

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
10-1
13-1
10-1
1/9
Collapse Dev. //1
Collapsa Deviation
10-2
Collapse Dev.
10-2
10-2
10-3
5(1
10
v1 7
x (20%)
10"1
10"
022
0.4
0.0
0.0
0.B
10
Normalized Compute
Figure 8: A quantitative explanation of how learning rate decay leads to supercollapse. Across schedules on CIFAR-5M, collapse deviation at normalized gradient flow time f follows the predicted /n (1 F) scaling, capturing the noise accumulated hetween that point and end of training. A schedule that decays faster has smaller 7) and (1 f) at fixed normalized compute or training step.
(1)) + (t)), where & is the relative fluctuation, we have ((t)= 200(1+000) ~ E(t)(1+ + (t) ()),), assuming ((+))1+0(0))
4. Discussion
Scale has enabled remarkable progress in machine learning, but thorough scientific understanding of scaling remains elusive. Key open questions include identifying robust prin- ciples that guide general hyperparameter transfer and char- acterizing scaling limits under realistic scaling ladders. Our discovery of supercollapse provides empirical evidence that model-size and data joint scaling limit generically exists in the compute-optimal regime, and that the scale-invariance of the training dynamics revealed by the collapse can diagnose proper hyperparameter configuration. We believe further investigation of these phenomena holds great potential for advancing the science of scaling.
* < 1. Therefore,
(t ~ E(((t) )(t))
(23)
We see that what controls the relative variance in E(t) is not (t) but the difference (t) (tt)), which roughly cap- tures the amount of optimization noise accumulated between time / and time '* Since the optimization noise per step scales with the instantaneous learning rate, decaying the learning rate over time will precisely serve to decrease the variance in f. By contrast, the squared per-model noise floor (t) is simply [y (t)), which captures the total cumula- tive optimization noise. Importantly, had we normalized by the expected rather than the empirical final loss in f, (t) would reduce to (t). Normalizing by the stochastic final loss is essential for supercollapse, where it acts as a control-variate (Glasserman, 2004), leveraging the strong time-correlation of stochastic fluctuations along the opti- mization trajectory to cancel much of the shared noise and thereby sharply reduce the variance of the collapsed curve.
We see many exciting extensions to this work. Empirically, our small-scale experiments provide a proof-of-concept. While small-scale proxies capture certain behaviors in larger systems (Wortsman et al., 2023), validating at larger scales and with practical scaling ladders, where width, depth, batch size, and weight decay are co-scaled (McCandlish al., 2018; Wang & Aitchison, 2024; Dey et al., 2025; Bergsma et al., 2025), is important and may yield new insights into optimal scaling and hyperparameter transfer. Scaling col- lapse beyond the form we studied here can be a general tool to study other scaling relations (Tamai et al., 2023).
Quantitatively, we can estimate Å under the quadratic model in Section 3.2. Let (w(() and AC( be the fluctuations of the parameters and loss from their means. We have Aw(() = Jo ds e -) (s)) and AL(T) = g(T)T (w(T) to first order in (W(T), where g(T) is the ex- pected gradient. Close to the end of training, for 7 = -5T where 7* is the final gradient flow time and 87 > is small, direct calculation shows (Appendix H)
While we have identified the key ingredients underly- ing supercollase--poeert power-law scaling and learning rate- dependent noise dynamics- analysis relies multiple approximations and takes power-law scaling given, sug- gesting deeper theoretical principles may be at work. Taking collapse as starting point instead may provide an alterna- tive route to understanding scaling laws, analogous to how in physics the renormalization group was developed to provide a unified set of principles explaining both universality and its associated power laws (Wilson, 1971). Finally, it would be interesting to understand why simple noise model predicts the impact of learning rate schedules on real models so effectively, compare it with alternative models such as Schaipp et al. (2025), and to test its predictive power for optimizing schedules, learning rates, and training horizons.
((())=->(())o) ())) - 0(8r"), (24)
In linear regression, x L and g x L, so we estimate (7) & (T))ōT to leading order. Since this relation holds for each model size p. predict (+) n(f)(1 -), where i T/T" denotes the normalized gradi- ent flow time. Figure 8 shows this form fits the observations well, with A(F)/ /f(f) approximately following the same vI-f scaling across many schedules, quantitatively ex- plaining how learning rate decay leads to supercollapse.
(+) x

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
Acknowledgements
Bordelon, B., Noci, L.., Li, M. B., Hanin, B., and Pehle- van, C. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. arXiv preprint arXiv:2309. 2023.
We thank Courtney Paquette and Zixi Chen for helpful com- ments on an earlier version of this paper. SQ was sup- ported by Google's TPU Research Cloud (TRC) program: bctps://siteet www.eesrreserrh... www.:eaacc...ggll/eeee..........
Bordelon, B., Atanasov, A., and Pehlevan, C. A dy- namical model of neural scaling laws. arXiv preprint arXiv:2402.01022 2024a.
Contribution Statement
SQ designed and conducted the majority of experiments, led the theory development, and wrote the paper. 1.X initially observed supereollapse, contributed to theory, experimental design, and writing the paper. AGW advised SQ and edited the paper. JP contributed to theory, experimental design, and writing the paper. AA managed the research project, proved some theorems, guided the theory development, contributed to experimental design, and helped write the paper.
Bordelon, B., Atanasov, A., and Pehlevan, C. How feature learning can improve neural scaling laws. arXiv preprint arXiv:2409. 17858, 2024b.
Bordelon, B., Chaudhry, H., and Pehlevan, C. Infinite limits of multi-head transformer dynamics. Advances in Neu- ral Information Processing Systems, 11135824-55877 2024c.
Cohen, J. M., Kaur, S., Li, Y., Kolter, J. 7., and Talwalkar, A. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065 2021.
Impact Statement
This paper presents work whose goal is advance the field of Machine Learning. There are many potential societal consequences of our work. none which we feel must be specifically highlighted here.
Cohen, M., Ghorbani, B., Krishnan, S., Agarwal, N., Medapati, S., Badura, M., Suo, D., Cardoze, D., Nado, Z.., Dahl, G. E., al. Adaptive gradient methods the edge of stability, arXiv preprint arXiv:2207. 244844447.1448 2022.
References
Achiam, J., Adler, S., Agarwal, S., Ahmad, L.,, Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08771 087744, 2023.
Cohen, M., Damian, A., Talwalkar, A., Kolter, Z., and Lee, J. D. Understanding optimization in deep learning with central flows. arXiv preprint arXiv:2410.24206, 2024.
d' Ascoli, S., Refinetti, M., and Biroli, G. Optimal learning rate schedules in nigh-dimensional non-convex optimiza- tion problems. arXiv preprint -xxi:2202.0050000 2022.
Agarwala, A. and Fisher, D. S. Adaptive walks on high-dimension fitness landscapes and seascapes with distance-dependent statistics. Theoretical population bi- ology, 130:13 13:49, 2019.
Dey, N., Zhang, B. C.,, Noci, L., Li, M., Bordelon, B., Bergsma, S., Pehlevan, C., Hanin, B., and Hestness, J. Don't t be lazy: Completep enables compute-efficient deep transformers. arXiv preprint arXiv:2505.01618, 2025.
Agarwala, A. and Pennington, J. High dimensional analysis reveals conservative sharpening and a stochastic edge of stability, 2024. URL htttts://aaxiii org/ abs/ 2404. . 9261.
Diederik P. Kingma, J. Adam: A Method for Stochastic Optimization. International Conference on Learning Representations (ICLR), 2015.
Bergsma, S., Dey, N., Gosal, G., Gray, G., Soboleva, D., and Hestness, J. Power lines: Scaling laws for weight decay and batch size in I1m pre-training. arXiv preprint 13335 2025.
Everett, K. Xiao, L., Wortsman, M., Alemi, A. A., Novak, R., Liu, P. J., Gur, I., Sohl-Dickstein, J., Kaelbling, I. P., Lee, J., et al. Scaling exponents across parameterizations and optimizers. arXiv preprint arXiv:2407.05872, 2024.
Binder, K. Finite size scaling analysis of ising model block distribution functions. Zeitschrift für Physik Condensed Matter, 43:1 19-140, 1981.
Fisher, D. S. Asexual evolution waves: fluctuations and universality. Journal of Statistical Mechanics: Theory and Experiment, 2013(01):P010N) 2013.
Bordelon, B. and Pehlevan, C. Self-consistent dynamical field theory of kemel evolution in wide neural networks. Advances in Neural Information Processing Systems, 35: 32240-32256, 2022.
Glasserman, P. Monte Carlo methods in financial engineer- ing, volume 53. Springer, 2004.
10

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
Hallatschek, O. and Fisher, D. S. Acceleration of evolution- ary spread by long-range dispersal. Proceedings of the National Academy of Sciences, 111(4)):E4111-449) 2014.
Nakkiran, P., Neyshabur, B., and Sedghi, H. The deep boot- strap framework: Good online learners are good offline generalizers. arXiv preprint arXiv:2010.08127, 2020.
Noci, L.,, Meterez, A., Hofmann, T., and Orvieto, A. Super consistency of neural network landscapes and learning rate transfer. Advances in Neural Information Processing Systems, 37:10269--10274. 2024.
Hendrycks, D. and Gimpel, K. Gaussian Error Linear Units (GELUs). Preprint arXiv 1606.08415, 2016.
Hestness, J., Narang, S., Ardalani, N., Diamos, G., Jun, H., Kianinejad, H., Patwary, M. M. A., Yang, Y., and Zhou. Y. Deep learning scaling is predictable, empirically. arXiv preprint ar AX12:0712000400 2017.
Paquette, C., Lee, K., Pedregosa, F., and Paquette, E. Sgd in the large: Average-case analysis, asymptotics, and stepsize criticality. In Conference on Learning Theory, pp. 3548-366 3626. PMLR, 2021.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L.., Hendricks, I., A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022.
Paquette, C., Paquette, E.,, Adlam, B., and Pennington, J. Ho- mogenization of sgd in high-dimensions: Exact dynamics and generalization properties. Mathematical Program- ming, pp. 1-90, 2024a.
Paquette, E., Paquette, C., Xiao, L.., and Pennington, J. 4+ 3 phases of compute-optimal neural scaling laws. arXiv preprint arXiv:2405. 15074, 2024b.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836/ 2020.
Pearce, M. T., Agarwala, A., and Fisher, D. S. Stabi-
lization of extensive fine-scale diversity by ecologically driven spatiotemporal chaos. Proceedings of the National Academy of Sciences, 117(25):14522-1453, 2020.
Levy, S. F., Blundell, J. R., Venkataram, S., Petrov, D. A., Fisher, D. S., and Sherlock, G. Quantitative evolutionary dynamics using high-resolution lineage tracking. Nature, 519(7542):181-188 2015.
Pearce, and Song. Reconciling kaplan and chinchilla scaling laws. arXiv preprint arXiv:2406. /2907, 2024.
Li, Q., Tai, C., et al. Stochastic modified equations and adap- tive stochastic gradient algorithms. In International Con- ference on Machine Learning, pp. 2101-2110. PMLR, 2017.
Rosen, M. J., Davison, M., Bhaya, D., and Fisher, D. S.
Fine-scale diversity and extensive recombination in a quasisexual bacterial population occupying a broad niche. 2015.
Science, 348(6238)
Luo, K., Wen, H., Hu, S., Sun, Z., Liu, Z., Sun, M., Lyu, K., and Chen, W. multi-power law for loss curve prediction across learning rate schedules. arXiv preprint arXiv:2503. 12811, 2025.
Schaipp, F., Hägele, A., Taylor, A.. Simsekli, U., and Bach, F. The surprising agreement between convex optimiza- tion theory and learning- scheduling for large model training. arXiv preprint arXiv:2501. 18965, 2025.
Malladi, S., Lyu, K., Panigrahi, A., and Arora, S. On the sdes and scaling rules for adaptive gradient algorithms. Advances in Neural Information Processing Systems, 35: 7697-7711, 2022.
Ser-Giacomi, E., Zinger, L., Malviya, S., De Vargas, C., Karsenti, E., Bowler, C., and De Monte, S. Ubiquitous abundance distribution of non-dominant plankton across the global ocean. Nature ecology & evolution, 2(8):1243- 1249, 2018
Marchenko, V. and Pastur, L. A. Distribution of eigenvalues for some sets of random matrices. Mat. SbS) 72(1 14): 4, 1967.
Sharma, U. and Kaplan, J. Scaling laws from the data mani- fold dimension. Journal of Machine Learning Research, 23(9):1-34, 2022.
McCandlish, S., Kaplan, J., Amodei, D., and Team, O. D. An empirical model of large-batch training. arXiv preprint arXiv:18/2.06166 2018.
Tamai, K., Okubo, T., Duy, T. T., Natori, N., and Todo, S. Universal scaling laws of absorbing phase transi- tions in artificial deep neural networks. arXiv preprint wrXiv:230.00224. 2023.
McLeish, S., Kirchenbauer, J., Miller, D. Y., Singh, S., Bhatele, A., Goldblum, M., Panda, A., and Goldstein, T. Gemstones: A model suite for multi-faceted scaling laws. arXiv preprint orXiv:2502.0m 06857, 2025.
Tissue, H., Wang, V., and Wang, L. Scaling law with learn- ing rate annealing. arXiv preprint arXiv:2408.11029, 2024.

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
Venkataram, S., Dunn, B., Li, Y., Agarwala, A., Chang, J., Ebel, E. R., Geiler-Sameroott, K., Hérissant, I.., Blundell, J. R., Levy, S. F,, et al. Development of comprehensive senotype-to-fitnes. map of adaptation-driving mutations in yeast. Cell, 166(6):1585-))
2016.
Vyas, N., Atanasov, A., Bordelon, B., Morwani, D., Sainathan, S., and Pehlevan, C. Feature-Icamingg net- works are consistent across widths at realistic scales. Ad- vances in Neural Information Processing Systems, 36: 1036-1060, 2023.
Wang, X. and Aitchison, L. How set adamw's weight decay as you scale model and dataset size. arXiv preprint arXiv:2405. 13698, 2024.
Wen, K., Li, Z., Wang, Hall, D., Liang, P., and Ma, T. Understanding warmup warmup-stable-der stable-decay learning rates: A river valley loss landscape perspective. arXiv preprint arXiv:2410.05192, 2024.
Wilson, K. G. Renormalization group and critical phenom- ena. i. renormalization group and the kadanoff scaling picture. Physical review B, 4(9)31174, 1971.
Wortsman, M., Liu, J., Xiao, L., Everett, K., Alemi, A., Adlam, B., Co-Reyes, D., Gur, I., Kumar, A., Novak, R., al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprini 2.0i::2309.13222 2023.
Xiao, Rethinking conventional wisdom in machine learning: From generalization scaling. arXiv preprint 2rXiv:2409.15156 2024.
Yang. G. and Hu, Feature Learning in Infinite- Width Neural Networks. International Conference Machine Learning (ICML), 2021.
Yang, G. and Littwin, E. Tensor Programs IVb: Adaptive Optimization in the Infinite-Width Limit. International Conference on Learning Representations (ICLR), 2023.
Yang. G., Hu, E. J., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, Tensor Programs V: Tuning Large Neural Networks via Zero- Shot Hyperparameter Transfer. Advances in Neural Infor- mation Processing Systems (NeurIPS), 2021.
Zhang, B. and Sennrich, R. Root mean square layer nor- malization. Advances in Neural Information Processing Systems, 32, 2019.
Zhang. G., Li, L., Nado, Z., Martens, J., Sachdeva, S., Dahl, G., Shallue, C., and Grosse, R. B. Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model. Advances in neural information processing systems, 32, 2019.
12

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
A. Experiment Details
Transformer Architecture. We use GeLU activations (Hendrycks & Gimpel, 2016), RMSNorm (Zhang & Sennrich, 2019), and learned positional embeddings. We untie the embedding matrix from the output head and do not use bias anywhere. The readout layer is always zero-initialized, as suggested by Yang et al. (2021). We denote the embedding dimension with D. We set the intermediate dimension in the feedforward layers D instead of the usual 4D, which enables to explore larger widths more efficiently. The head dimension is set to 64.
CIFAR-5M. We use the CIFAR-5M dataset (Nakkiran et al., 2020) of 6 million CIFAR-like images. We convert the 32 x 32 x 3 images to greyscale and flatten them into sequences of length 1024. The model autoregressively predicts the pixel intensities in raster-scan order. The vocabulary is the set of pixel intensities {o.... 255). Following jaP we parameterize the learning rate for each weight matrix as 7) = Thasee D where d is the model dimension, except for the embedding matrix which has 1) - Thase- We use a parameter multiplier on the embedding matrix. We use ?/sas 4 and 0.1 as they led to good performance in our early experiments. We initialize the embedding matrix Wijhh ~ (O 1), the output head as Whead 0, all other non-readout matrices W as Way (O 1/D). These hyperparameters were determined with small amount of tuning in early experiments. We use batch size of 256 images. We use linear warmup for 1000 steps.
For the experiments in Section 3.2.2, we used a slightly different setup due to switching to a new codebase in the middle of the research project. We use μP where the base embedding dimension is 128 and base learning rate is 0.01. We initialize the embedding matrix with a standard deviation (std) of 0.1 and multiply its learning rate by 10 relative to the base learning rate. The output projection of the feedforward and attention layers are zero-initialized. All other non-readout matrices are initialized with std 1/VD. We use batch size of 65536 tokens or 64 images. We use linear warmup for 10M tokens.
Chess. We run our experiments on the Lichess dataset available on Hugging Face https: / / huggiigg face. co/ games. We used character-level tokenization and context length of 128.
datasets / Li chess //s/staanne sstandardddchee We use juP where the base embedding dimension is 128 and base learning rate is 0.01. We initialize the embedding matrix with std 0.1 and multiply its learning rate by 10 relative to the base learning rate. The output projection of the feedforward and attention layers are zero-initialized. All other matrices are initialized with std 1/VD. We use a batch size of 65536 tokens. We use linear warmup for 10M tokens.
MLP Experiments. Our MLP architecture is identical to the transformer with attention layers removed and the token and position embedding layers replaced by linear layer. We use μP where the base width is 128 and base learning rate is 0.001. The output projection of the feedforward layers are zero-initialized. All non-readout matrices are initialized with std 1/vD. We use a batch size of 4096 examples. We do not use warmup.
The target function is defiend as x(z) = ΣM w.v22 cos (2rk x + b,). with x E R w; ,, ~ (O 1),b;
A [Orrnoull(0.5), k = round(s,u,) where s; is a scalar sampled from a power law with support [1, 2o) and exponent -2, n/ is a random unit vector, and round rounds to the nearest point in Z During training, x is sampled uniformly from 1-0.5, 0.5]s, making the Fourier features orthonormal over the data distribution. We suspect the details here are not necessary for generating power-law scaling laws beyond the power-law spectrum.
For Figure 4a, the constant learning rate scaling ladder matches the jP learning rate at D = 384 (smallest model).
Mcasuring Gradient Covariance Trace As mentioned in Section 3,2.1, we use the preconditioned gradient covariance E instead of the raw gradient covariance due to the use of Adam (and μPP). is defined as P p-1/221p-117 where Σ is the raw (mini-batch) gradient covariance, and P is the preconditioner. See an explanation for this definition in Appendix G. Since jPP uses layerwise learning rate, the preconditioner is defined as P-1 diag
1/ vez+r) where no is the vector of and is the
per-parameter learning rate (peak learning rate, before applying the schedule), x2 is the Adam Adam e. As the peak learning rate is absorbed into the preconditioner, any occurrence of the instantaneous learning rate (!() or (T(T) in Section 3.2.2 reflects only the schedule and takes on values in [0. 1]. This definition mirrors what is done in Noci et al. (2024); Cohen et al. (2024).
Training for More than One Epoch. The CIFAR-5M dataset, tokenized as 32 x 32 greyscale images, has about 5B tokens. Therefore, most models needed to be trained for more than 1 epoch to reach the compute-optimal training horizon. Up to the scales we tested, we did not observe a significant difference between the train and test loss. The chess dataset has about 20B tokens, which also led to data reuse for some models, but did not lead to significant overfitting. As we only processed subset of the full Lichess dataset, this can be avoided by processing a larger subset if desired.
13

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimall Trained Neural Networks
On Random Seeds. In the CIFAR-5M experiments in Figure 1, the random seed controls both the initialization and data ordering. In the other experiments (Transformer on chess and MLP regression), the random seed only controls the initialization while data ordering is held fixed, is often done in practice. Fixing the data ordering (no shuffling) had the advantage of speeding up data loading. We found that supercollapse occurs regardless of whether seeds affect data ordering, This makes sense: even with fixed ordering, different model sizes process different data due to varying training horizons. More fundamentally, supercollapse should be robust to which training components are randomized, as the variance reduction arises from strong noise correlations along individual trajectories rather than specific noise sources (Section 3.3).
B. Scaling Collapse Across Transformer Depths
For scaling depth, we additionally apply branch multiplier of 3/depth on the output of every feedforward and attention layer, as suggested by Bordelon et al. (2024c). We find a decent degree of collapse in Figure 9 when training on chess data. There is a small shift in the normalized curves, though we are unsure if it is simply finite-size effect.
ac
1.5
Depth
0.72
= 0.355
Normalized Loss
0.164
0.545
1.4
16
0.69
20
24
0.56
1.2
28
36
0.63
1.0
10a
10a
10a
0.0
0.2
0.4
0.6
0.8
1.0
Compute (PetaFLOPs)
Normalized Compute
(a) Scaling Law
(b) Collapse
Figure 9: Depthwise scaling collapse for transformers trained on chess.
C. Estimating Compute-Optimal Training Horizon
To estimate the optimal compute for training each model, we perform the following steps in each experiment:
We trained each model withou learning rate decay but keeping the initial warmup. We chose a large enough number of steps that the largest model could reach the compute-optimal loss frontier. We average the loss curves from 5 seeds.
We numerically computed the compute-loss Pareto frontier to obtain an estimate of c²) the optimal compute for each model size p. We use logarithmically spaced points for and find the p that achieves the best loss given e training FLOPs.
We fit a power law (p)) = kp!+> where x and y are fit parameters. The optimal number of training tokens is then
1' (p) = e* (p)/(6p), which scales as p". We remove outliers in this by dropping points from the smallest and largest model.
Using a constant learning rate schedule allows us to measure the loss of one model at different token budgets with a single training run, an approach also used in McLeish et al. (2025), rather than one training run per token budget as origianlly done in Hoffmann et al. (2022). Figure 10 illustrates this procedure.
14

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
3.27
Compute [PensFlops] Flops)
134925161 a 1.04
533620.12 12
10?
100
1227
32n1
13n11
12
122
10'
32'
Compute (PetaFlops)
Farameters
Compute (PetaFlops)
Parameters
(a) MLP Pareto Frontier
(b) MLP Data Exponent
(c) CIFAR-5M Pareto Frontier (d) CIFAR-5M Data Exponent
Figure 10: Estimating compute-optimal data exponent the MLP regression and Transformer CIFAR-5M experiments.
D. Universality and Scaling Collapse in Other Sciences
The simplest versions of collapse come from statisties and probability, where entire distributions of random variables show universal behavior between systems of different types and scales. The most well known is the central limit theorem which predicts universal Gaussian form for the sums of random variables with appropriately bounded moments (and Levy distributions for heavy tailed distributions). In random matrix theory, there are similar effects when studying the limiting empirical distribution of spectra, most famously the Marcenko- Pasttr distribution of Wishart matrices (Marchenko & Pastur, 1967). In all of these examples, showing universal behavior of single moment analogous to showing predictability the Pareto frontier in our work, while showing the universality the whole distribution is analogous to our statements about the entire loss curves (where e.g. the CDFs of different problems are converging).
Scaling collapse is ubiquitous in physics well. There are again distributional collapses like the famed Maxwell Boltzmann distribution first used to describe idealized gasses, but also functional relationships like the universal magnetization- temperature curves of near-critical Ising lattices of different sizes (Binder, 1981). In the Ising example, changes to the lattice topology can change the algorithms lead to different universal curves in our study. more general theory unifying and explaining the existence of universality physics arises from the renormalization group (Wilson, 1971).
peneriiaation-emperatttte curves, similar to how different datasets, architecture, and training
More recently, scaling collapse has been used to describe dynamical systems in biological contexts. Advances in genomies have led the rise of experimental microbial evolution with rapid timecourse data (Levy et al., 2015; Venkataram et al., 2016). Analysis of this data relies on quantitative modeling of evolutionary dynamics. Most of these models show universal scaling dynamics, in situations from rapid evolutions of diverse populations (Fisher, 2013), populations evolving under changing fitness conditions (Agarwala & Fisher, 2019), and populations expanding in space (Hallatschek & Fisher, 2014). In these settings the timecourse of key observables can be described with dynamical curves which can be rescaled to universal forms with transformations depending on population size, mutation frequency, and statistics of the fitness landscape. In ecology, the ubiquity of fine-scale diversity and the seemingly universal, power-law nature of species rank- aauaaabundance curves (Rosen al., 2015; Ser-Giacomi et al., 2018) can be explained using dynamical models which themselves show universal scaling behavior over ecosystems of different sizes (Pearce et al., 2020).
Scaling collapse has been used to study certain scaling relations in machine learning. Kaplan et al. (2020) identified universal scaling of overfitting, showing a collapse of the rescaled excess loss vs rescaled parameter count across dataset sizes. Tamai et al. (2023) used scaling collapse to establish universal scaling laws the forward signal propagation dynamics MLPs near the grder-to-chaos transition.
E. Power-Law Pareto Frontier is Necessary for Collapse
Recall (p) is the optimal training horizon for model size i.e. (((*(p)))) = mine 3/:0///((ppp I(L,p'). Let c*()) (p)p be the optimal compute for p. In what follows, rather than writing Ltt,p), we will find it convenient to express the loss curves in terms compute and model size. Letting L(c,p) p) be the loss curves expressed this way, we have the following theorem:
Theorem E.1. Let L(c, p) be C2 for all positive e, p and let e² (p) denote the compute-optimal budget for model size p.
15

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimaly Trained Neural Networks
Write [(c,p) [c,,)) - L for some offset L (e.g, L Lo the irreducible loss). Define the normalized loss curve
((x,p) = [(xc IE [0.
(25)
(p)
Then, 1. Necessity. If is independent of p (collapse), then the Pareto frontier of (f(e,
L (c) min [(e,p) L(c (pc) a power law (c(c)=ac" for some constants a,
(26)
P
Sufficiency at first order. Conversely, suppose L (c) = ac- .thn
d((x, p) -
(27)
dx
independent of Hence all curves share the same first-order behavior around x 1, i.e., they collapse to first order around
Proof. First, have the following identity for derivviivasives for a general differentiable function u(e)
dlog u v du d dog du
(28)
Applying this to our normalized curve from Equation (25) and using the chain rule:
dl(x,p)
c² dC(c,p)
(29)
dr
C(et,) dc
l==1
l==o
i log L(c,p)
(30)
Öloge c
where the second equality follows from Equation (28).
Necessity. If & is independent of (collapse), then dd(x.p) Equation (30),
is the same for all p. Set this common value to -6. By
dæ 1=1
i log L(c,p)
for every p.
(31)
bloge c
(p)
Since (c* (p). lies on the Pareto frontier and L. is C*,, the envelope theorem states
(c)
aC(c. p)
(32)
de
dc
leme((p)
la=e" (p)
ie. the loss curve is tangent to the Pareto frontier at the compute-optimal point. Applying Equation (28) to the frontier
C*(c):
dlog (c)
ct(p) dC (+(et(p)) de (emet())
(33)
dlog
(ame+(p)
c²(p) aL(c,p)
(34)
(pc(()).p)
dc
(me (p)
bog L(c,p)
(35)
aloge 0,
lose (p)
(36)
16

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
where used Equation (32) and Equation (31). This means the log-log slope of the frontier is constant, which means is power law L (c) ac
Sufficiency at first order. Assume (c) ac For any model size the envelope theorem gives the tangency condition at c*(()) (p):
@L(e, p)
(c) dc
-da(c")- -8-1
(37)
dc
lome
come
Applying Equation (30):
dl (2. p)
i log L(c,p)
(38)
dx
d logg e
m=1
lo=e (p)
8.
(39)
Therefore, all curves collapse first order around in Equation (27).
Remark. A power-law Pareto frontier is not only necessary for full collapse but also already enforces weaker, first-order form of collapse. Theorem E,1 assumes the compute- optimal point lies in the interior of each loss curve. This condition can fail for learning rate schedules that reach boundary of the curve, where the envelope theorem tangency no longer applies. Such schedules are used throughout our experiments and are common practice. Extension the proof to handle these boundary-optimal schedules would be interesting.
after finitely many steps, because the optimum may then coincide with the
F. Collapse for General Sum-of-Power-Laws Loss Curves
Theorem F.1. Suppose the loss curve is given by
m
L(t,p) Lo + at Mipti
Lat
a, > 0, p4, 20, 0,
(40)
with least one of pi, Vs positive for every i (else absorb the term into Lo). Let 1.* (p) kp with K > 0, 0 be the asymptotic compute-optimal training horizon, and define the total exponent Be: == P4N + U; and by a,K 14, Without loss of generality, assume 3, 's are sorted in non-decreasing order. Then,
1. Compute-optimality forces a tie. At least two Br share the minimum:
B1 = B2 = ... Bx < But1 S... S Bmn k 2.
(41)
2. Asymptotic collapse. The normalized loss curve
(xxp)) L(st" (p).p) Lo [((x()),P) Lo
(42)
is given by
Eis bir Hi
((x, p)
+ +0(p-),
c: Bx41 - 31 >0,
(43)
Eim1b
independent of p up to finite-size error that decays as O(p If m, is taken to be do (perfect finite-size collapse).
3. Locally fastest decay of finite-size error. Locally, " is the data exponent that achieves fastest decay of the finite-size error as measured by In particular, O()1) for any other data exponent o/ = ō with 8 + O, leading to more slowly decaying finite-size error and therefore a worse collapse.
Compute-optinett up to a constant suffices. Any training horizon that is a constant multiple of preserves the collapse, only changing the constants by in Equation (43).
17

---

Scaling Collapse Reveals Universal Dynamics in Computt-Optimally Trained Neural Networks
Proof. op forces tie. Fix the compute budget == 6lp and note l(p) c/(6p) so that de =-1pp
dp
With 8 == poy + Vi and by- := asKTMi
dL dp
dt à
a1
t Piip"
(44)
8p
dp d1
14 p
H
(45)
p
as(11 - vi)t "1pp""pp" p
(46)
m be(p) v1)pp= 9.
(47)
p
If [, < $2, the leading term bi (u1 v v1)p- 8: cannot cancel the rest for asymtotically large y, contradicting di by compute-optimallty. Hence at least two indices share the minimum exponent, yielding Equation (41).
0 required
dp
Asymptotic collapse. We compute f(z. p) explicitly. First, evaluate the loss the optimal horizon:
[(t" - Lo a at (p)) Uip
as (
(48)
a₂K 1iip-Mi-- 117-1i
6.p Bi
(49)
2=1
Since 31 = 32== =....--= =Bx < Bk+1 S S Bn, we can factor out p
-1;
m
[(L()))))) Lo = p 81
Σ +
Σ
bip-18: (8, -Ba)
bip
(50)
im
= p -1
Σ
(1+ -(844 (1+00=1(BAA1)111
(51)
Similarly, for t = zt" (p):
L(xt (p),p) Lu Σ
(xtx (p)) -Mip-
Σ
a,T 1i (t" (Mi)
(52)
i=1
m
Σ b Mi ) (1 +
b
Mi p
Br = =p -0
(53)
Taking the ratio gives:
p 91 E/_ bi 14 (1+ (1611-81)) (1+101-(0)111)))
Σi bar Hi
+ (1 (0)
(x.p)
(54)
p 91 Ex=1 b.) (1 + Opp (01+)-91)) (1+00-(0+-)1i1)
Σ Ei-1b bi
This produces Equation (43) with e = Bk-1 31 > 0.
Locally fastest decay of finite-size error. Let 7 be the optimal data exponent and perturb it by ô, writing 7' = + 8. For a small enough |8] > 0, the previously tied lowest exponents , .& split into distinct values Bi.... 8y which remain
Herr we assumed not all μ4- * 0 for i =,.... ,, kt else these terms would not affect dL./dp. If this is not true, then the loss T.(f. p) not interesting because the would asymptotically he function of compute 6tp alone, independent of how we allocate between and p.
18

---

dr dul 10-((/ Scaling Collapse Reveals Universal Dynamics in Compute-Optimall Trained Neural Networks
the lowest & exponents (since 5 is small), and the gap between the lowest and second-lowest grows as O(1))), which is strictly smaller than the previous gap < 3k-1 31 > 0 for sufficiently small 6. Therefore, locally ? maximizes the decay exponent of the finite-size error, i.e., it gives the best collapse locally.
Compute-optimality up to a constant suffices. Replacing t" by At" multiplies each bj by A- and leaves the rest of the proof unchanged.
Remarks.
1. Compute-optimal data exponent implies asymptotic collapse, but the converse is not necessarily true when m 2 3, since there can be multiple choices of y that lead to balanced dominant power laws, which imply collapse, but only one of them can be compute-optimal.
2. In general, when > 2, asymptotic instead of exact collapse is the best we can hope for. But asymptotic collapse alone is not that interesting, since under any choice of the data exponent yy only the terms with the lowest 3; enter the asymptotic normalized loss curve. For example, if L(t, p) = t" +p" any >> 2/μ will cause only the p- term to dominate, leading to ((x,)) p) + 1, whereas any 7 < 1/μ4 will cause t~ to dominate, leading to f(x, p) + x- The latter case is similar to the anfinitewwitth limit in neural networks, where under / (1), the loss curves become bottlenecked by training time alone and not model size (Vyas et al., 2023; Bordelon Pehlevan, 2022; Yang & Littwin, 2023). What is interesting about the collapse that happens under compute-optimal training that y is tuned (to // in this example) that more than one such term exists, so the collapse reflects a balanced scaling of both training time and model size. This fine balance is also why minor perturbations to from the optimal value can significantly disrupt the collapse, which is not true if there is only one dominant power law.
0.800
54655 6.0
pria p t.0
(998] ScomptteQppiinn
0.900) Compute Compute-Cpiina
0.775
0.750
0.725
0.04
0.700
0.675
0.02
10a
10a
10
10
10
10
10a
10a
Compute (PetaFLOPs)
Compute (PetaFLOPs)
(a) Transformer on Chess
(b) MLP
Figure II: cum-of-power-lawe fit on additional datasets. Both tasks have loss curves that can be approximated by the sum of two power laws when using constant learning rate schedule. Fitted constant multipliers are not shown in the legend. To not fit to early-time transients, we fit all steps after the first O. tokens for MLPs, and all steps after the first 0.1x the compute-optimal horizon for chess. The fit for chess is worse than the other datascts.
G. A Perturbative Model of Learning Rate Schedules
Let w' denote the parameter trajectory under the influence of gradient noise. The dynamics of stochastic gradient descent in gradient flow time are given by
(55)
with noise correlation (E)(())(+))T] = n1(7 7(T))(T 7'). For convenience, we rewrite this using E(7) ()))(())((T) so that
du' = - -VVVV(u') + (22)
(56)
dr
where now ([((()((T)T) = 5(T ((T--)). 7)).
19

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
Our strategy is to solve w' as w(r) + ów(() where w(T) is the deterministic trajectory satisfying du leading order in the gradient noise scale r)>
-VI((). up
Letting dw = w' g, = V and taking the difference of the two differential equations:
d(ów) (s(w') - g(w) + (1d))))
(57)
dr
At first order,
g(wo) g(w) + H(w)ow
(58)
where ((w) = L( is the Hessian.
Our SDE for dw becomes:
d(ów) H(w)ōw (1)))
(59)
We define the propagator G(T, 8) that satisfies:
dG(T, -((W(T))((T,s)
(60)
dr
with Gs) I.
For time-dependent (w(()), the propagator is:
= TT (-f, ddd((w(A)))
(61)
where T denotes time-ordering.
Assuming the initial perturbation (0) = 0, the solution for dw is:
8w(() (u(())=-( - Sf ds G(T, ()))))
(62)
Now expanding L(w') L(w 5w) to second order in dw gives
SL(T) (())(7)) [(w(TT)
Sw(()] H (wo()) 8w((). = ((w(T)) (w((T) + 2
(63)
Since E(ī(s)) = 0 and dw is linear in द, [[] = 0, so
E (1(w(())T66(7)) ow(T)] = 0.
(64)
Thus the leading non- non-vanishing contribution expected loss shift comes from the quadratic term.
Using the solution for ow,
(w()))))(TTT)))))))))))))))))))) = = (-- ds [ du G(r, "22
(65)
Taking the expectation with E((())((u)T] 5(s u) gives
E L ds G(T,)) ((s))((w(ss) s)
E
(66)
20

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
Substituting this into the quadratic term gives
// ds Tr (H
B[) = E
(67)
2
Using Tr(ABC] = Tr(CAB),
1 B[ő - [ ds Tr (
(68)
This equation is the exact leading-order expression for the noise-induced change expected loss, which the equivalent of (T) in Equation (10). The deterministic loss (ww(r) now plays the role of F(r) in Equation (14). Conceptually, the derivation shows that although the full dynamics non-linear and the loss is not assumed quadratic- the perturbation generated by small gradient noises behaves in simple, linear-quadrate fashion: the weight perturbation Ow is linear in the injected noise, and the resulting loss shift 5T. is quadratic in that perturbation.
Consequently the derivation and final formula completely mirrors the familiar quadratic-loss result, the only difference being that the constant Hessian is now replaced by the dimendepnddent Hessian H (w(()) carried along the deterministic trajectory. In other words, small gradient noise "sees" the network through an instantaneous linearization, so all schedule effects enter through the propagator G),,
s). the local Hessian, and the noise covariance, exactly in the linear case.
Relation to Stochastic Asymptotic Expansion. Li et al. (2017) used stochastic asymptotic expansion expand the dynamics in orders of ry1/2, but they treat the gradient-noise covariance differently. The stochastic asymptotic expansion in Li et al. (2017) first expands the diffusion term along the deterministic path, so E can be propagated analytically order by order. Our derivation instead keeps the exact E(w' (s)) inside the leading-order integral, allowing an empirically measured covariance to be inserted without further approximation. at the cost of analytic closure.
Varvatiittion and Late-Time Limit. As in the quadratic loss case, we can simplify the result under an adiabatic approximation where the Hessian, schedule, and noise covariance changes slowly compared the time-scale set by the instantancous Hessian. Specifically, if H (w(t)) x (w(()) == H over the support of G(r,s) (w((T) G(T, s), then G(T,s) 2 e- --((r=) and G(T, s) HG(r, s) z H e 222((7-)), Assuming the exponential decay is fast compared to the variation of the noise scale n,,, and taking + oo, we have
1 4 Tr (+))
B[oL(T)]
(69)
which agrees with the expression for {T) in Equation (14) for the quadratic loss setting.
Adaptive Optimizers. When using adaptive optimizers with a preconditioner P(t), the SDE becomes dw'
dt (V<0))=-(((((+++++""><((<>>"
(70)
The gradient flow time SDE in Equation (56) becomes
dw' -
(71)
dr
where abbreviated [ww") Σ' If the preconditioner varies slowly, the dynamics be treated if there is no preconditioner, but in transformed coordinate system:
= (2/2(()))
(72)
Differentiating and neglecting the O(P) term gives
dw'
p1/2 dw'
= 1P21/2 (V((w)) + +1 1/122211222)
(73)
dr
dr
=o(""") (1P-1
(74)
21

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
The deterministic trajectory in this coordinate system (1²(1) (1/2(() w(()) satisfies
du
-(.).
(75)
dr
To dervie the SDE for dw == w - w. at first order, we have
5( - g ④ = 1/2 How = How p-12==--/ppp-122000 = Hdw,
(76)
where H is the preconditioned Hessian. Therefore,
dôw H88 (1P-1
(77)
dr
is also easy to show that the preconditioned Hessian indeed governs the leading order perturhation to the expected loss in the transformed coordinates
1 =Bl6 HōØ)
E[öL]
(78)
Given Equation (77) and Equation (78), all steps in the previous derivation now applies after swapping H - H and E' + p-1/22'p-1/22 = E, The final result for the noise-induced loss perturbation in the slow-variation and late-time limit is
1 Tr [n(T) 4
(79)
Limitations. It worth highlighting some limitations of our analysis. First, due to the non-linear nature of neural networks, it is known that gradient flow can fail to model full-batch gradient descent with a finite step size, which exhibits effects such as the Edge of Stability (Cohen et al., 2021; 2022; 2024). Similarly, in the stochastic case, there a strong coupling between learning rate, batch size, and the Hessian spectrum (Agarwala Pennington, 2024). These phenomena suggest that is unlikely that we can fully model the effect of a time-varying learning rate schedule simply as injecting a schedule-dependent noise component on top of the deterministic trajectory that itself independent of the schedule, as changing the learning rate also pushes the model into different regions of the parameter space based on curvature (though it is possible that the leading-order perturbation theory can already capture this effect to some extent).
Second, when dealing with adaptive optimizers, we assumed the preconditioner stays the same (as a function of 7) between the deterministic and stochastic trajectories. For typical optimizers, such as Adam (Diederik P. Kingma, 2015), this assumption not correct as P depends on the gradient covariance, which can differ between the two trajectories, i.e., E(w(T)) # [(ww())) This can introduce an additional term in the SDE for ow, present even first order, which we did model.
The empirical effectiveness of our model for predicting the loss curves under different schedules suggests it is nevertheless on the right track, and that there may be other ways to derive similar or improved predictions with more accurate assumptions.
H. Computing Á to Leading Order
Let ((r) ((r))_((r) and define E(T)
A? - +(71))
(80)
with -- = 85 and 8r < 1. Because -( [(7)-1 0(77),
(²(7==(-9) E (AL(T)
AC( (T((())(((C(TT)) - (
(81)
where AC(T) [(T) C(r). For the quadratic model.
AC(T) = ( Aw(T)
(82)
22

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
first order in
Aw(T)= /'' ds (sc-
(83)
with
B[{(s)[(s')] = (7))0(ss =(5))s -
(84)
- Splitting the upper limit r and expanding 9(7")= + 0(7) gives
AC(T) - - [- dse H(T (20)) +0(67))
(85)
O(671/2)
where the remainder collects two subleading (δr) contributions:
R g(T) l/' ds (c-H(('-) -() - -e-H(r-s) Σ (s)
(86)
R2 = (a([T) - g(T)l Aw(T) = ](T)T ( 6r 0(679).
(87)
Using e (((--s) = I m()) ((r) + 0(T)), and E(s) = [(T + 0(8T) inside the leading-order integral.
E - ((AA(())A((r))))) 9(T)' [. ds ds)(s) 2H( = E(s)((T) + 0(672) [(()] (T))) +
(88)
(89)
Substituting this into the expression for (7) yields the desired result
= ()
(90)
I. Additional Results on Learning Rate Schedules
MIP fits. Figure 12 shows predictions for MLP loss curves. With single 0.26 (very close to 1/4), obtain excellent fits across schedules, model sizes, and training horizons.
p 10M. T 1B
(3M, 10M, 42M). T 1B
10M, T (0.5B, 1B, 2B)
0.040
BaD caattrx)) (25%)
0.048
031 +-me/ex))
Obeered
(199)
0.032
(72%) 0.511 con(3ex)) (00)
0.035
0.11 uo(3.xt)
0.040
355 Bmx]] (20%)
Loss
Loss
0.028
0.080
0.052
0.024
0.025
0.0
0.2
0.5 Tokens (Billion)
0.4
0.8
1.0
0.0
0.2
0.4
0.6 (Billion)
0.8
1.0
0.0
0.5
1.0
1.5
2.0
Tokens
Tokens (Billion)
(a) Vary Schedules
(b) Scale Model Size
(c) Scale Training Horizon
Figure A simple model predicts MLP loss curves trained across learning rate schedules, model sizes p. and training horizons on the synthetictic regression task. Dashed curves show the predicted loss as L' ('( = L(T) + a (r((r) Tr(E' (T)) (Equation (18)). a is the only free parameter and is set to 0.26. Each curve is smoothed with an exponential moving average with half life equal to 1% of total steps.
Effect of Dropping 77 Tr(8E). Figure 13 (top row) shows 6n Tr(E') is typically 3 to 10 times than 17 Tr(3E) in absolute value both CIFAR-5M transformer and MLP regression experiments. Since we decay the learning rate to zero (5q
23

---

Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks
comparable 11), this means the gradient covariance does not change much between the constant learning rate and other schedules, which can happen if the Hessian trace does not change much. The fact that this ratio is only moderately large shows that the gradient covariance or the Hessian trace did change considerably due to decaying the learning rate, which is what we expect due to generically inverse relation between learning rate and Hessian eigenvalues (Cohen et al., 2021; Agarwala & Pennington, 2024). However, somewhat puzzlingly, we found that including the smaller term vj Tr(6E) can produce worse fits, particularly for the slow oscillatory schedule (0.5 cos(2))), and makes the optimal constant a more schedule-dependent.
P 20M. T- 10B
P 10M. T 1B
L-X
2-
10
0.50 ontea]
10
2,501 cotteg))
Rabio
Rabio
(20Ng
0,0 ca/ (20%)
2,00 cor (205)
100
10
10-1
0.0
0.4
0.6
1.0
Tokens (Billion)
Tokens (Billion)
(a) (6rr Tr(2")) (Txo2))) CIFAR-5M
(b) [(ôn (1((2)) MLP
20M. T 10B
10M, T 18
3.20
[o- 0.32)
0.040
(o 0.33)
2.51 a ookevil) (o 0:31)
0.01 a catfes]) to 0.34)
2.512 sactmol) - 1.20)
0.035
0S01 cadteal) to- 0.22)
--x (205)) (o 2.20)
1- (20%) 6.20)
03)
cattaal) 324)
3.18
0.030
0.025
10
0.0
0.2
0.6
0.8
1.0
Tokens (Billion)
Tokens (Billion)
(c) CIFAR-SM with both terms
(d) MLP fit with both terms
Figure 13: Out of the two terms that make up (r))() the term 17 Tr(EE) is typically 3 10 times smaller than dn Tr(E') (top row). Moreover, including it sometimes produces worse fit and make the optimal vary more schedules (bottom row). We determine for each schedule by matching the prediction with the observation at the end point.
1 X
0.5(1 + cos(33xx) (3xx))
Tr(Z)/C
1.0
1M 4M
17M
67M
Tokens (B)
Tokens (B)
/
1.0
0.5
1.0
Normalized Compute Normalized Compute
Figure 14: Universality of gradient noise in MLPs. Fixing a learning rate schedule, the ratio Tr(E) L approximately function of normalized compute alone, independent of model size. this regression task, the estimated irreducible loss is negligible so L.
24


Reviews:
[{"review_id": "Q0c2qs4I0U", "reviewer": "Reviewer_MwGt", "summary": "- When neural networks are trained under compute optimality, their loss curves across model widths collapses to a single universal curve under a simple affine rescaling\n- The authors call this \"supercollapse\" because deviations between curves are smaller than fluctuations from multiple training runs (where the noise arises from random initialization and data ordering)\n- Supercollapse occurs across various learning rate schedules (so long as they're decaying) and architectures, but requires accurate estimation of the compute-optimal data exponent and irreducible loss\n- The authors prove that at power-law Pareto frontier is necessary (and in some cases, sufficient) for supercollapse to emerge\n- Several other complementary results\n\n## Update After Rebuttal\n\nI think this is a great paper. I don't think it deserves to be an Oral (in my Claims and Evidence, I pinpointed two steps that I think this paper needed to do in order to be that outstanding) but I think the authors have done a wonderful job.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive feedback! We provide several additional results and clarifications here and will include them in the updated paper. Corresponding figures are in [the linked PDF](https://drive.google.com/file/d/1ZkobNTqh90nnUcunKqUT2Dyx3T4zAb5O/view).\n\n**Further evidence for the generality of supercollapse.**\nWe focused on width scaling as it's the most well-studied dimension with extensive literature [1,2,3,4] on optimal learning rate and initialization scaling, which made finding the proper hyperparameter scaling needed for supercollapse much easier. Inspired by your question, we conducted a depth-scaling experiment by adding residual connections to our MLP (creating a well-defined depth scaling limit) and using depth-µP [5] to scale learning rate and initialization. Figure 1 (left) in the linked PDF shows depth scaling also produces supercollapse, confirming it's not width-specific; this plot is consistent with our theory linking supercollapse to power law compute-loss pareto frontiers.\n\nIn addition, Figure 1 (middle, right) demonstrates supercollapse occurring in two additional data modalities: chess game string representations and chemical structures, further establishing its generality.\n\n**Realistic scaling recipes that don’t admit supercollapse.**\nWe expect supercollapse to break down when hyperparameters such as the learning rate are not properly scaled with the model. In Figure 3 of the linked PDF, we ablate µP while keeping learning rate constant for all widths, which makes the learning rate increasingly suboptimal as width scales [4]. Without µP the rescaled curves don’t collapse. \n\n\n**Practical relevance of supercollapse.**\nExploring supercollapse's practical applications is an exciting direction for future work! Our primary goal in this work is to establish supercollapse as a novel phenomenon across multiple data modalities and architectures and to show how studying loss curve universality enhances our understanding of scaling.\n\nThat said, we believe this work already provides compelling evidence of supercollapse's practical value. Figures 3(a) and 4 in the paper show how absent supercollapse reveals errors in estimated data exponent $\\gamma$. Similarly, MLP without µP (Figure 3, linked PDF) shows comparable performance to MLP with μP on the loss-compute plot except for the largest model, but the absence of supercollapse reveals scaling errors. Supercollapse is a powerful diagnostic tool, making errors more detectable than when evaluating final losses alone.\n\n**Methodology for fitting optimal training compute.**\nWe estimate optimal compute by training models for a fixed number of steps (7e5 for transformers, 1e6 for MLPs) and identifying the compute-optimal point. We find models with lowest loss at each compute value on a logarithmic grid to create the compute-loss Pareto frontier (Figure 2). Our discrete model sizes mean each model is optimal for a range of compute values, producing multiple y-values per x-value in Figure 7 (Appendix C.1). We fit a power law to these points using log-log linear regression to estimate optimal compute per model size, similar to Chinchilla’s Figure 1 [6].\nWe use constant schedule to efficiently sweep compute/token budgets, rather than running separate experiments for each budget as in Chinchilla. This reduces experimental costs, a strategy also used in [7] for LLM scaling law studies.\n\n**Additional clarifications**\n\n* $\\theta$ refers to the fraction of optimal compute.\n* Figure 4 in linked PDF shows collapse deviation std dev across 5 trials. Colors indicate width (now in the legend).\n* $\\Delta < \\sigma$ for an $O(1)$ interval isn't trivial as $\\Delta$ can grow quickly away from $\\theta=1$, as we showed for constant schedule in Figure 6 of the paper.\n* SGD noise scales as $\\Theta(\\eta/B)$ (see Eq 33, Appendix E). Annealing LR thus reduces noise similar to increasing batch size.\n* Intuition for L415-420: variance of noise difference between timepoints is smaller than individual variances when sufficiently correlated.\n* We expanded $P_0$ range for sensitivity analysis in Figure 5 (linked PDF).\n* We suppressed $s$ in Eq 2 for brevity.\n* Outliers in Figure 9(a) occur when $\\Delta\\eta$ approaches 0 in multi-cycle cosine schedule, making $\\hat{K}$ divergent and not well-defined. \n\n[1] Yang et al. \"Feature learning in infinite-width neural networks.\"\n\n[2] Yang et al. \"Tensor programs ivb: Adaptive optimization in the infinite-width limit.\"\n\n[3] Bordelon et al. \"Self-consistent dynamical field theory of kernel evolution in wide neural networks.\" \n\n[4] Everett et al. \"Scaling exponents across parameterizations and optimizers.\" \n\n[5] Yang et al. \"Feature learning in infinite-depth neural networks.\"\n\n[6] Hoffmann et al. \"Training compute-optimal large language models.\"\n\n[7] McLeish, Sean, et al. \"Gemstones: A Model Suite for Multi-Faceted Scaling Laws.\"", "claims_and_evidence": "- Overall, I think this is a solid and thorough paper.\n- In my review, I scrutinize some specific components. I welcome the authors to improve the paper based on my comments, or to convince me that my comments are mistaken. Given compelling evidence for either, I will increase my score.\n- I think this paper could be an awarded paper if its practical relevance was increased. One way would be to show that supercollapse holds for language model ladders pretrained at scale, but this is likely an expensive request that the authors may be unable to afford. Another way would be to quantitatively compare how a “supercollapse-based” scaling law predictor compares to standard scaling law predictors, ideally in a backtested manner; the authors hint at this (“suggesting that supercollapse itself can be used to improve scaling predictions”) but stop short of actually applying this claim practically. I’m also unsure of whether a supercollapse predictor is feasible since Figure 3 suggests that $\\gamma$ and $L_0$ are the most sensitive, and I would intuitively guess that these two parameters require large models to accurately estimate, but I could be wrong!\n- To me, the point of highest uncertainty in this paper is its generality. My doubt stems from the fact the authors only swept fixed-depth networks. I couldn’t find a reason, I don’t understand why the authors did this, and there is now a question in my mind of whether supercollapse is as general as the authors claim it is.", "methods_and_evaluation_criteria": "Yes. I add specific comments in later sections.\n\n**Edit: Based on the authors' response, I'm upgrading my score to a 5**\n\n> $\\theta$ refers to the fraction of optimal compute.\n\nTo clarify, I was suggesting that $\\theta$ should be defined in the main text, likely in Section 2.2. It's possible I'm missing its definition in the main text, but I couldn't find one (and while $\\theta$'s meaning can be inferred contextually, I think one should prefer being explicit).", "theoretical_claims": "- Theorem 3.1 is a nice result.\n- Section 3.1 Sum of Power Laws Collapse Exactly: To clarify my understanding, Equations 7-9 (inclusive) hold only for the PLRF model? If so, could the authors please help me understand why a power law compute-optimal Pareto frontier is only sufficient for collapse near $\\theta=1$?\n- Section 3.2 Line 321 Right Column: The manuscript states, “This result suggests that the primary effect of learning rate decay is to accelerate the loss curves by annealing the noise introduced by SGD.” I’m confused by this claim. I would think that the noise introduced by SGD is inherent to the dataset, the batch size and the model. How is the learning rate able to suppress noise in SGD?\n- Equation 2: Since $s$ appears in the expectation and variance, I would expect it to appear inside the squared brackets. Is $s$ suppressed for brevity? Or have I misunderstood something?\n- Section 3.3 Lines 415-420 Left Column: “At late times… smaller than $\\sigma_D$.” I have no intuition for this result. Do the authors? If so, could they please add it?", "experimental_designs_or_analyses": "- I can’t tell whether I’m reading too much into this, but the sentence “We scale the model size by increasing the width D and keep depth fixed.” raised three questions: (1) Why did the authors not scale the models with depth? (2) How do the results change (if at all) when depth is scaled with width or independently from width? (3) More generally, are there realistic scaling recipes under which supercollapse does not appear?\n- Appendix C.1: I’m not quite sure I understand the methodology here. To confirm my understanding: (a) it appears that parameters were swept from ~7e6 to 7e7 for Transformer and from 2e6 to 7e7 for MLP? (b) For each parameter size, to vary the compute, presumably you swept a variety of tokens? What values of tokens were swept exactly, and why is the range of compute per parameter note the same? (c) Why were these models trained without learning rate decay? (d) In Figure 7, for each parameter value, I would expect to see some dot indicating the optimal compute for that parameter size, but I don’t; am I misunderstanding what this figure purports to show, or should the optimal compute at each parameter size be indicated by a unique marker? (e) I don’t know how much this matters, but how was the power law fit? Linear regression in log-log space, something more akin to Chinchilla with Huber loss, something else?\n- Section 2.2: What is $\\theta$? The fraction of optimal compute, yes? If so, it might be good to explicitly state this in the text.\n- Section 2.2: nit: Right column, line 135: Missing period after “training”.\n- Section 2.2: I feel like I’ve misunderstood something about the “last 80% of training”. Figure 1 (middle) suggests to me that the curves overlap for the last 90% of training. Am I misreading the data? Or where is this 80% figure coming from?\n- Section 2.3: Line 174 left column “In Figure 1 (right)”, am I confused or should this refer to Figure 2?\n- Figure 2: There appear to be no error bars or notions of uncertainty. Could the authors please repeat this experiment multiple times, or if that is too expensive, bootstrap it? I’m unsure whether the cross-D collapse deviation $\\Delta(t)$ is indeed meaningfully smaller than the cross-seed fluctuation.\n- Figure 2: What do the different hues correspond to?\n- Line 193-195, left column: “This shows that the end point of a specific loss curve efficiently encodes much of the randomness throughout the trajectory.” I’m not quite sure I understand this sentence. Could the authors please clarify?\n- Section 2.3: I'm uncertain whether the subsection's result is profound or trivial. We certainly expect this result at $\\theta=1$. The authors point out, the interesting aspect is that it stays smaller over an $O(1)$ fraction of training time, but I imagine one could handwave that we expect smooth(ish) scaling trends. \n- Figure 3 (c) Left: Could the authors please consider a larger range of $P_0$ values for the sensitivity analysis? I want to know how much $P_0$ can really be stretched before the collapse deviation increases substantially.\n- (minor) Figure 9 (a): Why is variance so large in the vertical direction?", "relation_to_broader_scientific_literature": "I think this paper contextualizes itself well w.r.t. large scale empirical scaling laws e.g., Kaplan, Hoffman, and analytically tractable models, e.g., Paquette.", "essential_references_not_discussed": "None that spring to mind.", "other_strengths_and_weaknesses": "See above.", "other_comments_or_suggestions": "- nit: “constan” on Line 533 should be “constant”\n- nit: On line 437 left column, should \"explaining why schedules at decay faster have...\" be \"explaining why schedules that decay faster have\" ?", "questions_for_authors": "See earlier boxes.", "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": {"ethical_review_concerns": "N/A"}}, {"review_id": "FuyX6n04v7", "reviewer": "Reviewer_hiQw", "summary": "The paper intvestigates the phenomenon of \"supercollapse,\" where loss curves from compute-optimally trained neural networks collapse to a single universal curve, after an affine rescaling. This universality is observed across different model sizes and learning rate schedules, and it is characterized by deviations smaller than those caused by stochastic optimization noise. The authors prove that training to a Pareto frontier is necessary for supercollapse and show that learning rate schedules deform loss curves in a predictable, scale-invariant way. They validate their claims using transformer models on CIFAR-5M and MLPs on a synthetic regression task, showing that supercollapse occurs across architectures and tasks. The paper suggests that the study of full loss curves, instead than just the final loss, can improve scaling predictions and provide insights into the training dynamics of large models.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful reading of our draft and supportive review. We will update the paper to ensure that all notations are clearly defined in the main text and fix the typos you identified. Specifically regarding the definition of $k_i$, we first sample a scalar $s_i$ from the power-law distribution defined in L140 and then sample a random unit vector $v_i$ and define $k_i$ as $s \\cdot v_i.$ This will be clarified in the text.\n\nWe have attached a PDF [here](https://drive.google.com/file/d/1ZkobNTqh90nnUcunKqUT2Dyx3T4zAb5O/view) with additional experiments requested by other reviewers to further demonstrate the generality of supercollapse. If you are interested, the details can be found in our other rebuttals. Let us know if we can clarify anything further.", "claims_and_evidence": "The claim are supported by extensive numerical evidence.", "methods_and_evaluation_criteria": "The methods are appropriate for the problem.", "theoretical_claims": "I check the proof of Theorem 3.1 and I have no issues to discuss.", "experimental_designs_or_analyses": "I have no issues to discuss.", "relation_to_broader_scientific_literature": "The key contributions of the paper relate to the literature on scaling laws and training dynamics in machine learning, expanding on previous work on empirically observed power-law relationships between compute, model size, and test risk. It goes further by uncovering a universal structure in full loss trajectories, moving beyond the usual focus on final loss values. The paper also links its empirical findings to theoretical research in high-dimensional optimization and random features.", "essential_references_not_discussed": "I am not aware of any essential references that have been omitted.", "other_strengths_and_weaknesses": "Strengths: The paper introduces the concept of supercollapse and extends the literature on scaling laws by revealing universal dynamics in the full loss trajectories. It combines a theoretical analysis (showing the necessity of a power-law Pareto frontier) and empirical validation across architectures (transformers and MLPs) and tasks (CIFAR-5M and synthetic regression), as a support to its claims. The findings and insights provided by this work have practical implications for scaling predictions and optimization strategies, suggesting that small-scale experiments can predict the behavior of larger models.\n\nWeaknesses: Some key quantities are not clearly defined in the main text (for instance $\\gamma$ and $P_0$ on line 147, $k$ in equation 6, and $M$ on line 136). While some of them can be deduced from the context and others are clarified in the supplementary material, I would still suggest that the authors clarify all notation directly in the main text to ensure accessibility for the reader.", "other_comments_or_suggestions": "- Typo on line 027: repetition of \"the\"\n- Typo on line 135, second column: missing period \n- Typo in eq. 2: there is a right square bracket without the corresponding left one\n- Could you clarify the recursive definition of $k_i$ on line 138?", "questions_for_authors": "I have no additional questions.", "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": "Xh8i9Tg1Mw", "reviewer": "Reviewer_NFpH", "summary": "This paper introduces the concept of \"supercollapse,\" where the loss curves of networks trained under compute-optimal conditions collapse to a universal curve after affine rescaling. The authors demonstrate this phenomenon across various architectures (transformers and mlps, with different dim sizes) and learning rate schedules, providing both empirical evidence and theoretical insights into why and how this collapse occurs. This study contributes to the understanding of scaling behaviors in networks and suggests new ways to optimize training processes.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review. \n\nYour point about the limited diversity of the datasets is well taken. To address this point, we conducted additional experiments on two non-image domains: the [Lichess chess games dataset](https://huggingface.co/datasets/Lichess/chess960-chess-games) dataset of 22M chess games recorded in algebraic chess notation, and [SMILES-250M](https://huggingface.co/datasets/HoangHa/SMILES-250M), a large collection of chemical structure representations. In Figure 1 (middle, right) of the [linked PDF](https://drive.google.com/file/d/1ZkobNTqh90nnUcunKqUT2Dyx3T4zAb5O/view), we show the next-token prediction loss curves of compute-optimally trained transformers and their rescaled versions. Our results demonstrate that supercollapse occurs consistently across both these datasets. In addition, we showed that scaling depth also gives supercollapse in the MLP setting similar to scaling width (Figure 1 left). Combined with our original findings on CIFAR-5M, MLP regression, and high-dimensional linear models, these results provide compelling evidence that supercollapse is a general phenomenon spanning diverse data modalities and model architectures.\n\nRegarding LLM pre-training, computational constraints prevented us from conducting meaningful experiments. The smallest model in the Chinchilla scaling law [1] study has approximately 80M parameters, similar to our largest tested model. Our limited computation meant that we could not perform high-quality scaling studies in this setting. We believe that our consistent observation about supercollapse across multiple architectures (transformer, MLP, linear model) and data modalities (images, chess, chemical structures, regression), as well as our theoretical analysis establishing the connection between supercollapse and sum-of-power-law loss curves motivates future studies on LLM pre-training, which we hope to pursue in followup work.\n\nWhile we recognize the potential practical impact of our findings for LLM pre-training, we would like to emphasize that the primary contribution of this work is establishing supercollapse as a novel and intriguing phenomenon across diverse settings, and demonstrating that studying universality in loss curves offers a promising avenue for advancing our scientific understanding of scaling. We will revise the paper to articulate this scientific focus more clearly.\n\nThank you again for your constructive feedback, which has helped improve the clarity and positioning of our paper. We will also correct the typos you mentioned in the revised version of the paper. We hope you will consider raising your score in light of our response and additional experiments.\n\n[1] Hoffmann et al. \"Training compute-optimal large language models.\"", "claims_and_evidence": "Pros: The study provides detailed experimental setups and rigorous theoretical analysis to support its claims, including:\n+ Supercollapse is observable across various architectures and learning rate schedules.\n+ Accurate estimation of compute-optimal parameters is essential for observing supercollapse.\n\nCons: However, some concerns remain. The training datasets are somewhat limited, as only CIFAR-5M was used, with next-pixel predictions as the employed task. This does not fully align with the motivation to explore the dynamics of large models, particularly large language models (LLMs) or multimodal language models, which are the primary focus. Ideally, both inputs and outputs should consist of text tokens. It remains uncertain whether the conclusions drawn from CIFAR-5M and the next-pixel prediction task are applicable to LLM pre-training.", "methods_and_evaluation_criteria": "The methods are appropriate and clearly described, but the evaluation is weak in terms of the dataset.", "theoretical_claims": "Theoretical analysis is provided to support the claims.", "experimental_designs_or_analyses": "The experimental design is solid and well-executed, both in terms of the model and the conditions. However, the inclusion of additional datasets, as previously mentioned, would further enhance the study.", "relation_to_broader_scientific_literature": "This paper provides valuable insights that could support a wide range of model training-related tasks.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "This paper attempts to address an important issue: predicting the universal curve of a model during training. It has the potential for a significant impact and could be highly beneficial for neural network training. However, due to the limited and insufficient diversity in the training data, the conclusions require further justification. \n\nAdditionally, there are a few typos and grammar errors:\n+ Line 135: \"80% of training\" should be \"80% of training.\"\n+ Line 156: \"denotes\" should be \"denote.\"", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "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": "gmGO1q9fzK", "reviewer": "Reviewer_oTD2", "summary": "Assuming that any model is trained for a number of iterations that are compute-optimal given its size, this paper shows, both empirically and theoretically, that the training curves of models of different widths are identical up to an affine transformation.\nDeviations due to randomness in the training procedure are larger than those across model width.\nFurthermore the paper shows, both empirically and theoretically, that the result holds across a variety of learning rate schedules.\nThese results can be used to improve scaling predictions.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and supportive review. We have attached some additional figures [here](https://drive.google.com/file/d/1ZkobNTqh90nnUcunKqUT2Dyx3T4zAb5O/view), and address your specific questions below. \n\n**On the limited range of widths in CIFAR-5M experiments.**\nWe acknowledge the limitations of this experiment, which was a combination of our computational constraints as well as the small size of the dataset itself. We set the maximum width to 2048 in the CIFAR-5M experiments to 1) limit the computational cost of our experiments, and 2) avoid excessive repetition of the training data. Based on the estimated compute-optimal training horizon, the width 2048 model already required training on the CIFAR-5M dataset (which only has 5M training images) for over 10 epochs. A larger dataset is necessary to train even wider models to compute-optimality without overfitting effects. \n\nWe note that our MLP experiment does not suffer from this data constraint, where we scaled up the width by 8x from 512 to 4096, corresponding to a 64x scale up in model size, and demonstrated supercollapse across all models.\n\nIn addition, we demonstrate in Figure 1 (middle, right) [here](https://drive.google.com/file/d/1ZkobNTqh90nnUcunKqUT2Dyx3T4zAb5O/view) that supercollapse occurs in two additional data modalities: chess game string representations and chemical structures, and along one more scaling dimension: depth, further establishing the generality of our observations.\n\n**On the utility of Equation 6.**\nWe emphasize that our claim regarding the dominance of two terms in Equation 54 is mathematically valid in the asymptotic limit of large $D$, not merely a convenient approximation. This generalized analysis of power law loss curves extends the connection between supercollapse and compute-optimal scaling beyond the simplified form in Equation 7. This extension is particularly important because recent theoretical work [1] demonstrates that this broader class of loss curves naturally emerges in high-dimensional regression models, whereas the two-term power law in Equation 7 fails to capture significant regions of the phase diagram (Figure 2, [1]).\n\nTo empirically validate our theoretical findings, Figure 6 attached via [this link](https://drive.google.com/file/d/1ZkobNTqh90nnUcunKqUT2Dyx3T4zAb5O/view) presents additional experiments using the Power Law Random Features (PLRF) model introduced in Section 3.1. Unlike our setup in Figure 4, we specifically selected values of $\\alpha$ and $\\beta$ to produce loss curves asymptotically described by three power laws (detailed in [1]), which cannot be accurately represented by Equation 6. These results confirm that supercollapse occurs exactly as predicted by our analysis in Appendix G. Without our generalized framework, we would lack the theoretical foundation to explain supercollapse in these more complex settings.\n\n**On determining the optimal decay schedule**\nFor a fixed $\\hat{K}$, equation 10 suggests that decaying the learning rate as late as possible to maximize total gradient flow time is best. However, as we show in Figure 5 and Figure 9 in the paper, the best fit $\\hat{K}$ has non-trivial dependence on the schedule and correlates with the preconditioned NTK trace, which tends to be higher for schedules that start decaying earlier. This gives a tradeoff between progress in gradient flow time (favoring later decay) and more favorable geometry (favoring early decay). Understanding this tradeoff requires a deeper understanding of $\\hat{K}$, its dynamics, and its relationship to the schedule. We believe this is an exciting direction for future research on optimizing learning rate schedules.\n\nThank you again for your review. We hope we have addressed your questions. Please let us know if we can clarify anything further.\n\n[1] Paquette et al. \"4+3 phases of compute-optimal neural scaling laws.\"", "claims_and_evidence": "This work can potentially help scaling up models more accurately.\nThe combination of empirical and theoretical results, and the agreement between the two, is compelling.\nThe density and significance of results is high.", "methods_and_evaluation_criteria": "In CIFAR-5M experiments, the model width varies very little, less than a factor of three, which corresponds to a one order magnitude variation of the parameter count. Fitting power laws with such small variation is unreliable.", "theoretical_claims": "It remains unclear how much useful is Eq.6. The proof of collapse in Appendix G assumes that the sum in Eq.6 can be approximated by two terms. Eq.7 is also a special case of Eq.6 with two terms of the sum only. So it remains unclear whether we need Eq.6 at all.", "experimental_designs_or_analyses": "OK", "relation_to_broader_scientific_literature": "OK", "essential_references_not_discussed": "OK", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "If it’s true that the primary effect of learning rate decay is annealing SGD noise, then can we determine what is the best way of decaying the learning rate?", "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:

Meta-review: The paper introduces a novel phenomenon called supercollapse, where loss curves from compute-optimal neural network training collapse to a single curve (after an affine rescaling).

The paper provides a theoretical analysis (e.g., Theorem 3.1, analysis of power-law Pareto frontiers) to explain why supercollapse occurs, as well as experiments using transformers and MLPs across different learning rate schedules.

The reviewers raised some concerns over the limited variation in model widths and noted some minor issues but they generally regard the paper as solid and thorough, with most scoring it in the strong accept to accept range. I therefore recommend acceptance.
Decision: Accept (oral)