On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling
Authors: Moritz Haas
Sebastian Bordt
Ulrike von Luxburg
Leena Chennuru Vankadara
Venue: NeurIPS 2025 spotlight ()
Keywords: Width Scaling
Feature Learning
Large Learning Rates
Maximal Update Parameterization
Standard Parameterization
Initialization
Scaling Theory
Tensor Programs
Primary area: deep_learning

Abstract: Scaling limits, such as infinite-width limits, serve as promising theoretical tools to study large-scale models. However, it is widely believed that existing infinite-width theory does not faithfully explain the behavior of practical networks, especially those trained in *standard parameterization* (SP) meaning He initialization with a global learning rate. For instance, existing theory for SP predicts instability at large learning rates and vanishing feature learning at stable ones. In practice, however, optimal learning rates decay slower than theoretically predicted and networks exhibit both stable training and non-trivial feature learning, even at very large widths. Here, we show that this discrepancy is not fully explained by finite-width phenomena.

Instead, we find a resolution through a finer-grained analysis of the regime previously considered unstable and therefore uninteresting. In particular, we show that, under the cross-entropy (CE) loss, the unstable regime comprises two distinct sub-regimes: a catastrophically unstable regime and a more benign controlled divergence regime, where logits diverge but gradients and activations remain stable. Moreover, under large learning rates at the edge of the controlled divergence regime, there exists a well-defined infinite width limit where features continue to evolve in all the hidden layers. In experiments across optimizers, architectures, and data modalities, we validate that neural networks operate in this controlled divergence regime under CE loss but not under MSE loss. Our empirical evidence suggests that width-scaling considerations are surprisingly useful for predicting empirically maximal stable learning rate exponents which provide useful guidance on optimal learning rate exponents. Finally, our analysis clarifies the effectiveness and limitations of recently proposed layerwise learning rate scalings for standard initialization.

On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling
Moritz Haas'
Sebastian Bordt
Urriee von Luxburg'
Leena Chennuru Vankadara 2
University of Tübingen, Tübingen Al Center
(mo. haas, sebastian. bordt, ulrike. luxburg eetee/iiiiiiingneeeeeeee de
Gatsby Computational Neuroscience Unit, University College London 1. aat.aaddaaaull.ac..o..
Abstract
Scaling limits, such as infinite- width limits, serve as promising theoretical tools study large-scale models. However, is widely believed that cxisting infinite-width theory does not faithfully explain the behavior of practical networks, especially those trained in standard parameterization (SP) meaning He initialization with a global learning rate. For instance, existing theory for SP predicts instability at large learning rates and vanishing feature learning at stable ones. In practice, however, op- timal learning rates decay slower than theoretically predicted and networks exhibit both stable training and non- coiviaivial feature learning, even at very large widths. Here, we show that this discrepancy is not fully explained by width phenomena.
Instead, we find a resolution through a finer- mner-grained analysis of the regime pre- viously considered unstable and therefore uninteresting. In particular, we show that, under the cross-entropy (CE) loss, the unstable regime comprises two distinct sub-regimes: catastrophically unstable regime and more benign controlled divergence regime, where logits diverge but gradients and activations remain stable. Moreover, under large learning rates the edge of the controlled divergence regime, there exists well-defined infinite width limit where features continue to evolve in all the hidden layers. In experiments across optimizers, architectures, and data modalities, we validate that neural networks operate in this controlled divergence regime under CE loss but not under MSE loss. Our empirical evidence suggests that width-scaling considerations are surprisingly useful for predicting empirically max- imal stable learning rate exponents which provide useful guidance on optimal learn- ing rate exponents. Finally, our analysis clarifies the effectiveness and limitations of recently proposed layerwise learning rate scalings for standard initialization.
Experiment Code
Refined Coordinate Check Package
Introduction
Scaling has become the dominant paradigm in building ever more capable vision and language models (Brown et al., 2020, Dosovitskiy et al., 2021, Kaplan et al., 2020, Hoffmann et al., 2022, Grattafiori et al., 2024). Infinite-width limits have served as a crucial theoretical tool for understanding large models, providing valuable insights into their optimization and generalization behaviour (Jacot et al., 2018, Du et al., 2018, Allen-Zhu et al., 2019, Arora et al., 2019, Mei et al., 2018, Rotskoff and Vanden-Fijnden, 2022, Sirignano and Spiliopoulos, 2020, Lai et al., 2023). Nevertheless, it is now widely believed that existing infinite-width theory, particularly in the kernel regime, does not
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

serve as faithful proxy for practical neural networks, as it fails to capture fundamental aspects of
their training (Sohl-Diae -ick-Dickttein et al., 2020, Lee et al., 2020, Vyas et al., 2022, Wenger et al., 2023).
This disconnect is especially prominent for the dominant training practice, standard parameterization (SP): He ini- tialization (He al., 2015) with a single global learning rate (OLMo Team et al., 2024). For example, infinite- width theory predicts that under SP, network dynamics should become unstable with learning rates scaling larger than (1/n) (where is network width), and that feature learning vanishes with (11/n) learning rates, causing the models enter kernel regime (Sohl-Dickstein et al., 2020, Yang and Hu. 2021). Empirically, however, net- works trained in SP exhibit stable feature learning and excellent generalization performance, often with optimal learning rates decaying much slower than theoretically predicted (commonly around ( /T)) ). This is depicted in Figure where we see that the optimal learning rates (solid lines) for different models trained in SP decay much slower than the theoretically predicted maximal stable seal- ing law (dashed gray lines). These observations represent fundamental puzzle and motivate two crucial open questions:
Optimal learning rate
-0,38
MLP SGD
10"
GPT SGD
GPT ADAM
10"
-0.46
00.42
10"
10
258
1024
4098
16384
Width
Optimal learning rate ex- ponents exceed the theoretically pre- dicted stability threshold. For ML.s on MNIST and GPT on language data, optimal learning rates in SP decay slower than the theoretically predicted maximal stable Th. = ((n in gray.
Figure 1:
Why does SP remain stable and effective at large learning rates, despite the theoretical predictions? And does there exist an infinite-width limit that corresponds more closely with the behaviour of practical finite-width networks?
In this work, we investigate these questions and find a resolution through a finer grained analysis of the regime previously dismissed as unstable. Consequently, we provide the first infinite- width proxy that corresponds more closely to practical, finite- wiith networks. Our main contributions are:
(a) We investigate and rule out finite-width effects as the key explanation. Plausible explanations to explain these discrepancies include finite-width effects, either accumulated over large depth or longer training times. We show these explanations are insufficient, as the gap is surprisingly more pronounced in shallow (2-layer) MLPs in a single pass (Figure F.15). Furthermore, we investigate other known winth dynamics, like the catapult regime (Lewkowycz et al., 2020) or the edge of stability (Cohen et al., 2021, 2022) in simplified linear models, and show that these mechanisms alone also cannot explain the stability of large learning rates in SP.
(b) We validate infinite-width alignment predictions. Contrary to what was hypothesized in previous work (Everett et al., 2024). we show that the infinite- width alignment predictions between weights and incoming activations indeed hold at moderate width when measured with sufficiently refined coordinate checks (RCC). This is crucial finding, since it confirms the theoretical prediction that logits do diverge at sufficient width in SP under empirically optimal learning rates. (c) The resolution. Instead, we find a resolution to these discrepancies through fine-grained analysis of the regime previously considered unstable and therefore uninteresting, In particular, we show that, under the CE loss, the unstable regime comprises two distinct sub-regimes: catastrophically unstable regime and a more benign controlled divergence regime, where logits diverge but gradients and activations remain stable. Moreover, at the edge of the controlled divergence regime, which corresponds scaling the learning rate as 11/2 for deep MI.ss under SP, there exists a well-defined infinite width limit where features continue to evolve in all hidden layers, which could partially explain the practical success of SP. To the best of our knowledge, this provides the first practical infinite-width limit in the feature- regime for SP.
(d) Empirical validation. We show that our width-scaling considerations provide surprisingly accurate predictions of maximal stable learning rate exponents, which often dominate optimal learning rates, particularly in Transformers (Vaswani et al., 2017). However, while output-layer divergence under CE loss remains benign with respect to the stability threshold, it occasionally influences the optimal learning rate choice. As one important example, previously observed learning rate transfer under layerwise learning rates breaks on all considered image datasets. Avoiding logit divergence while recovering feature learning with jP, on the other hand, opens up more loss functions such MSE loss as competitive alternatives.

---

Taken together, our results deepen the theoretical understanding of why SP remains effective at large scales, provide thorough empirical validation for critical assumptions in infinite-wite widhh theory, and offer practical insights into stable hyperparameter transfer for scaling neural networks.
2 Background: Width-scaling arguments from Tensor Program theory
Before exploring plausible explanations for the empirical width-scaling properties of neural networks, we first define used notation and distill necessary width-scaling arguments from Tensor Program (TP) theory (Yang and Hu, 2021, Yang and Littwin, 2023). We provide a more detailed introduction to TP scaling arguments in Appendix C.I, I, and a detailed account of related work in Appendix A.
Setting and Notation. We define an (T. + Jl-yee MIP of width n iteratively via h' (E) = L',,
z'(5)== (E),,
(S) == w (y(+12'(E)) x'{), f(E) = w
for inputs & E Rdin with trainable weight matrices W E Rxdd... W E Rux= for I E [2, L). and WT+1 E Rduxxn We call h' preactivations, x'' acctvvttions and f(E) output logits. Training the MLP with Stochastic Gradient Descent (SGD) with global learning rate
L TR x Rhe - R with labelled training point (Ee, E Rd x R > 0 under loss function is defined as W = w - ((L(E)) Mt). We denote updates accumulated over all time steps by A M2 - hy and the change in a single update step by 5h = h4 ---h h4-1 The fan-notation has the purpose of unifying all weight matrices and simply means W E @gfan_outxfan_.nn
In this paper, we define standard parameterization (SP) mean He initialization (W)), NO, faan/annnaa/a/ in trained with SGD or Adam with a single possibly width-dependent learning rate 77n = 7) n²" a E R. for all trainable weights {W:) [t[[++]] This models the typical practice, in which a global learning rate is tuned at each model scale. We denote the softmax function by af), = exp(Si) (jj((t.)) exp(f)))=1 In this paper, by CE loss, we refer to the concatenation L c a of the cross-entropy loss func- tion v( = y log(f) and the softmax as is the dominant practice implemented in torch. CrossEntropyLoss. For naturally measuring the average scaling of entries in vectors Rd,
the moo--mean- norm ||a]] RMS == d-1/2 ||o|||2 as the standard vector norm. For matrices W, we write ||W|| for the Frobenius norm and measure entry-wise scaling with the
RMS norm || WRMMS := (F 1/2 |L. The operator norm w.r.t. the RMS-normm is de- fined as || Wl == ||||MMssS == sup ( [ z:| iass// / [[=)) We use Bachmann- Landau notation O, G,Ω that purely tracks dependence width and omits all other dependencies.
Effective and Propagating Updates. When training neural networks, weights Wt of layer evolve from their initialization Wi through updates AW/, such that W W + AWW Although directly control the scaling of these initial weights and updates, we are ultimately interested in their impact on subsequent activations in the network. For standard architectures, including convolutional networks and Transformers, weights typically act linearly on incoming activations. Thus, for weights W and incoming activations 2 the change in the next layer's pre-activations An; can be decomposed into two distinct contributions: the effective updates arising directly from the change in weights AW? of the current layer, and the propagating updates, arising indirectly from activation changes Ar, A in preceding layers:
AlG = (AW))a + Wo(Ao!-1).
(RCC)
Ellective Updates
Propagating Updates
We say a layer admits maximal stable feature learning if both the effective updates and propagating updates remain width-independent network width n + co, that is (I(AW)) RI = @(1) and |[W] [[W -- rMs O(1) This definition of feature learning as non-vanishing effective updates formulates a necessary but not sufficient condition for learning well well-generalizing features at large width, common in related literature (Yang and Hu, 2021, Vyas et al., 2024, Bordelon al., 2025).
Identifying the correct scaling exponents. In the spirit of Everett al. (2024), we use pı and a denote the width-scaling exponents of the alignment ratios of the pairs (AW), [. xd11) and (W.) A respectively, that is,
|/WW/ 4-1 RMS
I
Art
RRMS
O(PP))
O(n
(a-rms)
|A AWW/IRMs ||a
W
RMS ||r RMS
3

---

LAWa, || Rusllll))
||AWsal/[[AWW]l..... (x
|Amll)))
10
10"
10"
-0.0006
-0.10
-0.04
leyer
width
-0, 0.065
layar
-0.16
bod UM
10"
254
0.26
und LN
10
ntp haddn
1104
TTp mophidden
4000
ouEaver
16
10"
10'
10
10'
10"
10
10' Training Step
Training Step
Training Step
Figure 2: Aligument has minimal width-dependence. Alignment ratio between accumulated weight updates AW; and incoming activations xi in RMS norm (left) and operator norm (center) as well as between initial weights Wo and activation updates Ar; in operator norm (right) for the last layernorm layer, the first MLP layer in Transformer block 2 and the readout layer. RMS norm may be confounded by accumulated rank over the course of training (e.g. compare (AW,, Zx) values for last LN). While operator norm alignment tends to decay over the course of training, it does not display strong width- deeendeeeen,, even after 2000 batches (see annotated width-dependent exponents).
A key insight from Yang and Hu (2021) and Yang and Littwin (2023) is that during training, correlations can emerge in certain layers between the two quantities in each pair in (RCC), causing them to become aligned the infinite-width limit and thereby inducing Pt and q: = due to a law of large numbers effect. If, instead, these quantities were uncorrelated, their product would exhibit smaller scaling exponents (pi = 1/2 and q: = 1/22 due to a central limit effect. In particular, inftn theory predicts the exponents PI.L+1 1. 91-L 1/2, and 9L+1 1. The alignment exponents Pt, qi are a consequence of gradient- parameterization used (e.g., SP, NTP, IaP).
training and do not depend on the specific
By adjusting the initialization variance, which controls the scale of initial weights Wo. and the learning rate, which governs the magnitude of updates AW, we can ensure that both contributions in (RCC) remain indth-independent as the network width grows. The corresponding choice of hyperparameter scaling defines the Maximal Update Parameterization (nP). As we will discuss in Section 4, under the theoretically predicted alignment exponents. SP with ((1/n) learning rates leads to vanishing activation updates in all layers |Ar,]I = o(1) and choosing the learning rate w(1 (11/)) leads to logit divergence in the infinite-widde limit.
3 Finhte-width distortions and long-training dynamics alone do not explain the stability of large learning rates in SP
Discrepancies between finite- vs infinite-width width networks is often attributed to finite-width effects. Plausible explanations include dynamics induced by (1) large depth or (2) longer training time which may induce accumulation of knite-width effects over multiple steps (e.g.. when T > n). Here, we show these explanations are insufficient. In Figure F.15, we find this discrepancy is surprisingly more pronounced in shallow (2-layer) MLPs in a single pass under SP. This finding-than that the issue persists even without large depth or multi-epoch accumulation other mechanisms. In this section, we explore two such explanations.
investigating
3.1 Update alignment between weights and activations is barely width-dependent
Everett et al. (2024) highlight that at finite width and over extended training times, a priori unclear whether the pairs (AW!, xz ) and (W--1. Ax,") remain strongly correlated or whether their alignment exponents (PI-L+1, 11+1) should rather be thought of as dynamically changing over the course of training. If the alignment exponents instead transition towards the central- limit regime and in particular if P1:Z+1 1/2, this could explain the observed v gap between theoretically predicted and empirically observed optimal learning rate scalings.
Since our objective is measuring with which bidth-scaling AW! propagates incoming activations xt - forward, but rank accumulation can decouple the RMS norm |AW||| RMS from this objective, we measure alignment with operator norms (Yang et al., 2023a),
|| Axr ||MM
A.x
(a-op)
A |RRMS || RMS

---

Specifically, if the alignment exponents P1-L+1 = 1, 91: 9L+1 hold, both contributions in (RCC) must propagate signals forward maximally as function of width with alignment metrics "awi - and Ow, of order (1) unifying all layers. In Appendix C.2, we explain our alignment measurement considerations in more detail.
In Figure 2, we plot the alignment metrics at varying widths over the course of Transformer training with AdamW in SP. The figure shows that while alignment can decrease over the course of training, exhibits minimal dependence on network width. Even after accumulating approximately 2000 batches of training, the width-scaling exponents are much closer to 0 than to -0.5, indicating that unfnite-widtth alignment predictions hold reasonably well. Hence a lack of alignment alone cannot explain the large optimal learning rate exponents observed in practice.
3.2 Does a catapult mechanism in the first update steps stabilize large learning rates in SP?
As another plausible explanation, initial divergence under large learning rates may be stabilized over the course of training at finite width. Unlike at infinite width, where there only exist divergent regime and a lazy regime without feature learning, an intermediate catapult regime was identified by Lewkowycz et al. (2020), for SGD training with MSE loss in Neural Tangent Parameterization (NTP) at finite width. They provide theory for 2-layer linear networks. Under small learning rates v) S 2/xo. where d denotes the largest eigenvalue of the Hessian at initialization, the network monotonically converges to minimum. Under large learning rates > 4/ /o, training diverges. But in an edge of stability regime (Cohen et al., 2021, 2022) of intermediate learning rates, the loss increases in the first (log(n)) update steps while the sharpness A decreases. Once the sharpness lies below the edge of stability 2/m, the loss decreases and the final learned function may generalize better as the solution lies in a basin with lower sharpness. But existing work does not study width-scaling with SP. May similar initial training dynamics be at play here?
In Appendix C.4 we analyse the 2-layer linear network model from Lewkowycz et al. (2020) in NTP, SP and μP trained with SGD under MSE loss, and provide loss and sharpness increase characterizations in Proposition C.18. In jPP the update equations of the learned function fe and the sharpness A are fully mndth-independent. which allows width-indepente least the conditions for loss and sharpness reduction are approximately width-indepee width-independent. In SP. on the other hand, sharpness increases \ 2 A iff X Z 707,,
learning rates. In NTP, at
4 (1 y. requiring 1,, O(n
to avoid sharpness (as well as loss) divergence the first update steps. The simulations shown in Figure C.2 validate the maximal stable learning rate scaling 1) O(n 1). Hence catapult dynamics
alone do not suffice for explaining large learning rate stability in SP.
4 Cross-entropy loss enables stable feature learning under large learning rates in standard parameterization
First, let us briefly recall why infinite-widie theory predicts divergence under SGD training in SP with learning rates 7/, = nn for a < 1.
Recall that the alignment exponents in (a-rms) satisfy Pi.++1 = 1. In particular, for the output layer, we have || of RM = O(n ||AW/++1 |RMS ||a| lRRM)). For SGD, the weight update of the last layer after update step is given by AW1+11 = -1 n Xo (a where Xo ((o(ff(E))) yo). Under SP, at initialization, both |= |=fllrss = @(1) and
O(1).
This implies logit divergence after step of SGD with learning rates yn = w[1nn):
| IIMss = @(1), HWWL+IIlaas O(n = AWWLt10-RRME O(n2-a)
So, why do larger learning rates remain stable and even effective, despite logit divergence?
Here, we demonstrate that simple yet fundamental aspect of training, the choice of loss function, resolves the large learning rate puzzle, and enables a well-defined and practical infinite-wiie limit that allows feature learning under SP. The key insight is that, under cross- cross-entropy (CE) loss, the logits never directly appear the training dynamics; instead, the effective output function is off). Unlike the destabilizing logit blowup encountered under mean squared error (MSE) loss, under CE loss, logit growth has harmless effect on training stability. Therefore, CE loss introduces an intermediate controlled divergence regime that is absent for the MSE loss (Figure 3).
5

---

Figure 3: Learning rate regimes for SGD in SP. Under MSE loss, training a deep MLP either remains stable (ax 2 1) or logits and hidden-layer activations diverge (a < 1) in the wifth limit. Under CE loss, controlled divergence regime a E [1/2, 1) emerges where logits di- verge, but training does not diverge. At 1/2, hidden layers learn features width-ind aeatt-inddeeddently.
MSE loss
Stable
Catastrophic Instability
Hidden-layer feature learning
Controlled Divergence
Catastrophic
CE less
Stable
Instability
T1
1
Y1
1/2
Learning Rate 1/n
Definition 1 (Learning regimes). Fix t E N. We say that training lies in the stable regime iff the activations || Rars of all layers remain ((1) and the logits [ RMS remain ((1) after update steps. We say that training is catastrophically wnstable iff the logits and activations in at least one layer diverge after steps, that is [[fill -> oo and 3 I E [7] such that ||=|]] nMS - so. If training neither lies in the stable nor in the catastrophic regime, it lies in the controlled divergence regime.
Proposition shows that there exists a non-vanishing controlled divergence regime under the CE loss.
Proposition 2. (Asymptotic regimes in SP, informal) fixed I, 2 2, 1 1, 0, a E R, consider training a (L + 1)-layer MLP of width n in SP with SGD and global learning rate The = rj n for t steps. Then the logits fú loss- derivatives Xt (f((((( yt), loss-weight gradients V = V Ww ((f((E) Vt) and activations zhy I E [L], after training scale as follows n 00.
Under cross-entropy (CE) loss, three qualitatively distinct regimes arise:
(a) Stable regime (a 1): Logits, gradients and activations remain stable, that is
[lfel RMS O1)), lixelRMs = O(1), [|V RMS Onn and ||a RR = (1) for all 1 E [L].
(b) Controlled divergence (1 a < 1): Logits diverge [lfell RMS = O(n²-a). but gra- dients and activations remain stable, that is || = @(1, llxe R = 0(1) and
V Villums- = On for all 1 E [T..
(c) Catastrophic instability (a < 4:: Logits, activations and weight gradients diverge, that is [ffell RMS - cc. ||
and | + oo, I L].
Under mean-squared error (MSE) loss, a stable regime as in (a) above arises i 1. If a 1, training is catastrophically unstable as in (c) above and, in addition, llxelRRMS + DC.
Remark 3 (Maximal-stable learning rate). In SP, different layer types learn at different rates. For
SGD, the output layer logits remain stable iff The O(n The input layer. biases, Layernorm gains and embedding layers all behave input-like and remain stable iff Tn = ((1). All other trainable weights in typical CNNs and Transformers behave hidden-lib and remain stable iff 7m = O(n 1/2))
If we define the maximal stable learning rate (max-stable LR) scaling of a network as the edge to the catastrophic regime, and the max-stable LR scaling of each layer that above which the layer's output diverges (Definition C. 1), then Proposition 2 shows that the max-stable LR of a network can exceed that of individual layers. Since CE loss does not require output layer stability, the proposition naturally extends to 2-layer MLPs, which do not contain a hidden layer, so that their controlled < (Remark C.14).
divergence regime under CE loss extends
The formal statement together with a proof can be found in Appendix C.3. For an intuitive under- standing of this result, note that the only effect that the choice of loss function L(,, y) has on the final learned function is through the loss-logit gradients X == (fr y/) over the course of training. Under MSE loss, the loss gradients are given by the residuals X fe(E.) Yt- But CE loss induces loss gradients X (ff(E))) - y:- Crucially, it is the correct choice of loss function to effectively view (f)) as the output of the network instead of the unnormalized logits f. If one were to use MSB(a(f)) (),y) y) as a loss function instead, additional derivative terms can induce vanishing gradients under exploding network output and not increase the optimal learning rate exponent (Appendix F.4). Under CE loss, the effective network output (f) at most converges to one-hot predictions when the logits diverge, and with increasing width training points are sharply memorized after a single update step. At large learning rates n) Onn-1/1 1/2), training points are not just memorized in last-layer weights, but feature learning is recovered in the infinite-width limit:

---

Input layer
Hidden layer
Output layer
|A RASS
update step
0.5
10"
10
10
10
64
128
206
012
1024
2048
4050
B182 16384
154
128
250
512
1024
2048 4060 815210334
128 206
012
1024
2048
4090 8182
Width
Width
Width
Figure Hidden-layer feature learning albeit logit divergence SP under large learning rates. Effective (-th layer update scalings a |RMS of MLPs trained with SGD in SP with Th = 0.0001 (n/226)-1/2 -1/2 on CIFAR-10 10 under CE loss. Our TP scaling predictions are accurate: Hidden layers learn features width-independentll, and input layers have vanishing feature learning. The update scaling exponents can already be accurately estimated at small width S 512.
Proposition& 4 (Under CE loss, SP with large learning rates learns features at large width, informal). Consider the setting of Proposition of training a (L+ -layer MLP with SGD SP with global learning rate Tn. 1) n a E R, in the infinite-width limit n. - 3c.
(a) Under both MSE and CE loss in the stable regime (a: 2 1). feature learning vanishes in all layers E [L] that is || RMS = =O (12).
(b) Under CE loss in the controlled divergence regime (2 S 1), input layer feature learning
vanishes at rate | RMS θ (n-1/2-0), and hidden layers I [2, L] learn features at rate || RMS = θ (n)/2-a). In particular, when a = 1/2, the weight updates of all hidden layers induce width indth-inddeendent activation updates, that is || A = (1).
RMs
To the best of our knowledge, this provides the first infinite- width limit of SP in the practical feature learning regime. Figure 4 empirically validates that the predicted width-scaling exponents that induce maximally stable feature learning despite logit blowup under T/s = r n- 11/2 are already accurate moderate width 512. Appendix E.4 shows that effective update predictions also hold accurately in Transformers trained with Adam. In the next section, we discuss the implications that training stability despite logit blowup has on learning rate scaling exponents in practice.
Consequences of training stability under logit divergence
In this section we perform extensive experiments to empirically evaluate the implications of the stability and feature learning predictions of our infinite-wide theory from the previous section.
Experimental details. We train MLPs of varying depth up to width 16384 with plain SGD and Adam on CIFAR-10, MNIST and a generated multi-index model (reported in Appendix F). We also train Pythia-GPT's with warmup and cosine learning rate decay on the DCLM-Baseling dataset (Li et al., 2024) up width 4096 or 1.4B parameters using both Adam with decoupled weight decay (Loshchilov and Hutter, 2019) and SGD (reported in Appendix F.2). If not stated otherwise, we consider SP with global learning rate. In this paper, we train for a single epoch to prevent overfitting effects. We leave a systematic study of the multi-epoch setting to future work. All details can be found in Appendix D. Open-source code to reproduce our experiments is publicly available.
Refined coordinate checks. Our results demonstrate that disentangling the width dependence of effective updates and propagating updates in each trainable weight in refined coordinate checks (RCC) can serve as an useful diagnostic tool for understanding and correcting layerwise update signal propagation, for improving training stability and performance at scale. Our implementation is easy to adapt and publicly available.
5.1 Infinite-widie theory is a useful predictor of empirical optimal learning rate exponents
While, in general, the optimal learning rate depends on the architecture and dataset (see Appendix F. it often saturates at the accurately predicted maximal stable learning rate in deep non-linear networks. We hypothesize that maximal stable feature learning in all layers induces optimal performance at large width. However, since different layer types require different maximal stable learning rate exponents, the single global learning rate under SP is subject to opposing forces for recovering feature learning under the constraint of training stability. We now evaluate several instantiations of this hypothesis.
7

---

SP, 5GD
SP. ADAM, train LN
SP ADAM fixed LN
.0.5
Train
width
256
4.3
4.3
-.0
.0.5
1000
bwil
0.42
4006
35
35
Max LR (width250000.55
LR (widiv258)"0055
10"
Max (widh:2566)11
10¹ 13
2M6
-0.56
With
Figure 5: Approximate learning rate transfer for GPT in SP. Left to center-right: Width-scaled learning rate versus training loss for GPT trained with SGD, Adam with trainable Layernorm parameters and Adam without trainable Layernorm parameters. Right: Corresponding optimal (solid) and maximal stable (dashed) learning rate exponents. For SGD, hidden-layer stability 7/ Onn-nnnn 1/2) clearly dominates the maximal stable as well as optimal learning rate scaling. For Adam without Layemorm parameters, hidden-layer stability induces a stability threshold 7h = O( ")) Trainable Layernorm parameters further stabilize large learning rates and induce larger optimal learning rate scaling Th R O(n toward preserving input-layer feature learning at scale.
MLPs and Transformers with SGD. Figures 5 and 6 show that the empirical maximal learning rate exponents under CE loss closely follow a = 1/2 for both MLPs on vision data and for GPT on language data. The x-axes scale the learning rate with the closest clean exponent from {0, 0.5. 13 to show that approximate empirical transfer is often enforced by the stability threshold O(n (1/2). While the theory only predicts the maximal stable exponent, Proposition 4 suggests that the optimal learning rate may follow the maximal stable exponent a = 1/2 since it is the only choice under which feature learning is preserved at large width in all hidden layers. When optimal learning rate exponents exceed the maximal stable exponents, as for CIFAR- 10, optimal learning rates saturate at the maximal stable learning rate at realistic width > 16384 on all considered datasets (see Figure F.11). The maximal stable learning rate under MSE loss also consistently scales as its infinite-width prediction (n= and optimal learning rates closely follow this exponent, as under smaller exponents a >> 1, not even logits are updated || fell RMS -> O. However, this approximate learning rate transfer under MSE loss is not useful, since the loss also converges fast and does not monotonically improve with scale. Overall, this shows that cxisting infinite- width theory was indeed predictive of the maximal stable learning rate exponents under MSE loss, but that CE loss induces qualitatively more favorable behaviour that is only captured by a finer-grained analysis.
MLPs with Adam. Adam approximately normalizes the gradient and therefore further stabilizes training against misscaled gradients beyond the effect of CE loss. W is effectively updated if the learning rate scaling counteracts the scaling accumulated in the inner product between normalized weight gradients and incoming activations. This leads to the ideal (μP) learning rates ry(W)) = D/fan_in(W'' (.').. Thus Adam in SP with n/e ((n=)) induces width-indeepndder updates, except for vanishing input layer feature learning and logit divergence through W Aut. For deep MLPs on image datasets, Figure F.25 shows optimal learning rate exponents close to 7/s- = O(n for both CE and MSE loss, suggesting a controlled divergence regime also arises from stabilized backward pass and that stable hidden-layer feature learning dominates the optimal learning rate scaling.
Transformer training with AdamW. In Transformers with trainable Layernorm parameters, which scale input-like, training is stabilized, and the exponent is increased toward input layer feature learning. Without trainable Layernorm parameters, in contrast, only the embedding layer scales input-like so that training becomes approximately width- aaath-independeett under nn = O(n "). Figure 5 shows that the max-stable and optimal learning rate exponents shrink from - 1/2 toward - 1 if we remove the trainable norer-norm parameters. This suggests that trainable scale parameters in normalization layers play essential role in maintaining high learning rates in Transformers, which could explain why they are almost unanimously used in modern architectures (OLMo Team et al., 2024, Grattafiori et al., 2024, Gemma Team et al., 2024). Moreover, input layer learning vanishes at scale in SP, which may explain techniques like removing weight decay in the embedding layer (OLMo Team et al., 2024). Logit divergence under large learning rates may be a reason for regularizing techniques like the z-loss (Chowdhery et al., 2023, Wortsman et al., 2024, OLMo Team et al., 2024).
Taken together, our empirical evidence suggests that infinite-width theory may serve as helpful proxy for understanding practical neural networks at finite width. Since training divergence imposes a hard constraint on the optimal learning rate and activation divergence in multiple layers becomes harder to stabilize, width-scaling predictions seem to hold even more accurately on deep and sensitive architectures such as Transformers.
8

---

controlled divergence. Closest clean ponent of max-stable LR follows theoreti- cal prediction -1 for SP under MSE loss, -0.5 for SP under CE loss and 0 for SP- full-align under CE loss. Max-stable LR often dominates optimal LR at sufficient width SP. Corresponding learning rate Figure 6: LR exponents at the edge of
0.00
LR Exponent
-0.25
0.50
-0.75
-1.00
Max
SP full align SP GE loss SP MSE loss
-1.25
X
Optimal LA
Theory
curves are provided in Appendix F.2 for GPT and in Appendix F.3 for MLPs.
-1.50
DLCM- Baseline
MNIST
CIFAR-10
Fashion MNIST
Tiny ImageNet
5.2 A novel understanding of standard initialization with layerwise learning rates
Everett et al. (2024) perform extensive Transformer experiments, and recommend training with Adam in SP with yP learning rates (SP-full-align) as the overall best performing parameterization in terms of validation loss, learning rate transfer and learning rate sensitivity. This parameterization only differs from μP through the larger last-layer He initialization W6+1 N N(., 21)). While the authors attribute the success of SP-full-align to a lack of alignment between W and Azt, they only measure the joint alignment between W; and De for each layer, which confounds the individual alignment exponents of (AW,, 24) and (Wo, Ar;)) from (RCC). We provide detailed explanation in Appendix C.2. Our empirical alignment reevaluation in Figure 2 and Appendix E.4 does not support the hypothesized lack alignment. Thus, at sufficient width, logits diverge through W Axx soon feature learning does not vanish. Instead our theoretical results in Section 4 show that logit divergence does not harm training stability under CE loss. Just like SP with 7, O(n -1/2), SP-full-align with 7 = (1) lies at the feature learning edge of the controlled divergence regime. Learning rate transfer of SP-full-align breaks on image datasets. Due to width-independent alignment between W and Art, logits diverge with width in SP-full-align at sufficient width. We validate this claim for CIFAR-10 moderate width in Figure F.32. This introduces width- dependent training dynamics. Consequently our single-pass experiments in Figure 6 and Appendix F.7 consistently show decaying optimal learning rates in SP-full-align for both SGD and Adam on common image datasets and generated multi-index data. We also observe that the maximal stable learning rate remains remarkably indeeedendeenn as our theory would predict. This constitutes our only experiment in which the maximal stable learning rate scaling is suboptimal in deep nonlinear networks. We leave fully understanding the driving mechanism to future work.
Large output dimension may explain learning rate transfer of SP-full-align on language data. On language data, Everett et al. (2024) report learning rate transfer of SP-full-align To understand this difference to the vision settings above, we measure the individual contributions to the only source of width dependence W Aaf in this parameterization.
By empirically verifying the predicted width-independent incoming updates ||
O(1)
and alignment ratio "w Aut @(1) from (a-op) (Figures 1. and E. 18), the scaling of the
propagating updates | Art |RRS is determined by the initial operator norm
| IRMS-RMS se d 1122(2 (d + n 12 1 + (n/ (emu 1/2
under standard initialization (Vershynin, 2010). Now, in the large width regime n >> dout the second term dominates, but in the regime dou the first term dominates and induces approximately width-indeppnddnt effect on the logits (verified in Figure E.). From this perspective, standard initialization is the approximately correct initialization in the regime dout >> n, but transfer should eventually break at sufficient width n ~ dan as logits start to diverge.
Initialization for transfer at all scales. To ensure that updates remain non-asymptotically scale- preserving, the above arguments motivate choosing the last-layer initialization variance OL+1 fan - //aannaingg (n) ! which transitions from SP to μP with increasing width.
/Tan_ouu

---

MNIST
CIFAR-10
FASHION MNIST
40
Train accuracy
97.9
loss
46
84
62.0
42
82
67.5
80
10384 200
1024 4090
10384
4096
Width
Width
Width
Figure 7: JuP enables other loss functions to be competitive. Optimal training accuracy of 8-layer MLPs trained with SGD on MNIST (left), CIFAR-10 (center) and FASHION MNIST (right). In SP, feature learning is lost under MSE loss, inducing highly suboptimal performance. The layer-balanced leaming in μP allows loss functions beyond CE loss to perform well.
5.3 A scaling- view on the practical success of CE loss in deep learning
Many success stories in deep learning. from computer vision to natural language processing, use the cross-entropy loss. We propose a scaling-theoretic explanation for this practical dominance. Our results show that networks trained under CE loss allow stable optimization at significantly larger learning rates in SP than under MSE loss, which recovers feature learning at large widths and consequently improves generalization. To empirically investigate this hypothesis. we compare the performance of CE and MSE losses under both SP and jPP. Since jP admits asymptotically stable dynamics, both losses exhibit similar limiting behaviours. Thus we predict that CE loss only signif- icantly outperforms MSE loss in SP, but not in μPP. Figure 7 confirms this prediction, which suggests that MSE loss may deserve renewed consideration as a practical choice under stable parameterizations like μP, especially given its theoretical simplicity and widespread use in theoretical analyses.
Discussion and future work
On the theoretical side, we have provided the first infinite-width proxy model for finite neural networks, as they are initialized and trained in practice. On the practical side, we have seen that fifinite-width feature learning and instability predictions surprisingly predictive indicators for empirical width-scaling exponents, in particular for deep Transformers.
Understanding of the controlled divergence regime. As practical neural networks operate at the edge of the controlled divergence regime, better understanding parameterizations beyond the stable regime is paramount. We believe that many conclusions in previous work about the properties of wide networks in SP or neural tangent parameterization (Lee et al., 2020, Wenger et al., 2023) change when studying the optimal learning rate scaling at the edge of the controlled divergence regime. Since the NTK diverges in SP with 7 ( -1/2), studying this limit is subtle. However, investigating the rescaled NTK might still be useful tool better understanding this limit. While width dependence is undesirable from transfer perspective, fast memorization under logit blowup may improve learning speed. How is generalization affected? Logit blowup may partially explain overconfidence in neural networks in SP, and suggests that wide networks in jP may be more calibrated.
Numerical considerations. In this paper, we consider the regime of sufficient numerical precision. From a numerical perspective, signals that diverge fast can leave floating point range at moderate widths. Hence implementations that ensure minimal accumulation of width-dependent factors in SP akin to Blake et al. (2025) could stabilize large-scale model training in practice.
Understanding optimal learning rate exponents. The exact conditions that induce hyperparameter transfer are still poorly understood. Without full width wudth-inddpenddnee, the optimal learning rate scaling cannot be predicted with certainty, and rigorous statements about optimal LR exponents likely require strong architectural and distributional assumptions, akin to neural scaling laws (Hoffmann et al., 2022, Bachmann et al., 2023) (see Appendix F.1 for more details). Both vanishing feature learning in input-like layers and logit divergence can induce strong finite-width effects, so that we would still recommend μP learning rates over SP from a width-scaling perspective. Similar to CE loss, normalization layers correct scaling in the forward pass. In combination with Adam which stabilizes the backward pass (Figure F.25), such stabilizing components can correct most misscaled signals. Deeply understanding their interplay and effect on optimal learning rates remains an important direction for future work.
10

---

Acknowledgments and Disclosure of Funding
This work has been supported by the German Research Foundation through the Cluster of Excel- lence "Machine Learning New Perspectives for Science" (EXC 2064/1 number 390727645), The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Moritz Haas. Leena Chennuru Vankadara is supported by the Gatsby Charitable Foundation.
References
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference Learning Theory (COLT), 2023. Cited on page 19.
Atish Agarwala, Samuel Stern Schoenholz, Jeffrey Pennington, and Yann Dauphin. Temperature check: theory and practice for training models with softmax-cross-enlinn cropscsoss- losses. Transactions on Machine Learning Research (TMLR), 2023. Cited on page 18.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning (ICML), 2019. Cited on page 1.
Maksym Andriushchenkko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flam- marion. A modern look at the relationship between sharpness and generalization. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference Machine Learning (ICML), 2023a. Cited on page 19.
Maksym Andriushchenko, Aditya Vardhan Varre, Loucas Pillaud-Viviin Vivien, and Nicolas Flammarion. Sgd with large step sizes learns sparse features. In International Conference on Machine Learning (ICML), 2023b. Cited page 19.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine- Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. International Conference on Machine Learning (ICML), 2019. Cited on page 1.
Alexander Atanasov, Alexandru Meterez, James B Simon, and Cengiz Pehlevan. The optimization landscape of SGD across the feature learning strength. In The Thirteenth International Conference on Learning Representations (ICLR), 2025. Cited on page 18.
Gregor Bachmann, Sotiris Anagnostidis, and Thomas Hofmann. Scaling mlps: A tale of inductive bias. Advances in Neural Information Processing Systems (NeurIPS), 36:60821-60840, 2023. Cited on page 10.
Jeremy Bernstein and Laker Newhouse. Modular duality in deep learning. arXiv:2410.21265, 2024. Cited on page 25.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning (ICML), 2023. Cited on page 34.
Johan Bjorck, Alon Benhaim, Vishrav Chaudhary, Furu Wei, and Xia Song. Scaling optimal LR across token horizons. In The Thirteenth International Conference on Learning Representations (ICLR), 2025. Cited on page 19.
Charlie Blake, Constantin Eichenberg, Josef Dean, Lukas Balles, Luke Yuri Prince, Bjöm Deiseroth, Andres Felipe Cruz-Salinas, Carlo Luschi, Samuel Weinbach, and Douglas Orr. u-pp: The unit- scaled maximal update parametrization. The Thirteenth International Conference on Learning Representations (ICLR), 2025, Cited page 10, 20.

---

Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Advances in Neural Information Processing Systems (NeurIPS), 35: 32240-32256, 2022. Cited on page 18.
Blake Bordelon and Cengiz Pehlevan. Deep linear network training dynamics from random ini- tialization: Data, width, depth, and hyperparameter transfer. arXiv:2502.0253/ 2025. Cited on page 19.
Blake Bordelon, Hamza Tahir Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head trans- former dynamics. In The Thirty-eighth Annual Conference on Neural Information Processing Systems (NeurIPS), 2024a. Cited on page 18.
Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyper- parameter transfer in residual networks: Dynamics and scaling limit. In The Twelfth International Conference Learning Representations (ICLR), 2024b. Cited on page 18.
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. How feature learning can improve neural scaling laws. International Conference on Learning Representations (ICLR), 2025. Cited page 3.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems (NeurIPS), 13:1877-1901, 2020. Cited on page 1, 18, 52.
Yuhang Cai, Jingfeng Wu, Song Mei, Michael Lindsey, and Peter Bartlett. Large stepsize gradient descent for non-homogeneous two-layer networks: Margin improvement and fast optimization. Advances in Neural Information Processing Systems (NewrIPS), 37:71 130--7155,, 2024, Cited on page 19.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- wee-parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018 Cited on page 18.
Lénaïc Chizat and Praneeth Netrapalli. The feature speed formula: flexible approach scale hyper-parameters of deep neural networks. In The Thirty-eighth Annual Conference on Neural Information Processing Systems (NeurIPS), 2024. Cited on page 19.
Lénaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems (NeurIPS), 2019, Cited on page 18.
Lénaïc Chizat, Maria Colombo, Xavier Fernández-Real, and Alessio Figalli. Infinite-width limit of deep linear neural networks. Communications on Pure and Applied Mathematics, 77(10): 4958-4007, 2024. Cited on page 18.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung. Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 2023. Cited on page 8, 21.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations (ICLR), 2021. Cited on page 2, 5, 19.
Jeremy M Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, et al. Adaptive gradient methods at the edge of stability. arXiv:2207. 14484, 2022. Cited on page 2, 5, 19.
Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization:: The implicit bias of gradient descent at the edge of stability. In ICLR, 2023. Cited on page 19.
Francesco D' Angelo, Maksym Andriushchenko, Aditya Varre, and Nicolas Flammarion. Why do we need weight decay in modern deep learning? In The Thirty-eightth Annual Conference on Neural Information Processing Systems (NeurlPS), 2024. Cited page 20.
12

---

Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-142, 2012. Cited on page 33, 73.
Nolan Dey, Bin Claire Zhang, Lorenzo Noci, Mufan Li, Blake Bordelon, Shane Bergsma, Cengiz Pehlevan, Boris Hanin, and Joel Hestness. Don't 't be lazy: Completep enables compute-efficient deep transformers. arXiv:2505.00 0161,, 2025. Cited on page 18.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR), 2021, Cited on page
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes 2018. Cited on page 1.
over-parameterized neural networks. arXiv preprini
Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander A Alemi, Roman Novak, Peter J Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee, et al. Scaling exponents across parameterizations and optimizers. arXiv:2407.05872, 2024. Cited on page 2, 3, 4. 9, 20, 21, 25, 48, 62, 63, 65, 66, 67, 68.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv:2403.08295, 2024. Cited page 8.
Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Cardoze, George Edward Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective on training instabilities of deep learning models. In International Conference on Learning Represen- tations (ICLR), 2022. Cited page 19.
Eugene A. Golikov. Dynamically stable infinite-width limits of neural classifiers, arXiv:2006.06574, 2020. Cited on page 18.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The Ilama herd of models. arXiv:2407.21788 2024. Cited on page 1, 8.
Moritz Haas, Jin Xu, Volkan Cevher, and Leena Chennuru Vankadara. μP²; Effective sharpness aware minimization requires layerwise perturbation scaling. In The Thirty-eigtth Annual Conference on Neural Information Processing Systems (NeurIPS), 2024. Cited on page 18.
Soufiane Hayou and Greg Yang. Width and depth limits commute in residual networks. In Interna- fional Conference on Machine Learning (ICML), 2023. Cited on page 18.
Bobby He, Lorenzo Noci, Daniele Paliotta, Imanol Schlag, and Thomas Hofmann. Understanding and minimising outlier features in transformer training. In The Thirty-eighth Annual Conference on Neural Information Processing Systems (NeurIPS), 2024. Cited on page 20.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In IEEE international conference on computer vision (ICCV), 2015. Cited on page 2, 33.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. An empirical analysis of compute-optimal large language model training. Advances in Neural Information Processing Systems (NeurIPS), 35:30016-30030, 2022. Cited on page 1, 10, 19.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and gener- alization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2018. Cited on page 1, 18.
Dayal Singh Kalra and Maissam Barkeshli. Why warmup the learning rate? underlying mechanisms and improvements. Advances in Neural Information Processing Systems (NeuriPS), 37:111760- 11 1801,, 2024. Cited page 19, 60.
13

---

Dayal Singh Kalra, Tianyu He, and Maissam Barkeshli. Universal sharpness dynamics in neural network training: Fixed point analysis. edge of stability. and route chaos. In The Thirteenth International Conference Learning Representations (ICIR), 2025. Cited on page 30.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. urXiv:2001.08366 2020. Cited page 1, 19.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. Interna- tional Conference on Learning Representations (ICLR), 2017. Cited on page 19.
Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay balances learning across neural networks. In Forty-first International Conference Machine Learning (ICML), 2024. Cited on page 20.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009. Cited on page 33, 73.
Daniel Kunin, Allan Raventos, Clémentine Carla Juliette Dominé, Feng Chen, David Klindt, An- drew M Saxe, and Surya Ganguli. Get rich quick: exact solutions reveal how unbalanced ini- tializations promote rapid feature learning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. Cited on page 19, 33.
Jianfa Lai, Manyun Xu, Rui Chen, and Qian Lin. Generalization ability of wide neural networks on R. arXiv:2302.05933. 2023. Cited on page 1.
Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study, Advances in Neural Information Processing Systems (NeurIPS), 33:15156-15172, 2020, Cited on page 2, 10, 18, 19.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. wrXii:2003302218, 2020, Cited on page 2, 5, 19, 29, 30.
Jeffrey Li, Alex Fang, Georgios Smyrnis, Maor Ivgi, Matt Jordan, Samir Yitzhak Gadre, Hritik Bansal, Etash Guha, Sedrick Scott Keh, Kushal Arora, et al. Datacomp-lm: In search of the next generation of training sets for language models. Advances Neural Information Processing Systems (NeurIPS), 37:14200-144 14282, 2024. Cited on page 7, 34, 73.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. Advances in neural information processing systems (NeurlPS), 32, 2019. Cited on page 19.
Lightnnng AI. Litgpt. https://githu/: htw.://ww....mm.gggannnng//////- 73.
2023. Cited on page 34,
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations (ICLR), 2020. Cited on page 20.
Peter J. Liu, Roman Novak, Jaehoon Lee, Mitchell Wortsman, Lechao Xiao, Katie Everett, Alexan- der A. Alemi, Mark Kurzeja, Pierre Marcenac, Izzeddin Gur, Simon Kornblith, Kelvin Xu, Gamaleldin Elsayed, Ian Fischer, Jeffrey Pennington, Ben Adlam, and Jascha Sohl-Dickstein. Nanodo: A minimal transformer decoder-only language model implementation in JAX. GitHub repository, 0.1.0, 2024.
www. inneeppddiddaaooo...omo Cited on page 20,
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations (ICLR), 2019. Cited on page
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. empirical model of barch-batch training. arXiv:1812.066 2018. Cited on page 19.
14

---

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences. 115633).77655-7777 2018. Cited on page 1, 18.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems (NeurIPS), 35:27198- 2022. Cited on page 19.
Lorenzo Noci, Chuning Li, Mufan Li, Bobby He, Thomas Hofmann, Chris J Maddison, and Dan Roy. The shaped transformer: Attention models in the infinite depth-and-wide limit. Advances in Neural Information Processing Systems (NeurIPS), 36, 2024a. Cited on page 18.
Lorenzo Noci, Alexandru Meterez, Thomas Hofmann, and Antonio Orvieto. Super consistency of neural network landscapes and learning rate transfer. In The Thirty-eighth Annual Conference Neural Information Processing Systems (NeurIPS), 2024b. Cited on page 18.
OLMo Team, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bhagia, Yuling Gu, Shengyi Huang, Matt Jordan, et al. 2 olmo 2 furious. arXhv:2501.00636 2024. Cited on page 2, 8.
Elliot Paquette, Courtney Paquette. Lechao Xiao, and Jeffrey Pennington. 4+3 phases of compute- optimal neural scaling laws. In The Thirty-eightt Anmual Conference Neural Information Processing Systems (NeurIPS), 2024. Cited on page 19.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems (NeurlPS), 2019. Cited on page 33, 73.
Dan Qiao, Kaiqi Zhang, Esha Singh, Daniel Soudry, and Yu-Xiang Wang. Stable minima cannot overfit in univariate relu networks: Generalization by large step sizes. In The Thirty-eighth Annual Conference on Neural Information Processing Systems (NeurIPS), 2024. Cited on page 19.
Yunwei Ren, Eshaan Nichani, Denny Wu, and Jason Lee. Emergence and scaling laws in sgd learning of shallow neural networks. arXiv:2504. 1144/19999 19983, 2025. Cited on page 19.
Grant Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, 75(9): 1839-1935, 2022. Cited on page 1.
Christopher J Shallue, Jachoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. /rXiv:/8/1.036000 2018. Cited on page 19.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal on Applied Mathematics, 80(2):22--522 2020. Cited on page 1.
Jascha fohl-Dickstein, Roman Novak, Samuel S Schoenholz, and Jaehoon Lee. On the infinite width limit of neural networks with a standard parameterization. arXiv:2001.07301,, 2020. Cited on page 2, 18.
Teodora Srećković, Jonas Geiping, and Antonio Orvieto. Is your batch size the problem? revisiting the adam-sgd gap in language modeling. arXiv:2506./233 2025. Cited on page 20.
Alexander Tsigler, Luiz FO Chamon, Spencer Frei, and Peter I. Bartlett. Benign overfitting and the geometry of the ridge regression solution binary classification. arXiv:250.00966. 2025. Cited on page 19.
Leena Chennuru Vankadara, Jin Xu, Moritz Haas, and Volkan Cevher. On feature learning in structured state space models. In The Thihtt-eighth Annual Conference on Neural Information Processing Systems (NeurIPS), 2024. Cited on page 18, 35.
15

---

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems (NeurIPS), 2017. Cited on page 2, 34.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. 1rXiv:0011.300 2010. Cited on page 9, 45.
Nikhil Vyas, Yamini Bansal, and Preetum Nakkiran. Limitations of the ntk for understanding generalization in deep learning. arXiv preprint (tt::2200.10012, 2022. Cited page 2.
Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan. Feature-learning networks are consistent across widths at realistic scales, Advances in Neural Information Processing Systems (NeurIPS), 36, 2024. Cited on page 3, 18, 19.
Jonathan Wenger, Felix Dangel, and Agustinus Kristiadi. On the disconnect between theory and practice of overparametrized neural networks. arXiv:2310.00137 2023. Cited on page 2, 10, 18.
Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie E Everett, Alexander A Alemi, Ben Adlam, John D Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Kelvin Xu, Jachoon Lee, Justin Gilmer, and Simon Kornblith. Small-scale proxies for large-scale transformer training instabilities. In The Twelfth International Conference on Learning Representations (ICLR), 2024. Cited on page 20, 21, 34, 49,50.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization the transformer architecture. In International Conference on Machine Learning (ICML). 2020. Cited on page 20.
Greg Yang. Wide feedforward recurrent neural networks of any architecture are gaussian processes. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019. Cited on page 18.
Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks In International Conference on Machine Learning (ICML), 2021. Cited on page 2, 3, 4, 18, 22, 23, 25, 28, 31, 62.
Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the infinite-width limit. arXiv:2308.01814. 2023. Cited on page 3, 4, 18, 26, 28.
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv:2203.03466. 2022. Cited on page 18, 19, 20, 34.
Greg Yang, James B. Simon, and Jeremy Bernstein. A spectral condition for feature learning.
arXiv:2310. /7813, 2023a. Cited on page 4, 18, 22, 35.
Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs vi: Feature learning in infinite-depth neural networks. arXiv:2310.02244, 2023b. Cited on page 18.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations (ICLR), 2020. Cited on page 20.
Zixiong Yu, Songtao Tian, and Guhan Chen. Divergence of empirical neural tangent kernel in classification problems. International Conference on Learning Representations (ICLR), 2025. Cited on page 18.
Ruiqi Zhang, Jingfeng Wu, Licong Lin, and Peter L Bartlett. Minimax optimal convergence of gradient descent in logistic regression via large and adaptive stepsizes. arXiv:2504.044 04100, 2025. Cited on page 19.
16

---

Appendices
Appendix Contents.
A Detailed Related Work
18
Take-aways for practitioners
20
B.1 A practitioner-oricnted introduction to width scaling
20
B.2 Practical implications of our results
C Theoretical considerations
22
C.I Distilled TP scaling arguments
................................
22
C.2 Measuring Alignment
................................
25
C33 Formal statements and proofs of Propositions and ................................
25
C,4 Scaling dynamics 2-layer linear networks
29
D Experimental details
33
D.I MLPs
33
D.2 Multi-index data
33
D.3 Language modeling
34
D.4 Figure Details
Refined coordinate checks
34
E.1
SGD
35
E.2 Adam
41
E.3 Normalization layers and Adam provide robustness to miss- intisiinitializttinn
44
E.4
Alignment and update scaling in Transformers
45
Empirical learning rate exponents
48
F.1
Summary of the MLP experiments in this section
48
F2
Transformer experiments
49
F3
Learning rate sweeps corresponding to 8-layer MLPs
52
F.4
MSE loss with softmax does enable feature learning, but Adam does
54
F.5
SGD
57
F66
Adam stabilizes updates similar to CE loss
59
F7
Effective update parameterizations beyond juP
62
17

---

A Detailed Related Work
Here we provide more detailed account of related work than what is possible in the main body of the paper.
Strongly related, Atanasov et al. (2025) study the 2D plane of learning rate and feature learning strength, meaning the output multiplier, in uP and find related pseudo-catapult regime in CE but not in MSE loss at fixed width. They study a simplified one perapmetteer model and provide several insightful empirical evaluations in MLPs and CNNs. Their results highlight the importance of tuning the 'feature learning strength' multiplier in the output layer, which can also be interpreted as the softmax temperature (Agarwala al., 2023). As SP and NTP are width-dependent parameterizations, the feature learning strength is implicitly altered when scaling width.
Neural networks in the infinite-width limit. Past work has extensively analysed the Neural Tangent Purameterization (NTP) (Jacot et al., 2018) due to its tractability. But due to lacking feature learning in the infinite-width limit, finite networks in NTP behave qualitatively differently and hence NTP is not the ideal model for understanding finite neural networks. Finite-width deviations already accumulate after few steps of training (Wenger et al., 2023), in particular under CE loss (Yu et al., 2025). Much feature learning has been empirically shown to improve generalization over the infinite- width kernel regime Chizat al. (2019), Lee al. (2020), Agarwala et al. (2023). Considerable effort has been invested in finding a descriptive infinite-wide width model for SP. Sohl-Dickstein et al. (2020) note that the NTK diverges under large learning rates the w(n ') in SP, which motivates them to consider a different parameterization which preserves a finite NTK in the dlfnite-widht limit, but consequently does not correspond to SP anymore. Golikov (2020) studies class of "dynamically stable' parameterizations, allowing large learning rates under a variant of SP, they call 'sym-default' parameterization, which again is not equivalent SP. Another popular width-dependent parameterization is the Maximal Update Parameterization (μPP) It achieves a widtth-nndependen effect of updates in all trainable weights on the output function. Its infinite-width limit has been observed to closely track finite networks in sP well over long periods of training in, for example, feature learning strength, the learned function, or gradient and Hessian statistics (Vyas al., 2024, Noci et al., 2024b). As an important practical consequence, it allows to tune small proxy models and train the large model only once with the optimal HPs (Yang et al., 2022). μP was derived using Tensor Programs (TP) framework (Yang, 2019, Yang and Hu, 2021, Yang and Littwin, 2023) that, in theory, allows to exactly track the learning dynamics of many popular architectures like MIPss, ResNets and Transformers trained with SGD or Adam in arbitrary parameterizations in the anfinite-width limit. Haas et al. (2024) derive a pP-like parameterization for sharpness aware minimization algorithms achieving transfer of the optimal learning rate and perturbation radius jointly by showing that perturbations should be scaled like updates in j.P. Vankadara et al. (2024) derive an initialization and learning rate scaling rule that achieves width-independent training dynamics for the state-space model Mamba, which shows that the spectral condition on the weights and weight updates in every layer for achieving μP provided by Yang et al. (2023a) does not apply to arbitrary architectures. At sufficient numerical precision, the mean-field parameterization (Mei et al., 2018, Chizat and Bach, 2018) is equivalent to JPP. While it was initially restricted to shallow neural networks, the dynamical mean-field theory (DMFT) by Bordelon and Pehlevan (2022) generalizes it to more complex architectures, including Transformers (Bordelon et al., 2024a). Although still expensive, the approximate solvers from DMFT are more computationally feasible than iteratively solving the exact TP limit equations. Chizat et al. (2024) studies deep linear networks in μP and shows convergence of gradient flow to a minimum 12-norm solution.
Other neural network scaling limits. Beyond width scaling, depth scaling L + oo has been studied in detail. For ResNets, Yang al. (2023b), Hayou and Yang (2023), Bordelon et al. (2024b) show that L 1L-1/2-scaling of shallow residual blocks induces depth-indeeenndnce and this limit commutes with width scaling, implying that depth can be scaled independent of width. Using approximative DMFT theory, Bordelon et al. (2024a) suggest that L --d-depth scaling may be necessary preserve feature learning in attention blocks although they consider a pure depth limit. Dey et al. (2025) confirm - -lock scaling to be the 'correct' scaling by providing additional desiderata and empirical evidence on Transformers. Bordelon et al. (2024a) also show that the infinite within-head dimension limit effectively leads to a single-head Transformer, an the infinite number of heads limit concentrates by aggregating over the coordinate distribution at fixed within-head size, closer to how scaling is typically performed in practice (Brown et al., 2020). Noci et al. (2024a) study joint width and depth limit close to initialization for Transformers with the goal of preventing rank collapse. Long
18

---

training time is much less understood. Bordelon and Pehlevan (2025) study the training dynamies of deep and wide linear networks trained on structureless Gaussian data. Chizat and Netrapalli (2024) considers the angle between activations and gradients to give scaling rules for hyperparameters toward automatic HP scaling. They correct output layer scaling of ML.s in jP depth-dependenntl. only for SGD.
Scaling laws. Robust compute-optinal scaling laws in LLMs were reported by Kaplan al. (2020), Hoffmann et al. (2022). Paquette et al. (2024) provide theory on random feature models trained with one-pass SGD and identify 4 phases and 3 subphases depending on properties of the data and the target. Bjorck et al. (2025) observe no transfer across token horizons but a predictable scaling law with exponent -0.32 on LLama. McCandlish al. (2018) suggests that the optimal learning rate scales as a/(1 + b/balchsize) with large batch size the optimal learning rate is roughly constant, which is in line with the empirical observations by Shallue et al. (2018), Yang et al. (2022). Ren et al. (2025) study SGD training of 2-layer MLPs on isotropic Gaussian data under MSELoss and find that different teacher neurons are abruptly learned at different timescales leading to a smooth scaling law in the cumulative objective. Further work toward assessing the compute-opting optimal and data-optimal Pareto frontiers under realistic assumptions remains an important and challenging task for future work.
constants a, b. Hence for sufficiently
Finite width training dynamics. Understanding finite-width training dynamics complements infinite- width theory very well, as the former line of work operates at fixed width, while the latter ask what changes with increasing width. From practical perspective, scaling networks with μP appears to preserve the properties from base width (Vyas et al., 2024, Noci et al., 2022). Deep understanding of neural network training dynamics is still limited to 2-layer nonlinear MLP (Ren et al., 2025, Zhang et al., 2025) or (deep) linear MIP (Kunin et al., 2024, Tsigler et al., 2025) toy models under strong distributional assumptions. Lee et al. (2020) find that large learning rates cause differences between finite and infinite-width networks.
Kunin et al. (2024) explain for 2-layer networks that varying layerwise initialization variance and learning rate scaling induces differing learning regimes: fast feature learning in balanced parame- terizations (desirable for linear nets), faster learning of earlier layers in upstream parameterizations with small parameter movement (desirable in nonlinear networks, as it reduces time to grokking and sample complexity of hierarchical data structures), faster learning of later layers in downstream initializations (that initial lazy fitting followed by slow feature learning). Abbe et al. (2023) show that, opposed to lazy networks, feature learning networks can learn low rank spikes in hidden layer weights/kernels to help with sparse tasks. Qiao et al. (2024) show that large learning rates induce sparse linear spline fits in univariate gradient descent training by showing that all stable minima are flat, notrppterpollaing and produce small first order total variation, hence avoid overfitting and learn functions with bounded first order total variation
Edge of stability. Large learning rates have broadly been observed to induce optimal generalization. While the empirical connection between optimal and max-stable learning rates is well-documented, the mechanisms remain poorly understood, independent of model scale. Our findings offer novel perspective on the advantages of large learning rates for facilitating feature learning at large model scales. Previously suggested explanations include improved generalization through reduced sharpness (A (Andriushchenko et al., 2023a), shift in the learning order of patterns (Li et al., 2019), enhanced SGD noise (Keskar et al., 2017), and implicit bias towards sparsity (Andriushchenko al., 2023b).
Lewkowyez et al. (2020) observe that under large learning rates at the edge of stability, 2/Ag < " < Core/ oA (where Care = 12 for ReLU nets) an initial blowup at training time at least log(n) induces a bias towards flatter minima. Cohen et al. (2021) find loss spikes during training, but that training self-stabilizes through sharpness reduction. Damian et al. (2023) and (Cai et al., 2024) develop some understanding of the mechanisms behind EOS dynamics. For Adam, the preconditioner matrix provides additional mechanism by which stabilization occur (Cohen et al., 2022, Gilmer et al., 2022).
Warmup. Warmup allows stability under larger learning rates via slow sharpness reduction (Kalra and Barkeshli, 2024). Kalra and Barkeshli (2024) also shows that warmup allows using larger learning rates than otherwise stable by constantly operating at the edge of stability. Warmup does not improve performance but stabilizes training: by allowing training with larger learning rates, these often induce improved performance. Large catapults harm Adam by persisting in its memory. Hence Adam's ' optimal learning rate is further away from the failure boundary than for SGD, and Adam benefits
19

---

more from longer warmup, Above the optimal learning rate, Adam has a regime of training failure, where early catapults persist in the second moment and prevent learning. Warmup also widens the regime of near-optimal learning rate choices. Liu et al. (2020) find that particularly Adam needs warmup due to large initial variance.
Effective learning rates. Kosson et al. (2024) study the effect of weight decay on rotation in weight vectors, which influences the effective learning rate. Also see references therein for literature on effective learning rates, which is related to the alignment discussion in this paper.
Stability of Transformer training. More empirically, a plethora of works study the training stability of large-scale Transformers with respect to warmup, weight decay (D' Angelo al., 2024), batch size (You et al., 2020), the optimizer (Kosson et al., 2024), the position of normalization layers (Xiong al., 2020) and their interplay with the parameterization and numerical considerations (Wortsman et al., 2024, Blake et al., 2025, Everett et al., 2024). Wortsman et al. (2024) find that ga-Leoymn stabilizes Transformer training beyond the stabilizing effect from using μ.P Xiong et al. (2020) propose pre-LN for enhanced training stability requiring less warmup. He et al. (2024) observe that outlier features (=extremely activated coordinates in activations) emerge quickly in Transformer training with Adam W and that rank-collapse under strong correlations between inputs is correlated with more outlier features. Non-diagonal preconditioning like SOAP and Shampoo resolves the issue.
Srećković et al. (2025) find that SGD with momentum almost performs en par with Adam in LLMs under small batch size, so that understanding small-batch SGD (as do) is practically relevant.
Most relevant our work, Everett al. (2024) perform extensive and insightful experiments for Nan- oDO decoder-only Transformers (Liu et al., 2024) in SP, μP, NTP and mean field parameterizations with corrected layerwise learning rate scalings. questioning the infinite-width aligmment predictions between weights and incoming activations at finite width over the course of long training. They recommend SP with ADAM in conjunction with yP-learning rate scaling (they call SP-full-alige align) as the best-pert best-performing empirical parameterization in terms of generalization, learning rate transfer and learning rate sensitivity.
B Take-aways for practitioners
B.1 A practitioner-oriented introduction to width scaling
In this paper, the term parameterization refers to bidth--e scaling of initialization and learning rate of each trainable weight tensor. Studying parameterizations then means applying a scaling rule for layerwise initialization variances and learning rates and understanding how relevant quantities such update scaling in activations and logits evolves, and where instabilities may arise at large widths. At some fixed base width, all parameterizations can be considered equivalent, if we allow tuning constant multipliers.
For properly comparing the performance of parameterizations, constant weight and initialization multipliers should be tuned at some fixed base width (Yang et al., 2022). This adjusts the layerwise activation and gradient size finite width for all parameterizations at once. The parameterization then prescribes the rule, by which the layerwise initialization and updates are rescaled when changing the width in relation to that base width width/ width/base, btspb/idsewddeeeeaeeeeae width. Alternatively, multipliers can be tuned at larger width for each parameterization separately, but this quickly becomes computationally infeasbile. The extensive LLM experiments in Everett et al. (2024) suggest that the advantage of large last-layer initialization may just be an artifact of the community extensively tuning performance in SP, and after also tuning all layerwise multipliers for μP, the performance difference vanishes.
While SP performs better than naive theory would predict, and can learn hidden-layer features mndth-indepennennly under CE loss, feature learning still vanishes in input-like layers like embedding or Layemorm layers under both SGD and Adam. Still only μP learning rate scaling effectively updates all layers. For AdamW without using width-dependent weight multipliers, layer-balancing μP learning rates are simply given by the learning rate scaling (WW 7/ in(W). Here, all biases as well as normalization layer weights should be understood as weights to the one-dimensional input 1, hence fan_in = 1. For recovering width-independent weight decay, weight decay requires the inverse scaling wd- fan_ (W).
TP-like width scaling arguments are very useful for identifying sources of divergence or shrinkage with scale, and architecture components such normalization layers and training algorithms such
20

---

as Adam correct most but not all divergent or vanishing scalings in the forward and backward pass, respectively. Of particular importance for evaluating the deetndde signal propagation is the refined coordinate check (RCC) for disentangling effective updates in the current layer from updates propagating forward through the network. Ideally, all W and AWr should remain width-independent, which is only guaranteed in μP at sufficient width.
B.2 Practical implications of our results
Constraining the search space for the optimal learning rate. Our experiments strongly support the prediction from Proposition that the maximal stable learning rate scales as Th2 17 n when training deep non-linear networks with SGD in SP under CE loss. This result provides a concrete constraint on the learning rate search space, significantly reducing the range practitioners must consider. Rather than exhaustive or heuristic searches, practitioners may narrow their search around the theoretically justified scaling, leading to substantial computational savings and more efficient hyperparameter tuning. Furthermore, we observe that, particularly in deep nonlinear architectures, the optimal learning rate often saturates at this maximal stable scaling (see e.g. Figures 5 and F.10, Figures F.4 and F.5, Appendix F.3). Thus, employing the correct scaling often effectively enables approximate transfer of optimal learning rates, even in SP.
Potential performance gains at scale using alternative loss functions with J.P. Correctly attributing the observed performance gap between MSE and CE loss in SP to scaling considerations reveals concrete recommendation: Practitioners could leverage stable parameterizations like jaP for image datasets SP-full-alige for language datasets to effectively utilize loss functions beyond cross-entropy at large scales. In Figure 7, we consistently observe that CE loss greatly outperforms MSE under SP, whereas MSE consistently outperforms CE loss under μ.P, with differences becoming particularly pronounced at large widths. These results highlight the potential for substantial performance gains from further investigating the interplay between the loss function and parameterization. They also suggest that it is worth exploring the use of further loss functions in conjunction with μcP.
Identifying the correct scaling mechanisms enables principled improvements of scaling practice.
Identifying the correct causal mechanism for training stability under large learning rates enables principled and informed exploration of the extended search space of practically relevant potentially best-performing parameterizations for efficiently finding improvements in model scaling practice: The controlled divergence regime had previously been neglected, even though common scaling practice as well as the best-performing parameterization fullull-aiggn from Everett et al. (2024) exactly operate in this regime. We now detail two concrete implications together with exciting avenues for future work.
Identifying the correct mechanism for training instability at large scales. Training instability due to logit divergence at large scales is a widely-established empirical phenomenon under SP which has motivated popular interventions like the z-loss (Chowdhery et al., 2023, Wortsman et al., 2024). However, only understanding the underlying causal mechanism and providing the exact width-dependent scaling exponents enables design principled interventions. Both our theoretical and empirical results suggest that practically relevant large-scale neural networks naturally operate at the boundary of the controlled divergence regime, allowing persistent hidden-layer feature learning despite logit divergence. Consequently scaling up networks in SP inherently causes logit divergence, leading to numerical instability. As principled intervention, stable parameterizations like jP should not suffer from such systematic divergence and potentially eliminate the need for ad-hoc interventions like the z-loss.
Logit divergence induces overconfidence. A similar argument applies to uncertainty calibration. Our theory also explains why we should expect predictions to be increasingly overconfident with increasing model scale and suggests that this may be partially mitigated by considering stable parameterizations like pP. However there may be inherent trade-off's between miss-calibration and faster memorization due to logit divergence. Their effect on performance and designing the ideal intervention deserves more thorough investigation beyond the scope of the current paper.
Understanding the transfer of SP-full-alig motivates a novel last-layer initialization. Small last-layer initialization recovers width-indeppndenn and hence predictable scaling dynamics un- der sufficient precision in the regime n >> duat whereas standard last-layer initialization induces logit blowup at sufficient width, which is not necessarily harmful for generalization but reduces
21

---

predictability as scaling is not fully irdth-independent. Standard initialization with μP learning rates (SP-full-align) can induce 'practical transfer' and empirically update all weights effectively without logit blowup at moderate scales n < dom in the regime where the width is much smaller than the output dimension, as is relevant for NL.P settings, but likely exhibits unexpected changed behaviour at sufficient scale, when logits start to diverge due to the last-layer term W Aay This can be read off from differing dominating terms in |[W || independent alignment Ow!++" Azt (1) (as verified in Figure E8) For uniform non-asymptotic transfer for both n. > dout and n & dout, the same argument motivates a last-layer initialization that transitions from SP initialization in the regime n < dout to uP initialization in the regime n doat: similar to the one proposed in Yang et al. (2023a). Verifying improved transfer under this initialization is left to future work, since it requires scaling to widths beyond our computational constraints.
assuming width-
0L+1 = 'aaan_out + Vfan_in)
fan_in
Faithful evaluation of the accuracy of infinite-width theory. While in principle the refined coordinate check (RCC) was previously known (Yang and Hu, 2021, Appendix H), we find that the predicted update exponents are surprisingly accurate in this decomposition already at moderate width and over the course of training (Figures 2 and 4), even in parameterizations with width-depender dynamies like SP. Consequently, optimal and maximal stable learning rate exponents in SP also follow the predicted width-dependent exponent in realistic settings (see e.g. Figures 5 and F.10, Figures F.4 and F.5). To lower the hurdle for practitioners to incorporate the (RCC) in their workflow as a diagnostic tool, we have made our fine-grainet and easily adaptable implementation of the wwwbinnenntrrcc-mmollee---
RCC using LitGPT publicly available at https://gitub..om//n Understanding and correcting miss-scaled update signals through each weight matrix has impactful consequences on performance and trainability at large scale.
C Theoretical considerations
C.1 Distilled TP scaling arguments
Here we aim to provide a more detailed, comprehensive introduction to the essential width-scaling arguments inspired by Tensor Program (TP) theory.
Effective Updates. When training neural networks, we have control over the initial scaling Wo and update scaling AW, of trainable weights W.- = Wo + AW,, but we are ultimately interested in their effect on the activations in the following layers. In standard architectures (including convolutional networks and Transformers), weights typically act linearly on the incoming activations. For such weights W2 and incoming activations 2t, we can decompose the next layer's (pre-)ace pare-vactivation updates Aht into effective updates of We and activation updates Ar, propagating forward from previous layers. Evaluating the contributions of both terms separately yields a refined coordinate check,
sh.- = (AW;) + Wo(Ax)).
(RCC)
Note that updates of previous layers can propagate forward through the term Wo even when the current layer's 's effect on the output vanishes AW,,; -> 0 as width n - oc. Hence, we say that the weight W; is effectively updated only if AWar contributes non-vanishingly. Plotting the width-dependen scaling of || ((AW,) A. RMS and | (A RMS as a refined coordinate check, has been very useful for us to gain insights into the network internal signal propagation. The usefulness of (RCC) for effective update scalings is illustrated in Figure C. While the activations and activation updates in a Layernorm layer evolve width-independenn
when training GPT with Adam in SP and global learning rate scaled as Th = 17 n-1 the refined coordinate check reveals that the effective updates in the current (input-like) layer and the activation update scaling instead stems from effective updates propagating forward from previous (hidden-like like) layers.
By choosing layerwise initialization variances and learning rates according to the Maximal Update Parameterization (nP), both terms in (RCC) become width-independent in all layers in each update step. Consequently, width-scaling becomes predictable, stable and feature learning is preserved even at large width. Starting from SP, μP can be realized with smaller last-layer initialization [|W/++1| |RRMS = O(n 11, larger input layer learning rate 7/w1 = 77 n and smaller last-layer learning rate 7W1+1 17 for SGD.
22

---

Predicting scaling exponents. While the TP framework formally requires writing out all forward and backward pass computations performed during training and provides the exact unfinite-widde limit objects of output logits and activation coordinate distributions, we simplify its implications on width-scaling exponents for practical purposes as follows. A linear transformation either maps fixed to width-scaling dimension (input-like), width-scaling to width-scaling (hidden-like) or width-scaling to fixed dimension (output-like). Here, all bias vectors and normalization layer weights can be understood as input-like weights to the one-dimensional input 1. Any sum of length n + oo that occurs in individual terms in (RCC) either accumulates a factor n 1/2 under sufficient independence of O-mean summands (CLTT-like behaviour) or a factor n when the summands are correlated or have non-zero mean (L..NIlike behaviour). Crucially, not any sum may be evaluated with this heuristic but only weight and activation (update) pairs as in (RCC) (see Yang and Hu (2021, Appendix H)). If, for example, we considered the confounded term (Wo + AW the initial part Wozo clearly scales CLT-like but AW, scales LLN-like; evaluating the scaling of their sum might result in wrong scaling predictions.
At sufficient width, all width-scaling inner products (4) from (RCC), however, are expected to behave LLN-like, that is || RMS = O(n | || RMS).
Add
Theiing Seo
With
WeB
Figure C.I: (Tracking effective updates requires refined coordinate check) Activation norm [Waaall RMMS, activation update norm (W propagating update norm [[a, laarss and effective update norm | RMS for the last normalization layer in GPT trained with Adam and learning rate scaling 7,, = 0.01 n for width-independent hidden layer feature learning. While activations and activation updates appear width-independent due to propagating updates, our refined coordinate check (RCC) reveals that Layernorm weight updates have vanishing effect in SP. Over time, effective updates accumulate effective rank, but do not lose alignment with width (Figure 2).
Concrete examples. Complementing the more generic introduction to TP scaling arguments above, we now provide more concrete examples for illustrating how weight updates affect activations in subsequent forward passes. Consider the Tensor Program for training MLPs with SGD from Yang and Hu (2021, Appendix H). We restate the relevant update scalings when using large learning rates in SP that induce output blowup. Since divergence is not allowed in the TP framework, it does not formally cover the unstable case, but we can still heuristically write down the scaling predictions, assuming that correlations still induce L.N--like exponents and independence still induces CL.T-like exponents, as we have measured to hold empirically. The crucial insight is that training with cross-entropy loss effectively means that we are considering fe(E) = (W..+11p =((W=++1/(EE) as the output function and the loss derivative also becomes Xt == 8C. = a (W,++10++
yt.. Hence, from a stability point of view,
af
we can allow f. := W1++ x² > oo, which results in a saturated softmax. Under one-hot labels y {0, 1] with E. Ve 1, this means fast memorization of training points (z). For width- independent hidden-layer feature learning, we may still require activations to have width- indth-independent coordinate-scail scaling,, but let the output function be arbitrary, since the softmax renormalizes.
Definition C.1 (Activation stability). A parameterization activation stable iff | = (L) for all times t 2 0 and all layers 1 E [Z].
We now show heuristically that MLPs trained with SGD in SP are activation stable and feature learning under global learning rate scaling 7 = O(n -1/2).
Backward pass. Here we denote the entry-wise scaling in width-scaling vectors as v O(n)), meaning [| = O(n"). Assuming || [k) = O(1) as for ReLU (otherwise we would get vanishing gradients), the entries of the following width-scaling vectors scale
0f 0x.
af af
0f
w = W - AW = O(n Onn-1/3). -1/3),
o o = 0
)
0h
dx
Oz
23

---

0f 0zt
(Wh)² T 0f ah!
)T
8f oh!
of Ohg
of
1/0w
X.
1-1
O (max
n
)2 1 - O(n (12). Dh!
==0
Note that any larger learning rate scaling would induce exploding gradients. For example, 1j (1) induces 6W (1), so Øf e (1) and /
ăuद On )z ( ((n²k-1) for k 2 This
of
results in exploding activations in the next forward pass, and even larger gradients in the following backward pass.
We therefore continue with 1) = ((()) and get the activation updates
0f (L: 1) T = On 1/2 n 1/2 1) 0
8h
- -IX 8
8h
W W2_dox 11 + &
0ʃ
Whook n/ww Σx= Xo-1
(2) 1-1 ( Took
Oh
((1)
O(n) 1/2
0f
-7xt 10--101m fl
(-
zt O1).
ah
((n)
8/n-1/2)
The output updates are (E)
SW 1+1 (7) + (W + (o
IX x 11-1-/10
(E)
W ( + AW = 0
(())
9(1)
((2
(E)
off +f.1+ ofr-1) = O(1)
2 layer networks. Observe that in 2 layer nets, there are no hidden layers, so that a larger learning rate can be chosen. Let r7 O(n)). Then in the first step, oh (1) Of Onn' ah as W L+1 = (nn)) so that
But note that the gradient scaling may grow after the first step,
def
8hg = (n²nn) Hence activation stability requires 1j = ((1), which results in feature learning after 2 steps _z1 = O11. Then f. = ((n(=/_))T
= (n)).
Random feature models. In random feature models, we only train the last layer and keep all other weights fixed W Wi for all 7 < T.. There, by definition, we do not get feature learning and the backward pass does not matter. The only gradient that matters is the last-layer gradient which has fixed scaling + 0xx--1af_1 1) = @(1) at all times / 2 0. The function update becomes 5W-+1y (E) =
-1 (xx-1 x (x) Onn"n). where the inner product between activations converges to the NNGP kernel in the infinite-width limit. Hence large learning rates 1 w(n -1 result in immediate extreme memorization of the training points fi(f:-1) - one- hot( as n + co, and f] = O results in fully indth-independent dynamics.
Adam. Adam with small enough normalizes the gradients each layer before updating the weights. of Since the gradients V W = aRt (T' are generally correlated with the incoming activations x'-1, their inner product accumulates @(fan_ @ffan_in). Non-vanishing correlation persists when only recovering the signs of the gradient. Hence for width-inneependett effect on the output of the current layer, the learning rate should always be chosen as ry(W)= Faa_in(WT Since both hidden and output
layers have fan_in = n, activation stability requires a global learning rate r) On--)). which results in effective hidden and output layer learning, but vanishing input layer updates. Networks recover input layer feature learning under 1 = @(1), where f. (n). In random feature models, just determines the extremeness of memorization of the training labels, where rj = O(n induces width- aatt-inndependenee and w1n increasing memorization.
24

---

C.2
Measuring Alignment
Everett et al. (2024, Fig. 2) provides RMS-al AMM-aalinmmnt exponents between weights W; and incoming
activations ztz But only measuring the alignment between A and &, as well as Wo and Ar, from (RCC) separately allows to evaluate the width-scaling predictions from Yang and Hu (2021).
For example hidden layers in yP scale as (W)), = On (1221/2 at initialization, as 0-mean indepen- dence induces CLT-llke scaling
Wo On [Woll RAVS 11=0- ([M)).
But updates are correlated with incoming activations, so that
AW, O ||W, |
which necessitates || -(n This implies that the entry size of W, Wo + AW, is dominated by the initialization and confounds |[WelRRMs for accurately measuring the alignment exponent of the layer's 's updates AW, For correct width-scaling of the layer's 's learning rate. the influence of Wo is irrelevant so that the joint alignment between W, and Iz does not reveal the alignment exponent that is relevant for correct learning rate scaling.
Additionally, replacing the RMS-norn ||A|RMs by the operator norm ||A]] provides a more natural measure alignment (Bernstein and Newhouse, 2024), since the RMS-norm is confounded by accumulated rank whereas under maximal alignment for the operator norm it holds that A (x WW RM = alignment. Under perfect alignment we expect the ratio I. - independent. We are not interested in constant prefactors, but only width-dependent scaling.
| JMs, and the left-hand side is smaller under less JAW,,]lRME
to remain width-
C.3 Formal statements and proof's of Propositions and 2
Before providing the full formal statements of Proposition 2 and Proposition 4, we formally introduce all definitions and assumptions.
C.3.1 Definitions
In this section, we collect all definitions that do not appear in the main text. We adopt all definitions from Yang and Hu (2021), up to minor modifications. If not stated otherwise, limits are taken with respect to width n + 00.
Definition C.2 (Big-O Notation). Given a sequence of scalar random variables = fcr E [] we write
(n (0) if there exist constants A, B > 0 such that for almost every instantiation of n=1 for n large enough, An- < |c,. < Bn Given sequence of random vectors
c = {c,, E R 2o
x [ E IR" J==1" we say x has coordinates of size θ (n and write x θ (n -9) to mean the
scalar random variable sequence {
{VI=Sal /n is θ (n "). For the definition of = O(n
and c = O(n ") adapt the above definition of c = O(n ") by replacing An - |col S Bn
with |c,l S Bn a and An-a < lc,// respectively. We write In o(n ") if na v | 11 /nn
almost surely.
Definition C.3 (SGD update rule). Given (L+1) MLP with layerwise initialization variances
(athe(t++) and (potentially) layerwise learning rates {nw* he(r-+)) we define the SGD update rule as follows:
(a) Initialize weights iid (W)).z (O,
(b) Update the weights via
W! = W! Tw Vw L (fe
(E.)
Definition C.4 (Parameterization) We define width-scaling parameterization as collection of
exponents (br U (cthe that determine layerwise initialization variances of = C n -b
and layerwise learning rates 1/I 1) a, with width-independent constants Cr, y for all
I t [L +1.
25

---

Weight-mulliplier version
Weight-multiplier- free version Ouiput-like
Input-like Hidden-like Output-like
Input-like
Hidden-like
01 W. of),
Oy
SP
07 x
11 1/2 n 1/2
Th Virc or W
1)) X
72
n
Oy
1
n-1/2
n-1
NTP
of).
07 x
1
1
y1 1/2
1/2
Th VwiC.
1]i xc
1
n2"
ar W of)
ay
1/2
1/2
072 X
n-1/2
n
n 1/2
n 1/2
1/2
Th Vir.. 1/ OC
1
Table C.I: (Common ade-parramtterizaionn) Here, we collect standard parameterization (SP), neural tangent parameterization (NTP) and the maximal update parameterization (nP) for SGD in their multiplier version which purely adapts the architecture and allows width-independee global learning rates (left) and in their weight multiplier- free version (right). Parameterizations differ in their layerwise choice of width-depent -idth-dependent weight multipliers a,, initialization variances or and learning rates '/- Weight multiplier-fres representatives of an abe-equivalence class purely adapt the optimization algorithm highlighting the fact that parameterizations effectively only induce layerwise learning rates. Knowing that uP correctly scales the updates in all layers, observe that the input- and hidden-layer learning rates in NTP induce vanishing updates. The same holds in SP when choosing e 2 1 as is necessary for avoiding logit blowup in the infinite-width limit.
In Table C.1, we summarize the three most common parameterizationn SP, NTP and μP (equivalent to the mean-field parameterization).
Definition C.5 (Training routine). training routine a combination of base learning rate 1j 2 0, training sequence heN and continuously differentiable loss function C(f(E), y) using the SGD update rule.
Definition C.6 (Stability). We say parametrization of a (I. + 1)-layer MLP is stable if
1. For every nonzero input & E Vd..
Ng- =0 (1),, AI E [L], and (oll))==0000
=
where the expectation is taken over the random initialization.
2. For any training routine, any time & E N, / E [4]. E 1dd.., we have
RE(E) - I6( (E) - =5((E) = O1(1). and S((E) = O1(),
where the hidden constant in O., can depend on the training routine, t, E, I and the initial function fo-
Definition C.7 (Nontriviality). We say a parametrization is trivial if for every training routine, fe(E) - fo(E) - 0 almost surely for n oc, for every time and input & E Rdd.. Otherwise the parametrization is nontrivial.
Definition C.8 (Feature learning). We say a parametrization admits feature learning in the I-th layer if there exists training routine, a time & > 0 and input & such that xi(E) - =b(()) = Ω2 (1), where the constant may depend on the training routine, the time f, the input & and the initial function fo but not on the width n.
Definition C.9 (o-gelu). Define o-gelu to be the function
(1 + erf (o-1z)) to
2/a
In order to apply the Tensor Program Master Theorem, all Nonlin and Moment operations in the NE@ORT program (Yang and Littwin, 2023), which do not only contain parameters as inputs, are required to be pseudo-Lipschitz in all of their arguments. For training with SGD, this is fulfilled as soon as ø' is ggeudo-Lipschite a-gelu fulfills this assumption.
Definition C.10 (Pseudo-Lipschitz). A function f R* + IR is called aseudo-Lipschir of degree d
there exists a >> 0 such that [f(x) - f(y)] S Cll - v1l(1 + Ex= |o,] + yil ).ld). We say f is
pseudo-Lipschitz if it is so for any degree d.
26

---

Finally, we differentiate relevant notions of maximal stable learning rates. Definition C.11 (Maximal stable learning rates). Fix & and consider any layerwise learning rate scaling rule 77, T/o n n for I E [L + 1] with prescribed width-scaling exponents {cr (e)))(r)+)) and tunable global learning rate 7/2 (n=")
scale width +
(a) Theoretical max-stable learning rate of network. We define the theoretical maximal stable learning rate exponent ā of the network as the boundary to the catastrophic regime, that is there exists training routine and Ç E gdd such that for all a
training with T/e = -) implies [[fel RMS + oc and | RMS + oc a.s. after i steps for all | E [L +
(b) Theoretical max-stable learning rate of layer. We define the theoretical maximal stable learning rate exponent âi W.Z.f layer / as the global learning rate scaling above which the layer's effective updates diverge, that is, given stable inputs ||=]|
(g(1) with maximal alignment Oaw, (1)) there exists a training routine such that for all < or training with a global learning rate T/, 17 - n -a implies | WlaaS + 0o a.s.
(c) Empirical max-stable learning rate. We softly define the empirical maximal stable learning rute exponent ā as the global learning rate exponent above which training diverges and the network's 's predictions degrade to random guessing.
C.3.2 Full formal statements of Propositions 1 and 2
Assumptions. For all of the results in this section, we assume that the used activation function is a-gelu for a > 0 sufficiently small. For small enough a > 0, a-gelu (Definition C.9) approximates ReLU arbitrarily well. We assume constant training time > 1 as width n. co. We assume batch size 1 for clarity, but our results can be extended without further complications to arbitrary fixed batch size.
Proposition C.12. (Asymptotic regimes in SP) For fixed L, 2, 0, a E TR, consider training a (I. + 1)-layer MLPP of width n in SP with SGD and global learning rate On = =n.n n for t steps. Then the logits fú, training loss ((fe(E)), 2t)) loss-logit derivatives X [f vc). loss-weight gradients V : WWi (,We),,))) We) and activations xt, I E [L]. after training scale as follows in the infinite-width limit n + 0o. The hidden constants in O., lt and w. below can depend on the training routine, t. E, L and the initial function fo-
Under cross-entropy (CE) loss, three qualitatively distinct regimes arise:
(a) Stable regime (a 2 1): For any training routine, all E [L] and any & Rdd,,, holds (1),, lxeel O.(1), ||
that ()) nass (1), [
(.(n (12) and (l = (1).
(b) Controlled divergence (7 S 0: < 1): For any training routine, all E [L] and any
& E Rdd., , it holds that ||na-1, fe(E)ll = =.(1), [=((E) - =xl(E)]l = O.(1),
[(fe(()). = (1), lxxellas there exists a training routine and input such that ||n f((E)) RMS = (1).
O1(1) and V VillaMs = On 12). In addition,
(c) Catastrophic instability (a < 2)): For any 1 E [L]. there exists training routine and
& E IR IRdi,, such that [[fx(E)] RMS (1), (E)) RMS (1) and [|V/] RMS (1).
Under mean-squared error (MSE) loss, stable regime in (a) above arises if 1. 1f a 1, training is catastrophically unstable as in (c) above and, in addition, there exists a training routine such that (C()(E))) wt)1 (1) and l|x|| RMS = (1).
Proposition C.13 (Under CE loss, SP with large learning rates learns features at large width). Consider the setting of Proposition of training a (L+1)-layer MLP with SGD in SP with global learning rate Th 1) naa a E R. in the infinitewwidth limit + co. The hidden constants in O., Ωt and w. below can depend on the training routine, I, E, and the initial fanction fo-
(a) Under both MSE and CE loss in the stable regime (a 2 for any training routine, I E [L] and É E gdin it holds that || [ = On 112)..
(b) Under CE loss in the controlled divergence regime (2 S < 1). for any training routine,
Le [L] and & E Bdm it holds that |/ (E)
O. (n- 11/2-a)), and |Ax (E) RMS
O, (n!/2-0). For any I E [L], there exists a training routine and & E Rd such that
(E)lRMMS = Qh (n-1/2-a), and (E)l Ωr (2)
27

---

Remark C.14 (2-layer networks recover stable training dynamics and width-independent feature learning at a 0). Similarly, it can be shown that 2-layer MLPs remain activation stable under width- indth-independent learning rate scaling ?h, ((1). The controlled divergence regime is given by 0 S a < 1/2, with width-independent input layer feature learning at a 0. Experimental evidence in Appendix F.4 supports this maximal stable learning rate scaling prediction.
Remark C.15 (Adam recovers stable training dynamics and width-independdnt hidden-layer feature learning at a = 1). For Adam and L 2 2. an analogous NE@OR -based proof (Yang and Littwin, 2023) would show that 7 = (n= recovers feature learning all hidden layers E [2, L., stable activations and loss-logit gradients, while logits blow up only through W Aat = 0(n 1/2. To avoid logit blowup, 7n = O(n -3/2) would be necessary. In that case, only the term W "A would contribute aonvaniihingly to the logit updates. Hence, for Adam under CE loss, the controlled divergence regime is given by v a < 3/2, with hidden-layer feature learning at a = 1.
C.3.3 Proof of Propositions and 2
The proof in Yang and Hu (2021) for general stable tbceaaramtterisaiions directly covers the stable regimes of both losses, showing a kernel regime and vanishing feature learning for 2 1.
For the controlled divergence regime under CE loss, however, note that the TP framework does not allow diverging computations. Here, we need to replace the logit updates by rescaled logit updates, before computing the softmax limit outside of the TP framework.
Formally, under standard initialization, W ~ NO0, 1/n) is replaced in the TP by W con- structed via Nonlin, conditioning on fo(E) (see Yang and Hu (2021, Appendix H) for all details). For stable parameterizations, the function updates are defined in the TP as
SWL+1sb
W
ot + +
8fc- = 1
n In the controlled divergence regime < 1, we now
where 0/,+1 n--a and 0/,1 n b141
define rescaled logit updates in the TP as
8fc = 01+1 SWL+11/ 1/ + HOs
W 8ut
[6=0011
n
n
by replacing 4/ by 0 == 0.0/ and replacing f/,, by θLf == BabLy, where 0g == n- The adapted pre-factors ensure that 5f remains (1) for well-defined TP. The TP master theorem now implies almost sure convergence of the rescaled logit updates of. + fft E R a.s.
Now we compute the softmax limit outside of the TP framework, as we want to recover the softmax values of the original diverging logits. Thus, given the convergent sequence ff. + df. E &d a.s., due to the smoothness and saturation properties of the softmax it follows that there exists a X Rdo such that a a0=² ff.)) - 3/+ - X a.s. Since [00 5f:)
yel + |yr| and I&el < 1 [uel. this sequence can again be used as a TP scalar. Now the last-layer weights are TP vectors updated with SW!++1 = Tn X which do not change the scaling of W = W + 00+11/5 -W, with 0/+1// I 1 as long as a 2 1/2. Thus the backward pass scalings are not affected and the rest of the TP can remain unchanged.
For larger learning rates a < 1/2 under CE loss, we provide heuristic scaling arguments. Observe afg
that preactivations diverge after the first update step 5hî = "mh ah (=d) x = (nn The 0
updates of the next hidden layer's preactivations scale even larger, that is 8h, = -mn oh (=²) of =
(122(122a) In this way, the exponent growth continues throughout the forward pass. But even if there is only a single hidden layer, the scaling of the backpropagated gradient is increased after or = Wn Tnxo0 - x = 22 ) = ((m1/2-20) = w(n at사 (x²) 11 = N(n df
the second step, 0iz
This, in turn, increases the preactivation update scaling 5hद - 7/o . ((13(122-a)) which in turn increases the gradient scaling in the next step, inducing a feedback
loop of cascading exponents between diverging activations and gradients, inducing fast training divergence.
Under MSELoss, observe how, already for < 1, diverging logits ōff W.[++ 5z (xx) 2/6 O(n increase the gradient scaling through x1 = fi y1 O(nt-a) which in
28

---

turn increases the activation as well as logit scaling in the next step, and induces a divergent feedback loop even worse than above.
C.4 Scaling dynamics in 2-layer linear networks
Here, we rederive the training dynamics of the minimal model from Lewkowyez et al. (2020) that shows initial catapult mechanism in NTP. They observe that the training dynamics of repeatedly updating a 2-layer linear network in NTP on the same training point is fully captured by update equations of the current function output and the current sharpness X
C.4.1 Deriving the update equations for SP, NTP and jP
NTP. The original model by Lewkowycz et al. (2020) given by
f = n - cux,
where u E Raxd v E IR" are initialized as uzj vs (O, 1) and trained with MSE loss L(J, (f(x)) - 2)², loss derivative X [x) - y and global learning rate r/-
Repeated gradient descent updates using (r, y). then results in the update equations,
J++1 A+11 =
f&(1 + n 1 x/xx 11=11 - mxxdee A + n 1x
4fc),
where the update kernel is defined as
((r,z')= = ([]u]] +1111)))
Note that the width-dependennc in ft and X results in qualitatively different behaviour in the infinite- width limit, In particular, in the limit, the sharpness cannot evolve over the course of training, X₂ Xg-
Maximal Update Parameterization. We define a 2-layer linear network in pP with arhitrary weight multipliers
f = eūx,
with reparameterization-iinue invaaaiiiiiiiniiinnnnna weights ūij 1/d6) and üs N(0, 1/n²), th u. u, and the original weights rates T/u = 1yne 1-2a. and Th = 1pm -142a, which results reparameterization invoorianni layerwise learning rates fh 7p and 77 11-
trained with MSE loss and layerwise learning
and but can work with ū and ö instead. For gradients holds that of 8 Au df nu this width scaling has be accounted for when transitioning between representatives of the yP equivalence class. For updates Th, Th should be used instead of Mas Tus the layerwise learning rate rescaling exactly chosen cancel out the effect of the weight rescaling.
Formally,
perform updates
0f du
0f
0f
Ue+1 0e = -n ou Th Du Du
n
20 Ma
1/1 Oū.
du
The derivatives for backpropagation given by,
aL = f(x)) ye,
0f = dv
0f
Xt = 0f
x T TT,
8 x
Oū
The updated weights are then given by
1 0, Tyn -1 'xx u
Me+1 0t I TT.
In the case din = 1, the updated function is then given by
f+1== B4114112 = ft + 77 x2
ग. x
x
xT nxx TTaa ur
fi(1 + x22211=113) nx. (Λ||a||2 + -n'n==1l''''
*| ||
29

---

fi(1 + + 0z)
where we call ① the reparamettrizzationiimee invariaan update kernel defined as
"I (z) 0f(x') ow aw! = =n -/n-taa]] 2|ll n
(xx,z') = Σ
The update kernel evolves via the reparameterization-nnaaatt invvaiaiti update equation
Ae+1
( x) = ||>|| [. + n || || + 1ma? 00
+ 11+
Enij.xtua + n Tu ūz 2n -1 Tu 0ūx
x + (r' (m2 1 ||||
+ =1| ) 411x
x + || (r²x 4nxxfi)
x, + m x (xxede 4ft))
Now note that, even under fa = 0. we get non-trivial, width-independdnt dynamics. Due to the LLN, initialization, we have n n ||ūal | and n [oll z 1 (=n times sum over n iid x² variables), hence Ào ~ 2.
To conclude, the training dynamics for repeatedly updating with the same training point (ar, y) are fully described by the update equations,
f41 =
fe - x²x2l|=(1)) mxdde.
(C.1)
A411 =
x + | (Txcde - 4fr).
(C.2)
This can be rewritten in terms of the error (or function-loss derivative under MSE loss) Xt St --y- y, akin to Kalra et al. (2025), as
X++1 =
X (1 - nd + +17 1 xe + t)), de + ||=|1 (TXed== - 4(xe + v))
(C.3)
A411
(C.4)
First observe that all terms in the update equations become width-independent in pP. Only the initial conditions are width-dependent with vanishing variance, fo (/2=1/2). As opposed to NTP, the sharpness update term v x22 |ax]| is not vanishing anymore. While Lewkowyez et al. (2020) simply use labels y 0, non-trivial dynamics in jP require y * fo -> O.
Importantly, ')x always appear jointly, so that interpolation effectively reduces the learning rate.
Remark C.16 (Characterizing sharpness increase: Critical threshold depends the labels.). When both sharpness and the loss increase, then training diverges as the learning rate lies even further from its edge of stability. In μP, since fo + 0, A, will grow the first step. For subsequent steps, the sharpness update equation (C.4) implies that sharpness increases (Ae+1 2 Ar) if and only if
A 2 4 (1 + x) 1 fz. Kalra et al. (2025) provides a more extensive analysis of the dynamics and fixed points of this model in jaP.
x
Remark C.17 (Weight multipliers). A natural choice of weight multipliers for JPP can be considered be a; 1/2 1([ L+ 1) 1/2 1(1 1). as this choice allows using width-indeepnnent global learning rate
nt. and the update kernel does not require width-dependent scaling factors, 0(r,x') = x [I|=||²²+ + )a' In other words, under these weight multipliers, width-independence in parameter space translates into hidth-independdece in function space.
Standard Parameterization. We define training 2-layer linear network in SP with global learning rate scaling n -
f Tüx
30

---

with initialization ū ~ ((0, 1/dm),, ((, 1/n) and global learning rate 7u = f mj : Parameter multipliers affect all scalings in the same way for p... Only the learning rate has different prefactor, and the last layer has larger initialization. The adapted update equations become
fu+1=
fr(1 + 2 n2 x1==)))))) n' m d + 1=11111 " "mxede 4n 4n Se),
A++1
where we define, as for NTP.
6(r,z))=n =n (
where n - ||00| R n=1 at initialization (n times sum over n iid x² -ariibbal with positive mean).
Choosing c < 1 results in output blowup of the term n n/- While this can in principle be counteracted by shrinking x at finite width, a well-defined stable and non-trivial infinite-width limit is only attained at c c= 1, where fe+1- = ft- x and A++1 X We now show that also at finite width, stable training with a constant learning rate in SP requires = (n
C.4.2 Finding the maximal stable learning rate scaling by characterizing the conditions for loss and sharpness decrease
The following proposition characterizes the choices of 77 that result in decrease in loss at any present state.
{ n, in SP: n, in NTP,
Writing Map
Kntp
1, else, we can write the update equations of parameter-
1, else,
izations jointly as
xi(1 nop + "x|=="x1===1"" + y)). A + 7/x (xxe|a=|13/1 ma (mxede - 4m (xt + y)).
X++1
Ax+1
Proposition C.18 (Characterizing loss decrease in SP and NTP). Let nj 2 O. For Dap or Thaap large enough, we write the update equations of repeatedly updating the wo-model with SGD on the training point (x, y) with | = in SP or NTP jointly as provided above. The loss decreases at any step, omitting time t,
1. in the case f(f -)) 2 O, if and only in S
2 nay + (1² or
PopBedgd 2 a
xx
(n (n-pnn71) mndg)) Rap
||=]|px]
2. in the case, f(f y) O, and only if S Aapd - (n. 1
holds that A4+1 2 A if and only if X Z Map'le (1+ 111 M)
Remark C.19 (Instability in SP). The crucial insight from Proposition C.18 for SP that both loss and sharpness increase early in training as soon as w(n unless an extensively large learning rate that depends on the current sharpness, training point and output function, is accurately chosen at each time step in slim interval of benign large learning rates, which is unlikely to hold in practice. Figure C.2 shows that in simulated training with constant learning rates, the maximal stable learning rate indeed scales as (n--1). This instability prediction is in line with the infinite-width prediction from Yang and Hu 1)021, and hence does not explain large learning rate stability in SP in practice. Figure C.3 shows that the learning rate scaling 1/ Too n induces asymptotically width- width-indeepnndet instability threshold for 7o,, asymptotically constant sharpness in the kernel regime, and that small widths diverge at smaller 7/ due to diverging sharpness.
Proof. From the update equations, it holds that Act1 2 Xe if and only if XtmmapXeeUuxA Rap
4
st)
0 if and only if X 2 7a x
Observe that the loss decreases if and only if Ix+1l < Ixel, which holds if and only if [1 7pp + v/ m (xx v)1 < 1, which can be written as, omitting all subscripts 't.
n |= Mmpah 0].
31

---

Assuming 2 O, the above holds if and only if (nxf v
PopPlated [iell
and n n
maa 2
-2). The first constraint is a mild one that states = Onx). We will now focus on the second one. Solving for the roots of this polynomial in 77, we get
(neepBopd I V m (ntp(f)
11. 2||x|3XX
Assuming n Bntp2 12 > 8||x112xf =: C, we get Dap V " Bndy
DspBntp221
Dnspp 1 "app )2). In that case 1/1 Rs
C
2 and 1/2 ~
neptinend
2
[==|'pX
Hence, if xf 2 O, we get loss decrease if n Om 3n mmtp): PapPaatd T
2 nap
O m mtp)
Rapd |=
2 Map
O " }. where the left end of the interval is negative. The upper end resembles the edge of stability that vanishes as n for SP but not for NTP.
xf we get loss decrease if 17 E RaaRa ||||">
Om k.tp)) mep
Hep
Note the interesting slim regime of benign large learning rates 7) f(ʃ - y) > 0. As all of the involved quantities are known at training time, an adaptive learning rate schedule may significantly speed up training by stable learning with excessive learning rates.
n1 x=
O(n) when
Kap
However it remains unclear whether similar regime exists in practical architectures under CE loss. In that case, the sharpness computation would also be much more computationally expensive.
Frai lusstinn 15)
Lambbda Dynamics (h)
Largest stable Learning Hate
withi
Figure C.2: Stable SP. fo (left) and A (center) after training to convergence for several widths. The largest stable learning rate for SP indeed scales as n (right). When lines end, training diverged for larger learning rates. The first subplot shows that training has succeeded in memorizing the training label 1 at the optimal learning rate at all widths. The second subplot shows that the randomness X due to random initial conditions vanishes with increasing width, as SP is approaching its kernel regime.
32

---

Dynamics 16, A.A etael
Oynamics S. eta-1
Dynamics Ov-Als elted 1.
Do
140
113
-10
Eordiion
52000
1.4
044
04
Beration
Figure C.3: Stable SP dynamics. Jz (top) and Xe (bottom) over the course of training in SP with "nn = 1/o n=1 for several widths, where 1/o = in the gradient flow regime (left), Yo 1.9 close to the edge of stability (center) and 770 2.1 above the edge of stability (right). Smaller widths tend to diverge at smaller learning rates due to sharpness divergence due to the non-vanishing term in the sharpness update. The learning rate scaling The On leads to asymptotically consistent dynamics. This shows that larger learning rate scaling would induce divergence. Asymptotically SP with ?n = O(n lies in the kernel regime and the sharpness stays constant over the course of training. Edge of stability oscillations do not enable larger learning rate scaling under MSE loss.
D Experimental details
If otherwise specified, train a single epoch to prevent confounding from multi-epoch overfitting effects.
D.1 MIPs
We implement our MLP experiments on MNIST (Deng, 2012) and CIFAR-10 (Krizhevsky et al., 2009) in PyTorch (Paszke et al., 2019). We train ReLU MLPs with the same width n in all hidden dimensions with plain SGD/Adam with a single learning rate for all trainable parameters, batch size 64 without learning rate schedules, weight decay or momentum to prevent confounding. We use Adam with the Py Torch standard hyperparameters. By standard initialization mean He initialization variance c fan_ in with Co for the ReLU activation function (He et al., 2015).
D.2 Multi-index data
For some experiments in Appendix F, we generate multi-index teacher data, inspired by Kunin et al. (2024), but setting deterministic teacher for ensuring a balanced classification task.
We draw the covariates & ~ U(( i.i.d. from the uniform distribution on the unit sphere in Rd... with input dimension d,, 100. The training set consists of 10 training points. We also draw test set consisting of 10t test points.
For the target function f", drawing random directions as Kunin et al. (2024) results in heavily unbalanced classes and f O on large part of the support with high probability. Instead, set 4 teacher neurons deterministically for less noisy results. The teacher net is a shallow ReLU
network given by f*( sign Di s.o(w, {)) with unit vectors
= é1, w4 -E2 and signs 81 53 +1 and 82 S4 = -. This results in the nonlinear target function (E) = Sign(E - {2) for all g E Ri with & > 0 or &2 > 0, but (*(() sign(E2 {1 for all label noise.
W3 -é1,
E E (-oo, 0) (-oo, 0). We do
This dataset requires learning to align with the first covariate dimensions (Ei, E2), where of signal for the labels (E) lies. If the input layer does not learn align with these dimensions, the sparse signal obscured in the activations (random features) after the first layer due the large variance in the remaining covariate dimensions.
33

---

D.3 Language modeling
We train small Transformer models (Vaswani et al., 2017) using LitGPT (Lightning A1, 2023). We adapt the Pythia (Biderman al., 2023) architecture with 6 Transformer blocks, standard d_head 1/2 attention scaling. pre-attention and ok--Layrnorm (Wottsman et al., 2024). We purely scale width, proportionally scaling the number of attention heads and the MLP hidden size while keeping the number of layers and head dimension d_head= 32 fixed. For widths 256, 1024 and 4096, this results in 8, 32 and 128 heads per Transformer block and a total of 30M, M and 1.4B parameters.
Standard training means A with a single, tuned maximal learning rate. (31, (2) (0.9, 0.95). = 10 12, sequence length 512, batch size 256, 700 steps of warmup followed by cosine learning rate decay to 10% of the maximal learning rate, weight decay 0.1, gradient clipping. We train for 10681 steps in mixed precision on the on the same number of tokens to prevent confounding effects from increased training time.
dataset (Li et al., 2024). We train all models
D.4 Figure Details
Figure I: The training accuracy of 8-layer MLPs is averaged over runs to reduce noise from random initialization. The training loss of GPT trained with SGD is averaged over 3 runs. GPT with Adam was only run once.
Figure 2: In the left subplot, Qj = for normalization layers, since they act like diagonal matrices and do not accumulate fan_in-depp in-deeendentt scaling. For all other layers, we expect LLN-like scaling c f fan_ which is the length of the inner products between weight updates and incoming activations. For determining the correct layerwise learning rate scaling, what matters how much of the expected alignment exponent is lost at finite width in single update step. For the first step, we measure the barely width-deeendent exponents 0.01, -0.08 and -0.08 for the last Layernorm, the ML.P and readout layer, respectively, suggesting that maximal width-dependent alignment approximately holds, so that infinite-width scaling predictions are useful for determining layerwise learning rate scaling rules. The readout layer and last Layernorm layer are chosen due to their particular importance for logit blowup. The MLP layer was chosen to add a layer that scales hidden-like. This layer was not cherry picked. We observe other MLP layers to have similar scaling properties.
Figure 4: 3-layer MLP trained with SGD CIFAR-10 with width-dependent learning rate Th. 0.0001 . 0.5 Averages over 4 random seeds.
Figure 5: Learning rate sweeps for GPT. SGD runs are averaged over 3 random seeds, due to noisy individual outcomes. The results for Adam stem from a single random seed due to limited computational resources. Minimal unstable learning rates are defined as the smallest learning rates to produce loss worse than (optimal CE loss +1) at each width. The x-axes showing learning rates are scaled as (n/256)". In this way, the learning rate at base width 256 remains the same for comparability of the constants. If the optimal or maximal stable learning rate indeed scales as Th- -]--12" n then the width-dependent scaling of the x-axis 7/. n1 shows learning rate transfer.
Figure 6: Exponents are based on the learning rate curves provided in Appendix F.2 for GPT on and in Appendix F.3 for 8-layer MLPs on image data sets. For MLPs on MNIST and CIFAR-10, we measure minimal unstable learning rates the smallest learning rate larger than the optimal one to produce NaN entries when using MSE loss, or accuracy < 20% under CE loss. For GPT, measure minimal unstable learning rates as the smallest learning rates that are larger than the optimal one to produce loss worse than (optimal CE loss +1) at each width.
Figure 7: Shown is the training accuracy at the end of training for one epoch for the optimal learning rate at each width.
E Refined coordinate checks
The standard coordinate check as provided in the readme of the mup-package Yang et al. (2022) may be considered the plot of activation norms | after steps of training for all layers I and the network output norm | with f = WhL+lzf as a function of width. Completely width-independen dynamics under μP then result in an approximately andth-independent coordinate check of all layers. However, width-dependence in the activations of previous layers would confound the (-th layer activation scaling, so that measuring the effective (-th layer updates requires measuring
34

---

||WW]all RMS in each layer, where one may be interested in the weight updates accumulated over the entire course of training AW; = W W or the update a single step oW? = W - In standard architectures, one can equivalently measure the operator norm of the weight updates (alaanin(()///a [ 22:
_ (11) (Yang et al., 2023a); however in non-standard architectures such as Mamba this spectral condition has been shown to fail, so that, in the general case, care should be taken in how exactly weight updates affect the output function (Vankadara et al., 2024). The difference between W || and the preactivation updates [[ || precisely [W,r HMMS which measures the effect of updates propagating from previous layers.
All coordinate checks run over 4 random seeds either small learning rate or the optimal learning rate 1/566 at base width 256 (after epoch of training) on CIFAR-10. The learning rate is then scaled in relation that base width t/2 77 (n/256)a with clean exponent a -01,-
E.1 SGD
Figure E.1 shows the refined coordinate check for 3-layer MLP in SP with global learning rate scaling 7h, 17 n 1/2 As predicted by Proposition 4, the input layer updates decay as n=1 the hidden layer learns features width baath-independently. and the output scales n which results in one-hot predictions after the softmax wide models, but not necessarily unstable training dynamics.
Both ||WWaad and || CIMMs- equivalently and accurately even narrow MLPs of width 64. Naively tracking the activation updates Azt = o - z however is confounded by non-vanishing feature learning in narrow models, and only shows the correct hidden- and last-layer scaling exponents for 4000, even after only a single update step.
measure the effective update effect in the I-th layer
SGD, 0.0001"n --.5.
sus love
Hidcum loys
Calool leyer
CONA
0,5
1004
2013
1001
6130
10354
135
1004
2003
15254
20.205
Wlcch
Width
SGD, 000001"0"'-..5 3-kyer
Input
Outout
12-0
SGD, 000001"n" O.S, layor
bput
layer
Cutpot layar
0.5
1024 2018
4150
142M4
125
RIC 1975E
125
512
1004
2045
110
Wichh
Figure E.1: (Hidden layer feature learning in SP under intermediate learning rate scaling) Effective (-th layer update scalings |AW! - |I
RRMS (top), weight update spectral norm |AW! [] (2nd row) and activation updates 5x (bottom) of MLPs trained in SP with small learning rate Mh 0.0001 (n/2256) -1/2 scaled preserve hidden-layer feature learning. The TP scaling predictions are accurate. Hidden layers learn features width-indepenndenlly and input layers have vanishing feature learning. At moderate widths, activation updates are confounded by previous layer updates, and thus do not provide an accurate metric for effective update scaling.
35

---

Figure E.2 shows a refined coordinate check for a 3-layer MLP in SP with width-independenn global learning rate scaling 7/, 0.0001 n2. While infinite-width theory predicts the input layer to learn width-independently and the hidden layer to explode as ((n), both empirical exponents are n- smaller, so that the input layer has vanishing feature learning and the hidden layer is still exploding. This ostensible contradiction is resolved when repeating the coordinate check but initializing the last layer to 0 (Figure E.3). Now the predicted scaling exponents are recovered, already at small width. The reason for this subtle but important difference is that the gradient that back-propagated is given by the last-layer weights, 0. 0f/00/ = W.L+1 W + AW?+1 Under standard initialization at the optimal learning rate, the initialization W = 1/22 still dominates the updates AW = (n)) in absolute terms after a few update steps at widths up to 16384. Comparing the absolute scales of || RMS RMMS or [W in both figures confirms this hypothesis. The pure update effects in Figure E.3 have lower order of magnitude in the constant before the scaling law, but follow clear scaling exponents. Therefore the faster scaling law under last-layer zero initialization can be extrapolated with certainty to induce phase transition under standard initialization around width 4 10". We do not have sufficient computation resources to validate this but arrive at this order of magnitude irrespective of whether extrapolate the scaling laws of ||WW/1|| RMS s or || ||ll, Il, as well as of the input or hidden layer laws. For base width no and width-dependent statistics A, and A? with differing scaling exponents eı and c2, A and A WWMS
intersect width no / ( 1/(er-ozz)
This consequential difference in empirical scaling exponents realistic widths due a subtle dif- ference in last-layer initialization highlights the attention to detail that is required to make accurate scaling predictions from infinite-widhh width limit theory, but, as we show in this paper, apparent contradic- tions can often be reconciled with enough attention to detail. and the clean scaling laws we arrive at as a result already hold at moderate scales and prove the usefulness of investing this extra effort.
Hence, one reason why scaling exponents in SGD can be larger than predicted up to very large widths, is due to differing orders of magnitude in the constant pre factors in the initialization versus update terms in the backward pass. Without our refined coordinate check, the phase fransition around width 10? is hard to predict.
As predicted, the width-exponents of 2-layer nets behave like the input and output layer in 3-layer nets (Figure E.4).
When choosing the optimal learning rate 7256 = 0.03 at width 256, stronger finite-width effects due to non-vanishing input layer feature learning already occur after a few steps and make the update scaling exponents after 10 steps only visible at larger width n 2 2048 (Figure E.5). As long as divergence is prevented in the first few steps, seff-stabilization mechanisms such as activation sparsification can quickly contain the initial catapult (Figure F.6). In deeper networks, explosion of several hidden layers is increasingly difficult to stabilize, and finite width effects are reduced.
Figure E.6 shows the effective update rank and the alignment between activations at initialization versus at time t. for the same input training points under unstable width-independent learning rate scaling. The updates each layer are remarkably strongly dominated by a single direction. As hidden-layer activations are slowly diverging. their alignment only beginning to decrease at large widths 2 4096. The beginning instability of RMS will eventually induce training instability and suboptimal accuracy at large width, which is hard predict without tracking the layerwise effective update scaling across widths.
36

---

SGD, Con001'n""On 3-byer
layer
0,5
45
[DDetaly
updat ctop
2010
6110
18241
2001
9GD, 10.001"n""O.0, 3-layer
Ippu leye
Hiddun lost
Cutoul lape
[[DNTuk
13-2
$1400
0.02
107M4
5340
1.04
Wiith
W5rth
V0dti
SCD, 0000011'"'0.0 3-kyee
Cutput
Figure E.2: (Inaccurate exponent predictions under standard initialization with large learning rate scaling) Effective 1-th layer update scalings LAW.=!-1| nmss (top), weight update spectral norm || (2nd row) and activation updates dx (bottom) of 3-layer MLPs trained in SP with width-indeepndeett 7/, = 0.0001. Hidden layer activation updates explode, and input layers have vanishing feature learning. By TP scaling predictions, however, the input layer should learn features width-indeeendeetlyy Instead, the TP scaling exponents are only accurate under last-layer zero initialization, not under standard initialization (see Figure E.3 for last-layer zero initialization) as the initialization scaling W O(n 9(2) still dominates the update scaling AW,++ (())) at realistic widths after a few updates under the optimal learning rate. Hence, the backpropagated gradient 0 0f/Ozt W.+++ relevant for the hidden and input layer updates, behaves for several steps like it should only behave in the first step. By comparing the absolute scales here versus those in Figure E.3 becomes apparent that this indeed a finite-width effect, as the absolute scale of ||Wall2 here is on the order 10-1 and 10 for input and hidden layer, respectively, whereas the pure update effects under last-layer zero initialization are of at most order 10-4 for both layer types. Clearly for sufficient width, the differing scaling exponents will induce phase transition toward the predicted scaling exponents. While the input layer learns features width-independennly under last-layer zero initialization, as predicted TP theory, this is not the case realistic scales under standard initialization. The qualitative statement that standard parameterization with width- independent learning rates activation stable in deep networks still accurate at moderate width.
37

---

SGD, Ir n.9001"n"6... 3-jyer, Bnd)
Ypur
SpOut layar
100
2013
11100 18291
2701
SGD, 0.0001"n"00. E-O
Ippi leye
Hiddun losn
Outoul lape
D000r skp
[1Diiuk
0102
102M
6104
3000
Width
V0dth
8CD, 0.0001"n""0.0 3-ayer,
Ippit
Cutput
[[Duttal
Figure E.3: (Accurate exponent predictions in SP with last-layer initialization under large learning rate scaling) Same as Figure E.2 with width-independdet 7/, = 0.0001 but initializing the last layer zero. Here, the TP scaling predictions are accurate. Hidden layer activation updates explode as n², and input layers learn features width-indeeendently. Observe a smaller absolute scale of the pure update effects here versus Figure E. that explains the differing exponents there. The updates in the input and hidden layers vanish in the first step, the gradient for backprop is W r++
Vich
V4275
SGD, F00000 VIT DAL PANN
500, DOO YOO 2-ayyy B-)
OuB 10y1
D00IYY70 Blayer
Figure E.4: (Shallow nets learn features width-inddependently under large learning rate scaling) Same as Figure E.1 but for 2-layer MI.Ps trained SP with width-independent 7), 0.0003 with standard initialization (left) and last-layer initialized to 0 (right). The input layer and output layer scalings behave in the 3-layer nets. Since there no exploding hidden layer, activation stability preserved in 2-layer nets under O(1).
38

---

BGC, Ir 6.03''n"0.. 3-ayer
06
[Detaaly
1034
SGC, 0.03"'""0.0 3-layer
Irped boat
Hädden bown
Oupai lowt
[1Diiuk
3543
0192
1354
0349
676/91
100th
WWich
SCC, 0.03""""0.. 3-layer
Cutput
1.5
[[Duttaf
14
Figure E.5: (Large finite-width effects optimal learning rate in shallow 3-layer MLPs) At the optimal learning rate 7/256 = 0.03 with indth-inddeendentt scaling, non-vanishing input layer feature learning confounds the scalings after few update steps up to moderate widths S 1024, similar to Adam (Figure E.9).
39

---

BGC, Ir 6.03''n"O.. 3-aayr
[[Dwtal =/1wwwil
1.5
1.0
2003
1001
312
1300
1031
2043
SGC, C.03"""00.0 3-layer
puput lapet
Hobben lapel
Outpu Hape
50
1424
W95th
65507h
SCC, Ir 0000""""0.0, 3-layer
1011
1/
7
SGC, Ir C.03"n""O.0 "0.0, layer
Ingul bayer
Midden layer
Gubput layer
updalo stop
1004
20.5
4505
5700
1030
130
206
612
1024
2648
4608
5.02
1004
250
612
1024
2560
110
0825h
W0th
Figure E.6: (Activation sparsification at the optimal learning rate) Effective I-th layer update ranks
activation sparsity and cosine similarity between activations each layer comparing time 0 and time t on the same input training point and on differing training points in the same batch of 3-layer MIPs trained with SGD in SP with width-indeppndent learning rate Ya- = 0.03 as in Figure E.5. As opposed to the gradient flow regime, at the optimal learning rate, there significant self-stabilization effects large width already after 10 steps through activation sparsification but less through activation rotation.
40

---

BGC, Ir 6.03''n"0., 3-ayer
updale stap
[Dettak >
1081
120
2010
Widh
SGC, C.03""""0.0 3-layer
Ippul leyer
Hiclln hape
0.4pui
1024
W4lth
Witth
Wktth
SCC, Ir (.03""""0.0 3-layer
0.1
7 39
Figure E.7: (Full width-independencce in μPP Effective 1-th layer updates (top), effective update ranks | |||l|/////////// AW; (second row) and cosine similarity between activations to each layer comparing time 0 and time on the same input training point (bottom) of 3-layer MLPs trained with SGD in μP with width-indeeendent learning rate 1h = 0.03. As expected, all statistics behave width-indeeendeenly. The effective update rank is remarkably small, for SP. The activation rotated quite quickly.
E.2
Adam
With ?h = 1 the optimal learning rate scaling for 3-layer MLPs with Adam on CIFAR-10 is larger than predicted (Figure F.30). Figure E.9 shows that this may be due large finite-width effects for Adam at optimal learning rate multiplier 7/256 = 0,0003 and moderate width < 8192. While the weight update spectral norm scales as predicted, the input-layer gets large updates moderate width (Figure E. 10) and induces a strong rotation of the activations. As result, the activation explosion only sets in large width n 2 8192. This qualitative change toward vanishing input layer feature learning will result in phase transition toward unstable scaling large widths which is hard to predict at small scale from measurements alone, except when measuring both |/WWI]) and alignment || A -
As opposed to SGD, observe large finite-width effects in the activation updates even under small absolute learning rate 10-6 at moderate width 8192 (Figure E.8).
41

---

ADAM, 1e:0E"n"--... 3iayer
layar
12"1
14
[DDetaly
16
2013
18201
2001
ADAM, 1e-06"n"-00 3-toye
Impul loye
Hiddun kost
Outocl laper
45
Dcce sepp
Us
0102
187M
6.04
3000
Wilth
10-06"n""-0 3-Rtyr
Cutest
[[Outaa
Figure E.8: (Large finite-width effects from input-layer updates in Adam) Effective I-th layer up- date scalings | IRMs (top), weight update spectral norm AWll. (2nd row) and activation update norm ||ōxa JRMS (bottom) of 3-layer MLPs trained with Adam in SP with 7), 10-6. nn11/2 1/2 Observe the theoretically predicted exponents in ||W' || do not transfer to the activation updates at moderate width n < 8192 due to large non-vanishing input layer updates at moderate width. Even the effective updates |/ || RMS do not perfectly align with the scaling law at infinite width, indicating that the alignment between AW; and 2; evolves non con-trivially across width and that the spectral norm AW and pure infnnit-wwdtth predictions are less useful for explaining the behaviour of Adam at moderate width.
42

---

ACMM, 0.0000 "1" 3-ayer
upodl Aap
10-1
1080
2013
1081
Widh
ADAM, 013003""1"" (1000""""""000000000 3-baye
Impul lowe
Hiddun kost
Outoul lape
055
Us
135
0.02 100M
4006
1004
6.04
W5hth
W5rth
V0dth
0.0008"""" "1"0000, 3-leyer
Cutest
Figure E.9: (Large winit effects from input-layer updates in Adam) Effective I-th layer update scalings ||WW IIMs (top), weight update spectral norm ||W! (2nd row) and ac- tivation update norm [| (bottom) of 3-layer MLPs trained with Adam in SP with large "h = 0.0003 -112 Observe the theoretically predicted exponents in || ||W*), do not transfer to the activation updates at moderate width n < 8192 due to large non-vanishing input layer updates at moderate width. Even the effective updates / |RRM do not perfectly align with the scaling law at infinite width, indicating that the alignment between AW? and 2!-1 evolves non-trivially across width and that the spectral norm |AW and pure infinite-width predictions are less useful for explaining the behaviour of Adam at moderate width.
Ir 0000009""""005 -0.5, 3-Ryer
Trpu: layer
Höklen koye
Ougpui kyww
updaae stap
1.00
1.5%
1,000
1001 Width
2014
6.00
116r
16/9
512
100
2010
4200
1130
10091
3716
Wicch
Width
11" layer
Output
10
10P
00
Figure E.10: (Strong activation rotation under Adam at moderate width) Effective I-th layer update ranks ll comparing time 0 and time the input training point (bottom) of 3-layer MLPs trained with ADAM in SP with large 7/2 = 0.0003 n 11/ The effective update rank is mostly growing in time in the input layer. Already after a few steps, the first-layer activation coordinates are drastically rotated at moderate widths. This induces u-curve in the hidden-laye -laer activations that inherit large rotation from the input layer at moderate width and update too much at large width under ?h ((n
(top) and cosine similarity between activations to each layer
43

---

E.3 Normalization layers and Adam provide robustness to miss-inttialization
For MLPs trained with SGD, initialization greatly impacts the training dynamics as both the forward and the backward pass are affected (Figure E.11). Large input layer initialization induces update instability large width, which is stabilized by extreme activation sparsification (Figure E.13).
By adding normalization layers, the forward pass can be enforced scale codth-inndeendently. This may affect the gradients. But the gradient norms become irrelevant under Adam with sufficiently small Adding both normalization layers and Adam to MLPs, observe that initialization barely relevant for update scalings (Figure E. 12), and other downstream statistics such as activation sparsity (Figure 14). Here we use RMSNNrm to fairly compare activation sparsity, but we expect LayerNorm to induce the same scaling behaviour.
SGC, Ir C.03"'"'..,, layer
upde
6,5
SCC, 0.03""*"000 3-layer
Irput boot
Hokiun koyn
Oupas loow
137"
377
12:0
200
2000
10:30
1201
2018
0380
1229
Width
Width
Dlich
Figure E.11: (Initialization matters MLPs with SGD) SP (top). SP with large input layer variance 2 (bottom). The initializations induce significant differences in the training dynamics. Large input layer normalization becomes unstable at large width.
10.000"""""-100
Hicdhnn
Catast layer
[Dotua
Wicch
Withh
Videi
ACAM, ".0000"'"""1111 layor
Figure E.12: (Differing initialization barely matters with normalization layers and Adam) Update spectral norms of MLPs with the most basic normalization layer RMSNorm after every layer trained with Adam and initialized with SP (top) versus SP with large input layer variance (bottom). Here, initialization barely impacts the update scaling.
44

---

BGC, Ir 6.03''""0.0 3-aayr
Soper
2041
1034
Widh
SGC, Ir 0.03"""'..0,
Tpul lout
Hikden loget
Outpu Lape
slep
1004
600th
400t
6000h
Figure E.13: (Big difference in activation sparsity under SGD) SP (top), SP with large input layer variance 2 (bottom). Large input variance has be stabilized by increased activation sparsity.
ADAM, 0.3008"'""11.. 3- 3-kyer
Inpur Styer
Holdon byyer
Ourgit bye
17.
upok2o Stop
015
1004 Wi-th
2012
312
1001
2010
110
16/0
31.
110
Wich
0.0003"n""11.0 3-lyer
Irpui kayer
Hidden
Ougsst lapor
12
W8tth
Figure E. (Activation sparsity barely affected under normalization) Same Figure 12 but showing the fraction activation entries that equal Both initializations do not significantly sparsify activations beyond 50%.
E.4 Alignment and update scaling in Transformers
For GPT in SP, alignment and update scaling follows the theoretical predictions: Hidden and output layers diverge with width constant learning rate (Figure 1. top). At the optimal learning rate where 7/n - 0 (here only 7/- ((n is shown), embedding and normalization layer updates vanish with width (Figure 16 and Figure 15, bottom).
SP-full-align achieves approximately width-indepenin doa > Since we measure width-independent alignment Ow,"-' @(1) (Figures 2 and E.18), under large output dimension duu > n, the initial output layer operator norm II approximately scales ((1) (Vershynin, 2010), as opposed to (/2) at sufficient width > dou- The term W Auf therefore induces approximately width-indeepndent logit updates even under standard last-l initialization, in the regime dlat > (cf. Figure 1717), but induces logit divergence at sufficient width d <
signal propagation in the language setting
45

---

Trating
Trahing
Treting
Figure E.15: (Effective updates follow predictions) Effective updates ||W, for constant learning rate scaling n,, 0.01 (top) and stable learning rate scaling 7/,, 0.01 (n/256) 25/)) (bottom) in GPT models of varying width (the darker, the wider) for the embedding layer, the first MLP layer in the Transformer block 2, the last Layernorm before the readout layer and the readout layer (from left to right). At constant learning rate, hidden and output layers diverge with width. At optimal learning rate, embedding and normalization layer updates vanish with width.
(105 IMCallas
0/003
0.068
dellles
20007
3885
0.006
0.015
-0014
3555
With
0g49
-0:022
([WG laus
38895
331
8:88
00001
0.076
10"
0.008
-0:278
0.000
226
1634
4006
1094 Withh
(22061 With
4000
901 Witth
100
with
-0978
[A/WGr]]lao
-100
250
-04800
W996
Wi2th
Widh
Figure E.16: (Refined coordinate checks for GPT in SP with Adam and 77, 0.01-n From left to right: Activation norm, activation updates, effective updates || AW propagating updates [|WoAeellRa after 2, 10, 100 and 700 batches of training (the darker, the more batches). Layers from top to bottom: readout, last Layernorm, first MLP layer in Transformer block 2, embedding layer. Infinite width-scaling predictions are accurate in all effective update terms AW,,, R Embedding and Layernorm layers scale input-like and their updates vanish as Onn='), all hidden and output layers are effectively updated width-indepenndently Against the infinite-width prediction, logit updates do not explode, not because of miss-alignment but because output dimension is much larger than width dam >> n, which changes the approximate scaling of [[W]+1]] [| from (n in the indinite-width limit, to ((1) in the large output dimensional regime.
46

---

945
(Maollau
[A
1
0.005
Suges
40201
00002
0,30
With
-0.15
0401
0:01
[[M Misillos
[b] micheilless
-0.29
40.01
0.01
0.076
o.gno
-1045
Aeillou
025
0. 1.5
40.022
-0.04
0.004
-0.001
240
16M
426
X44
304
with
With
With
Witth
0.002
000
0.000
10
Wiah
Wkth
4000
Figure E.17: (Refined coordinate checks for GPT in SP-fullalign with Adam and width- independent Th 0.003162) Same as Figure E.16 but for SP-full-align. The propagating update and effective update terms in all layers scale approximately width-independennly. The output layer propagating updates do not diverge as in the large width regime n " dast (Figure F.22). Here, instead, due to the large output dimension dom > n, |W [lRMs--AM is dominated by its other summand, which induces approximately width-independee signal propagation at realistic widths.
Wo ([Wolll(((
10
-0.064
Width
256
1024
4096
10
10
10
102
10
Training Step
Figure E.18 (Updates propagate maximally in the readout layer in SP-full-align) The operator norm ratio for propagating activations in the readout layer for training GPT with AdamW in SP- full-align with near-optimal learning rate Th 0.00316. The ratio is barely width dependent so that propagated activations can be computed when knowing both ||alloo [WolllMs--ae and | RMM-s
47

---

F Empirical learning rate exponents
F1
Summary of the MLP experiments in this section
In general, the optimal learning rate exponent appears to be architecture- well as data-dependent. We conjecture that the optimal learning rate scaling is subject to opposing objectives. Ideally, the effective updates in all layers scale width-independently. Since this cannot be achieved with single learning rate for input, hidden and output layers, the layer types act on the optimal learning rate scaling opposing forces.
SGD under MSE loss. For SGD under MSE loss, output blowup results in unstable training dynamics so that the maximal stable and optimal learning rate robustly scales as 7,2 = =0 across architectures and datasets. As consequence of vanishing feature learning, neither training nor test loss monotonically improve with scale under MSE loss.
Random feature models. When only training the last layer, fully width intthinndependent training
dynamics are achieved with 7/n- = 17 n n! Figure F.19 shows that this exponent clearly results in learning rate transfer for 2-layer ReLU random feature networks on CIFAR-IO. Also observe that since all learning rate scalings recover activation-stability, larger than optimal learning rates still result in non-trivial classification accuracy.
Deep MLPs. With increasing amount of hidden layers, their width-inddeendence eventually outweighs input layer feature learning in vision datasets. For at least 6 layers, we see approximate learning rate transfer under 7/ (n
-1/2) -1/2) for SGD and Th = O(n (1) for Adam as predicted for width-independern hidden layer feature learning for both CIFAR-10 and MNIST.
Shallow ReLU MLPs at moderate width and (deep) linear networks are not useful proxy models for deep nonlinear networks. For shallow MLPs, we often observe stronger finite-width effects than for deeper networks causing larger than predicted optimal learning rate scaling moderate width, as divergence in fewer hidden layers can be stabilized over the course of training up to larger widths (cf. Appendix E). In linear networks, on the other hand, feature learning not essential as the learned function always remains linear. Consequently we often observe that optimal learning rates decay faster than maximal stable learning rates in (deep) linear networks even under CE loss (Figures F.20 and F.40). These differences between deep non-linear networks and toy architectures suggest that shallow MLPs and deep linear networks do not serve as useful proxy models for practical non-linear networks in terms of optimal learning rate exponents at moderate width.
Input layer task. Under multi- index data with sparse signal and high- ditbnsiinnii isotropic covariates (explained in Appendix D.2), learning the two signal input dimensions is particularly useful for good generalization. Appendix F.4 shows the predicted exponent 77, 17 -n for input layer learning in 2-1ayer MLPs. Deeper MLPs recover hidden layer stability with optimal learning rate scaling 7/1 = (n 1/2).. Observe that generalization suffers when the input layer does not learn to align with the signal dimensions, that only the 2-layer MLP with CE loss generalizes well at large width.
Standard initialization with jPP learning rates (SP-full-align). While Everett et al. (2024) re- port good transfer properties of SP-full-align, Appendix F.7 shows that the optimal learning rate clearly shrinks across image datasets and our multi-index data. We also introduce a variant of this parameterization that matches the 1/2 logit blowup rate from the term W Ar, in the effective last-layer updates by increasing the last-layer learning rate. This variant performs similarly well as SP-full-align. In particular, both variants seem to be less learning rate sensitive than J.P.
Adam learns features with 7h 17. n= Adam simplifies the learning rate scaling for weight W to Mj // aaaaaannnnnaaaaaaaaaanaaaaaeeeee (W), because the gradient normalized but still correlated with the incoming activations since the sign is preserved in each entry. Thus 7. = 7/ is expected to induce width- independent hidden- and output-layer learning, but vanishing input-layer learning since here fan_ in is fixed and hence would require constant learning rate scaling. As for SGD, we still observe the optimal learning rate scaling 7n = 17 n in deep MIPs on MNIST and CIFAR-10 (Appendix F.6), indicating that width aatth-independennee in hidden- and ouyput-layer dominates input layer feature learning.
48

---

F22 Transformer experiments
As we consider single-pass training, training and validation loss approximately coincide, so that statements about the training loss transfer statements about the validation loss irrespective of the optimizer. All figures in this section show training loss on the left and validation loss on the right.
Stabilizing techniques like gradient clipping can improve the absolute learning rate multiplier in front of the scaling law, but do not seem to change the width-scaling exponent for SGD (Figure F.4 vs Figure F.5).
Pythia 14M, mixed precision, SP
Pythia 14M. mixed precision, SP
Final Training Loss
Final Validation LoGs
3.8
256
268
1004
azue
3.6
warmpp
3.8
warmup
400
400
700
3.4
700
10
10"
Max Learning Rate w
Max Learning Rate width?"1
Figure F.1: (Instability without gk-Layernorm) Train loss (left) and validation loss (right) of single-pass AdamW W training without gk-Layemorrm Training and validation loss approximately coincide. Optimal learning rate scaling is dominated by the maximal stable learning rate scaling that is most ((n But without qk-l.ayernorm, the stability threshold is decreasing faster than O(n-1) even when increasing warmup length, so that it may be that the instability threshold would decay beyond the ideal learning rate and performance suffers. our computational budget does not allow to scale further, this setting remains inconclusive.
SP, ADAM, ck-layemonn
SP. ADAM, ck-layernorm
Final Training Loss
Final Validalion Loss
4.2
width
width
3.8
256
3.6
256
3.6
10
10a
10*
10"
10
10
10
10"
10
Max Learning Rate (width/256)-0055
Max Learning Rate (width/256)
Figure F2: (Large learning rate stability with gk-Layernorr) Same as Figure F. but with qk- Layernorm as recommended by Wortsman et al. (2024). Training and validation loss approximately coincide. The optimal learning rate seems to approximately transfer under r) n -1/2, so the added Layernorm appears to stabilize learning at larger learning rate scaling, similar the softmax in CE loss.
49

---

SP. ADAM. qk-layemomm, fixed LN
ADAM. gk-layemomm. fixed LN
44
4.4
Final Training Loss
Final Validation Loss
3.8
3.6
width
width
258
1024
7024
4006
10"
Max Learning Rate (widhh2566)11
Learning Rate (width2256)"1
Figure F.3: (Smaller optimal learning rate exponents under fixed Layernorms) Same as Fig- ure F.1 with gk-Layernrrm as recommended by Wortsman et al. (2024), but all trainable Layernorm parameters fixed to initialization. Here only the embedding layer behaves input-like, so that all other parameters learn width-inndeeendenll under learning rate scaling (nn-1). While the optimum is drifting toward larger learning rates, an increasingly large plateau of near-optimal learning rates emerges at large width. (n==1) still approximately captures the maximal stable learning rate scaling.
SP, SCD, no gradient clipping
SP, SGD, no gradient clipping
7.0
7.0
Fihal Training Loss
Final Validatin Loss
8.5
6.3
8.0
6.0
width
width
256
4096
4096
10"
10"
10"
16"
Max Learning Rate (width/256)"005
Max Leamingg Rate (width/256)"0.5
Figure F.4: (GPT trained with SGD has 1/2-1 -eaaarinnn rate scaling) Train loss (left) and validation loss (right) of single-pass pass SGD training (averaged over 3 random seeds affecting weight initialization and data shuffling). Training and validation loss approximately coincide. Hence also oopiiaaionooptimlll learning rate scaling is dominated by maximal stable learning rate scaling -1/2) for hidden-layer stability.
SP. SGD, gradient clipping
SP. SGD, gradient clipping
Final Training Loss
8.5
Final Validalion Loss
5.5
5.5
3.0
width
width
256
256
4088
4.0
4.5
10"1
10'
10'
10"
:6"
10'
Max Learning Rate (width2256)-005
Max Learning Rate (width/256)
Figure F.5: (Larger learning rates remain stable under gradient clipping) Same as Figure F.4 with gradient clipping. Performance is significantly improved larger learning rate constants are stable (observe similar performance as without gradient clipping at the same learning rate). Optimal learning rate scaling is still dominated by the maximal stable learning rate scaling Tn = m 1/2 for laye stability.
50

---

SP-full-align,
SP-full-align, SGD
Final Train Loss
Final Valid. Loss
width
width
256
256
1024
1024
4096
4096
10
10"
10'
10
Max LR (widtth2266)00
Max LR (width/256))00
Figure F.6: (No clean transfer under SGD in SP-full-align) Same as Figure F.5 but in SP-full-alige The optimal learning rate tends to grow and saturate at the maximal stable learning rate which appears to be roughly width- aetthinndependent predicted, but transfer is not ideal.
SP-full-align, ADAM, qk-layernorm
SP-full-align. ADAM, qk-layernom
4.50
4.50
Final Train Loss
4.25
Final Valid, Loss
4.25
4.00
4.00
width
width
3.75
256
3.75
256
1024
1024
3.50
3.50
4096
4096
10
10"
10
10"
10
10"
10"
10
10"
10
Max (width/256)""0
(width/256)"00
Figure F.7: (Optimal learning rate shrinks at large width under A in SP-full-alinn) Same as Figure F.2 (AdamW and trainable qk-LLayrnorm) but in SP-full-align. The optimal learning rate initially transfers but starts shrink sufficient width. In Appendix F.7, the optimal learning rate already decays at small width on image datasets with small output dimension.
Final Valid. CE Loss
SP, MSE loss
width
256
6
1024
4096
5
10
10
-2
10
10
Max LR
(width/256)***
Figure F.8: (Optimal learning rate exponent 1 predicted under MSE loss) Same Figure F.5 (SGD in SP with gradient clipping) but under MSE loss. The optimal learning rate appears follow 7n = 17 n expected. Loss is worse than under CE loss.
51

---

Learning Rate
2
0,648
d.model
Figure F.9; (Large learning rate exponent in original GPT paper) Just plotting the reported learning rate and d model values from Brown al. (2020) results quite stable scaling law with exponent -06488, which is larger than -1 required for hidden-layer stability but significantly smaller than required for vidth-indeppndent embedding and normalization layer learning. But note that jointly increasing batch size, l_layers and heads might be confounding factors here.
F33
Learning rate sweeps corresponding to 8-layer MILs
Here we provide all learning rate curves that correspond to the learning rate exponent tables 8-layer MLPs trained on image datasets (Figure 6). The curves corresponding to GPT on
are provided Appendix F.2. For SP under CE loss and MSE loss, Figures F.10 and F.1 show that maximal stable learning rates approximately follow the predicted exponent -0.5 for CE loss and - for MSE loss, and even for the worst worst-fitting dataset CIFAR- 10 the optimal learning rate saturates at the maximal stable learning rate at realistic width 16384 (Figure F. F.13). Thus, for CIFAR-10 10, we expect a regime transition when scaling further where maximal stable learning rate constrains optimal learning rate to scale with exponent close to -0.5.
In all datasets, performance improves much under which feature learning is lost and the accuracy tends to asymptote at moderate scale. For GPT the DCLM-Baseling dataset with MSE loss, gradient clipping and normalization layers might stabilize larger learning rates, but optimal learning clearly requires exponent -1. For SP-full-align under CE loss, Figure F.12 shows that the maximal stable learning rate remains remarkably width- independent (as predicted) on all four image datasets, but that the optimal learning rate decays with differing exponents.
with scale under CE loss than under MSE loss,
MSE SGD MNIST
CE SGD MNIST
100% Train accuracy
10a
10a
width
256
1024
10
10'
4096
16384
10a
10a
10
10 3
10
10
10
10
10
10
10 1
10
LR (WIDTH/256))
LR
(WIDTH/256))
Figure F.10: (Learning rates decay slower under CE loss than under MSE loss) Width-scaled learning rate versus training error for MNIST showing approximate transfer with Th 11 n under MSE loss versus with 7) n under CE loss.
52

---

100%. Train accuracy
3 1'x
width
100% Train accuracy
xxx
100% Train accuracy
10'
2x
1024
widti
width
4006
x 100
256
84
256
1024
1024
x
16384
8 10
10
10"
10
10"
10
10"
10"
10"
10"
LR (WIDDH2258)"000
(WIDTH2256)"10.3
(WIDTH22566"100..
100% Train accuracy
100% Train accuracy
8xx 101
100% Train accuracy
9.055
9.8
width
8x
width
8.84
width
4 40
0.988 10
10
1024
2xx 10'
1024
1024
4286
4000
4095
2x
16384
x 10'
182994
0.52
R.15
10"1
LR (WIDDTH2266
LR (WIDTH/2566)
LR (WIDTH(256))(1100
Figure F11: (Maximal stable learning rate dominates optimal learning rate suflicient width) Same as Figure F.10 but for FashionMNIST (left), CIFAR-I0 (center) and linylmagenet (right): Training accuracy as function width-scaled learning rate of 8-layer ReLU MLPs trained with SGD in SP under CE loss (top) and MSE loss (bottom). The maximal stable learning rate approximately follows the predicted exponent -0.5 for CE loss and for MSE loss, and the optimal learning rate saturates close to the maximal stable learning rate at sufficient width. MLPs barely learn on Tinylmagenet when using MSE loss.
19%6 Trai 5000 scourai
100% Them pamaaconee
Tain securany
225 12!
4200
23:10
8.5
8s 12'
12'
100%
4090
Ex:00
10034
636
375: 12
as 12'
CC90
UR WM0TH2250"100
LR INDTH
LA (NDDH4256)"10
(WND) (MDTHH200"140
Figure F.12: (Optimal learning rate decays in SP-full-align) Training accuracy function of learning rate of 8-layer ReLU MLPs trained with SGD in SP-full-align under CE loss on MNIST, MNSSi CIFAR-10 and Tinylmagenet (from left to right). As predicted, the maximal stable learning rate remains strikingly width-independent in all datasets, but the optimal learning rate decays in all datasets.
SGD, MSE
SGD, CE -0.54 -0.37
Learning Rate
mnist
-..46
0.07
10
cifar10 fashionmnist
10"
X
-0.9
-0.77
-0.89
10
-0.7
10
-0.98
-0.9
-0.78
-0.55
256
1024
4096
16384
256
1024
4096
16384
Width
Width
Figure F.13: Learning rates decay slower under CE loss than under MSE loss. Optimal learning rate (solid) and minimal unstable learning rate (dashed) for 8-layer MLPs on CIFAR-10, 10, MNIST and Fashion-MNN MNIST corresponding to Figure F. Optimal learning rates are often close to max-stable learning rates. Theoretical instability predictions y = O(n -1) for MSE loss and 7/- O(n 1122 for CE loss are surprisingly accurate. Even for CIFAR-10, 10, the optimal learning rate saturates at the maximal stable learning rate at realistic width 16384, so that it is expected follow the maximal stable learning rate scaling at larger widths.
53

---

MNIST
Fashion- MNIST
CIFAR-
Tiny ImageNet
DLCM-
Baseline
CE loss
Max-stable LR expon.
-0.54
-0.7
-0.37
-0.33
-0.38
Optimal LR expon.
-0.46
-0.55
0.07
-0.33
-0.38
MSE
Max-stable LR expon.
-0.9
-0.89
-0.77
-0.50
Optimal expon.
-0.98
-0.83
-0.9
-1.05
SP-full-abee -aignn Sta-ee LR expon.
0.0
0.0
0.0
0.0
0.09
Optimal LR expon.
-0.25
-0.35
-0.33
-1.45
0.21
Table F.1: Learning rate exponents at the edge of controlled divergence. Same data as in Figure but as table. We do not compute exponents on TinylmageNet under MSE loss because the accuracy of MLPs remains too close to random guessing.
F.4 MSE loss with softmax does not enable feature learning, but Adam does
CE versus MSE versus MSE+softmax. Here we train 2-layer and 3-layer ReLU MLPs on generated multi-index teacher data as detailed in Appendix D. These data crucially differ from the other considered datasets in that the target function only depends on the first 2 input dimensions. Due to the isotropic covariate distribution, input layer feature learning is necessary for good generalization. Hence we observe an approximate 7/1 R ((1) scaling for 2-layer MLPs with CE loss, necessary for preserving input layer feature learning (Figure F.14, left). 3-layer MLPs attain the maximal activation-staie exponent ?he = (n 1 in CE loss (Figure F. F.16, left). 2-layer MLPs preserve better validation accuracy compared to their training accuracy than deeper nets. as input layer learning gets increasingly inhibited by ) (leannnn rate instability in the presence of hidden layers. Both for shallow and deeper MLPs with MSE loss, we lose feature learning under the maximal output-stable scaling Tn = O(n ((n=)), as expected.
0.2780
050025
1304
1104
1504
00/85
04190
0.19
10-4
121
80000
12:11
127
0.070
12"
LR (MIDTH229011000 (.0)
(WCTH2DR//118
LR (K10TH-250010111
0.544
Figure F.14: (CE loss increases maximal stable learning rate scaling @(1) in 2-layer nets) Training accuracy (top) and validation accuracy (bottom) for 2-layer MLP on generated multi-indes teacher data (mean
4 seeds) with CE loss (left), MSE loss (center) and MSE loss with softmax (right), The x-axis scales the learning rate with width- indth-dependent exponents; observe approximate transfer under (1)) On -1 and O(n 1) scaling, respectively. the MSE plot, ending lines indicate divergence for larger learning rates. MSE loss with softmax on the output does not increase optimal learning rate scaling due vanishing gradients and gets worse due a lack of input layer feature learning.
54

---

SGD, CE, teacher, layers
SGD, MSE, teacher, 2 layers
Learning Rate
Learning Rate
10"
-1.12
maxstable LR
optimal LR
10"
-1,08
maxstabie optimel LR
-0.25
256
16384
Width
Width
Figure F.15: (CE loss increases maximal stable learning rate scaling to @ in 2-layer nets)
Optimal and maximal stable learning rate exponents corresponding Figure F.14. Since 2-layer MLPs do not contain hidden layers, the maximal stable learning rate under CE loss is determined by input-layer stability 7/. = ((1) (see Remark C.14). Under MSE loss, the maximal stable learning rate is constrained by logit stability Th. x 1).
3.233
2014
with
1051
0.09
1004
Valvidon
0.19
18 1000017750/000
13 (0017--75E7115
1R (W0CT4556)11 2g
Figure F.16: (Cross-entropy loss increases maximal stable learning rate scaling to approximately ( in 3-layer nets) Same as Figure F.14 but for a 3-layer MLP. The x-axis scales the learning rate with width-dependent exponents; observe approximate transfer of the maximal stable learning rate under indicate divergence for larger learning rates. Observe that wider networks generalize worse with scale as they lose input layer feature learning.
O(n and ((n-1) scaling, respectively. In the MSE plot, ending lines
In this setting, it becomes particularly apparent that using the MSE loss with softmax applied the output of the network is not desirable (Figures F.14 and F.16, right). Ultimately. the only difference to CE loss is that the loss derivative with respect the network output f(E) WL x xz)(E) becomes
Oc oc) af
- - -())))))))())))..)))))))))))))))
o())),
SE(C)
where the inner derivative of the softmax (6, -y) vanishes as soon as the outputs diverge [i({)) [E) - oo training point Ç. Hence, while the softmax still mitigates output blowup in the forward pass, the gradients vanish under output blowup. The CE loss, on the other hand, exactly the correct choice of loss function to cancel out the inner derivative of the softmax and effectively
view (f)) as the output of the network, resulting in
ac 37
= of)); Yj-
Here vanishing gradients under output blowup the MSE+softmax setting is so severe that output blowup prevents learning under large learning rates and the optimal learning rate scales as ((n -1)-
55

---

width
30
width
40
250
Training accuracy
25
4086
Training accuracy
4080
16384
20
25
15
10
10
10
10-22
10" 10 10" 10P
10--5
10
10-23
10--2
10-11
10" 10'
LR (WIDTH22566"00.5
LR
(WIDDTH2566)
Figure F.17: (Layernorm on logits also stabilizes large learning rates under MSE loss) 8-layer MLPs trained with SGD (left) and Adam (right) on CIFAR-10 with MSE loss with Layernorm applied the logits. The Layemorm has similar stabilizing effect as CE loss and allows learning with logit blowup under Th On (n12),,, but performance does monotonically improve with width. For Adam, hidden and output layers learn width-independenttl with T/n = O(n
param,
40
Tus,
Training aczuraoy
aczuraay
mos
44
param, loss
40
MAD
218
40,
m100
255
1024
4026
258
98
98.5
param, loss
Training socuracy
Training accuracy
muc,
84
param, loss
87.5
up mos
ep.
map:
ap.
5p. rise
915.5
86
85.0
258
1024
4006
256
1024
4008
width
Figure F.18: (Large performance difference between losses for SGD in SP but not Adam) Optimal training accuracy 8-layer MLPs trained with SGD (left) and Adam (right) on CIFAR-10 10 (top) and MNIST (bottom) with MSE loss (dashed lines) and CE loss (solid lines) in μP (blue) and SP (orange). For SGD in SP, CE loss performs much better than MSE loss as large learning rates recover feature learning at large widths. The performance in juP depends much less on the loss function since features are always learned budth-independently In μPP MSE loss slightly outperforms CE loss, suggesting that μP might benefit from exploring other loss functions. For ADAM, The O(n in SP remains stable even under MSE loss and recovers hidden-layer feature learning so that the difference between losses is much smaller.
In Figure F.17, we apply a Layernorm to the logits instead of softmax. The Layernorm allows learning despite logit blowup under 7)- = On -1/2), but performance does not monotonically improve with width either.
Adam reduces sensitivity to loss function. Figure F.18 shows that big performance difference between MSE loss and CE loss for SGD in SP is reduced if either (a) jP is applied to ensure width- independent feature learning or (b) Adam SP used, which remains stable with 7), = ((n-1) due to stable update scaling even under MSE loss, allowing width-indeeendent hidden-layer feature
56

---

learning. Observe that MSE loss outperforms CE loss in both jPP and Adam in SP. This suggests that exploring loss functions beyond CE loss might be promising when using uPP and/or Adam. more detailed evaluation of Adam is provided in Appendix F.6.
F.5
SGD
F.5.1 Random feature models are optimal under stable learning rates
Figure F.19 shows that 2-layer random feature ReLU MLPs transfer under 7/o, = 17 n= learning rate scaling under any loss. Under CE loss, also larger learning rates result in non-trivial learning as saturating the softmax does not harm training stability. Under MSE loss on the other hand, training diverges above the edge of stability and results in trivial accuracy of 10%. Under MSE with softmax on the output logits, a exploding logits induce vanishing gradients which also inhibits learning (see Appendix F.4 for more details), and results worse accuracy than under CE loss.
w92h
wilh
75
LRT (MDTH2380111)
UR (ADTHOMP"11))
LR (MD (WDTH236/7013
Figure F.19: (Random feature models approximately transfer under 7mn O(n for SGD) Training accuracy after one epoch of only training the last layer of 2-layer MLPs on CIFAR 10 with CE loss (left), MSE loss (center) and MSE loss with softmax (right). Observe approximately width-independent dynamics with Note that also larger learning rates result in non-trivial generalization because there is instability caused by activation blowup. The larger learning rates are not optimal, because the usual benefits of larger learning rates like increased feature learning or activation sparsity do not apply random feature models.
independent of the loss function or architecture used.
F.5.2 Linear networks are optimal under stable learning rates
Linear networks lack the ability to learn non-linear features for improved generalization at large learning rates. Hence, small learning rates 7,, O(=-1) are optimal for MNIST (Figure F.20), where feature learning is lost, but activations and logits remain stable. Observe that also deeper linear MLPs are only stable under (n-1/2) as theoretically predicted.
Thaking aoouary
8
1ich
Trikkgg SNIMEEY B
18"2 LR 1VIETT475500
LR (NDDTH"
LR (NDTTT"T"
Figure F.20: (In linear nets on MNIST, the optimal learning rate shrinks faster than the maximal stable learning rate) Same Figure F.22 but for linear nets. The maximal stable learning rate scales similarly as for the non-linear nets, but the optimum approximately follows (nn Irrespective of the depth, linear MLPs can only learn linear transformation; hence under sufficient width, feature learning under large learning rates does not provide benefit over mere last-layer learning.
57

---

F.5.3 Small optimal learning rates under MSE loss
With MSE loss, observe a clear 7/ = On optimal and maximal stable learning rate exponent irrespective of MLP depth (Figure F.21), Any blowup induces catastrophically cascading updates, so that the stability threshold O(n is strictly enforced at all depths, and the loss converges to its kernel limit moderate width for MLPs with at least layers.
with
neh
with
12
12-3
13
30-2 3017
12:3 00-4
13-2
18:22 LR (WDDDTHT710))
UR (NDTH711)))
L8 1MDTHT1110
Figure F.21: (MSE loss on MNIST approximately transfers under ((==1) learning rate scaling) Both the optimal as well as the maximal stable learning rate approximately transfer under global (n==1) learning rate scaling when training 2, 3 10 layer MLPs (from left to right) with MSE loss on MNIST. Loss is not improving as feature learning is lost under O scaling. Especially layer nets, the input layer is learning features at small width, but not at large width, so that the loss worsens with width. In 3 and 10 layer MLPs, on the other hand, the loss quickly converges with width, which suggests fast convergence to the kernel limit under MSE loss.
F.5.4 Large optimal learning rates under CE loss
As theoretically predicted in Remark 14, 2-layer MLPs transfer the optimal and maximal stable learning rate ?)n 2 (n2²), where the input layer learns width-independently
100
Tuknioo 600 100
Thaning
18-4
18
12-1
1pdd
13"
13n4
1 0
UR WDDTH)C0
LR WWITT"HHH0
Tornin acCem001
(MDTH22567008
Figure F.22: (Deeper nets follow infinite width theory increasingly accurately) Training accuracy after epoch of training MLPs with 2, 3, 4, 6, and 10 layers (from top-left to bottom right) on MNIST. While 2, and 4 layer MLPs staf-stabilize under large learning rates (nn)) and approximately transfer the optimum as well as max-stan max-stable learning rate, in 6, 8 and 10 layer MLPs it becomes increasingly apparent that the maximal stable learning rate transitions towards ( (n111) to prevent hidden layer blowup, which also forces the optimal learning rate to be (n2 for at least feature learning in the hidden layers. Hence the theoretical activation stability predictions hold more accurately in deeper nets, with too many hidden layers to stabilize blowup in all of them.
For 3-layer RELU MLPs with CE loss on CIFAR-10, the maximal stable learning rate and the optimal learning rate transfer over many widths before beginning to shrink at width 16384. At moderate width and depth, strong self-stabilization mechanisms such as activation sparsification stabilize an initial catapult at large learning rate (Figure E.6).
In increasingly deep networks, the maximal stable learning rate scaling 7, O(n11/2) becomes increasingly pronounced, as it becomes increasingly difficult to stabilize activation blowup in an
58

---

increasing amount of hidden layers, and width-indeepedent learning of increasingly many hidden layers dominates (MNIST Figure F.22, CIFAR-I0 Figure F.23).
Observe in all figures here that the optimal learning rate saturates at the maximal stable learning rate at sufficient width.
JAT 4071100007100
Figure F.23: (Hidden-layer stability determines learning rate scaling in deep MLPs for SGD on CIFAR-10) MLPs trained with SGD on CIFAR-I0 10 with 4, 6, 8 and 10 layers (from left to right). First two x-axes are width- indth-independent, last two scaled by n²/² While 4- and 6-layer MLPs self-stabilize sufficiently for approximate transfer under width- hidth-inddeenneet learning rate scaling, 8- and 10-layer MLPs have max-stable learning rate scaling 7n (m 1/2).. The optimal learning rate saturates at the maximal stable learning rate at sufficient width.
F.6
Adam stabilizes updates similar to CE loss
We first show juP baseline for learning rate transfer in MLPs on MNIST (Figure F.24).
1.00
100
width
width
Training accuracy
40916
Training accuracy
98
1090
16384
1C304
98
64
94
10
10
10
10-2
(WIDTH/256)"10)
(MDDTH256))10)
100
100
width
wieth
266
Training accuracy
90
1024
4095
Training accuracy
41245
13184
142114
96
94
(WIDTH/256)110)
(WIDTH258))(0)
Figure F.24: (uP baseline for transfer) 8-layer MLPs trained on MNIST with SGD (top) and ADAM (bottom) under CE loss (left) and MSE loss (right). No systematic learning rate shifts in μP; saturating drifts may occur. Transfer and monotonic improvement looks less noisy under MSE loss.
Controlled divergence through stabilized updates. The closest clean optimal learning rate exponent of 8-laaye MLPs trained with Adam under both MSE loss as well as CE loss is - for most of the evaluated image datasets (Figure F.25). Validation-optima learning rates tend to be larger than train-optimal learning rates, suggesting a gellraenera bias of large learning rates. The fact that Adam with MSE loss (Figure F.26) can have optimal learning rates as large as CE loss indicates that the crucial effect of CE loss in SGD stabilizing the updates. A crucial difference to SGD that activation blowup does not affect the updates in Adam since the gradient is normalized. For SGD, exploding gradients induce even larger explosion in the next forward pass, which in turn induces even larger explosion in the next backward pass. Hence, without activation stability, even the divergence
59

---

exponent grows over time in SGD resulting in catastrophically cascading updates. For Adam, on the other hand, gradients are normalized, so that the forward pass always accumulates the same width-dependent exponent that is stabilized when passed through the softmax. Thus under sufficient numerical precision, from a stability point of view, Adam can even tolerate larger learning rates than the hidden-lan layer feature learning ?/, = O and the optimal learning rate may also be pushed toward input layer feature learning.
In shallow networks the input layer enforces slower optimal learning rate decay. As for SGD, with increasing depth, the optimal learning rate decays faster, from around 7h, O (m/1/2) in shallow networks to small 17, = (n ') in deep MLPs in both MNIST (Figure F.27) and CIFAR-10 (Figure F.28), and in both train and validation accuracy (Figure F.29). Figure F.30 confirms the hypothesis that the optimal learning rate is subject to opposing objectives for width-independent learning of input-like and hidden-like layers by showing that 3-layer MLP trained with Adam with fixed input layer follows the clean exponent T7 R (n 1), where the hidden and output layer learn fully -ndte-inddeendntyy
Less clear maximal stable learning rate. For Adam, the optimal learning rate typically does not saturate at the maximal stable learning rate. Instead, a regime of suboptimal large learning rates emerges where its moments are already harmed (Kalra and Barkeshli, 2024), and the maximal stable learning rate threshold is often less clear cut compared to SGD. For all depths, the instability threshold appears to scale around Th, x O(n O/2-
ADAM, MSE
ADAM, CE
10
Learnnng Rate
-0.93
-1.02
cifar10 mnist
-1355
-1.27
fashionmnist
finyimagenet
10
train-optimal
-0.93
-0.88
-0.93
-0.56
0082
val-optimal
-1.27
11003
-1.02
-0.79
-1.24
4096
16384
256
1024
16384
Width
Width
Figure F.25: (MSE and CE loss share similar optimal learning rate exponents under Adam) Train-optimal (solid) and validation-optima (dashed) learning rate as function of width for several image datasets for 8-layer MLPs trained with Adam under MSE loss (left) and CE loss (right). Generally observe exponents around -1. The MSE-optimal learning rate does not decay faster than under CE loss, which indicates that Adam's 's parameterwise normalization of the gradient stabilizes exploding updates, recovering feature learning with n, z O(n=) under both MSE and CE loss.
Mm
Figure F.26: (Adam stabilizes the backward pass even under MSE loss) MLPs trained with ADAM on CIFAR-10 under MSE loss with 2, 3, 6, layers (from left to right). 2 and layers show approximate transfer under n 11/ learning rate scaling, 6 and 8 layers show approximate transfer under n
60

---

1101
10
1
MO0TH-000
LR "NUDT-"(U0
67
Teaking 20000000
Figure F27: (Learning rate transfer in deep MLPs for ADAM on MNIST) MLPs trained with ADAM on MNIST with 2, 3, 4, 6, 8, 10 layers (from top left to bottom right). In the first row, the x-axis is width-dependentll scaled to show approximate transfer under n n/1 learning rate scaling. In the bottom row, the x-axis is width-dependently scaled to show approximate transfer under 7/, ~ (n )) Observe the optimal learning rate scaling transitioning from larger than O (n=11/2 in 2-layer MLPs toward at most (n with increasing depth.
Figure F.28: (Learning rate exponent 7/, O(n 1) for ADAM deep MLPs on CIFAR-10) MLPs trained with ADAM on CIFAR 10 with 2-layer random features, 2, 8, 10 layers (from left to right). The first 2 x-axes show approximate transfer under n 112 learning rate scaling, the last 2 under n- learning rate scaling. As for SGD, in deeper nets hidden-layer widdh-iindependence dominates input-layer width-independenne and induces optimal learning rate scaling Th & O (n)
LR D.MMTTOON'SS
L11 11:3007-00000000
Figure F.29: (Transfer in validation accuracy in deep MLPSs for ADAM on MNIST) Validation accuracy of MLPs trained with ADAM MNIST with 2 layer random feature, 3, 8, 10 layers (from left to right). Validation-optimal learning rate deep MI.Ps scales as 7h Onn-1). layer RF and layer nets appear to approximately transfer under 7n zs ((n but lose monotonic improvement and predictability at scale.
61

---

80
width
widtt
Training accuracy
Training accuracy
1024
4090
2048
113004
4036
5102
18994
48
10-3
10
10-3
10-2
101
10-5
10-5
10-
10-33
10-22
10-1
10'
(WIDTH2588)
(WIDDTH/56))111)
Figure F.30: (Trade off between input- and hidden-layer widtt--ndeependenc)) 3-layer MLPs trained with ADAM on CIFAR-10 10 (left) and not training the first layer (right). 3-layer MLPs approximately transfer under 7n = Ön 1/2),, being pushed toward input-layer feature learning. As there are conflicting goals like preserving input layer feature learning, 3-layer MLPs with fixed input layer follow the widtt--independent exponent 7/2 = O(n (() that yields hidden-and output-layer width-independent feature learning.
F.7 Effective update parameterizations beyond JP
The logit updates be decomposed into
fe(E) fo(E) = W (E) + AW/+t of
for arbitrary inputs & E dRdio and - Xv (Et))
In this section, we consider vision and generated data sets in the regime > dost First note that under large last-layer initialization (W );g- ~ NO,, 1 as in SP, fully width-indeepnddent training dynamics are impossible, since width-independent feature learning ArF = O(1) implies logit blowup through the term Wh+*1 2 = (n2 for both SGD and Adam. The fact that logit blowup does not prevent stable training under CE loss explains why we can achieve non-vanishing feature learning under SP last-layer initialization. When dropping the logit stability constraint, we can ask which is the optimal layerwise learning rate scaling under standard last-layer initialization. Following the μP desiderata, we still want to effectively update all layers, meaning non- vanishisiiin effect of the weight updates in each layer on the output function. With the correct choice of layerwise learning rates, we can still satisfy these desiderata for all scalings of last-layer initialization variance, which also implies that there is not a unique abc-equivalence class fulfill these effective update desiderata when not requiring logit stability. We will see that SP full-align in Everett et al. (2024), which just uses the μP layerwise learning rates for SP initialization (which they promote as their overall best-performing parameterization without identifying stability under logit blowup as the key mechanism), fulfills these desiderata, except for vanishing last-layer update effect on the output function. We will introduce another variant with larger last-layer learning rate that recovers effective updates of all layers. For avoiding confusion with SP, meaning using a global learning rate, and with pPP meaning also achieving hidth-independdnce in the logits, call this last variant Maximal Update Parameterization under Standard Output-laye layer Initialization (MUSOLI).
For deriving the optimal layerwise learning rate exponents, first consider the scaling of hidden-layer pre-activation updates 5h,, E [2, I.), and input-layer pre-activation updates ōh, (Yang and Hu, 2021, 51),
(E)
θ (WWjdz! 1-1 + +mxi-1 oh!
((x)
8h (E)
0 @ 0h
0f
(f1-1) (
0(1)
where holds that af/ah/ 0(05/ax4) W = @(W T 11111(x*)) (at latest in the second step) ( (Yang and Hu, 2021, p. 52). Hence the correct /-th layer learning rate yp for
62

---

achieving a width- -edth-indeeendent effect on the next layer's pre-activations needs to cancel out the backpropagated gradient scaling O /h and for hidden layers additionally the LLN-like scaling from the inner product between activations. As we still require activation stability 2f = ( we have 0 /0 ( min(bb)11ett+)), While under standard jPP. it holds that 0f/Ox² O the changed gradient scaling must be counteracted by choosing hidden layer learning rate Y mm 1 12)_11)))...).)) I [2, L], and input layer learning rate 7n O(n 1 (1)). In words. under larger last-layer initialization or learning rate, the hidden and input layer learning rates should be scaled down by the same amount. Finally, SP-full-align achieves width-independent effect of the last-layer weight updates the logits. But as the width-independent feature updates Az = (1)) induce logit blowup W6 Aaf (m22/2), the effect of the last-layer weight updates the softmax output actually vanishing. For last-layer weight updates to affect the softmax output in the same scaling the updates propagated forward, the last-layer learning rate needs to be 77++1 O On--)z))) bz-1),
hence CL++1 b Hence MUSOLI is defined as SP-full-align but setting 9/++1 = Om
This last-layer learning rate is larger than jP or Everett et (2024) under standard last-layer initialization 1 1/2, but necessary for fulfilling the desideratum that the weight updates in all layers affect the output function non-vanishingly.
SGC, C.03"""'..0. layer
Input taper
|iolden Iaper
Cltpol layer
updito stop
1824 Witth
2549
1024
2641
1024
2040
Witt
10dth
SGC, Ir 0.03""""0.00 10.0, 3-layer
12
Figure F31: (Coordinate check for pP for SGD on CIFAR-10) yP induces fully width-indeeendent update dynamics.
SGC, C.03"'"'0.0. 3-laer
I'puut laper
Liddtin loyer
Ourput byw
updio stop
1019
1079
13"4
1524 Witt
Wath
With
SGC, Ir 0.03""""0.0, layer
DDttur
3
11
Figure F.32: (Coordinate check for SP-full-align for SGD on CIFAR-10) Effective updates |AW''' | and activation updates ||Ar" ||MM as a function of width. The theoretically predicted scaling exponents hold: All layers update widthh-independentl. but due to the large last- layer initialization, the activation updates correlated with W propagated forward induce output logits exploding as W Sx. O(n /2). This motivates increasing the last-layer learning rate to ML+1 (nn-1/2) so that last-layer updates contribute with the same scaling. Note that absolute terms, the updates are much smaller than under μP (Figure F.31).
63

---

By definition, the effective update and propagating update terms in all layers scale width- independently in μPP (Figure F.31). For SP-full-align, Figure F.32 shows that indeed all weight updates behave width-independently, but the output logits are dominated by the activations propa- gated forward as W = ((n 1/2), since óx and W are highly correlated. Consequently, the last-layer updates have vanishing effect on the output function, which induces width dependence. By additionally scaling up the last-layer learning rate 1)r++1 (n-11/2), the logit scaling exponent ((n
in the term W óx (22 is matched in the last-layer update term AW of so that 61+1 = 1/2 and CL+1 = 1/2 recovers a balanced influence of all layer updates in the softmax output.
Figure F.33 and Figure F.34 show that after single-pass SGD or Adam, for both SP-full-align and MUSOLI the optimal learning rate shrinks with width for both generated 2-class multi-index teacher data as well as MNIST. The optimal learning rate exponent is often closer to -0.5 as we consistently observe under MSE loss, preventing logit blowup. Figure F.35 shows the same for CIFAR- 10. This behaviour persisting across data sets suggests that neither SP-full-align nor MUSOLI can be expected to transfer the optimal learning rate in general. An interesting question for future work remains why logit divergence introduces a width-dependence in the optimal learning rate these parameterizations.
1.20
3000
Figure F.33: (Effective update variants do not transfer optimal learning rates on multi-index data) Training accuracy of 8-layer MLPs trained for 1 epoch multi-index teacher data under CE loss (top) and MSE loss (bottom) with SGD in SP-full-align, SGD in MUSOLI, Adam in SP-full- align and Adam in MUSOLI (from left to right). In all cases, logit blowup is avoided by optimal learning rates shrinking n/, O(n -1/2). Under CE loss the maximal stable learning rate remains width-independe for SGD under MSE loss the maximal stable learning rate decays as n- 1/2 -1/2 as necessary for stability.
Figure F.34: (Effective update variants do not transfer optimal learning rates on MNIST) Training accuracy of 8-layer MIPs trained for epoch on MNIST under CE loss with SGD in SP-full-align SGD in MUSOLI, Adam in SP- and Adam in MUSOLI (from left to right). In all cases, the optimal learning rate decays with width, while the maximal stable learning rate stays constant.
64

---

Wdth
Tral 9R OACORLS
2.4 13-1 6x 1911
2x
10
12
12
12"
0R (M10THO90)16.00
LR (METH266)"10.
LR (MITTHOSR""-00)
250
m03h
190
1.01
1004
Toa 1a1g ocuniie
2008
185
1RT (NIDT12200 12.
LB (WITTEE0000055
ER (M0TTE0500010055
Figure F335: (Effective update variants for SGD on CIFAR-10) MLPs with 2, an 6 layers (from left to right) trained with SGD on CIFAR-10 in pP (top) versus SP full-align from Everett et al. (2024) (2nd row) versus SP full-align with larger last-layer learning rate (MUSOLI) (bottom row). While yP transfers with low variance as expected (left), μP with large standard last-layer initialization b = 1/2 and large last-layer learning rate CL+11 = 1/2 (right) have a non-trivial optimal learning rate scaling between (n21/ and ( while the maximal stable learning rate scales width indth-independently.
As expected from parameterizations in the controlled divergence regime, Figure F.35 also shows that the maximal stable learning rate scales budth-independentte since activation and gradient stability preserved. Over the course of 20 epochs, the training dynamics under large learning rates in MLPs with at least layers are stabilized and the optimal learning rate indeed scales aedth-inddeendeentl under standard last-layer initialization. Hence width-deppnndnnce in parameterizations can induce optimal learning rate scaling that varies over course long training. But often the optimal learning rate scales like the maximal stable learning rate. In such cases our theory is predictive. For MSE loss, Figure F.36 shows that logit divergence needs to be avoided through optimal and max-stable learning rate scaling Th O(n - 1/2), irrespective of single multi-epoch settings, because training diverges in the first steps under larger learning rates. This shows that /PP is necessary for making MSE loss a viable alternative to CE loss, avoiding logit divergence while recovering feature learning. Under CE loss, SP full-align and MUSOLI are more robust to poor tuning of the learning rate than μPP both in terms of training and test accuracy (Figures F.37 and F.38). We leave closer analysis of the multi-epoch setting to future work.
65

---

LAT DRIIHDMTTEE
LAT DRRIHNMTTTE
Figure F.36: (Effective update variants with SGD under MSE loss avoid logit blowup) Training accuracy of 2-layer, 3-layer linear, 6-layer and 8-layer MLPs (from left right) trained with SGD for I epoch (top) and 20 epochs (bottom) on CIFAR- 10 SP full-align from Everett et al. (2024). Optimal learning rates shrinking as 7/ = O(n ((/)) persists, avoiding logit blowup through w/L+1 Az Only in 8-layer MLPs is the optimal learning saturating the width- width-independent stability threshold.
with
60
(MDDTH226)0000
(WIDTH2W0""00
L8 "MMDTH256)
with
2011
Traking
10-
10-2 18-2
16-4
8'-3
19.7
UR (WD1H2266935)
1R (WDD12250010
111 (W10 (M1250/220000 0.01
The tha MOCASCC
w0th
1014
13:2
Figure F.37: (Effective update variants for SGD on CIFAR-10 after convergence) MLPs with 2, an 6 layers (from left to right) trained with SGD in uP (top) versus SP full-align from Everett et al. (2024) (2nd row) versus SP full-align with larger last-layer learning rate (MUSOLI) (bottom row) as in Figure F.35 but trained for 20 epochs. After sufficiently long training the large learning rate dynamics stabilize MUSOLI so that the optimum indeed scales brdth-indeeeeneenlyy MUSOLI strictly dominates original μP in training accuracy, and robustness to badly tuned learning rate is strongly improved under SP last-layer initialization compared to original μP. In sufficiently deep MLPs, the larger last-layer learning rate barely matters, but in 2-layer nets SP-full full align avoids output blowup and feature learning by transferring under 7) O(n 1/2).
66

---

with
wit
Vdlitation Lecumary
12-
10.
12
12
30
LR. (MDTH4000"000)
18 "MCTTH2266
L3 (M10THO5""
455
1904
1904
Wildoonn
8
100M
with
VAGOOA on MCD01000
-013
10 (MID 1-30)
1.2 (MICTAA5R))00
1.3 (XTTTTPPS0)"C0
Figure F.38: (Test accuracy of effective update variants for SGD on CIFAR-10 after convergence) Test accuracy 2-layer, 3-layer and 6-layer (from left right) MLPs trained with SGD for 20 epochs on CIFAR- 10 in μP (top) versus SP full-align from Everett et al. (2024) (2nd row) versus SP full-align with larger laye-l learning rate (MUSOLI) (bottom row). The validation-optimal learning rate scales width-independently in all cases. Observe that, while all variants generalize similarly well, the susceptibility to poorly tuned learning rates much larger jPP than under parameterizations with large last-layer initialization.
For ADAM, the gradient normalized in the backward pass, so that input- and hidden-layer learning rates remain the same in jP under large last-layer initialization. This is again equivalent to the SP full-align parameterization from Everett et al. (2024). The logit update term W Ant ((n should again be balanced with larger output layer learning rate M)++1 ((n-1/2) if the weight updates all layers should have non-vanishing effect on the softmax output in the infinite- width limit (MUSOLI). Figure F.39 shows that nonlinear networks trained with Adam and large last- layer initialization already tend to transfer better under MUSOLI than under SP full-align after epoch. Linear networks again have smaller optimal learning rate exponent, indicating that avoiding logit blowup improves over feature learning in this case, where feature learning does not even add expressivity. Generalization, learning transfer and learning rate sensitivity after 20 epochs tends to be similar in all 3 considered parameterizations in deep ReLU MLPs (Figure F.40), showing again that parameterizations with logit blowup are a viable alternative.
Especially in deep ReLUU U MIPs, the last-layer learning rate does not seem have a big impact, and SP full-align and MUSOLI overall behave similarly for both SGD and Adam.
67

---

14 11:JJ01--001
L4 N.--WAreeeeer
100
Figure F.39: (Train accuracy of effective update variants for ADAM on CIFAR-10) Train accuracy of 2-layer, 3-layer, 6-layer and 3-layer-linear MLPs (from left to right) trained with ADAM for epoch CIFAR-10 in juP (top row) versus SP full-align from Everett et al. (2024) (2nd row) versus MUSOLI (bottom row). The learning rate transfers irrespective of the architecture in jPP. Large last-layer learning rate improves transfer in MUSOLI over SP full-align. The optimal learning rate scales as (n (//3) in both parameterizations with large last-layer initialization, as feature learning does improve expressivity.
V000
Uf-re 20160/710
Un Un-WoTiikrn
07 00-180TH50005
15 10
Figure F.40: (Test accuracy of effective update parameterizations for ADAM on CIFAR-10 after convergence) Test accuracy of 2-layer, 3-layer, 6-layer and 3-layer-linear MLPs (from left right) trained with ADAM for 20 epochs on CIFAR-10 in μP (top row) versus SP full-align from Everett al. (2024) (2nd row) versus MUSOLI (bottom row). The validation-optimal learning rate scales width-indepen bndth-independently in all ReLU MLPs with at least layers. 3-layer linear networks clearly transfer under 7/, = (1/2 in SP full-align and MUSOLI, as for sufficient width learning features does add expressivity, and instead avoiding logit blowup dominates the learning scaling.
68

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In the abstract and introduction, we tried summarize and contextualize our
contributions while properly acknowledging related work to the best of our ability.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
Justification: We clearly state the caveats of every reported result. We mostly diseuss limitations and avenues for future work in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms
and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofss
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We detail all assumptions and the proofs of our theoretical results Ap- pendix C.
Guidelines:
The answer NA means that the paper does not include theoretical results.
69

---

All the theorems, formulas, and proofs in the paper should be numbered and cross
referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes)
Justification: We detail all necessary information in the main paper and Appendix D. Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (c.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-souree dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results. 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Upon acceptance, we will release open source code to reproduce all of our experiments on GitHub.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cd
iol for more details.
70

---

reproducible, they While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ //nips. cc/public//aeee sellic//uuides gewguuie//aee "CodeSubmisse ssionpolicyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We detail all necessary information in the main paper and Appendix D. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core of the paper to level of detail
that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: If not stated otherwise, error bars denote the region between the empirical 0.025- to 0.975-quantiles from 4 random seeds (affecting initial weights and data shuf- fling/generation). Due to computational constraints, we only train neural networks on several random seeds for subset of the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula,
call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2- 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
71

---

For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [No]
Justification: Single training runs of 8-layer MLPs of width 16384 including tracking all relevant statistics as well as our 1.4B GPT model of width 4096 run within less than 24 hours on single Nvidia A100 GPU. We typically trained MLPs up width 4096 on single Nvidia Geforce Rix 2080 Ti within less than 24 hours. Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper). 9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https / /ewrris...pp..e/ Answer: [Yes]
Justification: The research in this paper conforms in every respect with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid-
eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed? Answer: [Yes]
Justification: This paper aims to advance the understanding of standard deep learning practice. We do not foresee any societal impacts beyond the dual use considerations that apply to the whole field of machine learning research.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), (rouups).rrivcyy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
72

---

11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper studies pre-existing training procedures on established datasets, We do not release any data or high-risk models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the standard CIFARIO (Krizhevsky et al., 2009), MNIST (Deng, 2012) and DCL.-Basellng
al., 2024) datasets following the standard practice. We also cite the Python assets Py Torch (Paszke et al., 2019) and LitGPT (Lightning AI, 2023) that use as a basis for our experiments. Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include a URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi twroiethoodeee.e com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mooell as part of their
submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
73

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLMs were not used in a way that impacts the core methodology, scientific rigorousness, or originality of the research. Guidelines:
The answer NA means that the core method development in this research does not involve LLMMs as any important, original, non non-standard components. /205LLLM))
Please refer to our LLM policy (https://nn) //te://neeuraae for what should or should not be described.
74


Reviews:
[{"review_id": "d8ReGzS9pG", "reviewer": "Reviewer_YurV", "summary": "The authors try to answer why standard parameterization remains stable and effective at large learning rates, despite the theoretical predictions. While the catapult effect and edge of stability theory does not fully explain this phenomenon, the authors find a reason from the choice of cross-entropy (CE) loss as a loss function. The authors insist that CE loss allows stable training under large learning rates even when the output logits diverge.", "questions": "1.\tIs inducing the maximal stable learning rate from an optimal learning rate a frequently exercised custom or wisdom in this area?\n2.\tPage 30 mentions that the authors train a single epoch to prevent confounding from multi-epoch overfitting effects. Could you explain more on why this is necessary for the experiments?\n3.\tOne question for my sheer curiosity. It is known that a cross-entropy loss leads to a max-margin solution. What would be the relationship between the learning rate and a converged solution? Beyond a maximal stable learning rate, would the parameter not converge to a max-margin solution?", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n1.\tThe paper is easy to follow even for those who are not familiar with the subject. The informal explanations of the propositions and theorems as well as the proof sketches highly enhance the readability and help the readers understand the intuition behind the theories.\n2.\tThe paper has strong theoretical contributions. The focus on the cross-entropy loss seems very interesting to me.\n3.\tThe experimental results seem to support the theoretical results as well. The experiments seem to be extensive and include examples on large models like GPT.\n\nWeaknesses\n\n1.\tI am not very much fond of the idea of inducing the maximal stable learning rate from an optimal learning rate. The authors briefly explain that Proposition 2 suggests that the idea is true since it is the only setting under which feature learning is preserved at large width in all hidden layers. I only partially agree with the authors’ opinion, and I think authors need to provide a more rigorous explanation to confirm this. The result is practically useful, though.\n2.\tIn Page 30, it says that the authors train a single epoch to prevent confounding from multi-epoch overfitting effects. I personally cannot find a particular reason why this is necessary, since feature learning often occurs after a certain period of training time, especially when the data is complex and grokking occurs.", "was_revised": true, "final_justification": "I am satisfied with their clarification in the rebuttal. So I raise the score.", "rebuttal": "Thank you for your thoughtful review and constructive criticism. We are thrilled that you find our paper **easy to follow even for readers unfamiliar with the theory, with strong theoretical contributions elucidating the important effect that cross-entropy loss has on training stability and with extensive experiments also covering practical LLM settings**. We now address your main points of criticism and answer your questions.\n\n### **Is optimal LR scaling induced by the maximal stable LR scaling?**\n\nFirst, let us clarify the definitions of maximal stable learning rates (**max-stable LR**):\n\n- **Theoretical definition:** Historically, the maximal stable LR for a neural network was defined as the boundary beyond which either the logits or at least one layer's activations become unstable with width. A layer-specific variant defines this as the LR threshold at which that particular layer’s activations diverge.\n- **Empirical definition:** Empirically, the maximal stable LR is naturally defined as the largest learning rate beyond which training diverges, precisely quantified by the learning rate threshold beyond which the network’s predictions degrade to random guessing (e.g., within ~1% of random guessing accuracy).\n\nPrevious work (Yang et al., 2022; Dey et al., 2023; Everett et al., 2024) extensively demonstrates that *when all layers share the same theoretical maximal stable LR exponent, both empirical max-stable-LR and optimal LR scaling closely matches the theoretical max-stable-LR scaling***.** This phenomenon is robustly established through numerous experiments, particularly in the context of $\\mu$P. However, under SP, this discussion becomes much more nuanced since 1) Different layers exhibit distinct theoretical max-stable LR exponents and it is not obvious how to define the theoretical max-stable LR at a network level. 2) It is also unclear what the scaling exponents of empirical maximal stable LR and optimal LR look like.\n\n**Key contributions of our work:**\n\n- **Breakdown of correspondence:** We reveal that under cross-entropy (CE) loss, the empirical maximal stable LR does **not** correspond with the traditional theoretical maximal stable LR. Instead, we find that empirical maximal stable LRs occur at the boundary to the **catastrophically unstable regime where both logits and activations of all layers diverge**, rather than merely at the edge of the stable regime (as previously believed).\n- **Theoretical redefinition:** This observation motivates us to formally redefine the theoretical maximal stable LR explicitly as the boundary to this catastrophically unstable regime.\n- **Optimal LR scaling alignment:** To our surprise, our experiments show that when network dynamics are dominated by a certain layer type (e.g., deep nonlinear MLPs dominated by hidden layers), the empirical optimal LR scaling often saturates towards our definition of theoretical maximal stable LR scaling with respect to that layer type.\n\nWe intentionally **refrain** from making explicit claims about the exact scaling exponent for the optimal LR. Although we narrow its potential range and successfully predict optimal LR values in many scenarios, we acknowledge limitations and the difficulty of fully characterizing conditions under which this alignment breaks down. Properly accounting for finite-width effects remains essential yet highly nontrivial. We believe a rigorous characterization of the optimal LR exponent at finite widths remains challenging, as it likely requires strong assumptions about architecture or data distributions—analogous to assumptions required by neural scaling laws (Hoffmann et al., 2022; Bachmann et al., 2024). Hence, this paper intentionally emphasizes general insights on signal propagation and training stability, deferring deeper theoretical exploration of optimal LR exponents to future work (as explicitly detailed in the subsection *‘Understanding optimal learning rate exponents’* of our future work section).\n\nEmpirically, however, we emphasize that *the optimal learning rate consistently saturates at the maximal stable learning rate at sufficient widths across all our SGD experiments in SP, including GPT-like transformers and deep ReLU MLPs* (summarized in the table in the Response to Reviewer dfje). While the empirical connection between optimal and max-stable learning rates is well-documented, the mechanisms remain poorly understood, independent of model scale (see related *edge-of-stability literature* by Cohen et al., 2021, 2022; Lewkowycz et al., 2020; Cai et al., 2024; Damian et al., 2023, among others). Our findings offer a novel perspective on the advantages of large learning rates for facilitating feature learning at large model scales. Previously suggested explanations include improved generalization through reduced sharpness (Andriushchenko et al., 2023a), a shift in the learning order of patterns (Li et al., 2019), enhanced SGD noise (Keskar et al., 2017), and implicit bias towards sparsity (Andriushchenko et al., 2023b).\n\n### **Multi-epoch setting**\n\nThanks for raising this important question, we will provide this clarification in the main paper. First, we would like to clarify that the notion of feature learning commonly used in the context of grokking - i.e., alignment with task-specific features - differs significantly from our definition of feature learning in the context of width scaling, which specifically refers to feature updates that neither diverge nor vanish as width scales. Under this latter definition, feature learning in SP within the controlled divergence regime occurs after only one step of SGD, without requiring prolonged training.\n\nTo evaluate our theoretical claims, which currently address width dependence without exploring dependencies on training time or data size, we consider the online training setting to be most appropriate. Online training avoids introducing additional overfitting dynamics that have been theoretically shown to accumulate when data points are repeatedly presented (Bordelon et al., 2024). Furthermore, we focus explicitly on the optimization effects of  learning rate rather than its regularizing properties. Recent theoretical work indicates that SGD noise has negligible influence on learning-rate scaling in the online scenario (Paquette et al., 2024). Finally, we note that modern large language models are often trained online, underscoring the practical relevance of our analysis.\n\nBut we do agree that the multi-epoch setting is an interesting and important setting to study in the future, and *will also mention it in the future work section*.\n\n### **Relation to max-margin solutions**\n\nAs opposed to our more architecture- and data-agnostic stability considerations, the convergence results to a max-margin solution are more architecture (often only 2-layer networks) and data-dependent (see e.g. Lyu et al. (2021) for positive and negative examples), and they often study gradient flow. When a neural net converges to a max-margin solution under gradient flow, we expect convergence to continue to hold for ‘small enough’ LRs. Beyond the max-stable LR, the parameters would indeed diverge, because large updates induce divergence in the activations, which leads to a cascading feedback loop between exploding gradients and activations. To provide a concrete example for CE loss, under favorable separability assumptions, consider Theorem 3.2 in Cai et al. (2025) which considers muP with width-independent LR. Here, increasing the LR while keeping training time fixed implies divergence with increasing model scale.\n\nWe sincerely *thank the reviewer for stimulating several changes that will enhance clarity and provide more convincing evidence for our claims*. Overall, we see the revised paper as a necessary intermediate step with fundamental contributions that opens up several exciting avenues for future work. We hope that our changes and clarifications partially address your main concerns. In that case, *we kindly ask you to consider raising your score*. We are happy to elaborate further, if any questions remain unanswered or if there remain further fundamental concerns.\n\n\n###  **Added References:**\n\nAndriushchenko et al. “A modern look at the relationship between sharpness and generalization.” *ICML* (2023a).\n\nAndriushchenko et al. “Sgd with large step sizes learns sparse features.” *ICML* (2023b).\n\nBachmann et al. “Scaling mlps: A tale of inductive bias.” *NeurIPS* (2024).\n\nBordelon et al. \"A dynamical model of neural scaling laws.\" *ICML* (2024).\n\nCai et al. \"Implicit Bias of Gradient Descent for Non-Homogeneous Deep Networks.\" *arXiv:2502.16075* (2025).\n\nCai et al. “Large stepsize gradient descent for non-homogeneous two-layer networks: Margin improvement and fast optimization.” *NeurIPS* (2024).\n\nHoffmann et al. “Training compute-optimal large language models.” *arXiv:2203.15556 (2022).*\n\nKeskar et al. “On large-batch training for deep learning: Generalization gap and sharp minima.” *ICLR* (2017).\n\nLi et al. “Towards explaining the regularization effect of initial large learning rate in training neural networks.” *NeurIPS* (2019).\n\nLyu et al. \"Gradient descent on two-layer nets: Margin maximization and simplicity bias.\" *NeurIPS* (2021)."}, {"review_id": "2E0lOjCKwp", "reviewer": "Reviewer_dfje", "summary": "This paper investigates why standard parameterization (SP) with He initialization remains effective at large learning rates, despite infinite-width theory predicting instability. The authors demonstrate that cross-entropy loss fundamentally changes the stability landscape compared to MSE loss, enabling a regime where logits diverge but training remains stable. They provide theoretical analysis and validate their predictions empirically across small tasks on MLPs and Transformers.", "questions": "- How does numerical precision affect the controlled divergence regime in practice?\n- What are the practical implications beyond explaining why SP works?\n- What happens at very large widths or training times?", "limitations": "The authors briefly mention numerical considerations but do not adequately address how floating point limitations affect the practical applicability of their theoretical results.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strengths\n\n- The paper addresses a fundamental puzzle in deep learning theory - why standard parameterization works despite theoretical predictions of instability. \n- The key insight about CE loss enabling a controlled divergence regime is novel and provides a principled explanation for empirical observations.\n- Empirical validation across multiple architectures (MLPs, Transformers), optimizers (SGD, Adam), and datasets strengthens the theoretical claims.\n- The connection between width-scaling theory and practical learning rate selection cloud be valuable for practitioners.\n\n### Weaknesses\n\n- I found the paper a bit hard to read and missed a clear \"golden thread\".\n- Proposition 1 (line 208) is marked as \"informal\" but is central to the paper's claims. The formal statement in Appendix C.3 should be in the main text.\n- Numerical precision issues when logits diverge are mentioned only briefly (line 354). This is a critical practical concern that undermines the applicability of the controlled divergence regime.\n- Section 5.2 on SP-full-align feels disconnected from the main narrative and the claim about \"breaking learning rate transfer on image datasets\" (line 319) needs more thorough investigation.\n- While the paper explains why SP works, it doesn't lead to better training methods or actionable insights for practitioners. The contribution is primarily explanatory rather than enabling new capabilities.\n- In general, I found the scope of the paper a bit unclear and it lacks a strong empirical validation that the width-scaling considerations provide good predictions of the maximal stable learning rate and that the maximal stable learning rate often dominates the optimal learning rate.", "was_revised": true, "final_justification": "The authors addressed most of my critics and answered my questions. I'm still concerned about the weak empirical evaluation and the presentation (\"golden thread\").", "rebuttal": "Thank you for your thoughtful review and constructive criticism. We are delighted that you agree our results **resolve a fundamental puzzle in deep learning theory by providing a principled explanation through the controlled divergence regime, and that you also recognize their potential practical impact, validated by our empirical evaluations, on practical learning rate selection.** Below, we address your main points of criticism and answer your questions.\n\n**Scope.** As you acknowledge, our results resolve a **fundamental and long-standing problem in deep learning theory**: we want to understand finite nets in SP, but previous work had only proposed proxy models that are easier to study but differ qualitatively in their feature learning ability. As we show, practical neural networks operate in a understudied controlled divergence regime. Our theory resolves this fundamental puzzle, closes the gap between infinite-width theory and width-scaling practice and hopefully enables the theory community to generate more practically relevant insights. \n\n*“It is unclear that maximal stable learning rate often dominates the optimal learning rate.”*\n\nWe acknowledge that our discussion in the paper might have caused some confusion regarding the definitions of the maximal stable learning rate (**max-stable LR**). To address this, we clarify the definitions explicitly below, and will incorporate these clarifications into the revised manuscript.\n\n**Theoretical max-stable-LR:** Historically, the theoretical maximal stable LR of a neural network is defined as the boundary beyond which either the logits or at least one layer's activations become unstable with width.\n\n**Empirical max-stable-LR:** A natural empirical equivalent of a max-stable-LR of a network is the largest learning rate beyond which training diverges, precisely quantified by the learning rate threshold beyond which the network’s predictions degrade to random guessing (e.g., within ~1% of random guessing accuracy). **Under this definition, the max-stable LR of the network clearly provides a strict upper bound on the optimal learning rate.**\n\nWe appreciate the reviewer’s feedback, which helps us refine and clearly convey these foundational concepts in our paper.\n\n*“Empirical evidence for learning rate scaling in SP.”*\n\nThank you for encouraging us to provide a more accessible summary of our empirical results. For convenience, we have included tables of learning rate (LR) exponents below, and we will ensure they appear at least in the Appendix to provide an accessible overview, referenced clearly from the main paper.\n*Additionally, we now provide the optimal and maximal stable learning rate exponents on two further datasets: FashionMNIST and TinyImagenet. Notably, for every dataset considered, the maximal stable learning rate exponent consistently aligns closely with our theoretical predictions, namely -0.5 for CE loss and -1 for MSE loss.*\n\nThe optimal LRs generally follow the same scaling, though they occasionally exhibit strong finite-width effects. In our response to Reviewer YurV, we summarize relevant existing literature that establishes the connection between optimal and maximal stable learning rates.\n\n| Dataset | Architecture | Optimal LR exponent | Maximal stable LR exponent |\n| --- | --- | --- | --- |\n| tinyimagenet | 8-layer MLP | -0.3347 | -0.3347 |\n| fashionmnist | 8-layer MLP | -0.55 | -0.7 |\n| mnist | 8-layer MLP | -0.46 | -0.54 |\n| cifar10 | 8-layer MLP | 0.07 | -0.37 |\n| dclm-baseline | GPT | -0.38 | -0.38 |\n\nEven when the optimal LR exponent does not closely track the maximal stable exponent - as observed with CIFAR-10 - the learning rate curves (such as Figure F.19, summarized in the table below and averaged over four random seeds) demonstrate that the **optimal LR exponent saturates toward the maximal stable exponent at width 16384.** This trend is expected to continue at larger widths, where our theory holds even more accurately. All corresponding learning rate curves will be included in the revised Appendix.\n\n**CIFAR-10, SGD, SP, CE loss**\n\n| Width | 2.05e-02 | 4.10e-02 | 8.19e-02 | 1.64e-01 | 3.28e-01 | 6.55e-01 |\n| --- | --- | --- | --- | --- | --- | --- |\n| **256** | **42.57** | 41.51 | 40.95 | 38.26 | 23.62 | 10.00 |\n| **1024** | 44.33 | **44.54** | 44.45 | 42.81 | 10.00 | 10.00 |\n| **4096** | 46.13 | 46.57 | **46.69** | 10.00 | 10.00 | — |\n| **16384** | 46.10 | **46.92** | 10.00 | 10.00 | — | — |\n\n*“How SP-full-align relates to the main narrative.”*\n\nSP-full-align connects to the main narrative because, like SP with CE loss, **this parameterization also operates in the *controlled divergence regime*:** training remains stable despite systematic logit divergence. The key distinction from standard SP is that, due to its layerwise learning rates, all layers exhibit effective feature learning similar to $\\mu$P—a property often linked to improved performance (e.g., Bordelon et al., 2024; Kunin et al., 2024). Additionally, we note that SP-full-align was originally motivated by alignment considerations, which we disprove in Section 3.1. Instead, our work provides a corrected explanation for the empirical behavior observed in networks trained with SP-full-align.\n\n*“More evidence for the claim \"breaking learning rate transfer on image datasets under SP-full-align\" (line 319).”*\n\nThank you for suggesting this valuable clarification. *To further strengthen our claim, we conducted additional experiments training 8-layer MLPs on FashionMNIST and TinyImagenet.* Consistent across all our SP-full-align experiments (summarized in the table below), we find that while the maximal stable learning rate remains remarkably width-independent - as predicted by our theory - optimal learning rates decay across all datasets considered. We previously reported a subset of these results in Appendix F.8, but agree that presenting them in an accessible tabular form significantly enhances clarity. \n\n| Dataset | Optimal LR exponent | Maximal stable LR exponent |\n| --- | --- | --- |\n| fashionmnist | -0.35 | 0.0 |\n| tinyimagenet | -1.45 | 0.0 |\n| mnist | -0.25 | 0.0 |\n| cifar10 | -0.33 | 0.0 |\n\n**Provide formal result in the main paper.** We would also have liked to present the full theoretical results in the main paper. But unfortunately this requires introducing all of the Tensor Program definitions and machinery, which would take too much space and would significantly harm readability without significant information gain. We instead opted for distilling the necessary knowledge into a more accessible and practitioner-oriented introduction of the main scaling arguments.\n\n**\"Golden thread\" of the paper.** Our intention was precisely to provide such a cohesive narrative by first clearly formulating the puzzle of large learning rate stability in SP, demonstrating that existing explanations fall short of addressing the observed discrepancies, and subsequently identifying the underlying mechanism that resolves this puzzle. Finally, we discuss the fundamental implications of stable training despite logit divergence. Given that other reviewers noted the paper is already well-written and accessible, we would particularly appreciate additional concrete suggestions on improving clarity further, as ensuring accessibility to a broad audience was one of our main concerns while writing this paper.\n\n**Larger width and longer training time.** We emphasize that our experiments already span a sufficiently large width to be practically meaningful: the hidden dimension of the LLAMA-3 model with 405B parameters is 16384, exactly matching the width used in our MLP experiments. Due to computational constraints in our academic environment, we are unable to provide experiments at even larger scales. However, we expect our theoretical predictions would only become more accurate at widths beyond this already substantial scale. Indeed, **it is particularly noteworthy that our theory already yields accurate exponent predictions at these practically relevant widths.** We also expect our experiments to hold robustly over longer training time in the online setting. The discussion in the multi-epoch setting is more nuanced as detailed in our response to reviewer YurV. \n\n*“Contribution is primarily explanatory rather than enabling new capabilities.”*\n\nWhile we strongly believe that the explanatory nature of our work is valuable in its own right, there are several concrete *practical implications of our work. Due to space constraints we have provided a detailed response to this point in our response to Reviewer r9tC*.\n\n*“How does numerical precision affect controlled divergence?”*\n\nNote that our results highlight that *neural nets in standard practice operate in the controlled divergence regime* so we should expect training instability due to finite precision  — which is precisely what has been established empirically (Chowdhery et al., 2022; Wortsman et al., 2023). Our results precisely identify the mechanism responsible: practically relevant large-scale neural networks naturally operate at the boundary of the controlled divergence regime, allowing persistent hidden-layer feature learning *despite logit divergence*. (also see our response to r9tC where we highlight practical implications of this attribution).\n\nWe sincerely *thank the reviewer for stimulating several changes that will enhance clarity and provide more convincing evidence for our claims*. Overall, we see the revised paper as a necessary intermediate step with fundamental contributions that opens up several exciting avenues for future work. We hope that our changes and clarifications address your main concerns. In that case, *we kindly ask you to consider raising your score*. We are happy to elaborate further, if any questions remain unanswered or if there remain further fundamental concerns."}, {"review_id": "xbrUxzzQfW", "reviewer": "Reviewer_r9tC", "summary": "This paper tackles the discrepancy between theoretically predicted optimal (stable) learning rates (LRs) and empirically observed ones as model scale increases under standard parameterization (SP). The authors particularly highlight the role of cross-entropy (CE) loss, which in contrast to MSE loss allows stable training under large LRs even with divergent logits. Interestingly, as shown in the paper, neither long-training dynamics nor the catapult mechanism arguments are able to explain the stability of large LRs in SP. Instead, thanks to the properties of the CE loss derivative, training dynamics remains stable with LR exponent up to $-1/2$, which aligns well with practice. Experiments confirm the theoretical predictions and help further elucidate the intuition behind different choices of architectural elements, optimizers, and parameterizations.", "questions": "Please refer to Weaknesses.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "# Strengths\n\nThis paper provides a solid foundation for the empirically observed LR scaling laws, bridging the gap between theory and practice and enabling better understanding and more accurate prediction of optimal LR values at scale.\nThe idea to particularly consider the role of the loss function seems novel and interesting, which is a solid contribution of this work.\nIt is indeed counterintuitive to anticipate stable training dynamics and even feature learning when model outputs (logits) enter the divergence regime.\nThe authors did a great job of explaining the practical benefits of cross-entropy loss, and I specifically like the part demonstrating that MSE loss can possess similar benign properties under $\\mu$P, which makes it both theoretically attractive and practically usable.\n\nThe authors clearly position their work in relation to prior art, explicitly focusing on missing parts of the puzzle and suggesting possible explanations.\n\nThe theoretical results are supported by extensive practical experiments and detailed discussions.\n\nWriting is clear and consistent. \nThe paper is well-written, although it is still relatively hard to follow by a nonspecialist in deep learning scaling theory.\n\n# Weaknesses\n\nSince I am not an expert in deep learning scaling theory, I cannot fully assess the possible technical limitations of this work.\nFor me, the main weakness of this work is the unclear practical implications of the proposed theory, which has a descriptive rather than predictive nature.\nI would suggest that the authors pay more attention to this aspect in a future revision.", "was_revised": true, "final_justification": "In light of the authors rebuttal and other reviews, I decide to keep my initial high score as I consider this work a strong submission.", "rebuttal": "Thank you for your thoughtful review and constructive criticism. We are delighted about your positive evaluation of our work, acknowledging **clear writing and extensive experiments**, in particular that you find our **theoretical contributions novel and interesting with solid contributions and counterintuitive insights into stable feature learning despite logit divergence**. We were also surprised by this surprisingly simple resolution of a long-standing and seemingly hard puzzle, and are also excited about the potential for practically relevant theory in the controlled divergence regime as well as for the practical potential of using MSE loss with muP. We will now address your main point of criticism.\n\n### **Practical implications**\n\nWe are grateful for this criticism and agree that it is important to further clarify how useful and surprisingly accurate our width-scaling theory is for large-scale model training. We will state these implications more clearly in the revision. We believe that our work has the following practical implications:\n\n1. **Constraining the search space for the optimal learning rate.** \n\nOur experiments (pre-existing as well as novel experiments are summarized in the table in the Response to Reviewer dfje) strongly support the prediction from Proposition 1 that the maximal stable learning rate scales as $\\Theta(1/\\sqrt{n})$. This result provides a concrete constraint on the learning rate search space, significantly reducing the range practitioners must consider. Rather than exhaustive or heuristic searches, practitioners may narrow their search around the theoretically justified scaling, leading to substantial computational savings and more efficient hyperparameter tuning. Furthermore, we observe that, particularly in deep nonlinear architectures, the optimal learning rate often saturates at this maximal stable scaling (see e.g. Figures 5 and 6 (left), Figures F.4 and F.5, we will add all learning rate curves corresponding to the added learning rate exponent table to the appendix). Thus, employing the correct scaling often effectively enables approximate transfer of optimal learning rates.\n\n2. **Potential performance gains at scale using alternative loss functions with muP.**\n\nCorrectly attributing the observed performance gap between MSE and cross-entropy (CE) loss in standard parameterization (SP) to scaling considerations reveals a concrete recommendation: practitioners could leverage stable parameterizations like muP for image datasets or SP-full-align for language datasets to effectively utilize loss functions beyond cross-entropy at large scales. Specifically, your point about correct width-scaling via muP enabling alternative loss functions to excel at large scales is further validated by our extended experiments. *We will extend Figure 7 to width 16384 and add results for FashionMNIST* (which closely resemble CIFAR-10). In the extended experiments, we still consistently observe that CE loss greatly outperforms MSE under SP, whereas MSE consistently outperforms CE loss under muP, with differences becoming particularly pronounced at large widths. These results highlight the potential for substantial performance gains from further investigating the interplay between the loss function and parameterization. They also suggest that it is worth exploring the use of further loss functions in conjunction with muP.\n\n3. **Identifying the correct scaling mechanisms enables principled improvements of scaling practice.**\n\n*Identifying the correct causal mechanism for training stability under large learning rates enables principled and informed exploration of the extended search space of practically relevant potentially best-performing parameterizations for efficiently finding improvements in model scaling practice:* The controlled divergence regime had previously been neglected, even though common scaling practice as well as the best-performing parameterization SP-full-align from Everett et al. (2024) exactly operate in this regime. We now detail two concrete implications together with exciting avenues for future work.\n\n**Identifying the correct mechanism for training instability at large scales.** Training instability due to logit divergence at large scales is a *widely-established empirical phenomenon under SP* (Chowdhery et al., 2022; Wortsman et al., 2023) which has motivated popular interventions like the z-loss (Chowdhery et al., 2022; Wortsman et al., 2023). However, only understanding the underlying causal mechanism and providing the exact width-dependent scaling exponents enables to design **principled interventions**. Our results precisely identify the mechanism responsible: practically relevant large-scale neural networks naturally operate at the boundary of the controlled divergence regime, allowing persistent hidden-layer feature learning despite logit divergence. Accurately attributing this phenomenon to the correct mechanism opens the door for principled and theoretically grounded interventions. Concretely, our results clarify that practical networks trained under SP lie in the controlled divergence regime and therefore inherently cause logit divergence, leading to numerical instability. This provides a concrete hypothesis to test: instability issues due to logit divergence could be systematically resolved with stable parameterizations like muP potentially eliminating the need for ad-hoc interventions like the z-loss. A thorough investigation of this is beyond the scope of the current paper. \n\n**Logit divergence induces overconfidence.** A similar argument applies to uncertainty calibration. Our theory also explains why we should expect *predictions to be increasingly overconfident with increasing model scale and suggests that this may be partially mitigated by considering stable parameterizations like muP.* However there maybe inherent trade-offs between miss-calibration and faster memorization due to logit divergence. Their effect on performance and designing the ideal intervention deserves a more thorough investigation beyond the scope of the current paper.\n\n4. **Faithful evaluation of the accuracy of infinite-width theory.**\n\nAs a first contribution, we provide a refined evaluation of whether the width-dependent exponents of important quantities like effective updates or the maximal-stable learning rate scaling predicted by infinite-width theory accurately hold over the course of training. As opposed to previous work (Everett et al., 2024), we do find that *these predictions hold remarkably accurately already at moderate width ≤512 and over the course of training* (see e.g. Figures 4 and 5, F.4 and F.5), when evaluated with sufficient attention to detail:\n\nWhile in principle the refined coordinate check (RCC) was previously known (Yang and Hu, 2021, Appendix H), we find that the *predicted update exponents are surprisingly accurate in this decomposition already at moderate width $≤512$ and over the course of training* (Figures 2 and 4), *even in parameterizations with width-dependent dynamics like SP*. To lower the hurdle for practitioners to incorporate the RCC in their workflow as a diagnostic tool, *we will make our fine-grained and easily adaptable implementation of the RCC using LitGPT publicly available upon acceptance.* Understanding and correcting miss-scaled update signals through each weight matrix has impactful consequences on performance and trainability at large scale.\n\n**References:**\n\nChowdhery et al. “Palm: Scaling language modeling with pathways.” *JMLR* (2023).\n\nOLMo Team. “2 olmo 2 furious.” *arXiv:2501.00656* (2025).\n\nWortsman et al. “Small-scale proxies for large-scale transformer training instabilities.” *ICLR* (2024)."}]

Author rebuttal:

Meta-review: This paper tackles the puzzle of why standard parameterization remains effective at large learning rates despite infinite-width theory predicting instability. Among several potential explanations, the authors identify cross-entropy loss as a key factor, backed by a theoretical analysis based on Tensor Program. They also present a series of experiments across optimizers, architectures, and datasets to further validate the claim.

The reviewers considered the paper a strong contribution to the scaling theory for deep learning. Some concerns on the clarity, weak empirical evaluation and practical implications were raised, but the authors' rebuttal largely addressed these concerns. While some reviewers found the paper easy to follow, the presentation of the paper could be further strengthened for non-experts. Overall, I recommend acceptance.
Decision: Accept (spotlight)