The Optimization Landscape of SGD Across the Feature Learning Strength
Authors: Alexander Atanasov
Alexandru Meterez
James B Simon
Cengiz Pehlevan
Venue: ICLR 2025 Poster (2025)
Keywords: Feature Learning
Maximal Update Parameterization
Optimization
Edge of Stability
Catapult Effect
Silent Alignment
Primary area: optimization
TL;DR: We empirically sweep over the feature learning hyperparameter for models trained online, identifying a variety of novel scaling relationships with sharpness, optimal learning rate, and performance, and explaining these in a linear network model.

Abstract: We consider neural networks (NNs) where the final layer is down-scaled by a fixed hyperparameter $\gamma$. 
Recent work has identified $\gamma$ as controlling the strength of feature learning.
As $\gamma$ increases, network evolution changes from "lazy" kernel dynamics to "rich" feature-learning dynamics, with a host of associated benefits including improved performance on common tasks.
In this work, we conduct a thorough empirical investigation of the effect of scaling $\gamma$ across a variety of models and datasets in the online training setting.
We first examine the interaction of $\gamma$ with the learning rate $\eta$, identifying several scaling regimes in the $\gamma$-$\eta$ plane which we explain theoretically using a simple model.
We find that the optimal learning rate $\eta^*$ scales non-trivially with $\gamma$. In particular, $\eta^* \propto \gamma^2$ when $\gamma \ll 1$ and $\eta^* \propto \gamma^{2/L}$ when $\gamma \gg 1$ for a feed-forward network of depth $L$.
Using this optimal learning rate scaling, we proceed with an empirical study of the under-explored ``ultra-rich'' $\gamma \gg 1$ regime.
We find that networks in this regime display characteristic loss curves, starting with a long plateau followed by a drop-off, sometimes followed by one or more additional staircase steps.
We find networks of different large $\gamma$ values optimize along similar trajectories up to a reparameterization of time.
We further find that optimal online performance is often found at large $\gamma$ and could be missed if this hyperparameter is not tuned.
Our findings indicate that analytical study of the large-$\gamma$ limit may yield useful insights into the dynamics of representation learning in performant models.

Published as a conference paper at ICLR 2025
THE OPTIMIZATION LANDSCAPE OF SGD ACROSS THE FEATURE LEARNING STRENGTH
Alexander Atanasov":1.24 Alexandru Meterez James B Simon +3.8 Cengiz Pehlevan 2:4,5
1 Department of Physics, Harvard University
2 School of Engineering and Applied Science, Harvard University
3 Department of Physics, University of California, Berkeley
4 Center for Brain Science, Harvard University
5 Kempner Institute, Harvard University Imbue
(atanasov, cptheeovcceeeeeaaaaea haareaeeeeeegg.....rdeeeeed.
harvard.
ABSTRACT
We consider neural networks (NNs) where the final layer is down-scaled by fixed hyperparameter 7. Recent work has identified fy as controlling the strength of feature learning. As increases, network evolution changes from "lazy" kernel dynamics to "rich" feature-learning dynamics, with a host of associated benefits common tasks. In this work, conduct
including improved performance thorough empirical investigation of the effect of scaling across a variety of mod- els and datasets in the online training setting. We first examine the interaction of y with the learning rate 17. identifying several scaling regimes in the 7-17 plane which we explain theoretically using simple model. We find that the optimal learning rate 17" scales non-trivially with y. In particular, '7" о y² when << 1 and 17" & 72/2 when 7 >> 1 for feed-forward network of depth L. Using this optimal learning rate scaling, proceed with an empirical study of the under- explored "ultra-rich" > regime. We find that networks in this regime display characteristic loss curves, starting with a long plateau followed by drop-off, sometimes followed by one or more additional staircase steps. We find networks of different large y values optimize along similar trajectories up to a reparame- terization of time. We further find that optimal online performance is often found at large " and could be missed if this hyperparameter is not tuned. Our findings indicate that analytical study of the large-y limit may yield useful insights into the dynamics of representation learning in performant models.
INTRODUCTION
The study larth-width limits of neural networks (NNs) has led to a variety of insights about deep learning. As network width N tends infinity, networks in either "PyTored standard" parameter- ization or neural tangent kernel (NTK) parameterization converge to kernel methods (Jacot et al.] 2018 Lee et al. 2019; Sohl-Dickstein et al., 2020). Kernel methods have fixed set of features, and networks in this kernel limit have fixed representations at each layer over training. By contrast, realistic neural networks adapt their features during the course of training (Fort al., 2020 Vyas et al. 2022: Lee et al., 2020). Feature learning is widely believed to be intimately connected to the excellent performance of NNs on practical tasks.
If one takes a centered fx;; 0) in NTK parameterization with inputs æ and parameters 0 and scales the output as fxx 0) of (a;; 0), one can control the degree to which the NN learns features. Taking a - oo yields the "lazy limit," where the network behaves as kernel (Chizat et al. 2019). By contrast, if one sets a (vvNN) for fixed and simultancously scales the learning rate as n.N one obtains network whose hidden representations evolve even at infinite width (Geiger et al. 2020
Mei et al. 2018 Rotskoff & Vanden- Eijndenn 2018 Yang & Hu 2021 Bordelon & Pehle-
"Joint first authorship, alphahetized
"In this paper, we assume all networks are centered, meaning that f(z; 0) = 0 at initialization. This can be achieved by subtracting non-trainable copy of the network imitialization.

---

Published as a conference paper at ICLR 2025
van 2022), scaling termed the "maximal update parameterization" or "uP" by Yang & Hu (222 Although at practical widths and dataset scales, NNs in standard and NTK parameterization deviate from their infinite-widre kernel limit (Lee et al.
2020), networks in μPP achieve their infinite width Noci et al. 2024). Networks parameterized in jP be utilized effectively for zero-shot hyperparam-
behavior at realistic scales (Vyas et al. 2023a 3a
generally achieve better performance, and μP eter transfer from small models to large (Yang ct al. [021]] Bordelon et al., (2023). The width does not affect the degree of feature learning μP, so identify y as the feature learning strength or richness parameter (Woodworth et al. 2020).
This story leaves Y as free hyperparameter to be tuned, but despite its importance, little is practi- cally known about the effects of varying In this paper, we perform a thorough scaling analysis in which we train a variety of deep networks with y ranging across many orders of magnitude, with an eye towards both optimal tuning of y and cataloging the diverse dynamical phenomena that occur in different hyperparameter regimes. addition to studying the lazy (7 < 1) and rich (? ~ 1) regimes, we pay particular attention to what we term the ultra-rich (7 > 1) regime.
In this work, we focus on the online setting, where a fresh batch of data is given at each step and we train for fixed time set by compute budget. This is realistic of modern large-scale language and vision models and is scientifically clarifying, removing the confounding effect of finite dataset size.
Concretely our contributions are follows. We systematically vary Y over suite of models span- ning MLPs, CNNs, ResNets, and Vision Transformers (ViTs) on a variety of datasets. In all settings:
Sweeping over both y and the learning rate 1). we observe a characteristic phase portrait of model performance in the 7-7) plane, which depends only on the choice of the loss function and not the model architecture (Figure IJ. We show that the optimal learning rate scales as 7/- 72 in the lazy (7 < 1) regime and 1). 22/L, where L is the model depth, in the ultra-rich (y > 1) regime. Our phase portrait reveals an upper limit on y for a deep network be optimizable for a given compute budget.
After adopting the correct learning rate scaling, we find that large networks usually ex- ceed or match the performance of naive y = 1 networks if trained for sufficiently long (Figure 2). This disagrees with findings in the offline training setting (Petrini et al. 2022) where strong feature learning was found to degrade performance.
At small y. we observe the scaling of 1) with 7 predicted by kernel theory. To our knowl- edge, we are the first to observe that transformer- based architectures on large datasets can empirically reach the lazy limit. At -lightly-too-large nj in the lazy (? < 1) regime, we observe the "catapult effect" of Lewkowycz et al. 2020) in which the loss quickly grows large before gradually converging. We see catapults for 72 z 7) S y for cross-entropy loss, but only in narrow band around 7² for MSE. Smaller y leads to larger catapults.
At large Y, we observe silent alignment as in Atanasov et al. (202)): at early times, the kernel aligns itself to the task before the loss drops. During this early period, NNs exhibit strikingly similar dynamics upon rescaling - = 1/t/7, with the loss often exhibiting step- wise drops as in (Saxe al.; 2014, Simon et al., 2023b) even in very realistic networks. We further observe that the loss drops at the end of silent alignment coincide with the Hessian growth predicted by the edge of stability (EOS; Cohen et al. (2021)).
At optimal learning rate, we see a surprising agreement in the loss trajectories of large- networks at late times in training. This suggests that large-ry networks are learning similar features and functions past certain threshold. We verify this similarity by comparing losses, accuracies, learned functions, and internal representations.
Finally, we reproduce and analytically derive our phase portrait and many of the above behaviors in the simple setting of a linear neural network model (Section 4). There, we show the modified 1) scaling in the large 7 regime is due to progressive sharpening effect.
We review the catapult, silent alignment, and progressive sharpening effects in Appendix C
1.1 RELATED WORKS
Chizat et al. (2019) introduced the laziness parameter a in their seminal work. Mei et al. (2018); Rotskoff & Handen-Fijndes (208)) introduced networks in mean-field parameterization. Geiger
2

---

Published as a conference paper at ICLR 2025
et al. 2220)) introduced the correct scaling of a with aas (cy/N)=' to achieve feature learning at infinite width. They also highlighted that ry can be identified as the parameter that controls feature learning. Yang & Hu (202I) expanded this idea to deep networks of arbitrary architecture and coined the term jdP. [Bodeloon & Pehlevan (2022) described this limit in terms of dynamical mean field theory. Extensions to infinite depth were performed in [Bordelon et al. (202)) and Yang & Littwin (2023) and to infinite attention in [Boddllon et al. (2024c). We give review of network parameterizations in Appendix B key motivation for this work is that y is the main free parameter
in the DMFT description of inhnite-width feature learning networks, warranting more thorough empirical investigation.
The question of optimal was explored in the paper of Petrini et al. 2022). There, for several models in the offline setting, they showed large Y networks performed worse. In a linear network toy model. Pessm et al. (2021) showed that SGD noise could be absorbed into a rescaling of a. Selocchi et al. (2023) performed a large set of thorough experiments to study the interplay between the SGD noise in different regimes of y and find that SGD noise can hurt or help in the offline setting. By contrast to these papers, we focus on the online setting. We comment on this further in Section 2.1] Under cross-entropy loss, [Agarwala et al. (020) performed extensive sweeps un-centered networks over 7, 77 identifying y as the soft-max temperature, similarly finding that slightly larger 7 achieves top performance. Our cross-entropy results overlap with this work and recover their findings, but we additionally identify a different way to scale 1) at large y, allowing us to probe the very rich limits without training instability.
Various analytical studies of the large-"y limit have been performed in the setting of linear networks (Atanasov et al. [022 Tu et al. 2024). with Jacot et al. (2021) terming this limit the "saddle-to- saddle" regime. The feature learning parameter has been shown to play a key role in the mechanism
for grokking (Kumar al. 2024). The paper of (Kalra et al. 2023) performed thorough analysis of the causes of progressive sharpening and catapult behavior in networks across μP and SP in the offline setting. the online setting we show that all effects are present in juP and are modulated by the y parameter. In the setting of linear network, (Marion & Chizat (024) also found nontrivial and depth-dependent sharpening effect. Earlier works studying the effect of varying the 7 parameter (Selocchi et al. [2023] Bordelon & Pehlevan) 2022 2023) focused on the offline setting of training with repeated data until some convergence criterion is satisfied.
2 SETUP AND NOTATION
We train neural networks f(x;;) parameterized using uP scaling, with the output further down- scaled by factor of 7 We denote the width of a given network by N and the depth of a feed- forward network by L We train online on batches B, = f Bu)ر= batch with size B for
t e {I...., } time steps. The loss of each batch is given by mean of per-example losses:
1 E E (H)( (g) = B
(1)
MEB,
We will consider both mean- rean-squared-eree and cross-entropy loss for f. In this work we focus on optimizing the parameters by vanilla SGD, as even in this simple setting, the effect of y is not well- understood. For small batch sizes, the the SGD effects are controlled by the ratio // B (Jastrzebski
et al. 2017 Smith et al. 2020 Sclocchi et al. 2023). We verify this this in Appendix K
We study networks trained on the datasets, MNIST, CIFAR and TinylmageNer. Our motivation is to study networks training in the online setting over several orders of magnitude in time. To this end, we adopt larger versions of these datasets: "MNIST- data augmentation to Iingemtgenet. We generate MNIST- 1M 1M using the denoising diffusion model (Ho et al. (220) in |Pearce| (2022). We use CIFAR-5M from Nakkirrn et al.) (2021). Earlier results in Refinetti et al. (2023) show that networks trained CIFAR-5M have very similar trajectories to those trained on CIFAR-10 without repetition.
and CIFAR-5M, and apply strong
2.1 WIYY TRAIN ONLINE?
In the modern era of large models, the full training set is seldom repeated (Kaplan et al., 2020: Muennighoff et al. 2023). Online, the effect of SGD noise was seen to have negligible effect on
We verify our μP implementation, confirming consistency of dynamics across width in Appendix Here, depth counts the number of weight matrices, e.g. f(x) W2a(Wxx) has depth I. 2.
3

---

Published as a conference paper at ICLR 2025
MSF loss
cross- loss
laay
rich
ulbrariia rich
lasy
rich
ultra-rich
7
7-1
7 2.
divergence
divergence
Fyrig 3/1
log !7
log
9:
catepult
Nonvergenee
catapult
convergunce
no training
no training
log Y (a)
log y (b)
N-64, L=3, B=128, 1=1000, loss oss
Accuracy for One Pass MNIIT-IM IM MLP
ConvNet N-64, Lw4, B- 128, 1w4600, loss-xant
Augmentation CIFAR-10
0.6
0.8
0.5
0-Y
0.
0.0
51
V
(c)
(d)
Figure 1: Phase portraits of the 7 1) plane for MSE and cross-entropy losses. a.b) Schematics of the regimes of network training for both losses. As the number of gradient steps increases, the lower boundaries of the convergent region descend, and the "no training" region shrinks. L here denotes the network depth. All scalings are obtained analytically in Section 4] c.d) Final accuracies of deep networks trained across grid of values of y, 17 for e) and MLP on MNIST-IM with MSE and d) CNN on CIFAR-5M with cross-entropy loss. As shown by dashed lines, these empirics agree with our analytical diagrams. These results are robust to the choice of model and task.
scaling both empirically (Vyas et al. 2023b) and theoretically (Paquette et al. 2024). The online setting thus allows to isolate properties of gradient flow that SGD approximates. This removes the complicated overfitting effects that have been theoretically shown to build up when data is repeated (Bordelon et al. (024a). as well the many regimes of SGD under an interpolation constraint (Sclocchi & Wyart] 2024).
optimization hyperparameter for the population
3 EMPIRICAL RESULTS
In this section we consider variety of architectures: MLPs, CNNs, ResNet- 18s, and Vision Trans- formers (ViTs), trained on variety of tasks. We first discuss the different regimes observed in 7/ン plane across architectures. We then discuss the properties of the Hessian, and identify spe- cific dynamical phenomena. Finally, focus the large limit and show that learned functions learned representations largely agree across large y networks.
3.1 PHASE PORTRAIT OF WITH ?
We begin by considering the 7-7 plane for variety of architectures and datasets. We sweep jointly over every pair of 77, y in log-spaced grid running from y = 10-5 to 10a. For each 7, we sweep from 7) 10 1) 10-12 downwardy in log-spaced grid until the first convergent 1) is reached.
4

---

Published as a conference paper at ICLR 2025
Online losa for CITRL-NN CMN
"ustt Acouracy, CIAARSH CNN
Be225,
B-236, 1025
1- 55-2 3e-3
t=0c+00
1=20+000
1- Se-04
t=le+01
3.20
y- 1e01..--3-00
1-5e+-0
y 1e400000 7e0000
1-101110g0- g-3em01
(=1o+23
y=1=+02.q
1-Se+21
y=1e105.0 g-2e03
1-2e+ 04
y- 101104.0-2erg
(=1o+05
y -=1+405.0p (33e003
02
1=5e+25
Lasy Scaling
2x30
Rich Scaling
12 181 28 14"
134 36-
30'
382 100
(a)
(b)
Figure 2: a) Online loss curves of CNN trained on CIFAR-5M across 7, at the second largest convergent y). 1) scales as 722 for 7 and 72/L for y > 1. We also observe that very large y have long period of flat loss before a drop. We overlay dashed lines to highlight the different power law scalings of loss with training time observed in the lazy and rich regime. b) Final test accuracy a function of y. We that, long as one trains long enough, larger y yiclds equal or better generalization to 1, but the returns are marginal past some point. We verify this for other networks in Appendix Error bars show the small variance over three initializations.
After finding the first convergent 7/, we sweep over a range of ry below it. We consider both MSE and cross- cross-entropy losses. Further details are given in Appendix A
At small 7, 1) x y for lazy networks We consider the limit of small 7, 7 - 0. This recovers the lazy limit of Chizat et al. (019). The (empirical) neural tangent kernel of the network is K KNTK = Vf f O ") In kernel regression under MSE loss, the maximum achievable learning rate proportional to the top eigenvalue of the Hessian. This sets the upper bound to be 7) y². More generally at small y the Hessian is dominated by the Gauss-Newtor term, as we will explicitly write in Section 3.2, equation 3 This term scales as (()-?). Moreover, in order to achieve trainable network as 7 + 0, one needs the loss to have a nonzero change in L. in that limit. We have:
L441 Lu - 17 B Σ L e (Erlemn
())
By centering, near initialization e' is (1) The left hand side is thus O(r-2), so the minimum learning rate to see the loss drop in T steps scales as 1j ~ >²//T. We see this on the left region the plots in Figure I(c) and I(d) We verify that networks of different but sufficiently small 7 have the same loss dynamics in Appendix
indicating that the lazy limit is reached by realistic networks.
At small y and larger 7,, catapults occur For 1) larger than the sharpness allowed by convex optimization, the loss explodes early on in time. For a range of 7) above this, the loss eventually drops again. This is the catapult effect of Lewkowyez al. (2020). The maximal learning rate for these scales as n2 under MSE loss and as under cross-entropy loss. We revisit this in Section 3.3 and Figure 4(a)
At large % there are two learning rate scalings Empirically, in Figures I(c) and I(d) we ob- serve a "triangle of optimizability" in the large 7 limit. To our knowledge prior works have not highlighted this fact. The upper boundary of the maximal learning rate before the loss explodes is found empirically to scale as 17 ~ 2/1 On the other hand, at large y, activation movement precedes function movement, by contrast the lazy limit. The minimal learning rate scaling in order for the activations to move by (1) 1) is 1)
as in (Bordelon & Pehlevan 2022). 1) 7/T under a budget of T' training steps. Taking C' x TB to be the compute budget, we see that generally increasing C' increases the optimizable region for large 7,
Sufficiently large y NNs see improved scaling laws We plot the loss in time across y in Fig ure 2(a) We overlay different scaling curves dashed lines. We the slopes between the
5

---

Published as a conference paper at ICLR 2025
Final Hossiar Eigervalues, C CNN N-96, Lmd, Bm256, lossemse
Sharpn oss of Less Landscape. CNN, CIFAR-SR SML.
N-96. B-256.
37
16-44
1e
Qharppan
y
12-4
y- 1c400
1=401
28-
1e404
224
1e+05
(a)
(b)
Hessian Spectrum, y- CNN, CIFAR-SM, B=256. loss-mse
Hessian Spectrum, V- 1e44 04 CNN, CIFAR-SM, N=96. L=4, B-256, loss-mse
6.084
0.055
6.081
8000
0.002
60000
6.011
165
(c)
(d)
Figure 3: a) The top eigenvalues of the Hessian at the end of training vs y for an NN trained with MSE. We clearly see two regimes. For small 7 we see that Amax b) The eigenvalue vs time across several value of 7. c) The top eigenvalues across time for a lazy network. see ten outliers, equal the number of classes. The Hessian otherwise does not change in MSE. d) The same for rich network. At late time, rather than seeing a small set of outliers, we see many eigenvalues grow to sizeable range. Further plots given Appendix M
the lazy setting, we
small lazy networks and the large rich networks appear substantively different. The former slope can be predicted from kernel theory, as in Pillaud- Vivien et al. (2018) and Bordelon & Pehle- van (202T). Increasing y beyond some threshold does not appear to improve the scaling law. This echoes similar findings found by the recent model of Bordelon al. (2024b).
Large-y NNs achieve good generalization, given sufficient training time We plot the accuracy of NNs across 7 in Figure 2(b) We see that generalization improves with larger 7, eventually reaching roughly constant value. This contrasts with studies in the offline setting (Geiger et al.) 2020 Sclocchi et al. (2023) where performance decreases at large 7, We highlight that if networks do not train long enough, non-monotonicity is indeed present, but this goes away with longer online training.
3.2 HESSIAN SCALING
We next consider the scaling of the Hessian of the loss as function of y. On a batch B of B test points, this is given by Hg = V2LB. or:
1
Hs
(a (f(xx) (f(x)))))
2 f(xp)) ((((()))))))
(3)
B
B
Here, we have split the Hessian into a Gauss-Newton term & and loss gradient term R. The Hessian contains important properties about the optimization landscape. One particularly important

---

Published as a conference paper at ICLR 2025
Aligrment for VENIST-IM MIPP
Online loss LEAILSN CMN
N=56, L=d, 2-6 loss-it5
0.4
30.05 0- 18
Los
Melli
16:04 1- 1:
Algrment
36-04, 16
Sharpness
1e-03 16-05
y=1+401
1.0
36.033 12:04
A.M
1e403
1e104
14
y=10105
y- 9-12000
301
y- -11100.0-1ee000
17 141
13"1
(a)
(b)
Figure 4: a) At sufficiently small y, we observe that the loss catapults. Small y incur a larger catapult. Further catapult plots are in Appendix[N] b) At large y, we do not observe catapults. At the optimal learning rate, the loss stays near a saddle for an extended period of time, before suddenly dropping. We show that during this period, the alignment between the final-layer kernel and the task grows before the loss drops. Kernel-targe -target alignment is defined in Appendix [P] Further silent alignment plots are in Appendix Q
We also see characteristic sep-wise loss drops this regime.
quantity is top eigenvalue Amaz, also known as the sharpness. The sharpness governs the maximal allowable learning rate in a convex optimization problem.
Across all tasks we find remarkably consistent behavior of the Hessian spectrum, see Figure 3(a) For < 1, the network stays lazy and the Hessian does not evolve. There, the Hessian is dominated by the G term, which scales as 7 Once 7 exceeds a threshold however, we observe a distinct change in scaling the end of training. For feed-forward MLPs and CNNs, we verify scaling
going as y -2/L, In section we derive this scaling in linear network model.
We also track how the top eigenvalues of the evolve in time in Figures 3(b) to 3(d) For lazy networks predictably see no evolution. In the next section show that small y but larger 1) networks can catapult and lower their sharpness. For rich networks we highlight that a very extensive number of eigenvalues grow to large scale. This is in contrast with the commonly observed phenomenon that for a C'-class classification task there should be C' outliers. Here, in the lazy limit there are indeed 10 distinct outliers, but in the rich regime there are no distinct outliers at late times.
3.3 DYNAMICAL PHENOMENA
We here describe a catalog of different dynamical phenomena that we find in networks trained in the small- and large-y regimes. We give background on these effects in Appendix C
The catapult effect at small We see in Figure 4(a) that small networks exhibit large increases followed by decreases in the loss, with remarkably regular behavior, decreasing scale as y grows. In Appendix DD we leverage similar linear network model to arrive at the desired catapult scalings with y under both MSE and cross-entropy loss.
Silent alignment at large 7 We now consider the dynamics at large y. Empirically, we find that our sweeps do not detect any catapult effects once y is large enough. Taking the optimal learning rate for SGD at large 7 still yields a long loss plateau. During this plateau, the NN can be shown to be adapting its features, and appears to align its hidden-layer representations with the task. We illustrate this in Figure 4(b) This effect was identified in linear networks at small initialization in Atanasov al. (2022). and shown to empirically hold in restricted class of NNs on whitened data. Here, we show that it arises in a much broader class of realistic settings.
Stepwise loss drops at large At very large 7, we often observe loss trajectories following a characteristic staircase pattern, with one more plateaus punctuated by sudden drops. This is prediction of linear network theory (Saxe et al.]2011 2014 Simon et al.) 2023d that seems to hold even for very realistic models. We see these step-wise trajectories for ML.P in Figures 4(b) 3 and 41
7

---

Published as a conference paper at ICLR 2025
Conelations of Lugits, OUN, CHAAI-5M N-66, L-4, B-226, 105s- 2=20104
Alignment MNNSSMM MIP
L-4,
11-05
3e-02 16+02.
Scomdd
4.47
-1-01
Vo 10w53, 1o-32
-1c+00
=1e+01
=lc
0.
=3c+33
0.5
=2
-10+05
2.54
0.5
13"
Find V
60V
(a)
(b)
Figure 5: a) Function outputs between pairs of networks, as descrihed in the main text. We see that lazy networks have nearly identical function outputs, and that rich networks agree in their function outputs at the end of training as well. See Appendix[] [[ further plots. b) We study the kernel-target alignment of final-layer representations across y the end of training.
CNNs in Figures 31 and 40 and ViTs in Figure 16 We show that, upon rescaling time as T = nt//,, the early time dynamics coincide. The first loss drop goes as 7. We explain the origin this timescale in Section
Progressive sharpening at the end of silent alignment Across all tasks and models, we find that the loss drop that ends silent alignment always accompanied by a growth the sharpness to the final value of approximately 2/n. We show this in Figure 4(b) This is the large analogue of the
progressive sharpening effect observed in Cohen et al. (201)). This effect is not present at small y.
3.4 COMPARING LEARNED FUNCTIONS AND REPRESENTATIONS
Finally, we consider the question of to what extend different networks are learning the "same func- tion" and the same representations. To this end, consider two distinct networks A, and B well as a held out set of test points. For each point Tye. there is correct class c, so that "ylee = d as one-hot vector. We consider the value that a given network places on [yn]e- On one axis, we plot network and on the other we plot network B. Networks are learning the same function if and only this plot yields an exact diagonal line. In the lazy limit, we see in FFiuure 5(a) that indeed the two networks are learning nearly identical functions. We also see strong agreement between different network in the rich limit at the end of training. By contrast, rich and lazy networks not have substantial function agreement.
We then consider the learned representations in Figure 5(b) Here, we plot a version of centered kernel alignment (CKA) between the final layer network activation kernel K and the "task kerne"" yy for y a x C class label matrix. We see that, after adopting a suitable time rescaling, the alignment scores for variety of networks across y agree. Larger y networks retain slightly higher alignment scores.
4 A SIMPLE MODEL EXPLAINING OBSERVED SCALINGS
In this section, provide theoretical model that explains all observed scaling relationships em- pirically observed in Figure and analytically reproduces our phase portraits.
The following simple model will prove sufficient: deep linear network of width one, trained on a single example. The network function this case is f(x) = Wt. w We will initialize all weights to 1 at initialization, and due to the commutativity of multiplication they will receive the same gradients and remain equal throughout training. We are thus justified in replacing all weights with a single weight w and, letting æ we need only keep track of the function value = w"
A technicality: because the variable w is repeated L times, it will receive gradient T. times larger than of the original variables we. To recover the original training dynamics, we would thus need to downscale n
8

---

Published as a conference paper at ICLR 2025
MSE Loss
Cross-Entropy Loss
Lazy (? < 1)
Ultra-rich (Y > I)
Lazy (? < 1)
Ultra-rich (? >>>>>>>>>>> 1)
Mmin
Meri
Mmax
Mnin
Mmax
Mmin
Merit
1huax
Mmia
Mmax
7/2
72
22/L
72
7
7
72/L
Table 1: SGD learning rate scalings for our toy models for loss functions in the lazy and ultra-rich regimes. All values come out of our one-parameter model except Temi for MSE loss, which comes from the two-parameter model given Appendix D experiments.
Predicted scalings match NN
We divide the function value by y and center it by subtracting off its value at initialization before passing it to the loss function.
4.1 SCALINGS FOR MSE LOSS
We first consider MSE loss with target value y 1. All together. this yields loss equal to
CMSE (w) = 11)
(4)
with w at initialization. The loss is minimized by
1+ 1 for < 1, for > 1,
w. = (r (
(5)
7
where we use "z"" to indicate that we are neglecting higher-order terms in as appropriate to the regime. In the lazy 7 < regime, w, close the initial w, whereas the ultra-rich 7 > regime, the parameter must grow substantially. The Hessian of the loss at convergence is then:
w. 2L-2 722
La no
32
12
for < 1,
(6)
12--2111
2/L
for ?
Maximal convergent learning rates: In order to converge, the learning rate must satisfy 1) < 3/""(w") near convergence. Thus, from rate 7max will scale as
Equation (6) we see that the maximal convergent learning
Tmmax {
x² for y < 1, 2//2 for y > 1.
Minimal nontrivial learning rates: In investigating the minimal learning rates, we can safely take the gradient flow approximation. We define the eиrоr -L'(ww
3/2 (w -1)
When y < 1, we may linearize the network output with respect to w. and doing so we find that W ~ 1+ + (1 exp(- - -xp(--Tf))
)). The error is then given by A ~ exp(- -_-2m2) Therefore, in order for the loss to move nontrivial amount, require the scaling Tmmin ~ -//T.
When 7 >> 1, in order for the network to learn features, we must have that the hidden layer acti- vations move by (1). Because in this model these are given by w'x for each layer l, we see we require (1) weight movement. The weight updates are given by ie = -n][(w) = n - So in order for the network to learn features in T steps of training, we require > 1FL n/T.
Beyond constraining that the activations move, we also consider how long takes escape the initial saddle and lower the loss. Using that w'( Ry nL w6--1 during the initial period, find that
7
LaT Ee7
if L = 2,
wT( R
(8)
(1
Lait
if L 2 3.
by factor of T.. However, because assume I. ((1), neglect this for simplicity and our y, T scalings will be unaffected.

---

Published as a conference paper at ICLR 2025
Final MSE 1-puram toy model
w
Final xeni loss of param model f w
1.0
1.0
Tew 11
0.8
101
10'-
0.9
Learning rate
divergence
0.6
Leaming rate
divergence
10-2.
Conver convergence
10-2.
convergence
ploade chtipll
0.4
10-5
10-5
0.7
10-4
DO Uainng
0.2
10-
no training
0.6
10
10"
10-3
10a
Output multiplier
Output multiplier
Figure 6: Simulations of our one-oonee-eee aae-pareetter model with MSE loss (left; see seeatiuni (4) and cross- entropy loss (right; see EEquaion (9) with varying 7, recover our principal phase portrait. Simula- tions run with For simulations of two-parameter model which exhibits true catapults, see Figures 8 and
This phase ends when the function value whth - reaches order unity. In order for this to occur, we need " T log y if L 2 and 17 ~ 7/T if I. 2 3. Neglecting the logarithmic factor, find lower bound 7m ~ H. Note as consequence that the early-time dynamics are approximately 7-invariant after the rescaling 7 1///7. We verify this in Figure 4(b) 5(b) and Appendix L.2
Critical learning rate for the catapult effect: The above one-parameter model is sufficient to obtain all the coarse dynamical phenomena we see except for so-called loss "catapults." As described by (Lewkowyzz et al. 2020), network trained in the lazy regime with just-too-large learning rate will grow quickly loss, but then move to a flatter region of parameter space and converge.
The catapult effect generally requires two parameter dimensions to describe (Damian et al., 2023), so our one-parameter model cannot capture In Appendix [N] we describe and simulate a minimal two-parameter model that exhibits catapults. We find that for MSE loss, catapults occur in narrow band Meiit < 1) < T/max with Yerit ~ 72 just like 7mmax-
4.2 SCALINGS FOR CROSS-ENTRONE ENTROPY Loss
We now consider the toy cross-entropy loss
log(e((w) +
log(e -(w) +1)
Lxem( (w:) = 1 1
(9)
1
where we use the shorthand f(w) - - -w) for the centered and rescaled function output. This a binary cross-entropy loss with ground -ruth class probabilities = + -)- and p1 = 1 Po chosen so that the loss minimum occurs at ((w) = 1. The resulting loss landscape is, scaling-wise, the same as that for MSE, with one salient difference: when f(w) > 1, we have Cxem( (w) ~ [w(w] as opposed to LASE jP(w).
The main effect of this difference is seen the lazy (7 < 1) regime when rj too large for stability at the minimizer. Once w 1 has grown to order just larger than y (and we are still seeing linearized dynamics), it begins to experience constant-sizel restoring steps dw
nL sign(w 1) It is only 17 Y or larger that this step large enough to lead to a O, (1) change w, upon which we escape our initial linear region and can diverge. The result is that the maximum stable learning rate for our toy model in the lazy regime is 7eii 7² but the maximum nondivergent learning rate is 7max In our minimal two- parameter model described in Appendixx D we find that the loss does eventually settle back down and converge for Jen 1j < Muaax and that this triangle of the phase diagram is there true catapult regime. All the above scalings are summarized in Table
Contrrast this with the linear-sized restoring steps Ew
52 5m which with MSE.
10

---

Published as a conference paper at ICLR 2025
We perform simulations of our one- parameter model in a fashion analogous to our parameter sweeps for deep networks for T = 10 steps of gradient descent. The results, plotted in Figure 6 confirm analytical scalings and match our phase portraits for deep networks (e.g. Figure 1j. Finally, also provide a similar analysis in the case of the Adam optimizer in Appendix Appendix F
5 CONCLUSION
We have presented set of large-scale empirical sweeps across for deep networks trained on realistic data. We have empirically identified how the upper bound for optimizable 7j is determined across a host of different 7 in terms of the training time and the depth of the network. We see that, conditional on an adjusted learning rate scaling with y, larger NNs are competitive with or better than NNs after training long enough. We have seen very predictable scaling of hessian spectra with y and identified various dynamical phenomena whose appearance is governed by the size of y. By appeal to linear NN, have been able to explain many of observed scalings. In future steps, we hope to to extend our analysis to optimizers beyond SGD.
AUTHOR CONTRIBUTIONS
AA ran early empiries that revealed the phase plot structure, initiating the project, ran all MNIST-IM and most CIFAR-5M experiments, wrote the initial draft, offered feedback on the solvable model, and led the team logistically. AM ran all VIT experiments and experiments TinyImageNet, helped with CIFAR-5M experiments, observed the role of attention temperature in the final phase diagram, and offered technical feedback both experimental and theoretical fronts. JS proposed and de- veloped the solvable model exhibiting all phases, ran empirics on that model, and offered technical feedback all experiments. CP supervised the project and gave feedback and direction. All authors contributed substantially to the final draft.
CODE AVAILABILITY
The code to reproduce all figures in this paper can be accessed at:
BERNOWLEDGEMENT
The authors are thankful to Blake Bordelon, Daniel Kunin, Katherine Lee, Allan Raventos, Jascha Sohl-Dickstein, Jacob Zavatone-Vedts and Brian Zhang for helpful discussions. AA thanks Blake Bordelon and Jeremy Cohen for in-depth discussions that motivated many of the questions investi- gated in this paper. We thank Blake Bordelon and Jacob Zavatone- Veth for comments on a draft. AA was supported by Yaser Abu-Mostafa Fellowship from the Fannie and John Hertz Foundation. gratefully acknowledges support from the National Science Foundation Graduate Fellow Research Program (NSF-GRFP) under grant DGE 1752814. C.P. is supported by NSF grant DMS- DM3-2134157, NSF CA- REER Award II3-223778, DARPA grant DIAL-FP-038, Sloan Research Fellowship, and The William F. Milton Fund from Harvard University. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. Discussions contributing to this work took place at the Kavli Institute for Theoretical Physics' 2023 program on "Deep Learning from the Perspective of Physics and Neuroscience." AA and JS thank the builders of particular house in Reston, VA which, to their surprise, they discovered they had both lived in at different times.
AM acknowledges the support of a Kempner Institute Graduate Research Fellowship.
REFERENCES
Atish Agarwala, Jeffrey Pennington, Yann Dauphin, and Sam Schoenholz. Temperature check: theory and practice for training models with arXiv:2010.07344, 2020.
losses. arXiv preprint
Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-orde regression models exhibit progressive sharpening the edge of stability. arXiv preprint arXiv:22/0.0448 2022.
11

---

Published as a conference paper at ICLR 2025
Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks kernel learners: The silent aligmment effect. In International Conference on Learning Representations, 2022. URL htt.ttpss//ppeneeeeee.eeee htt.:e/e/ooommm11oollllooo....o.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Anima Anandkumar. signSGD: Compressed Optimisation for Non-Convex Problems, February 2018.
Blake Bordelon and Cengiz Pehlevan. Learning curves for sgd on structured features. arXiv preprint arXiv:2106.027133 2021.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. arXiv preprint arXiv:2255006633 2022.
Blake Bordelon and Cengiz Pehlevan. Dynamies of finite width kernel and prediction fluctuations in mean field neural networks. arXiv preprint arXiv:2304.03400 2023.
Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit, 2023.
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. A dynamical model of neural scaling laws. arXiv preprint arXiv:2402.01092 2024a.
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. How feature learning can improve neural scaling laws, 2024b.
org/ abs /tt.//449...188
Blake Bordelon, Hamza Tahir Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head trans- former dynamics. arXiv preprint wrXv:2240.115712 2024c.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing sys- tems, 31, 2018.
Lénaïcc Chizat and Praneeth Netrapalli. Steering deep feature learning with backward aligned feature updates. arXiv preprint arXiv:2311. 2023.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065, 2021.
Corinna Cortes, Mchryar Mohri, and Afshin Rostamizadeh. Algorithms for learning kernels based on centered alignment. The Journal of Machine Learning Research, 13(1):995-828, 2012.
Alex Damian, Eshaan Nichani, and Jason D Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. ICLR 2023, 2023.
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study loss landscape geometry and the time evolution of the neural tangent kernel. In Hugo Larochelle, Mare Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tices Lin (eds.), Advances in Neural In- formation Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NewrIPS 2020, December 6-12, 2020, virtual, 2020.
Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthicu Wyart. Disentangling feature and lazy training deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020 (11):113301, 2020.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemnl:: Convergence and gen- eralization in neural networks, Advances in neural information processing systems, 31, 2018.
12

---

Published as a conference paper at ICLR 2025
Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity, arXiv preprint arXiv:2 155.33,,,. 2021.
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influeneing minima in sgd. arXiv preprint arXiv:/711.06623 2017.
Dayal Singh Kalra and Maissam Barkeshli. Phase diagram of early training dynamics in deep networks: effect of the learning rate, depth, and width. In Proceedings of the 37th International Conference on Neural Information Processing Systems, pp- 51621-51662, 2023.
Dayal Singh Kalra, Tianyu He, and Maissam Barkeshli. Universal sharpness dynamics in neural network training: Fixed point analysis, edge of stability, and route to chaos. arXiv preprint arXiv:2311.02070 2023.
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. arXiv preprint arXiv:2001.08361 2020.
Diederik Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6909 2014.
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International Conference on Machine Learning, pp. 3519- 3529. PMLR, 2019.
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the tran- sition from lazy rich training dynamies. In The Twelfth International Conference on Learning Representations, 2024. URL https::/openrevieee net/ forum? ?id SmnLVIVo
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
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, 33:15156-15172, 2020.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.022/1 2020.
Pierre Marion and Lénaic Chizat. Deep linear networks for regression are implicitly regularized towards flat minima. arXiv preprint arXiv:2405. 13456, 2024.
Song Mei, Andrea Montanari, and Phan-Ming Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 11566)))77665--7777,, 2018.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, pp. 2388- 2464. PMLR, 2019.
Si Yi Meng, Antonio Orvieto, Daniel Yiming Cao, and Christopher De Sa. Gradient descent on logistic regression with non-separable data and large step sizes. arXiv preprint arXiv:2406.0500 05033,, 2024.
Niklas Muennighoff, Alexander M Rush, Boaz Barak, Teven Le Scao, Aleksandra Piktus, Noua- mane Tazi, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling data-constrained language models. arXiv preprint 2305 2023.
13

---

Published as a conference paper at ICLR 2025
Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi. The deep bootstrap framework: Good online learners are good offline generalizers. In International Conference on Learning Represen- tations, 2021.
Lorenzo Noci, Alexandru Meterez, Thomas Hofmaann and Antonio Orvieto. Why do learning rates transfer? reconciling optimization and scaling limits for deep learning. arXiv preprini arXiv:2402. 17457, 2024.
Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington. 4+ phases of compute- optimal neural scaling laws. arXiv preprini arXiv:2405. 2024.
Tim Pearce. Conditional diffusion mnist.
https://githai com/TeaPearce/
Conditional pineiitiioal_iiffssioooo..o. wwwfffiifoooooooe MNIST 2022. Accessed: 2024-05-14.
Scott Pesme, Loucas Pillaud- Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34:29218 0-22218--9230,
29230, 2021.
Leonardo Petrini, Francesco Cagnetta, Eric Vanden-Eijnden, and Matthieu Wyart. Learning sparse features can lead to overfitting in neural networks. Advances in Neural Information Processing Systems, 35:9403-9444 9416, 2022.
Loucas Pillaud-Vivien, Alessandro Rudi, and Francis Bach. Statistical optimality of stochastic gra- dient descent on hard learning problems through multiple passes. Advances Neural Information Processing Systems, 31, 2018.
Maria Refinetti, Alessandro Ingrosso, and Sebastian Goldt. Neural networks trained with sgd learn distributions of increasing complexity. In International Conference on Machine Learning, pp. 28843 -2866. PMLR, 2023.
Grant Rotskoff and Eric Vanden-Eijnde Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks. Advances in neural information processing systems, 31, 2018.
A Saxe, J McClelland, and S Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In Proceedings of the International Conference on Learning Represen- atations 2014. International Conference Learning Represenatations 2014, 2014.
Antonio Sclocchi and Matthieu Wyart. On the different regimes of stochastic gradient descent. Proceedings of the National Academy of Sciences, (119))2311331122 2024.
Antonio Sclocchi, Mario Geiger, and Matthieu Wyart. Dissecting the effects of sgd noise in distinct regimes of deep learning. International Conference on Machine Learning, pp. 30381-30405. PMLR, 2023.
James Simon, Maksis Knutins, Liu Ziyin, Daniel Geisz, Abraham J Fetterman, and Joshua Al- brecht. On the stepwise nature of self-supervised learning. In International Conference Ma- chine Learning, pp. 31852-31876. PMLR, 2023a.
James B Simon, Maksis Knutins, Liu Ziyin, Daniel Geisz, Abraham J Fetterman, and Joshua Al- brecht. On the stepwise nature of self-supervised learning. In International Conference Ma- chine Learning, pp. 31852-31876. PMLR, 2023b.
Samuel Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. In International Conference on Machine Learning, pp. 9058-9067. PMLR, 2020.
Jascha Sohl-Dicksteen, Roman Novak, Samuel S Schoenholz, and Jachoon Lee. On the infinite width limit of neural networks with a standard parameterization. arXiv preprint arXiv:2001.077 07301, 2020.
Zhenfeng Tu, Santiago Aranguri, and Arthur Jacot. Mixed dynamics in linear networks: Unifying the lazy and active regimes. arXiv preprint arXiv:2405. 100005..7788) 17580, 2024.
14

---

Published as a conference paper at ICLR 2025
Nikhil Vyas, Yamini Bansal, and Preetum Nakkiran. Limitations of the ntk for understanding gen- eralization in deep learning, arXiv preprint urXiv:2206.10012, 2022.
Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cen- giz Pehlevan. Feature-learning networks consistent across widths realistic scales. arXiv preprint arXiv:2305.184 2023a.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Gal Kaplun, Sham Kakade, and Boaz Barak. Beyond implicit bias: The insignificance of sgd noise online learning. arXiv preprint 20XX::23066085000 2023b.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference Learning Theory, pp. 1665-3673, PMLR, 2020.
Greg Yang and Edward Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pp. 11727-11737, PMLR, 2021.
Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the infinite-width limit. arXiv preprint arXiv:2308.018/ 2023.
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: //openreview www.fooummiideeeeeeeeee.mmmm Bex.rmmm22.-XBKK2X22222. 6UKDBMAAD
Greg Yang. Edward J Hu. Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466, 2022.
Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310. 2023.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney, Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pp. 581-590, IEEE, 2020.
15

---

Published as a conference paper at ICLR 2025
A EXPERIMENTAL DETAILS
We evaluate all our experiments on A100 and H100 GPUs, using PyTorch. In order to compute the Hessian eigenvalues we use PyHessian Yao al. 2020) evaluated on fixed batch of size 512 at every point in time. In the following subsections we provide individual experimental details for every model trained.
Across all models, we swept over y from 10-5 to 10 in log-spacing with two samples per decade. We swept over 1) from 10 down to lower limit that varied by experiment, as reported on the plots. For each run that caused a gradient explosion, the model was not saved. We only saved models with 7 less than or equal to the first convergent learning rate for each value of y. For the MNIST- IM plots and CIFAR-5M plots, at each value of we only kept the seven values of that are within factor of 10s from the top stable to make the plots as in e.g. Figure
In all cases, we created a CenteredModel class, ensuring that the output of the network is zero at initialization by making the following definition for the trained function f:
1 [f 0) f(x 0o)] f(x, 0) 7
(10)
Here, On is the initial setting of the parameters, and is not trainable.
In the of ViTs compare training with QK-LayerNorm and QK-Norm in order to account for the attention logit temperature.
A.1 DETAILS ON THE PARAMETERIZATION
On top of the centering operation explained above, ee parameterize all our networks under pP (and Depth- paPp respectively) parameterization. More concretely, we define our L layer neural network of width N f(x) with nonlinearity ((x) as:
1
h (')x) = VD
(11)
1º (x)(()) = Th((x)) + VNLa
W (')(h(())
(12)
1 w (xh((z))
f(x) 7N
(13)
where D is the input dimension, and for 0 and 0 recover sP, and for 1 and
we recover Depth-jP (Bordclon et al.) [023] Yang & Littwin 2023). When running SGD, we update the weights with global learning rate going as N... That is:
Wh W Vo
(14)
Similar modifications exist for Adam, see Yang & Littwin 2023).
In the case of Transformers, we additionally change the attention temperature d, following (Yang et al. 2022). We further detail and contrast this parameterization with the NTK parameterization in
Appendix B
A.2 COMMENTS ON NUMERICAL STABILITY IN THE LAZY LIMIT
Here we detail ubiquitous problem that arose in the training of lazy neural networks at small y. To our knowledge, this problem has not been highlighted in the literature. At very small values of y, approximately 10-3, the appropriate learning rate in the lazy limit scales 17 ~ 222 This leads to weight updates on the order
6 or less. Machine epsilon for float32 is near 10-7 Because the final output of the function involves the subtraction in equation 10 the weight updates matter in both relative and absolute magnitude. For this reason, we find that we needed to go to double precision to reliably obtain curves below 7 10 33 In practice, we find that value of 10 or even 10 -2 essentially enough to reach the lazy limit in the online setting for essentially all tasks listed. See A Appendix [.. for check of this.
16

---

Published as a conference paper at ICLR 2025
We also highlight, that is insufficient to simply zero out the last layer weights initialization way to center the network. The lazy limit a network with zero- initiiaiii last layer weights yields the final-layer conjugate kernel or NNGP kernel. This kemnel has very different properties from the NTK, especially in that its dimensionality scales as the width rather than the number of parameter of the network.
A.3 MNIST-IM
MNIST-1M example digits
2
3
4
6
Figure 7: Examples of digits from the MNIST-IM dataset, with one digit from each class generated.
A diffusion model generated version of MNIST with million images, which we will call MNIST- 1M, was used train these networks. We swept over a variety of widths, depths. batch sizes, and architectures. Each individual run took approximately 15 seconds on A100 GPU. For each setting of width, depth, batch size, and architecture, we swept over the fine grid of 7) stated above. This led to 7 runs per value of y or approximately 150 runs total for each phase plot. Most runs were repeated over three initialization seeds, and key quantities, such as loss and accuracy, were seen to minimally vary over the initialization. This is expected in μP, by arguments of approximate initialization independence at sufficiently large widths as in (Bordelon & Pehlevan 2023 Mei et al., 2019).
The architectures we used were MLPs of depths 2, 3, and a feed-foowarr CNN of depth 4 consisting of two convolutional and two feed-forward layers.
The corresponding test set for this dataset is the original MNIST test set.
Although it is quite easy for most networks get near 100 percent accuracy in one pass through this dataset, we believe that it fills an important gap in the space of datasets. It provides a task that is fast enough learn, allowing for quick iteration, yet avoiding repetitions of data as remain online. In the era of online training, is useful to have simple but large-scale datasets in order to test theory against. We do not claim that this dataset representative of realistic imagenet scale natural data, but it is encouraging that virtually all effects observed in the CIFAR-5M and tinyimagenet datasets also present this smaller setting.
We observe that training accuracy on MNIST- M meaningfully transfers to the original MNIST, and there is no observed overfitting effect the networks that trained.
A.4 CIFAR- 10 AND CIFAR-5M
We train on the full dataset of CIFAR-5M introduced in (Nakkiran et al. 2021). This includes both the train set of 5 million images as well the I million image test set, making for total of 6 million images. We stay in the one pass setting. We also trained on CIFAR-10 with heavy data augmentation. We swept over width, depth, batch size, and architecture above. In essentially all plots, we followed the same recipe for sweeping over y and n. For some of the ResNets and Vision transformers, larger sweeps were performed.
17

---

Published as a conference paper at ICLR 2025
The architectures used were feed-forwart CNNs of depth 4 as above, ResNets of varying depths as well as ResNet- 18s, and Vision Transformers.
The corresponding test set for this dataset is the original CIFAR-10 test set.
A.5 TINYIMAGENET
In order to train on Tinylmagenet in an online setting, we use random data augmentations: rotations and crops. All the experiments on this dataset have been run using batch size 128. Individual architectural details are offered for every plot. We train all the models under μP (or Depth-uP where specified), using stochastic gradient descent (SGD).
B REVIEW OF NTK PARAMETERIZATION AND J.PARAMETERRAATION
For detailed discussion of μP more generally, see (Yag & Hu 2021 Yang et al. 202 Geiger et al, 2020 Bordelon & Pehlevan 2022). The aim of this section conceptual clarity the difference between NTK parameterization and --paramtterization. as well as to give some degree of motivation of the latter. See also the recent accessible works of (Yang al., 2023) and (Chizat Netrapalli, 2023
There are several equivalent ways of parameterizing NNs that will yield the same dynamics as pPP In what follows, we focus on a scalar-output feed-forward network for simplicity. The key distinction between NTK and f parameterizations it that the hidden layer activations move n (1/1 (1/vN) in the former and O (1). One can straightforwardly extend the argument to more general architectures.
B. NTK
For simplicity, we take all layers have hidden width N and the input space have dimension D. We Given an example F. the pre-activation h E+1 in layer &+ is given by in terms of the activation oe in layer las
c+1 h
W N
(15)
where Φf is given by an element-wise nonlinearity an wiement-wise non-linearity acting on h The 1/2 factor enforces that f to be e(1) in root-mean root-mean-square norm at initialization as N 0o. The output of the network fn is given by
a VN
fp
(16)
In NTK parameterization a is taken to be order 1. It is the laziness parameter identified in (Chizat et al. 2019). The change the function is given by
dfp dt
7SP Kuvll (fo.
(17)
Here Ky Vofu Vofv is the NTK between data points μ and v. The NTK is easily seen to be order a² and l' is order small a. In order to have the change in the function be O(1) we MSP = =
By applying the chain rule, obtains that the preactivations move (Geiger al. 2020
Jacot
et al. 2018 Chizat & Bach, 2018)
dh 7SP
no0
(18)
dl
N
Thus, at large N and a = 1 the areaactiationss of this network evolve as ( As N + 00 this becomes kerne limit, given by the wiith NTK that does not learn features.
18

---

Published as a conference paper at ICLR 2025
B.2 μP
In order get (1) preactivation and activation movement, that we must take 1/TvN for a constant 7, fixing no to be -N-independent. This implies that we simply replace the final layer of the network by:
1 w o P.N
MSP Nn..
(19)
As the prior analysis shows, this still yields that df ~ Thetan (1) at initialization. Note that we have not considered how to scale 7y with y. The scaling is the subject of the empirical studies of this paper.
C REVIEW OF DYNAMICAL PHENOMENA
In this section review the essential dynamical phenomena studied in prior literature that arise in various regions of the r7-7 plane studied
C.I CATAPULT EFFECT
The original catapult effeett was studied in (Lekkwwyzz et al. 2020). There, in the NTK param- eterization, it was observed that at large but finite width, taking the learning rate 1 slightly above the bound given by convex optimization theory 2/Amax led to regime where the sharpness of the loss monotonically decreased. Follow up work has observed that the sharpness itself can undergo catapult behavior (Kalra & Barkeshli 2023 Kalra al. 2023). In (Bordelon & Pehlevan 2023 was argued and shown that in μP, one can observe catapults even at infinite width.
C.2 SILENT ALIGNMENT
In (Atanasov et al. 2022), it argued that linear neural networks trained on whitened data starting from small initialization undergo an alignment of their weights in the task-relevant direction before their function output scale has grown sufficiently. The kernel alignment thus grows well before loss drop, indicating that even at a loss plateau, the NN has learned a meaningful representation of the data. The results were empirically tested relatively small class of more realistic tasks. Here, we extend the class of tasks, and observe silent alignment phenomena for nonlinear networks on realistic and anisotropic tasks.
C.3 3 EDGE OF STABILITY
In (Cohen et al. 2021). in the setting of full-batch gradient descent, it was observed that the maxi- mum hessian eigenvalue grew to the scale of 2/n, where it stabilized. This effect further studied in (Damian et al. 2023). In the SGD setting, a similar effect persists, though the final sharpness holistically the spectrum. See Agarwala et al. 2022) for a theoretical study and discussion in this setting.
modified from the 2/n limit to something that depends
D A TWO-PARAMETER MODEL RECOVERING CATAPULT SCALINGS
In this appendix, describe two two-parameter model which recovers our main phase portrait and unlike the one-parameter model of [Section 4 - exhibits true catapult behavior with just-too-large learning rates the lazy regime. We will first give an analytical treatment, then conclude with simulations.
Our model and argument will be very similar to those used by Lewkowycz et al. 2020)) but with slightly different focus: rather than examining the effect of the width N on catapult behavior, study how the y parameter controls catapult behavior. When width and y are properly disentangled (that is, when recent work (Bordelon & Pehlevan 2023) has shown that deep networks can catapult even at infinite width. The catapult effect does not happen large 7, will ultimately focus small 7.
μP), the width of the network plays no role in determining catapults. Indeed,
19

---

Published as a conference paper at ICLR 2025
We begin with multi-paramert model which will shortly simplify. We consider two-layer linear network f()) = W2 w1 with width N trained on a single example æ,y. We assume that | l2 1, and then without loss of generality, the rotational symmetry of SGD. For MSE loss, we will assume the target has ly| = 1, and we may likewise treat the output scalar without loss of generality. For cross-entropy loss, we will simply assume the output is scalar.
may assume the input is scalar with = 1 due to
W E RN and Wi R- we now have a function value u² V. We now make two observations. First, when N is large, u and v will be orthogonal at initialization. Second, as noted by [ewkowyez et al. (2020), upon training by SGD, u and will each remain in the two-dimensional subspace spanned by the two vectors at initialization. We are thus free to rotate into this subspace and assume u, v E R.2.
Relabeling
We will assume that, at initialization, u == [1, 0) and = [0, 1]. Now observe that, for any loss func- tion C()) depending only on f. will be the case that VuL and VoL u. consequence, it will be the case that, at all later times, will will continue to have v swap(u), where swap is an operation that swaps the two elements of its argument. Denoting u(t) = [M) (1), ug (1)], we thus have v v() [uz(4), u1 (()]- We are down to just two parameters, Wy (1) and (t), with ui (0)
and u2 (0) 0.
We denote the downscaled network output as f f. Because fle=o = 0, we do not need to center it.
D.I ANALYTICAL ARGUMENT FOR CATAPULTS FOR MSE LOSS
In this section, we leverage an argument similar Lewkowycz et (200)) demonstrate how the y parameter controls catapult behavior for MSE loss. Let the loss be LMSE = [ 1)2. We then study the discrete-time dynamics of gradient descent with learning rate 17. Upon gradient updates, we have that
U1;0+1 12:4+1
U1;t U2:0
- f
U2.0 U1d
(20)
7
We then define
1 (Iull2 + 1013) K= 72
2
(uf
nz)::
A=1-f,
(21)
222
where K is the neural tangent kernel of the model and A is the residual. Using these substitutions, find that the discrete-time dynamics yield
J411 - Jt = =,,K,,KKK + Aij..
7
(22)
4f.
K++1 -ii--K
T/Ke +
72
A
Suppose now that we choose too large such that f begins to explode. When f is large, we have that f/A ~ -1. Observe now that, so long as nK, S 4, we find that K will decrease upon subsequent gradient steps, eventually stabilizing the dynamics of f. We thus see that the upper bound on v) in the catapult regime scales as 1) 72
D.2 CROSS-ENTROP LOSS
As with the one-parameter model of Section 4 we may also train our two-parameter model with the binary cross-entropy loss given in Eqquation (9) The dynamics for cross-entropy loss are sub- stantially more difficult to analyze than those for MSE in fact, as noted by Meng et al. (2024), even an ordinary linear model trained with cross-entropy loss will exhibit chaotic dynamic in certain parameter regimes! In actual fact, the dynamics are quite similar, scaling-wise, to those for MSE, except in critical "catapult triangle" in the lazy regime in which Mu < y < 7max where 7eii ~ 222 and 7maa 7y Recall that Our aoo-parmettrr model begins to show similar behavior, but, after sufficiently long training time, bounces to different part of the loss landscape and converges. The dynamics here are similar to
me-parameter model neither converged nor diverged in this regime.
20

---

Published as a conference paper at ICLR 2025
Final MSE of toy model /=u"
Max MSE during training
1.0
10
10
10
0.8
10c
10
divergeace
Leurning rate
divergenoe
10-2
convergence
10
0.4
eatapult
10a
10
1
no cataptit
10
10 -
no training
0.2
10
10-3
10
10
10
10a
10'
Output mulliplier
Output mulliplier
Figure 8: Simulations of our two-parameter model with MSE loss recover all our scalings, plus catapult dynamics. The left plot shows final loss, while the right plot shows the maximum loss reached during training. Grey pixels indicate divergent runs. Catapults occur along narrow band. Note that, for this model, L = 2, so Mmin ~ y in the ultra-rich regime.
Final xent loss of toy model f- n'v
Mix xenl loss during training
1.0
2
10
10
10
Ray oe qol
convergence
0.9
101
divergeace
10c
Leariing sate
divergence
10a
10
0.8
2
catopolt
catapuli
10a
10-5
10
0.7
catapolt
10
10-8
training
10-8
0.6
10"
Output multiplierry
Output multiplier
Figure 9: Simulations of our two-parameter model with cross-entropy loss recover all our scalings, plus catapult dynamics. The left plot shows final loss, while the right plot shows the maximum loss reached during training. Grey pixels indicate divergent runs. Catapult dynamics occur within triangular region. The black subtriangle within the catapult triangle for final loss is made of runs which would eventually converge given substantially more training iterations. Note that, for this model, = 2, so Imnn ry in the ultra-rich regime.
the catapult dynamics described by Damian al. 2023). except that the "unstable direction" sees a straight-walled loss basin instead of a quadratic basin (think L(x) = [x|] instead of [(x) = x²)) We anticipate an analytical treatment in the style of Damian et al. (2023) would be elucidating, but it could make up substantial part of paper in and of itself, so at this point we turn to numerics.
D.3 SIMULATIONS OF THE TWO-PARAMETER TOY MODEL
We numerically simulate our two-parameter model for both loss functions considered. For MSE, we train for T 10 steps. For cross-entropy, train for T 10* steps give the dynamics in the "catapult triangle" more time to converge.
The results shown in Figures and We plot both the final loss and the maximum loss any point in training. In both cases, we see phase portraits very similar to those for our one- paramereet model (Figure 6), but also see the loss exploding before convergence in the catapult phase.
21

---

Published as a conference paper at ICLR 2025
E SCALING RELATIONSHIPS ARE UNAFFECTED BY LEARNING RATE WARMUP AND DECAY
In this section, we empirically verify that adding a warmup and cosine decay to the learning rate, as common in more modern deep learning practice, does not affect the key scaling relationships observed. The models are trained using learning rate warmup for approximately 5% of the total number of steps, followed by cosine decay until the end of training. All models are still trained with SGD, In the next section, we analyze the effect alternative optimizers such as Adam and SignSGD.
08
0.7
06
logn
2
0.5
04
88
02
0.2
-10
-10
logy
logy
Figure 10: Empirical runs of a layer and readout ConvNet without residuals model trained with SGD with MSE loss (left; see Eqqutiin (4) and cross-entropy loss (right; see Eqqation (9) with varying 7; 7) on CIFAR-5m.
F SCALING RELATIONSHIPS FOR ADAM
All analysis thus far treats MSE loss. However, in practice, Adam is more common choice of optimizer when training large models, so is worth asking how our story changes when switching SGD to Adam. We repeat our analysis our toy model for SignSGD - which is equivalent to Adam without momentum - in im/ppendix F. We find a very similar story with different scaling exponents:
Tmin 7max Y when y < and mmin Maax 7 1//L when y > Remarkably, unlike with SGD, the upper and lower scaling exponents are the same in the ultra-rich regime (y > 1) due to the fact that Adam escapes the saddle near initialization in (1)) time. This suggests that Adam may be optimizer of choice for future investigations into the ultra-rich regime.
In this appendix, we obtain both empirical and analytic phase plots and scaling relationships between 77 and under the Adam optimizer. We present variant of the onc-parameter model in Section 4 for training with SignSGD (Kingma 2014). Adam can thought of as essentially SignSGD with momentum. The predictions of the one mer-aaraettrr model are then validated in via simulation.
F.1 THEORETICAL PHASE PLOT FOR SIGNSGD
Adam is essentially SignSGD (Bernstein et al. 2018) with momentum, and the effect of this mo- mentum will only change our story up to order-unity prefactors, so we instead study SignSGD. The update rule for SignSGD is
0 - 7 sign (VoL),
(23)
where 0 is vector of parameters and is the loss. As in Section wish to determine the viable range of learning rates for training the centered, rescaled target function f(w) 7 (w" 1) MSE and cross-entropy losses.
In the lazy regime (y < 1), wish the parameter to move by amount scaling like [desired update] ~ 7. This parameter feels a gradient scaling as [gradient] ~ - so with SGD, we required learning rate scaling as 1) [desired [pdate]/[gradient] ~ 722 However, with SignSGD,
22

---

Published as a conference paper at ICLR 2025
the magnitude of the gradient is irrelevant, and simply want r [desired update] ~ y. Any smaller learning rate will lead to no training, and any larger learning rate will lead to too large a parameter change, greatly overshooting the target.
In the rich regime (y > 1), we wish the parameter w to move by an amount scaling as [desired update] ~ 71/2, By the same argument, we desire v ~ 71/2, Unlike for SGD, the max and min viable learning rates are the same in the ultra-rich regime: instead of a finite "triangle of optimizability." the convergent region extends to arbitrarily large "y. This is because the model escapes the saddle at initialization immediately because the (effectivelly small parameters do not suppress each other's 's gradients. Because of this fast saddle escale, the silent alignment effect does not occur under the SignSGD Adam optimizers.
unknown to the authors whether catapults occur for SignSGD, and if so, whether there is a "catapult triangle" in the phase diagram seen with cross-entropy loss and SGD Figure l/b,d).
Results of simulations of our one-parameter toy model with SignSGD are shown in Figure 11
Final MSE of 1-puram model f w
Final loss of 1-param toy model f= u
Optimized with SignSGD
Optimized with SignSGD
10
10P
Al for 20
Fewa Cory 26
10*
Texa: Boos 7-t.
0.8
10a
Learning rate ff
0.9
10-1
farge
Learning rate
10-1
convergenca
convargence
10-4
10-6
0.8
0.4
10-11
10-16
20 training
Do training
0.6
10-6
10
10
10a
10-6
10
10
10
Output multiplier
Output multiplier
Figure 11: Simulations of our one mne-paraeeter model with SignSGD with MSE loss (left; see Equa- tion (4) and cross-entropy loss (right; see Equation (9) with varying 7, 1 exhibit predicted scaling relationships. Simulations with 7. 5. Unlike with SGD, SignSGD does not diverge; instead, points in the upper region oscillate forever.
F.2 EMPIRICAL PLOTS OF PHASE DIAGRAM ADAM
08
0.8
07
logn
04
0,4
66
-6
0.3
7-Y
-
6
03
0.1
-2
0
Figure 12: Empirical runs of layer and readout ConvNet without residuals model trained with Adam with MSE loss (left; see Equation (4)) and cross-entropy loss (right; see Equation (9) with varying Y, 17 CIFAR-5m.
We supplement the theoretical plots from Figure IT with empirical plots in a Convolutional Neural Network with 3 layers and readout layer, residuals, trained CIFAR-5m using Adam with
23

---

Published as a conference paper at ICLR 2025
MSE and Cross Entropy losses 12 The models are trained using learning rate warmup for approxi- mately 5% the total number of steps, followed by a cosine decay until the end of training.
24

---

Published as a conference paper at ICLR 2025
G ADDITIONAL RESULTS ON SINGLE-INDEX MODELS
In this section we empirically show that the results of the main text transfer to the commonly studied class of "single index models". Here, we train 3-layer MLP of width 200 with on single-index model of the form the same as in the main text, namely the MLP parameterized in pP and centered. We show the phase portrait of training this model in an online setting with SGD in Figure 13 This portrait is consistent with that predicted by our linear model.
(w- x² where wo ~ (0, I), (0, I), w E RR, x E IRD Our setting
Orline Single-Index 0 = N 200, L
Check Consistency across Wisths, Quadratic Task D 15, 2,
100
Ni-10h t0, y- qynte-01
y--1-002 meI1e=11 N-500, 2- oymieva1
0.0
Op =
-0.5
Lozs
10a2
-1.0
19a1
-.0
10a
-10
-2.5
-12
100
10
10a
10a
10a
-30
Time Step
Online Loss at t- 900, Single-index Model
Check Consistency across Widths, Quadratic Task D= 15, 3=3,
D 15, k 3, N 200, -
1.0
18
10-1
76
1.5
10-1
-2.0
Loss
10
-2.
-8
N-100, 1y-1e102, 2o-
-3.0
20200, 19-1e400, op- 10-01
-10
500, we
10-5
10'
10a
107
10f -12
-3.5
Time Step
Figure 13: Empirical runs of a layer and readout MLP model trained with SGD and MSE loss on quadratic single-index model y (w xx) a) Verification of consistency across widths in μP b) The phase plot observed as in Figure c) and d) illustrate the same, but for 10 (w x)33
H ADDITIONAL RESNET AND VIT PLOTS
Here we supplement the 77ー phase plots in the main text with additional plots across different ar- chitectures and datasets. We report results for ResNets and visiont transformers. Across all settings, we recover the phase portraits sketched out in Figures [
25

---

Published as a conference paper at ICLR 2025
Conuter Tharing Aruutry Tndtrogee mardue
JEVI Trankin Accanay Trylmngennt Aupmertatoon
Doeth ogph Ontine Accurccy SHAM
W-512, B-222) b--620 bes-sent
120
335
1977
Ugy
(a)
(b)
(c)
Figure 14: Training accuracy for (a) ConvNets and (b) ViTs (with QK-norm) trained TinyIma- genet, and (c) ViTs without LayerNorm trained on CIFAR-5M for classification using Cross Entropy loss, parameterized with uP. Removing LayerNorm causes sharp maximum learning rate cut-off in the training of ViTs. Hyperparameters: batch size 128, using random data augmentations for Tinylmagenet, head, 3 transformer blocks.
JP Temp=led Augmentation CIFAR-10
N#226, Lib,, blocks, Bir
0.22
022
3252
0.18
1225
2200
0.14
0.14
0.32
2102
12
0.10
Y
(a)
(b)
Figure 15: The effect of introducing a temperature 1/T parameter in the attention softmax calcula- tion. For large T. this makes the attention matrix more uniform which leads to easier optimizability in the large learning rate region. Training is done on CIFAR- 10 with random data augmentations, and with LayerNorm.
Tiritrapann with Data Augmentation 4--32. Lec. blocks B--288, losa-amm
54'-3
1.107
46.327
43x10"
Figure 16: Stepwise loss trajectories can be seen even in the very realistic setting of vision trans- former on TinylmageNet. The trajectories here exhibit two lower plateaus addition to the long plateau from t = 0 until the first drop.
26

---

Published as a conference paper at ICLR 2025
1 EFFECT OF WIDTH AND CHECK OF J.-ONSISTENCY
In this section we provide checks of μ&P working. We do this by both plotting in Figure 17 and Figure 18 as well as their CNN and croseentropy analogues. We observe the following:
The loss curves as function of time across widths begin to match as N grows, confirming that they are converging to an "infinite width" loss curve
The performance vs the learning rate 17 across width is stable, as expected in p.P.
Importantly, consequence of consistency across widths sufficiently large N, one does not need to additionally vary the width parameter in order to exhaustively sweep out the space feature learning NNs.
1.1
7-- D-DEENNDENCY OF μ-CONSISTENCY
Here, we study to what extent width-consistency is violated as a function of T.
Te-00 VNITT-IM L-3 pescitio
Orline 1e-33 HMSTT- MiP Lo Bu112 bot.unee
L- 3-120 basei7se
4--22
M-256
12
(a)
(b)
(c)
Figure 17: We sweep over N different values of y. We see that the dynamics across different N are closer at larger indicating they are converging faster to their infinite width feature learning limit than the lazy networks are. We further see that during training we achieve wider is better effect, expected from JPP.
Loceatt See22forr 7e-00 MASSTII MLP 1-3, 1.3 min
Loceate Jew33 1e-02 MNUTTIM MLP L-3 128, (cos
Lons Jeedatoryy Jee33 MAJGT-H 1-3 R-125, I065- mee
251
N-32
184
512
12-1
12-1
12
144
(a)
(b)
(c)
Figure 18: We sweep over N at different values of y and track the optimal 17 across N. Plotted is the loss at the end of training. We observe that the minimum essentially remains stable.
27

---

Published as a conference paper at ICLR 2025
Orline ye 2e-02 CHAR--M Iossemas
Oncine Losa 1:+00 CITAR-3N
Losa for ye 12402 CITAR-3N CNN
B- 256 Iospemmss
1-104
Ex
V-60
N-125
3x0m
121
V=122
1+255
14
30
(a)
(b)
(c)
Figure 19: The same as 17 but for a CNN on CIFAR-5M with MSE loss.
104
twleeeeeeeree OHAN Led, 2-020 MSSmIrse
2410
Led, be 25.2 500 Yeal
Led, Bw220. usserise
N-10
1--10
41184
V=20
14
41-22
4-255
3430
14-25
44-25
38
(a)
(b)
Figure 20: The same as 18 but for CNN on CIFAR-5M with MSE loss.
Loos of 3o 53 for you to 03 CIFAR 5M CNN
Loss 2t tw320-d0 2pe 20 CINAR-SH ChN Led., Bw235. kosmmsent
Loss at tt 05 for Sen CPAAI-MM ChN
Led, Bw235, kssawart
LOSS
N1-16
20
N=20
N--64
N-120
No 120
N-120
10.4
37
17
(a)
(b)
(c)
Figure 21: The same as 18
but for a CNN on CIFAR-5M with cross-entropy loss.
Online for yo 1e-33 C1R.-IIAA---MOOO CWN
On foretor Te400 CIFAR-SN ChN Ba2SE Wosssert
Online y=1e=03 CIFAA-SM CNN
Le4, B-235 lopmmont
LOSS
14x1P
M-10
14K17
-22
(a)
(b)
(c)
Figure 22: The same as
but for a CNN on CIFAR-5M with cross-entropy loss.
28

---

Published as a conference paper at ICLR 2025
J EFFECT OF DEPTH
In this section we vary the depth L. for both feed-forward ML.Ps and for ResNets. In the MLP setting, find a straightforward change in the 7 - 7 plane. In the rich regime, the maximal observable learning rate scales y consistent with our model in
J.1 MLPS
Accuracy WMGTTMM MLP N--256 L-2, D-228, b-le-3), Das-mme
Accurscy HUIST-1M
MMST-1M NUP
No-25. L-3, B-328 1-1e-23 Jose
N--26. L-4, B-128, t-1e- los
4-
A2
42
-10
0gy
9g
bagy
(a)
(b)
(c)
Figure 23: Accuracies after 1000 steps of MLPs trained on MNIST-IM of depth a) I. 2, b)
I. = 4. We see that only for the single hidden layer neural network arbitrarily large 7 values be optimized. As L grows in the feed-forward setting, the triangle of large optimizability shrinks.
Vil BO
(b)
We restrict ourselves the line 17 3 if y and 3/2// if 1 and only plot pixels in the 17-y grid that within 10% of this line.
Figure 24: The loss trajectories from selected pixels in Figure 23
29

---

Published as a conference paper at ICLR 2025
J.2 RESNETS AND DEPTH-uP
In this section we test the effect of the parameterization on deep neural networks with residual connections. Specifically, in Tigure[26Jan Figure 25 we evaluate the phase plots of neural networks with increasing depths parameterized in Depth-juP, as detailed in (Bordelon et al. Littwin 2023 and p..
2023 Yang &
Acaracy Bata Augmertaiion Resited Cantte.
Accurayy Sata Augmentras
Accuracy Data Hagmertarin OTAAU20 Restual Canvles,
Resitedd Caneles Leb, bucks $-020
h-120,
Insumant
M-220, frooo,, Bw328, foomeor.
0
(a)
(b)
(c)
Figure 25: We compare ConvNets with residual connections parameterized in the uP regime in the online setting by training CIFAR-5M and testing CIFAR-IO. In this setting the scaling seems to change as we increase the depth of the networks by adding multiple residual blocks. the very deep setting shown in (c), the very rich networks have flat 17 ~ 7 scaling, unlike in the Depth-uP case where networks across all depths have similar scalings.
Acoracy Bata Augmantarinn CTA--E b-112. L-3,1 Hbcks 9-123, 6-4600 lze-wont
Accusscy Data Augmentrtnn CRA-C0
Acorscy Data Augmantacion OTAR-1
K-122 L-d, bodk, 3 9-129 4-4000 lzer-owrt
N-328, L-12, bbcco,, 8-128, E-4005 jose-eart
65
4A
-1
-1
(a)
(b)
(c)
Figure 26: We compare ConvNets with residual connections parameterized in the Depth-ud regime in the online setting by training on CIFAR-5M and testing on CIFAR-10. Note that the scaling remains much more stable than in Figure 25 We hypothesize that this is due to the downscaling of the residual branch by VT.
30

---

Published as a conference paper at ICLR 2025
K EFFECT OF BATCH STZE
In this section, we verify that in the online limit, different batch sizes have a relatively negligible effect on the dynamics, The batch noise does however make it that smaller batches cause higher variance in the loss. This is the tradeoff for being able to train to a lower loss at fixed dataset.
Ontine 28:55 WART--1 MLY
Ondine Loga NIBST-IM
MMST-MM
18-
D:115
Dm-22
D--4
2-
184 2-2
(a)
(b)
(c)
Figure 27: The temporal dynamics of online loss versus time across different values of at fixed 77. Predictably, smaller cause larger batch noise variations, but the mean loss path is essentially identical across orders of magnitude in B. Upon fixing a data or compute constraint, smaller B networks can train to lower loss values. Larger B networks however are preferable under a wall- clock time constraint.,
Accuracy for One Pass MNIST-1M N=256, L=3, B-16, t=1e+03, loss=mse
Accuracy for One Pass MNIST-1M
N=256, L-3, B=256, t=1e+03, loss=mse
08
0.8
0.6
0.6
log
1900
77
V²
17 yd
0.2
0.2
77
-M
7 y2
-10
0.0 -10
0.0
log
log Y
(a)
(b)
Figure 28: The 1--y phase plot for two MLPs trained for 1000 steps MNIST-IM for B 16 and b) B = 256. The location of the dashed lines is unchanged from a) to b) We see that the scalings of with y are identical. The lazy regime observes essentially change. However, large y, the maximal optimizable is decreased when B shrinks, leading to a more constrained triangle of optimizability large 7.
31

---

Published as a conference paper at ICLR 2025
L CONSISTENCY ACROSS 7
In this section we consider how networks successive y behave in both the lazy regime y < 1 and in the rich regime Y >> 1. In the lazy regime, past certain threshold 7, often the order of 1e 3. 3, we see that for the datasets and models considered, the dynamics are nearly identical throughout training. This is strongly indicative that these values of y are small enough to
L.1 LAZY CONSISTENCY
In this section we verify that the lazy limit is indeed reached at y = One simple way to do this is to confirm that, conditional on scaling ~ 7² the dynamics agree across no matter how small becomes. We plot this in Figure 29
be
Verification that the lazy limit is reached in practice for CNNs and ViTs trained on CIFAR-5M under MSE and cross-entropy losses. We take 77 = y2 in all plots. Except for a few larger y networks, see that the loss curves agrec very sharply across arbitrarily small values of y. The majority of the small curves are on top of one another in these plots.
Figure 29:
L.2 EARLY TIME RICH CONSISTENCY
In this section, we verify that NNs in the large y limit exhibit similar behavior at early time when time is rescaled to be measured T = rpt/"y when 17 is scaled as 2/4. Thus T ty 1422/L This
reasoning is explain in Section We highlight this empirically in Figure 30
40V
N-512. Lec,
B-235 Insment
13:27
36-32
30 Sexcc
2018
20104 Se
2-3
6230
11:24
60y
Figure 30: Consistency of dynamics in the large early time setting for CNNs and MLPs trained on CIFAR-5M. We take 17 ~ 2/2 in all plots.
L.3 LATE TIME RICH CONSISTENCY
We next study the late time dynamics at large 7. We find that large y networks reach the same loss values networks with y ~ O(1) after sufficient time passes for them escape the plateau.
32

---

Published as a conference paper at ICLR 2025
Oekic hos for CRAN-MM
Mode Rodhs Iefovtio
187
Led
30
1.30
14.0
4X0
Deble woth MAME
2-258 Insomett
16-
T1103 Ioeas 30-00 Levi2 heeen bends
83, 1eme
Les 3 Seedd g-lmm
5430
30
Figure 31: Consistency of dynamics in the large y late time setting for CNNs and MLPs trained on CIFAR-5M. We take 7/ = 22/2 in all plots. In the inline, we zoom into the location near the first drop, and see that after adopting this time rescaling, all networks achieve the drop at the exact same time.
We also highlight how, after sufficiently long time has passed, all large y networks swept over achieve comparable final losses. One such example was shown Figure 2(b)
Accuracy,
Accuracy.
N-96, B-256,
N-96, B-256,
t=20+00
t-2e+00
t=le+01
tele+01
5=5e401
1=5e1001
L-2+402
t=Sc#03
t=
t-2c+04
t-26+04
tale+05
tale+05
0.2
t=5c+05
t-20+06
t-2e+06
12, L-3, B-258,
N--122 L-3, B-256, lowemeent
7.5
1=2+400
1-2++00
t-le401
1-1e101
t=5ee01
t=5ex01
b=2c+02
b=
taSe+03
ta
1=2e406
1-2++34
t=le405
[-1ei00
L=50105
t=2c+06
Final accuracies across ? for a broader range of networks. We see that for very large NNs that have not yet escaped the saddle point, the accuracy is not competitive. However, after enough time has passed, these networks also achieve comparable values of the accuracy. The longer optimization time makes such networks undesirable from practical perspective.
Figure 32:
33

---

Published as a conference paper at ICLR 2025
M FURTHER HESSIAN PLOTS
Hassian Egenvalues, CIFANSMM
120
N=512. 5-256. losd-mae
200
1w45
Eigeivsless
1604
Starsnics
1601
1e+00
30'1
3cr
10a
10m
30.1
Figure 33: MLP on CIFAR-5M under MSE loss. a) The scaling of the sharpness at the end of training function of y, exhibiting sharp transition between lazy and rich. b) The sharpness as function of time across At large "y and early times, calculating the sharpness runs into numerical stability issues that make unreliable.
Hessian CIFAR-SM CNN
Nom 9b, Bmm250, lossixent
Sharpness of Lose Landscape, CNN. GEARSSM.
100
N-25. B-255, losomaent
Eianvalies
167
331
y
3e-03
100
10
2
36'1
3+403
3e406
10-1
Figure 34: As in Figure 33 but for a CNN under crossentropy loss, and including the top eigen- values in a).
Fiial Hessian Bipervoluss, Dupthuu Reciliaal Convliitt
232
36-1
36-1
100
Figure 35: Depth-pP.
As in Figure 33, but for Residual CNN under crossentropy loss, parameterized with
34

---

Published as a conference paper at ICLR 2025
N
FURTHER CATAPULT PLOTS
Ondhe MRISTEN NLF
Ortin foee ON
N.256 0,0 100
180
36:07
16-07
36-55
10-7
26:04
16:06
1e:44
2604.
16:55
26.22
1844
10400 0=12401
13
y=11400. g=
(b)
Figure 36: The catapult effect along the 7) ~ 2² line for the MSE loss for MLPs on MNIST-IM and CIFAR-5M
bn jas for PAR-5H ONN
26.55
06.55
10-33
1e-33
16-01
2-20-0
(a)
(b)
Figure 37: The catapult effect along the line for the cross- cross-entropy loss for CNNs and ViTs.
O ALTERNATIVE SCALINGS OF 7) WITH 7
In this section, we briefly illustrate how taking alternative scalings of 1 different from 72/ can lead network that either undertrains or does not converge at large 7.
35

---

Published as a conference paper at ICLR 2025
Accuracy for One Pass MNIST-1M MLP Nm256, Lm3, Bm128. t=1e+03, loss=mse
Accuracy for One Pass CIFAR-SM CNN
Nm64, Lm4, Bm128, t=1e+03, loss
1.0
1.0
2.6
0.6
10g
-2
0.4
y
y 202
-10
0.0 -30
0.0
lag
log Y
Figure 38: Sketch of phase plots with alternative scalings of with y (overlaid dashed blue lines). We expect the ~ 23/ scaling to diverge at large 7 at any T' while the ~ 2/((++1) will remain convergent, but under-train compared to the optimal learning rate
the Loo IN
dedine Loot CHRA
30
1 m
2a-02
Jen 2e2ee
y-
164
2k20
LHd B-22E ReSeEE 2012
Bics Lod D-120 ba-raa 27
1-220:00
104
(weds, Soed?
B-256: LHD 8-228 Barrrus 2012
Reds Bo..0. botn
164
Sen 00 20
Figure 39: Alternative scalings of 17 7/3/ for a) an 3 MLP MNIST-IM and b) L CNN on CIFAR-5M. Large values not shown are those that have diverged. Indeed, for this scaling, values above a are found to have divergent loss dynamics. Similarly, we plot v -,2((2+1) in the same settings in c) and d). All of these networks train, but they do so less efficiently than the optimal scaling of 2 We illustrate the optimal scaling in figures e) and f), and observe that they achieve better final loss values than c) and d) while still allowing arbitrarily large values of y to train.
36

---

Published as a conference paper at ICLR 2025
P KERNEL-TARGET ALIGNMENT
In this section we define what we mean by kernel-target alignment. In the case a single-class task with labels y for each a, we consider kernel associated kernel K originating from an (initialized or trained) neural network. This can be either the neural tangent kernl,, [TK x') = f()) ( f(('), or the hidden layer kernel, K' (x x') ( We evaluate P Px kernel gram matrices K on a held out test set of P points. We stack the corresponding labels into a vector y Rp. The associated kernel-targer alignment given by:
y T Ky
KTA(K, y)
(24)
| 1/
Here ||||| some appropriate norm of the kernel matrix. There are many options, including Frobe- nius norm, operator norm (maximum singular value), and nuclear norm (sum of singular values). The alignment during the loss plateau is quite sensitive to the choice of norm on K,. In practice, we've found that the nuclear norm produces the best results, with Frobenius and Operator norm metrics sometimes resulting in dip of the alignment before the loss drop rather than a growth.
A related quantity, centered kernel alignment (CKA) Cortes et al. 2012), has been argued in (Ko- rmblith et al. 2019) to be good candidate for the comparison of neural representations. In practice, found direct CKA measures between NTKs of different networks to very noisy and sensitive. CKA on hidden layer activations yielded more reliable plots.
37

---

Published as a conference paper at ICLR 2025
Q FURTHER SILENT ALIGNMENT PLOTS
Allggemett CIAR--M CMN
Alignment for CIAAL-SM CNN N=-5,L-4, B-128, loss-mse
N- Len, S- 10s2 17.01000-1100
Lows
Alghment
Algrmeet
044
Sharpress
Sharpness
y=12101
1e-01
y=1c102
V- 1e+02
1e403
Ie+03
Te4fc
1e-04
042
15
397
124
300
402
237 10" 130
00y
Figure 40: Silent Alignment for CNN on CIFAR-5M. There is a subtle, but noticeable increase in the alignment metric before the loss drop. a) With 1 on the x-axis b) with nj// on the x-axis. For CNN, we decrease in the alignment metric after the loss drop. This is due smaller singular values growing, causing the nuclear norm in the denominator to grow faster than the numerator.
Algrmarn MAST
144E)
Figure 41: Silent Alignment Across Different Depth MLPs on MNIST- 1M
Algrmmntt for CAN Depth (he Hw522, de lossece
33
Figure 42: Silent Alignment Across ViTs on CIFAR-5M, evaluated on a single batch of CIFAR-10
38

---

Published as a conference paper at ICLR 2025
R FURTHER FUNCTION SIMILARITIES
K--12, L-3,0-25,, tos moe. 1-2e004
N--A Lind, D--254, (ows man 20443
2.2
Buth See 200
1.0
pio 237
Gae 1oe22
2-100
pi 2.00
Rist
512
Figure 43: Further comparisons of learned functions between lazy and rich networks.
losses in Tim ML3 CRAISM
Lesags Tims CNAL WAST-UM
Le 20:55
y 16-05 1 1o 16-00,
Compettor of Logit OUT CPARAN
Conmact of Lagts DWL PASTOM.
03
Cumpubb
Comsattton
Lacyllic
Lao-Rich
64
Figure 44: Correlations coefficients between learned functions outputs test set between lazy and rich networks time
WEIGHT AND ACTIVATION MOVEMENT ACROSS TIME AND RICHNESS
In this section, we confirm that larger values of y correlate directly to weight movement across variety of norms. We consistently verify that the weights deviate as |[W((t) W' 1/L across layers, times, and 7.
39

---

Published as a conference paper at ICLR 2025
Layen
1000
30
Figure 45: Weight movement (t) - W0) for a CNN trained on CIFAR-5M. We plot these across different layers different steps / of training and clear trend that larger value of 7 correlate with meaningfully larger movement across a vari- ctt of norms. Specifically, we see that all norms of matrices scale as 1/L after sufficient training steps. The dashed black line represents curve ///4 for reference. a) Frobenius norm, b) Spec- tral/Operator nom, c) Nuclear norm.
function of 7. We very sharp
40


Reviews:
[{"review_id": "ZeEh1pQ9vn", "reviewer": "Reviewer_aPNz", "summary": "The paper studies the scaling law between $\\gamma$, the feature learning strength parameter, and $\\eta$, the learning rate, both empirically in a diverse set of data (MNIST, CIFAR, TinyImageNet) and model architectures (MLPs, CNNs, Vision-Transformers), and theoretically in simple setting with deep linear network. They observed a characteristic phase portrait in the $\\gamma-\\eta$ plane, and showed that the learning rate scales as $\\eta \\sim \\gamma^2$ in the lazy regime ($\\gamma << 1$) and $\\eta \\sim \\gamma^{2/L}$ in the rich regime ($\\gamma >> 1$). They confirmed this phase portrait empirically by sweeping over every pair of $\\gamma, \\eta$ in a log-spaced grid running from $\\gamma = 10^{-5}$ to $10^5$ and analytically by theoretically analyzing a deep linear network models for both MSE loss and cross-entropy loss. They also observed and reported various dynamical phenomenon such as \"catapult effect\", \"silent alignment\", and \"progressive sharpening\".", "questions": "**Questions**\n1. [L137] Why further downscale by $1/\\gamma$? Isn’t the $\\mu P$ scaling already scale by a factor of $1/\\gamma$?\n2. I couldn't find any information in the paper about how many repetition were run for each empirical results? Would appreciate if you can provide some information about this.\n\n**Suggestion**\n1. [L068] Key contribution should be clear and emphasizing the key points (usually 3-4 key points). This current key contribution section is too long and difficult to follow. I’d appreciate if the authors can re-organize this section to make the key contribution clearer to the readers.\n2. [L068] Hyperlink each key contribution with the section and/or figure that demonstrate the key contribution so that readers have easier time to follow.\n3. [L215] Explicitly mention whether this scaling results are for MSE or cross-entropy loss, since they have different scaling. I was quite confused when the author states that “At small $\\gamma$, $\\eta \\sim \\gamma^2$ for lazy networks”, but Fig 2d shows that $\\eta \\sim \\gamma$ instead, but later I realized that the claim “At small $\\gamma$, $\\eta \\sim \\gamma^2$ for lazy networks” is for MSE, and the authors have another section for cross-entropy loss.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper provides extensive empirical study for the scaling law of $\\eta$ (learning rate) and $\\gamma$ (feature learning strength), and also provided an analytical explanation using deep linear network that matches with the empirical results.\n\n**Originality**\n1. The paper investigates the scaling relationship between feature learning strength $\\gamma$ and learning rate $\\eta$ in both lazy, rich, and ultra-rich regime. To my knowledge, I haven't seen work that extensively address this topic.\n\n**Quality**\n1. Empirical study: The paper provided an extensive empirical study across $\\gamma$ and $\\eta$ plane, ranging from $\\gamma = 10^{-5}$ to $10^5$ and $\\eta = 10^{12}$ to $10^{-12}$ with a diverse set of dataset (MNIST, CIFAR, TinyImageNet) and model architectures (MLP, CNN, Vision Transformer), and confirmed that the claimed scaling laws are observed across different datasets and models.\n2. Theoretical analysis: In a simple settings (deep linear network), the paper derived the $\\gamma-\\eta$ scaling law for both MSE and cross-entropy loss, and showed that the theoretical results match with the empirical observations. The authors also produced an analytical phase portraits in Fig 6 that match with the empirical phase portraits from parameters sweeping.\n3. Analysis of dynamical phenomenon: The paper also reported various learning dynamic phenomenon that was reported in earlier works, including catapult effect, silent alignment, step-wise loss drops, progressive sharpening, in various $\\gamma-\\eta$ region that they perform the  parameter sweeping.\n\n**Significance**\n1. The paper showed the scaling relationship between feature learning strength $\\gamma$ and learning rate $\\eta$ in both lazy, rich, and ultra-rich regime. This scaling law enables the authors to find the range of possible learning rate to train models in ultra-rich regime, an area that few previous work touched on. Using the suitable learning rate, the authors show that in models trained in the ultra-rich regime can exceed or match the performance of models trained in rich regime, which open a potential area to improve current model training dynamics.", "weaknesses": "While the paper addressed a significant topic in the feature learning area and provided reasonable empirical and theoretical results, I think the clarity of the paper still have lots of room for improvement to help the readers have an easier time to read the paper. Specifically, the authors can consider to try to re-organize the paper flow to focus on the key points: Since the key results of the paper is the scaling relationship between feature learning strength $\\gamma$ and learning rate $\\eta$, the paper can be reorganize so that these key theoretical and empirical results come first and emphasized in the paper, instead of leaving the theoretical results at the end of the paper, and mention many other detailed results in the middle of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "a72jYaofjc", "reviewer": "Reviewer_CuAx", "summary": "This paper systematically studies the effect of $\\gamma$, which is the scale of the output layer of the neural networks, on the optimization and generalization dynamics of neural networks through experiments. They recovered a couple of phenomena that have been analyzed in the literature, including lazying training, catapult effect, silent alignment, etc. They also observed that the generalization performance is similar when $gamma$ is greater than $1$ if the network is trained sufficiently long. Furthermore, they use a simple deep linear neural network with a single univariate data to explain the optimal learning rate scaling.", "questions": "1. I am curious how do you conclude that the optimal learning rate is $\\gamma^{2/L}$ for $L$-th layer neural network? I understand that for the deep linear network with a single univariate sample the optimal one is that value, but I am surprised that for deep nonlinear networks it still holds, and the power is quite specific. Does the generalization performance have a big gap when the learning rate is $3/L$ or $2/(L+1)$?\n\n2. In practice, the learning rate decay is useful to improve the generalization performance, and it may avoid optimization from divergence. Do you think the optimal learning rate will be different if you include learning rate decay?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. This paper is well-written and the logic is easy to follow. The experiments are nicely done and the theoretical results are well presented. \n\n2. There are some interesting empirical observations. For example, they observe that in the feature learning regime i.e., $\\gamma \\geq 1$, the generalization performance will end up being similar with different $\\gamma$, if the learning rate is appropriately chosen. I also find the connection between silent alignment and edge os stability quite interesting and worth further investigation.", "weaknesses": "1. One concern is the significance of this paper. Feature learning is an interesting topic, but I believe that this kind of parameterization is not commonly used in practice, i.e., large $\\gamma$ with the model being centered. Furthermore, the learning rate used in this paper is constant hence is far from practice. Therefore, I am not sure if the insights observed in this paper can shed light on practical network training.\n\n2. This paper may provide valuable insights for theoretical analysis. However, many of the phenomena discussed, such as lazy training, the catapult phase, and silent alignment, have already been extensively studied. Although the ultra-rich regime is analyzed in depth, it does not appear particularly interesting, as the generalization performance turns out to be similar to that of $\\gamma=1$. All in all, the paper presents a nice summary of the network training dynamics, but I do not find anything particularly new or significant.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qrXlnG1Vkh", "reviewer": "Reviewer_nULc", "summary": "The authors conduct a systematic and comprehensive empirical study on the effect of scaling the hyperparameter $\\gamma$, located in the final layer, which controls the strength of feature learning, across a variety of datasets and tasks. In this context, a small $\\gamma$ corresponds to the kernel regime, while $\\gamma = 1$ represents the mean-field regime, where feature learning can still occur even with infinite width. \n\nFirst, the authors examine the relationship between $\\gamma$ and the maximal learning rate, empirically demonstrating that the maximal learning rate scales nontrivially with $\\gamma$. Using this optimal learning rate scaling, they then explore the ultra-rich regime, $\\gamma \\gg 1$. In this large $\\gamma$ regime, the network exhibits consistent training dynamics across different values of $\\gamma$. Also, the model always achieves optimal performance when $\\gamma$ is sufficiently large.", "questions": "See the weakness part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The experiments are conducted systematically, making the results highly convincing. The paper is also well-written and clear.\n\n2. The observation that networks with a larger $\\gamma$ can often achieve the same or even better performance than those with $\\gamma=1$ after sufficient training time is intriguing. This suggests a potential new technique to improve performance in practice. It's also interesting to see that networks with large $\\gamma$ exhibit very similar training dynamics after an appropriate time rescaling.", "weaknesses": "1. I am somewhat unclear on how these results translate to practical applications. While the empirical findings in this paper are robust, the setup does not precisely match real-world scenarios (e.g. the architectures are different). I wonder what insights could be derived if we consider a standard practical setup.\n\n2. I believe the paper would benefit from a more detailed and mathematically formulated introduction to the networks’ structure and parametrization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NzlERzkoIy", "reviewer": "Reviewer_EZic", "summary": "The paper empirically studies the relationships among the scaling hyperparameter $\\gamma$, learning rate $\\eta$ in neural networks, and the overall performance of neural networks.", "questions": "1. In Figure 4, which illustrates the dynamical phenomena under different $\\gamma$ settings, the conditions for 4(a) and 4(b) differ (CCN for CIFAR and MLP for MNIST, respectively). Are there similar saddle point phenomena as observed in 4(b) under the conditions of 4(a), or catapult phenomena like those in 4(a) under the conditions of 4(b)?\n2. Are there differences in the weights of the first few layers under various $(\\gamma,\\eta)$ combinations, in addition to the weights of the final layer shown in Figure 5?\n3. What are the implications of this work on representation learning?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is clearly written and well-organized.\n2. The authors conduct a variety of experiments across multiple datasets and network architectures, exploring different combinations of \n$(\\gamma,\\eta)$. The figures presented are particularly effective in conveying the findings.\n3. The optimal scaling of $(\\gamma,\\eta)$ offers valuable insights for selecting hyperparameters in empirical studies.", "weaknesses": "1. Although the title emphasizes \"Feature Learning\", the paper does not adequately elaborate on this concept. There is insufficient empirical evidence regarding whether and how feature learning occurs when $\\gamma\\gg1$. I recommend that the authors provide more detailed results comparing the weights of networks that exhibit feature learning to those that behave according to the NTK regime.\n2. The theoretical framework relies on an oversimplified linear model with a constant target. To gain a better understanding of feature learning, I suggest the authors analyze more complex targets, at least simple single-index models such as $f(x)=\\langle a, x \\rangle^p$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)