Dimension-adapted Momentum Outscales SGD
Authors: Damien Ferbach
Katie Everett
Gauthier Gidel
Elliot Paquette
Courtney Paquette
Venue: NeurIPS 2025 spotlight ()
Keywords: power-law random features
scaling laws
momentum
random matrix theory
acceleration
stochastic optimization
compute-optimal
Primary area: theory
TL;DR: We prove that stochastic momentum can improve the scaling law exponents over SGD on power-law random features by selecting hyperparameters to properly depend on data dimension or model size.

Abstract: We investigate scaling laws for stochastic momentum algorithms on the power law random features model, parameterized by data complexity, target complexity, and model size. When trained with a stochastic momentum algorithm, our analysis reveals four distinct loss curve shapes determined by varying data-target complexities. While traditional stochastic gradient descent with momentum (SGD-M) yields identical scaling law exponents to SGD, dimension-adapted Nesterov acceleration (DANA) improves these exponents by scaling momentum hyperparameters based on model size and data complexity. This outscaling phenomenon, which also improves compute-optimal scaling behavior, is achieved by DANA across a broad range of data and target complexities, while traditional methods fall short. Extensive experiments on high-dimensional synthetic quadratics validate our theoretical predictions and large-scale text experiments with LSTMs show DANA's improved loss exponents over SGD hold in a practical setting.

Dimension-adapted Momentum Outscales SGD
Damien Ferbach"
Katie Everett Google DeepMind & MIT
Gauthier Gidel
Mila Université de Montréal
Mila & Université de Montréal gidelgau@mile la quebec
ferbach. damien@gmail. com
everettkógoogle. com
Elliot Paquette
Courtney Paquette
McGill University
Google DeepMind & McGill University courtney paquette@mcgill.ccomm
el elliot. pll..t.paquueteenngggl...cc
Abstract
We investigate scaling laws for stochastic momentum algorithms with small batch the power law random features model, parameterized by data complexity, target complexity, and model size. When trained with a stochastic momentum algo- rithm, our analysis reveals four distinct loss curve shapes determined by varying data-target complexities. While traditional stochastic gradient descent with momen- (SGD-M) yields identical scaling law exponents to SGD, dimension-adapted Nesterov acceleration (DANA) improves these exponents by scaling momentum hyperparameters based on model size and data complexity. This outscaling phe- nomenon, which also improves compute-optimal scaling behavior, is achieved by DANA across a broad range of data and target complexities, while traditional meth- ods fall short. Extensive experiments on high-dimensional synthetic quadratics validate our theoretical predictions and large- scale text experiments with LSTMs show DANA's improved loss exponents over SGD hold in a practical setting.
1 Introduction
When pretraining large neural networks, the loss typically scales like a power law with respect to the amount of data, number of parameters, and total amount of compute [54]. Scaling laws for the loss function (0) in their simplest form are (0,)) x 1- + d-r where {0, E Rd) is a sequence of iterates generated by stochastic algorithm." The exponents and of practical importance because they control the number samples and parameters needed to attain a desired loss value.
While model architectures and training methods have advanced rapidly, it has long been unclear whether innovations optimization algorithms could fundamentally change the exponents of these power laws [48]. Some evidence suggests that major advances like the Adam optimizer [56] primarily improve the constants in the scaling law rather than improving its exponent [48].
Moreover, recent work has extensively investigated how various algorithmic parameters such as learning rates [103] and batch sizes [70] should scale the model sizes and compute grow. However, momentum parameters are typically treated as fixed constants [93] rather than dimension/compute dependent quantities, despite their widespread use in large model training. This leads the question:
Corresponding author; website: https: https: httpe:://dmmieeffeeeeee www.fferrachh.ithht... git.ub. 1 'Canada CIFAR Al Chair
'The authors contributed equally to the paper.
'Notation: x y means there exist constants and C (both independent of such that eg S Cg. We use (x)+ to denote max(0, 20.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

41pha =: 1:25 bets 2.5. 5 $.5.00 d 109
19r = 1-, pets batch
2/hh 122 07
CEAAINNEE Thary
33:
Apl
1942
10
39-1
22
10 22 10
10 22
22 120
Mops
Rops
Figure 1: For empirical runs and deterministic ODE (7) simulations of PLRE, DANA outscales SGD while Schedule-Fres SGD and Adam do not. (left) SGD & DANA fixed d, single run. Gray dashed line indicates 1) shifts from SGD--ike behavior to acceleration. Deterministic ODE predictions (bold curves) match single runs of the stochastic algorithms (faded curves). (middle) Deterministic ODEs for SGD. DANA and Schedule-free SGD [33]. across d x L,.,, 10. DANA-decaying (rig = Za outscales DANA-constant, which itself outscales SGD. Schedule-free SGD scales similarly to SGD. (right) ODEs show Adam cosine decay [56] appears to have the same scaling law as SGD, See Sec. M for details. In all figures, batch size is 1.
transition (t x 1//3- d) where DANA-constant(x2
1)
Can one select hyperparameters of stochastic momentum algorithms as a function of model size and data to provably change the exponents in the scaling laws for the loss?
In this work, we answer in the affirmative, mathematically showing that one can indeed improve the scaling law exponents over standard stochastic gradient descent (SGD) on a simple power law random features (PLRF) model.
Outscaling. To address the question of scaling, let us consider the following learning problem,
min AER
(0) = E, [S (0;z))) x)]}: where H Rd - R,
where d is the parameter count, and where is drawn from unknown distribution.
training regime, l x d', l 0, is scaling of iterations (or samples) to parameters. There are many examples of training regimes, e.g., the proportional regime (l. = d) or the compute-optimal regime, in which one selects the é that yields the best loss under a fixed compute budget (see See. 5). Now suppose the loss under an algorithm follows the scaling law
(l, del o (0,, d) x 1 + and suppose 1. = d',
is training regime.
Then under this training regime, the loss satisfies Add d) x d' (in,), We call the absolute exponent on d, the loss exponent. For given training regime, we say an algorithm outscales another algorithm if the loss exponent is larger.
We emphasize that this notion of outscaling differs in one key aspect from the more traditional notion of "acceleration" from optimization theory. Acceleration is typically formulated for a fixed-
dimensional problem with constants that can have large d-dependence (e.g.. |00 113).. In other words, acceleration generally denotes outperformance when t. + co and d = O(1)
In this work, we interested in scaling laws of one-pass, mini-batch SGD with momentum with batch size B. At iteration 2 O, we generate independent samples { and update:
= (1 - -0)) m (1; d) Ei= (
(Gen-Mom-SGD)
0:41 0: 72( dt V V (0.; ra(1; d)yr,
where A()) : [0, oc) + [0, oo) is momentum hyperparameter and (;(1; d) : [0, oo) + [0, 2)) are
learning rates. This framework incorporates classical SGD and Mom-nuu by setting
7 (t) = 1, (2 d) = 7/2 32d " (t:: t d) = % 3gd (t) K; 2 0 (SGD-M)
as well as stochastic Nesterov, Schedule- Free SGD [33], and accelerated SGD [52, 94]. For a detailed discussion related work see Section A and Tables 2 & 3.
Main Contributions. In this work, we analyze Gen-Mom-SGH under the power law random features (PLRF)- four-parameter model with data complexity (1/a),, target complexity (1/B)²
SSee Assumption 2 for formal definition of and 8 in the context of the power law random features model.

---

60
387
10
20
10
12
32 12
2033
1633
Figure 2: DANA-decaying improves the loss exponent on LSTM language models. (left) Sweeping DANA- decaying K3 shows stability and divergence similar to PLRF (Fig. 3a). (center) DANA R3 = 0.7 maximizes the
loss exponent and outscales SGD. (right) Compute-optimal loss exponents (top) and validation loss for final iterate (bottom) DANA R3. All loss exponents (Fig. 22) have R 2 0.984 and vary smoothly across Ka, traversing the divergent, outscaling, and SGD-like regimes seen in PLRF (Fig. x-axis). Schedule-Free SGD 83 1.0 matches SGD loss exponent. SGD-M matches SGD (Fig. 24). See Sec. I..
model parameter count and hidden dimensionality v extensively studied for its scaling properties [11, 69]. We derive a scaling rule for hyperparameters of Gen-Mom-SGii that improves loss exponents over SGD across much of the (a, -phase plane under a variety of training regimes. We denote this scaling hyperparameter rule diieesiion-daapter Nesterov acceleration (DANA):
71 71 d) 1, d2) 32d=". ra(L; 53d N2 (1+1)- =1
(() 8(1 +1))', (DANA)
where K; 2 0. These hyperparameters will further be made explicit for the PLRE
To show this improvement in the exponents, we show that the loss curves for PLRF can be described exactly by system of differential equations, and derive precise theoretical scaling laws for SGD-M and DANA. Using dimension- and data-dependent hyperparameters. DANA outscales SGD (see Fig. 1) above the high-dimensional line (2a 1) while performing no worse than SGD elsewhere. In contrast, traditional SGD-M, with fixed, non-scaling momentum, A, produces identical scaling laws to standard SGD across all regimes (see Fig. 24). DANA--onstae (Æ² = 0) employs decaying momentum schedule (1 +)) with the momentum learning rate given as ?3 72- 1/d, and outscales
SGD for 2a > under many regimes. DANA-decayine (N₂ = K1 0, Kg > 0) replaces with
time-dependent effective dimension, resulting in a schedule for " that outscales SGD for all regimes and all 2a > 1, succeeding where DANA-constane cannot.
We then investigate the compute-optimal regime, explicitly deriving compute-optimal parameter, loss, and data exponents. While the empirical Chinchilla laws [50] suggest compute-optimaliig occurs at d > t, we show DANA-decaying (for 2a > 1) has different compute-optimal relationship between and emphasizing that outscaling can occur outside the Chinchilla regime. Moreover, for some (o.,3) (), DANA reduces the data exponent needed to reach compute-optimanl compared to SGD.
We perform extensive PLRF experiments in Sec. K to validate theory, showing excellent numerical agreement for loss curves and compute-optimal exponents with the analyzed ODEs. Finally, we train LSTMs on text data (Fig. showing the DANA loss exponents (Fig. & 22) vary smoothly over K3 and recover the divergent, outscaling, and SGD-like regimes predicted theoretically by Fig. 6.
While show the existence of stochastic algorithms that provably outscale for 2a it remains an open question as to whether outscaling SGD can occur in the high- digh-dimensional setting (2a < 1).
2 The power law random features model (PLRF)
In this work, we analyze four carrpparaetter model called power law random features (PLRF) (3) [11, 69, 80]. which exhibits rich behavior and phenomenologically captures many aspects of scaling law setups [16. 63, 80]. For a data vector E R" we embed this vector in Rd using a matrix W E Rdudd and construct noiseless targets by dotting fixed & E R" with the sample x. This leads to the formal problem statement:
min { (0) def 1 [(WW T x, 0) (z, 1))
(3)
BERd
3

---

The matrix W allows the model to have variable capacity (d) independent of the data dimension, and we choose the matrix W have entries distributed as [0, 1/d). The key structural assumptions are:
Assumption 1 (Data and targets, a and 3). The samples a E R" are distributed according to (2z) ( j= for all 1 S S and {2g}Ø=1 ~ N0,1)). The targets scalars constructed by dotting the sample x with a signal b R whose entries (b,)= j
Power law type data distributions are ubiquitous in language, vision, and many other tasks, and these distributions are largely responsible for making this model phenomenologically similar to scaling law setups [69, Fig.2.3), Without the random matrix W. (a,, (5) are related to what is known in the literature as source and capacity conditions [22, 23, 36, 81] (see Sec. A and Tab. for details).
The hidden dimensionality u is assumed to be large and proportionate to d, so that v/d + E (1, oo). In the case that 2o 1, this assumption or even w = oc. It should be noted that in many scaling law setups, such as [50], the task scales with the parameter count, so that it is natural to assume v grows d grows.
be relaxed, in that can take larger as function of
3 Continuized analysis of general stochastic momentum algorithms
Continuized frameworks widely used for the analysis of momentum algorithms (see especially [38, 78, 92, 99]). When run on the PLRF model, the algorithm class Bee--Mom-SGD) has a loss curve that can be described exactly by system of differential equations; Section C.I and (22) for details and derivation.
We use common probabilistic trick for continuizing discrete process called Poissonization in which we let (N, 2 2 0) a standard Poisson process, and then define Y, = UN, and O = ON,- the eigevvaluu-rigevveetoo pairs for K E Rdxd which is the covariance of the projected data W x, and let @* be the minimizer of (O)). We then introduce the system of variables, with the expectations taken over all randomness except W.
Now we let (Aj, و
(t) def E [(w,, θ ),)] (t) def [ww, ). and (t) def E [w 9, 0*) Y.)
(4)
In terms of these variables, we recover the expected loss by summing over the components of (t).
P(() E| [[ Σラ \npz(1).
(5)
Then we can derive a coupled linear system of differential equations:
A1: def (),(),
((; = 224, A3) x (1; P()g((())))
(6)
for explicit matrix Ω and vector g which are polynomials in (Aj, A, B, 71, 72, 7a) (see (22) for the full formula).
These equations are sufficiently complicated that we simplify the ODE before performing the analysis. This can be viewed as taking a "high- nigh-dimensiooal limit" (see Remark C.3), or, in short, dropping all terms from Ω and y which are more than first order in (A, (a) to produce:
ltt: X1) def 272
0
-2aa(() 27 (I)
(1)
-2((1)
and g def AB
(t(
(1)BA)
-1a(()
A(()) BA)
We formulate all results going forward for the simplified ODEs; see Sec. K and Fig. 8 for extensive numerical validation of the simplified ODEs as a model for the risk curves of Gen-Mom-SGDD We also note that these simplified ODEs correspond exactly to the risk evolution under the SDE model of Gen-Mom-SGD:
dY, - (00 X0X (K
(7)
dθ,- - BK (K) 7o(()Y dt.
'In the 2a < 1, there larger range of scalings of e that possible, with < 21122a)

---

Ow 17.5-11. Fill
Jyn3 1/ lets - - /n2 995
Ferry
argical
12"*
Daps
elgenvaluss
(a) (1) x (1 + f)
(b) (t) d' d22
(c) K4g versus JAy
Figure 3: (left) (center) DANA sweeps of ni, K3. A bigger than (stability) + divergence after some t; DANA-decaying w/ Ri = 1/(2)) the envelope the divergent algorithms. Ki = 1/(2a), R2 = is optimal in DANA. (right) Deterministic equivalent jer (8) vs empirical density estimate match. See Sec. M for details.
are independent standard d-dimensional Brownian motions. ? Using Itô's formula (see (43) for details), one then can easily derive the simplified ODEs for the system v(t; X²) def
Here d'B," and d (2)
4 Deriving the scaling laws of (Gen-Mom-SGD) on the PLRF model
Random matrix theory of the PLRE By solving these ODEs for terms their initial data, we represent the curve P entirely as function of the spectral data (Ag, (0)) A mathematically
convenient representation for this data is a pair of pure-point measures
To (dz) = = 5x, (dz)pp
(0) and (xx(dz) = Ein 8x, (dz)
(8)
terms of which P can be represented as a solution of an integral equation (see (36), (56), (57)).
Random matrix theory (RMT) gives prediction for these measures ( using the so-calla so-called deterministic equivalent; for details, see (33). Deriving such an equivalent a textbook tool in RMT. but the proof of equivalence in the case of PLRF falls outside of textbook RMT (see for example [62]). To limit the scope of the theoretical analysis, all the scaling law statements that follow are proven for the deterministic equivalent; the numerical agreement between jay and juy is excellent (see Fig. 3c).
Scaling laws for SGD on the PLRE In [80], using the deterministic equivalent (Jrr, pac), the authors show for SGD,
(deterministic loss curve) P(t)= x (((()) + T2 where F(t) 5 (t) + Fac (t) + Fo.
where (()) def 1 + 222 Bl. Here 3 9 can be viewed as the bias term and Kpp o I as the variance due to the stochastic gradients. Each of these terms has interpretation and explicit scaling law:
Component
Symbol
Scaling Law
Contributing Part of μ
Population bias
pt)(t)
x t- (20-28-1)/(20)
Spikes (pure point part) in prr
Model capacity
Fo(()
x d (2a+((11-))))
Irreducible loss level, (r({)))
Embedding bias
Fac(t)
x d dd11-1111(20)
Bulk (absolutely continuous part) of μe x d1 1-241/(200) Pure point part of μo
Variance
Ktp(t)
4 Distinct Phases. For different regions of (a, B). one or more of these four terms may be fully dominated by the others, yielding 4 different phases. See also Fig. 5 as an example. The different components of the loss curves do not change across the algorithms, but the scaling exponents (α,7) vary. The phase boundaries determined by major quantitative and qualitative changes in the problem geometry.
We have also used two Brownian motions for additional simplification of the analysis; this would correspond using independent stochastic gradient estimates in the two lines of (Gen-Mom-SSO) SGD).
5

---

Dominant terms
Phase I:
Ppp + Fo
Phase II:
Fpp(t) + Foe + 3o
Phase III:
(tp ato Fo
high-d
Phase IV:
Fpp(() + ppp(1) + +
IV
0.25
Each phase is characterized by differrnn dominant terms in the loss function (t).. The terms pp (t). pkp((t) Fee(t) and Fg represent different components that contribute to the overall loss behavior, with their relative importance varying across the parameter space.
0.5
Figure 4: Phase diagram and corresponding dominant loss terms.
The high-dimensional line, which occurs where 2a 1, separates gradient norms that grow with d from constant ones. The line 23 = 1 determines if the |10g 0 1|² grows with d. The line a = B determines if the target complexity or data complexity higher, which determines the relevance of gradient noise. Finally a = 1 determines where the gradient noise becomes high-dimensional and has aspects which are not power law.
Scaling laws for momentum methods. Our main result is the extension of these scaling laws SGD-M, DANA-constant and DANA-decaying in the small batch setting.
Theorem 4.1. (Summarized Version of Thum G.I (SGD-M), Thm H.3 (DANA-constant) Thm 1.2, (DANA-decaying)). Suppose that (a, (3 are not on any critical line, that 2a + 2B > 1, that a > 1' and that 3 S a Suppose that batch size B is fixed independent of d. Consider DANA-constant, DANA-decaying, SGD, and SGD-M. Define the time change
1 + 2(72 + Bt, BL + (Io (T 1 1 Bl
(SGD-M).
0(t) def
(DANA).
Then, provided the algorithm stable which is to say that for some constant sufficiently small, independent of d (and where (r ビグ= 2a x d(1=20)+)
< 2, (12 + m I emin { B3 1r 1),
(SGD-M),
c5 > 1,72 S cmin B 17 1: 7xd=== < C22
(DANA) with 2
2a (stability)
1
cb > 1,72 cmin B u. 3g X2 < 22nd2alkss 20,
(DANA) with K3
one has the following scaling law where & = for DANA and y = 72 2 for SGD-M
p(t) = ((((t)) (0xm))
(10)
For DANA- decayidc,,, we have the following additional requirements. We only consider 2a > 1, bounded below, and (1/2a) < 63 <
72
We remark that DANA-decaying/conse
appear to be the most interesting (and optimal) DANA cases (see Fig. 6). From (10), it is simple to produce explicit scaling laws for DANA
P(t; d) x d=
T +d-
(11)
population bias
model capucity
embedding bias
variance
see for example Figure 5. The parameters as and r; are explicit, depend on the algorithm, and vary continuously in (a, (). We relegate the proofs to Sec. G. H. and specific exponents to Table 4.
Understanding stability. The stability conditions listed in Theorem 4.1 are essentially sharp. that we can show the loss curves are unbounded in d if the hyperparameters have lower bounds that are on the same order as these upper bounds.
Stabiliity is proved in Cor. G.1 for SGD-M, Cor. H.2 for DANA-constant,, Prop. 1.2 for DANA-decaying

---

SCIDSCIL-M
Figure details: Here def 20-+3-11. Population bias, pp ((()) dominates at small times, then the loss slows down from embed- ding bias, Fhe (f)))), and then finally reaches the model capacity, Ft(t). In this phase, the gradient noise (variance, Kpp d) is always dominated by Fyp 0, unlike Phases III/TV. When t d. DANA- constant behaves like SGDISGD M. For all phases, see Fig. 7, 12 for phase diagrams and Fig. 13,14,15 15 for loss curves.
DAM-cmmdnn: DANAvkssyngt-
loy (Prrd)
pp
log()
Figure 5: Example of loss curves (Phase II,
B > 0.5).
In the case of SGD-M, above the high- digensinnnio line (2a > 1), the effective learning rate 72 + 's can remain constant, but below it must scale with the parameter count d. This transition is precisely captured by tr 5
For the DANA class, the simplest choice of n3 is to take it constant, (i.e, K3 0, DANA- DANA-constant). The largest constant stable choice for 7/ d'1 (see also a related algorithm [94]). This has key limitation: with such small ?3. the momentum term needs ((d) steps to reach the gradient magnitude and to have any effect on the behavior of the loss. Thus, for early iterations, S d. and SGD exhibit similar scaling behavior.
One can resolve this, and produce faster algorithm by observing that after iterations, only part of the d-dimensional feature space is being used, i.e. there is effective dimension of K at iteration t [14, 61, 101, 110]. Quantitatively,
(effective dimension at iteration = max Ay(K) 1/(LB)) = (1/))
If we replace the fixed dimension parameter d in DANA-constant with (IB)--1(22)) this yields the DANA-decaying algorithm (see also [107]) with the smallest possible Æ3- Fig. shows DANA- decaying with 1/(2a) appears to be optimal on PLRF when varying (H2, K3); see also Fig. 6.
The choice of 8 in A(t) 8(1 (11+t)- must be chosen sufficiently large to guarantee acceleration and stability. We summarize the recommended DANA-constan and DANA-dccaying hyperparameters
for PLRF in the table below, where ir = -E 20.
Algorithm
Hyperparameters ($
> (2 - K3) max 20 .4 11)
DANA-constant
2a
K1, Ki K2 = 1,72 1/(2tr), 733 1/5 72
2a
K1 = 20, 62 = 1 K1, K3 0
DANA-dccaying
2a >
K2 0, N3 = 1/(2a), 7g
Hyperparameters beyond PLRE In settings where (a, () are unknown or meaningless, we recom-
mend setting 7/2 to stable SGD learning rate, (t(t) = 72 x (1 f) + and 5 = 8. Then, sweep over N3; the minimum K3 where DANA DANA-decaying converges appears optimal (Fig. 2c & 20).
2a
Full DANA class. There are other stable scaling rules in the DANA class (above red line). DANA MANA-iec w/ K3 2a is optimal (e.g.. fewest iterations to optimum, best loss exponent) amongst the whole class. (left, above high-d line, иj = 0, B 1, 1.1): Plot loga to reach irreducible loss). DANA stability boundary (red) at
2.1
2.0-
soundary boundary
2.5
1.5-
2
La
1.0
DoNikconstan:
1.
22nx3 1 with divergence below; DANA takes same number of itcrations as SGD at (pink) line, X2 2 2a(1 -Ks). Darker color smaller number of steps. Iterations to reach irreducible loss, dla? /(4a-1) < d'+1/2 d (DANA-decay < DANA-constame < SGD). Stochastic Nesterov does not converge [38].
0.5
1.2
Schaedule free Scap
8.0
0.2
0.4
0.8
1.0
Ks
Figure 6: Full DANA class, time to reach irreducible loss.
7

---

Pie-se lo: 2.bth 2.4. beto - 2.5. 22.2
Phase lla: alph O. beda 100 +3.5
Plase (b: a/pha
Duce U.A, 22.22
dele ore
denddeeee ode 101, one
10-4
30-1
716
23
22-
10-10
10 ShiVe Care
31"0
195 cure
33
comp deci
107
10
10
10
10
12
12
12
10
22
25
flpss
flops
Maps
Phose lim: Abbha 1.4, 3., 5 5-5.57
Phose 1b; slonp
duay de
10
caracett
Outscales
spd
13
la
0.75
1lb
11th
high-d
line
10:19
lb
1aa
2-v v
10-12
10
18-A00ing OUF parg.a CFR
R
0.25
12
6p6 di Voria CaTc
12
30"
Ic
Same as SGD
Raps
tops
0.5
/
Figure 7: Comparison of SGD, DANA-constant, and Deca-dgg with compute-optimal curve predictions. Numerical set-up: d 100 x 2', " I,..., 15; Simplified ODEs (43) plotted for the scaling-lane equivalent model. DANA- -ecayynng outscales in all phases (Phase Ia. II III where 2a > 1. DANA-constan
outscales SGD in all phases 2x > 1 except Phase IIIb. Predictions for compute-optimality loss exponents match empirical results. (see Sec. M for details).
Using Momentum to outscale SGD
SGD-M fails to outscale. The typical approach to momentum is to use a constant momentum
(t) = 8, in practice usually set to 1 ó = 0.9. Regardless of the choice of fixed A, SGD-M produces the same scaling laws as SGD. In particular, most 'speed up' of SGD-M can be attributed to larger effective learning rate for SGD, resulting in improved scaling law constants, but in the loss exponent. There is an equivalence in risk dynamics when 7 72
a change SGD-M + > SGD-M /5; see
Rem. G.2 for details and Fig. 16 and 24 for empirical equivalence on PLRF and LSTMs respectively.
DANA-constant outscales SGD for most (a., 3) with 2a > 1. The limitation of DANA-constant's small i's learning rate is that the momentum term requires at least d iterations of the algorithm to become noticable in the loss dynamics. Thus, for S d, DANA-constan and SGD have the same scaling law (see Fig. 1a where gray line indicates > d). If one is able to reach Fo before d (which occurs for 2a > 1), DDAN--consaal will outscale SGD in training regime = d' with 1 < l < 2o. We note that once e = 2a, no algorithm outscales SGD, as the irreducible loss level is reached.
DANA-decaying outscales DANA-constait for all (a., 3) with 2a 1. For any regime t = d with < < 2a, DANA-decaying will outscale both SGD and DANA-constant, provided the former is not already at the irreducible loss level Fo. Hence DANA- will always be both more sample efficient and compute efficient than SGD and DANA-constant, One may look further at the full DANA class of algorithms (Fig. 6), which may potentially contain other interesting scaling rules.
Schedule-Free SGD [33], Nesterov [75]. AcSGD [94]. Accelerated SGD [52. 61] (see Tab. 3 for param. comparison). For these, deterministic ODEs (22) exactly describe the risk evolution. We give heuristics for the scaling laws of these algorithms (Sec. B.5). See Tab. 2 for sample complexity comparison and Fig, 10 for scaling behavior, in which DANA-decaying outperforms all (including Adam with cosine decay, see Fig. Ic).
Moreover, many stochastic algorithms fall into Gen-Mom-
In principle, one could also look at unstable scaling rules, which are stable in one training regime but eventually diverge. For example, one could sweep over the 73 in DANA-constant to find the 73 that minimizes P(I d); this produces schedule "aa(1; d), which on the PLRF is exactly DANA-decaying with K3 1/22)) (see Fig. 3b & Section J.3). See also Fig. 3a on PLRF and Fig. 2a on LSTM for related sweep over K3.
8

---

C4WMooecart to 122
23
MCD Sppatt
0.32
git pirsacally (xale))
23
13-4
22
Orpator HECE
0.59
oypaki 244
23"
35 10 325 850
0.35
1.66
125
136
3.5
2.30
3. 13
22
13
Figure 8: Loss curves match almost exactly between empirical PLRF experiments and ODE solutions. Compute-optimal loss exponents match approximately, (left) Loss curves for simplified ODEs (7) closely match empirical PLRF across 13 d-values up to d = 12,800. Empirical loss exponent fit by Chinchilla Approach 1 [50]. (middle) Loss exponents 1) in Pf/d""; dd) f for fixed و = 0.7 across a's. Theoretical predictions (solid lines, Tab. 5,6.7) match empirical (dots) within 0.09. (right) Compute-optinal parameter exponents { d' (f) = i match within 0.13 (top) and data exponents. Ç (bottom), for fixed a 1 across values. See extensive experiments Sec. across 84 values of (at, + discussion error sources. Fig. details in Sec. M.
Compute-optimal regime. One natural training regime for d is the compule-optimal training regime [16, 50, 63]. where for a given compute budget f measured in flops, one chooses d' (f) to minimize the loss. We measure f via:
Compute (flops f) (iterations of alg. (t) x batch size (B)) x parameters (d). (12)
We plot the loss curve (0,; ((t; ((: as a function of flops, and then we solve for the compute-optimal parameter size d' (f) = arg m = fs. With access to the explicit functional form of the loss curve (11), it is straightforward find d= We denote & the parameter exponent and v the loss exponent. Here 9 dz) = f- is known as the compute-optimal curve. The data exponent is 1 - since iterations times batch equals amount of data used.
Given the form of the loss curves (11), compute- compute-optimallyy must occur at intersection point of T Facr To. and Kyp Thus. the phases are further broken down depending on the tradeoff location We derive compute-optimal frontiers and present the exponents in Tab. 5,6,7, Sec. E for proof details. See also Fig. 12 (phase diagrams) and Fig. 13,14,15 cartoon plots of loss curves and compute-opimality. Fig. 9 gives a detailed example of the shape of the risk curves. Predictions of exponents (C. n,.) {) match theory- see Fig. 8, Fig. 26-33. For specific details about the different phases and algorithms, see Sec. E.5.
of compute-optimalie
Compute-optimality main takeaways. The empirical Chinchilla law [50] showed the optimal parameter count scales like d= (f) x f1/2 large language models while [80] observed theoretically that SGD on PLRF reproduces this behavior in Phase III. However, for 2a > 1. DANA-ddcaylle is never compute-optiman with d' (f) x fl/2; depending on the phase, the compute- optimal training regime of DANA-decaying requires undertraining or overtraining relative to the Chinchilla law.
Furthermore, the compute-optimal tradeoff point can change based on the choice of algorithm. For exam- ple, Phase Ial/IIla for SGD it is compute-optimal to stop training early, when the optimization enters the Fhe regime (which is slower to optimize). However, for DANA-decaying, it is compute-optimal to continue train- ing, and in fact, to continue training to the irreducible loss level F, because MANA-decaying is substantially more sample/comppte--ffii efficcieteeffintttteeeeee than SGD in this regime. Moreover, even when the tradeoff point is the same as SGD, such as in Phase la, DANA- -ANA-iecaying attains better loss than SGD for the same compute or sample budget.
SGDISGD (slope):
DANA-cosstaid (slope).
DANA-docyynng
log PIr
F
Compute Optimal
0.5
Figure 9; Compute-optimal scaling laws (Phase Ia, (a > (3 + /5)/4, > B 0.5)). We reparameterize Fig. 5 by fixing the compute budget and parameterizing the x-axis with E [0, 1), so that d = pi- and t= f.. This shows the relationship between compute bud- get and optimal parameter count.
LSTM results. We train 2-layer LSTMs (Fig. 2) on the C4 language dataset [84] and co-scale the embedding and hidden dimensions to sweep model sizes. This in-

---

duces a power-law regime for the compute-optimal fron- tier on intermediate model sizes with all R² 2 0.984, similar to Fig 7 in [54]. For DANA-decaying, the loss exponents (Fig. 22) vary smoothly with Kg and closely match PLRF behavior (Fig. 6): diverging for Ki < 0.6, outscaling SGD for 0.6 S K3 I 0.9, and matching SGD when Ki 2 1.0 where ri 1.0 is Schedule-Free SGD. We show equivalence for SGD-M and SGD in Fig. 24. For details see Sec.
Conclusion, limitations and future work. We have shown that DANA outscales SGD on PLRF in the 2a > 1 regime by properly scaling Nesterov-type momentum in data- and simension-dependen way. We validated Theorem 4.1 with extensive PLRF experiments; moreover, PLRF acts as useful proxy for LSTMs on text data where theoretical predictions for loss exponent improvements hold. A full convergence argument beyond quadratics for DANA-decay would be desirable, e.g. theory for cross-entropy might suggest different momentum scaling strategies. It is an open question if outscaling is possible for either the 2a < 1 case or for d-dependent batch sizes. Additional limitations of our analysis include the use of fixed features which does not incorporate any kind of feature learning, well as deterministic learning dynamics based on ODE description of the loss evolution.
Additionally, while we compare DANA against other non-adaptive momentum methods, most real- world problems benefit from adaptive methods such as Adam. Hence it would be interesting to have an analysis of DANA combined with Adam or other preconditioned methods. Finally, in the LSTM setting, the exact meaning of a is not clear, although empirically a 0.71 (corresponding to K3 = 0.7) appears near optimal. Defining and measuring this on real-world problems, particularly determining whether 2a 1, is an important direction for future work.
Acknowledgments and Disclosure of Funding
D. Ferbach is supported by Fonds de Recherche du Québec (FRQ) (DOI assigned: G. Gidel is CIFAR AI Chair, he is supported by a Dis- covery Grant from the Natural Science and Engineering Research Council (NSERC) of Canada and a Google x Mila research grant. C. Paquette is Canadian Institute for Advanced Research (CIFAR) chair, Quebec AI Institute (Mila) and a Sloan Research Fellow in Computer Science (2024). C. Paquette was supported by Discovery Grant from the Natural Science and Engineering Research Council (NSERC) of Canada, NSERC CREATE grant Interdisciplinary Math and Artificial Intelligence Program SETER-MATH-AJ) Google x Mila research grant, Fonds de recherche du Quebec Nature et technologies (FRQNT) NOVA Grant, and CIFAR AI Catalyst Grant. Research by E. Paquette was supported by Discovery Grant from the Natural Science and Engineering Research Council (NSERC). This research was enabled in part by compute resources provided by Mila (mila.quebec)) Additional revenues related to this work: C. Paquette has 20% part-time employment at Google DeepMind.
The authors would like to thank Jeffrey Pennington, Lechao Xiao, and Atish Agarwala for their careful reading and helpful feedback that improved the paper.
10

---

References
[1] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfel- low, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg. Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, llya Sutskever, Ku- nal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg. Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow Datasets, a collection of ready-to-use datasets. https: //ww:/ tensorflow org/datasets.
[2] Zeyuan Allen-Zhu. Katyusha: The first direct acceleration of stochastic gradient methods. The Journal of Machine Learning Research, 18(1):8194 8244, 2017.
[3] Zeyuan Allen-Zhu and Lorenzo Orecchia. Linear coupling: An ultimate unification of gradient and mirror descent. arXiv preprint
2014.
[4] Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high- dimensional mean-field dynamics to dimensionless odes: A unifying approach to sgd in two-layers networks. In The Thirty Sixth Annual Conference on Learning Theory (COLT), pages 1199-1227. PMLR, 2023.
[5] Søren Asmussen, Soren Asmussen, and Sren Asmussen. Applied probability and queues. volume 2. Springer, 2003.
[6] Mahmoud Assran and Michael Rabbat. On the Convergence of Nesterov's Accelerated Gradient Method in Stochastic Settings. In Proceedings of the 37th International Conference on Machine Learning (ICML), 2020.
[7] Alexander Atanasov, Blake Bordelon, Jacob A Zavatone- Veth,, Courtney Paquette, and Cengiz Pehlevan. Two-Point Deterministic Equivalence for Stochastic Gradient Dynamics in Linear Models. arXiv preprint urXiv:2502.05074 2025.
[8] Krishna Athreya, Peter Ney, and PE Ney. Branching processes. Courier Corporation, 2004.
[9] Needet Serhat Aybat, Alireza Fallah, Mert Gürrüzzallaaan, and Asuman Ozdaglar. Robust Accelerated Gradient Methods for Smooth Strongly Convex Functions. SIAM Journal on Optimization, 30(1):717-751, 2020.
[10] Francis Bach, High-dimensional analysis of double descent for linear regression with random projections. SIAM Journal on Mathematics of Data Science, 6(1):26-50, 2024.
[11] Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. Proc. Natl. Acad. Sci. USA, 121(27):Paper No. e23 1878121, 8, 2024.
[12] Zhidong Bai and Jack W Silverstein. Spectral analysis of large dimensional random matrices, volume 20. Springer, 2010.
[13] Prosanta Barai, Gondy Leroy, Prakash Bisht, Joshua M Rothman, Sumi Lee, Jennifer Andrews, Sydney A Rice, and Arif Ahmed. Crowdsourcing with enhanced data quality assurance: An efficient approach mitigate resouree scarcity challenges in training large language models for healthcare. AMIA Summits on Translational Science Proceedings, 2024:75, 2024.
[14] Peter I.. Bartlett, Philip M. Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proc. Natl. Acad. Sci. USA, 117(48):3003333007 2020.
[15] Raghu Bollapragada, Tyler Chen, and Rachel Ward. On the fast convergence of minibatch heavy ball momentum. IMA Journal of Numerical Analysis, page drae033, 08 2024.
[16] Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. A Dynamical Model of Neural Scaling Laws. In Proceedings of the 41st International Conference on Machine Learning (ICML), volume 235 of Proceedings of Machine Learning Research, pages 4345 4382. PMLR, 2024.
[17] Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. How Feature Learning Can Improve Neural Scaling Laws International Conference on Learning Representations (ICLR), 2025.
11

---

[18] J. Bradbury, R. Frostig, P. Hawkins, M. Johnson, C. Leary, D. Maclaurin, and S. Wanderman- Milne. JAX: composable transformations of Python+NumPy programs. 2018.
[19] Helmut Brakhage. On ill-posed problems and the method of conjugate gradients. In Inverse and Ill-Posed Problems, pages 165-175. Academic Press, 1987.
[20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-
1901, 2020.
[21] Bugra Can, Mert Gurbuzbalaban, and Lingjiong Zhu. Accelerated Linear Convergence of Stochastic Momentum Methods in Wasserstein Distances. In Proceedings of the 36th Interna- tional Conference on Machine Learning (ICML), volume 97, pages 891-901. Proceedings of Machine Learning Research (PMLR), 2019.
[22] Andrea Caponnetto and Ernesto De Vito. Optimal rates for the regularized least-squares algorithm. Foundations of Computational Mathematics, 7:331-368, 2007.
[23] Luigi Carratino, Alessandro Rudi, and Lorenzo Rosasco. Learning with sgd and random features. Advances in Neural Information Processing Systems (NeurIPS), 31, 2018.
[24] Chen Cheng and Andrea Montanari. Dimension free ridge regression. arXiv preprint
arXiv:2210.0857// 2022.
[25] Earl A. Coddington and Norman Levinson. Theory of ordinary differential equations. McGraw- Hill Book Company, Inc., New York-Toronto-London, 1955.
[26] Elizabeth Wooodin, Courtney Paquette, Elliot Paquette, and Inbar Seroussi. Hitting the high- digeeniimenii notes: an ODE for SGD learning dynamics on GLMs and multi-index models. Inf. Inference. 13(4)) No. iaac028, 107, 2024.
[27] Elizabeth Collins-Woodfig Inbar Seroussi, Begoña García Malaxechebarria, Andrew W Mackenzie, Elliot Paquette, and Courtney Paquette. The High Line: Exact Risk and Learning Rate Curves of Stochastic Adaptive Learning Rate Algorithms. In Advances in Neural Information Processing Systems (NeurIPS), volume 37, 2024.
[28] Romain Couillet and Zhenyu Liao. Random Matrix Methods for Machine Learning. Cambridge University Press, 2022.
[29] Hugo Cui. Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Generalization Error Rates in Kernel Regression: The Crossover from the Noiseless to Noisy Regime. Advances in Neural Information Processing Systems (NeurIPS), 34, 2021.
[30] Anh Quang Dang, Reza Babanezhad Harikandeh, and Sharan Vaswani. (Accelerated) Noise- adaptive Stochastic Heavy- Momentum. Transactions on Machine Learning Research (TMIR), 2025.
[31] Stéphane d' Ascoli, Marylou Gabrié, Levent Sagun, and Giulio Biroli. On the interplay between data structure and loss function in classification problems. Advances in Neural Information Processing Systems (NeurIPS), 34:8506-883 8517, 2021.
[32] DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Pe- ter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, lurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papa- makarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojeiech Stokowiee, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem. deepmiind, 2020.
[33] Aaron Defazio, Xingyu Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, and Ashok Cutkosky. The Road Less Scheduled. In Advances in Neural Information Processing Systems (NeurIPS), volume 37, pages 9974-10007, 2024.
[34] Leonardo Defilippis, Bruno Loureiro, and Theodor Misiakiewicz. Dimension-fred deterministic equivalents and scaling laws for random feature regression. In Advances in Neural Information Processing Systems (NeurIPS), volume 37, pages 104630-104693, Curran Associates, Inc.,
2024.
12

---

[35] Ofer Dekel, Ran Baad-Bac Ohad Shamir, and Lin Xiao. Optimal Distributed Online Prediction Using Sani-Bachhes. Journal of Machine Learning Research (JMLR), 13(1), 2012.
[36] Aymerie Dieuleveut and Francis Bach. Nonparametric stochastic approximation with large step-sizes. The Annals of Statistics, 44(4):1363 1399, 2016.
[37] NIST Digital Library of Mathematical Functions. httpp::/
/tlmf/dlmff//../////////////d/ /ttpp://dmmf of 2025-03-15. F. W. J. Olver, A. B. Olde Daalhuis, D. W. Lozier, B. I. Schneider, R. F. Boisvert, C. W. Clark, B. R. Miller, B. V. Saunders, H. S. Cohl, and M. A. McClain, eds.
mwst.tstt...... .oosttgooo. Release 1.2.4
[38] Mathieu Even, Raphaël Berthier, Francis Bach, Nicolas Flammarion, Hadrien Hendrikx, Pierre Gaillard, Laurent Massoulié, and Adrien Taylor. Continuized Accelerations of Deterministic and Stochastic Gradient Descents, and of Gossip Algorithms. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 28054-280666 2021.
[39] Katie E Everett, Lechao Xiao, Mitchell Wortsman, Alexander A Alemi, Roman Novak, Peter J Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jachoon Lee, et al. Scaling Exponents Across Parameterizations and Optimizers. In International Conference on Machine Learning. pages 17666611770 PMLR, 2024.
[40] Nicolas Flammarion and Francis Bach. From Averaging to Acceleration, There is Only Step-size. In Proceedings of The 28th Conference on Learning Theory (COLT), volume 40 of Proceedings of Machine Learning Research, pages 658-693 PMLR, 03-06 Jul 2015.
[41] Sèbastien Gadat, Fabien Panloup, and Sofiane Saadane. Stochastic Heavy Ball. arXiv preprints (eti://6000022255 2016.
[42] Cedric Gerbelot, Emanuele Troiani, Francesca Mignacco, Florent Krzakala, and Lenka Zde- borova. Rigorous dynamical mean-field theory for stochastic gradient descent methods. SIAM Journal on Mathematics of Data Science, 6(2):400-427, 2024.
[43] Saeed Ghadimi and Guanghui Lan. Optimal stochastic approximation algorithms for strongly convex stochastic composite optimization I: A generic algorithmic framework. SIAM J. Optim., 22(4): (144):14691492,, 2012.
[44] Saced Ghadimi and Guanghui Lan. Optimal stochastic approximation algorithms for strongly convex stochastic composite optimization, II: Shrinking procedures and optimal algorithms. SIAM J. Optim., 23(4):2061-2088 2013.
[45] Gustaf Gripenberg. On the resolvents of nonconvolution Volterra kernels. Funkcial. Ekvac.. 23(:):83-95, 1980.
[46] Kanan Gupta, Jonathan W. Siegel, and Stephan Wojtowytsch. Nesterov acceleration despite very noisy gradients. In Advances in Neural Information Processing Systems (NeurIPS), volume 37, pages 20694-20744, 2024.
[47] Walid Hachem, Philippe Loubaton, and Jamal Najim. Deterministic equivalents for certain functionals of large random matrices. Ann. Appl. Probab., (((3):875-930, 2007.
[48] Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically. arXiv preprint arXiv:1712.004000 2017.
[49] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9():1735-11880 1997.
[50] Jordan Hoffmaan,, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack Rae, and Laurent Sifre. An empirical analysis of compute-optimal large language model training. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, 2022.
[51] Hong Hu, Yue M. Lu, and Theodor Misiakiewicz. Asymptotics of Random Feature Regression Beyond the Linear Scaling Regime. arXiv preprint arXiv:2403.08
2024.
[52] Prateek Jain, Sham Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Acceler- ating Stochastic Gradient Descent for Least Squares Regression. In Proceedings of the 31 Conference On Learning Theory (COLT), volume 75, pages 545-604, 2018.
13

---

[53] Rafal Jozefowicz, Oriol Vinyals, Mike Schuster, Noam Shazeer, and Yonghui Wu. Exploring the limits of language modeling. arXiv preprint 1XX.::6602.021000 2016.
[54] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alex Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint vrXIv:2001.08311 2020.
[55] Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham Kakade. On the Insufficiency of Existing Momentum Schemes for Stochastic Optimization. In 2018 Information Theory and Applications Workshop (ITA), pages 1-9, 2018.
[56] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations (ICLR), 2015.
[57] Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint 1808.0880..622 2018.
[58] Andrei Kulunchakov and Julien Mairal. A Generic Acceleration Framework for Stochastic Composite Optimization. In Advances Neural Information Processing Systems (NeurIPS), volume 32, 2019.
[59] Maxime Laborde and Adam Oberman. A Lyapunov analysis for accelerated gradient methods: from deterministic to stochastic case. In Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics (AISTATS), volume 108, pages 602-612. Proceedings of Machine Learning Research (PMLR), 2020.
[60] Kiwon Lee, Andrew Cheng, Elliot Paquette, and Courtney Paquette. Trajectory of mini-batch momentum: batch size saturation and convergence in high dimensions. Advances in Neural Information Processing Systems (NeurIPS), 35:36944-36957, 2022.
[61] Xuheng Li, Yihe Deng, Jingfeng Wu, Dongruo Zhou, and Quanquan Gu. Risk bounds of accelerated SGD for overparameterized linear regression. arXiv preprint 1rXiv:2311.1422. 2023.
[62] Zhenyu Liao, Romain Couillet, and Mahoney Mahoney. A Random Matrix Analysis of Random Fourier Features: Beyond the Gaussian Kernel, a Precise Phase Transition, and the Corresponding Double Descent. arXiv preprint arXiv:2006 05013, 2020.
[63] Licong Lin, Jingfeng Wu, Sham M. Kakade, Peter L.. Bartlett, and Jason D. Lee. Scaling Laws in Linear Regression: Compute, Parameters, and Data. In Advances in Neural Information Processing Systems (NeurlPS), volume 37, pages
Curran Associates, Inc., 2024.
[64] Chaoyue Liu and Mikhail Belkin. Accelerating SGD with momentum for ever-parameterized learning. In Proceedings of the 37th International Conference on Machine Learning (ICML., 2020.
[65] Peter J. Liu, Roman Novak, Jaehoon Lee, Mitchell Wortsman, Lechao Xiao, Katic Everctt, Alexander A. Alemi, Mark Kurzeja, Pierre Marcenac, Izzeddin Gur, Simon Kornblith, Kelvin Xu, Gamaleldin Elsayed, Ian Fischer, Jeffrey Pennington, Ben Adlam, and Jascha-Sohl Dickstein. Nanodo: A minimal transformer decoder-only language model implementation in JAX., 2024.
[66] Yuanshi Liu, Haihan Zhang, Qian Chen, and Cong Fang. Optimal Algorithms in Linear Regression under Covariate Shift: On the Importance of Precondition. arXiv preprint arXiv:2502.09047, 2025.
[67] Nicolas Loizou and Peter Richtárik. Momentum and stochastic momentum for stochastic gradient, Newton, proximal point and subspace descent methods. Comput. Optim. Appl., 77(3):653-710, 2020.
[68] Brumo Lourciro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Mare Mezard, and Lenka Zdeborová. Learning eurves of generic features maps for realistic datasets with a teacher-student model. Advances in Neural Information Processing Systems (NeurIPS),
34:18137-18151, 2021.
[69] Alexander Maloney, Daniel A. Roberts, and James Sully. A Solvable Model of Neural Scaling Laws. arXiv preprint trXiv:2210.16852 2024.
[70] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of batch training, arXiv preprint arXiv: 2018.
14

---

[71] Song Mei and Andrea Montanari. Precise asymptotics and the double descent curve. Communicationns on Pure and Applied Mathematics, 75(4):667-766, 2022.
The generalization error of random features regression:
[72] Gabriel Mel and Jeffrey Pennington. Anisotropic random feature regression in high dimensions. In International Conference on Learning Representations, 2021.
[73] Depen Morwani, Nikhil Vyas, Hanlin Zhang, and Sham Kakade. Connections between Schedule-Free Optimizers, AdEMAMix, and Accelerated SGD Variants. arXiv preprint arXiv:2502.0244 2025.
[74] Yoonsoo Nam, Nayara Fonseca, Seok Hyeong Lee, Chris Mingard, and Ard A. Louis. An exactly solvable model for emergence and scaling laws in the multitask sparse parity problem. In Advances in Neural Information Processing Systems (NeurIPS), volume 37, pages 39632- 39693, 2024.
[75] Yuri Nesterov. Introductory lectures on convex optimizaaiio. Springer, 2004.
[76] Antonio Orvieto, Jonas Kohler, and Aurelien Lucchi. The Role of Memory in Stochastic Optimization. In Proceedings of The 35th Uncertainty in Artificial Intelligence Conference, volume 115 of Proceedings of Machine Learning Research, pages 356-366, 2020.
[77] Courtney Paquette, Kiwon Lee, Fabian Pedregosa, and Elliot Paquette. SGD in the Large: Average-case Analysis, Asymptotics, and Stepsize Criticality. Proceedings of Thirty Fourth Conference on Learning Theory (COLT), volume 134, pages
2021.
[78] Courtney Paquette and Elliot Paquette. Dynamics of stochastic momentum methods on large- scale, quadratic models. Advances in Neural Information Processing Systems (NeurIPS), 34:9229-9240, 2021.
[79] Courtney Paquette, Elliot Paquette, Ben Adlam, and Jeffrey Pennington. Homogenization of SGD in high-dimensions: exact dynamics and generalization properties. Mathematical Programming, pages 1-90, 2024.
[80] Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington. 4+3 Phases of Compute-Optimal Neural Scaling Laws. In Advances in Neural Information Processing Systems (NeurIPS), volume 37, 2024.
[81] Loucas Pillaud- Vivien, Alessandro Rudi, and Francis Bach. Statistical optimality of stochastic gradient descent on hard learning problems through multiple passes. Advances in Neural Information Processing Systems (NeurIPS), 31, 2018
[82] B.T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 04, 1964.
[83] Tomer Porian, Mitchell Wortsman, Jenia Jitsev, Ludwig Schmidt, and Yair Carmon. Resolving Discrepancies in Compute-Opting Scaling of Language Models. In Advances in Neural Information Processing Systems (NewrIPS), volume 37, pages 100535-100570, 2024.
[84] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
[85] Alessandro Rudi and Lorenzo Rosasco. Generalization properties of learning with random features. Advances in Neural Information Processing Systems (NeurIPS), 30, 2017.
[86] Othmane Sebbouh, Robert M Gower, and Aaron Defazio. Almost sure convergence rates for Stochastic Gradient Descent and Stochastic Heavy Ball. In Proceedings of Thirty Fourth Conference on Learning Theory, volume 134, pages 3935-3971. Proceedings of Machine Learning Research (PMLR), 2021.
[87] C.J. Shallue, J. Lee, J. Antognini, J. Sohk-Dickstein, R. Frostig, and G.E. Dahl. Measuring the Effects of Data Parallelism on Neural Network Training. Journal of Machine Learning Research, 20: 2019.
[88] Ohad Shamir and Tong Zhang. Stochastic gradient descent for non- non-smooth optimization: Convergence results and optimal averaging schemes. In International conference on machine learning, pages 71-79. PMLR, 2013.
[89] Utkarsh Sharma and Jared Kaplan. Scaling Laws from the Data Manifold Dimension. Journal of Machine Learning Research (JMLR), 23(9):1-34, 2022.
15

---

[90] Jack W Silverstein and Zhi Dong Bai. On the empirical distribution of eigenvalues of class of large dimensional random matrices. Journal of Multivariate analysis, 54(2):175-192, 1995.
[91] James B. Simon, Dhruva Karkada, Nikhil Ghosh, and Mikhail Belkin. More is better in modern machine learning: when infinite overparameterization is optimal and overfitting is obligatory. In International Conference Learning Representations (ICLR), 2024.
[92] Weijie Su, Stephen Boyd, and Emmanuel J. Candès. A Differential Equation for Modeling Nesterov's ' Accelerated Gradient Method: Theory and Insights. Journal of Machine Learning Research, 17(153))1-43, 2016.
[93] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In Proceedings of the 30th International Conference on Machine Learning (ICML), volume 28, pages 1139-1147, 2013.
[94] Aditya Varre and Nicolas Flammarion. Accelerated sgd for non-strongly-conve least squares. In Proceedings of Thirty Fifth Conference on Learning Theory (COLT), volume 135 of Proceedings of Machine Learning Research, pages 2062-2126, 2022.
[95] Aditya Vardhan Varre, Loucas Pillaud- Vivien, and Nicolas Flammarion. Last iterate con- vergence of SGD for Least-Squares in the Interpolation regime. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pages 21581-21591, 2021.
[96] Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and Faster Convergence of SGD for Models and an Accelerated Perceptron. In Proceedings of the Twenty- Second International Conference on Artificial Intelligence and Statistics (ICML),, volume 89 of Proceedings of Machine Learning Research, pages 1 1195-1204 PMIR, 2019.
[97] Pablo Villalobos, Anson Ho, Jaime Sevilla, Tamay Besiroglu, Lennart Heim, and Marius Hobbhahn. Will we run out of data? Limits of LLM scaling based on numan-generann data. arXiv preprint arXiv:221 1. 1.04325, 2022.
[98] Yichen Wang. Yudong Chen, Lorenzo Rosasco, and Fanghui Liu. Re-examining Double Descent and Scaling Laws under Norm-based Capacity via Deterministic Equivalence. arXiv preprint arXiv:2502.0 000225002.1883 5155, 2025.
[99] Andre Wibisono, Ashia C Wilson, and Michael I Jordan. A variational perspective on accelerated methods in optimization. proceedings of the National Academy of Sciences, 2016.
113(47):7331-
[100] Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie E Everett, Alexander A Alemi, Ben Adlam, John D Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, et al. Small-scale proxies for large-scale Transformer training instabilities. In The Twelfth International Conference Learning Representations, 2023.
[101] Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham Kakade. Last Iterate Risk Bounds of SGD with Decaying Stepsize for Overparameterized Linear Regression. In Proceedings of the 39th International Conference on Machine Learning (ICML), volume 162 of Proceedings of Machine Learning Research, pages 24280-2431 14. PMLR, 2022.
[102] Yan Yan, Tianbao Yang, 7he Li Li, Qihang Lin, and Yi Yang. A Unified Analysis Stochastic Momentum Methods for Deep Learning. In Proceedings of the Seeent-SSventth International Joint Conference on Artificial Intelligence, IJCAI-18, pages 2955-2961, International Joint Conferences on Artificial Intelligence Organization, 2018.
[103] Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference Machine Learning, pages 1 11727-11737. PMLR, 2021.
[104] Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs V: tuning large neural networks via zero-shot hyperparameter transfer. In Proceedings of the 35th International Conference on Neural Information P Processing Systems, pages 17084-17097, 2021.
[105] Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the infinite- width limit. arXiv preprint arXiv:2308.018144 2023.
[106] Dmitry Yarotsky. Corner Gradient Descent. anXiv preprint arXiv:2504.12579, 2025.
[107] Dmitry Yarotsky and Maksim Velikanov. SGD with Memory: Fundamental Properties and Stochastic Acceleration. International Conference Learning Representations (ICLR), 2025.
16

---

[108] Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sachdeva Sachdeva, George Dahl, Christopher Shallue, and Roger Grosse. Which Algorithmic Choices Matter Which Batch Sizes? Insights From a Noisy Quadratic Model. In Advances in Neural Information Processing Systems (NeurIPS), volume 32, 2019.
[109] Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Fos- ter, and Sham M. Kakade. How Does Critical Batch Size Scale in Pre-training?. International Conference on Learning Representations (ICLR), 2025.
[110] Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanguan Gu, Dean P. Foster, and Sham Kakade. The Benefits of Implicit Regularization from SGD Least Squares Problems (NeurlPS). In Advances in Neural Information Processing Systems, volume 34, pages 5456-
5468, 2021.
[111] Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham M Kakade. Benign overfitting of constant-stepsize SGD for linear regression. Journal of Machine Learning Research, 24(326):1-58 2023.
17

---

Supplementary material
Broader Impact Statement. The work presented in this paper is foundational research and it is not tied to any particular application. The primary set-up is on a simple well-studied random features model with synthetic data and solved using commonly deployed algorithms variants of stochastic gradient descent with momentum. We present (theoretical) compute-optimal curves for this model. While a common positive foresceable impact of algorithms with better scaling properties is that would allow more efficient model training and hence lower energy consumption of AI, it is important acknowledge that, many times across history, cost-lowering technological improvements have led to increase in consumption due to Jevons paradox. Regarding our experiments, we include results on language models trained on public language dataset that show the theoretical behavior holds on more practical settings. We do not release any pretrained models but because the dataset we used is based on common crawl which has potential issues of fairness, harmfulness, accountability, and transparency we may be contributing to reinforce these issues by supporting the making of the current C4 dataset as a standard NL., obfuscating these concerns. Finally, since our main results are theorems regarding scaling laws, we anticipate that one potential negative impact of such a line of research that fundamental advance in scaling laws will advantage large entities that have more computational resource, thus enhancing centralization of technological power.
Outline of the paper. The remainder of the article is structured as follows:
1. In Section A we discuss related works and compare previous convergence rates to ours.
2. In Section we provide more details on the algorithms under study, i.e., SGD, SGD-M, DANA-constant, DANA-decay. We additionally show that multiple existing algorithms can be re-framed within the class (Gen-Mom-SGD). We then use our results to precisely conjecture their scaling laws.
3. In Section C we derive the exact ODEs for the loss along with the exact Volterra equation. We discuss the simplification of this ODE and the resulting Volterra equation. We finally give some background Volterra equations and how to reduce their complexity.
4. In Section D we derive the random matrix analysis of the deterministic equivalent measures Hy,, p.x. Previous results from [80] estimated the deterministic equivalent on contour enclosing the spectrum of the data covariance matrix. We reformulate and strengthen these on the real line and provide general results for the integration of well-behaved real valued functions against Heg, Pxx.
5. In Section E we derive the compute-optimal scaling laws for SGD-M, DANA-constant and DANA-decaying.
6. In Section F we analyse SGD with a general learning rate schedule.
7. In Section G we derive the scaling laws for SGD-M, and prove that they are identical to SGD with precise correspondence between the learning rates of these two algorithms.
8. In Section H we analyse the DANA-constant algorithm. We use Frobenius method obtain asymptotic estimates for the solutions of the simplified ODE system (43) in Sections H.2 H.4. We use these estimates to derive stability conditions on the hyperparameters in sections H.8 and H.9 and compute the forcing and kernel functions.
9. In Section I we analyse DANA-decaying above the high- nigh-dimensional line. We give estimates on the solutions of (43) in Section I.1. We use these to derive the forcing and kernel functions, and the scaling laws. Finally, we discuss in Section 1.5 heuristic generalization to the algorithm class (DANA) with stability conditions and associated scaling laws.
10. In Section we study whether non-stable learning rates can accelerate the dynamics early in training. Especially we show that sweeping over the iy hyperparameter in DANA-constant recovers the DANA-dccaying schedule.
11. In Section K we discuss experimental details on the PLRF model and show numerical agreement between empirical and theoretical PLRF behavior.
12. In Section L we discuss experiments LSTMs language dataset.
13. In Section M report further experimental details for figures.
18

---

Notation. We use ( (t) when we want to emphasize the iteration counter t. We say (t,, A(t, v, d) for functions of (t, v, d),
v, v.d) 0 if for every 0 and for all admissible u and d, there exists i0, do such that for all d > do and 2 to
(1 - e), v, d) S of (1, v d) S (1 + + ( (t, v,d).
For some E NN we denote the diagonal matrix D Diag(j
i5)5s).
For A, B two matrices, comparisons of the form < B are understood coordinate-wise Additionally, ||A]] denotes the matrix of the absolute values of coordinates of A. For x E RR+, v 2 0 denote the positive square root of x and v-z def W
def
{= E C, Im(2) > 0} denotes the complex upper-half plane. For 1 > and u E we denote the Poisson kernel Pois,, def w² For R4, we denote r = max max [1,
For x E R, we use x4 max 0}.
Detailed proof sketch. We finally provide below a detailed proof sketch of our main result Theo- rem 4.1
1. The first step is to derive an evolution equation for the quadratic risk using the raw algorithm updates Equation (Gen-Mom-SGDD
We adopt in Section C.1 continuized analysis in Equation (Gen-Mom-SGGD by replacing the update times by the jumps Poisson process of rate 1 (Equation (18)). This allows us to reduce the analysis to ODE systems without sending the learning rates to O.
In Section C.I.1 we derive ODEs for the projected risk on each eigenvalue direction of the covariance matrix for two algorithms: the original one in Section C.1.1 and the coin- flip algorithm in Section C.2.2. The coin-flip algorithm uses two independent Poisson processes for the momentum and parameter updates in Equation which gives rise to much simpler ODEs.
The ODEs obtained are theoretically solvable but lot of small high-order terms make the analysis difficult. Hence we drop high-order terms in À in the exact and coin-flip ODEs to obtain the 'simplified ODEs' (Section C.2.2). We prove in Section C.2.1 that the 'simplified ODE' exactly represents the evolution of the SDEs presented in Equation (7).
Finally, using these ODEs, can derive Volterra equation for the full quadratic risk P(t) ((t)==Fp F + (KC * (t) where P(t) is the risk, F(t) is the forcing function, KC(t, s) the kernel function. This is done in Sections C.1.2 and C.2.3.
The rest of the proof focuses on computing asymptotics for F and IC (Equation (55)) by quantifying the empirical distribution of the covariance matrix spectrum and solving the ODEs.
2. Quantifying the spectrum of the data covariance matrix K.
We first rewrite J,F; ( in Equation (55) as integrals against two measures JAFF PK which are defined through the deterministic equivalent Equation (63).
Using Random Matrix Theory and previous results from [80], we provide upper and lower bounds on the deterministic equivalent measures fur, UK in Sections D. and D.2.
In Sections D.3 and D.4 use these bounds to define diverse components of the forcing function and kernel function (Equations (66) and (68)), formalized in Propositions 1.13 and 1.14.
3. We then solve the 'simplified ODE' Equation (54) for each class of algorithm in Equa- tion (Gen-Mom-SGD) separately.
For SGD-M we solve the ODE in Section G.I. These are the simplest as the ODE has constant coefficients, and the solutions are exponentials.
For DANA-constant, the ODE has time-varying coefficients. We employ Frobenius theory to obtain asymptotic solutions of the ODE in Section H.2 to Section H.6.
For DANA- Adecayin,,, we employ orthogonal polynomial theory to solve the ODEs in
Section I.I.
4. Finally, need to solve the Volterra equation for the risk.
19

---

For each algorithm to remain stable, the kernel norm has to be smaller than 1. This leads to stability conditions on the learning rates for SGD-M (Corollary G.1), DANA- constant (Corollary H.2) and DANA-decaying (Proposition 1.5).
Using this stability condition we can simplify the Volterra equation as P x F + using Kesten's lemma (Propositions G.8, H.20 and I.5).
then compute F, for each algorithm by integrating the previous solutions against JF and obtain explicit scaling laws for the loss P. This is done in Proposition G.4 for SGD-M, Propositions 13, H.14 and H.19 for DANA-constant and Propositions I.1,
I.3 and 1.4 for DANA-decaying:
Finally, we determine the compute-optimal regime in Section E by comparison of the different terms loss.
Contents
Introduction
2 The power law random features model (PLRF)
3
3 Continuized analysis of general stochastic momentum algorithms
4 Deriving the scaling laws of (Gen-Mom-SGDD on the PLRF model
5 Using Momentum to outscale SGD
A Related Work
22
Additional Algorithm Set-up
28
B.1
Stochastic gradient descent (SGD)
28
B.1.1
SGD with constant learning rate
28
B.2
Classic (constant) momentum (SGD-M)
29
B.3
CimeensionAAdapeer Nesterov Acceleration: DANA- -consacnnt
29
B.4
Data-Adapted Nesterov Acceleration (DANA-decaying) Time/Data-depen dependeen learn- ing rates
B.5
Comparison with Schedule-free. AcSGD, Stochastic Nesterov and conjectured scal-
ing laws
Deriving ODEs and Volterra equations
32
C.I
Derivation of exact ODEs and exact Volterra for the expected loss of meta stochastic momentum algorithm, (Gen-Mom-SGD)
32
C.1.1 Derivation of the Exact ODEs.
33
C.1.2 Derivation of the exact Volterra equation
36
C.1.3 Deterministic Equivalent of the Expected Loss
38
C.2 Simplified ODEs and simplified Volterra
39
C.2.1 Derivation of the simplified ODEs
39
C.2.2 Simplified ODE as the large time limit of coin flip algorithm
C.2.3 Derivation of the simplified Volterra equation
45
C33 Background on Volterra equations
46
20

---

C,4 Reducing the complexity in the Volterra equation
47
D Measure of the Deterministic Equivalent
48
D.I Estimating Wer
50
................................
D.1.1 Upper bound on po
................................
D.1.2
Lower bound μk
53
D.2
Estimating px ................................
................................
54
D.2.1
Upper-bound on Ax
54
D.2.2 Lower bound Pax
55
................................
D.3
Forcing function...
56
D.4
Kernel function
58
E
Compute-optimal curves - General
60
E.1
Stochastic momentum (SGD-M), compute-optimal curves
61
E,2
DANA-constant, compute-optimal curves
62
E,3 DANA-decaying, compute-optimal curves
64
E.4 Comparison of samples needed at compute optimality
67
E,5
Summary on compute-optimality results
68
Stochastic gradient descent (SGD)
74
F.1
Volterra equation for SGD
74
F.2
SGD with learning rate schedule
74
F.2.1
Simplifying the Volterra equation for SGD with learning rate schedule
75
E.2.2
Forcing function for SGD with decaying learning rate schedule
79
F2.3
Kernel function for SGD with decaying learning rate schedule
Classic Stochastic Momentum (SGD-M)
85
G.1 Solution to the ODE for classic stochastic momentum
85
G.2 Asymptotics of the kernel and forcing functions, SGD-M
85
H DANA-constant
95
H.I Simplification of the ODE
96
H.2
Getting asymptotic solutions of the ODE through Frobenius method
96
H.3 Fundamental solutions around zero and infinity
97
H.4
Behavior of Φ (1, 8) and (1,
100
H.5
Summary
106
H.6 Bounds the singular point 22 Bo [11 v) 478 B, (1 4)
107
H.7
Estimation of the forcing function
108
H.7.1
Asymptotics of Fo(().
108
H.7.2 Asymptotics of Fac and Tpp
"F..............................
109
H.8 Necessary conditions for stability
113
21

---

H.9
Sufficient condition for stability: upper-bound on the kernel norm
116
H.10 Upper-bound the kemel function
117
H.10.1 First term: SGD noise
118
H.10.2 Second term: momentum noise
122
H.10.3 Summary
125
H.11 Verifying the hypothesis of Kesten's Lemma
126
I DANA-decayine
132
L1
Solutions to the simplified ODE
134
1.2
Computing the forcing function
145
1.3
Stability condition for DANA decayine
146
1.4
Kernel function
148
1.5
Extended heuristics for general al algorihhm
155
1,5.1
Forcing function
................................
156
I.6
Stability conditions
156
1.7
Kernel function
158
J
Compute-optimality beyond stability and motivation for DANA-decaying schedule
158
J.1
Strategy
158
J.2 Stochastic Gradient Descent
159
J.3 DANA constant
159
J.4 DANA-decaying
160
K
Power-Law Random Features Experiments & Numerical Simulations
162
K.1
Power-Law Random Features Experiment Details
162
K.2 Measuring the Empirical Scaling Law Exponents: Chinchilla Approach
162
K.3 Computing the deterministic equivalent for K
163
K.4
Implementation of the ODE...
163
LSTM Language Model Experiments
164
L.1
LSTM Results: DANA-decaying across Kg
164
L.2 LSTM Results: Loss Exponents
166
L.3 Equivalent risk dynamics for SGD-M and SGD
169
L,4 Experiment Details
169
M Additional Information on the Experimental Set-Ups for Figures
171
A Related Work
PLRF Scaling laws on the PLRE Our work builds on the study compute optimality in large language models, introduced by [50] and [54], who performed empirical investigations of this phenomenon. The problem formulation follows [69], which examined data-limited scaling scenarios but did not address compute optimality algorithmic considerations. Related dynamical analyses
22

---

This work''
[DLM24]
[Bahri+21]*
[MRS22]
[BAP24]'
[Lin+24]o
Input dimension
d
d
d
M
M
M
# of features
V
P
Iterations/sampls
t
D
T
P
N
Capacity
20
1 +
+ a
a
Source
2a 20-2881
1 2
(1 - (
[1 )
6-1
b-1
>
1
4a
25
2a
Target decay
a + 3 aT
0
(in [22
[Paq+24] E. Paquette, Paquette, L. Xiao, J. Pennington. 4+3 Phases of Compute-optimal Neural Scaling
Laws. 2024
[DLM244] I.. Defilippis, B. Loureiro, T. Misiakiewicz. Dimension-free deterministic equivalents for random
feature regression. 2024
[Bahri21] Y. Bahri, D. Dyer, Kaplan, J. Lee. and U. Sharma. Explaining neural scaling laws. 2024.
[MRS22] A. Maloney, D.A. Roherts, J. Sully. A solvable model of neural scaling laws
[BAP24] Bordelon, A. Atanasov, Pehlevan, dynamical model of neural scaling laws. 2024.
[Lin24] L. Lin, J. Wu, S. Kakade, P. Barlett, J.D. Lee. Scaling Laws in Linear Regression: Compute, Parameters, and Data, 2024
Table 1:
Comparison of the source/capacity parameters across various related work. We note this table is taken from Table in [DLM24]] and Table 4 in [Paq+24]'. [Paq+24] also uses the same notation as this paper except samples are instead
through gradient flow in similar settings can be found in [16]. significant body of research has investigated scaling laws relating loss minimization to dataset size and parameter count across various settings (linear models, random features, deep networks). Notable contributions include [11]. [89], and [91], which explore the "hidden manifold" "hidden-manifold" data model for one-pass SGD.
Scaling laws and compute compute-optimalliy for the PLRF model under one- mne-pass SGD have been analyzed by [16, 63, 80] with an extension to feature learning in [17]. This work notably follows and extends the ideas of [80] to stochastic momentum algorithms. Particularly, [80] developed scaling laws for SGD under a deterministic equivalent and showed that there exists distinct phases. They used their scaling law to find compute-optimality, Related work on ridge-regression gradient descent includes [29] and [34]. Other theoretical guarantees for scaling laws beyond PLRF, typically for gradient flow, have been established in works such as [74].
Scaling hyperparameters in large models. There extensive literature studying the optimal scaling of optimizer hyperparameters in large neural networks, particularly the initialization and learning rate [39, 100, 103, 104], the batch size [70, 87, 109] and the epsilon hyperparameter in adaptive optimizers [39, 100, 105]. While the momentum hyperparameters for large models are often treated fixed constants rather than scaled quantities, the GPT-3 model [20] set B3 = 0.95 rather than the typical B2 0.99 0.999, which may improve stability with large batch sizes.
Random features and random matrices. Our work employs random matrix theory to examine random features problem that, from a statistical perspective, represents the generalization error of one-pass stochastic momentum algorithms. Random matrix theory has played an increasingly large role machine learning (see for [28] for a modern introduction).
For our random matrix analysis, we require sample covariance matrices with power law population covariance (i.e. linear random features). The analysis of sample covariance matrices precedes their usage in machine learning (see e.g. [12]). A detailed study of all parts of the spectrum of sample covariance matrices with power law population covariances appeared in [80] and has been subsequently used in [17] to study one-pass SGD and multi-pass SGD [7]. The study of ridge
regression has been extensively investigated (see for e.g. [10, 24)), and the work [34] provides complete analysis of the ridge regression problem under power law random features when 2a > 1.
23

---

There is a larger theory of nonlinear random features regression, mostly in the case of isotropic random features. For isotropic random features with proportional dimension asympptotics, this has been explored in works such as [71] and for some classes of anisotropic random features in [31, 68, 72]. We note that lots of the complexity of the analysis of power law random features arises from the analysis of self-consistent equations though the proof of these self-consistent equations falls outside the typical random matrix theory setting from textbooks. The use of self-consistent equations to study random matrix theory dates to [90], but the analysis of these equations (i.e, for K), as far as we can tell, dates to [80]). This strongly motivates non proportional scalings (which would be inevitable in power law random features with nonlinearities); in the isotropic case, the state of the art is [51].
Random features regression, 'source/capacity" conditions, and SGD. A large body of kernel regression and random features literature is formulated for "source/capccity" conditions, which are power law type assumptions that contain the problem setup here, when 2a >> 1 (the low-dimensional regime). For convenience, we record the parameters
2a:
233
Glgouree = 2a and Peapacity
4a
Here we have taken l'cpucityy the limit of those r's for which the source/capacity conditions hold (see Table 1). We note that in this language Tepaacity is often interpreted as 'hardness' (lower is harder), and that E (0, 0.5). E (0.5, 1.0) and r E (1.0, oo) correspond to 3 regimes of difficulty which have appeared previously (see the citations below); they are also precisely the 3 Phases Ia, II, and III.
Under source/capacity conditions, the authors of [85] establish generalization bounds for random feature regression with power law structures in 2a case for one-pass SGD. These bounds were sharpened in [29] and extended in [34] (see also [98]). An carlier work [22] shows kernel ridge regression is 'minimax optimal' under various 'source-capacity conditions', We give a comparison to these bounds in Table 2, but we note that the problem setup we have is not captured by 'minimax optimality' (in particular minimax optimality worst-case behavior over a problem class, and our problem setup is not worst-case for the traditional source/capacity conditions). Moreover, the authors [52, 94, 107] study stochastic momentum algorithms under source/capacity conditions (see below for details).
We note that this paper is fundamentally about scaling laws, but the novel mathematical contributions could also be recast in terms of generalization bounds of one-pass stochastic momentum algorithms. For SGD, the work of [23] compares SGD to kernel ridge regression, showing that one-pass SGD can attain the same bounds as kernel ridge regression and hence is another minimax optimal method (again under 'source-capacity" conditions). See also [36] which considers similar statements for SGD with iterate averaging and [81] for similar statements for multipass SGD; see also [35, 88] which also prove the single-bated versions of these. These bounds attain the dinimax-optimm rate, which are worse than the rates attained in this paper (see Table 2 for a comparison).
Stochastic momentum algorithms. Recent research has established convergence guarantees for stochastic classical momentum (SGD-M) in both strongly convex and non-strongly convex settings [40, 41, 76, 86, 102]". The latter references have established almost sure convergence results. For quadratic minimization specifically, SGD-M iterates converge linearly (though not in L²) under exactness assumptions [67], while under additional constraints stochastic gradient noise, [55] and [21] demonstrate linear convergence to a neighborhood of the solution. Batch size determination significantly impacts the convergence rate of both SGD and SGD+M. For small batch sizes, SGD+M does not necessarily outperform SGD [55, 78, 108], while some acceleration has been demonstrated for large batch sizes [15, 30, 60].
Convergence results for stochastic Nesterov's accelerated method [75] (SNAG) have been estab- lished for both strongly convex and non-strongly convex settings [6, 9, 58). Under stronger as- sumptions-such as the strong growth condition [96] or additive noise on stochastic gradients coveeconveeee to the optimum at accelerated rate can be guaranteed. As noted in [38, Thm. 7] (see also references therein), a naive implementation of stochastic Nesterov acceleration fails to converge in the non-strongly convex setting.
This is non-exhaustive list of work on stochastic momentum algorithms.
24

---

The absence of general convergence guarantees showing acceleration for existing momentum schemes in stochastic settings has prompted the development of alternative acceleration techniques, known as "accelerated SGD" [2, 33, 43, 44, 46, 52, 55, 58, 61, 64, 73, 94, 107], see (Gen-Mon-SGDD) Recent empirical success in large transformer models indicates benefits of accelerated SGD methods [83, 109], though theoretical understanding of hyperparameter scaling with model size [78, 94] and rigorous proof's of improved compute-optimal scaling exponents remain open questions.
The idea of accelerated SGD emerged from [3, 52, 61, 75, 94]. They observed acceleration can be obtained by coupling stochastic gradient descent and another update with aggressive stepsize. Consequently one simply has to scale down the stepsize in the aggressive update to make it robust to the gradient noise. In [61] (see also [66]). the authors produce an instance dependent risk bound for an accelerated SGD algorithm using constant (dimension independent) hyperparameters (at least above the nigh-dimensional line). We match the hyper-paranters of Accelerated SGD [61] to our set-up in Table 3. Another popular algorithm, Schedule-free SGD [33], was proved to achieve the worst-case rate (in the sense of classical analysis of algorithms) of SGD and can also be incorporated into various other algorithms (e.g., Adam[56]). On benchmarks, it performs remarkably well. After rewriting, we can express Schedule-free SGD in the form of MGD-)) see Table 3.
The algorithms AcSGD [94] and DANA in [78], are shown to exhibit acceleration for non-strongly convex quadratics by using momentum parameter (1 + t)- and dimension deeenddddndd learning rates on the order of 1/d. Specifically, DANA [78] showed acceleration in the -ighedimensional setting (samples and parameters are proportional) whereas AcSGD provides general bounds on quadratics in the setting that 2a > 1 (above the high-dimensiona line) under source/capacity conditions. AcSGI) achieves the rate EL| (0)) S min F: क/ } 1100 0* 1 + d 2a (1-28)+ Since ||0g - 0*1 carries a d d-dependence. it is difficult to obtain a scaling law from this. If [|0p - 1-2 x 1 (as anticipated in Phase II and III) the derived rate agrees with DANA-constant (with K2 = 1/d) at the time t x d that DANA-constar hits the irreducible loss level. Otherwise the derived rate is worse than the rate at any other time given for DANA-constan with 62 = 1: this makes sense as the AcSGD rate holds for more general data covariances and it was derived without taking into account scaling. Moreover we see that when t x d, AcSGD behaves as 1/t, or SGD similar to DANA- We match the hyperparameters in AcSGD with our set-up (see Table 3). Additionally, using our Theorem 4.1, We derive heuristically the scaling law for AcSGD (see Sec. B.5) which asymptotically matches DANA -oonstant with R2 = 1/d and we perform experiments showing the close relationship between DANA-constant and AcSGD, Figure 10.
(K2 = 1/d).
In a recent work, the authors [107] introduced SGD with 1-memory, a one-pass algorithm, analyzed on the (infinite) quadratic under the source/capacity constraints (2a > 1). In particular, the set-up infinite dimensional (d + oo) and does not contain the embedding matrix W in PLRF This algorithm has the same form as (Gen-Mom-SGD). The authors propose a variant of DANA-decaying (with K3- = 1/(2a)) (see Table 3). While not proven, they heuristically expect in Phase Ia/lI (signal-dominated regime) that the rate matches our F rate. We prove this rate, well as the rate for the noise-dominated regime (Phase III) (Theorem 1.2). The work of [107] is based off the deterministic result which looked at power law eigenvalue distributions on the conjugate gradient algorithm, see [19]. In a concurrent work [106), the authors used constant momentum and learning rates in (Gen-Mom Sinn--oo-SGDD to show acceleration in the infinite dimensional (d + oo) set-up but using oo-memory SGD. This would require storing multiple (in fact, infinite) vectors, each on the size of the parameter space. They propose a finite memory version that approximates the infinite version. Such a result is an interesting future direction to derive scaling laws.
In [40, 73], they showed many known algorithms could be rewritten in the form of (Gen-Mom-SOO SGD). We include Table 3 to provide an equivalence of hyperparameters of related work with (Gen-Mom-SGDD These include stochastic Nesterov, AcSGD [94], Schedule-free SGD [33], accel- erated SGD from [61], and SGD with 1-memory [107]. We perform scaling experiments with the PLRF model on these algorithms in Fig. 10 and Fig. 1c
Lastly, the notion of effective dimension has been used by [14, 61, 101, 110], but never quantified directly as a time-dependent learning rate.
Dynamical deterministic equivalents, Volterra equations and ODEs. Using the deterministic equivalents for random matrix resolvents [47], we in turn derive deterministic equivalents for the risk curves of stochastic momentum algorithms.
25

---

This work
This work
Algorithm
one-pass DANA- (R' )
one-pass DANA-constant (N₂ 1), n 2
Phase la
(n (fn-00/(2)))
(n (20-10022a))0-ss Vdd
((
Risk (n)
Phase
Vd (2-ns)
Vd 1/(20)))) 241//a)
(n (2 1/(22)22) (2 (2-15))))))
O(d 1/(22) 4+ v vd 1/32/( 2211/a
Phase III
Vd "n (2=
[VF22] (see also [52])
Algorithm
one-pass Accelerated SGD (n 2 d)
Phase
(dn 2100 0r 112 P)
Risk
Phase II
(dm 21100 0-0000'' 11 (2)
9 (n)
Phase III
Odn 31100 - 0- -0
This work [Paq+24)'
[DLM24)
Algorithm
one-pass SGD/SGD-M
RR O(1)-rigg
Phase la
(n (ln=) Pd
same as [Paq+24]
Risk (n)
Phase
(n -/(22) vd 1+1/(20) vd vd)
same as [Paq+24]'
Phase III
(n - 'd
20-1 20
Vat Vd-2ay 20)
O(n An 2n-1
vd-
Minimax optimal 789
[Lin+24]
one-pass SGD,
Algorithm
one-pass SGD, very small stepsize
very small stepsize
Phase la
0 (n=PP(204220)
(/d +ing w, (111/22)))
Risk
Phase
n (2(2))22))
does not cover
(m) Phase III
((=-Ax((4a+1)
does not cover
Carratino, Rudi, Rosasco. Learning with sgd and random features. 2018
Dieuleveut and Bach. Nonparametric stochastic approximation with large stepxizes. 2016.
Pillaud-Vivien, Rudi, Bach. Statistical optimality of SGD hard learning problems through multiple
passes. 2018.
Varre, Pillaud-Vivien, Flammarion. Last iterate convergence of SGD for least squares in the interpolation regime 2021.
Varre, Flammarion. Accelerated SGD for Non-Strongly-Convee Least Squares 2022.
Table 2: (Nonexhaustive) Comparison of samplc-complexit results. Let p del 2a 23 1. We sample size(t in our notation), d parameters. [DLM24] can also be done with RR+optimal-ridge, which yields in Phase Ia, but different in Phase II/III. [VPF21] obtain S < that is, they capture the Fpp. but not Fac. The minimax optimal SGD rates never achieve any the rates (always worse), which can be connected to overly conservative, small stepsizes. For derivation of the minimax rates, we used Cor. 2 from [DB18]'. [Lin+24] requires label noise order 1 and also very small learning rate. For [VF22]"0 we believe (though proven) after numerical experiments that ||0b- 8" 2 x d -229 in Phase la and otherwise constant in Phase II/III. As it takes on the order of +1/ to reach stationarity for DANA-come Aons-cont,nnt, we see that our bounds improve over [V but similar. In particular, the two results agree as one approaches stationarity. For all algorithms, set batch size
26

---

Equivalent Hyperparameters Alg. (Gen-Mom-SGD)
Good Choices for Hyperparameters
(1(1) a(411)
4.13.
â cı/(d Tr(D)), ca/TTDD) C1, 02 constants; See [94, Thm 3]
AcSGD[94] (ā,
72(1) B, 7t() = EFI*
A() [42
(t) M(++1)
1=1
S-Nesterov (5)
++2
1-2
c2 o2/Tr(D), c2 constant: See [40] for details
(2(t) 3. T3(1)
11
((1) 1+2
1
Schedule-Free
(1 1, 72(t) 7/1 8),
0.9, 7 = c2 (T)) constant:
SGD [33] (%,
7u(t) ++1 A(t)
58
1
See [73] for equivalence proof.
4+
(t)) a(ō - -1),
See [61] for choices of hyperparameters: Time-independent hyperparameters (f), â, .8), but some depend on 1/d.
ASGD [61] (5,6,3,0)
72(f) âá + (1 8)%
(t()) (1 () B),
(() (1 f)
(1(t) I.
SGD with
(f()) 90,
Heuristic only (no specific constants)
1-memory [107]
(t()) = (1+t) a)
0 à (2a)
(90, ã,
0 811
A() = (1 + 1)
Table 3: Other algorithms in the form of updates given by (Gen-Mom-SGID) and hyperpa- rameters. See [40, 73, 94] for details. Although not proven, AcSGD [94] should attain similar scaling laws as DANA-constant and ASGD [61] should attain similar scaling laws as SGD. We note that ASGD [61] is not quite in the form of (Gen-Mom-SGDD as the update in y uses a stale gradient V ( (0:- instead of V (0,; 1) for (Gen-Mom-SGD). SGD with I-Memory [107] heuristic algorithm independently developed at the same time as DANA-decaying and very similar; No proof is given. We believe our result proves this algorithm as well.
The method of analysis of the risk curves in this paper is by formulation of Volterra equation. For SGD and stochastic momentum, these Volterra equations were studied in the high-dimensional regime, see [60, 77, 78, 79, 80]. Particularly, system coupled difference equations for weights the residuals in the observed data covariance. This has been shown to generalizes beyond the least-squares context, at least under SGD, [26]; in isotropic instances, this simplifies to a finite-dimensional family of ODES [4]. This can also be generalized to momentum SGD methods [78] and large batch SGD methods [60]. Convolution- Volterra equations are convenient tools, as they are well-studied parts of renewal theory [5] and branching process theory [8].
use the formulation that derives the Volterra equation via
Another method analysis is dynamical mean field theory. The closest existing work to this one in scientific motivations is [16, 17], which uses this technique. This formally can be considered as type of Gaussian process approximation, but for a finite family of observables ("order parameters"). In instances of one-pass SGD (including in anisotropic cases), this is rigorously shown to hold in [42]. The analysis of the resulting self-consistent equations is nontrivial, and [7, 16, 17] does some of this analysis under simplifying assumptions on the structure of the solutions of these equations.
Besides these works, there is large theory around generalization error of SGD. The work of [95] gives a direct analysis of risks of SGD under "source/capacity" type assumptions which formally capture the Fpp parts of the Phase la/lI loss curves. The risk bounds of [111] give non-asymptotic estimates which again reproduce tight estimates for the Fpp parts of the loss (note that to apply these bounds to this case, substantial random matrix theory needs be worked out first); see also [63] where some of this is done.
27

---

B Additional Algorithm Set-up
Let 72 73 (0) - (0, oo) be learning rate schedules and A [0, [0, x)) be the momentum schedule. To solve the PLRF (3), we use a class of one-pass (mini-batch) stochastic momentum algorithms with batch size B. Let y-1 du Rd. At each iteration t 0, generate independent, new samples {2(+1)== and update by
B 1/1 =1 (1- (1-1(t1))) (11-((t))-+++all Σ
1 x/+1
x/41, (2/+1. 6))
(13)
B W x/41 ((W" 2 6 0/+1) ()) (rs(t) x yt.
B111 = 01 - (2(()) Σ
where (t), (t) are non-negative functions. We will consider multiple versions of the algorithm (13), that is, with different choices for the learning rates 7's and the momentum schedule Att). In Section 3, we summarize the different algorithms we consider based on hyperparameters as well as good choices for those hyperparameters.
B.1 Stochastic gradient descent (SGD)
An example of known stochastic algorithm that falls within the update rule given in (13) is the stochastic gradient descent (SGD) with learning rate schedule, (2()). This algorithm is determined by setting 7i, 73 0, A 1, and 72(1) [0, oo) (0) learning rate schedule in (13). Specifically, it updates by
0111 0, 72
T Fe (WW 12 (z)+1:
(14)
In Section F, we derive the deterministic ODEs and analyze SGD for compute-opiimalit when the learning rate schedule (2))) is a decreasing function.
B.1.1 SGD with constant learning rate
The compute-optimal curves for SGD when (2 a constant was studied extensively in [16, 63, 80]. In [80], necessary and sufficient condition for convergence of the algorithm based on the learning rate ?/2 was established (see [80, Prop. C.2]). Particularly, the condition was
2
and 1- (() (()) <
72
B
B+1
where H the kernel function for SGD (see Section F for a specific definition). Asymptotically, we know that X ([80, Cor. G.1)) satisfies
l
(t) dt
72
Σ ,
2a
2
j=1
Therefore, have two cases to consider: 2a (below high-dimensional line) and 2a: 1 (above high-dimennn nigh-dimensional line).
Remark B.1 (Stability conditions for SGD with constant learning rate). Let Tr(D) def ビー When 2a < 1, suppose that Y (1, oo). The stability conditions for SGD with constant learning rate ?/2 are
(2o > 1): (2a < 1) 7/2 <
72 <
B+T
2
and
72 < Tr(D)
2
B+1
2
and
"/2 <
Tr()
2 211-2a] d0120))
This means that for > 1, the learning rate is always constant (order 1) as the Tr(D) summable. Below the high-dimensional line, Tr(D) grows with the number of features d.
28

---

B.2 Classic (constant) momentum (SGD-M)
In this section, we consider the classical stochastic momentum where the learning rates "ýi 1, 7/2 52d=*1, and 73 "zd=*2, and momentum, A = 6, is constant (does not change with time and order 1 with respect to d). We call this algorithm stochastic gradient descent with momentum (SGD-M). In particular, we note that (stochastic) heavyball [82] occurs when 72 = 0 and "7 and A are specific constants determined by the largest and smallest (non-zero) eigenvalue of K. In this case, the updates in (13) with y-1 0g- = 0 follow
(1 8) + Σ
WT
oi r(WT 1141, (1 (1)+1,6))
==1
(15)
B
0ut1 0, 72
WT ('W
(2(+1: b)) - -is x yc.
We will show sufficient condition on the hyperparameters 72, '3 d in Cor. G. for the simplified Volterra equation (55) to remain bounded. However, we emphasis that there are discrepancies between this stability condition and the actual stability condition of the algorithm. Indeed, the simplified Volterra equation neglects A terms which makes any A >> 0 converge. On the other hand, it is possible to make similar analysis that case, we can derive an explicit condition for stability of the algorithm. Specifically, the stability condition takes the form 2)
different of ODEs (coin-flip ODEs (49)) with 72
In
73 0 with d (0, The condition å E (0,2) crucial to avoid exponential growth of the sequence (yr. 2 0). This shows limitation of the simplified ODEs (43) which may miss some stability conditions.
25(2-5)
B.3 Bimension-Adaple Nesterov Acceleration: DANA-constant
Another important algorithm we consider is DANA, introduced in [78]. where it was shown to accelerate SGD in the proportional d and sample setting (a.k.a. thermodynamic limit). The main distinction with (DANA) over SGD-M is the momentum schedule. Following a Nesterov style momentum, we set (t) = &_
(possibly dimension-dependent. T+7. DANA-constant, the learning rates are all set to be constant We will consider a more general setting for consacon in the appendix then what was introduced in the main introduction. Setting y-1 = Oo = 0, the DANA-constant algorithm updates as
01) - ( b)) y1 = (1 (()))(-1+ EB WT (( x(+1
(DANA-constant)
D++1 02 (2 EES WT ((WT o441, 0²) - (0t+1, b)) - ya(d) x yt-
constant d
constant
where A(I) = i ++1 1 1, (2
and
ra(d)
d=a
(16)
Here we have batch size B t where the exponents Kb, Ki, R2 2 introduced in the main introduction, we set K2 = K1 + and Kg = 0. At times throughout the appendix, we will reduce back down to DANA-constant with these specific parameters, but for the proof, see Section H, we will use this more general DANA-constad setting.
0. For the DANA-constant
Remark B.2 (Good choices for hyperparameters of D.NA-constant.). In Lemma H.5, we give sufficient conditions for stability of the solution to the Volterra equation (55) on the simplified system of ODEs (43) that accurately predicts the dynamics of the stochastic algorithm DANA-constant, For necessary conditions on the hyperparumeters (on the simplified ODEs), see Corollary H.2. Some good choices for hyperparameters of DANA-consten (in batch size B 1) are
12 + 2 1
C2 and Ys =
C3
8 > 4 X max
S
n 1 72
Tr(D)'
x Tr
71733 d
12 Tr(D) < 1.
satisfying
272
2
Here C2, e3 positive constants and the matrix D Diag(j 2a 15 و v).
29

---

1.000, batch
3pha
1.000. batch
1.4,
1,000, bakch
10a4
107
12-
12
12
30-2
ri6k
10m
13
12"
10
30
20
flops
flops
flops
(a) Schedule- free SGD, single
(b) AcSGD, single
(c) Stochastic Nesterov, single run
a/3ha = beta beta
1.0, batch
alpha = 1.4, beta 1.0, batch
10"
alpha 1.4, beta 1.0, batch
13-1
34-
10-
23
10
201
10a
18-33
100
10
30
10
100
18
100
13
10
107
127
30
310
490
10
17
flpss
flops
flops
(d) Schedule-free SGD, scaling be- havior
(e) AcSGD, scaling behavior
(1) Stochastic Nesterov, scaling be- havior
Figure 10: Comparison of SGD, DANA with other known algorithms (Schedule-free, AcSGD,
Nesterov). Numerical setup: SGD (blue curves) learning rate F22 = 0.5/Tr(DD), DANA-constant
(green) has (1) 1, (t) 0.5/ (T((),, 7a(()= = 0.1/d, (t) 5/(1 + t) where 5 = max
1/a, 23- +1; DANAA-eccying (orange) same Y1, 72, and A as DANA-constant,
0.1/(1 + t) 12(m)) Schedule-free SGD (red in 1st column) ,3 = 0.9 and Fy 0.5/ in [33];
AcSGD (red in 2nd column) ā = OrL/Tr(D) x DD and _ 0.4/Tr(D) as defined in [94]; Stochastic Nesterov (red in 3rd column) â = 3 O.1/Tr(D) 1.0/d defined in [94]. Top row: algorithms were run for 10 steps with df = 1000, e = 10000; Bottom row: algorithms run using the ODEs with equivalences given in Table 3: 10 iterations of algorithm, d (100 2 i 1, 1,..., 10 and e = 10 x d. (Ist column: Schedule-free SGD [33]) Schedule- free SGD (red) scales very closely with SGD (blue) for this (a. We both DANA-decay and DANA-constann accelerate. (2nd columm: AcSGD [94]) AcSGD (red) gives scaling laws similar to (green). Moreover, it has the same property of changing behavior at d. Although an upper bound was proven for AcSGD in [94], this is not optimal when applied to the PLRF and, in particular, one would need to understand 110, 0- 1|22. A similar technique as our approach to DANA-constant should apply to AcSGD, (3rd column: stochastic Nesterov) Stochastic Nesterov is known to not converge (see e.g.. Thm.7 [38] and references therein).
30

---

B.4 Data-Adapted Nesterov Acceleration (DANA-decaying); Time/Data--dependent learning rates
In the same spirit as (16), we consider a variation of case, we not only have momentum time-dependent schedule, but we also allow the learning rates to be time dependent. The motivation for decaying 73 schedule is, in light of Section 3, to avoid the trivial behavior of CANAsccnttant for small by using larger "/3 at the start, while decaying progressively to reduce the noise and preserve the stability of the algorithm. Therefore, we introduce the DeNA-dec algorithm given by
called DANA-decaying, In this
yt (1 ((1))r +n E W s ((WTx/-1 01) (0k-1: b))
(DANA-decaying)
0(+1 8, (2(d) WT x ((WT + O1) (xi+1 (a) ra(1: d)yl--
constant
where A(t) t 1 dt(
and (2 71 constants.
(1 + f)
(17)
Here the exponent Ki 0 and batch size B is independent of d. We only prove the result when 2a > 1 and (d) x 1, see Section I. In the setting of 2a > 1, this amounts to choosing (up to an absolute constant) the maximal learning rate. We provide some good choices for the hyperparameters that work well across scales on the PLRF
1). In Section I, we provide some heuristics as to the correct sufficient conditions for stability of DANA-decay, More precisely for (for DANA- decayyeca,nng, one should pick Kig
Remark B.3 (Good hyperparameters for DANA- with B =).
1 the conditions are stated as 2a
any
K3 > 2a
K3 8+
> (1
K3 2
max
2a 23
1 4
1
72(d) Tr(D)*
C2
2
a
(d)
and
ra(l; d) (1 + t)ms
Here the constants C22 C > 0 should be chosen small enough and the matrix D Diag(j- 20 1
j S v). Below the nigh-dimensional line (2a < 1). we are not certain if another choice of may accelerate.
In the rest of the appendix, we use the following.
Unless it is clear that Kz/Ra, are being used, e.g.. in the proofs of DANA--comstaa//AANN decaayin (Section and Section 1), DANA- constant refers (K2
Remark B.4 (DANA-deaayin//DANAA-oo
K1, K3 0) and DANA-decaying refers to (K2 0, K3 1/(2x)). In some cases, we will further reduce to the setting where batch size is order in We note when refer DANA-decaying we always assume that batch size is order 1.
B.5 Comparison with Schedule-free, AcSGD, Stochastic Nesterov and conjectured scaling laws
Other algorithms (e.g., stochastic Nesterov, AcSGD [94], Accelerated SGD [52, 61]) be written in terms of Gen-Mom-SGD. For the equivalence of hyperparameters of these algorithms to the hyperparameters in Gen-Mom-SGD,, see Table 3. See also Figure 10 for their scaling performance compared with DANA and SGD.
We conjecture that Theorem 4.1 can be extended to include time-dependent 71 d)) In some sense, the algorithm Gen-Mom- SGD over-parameterized, that is, for Gen-MommSSGD(1u 72, 73. x (1 + ([+t)=] 11) such that
(1 (ff:ddxxxxt:::: d) = (s((;d] d)
Gen-Mom-SGD(11, 72, 73. A (1
( DANA(12, 1z).
This certainly true when 71 constant (d-dependent allowed). It should also hold (at least asymptotically) when (1; d) and ra(1;d) are time-dependent provided that y (4; d) is growing and (()) is nice. Intuitively, this is true as O.+1 in Gen-Mom-SGID gets updated by approximately Ti (t; d) x va (t; d) x Ye after unrolling the recursion on y/ s.
31

---

Under this belief and using the equivalence with our parametrization in Table 3 we can very precisely conjecture their scaling law behavior.
Schedule-free SGD. In our parametrization, Schedule-free SGD uses the asymptotic learn- ing rates (tt d)n d)
40 and (2(t: d) -(11 3) with y Tr(D)' / x Therefore, we see Schedule-free SGD is approximately DANA-decaying with 72 d) Tr(D) and 1. Then from Theorem 4.1, we have that the
1+1
(a(t; d) TTXX1+1 or, if you like, function 0(1) x 1 + Bt. This leads the conjectured behavior:
Conjectured Scaling Laws for Schedule-free SGD: For any a > 0, 2a + 23 > 1, Schedule-free SGD has the exact same scaling laws as SGD.
AcSGD, our parametrization, AcSGD [94] uses the asymptotic learning rates (t))aa()) = ā and (2(t) 3 with à x ZxTr(D) and Tr()) We see that this exactly
corresponds for large time to the schedule for DANA-constant with 7/s(t; d) = 72 x J/- Using our result on
Theorem H.3, this leads to the conjectured behavior:
Conjectured Scaling Laws for AcSGD: For any 0, 2a 2 1, AcSGD has the exact same scaling laws as DANA-constant.
Stochastic Nesterov In our parametrization. stochastic Nesterov uses the asymptotic learning rates 7) (()) = 7 and 722 (4) = 5 with 5 x TH(D). We see that this exactly corresponds for large time the schedule for DANA-constant without the 1 scaling the ra(t) learning rate. Using our stability result on DANA-constan Corollary H.2, this leads to the conjectured behavior (see related results in e.g., Thm.7 [38] and references therein):
Conjectured Scaling Laws for Stochastic Nesterov: For any a >> 0, 2a 2/3 >> 1,
Nesterov Momentum diverges for large times.
C Deriving ODEs and Volterra equations
In this section, we derive Volterra equations that match the dynamics of the meta stochastic momentum algorithm updates in (13). We develop two systems of ODEs that culminate in two (related) Volterra equations. In Section C.I, deduce system of ODEs and Volterra equation that, in expectation, matches the expected loss under the updates (13). We denote these as the exact ODEs and the exact Volterra equation, resp. While numerically solvable, analyzing the exact ODEs presents significant challenges. Therefore, we introduce a system of SDEs derived from studying stochastic momentum algorithms in the high-dimensional optimization framework [78]. In Section C.2, we derive the ODEs and Volterra equation for the SDE framework (simplified ODEs), resp. We empirically demonstrate that this alternative simplified ODE system accurately model the behavior of SGD/SGD-M, DANA- constant, and DANA-decaying (see Fig. 26, 27, 28, 29, 30, 31, 32, 33), and are considerably simpler for analyzing scaling laws. Throughout our analysis, we clearly indicate which system of ODEs/Volterra equation we are examining for each algorithm.
C.1 Derivation of exact ODEs and exact Volterra for the expected loss of meta stochastic momentum algorithm, (Gen-Mom-SGD)
In this section, we derive a Volterra equation for the dynamics of the meta stochastic momentum algorithm which updates using (13). Let N₂ be iid Poisson process. At jump of N,, we generate {x²), ) new data points whose coordinates follow an a-power law and update by
Y - (()) +n((t)
W x'y. (T -6) def Y + 0,_))
B
, = 0 - 72(1)=
W s'y, (a'x.) y (WO,- b) 71)( Y. del O, O.-))
(18)
32

---

Here we naturally extend time t to be a continuous parameter with Y = Y--1 and we start this process so that O do and Yo- = y-1, that is, at initialization the two stochastic algorithms agree. This is standard way to embed discrete processes into a continuous process.
For any (O, Y) (( E R, the mean behavior is given by
d L [f 0.)l E[f(Y AY O,), 9r (.Y, f(Y, O.) dt
(19)
C.1.1 Derivation of the Exact ODEs.
For this, need to introduce statisties that are closed under differentiation defined in (19).
To do so, we begin by writing R" = lm()) @ W Thus, there exists a E Rd and b e R" such that one can write b = Wb + b, that is, we can decompose b as element in the image of W and an element the co-ker of W. Formally, we have that
= Wbb bb, where T Db = 0.
We will now choose some specific functions/statistics that will close under differentiation. Let K def WDDW E Raxd and let 4 by the eigenvalue-eigenveeoo pairs for K. We will consider the follow special statistics/functions for each eigenvector of K
()) = [[ (0: - )) W], (t) = E [u, Y, ,
(20)
and X²(1) E [( (0, - b) Y.)
We recover the expected loss function from the statistics (20) by
E [[ (0x=)) Σ>pp}(t), where F [9 lim E (}.)))
Therefore, knowing the evolution of (t) and E [[)P(0>)) for every allows us to recover information about the evolution of the loss curve under O. Moreover, we will show that the functions E2 close, that is, there is closed system of ODEs that defined their behavior.
First, we observe that we will need moments of Gaussians via Wick's formula. In particular, for fixed vectors E R", = 2, 3, 4 and 2, = j 3-2, zy where zj (),
E (2 vi)l = Do, E Ez (x, 02) (zx an) (z, va)l (D, un @ (D, vs @ ml) + (D,01 (D) 0
+ (D, vn @ va) (D, U2 & 03)-
The (v ma-matrix D Diag(j-2o S S v) is the covariance of x. Using these moment computations, we can compute dr for the functions in (20).
03 functions: Using (19) applied to E [S(0,, Y,) - (1) yields,
d 000(()) 2E (w ( (2()) DL ['=((x)) - b)) 8 (0, (0:-bb) b)) [W]
+ 2F
[w / 02 (1) A(1) (0: -)) W]
02 (
+ 2E
b)) @ (0: b)} W|
+ 2E
fu (- ("(d)) (1) A(t)) w]
2E (u02, (1)2(1) Wodc)""" b))
@
(WW b))) W
+ 2F (2 (1( - ((1))) Σ WT (xt) (WQ, b))) W]
+ [w
(- (->2
[
+ [w (-1a(0)( - -())) 89 W
+ [w (->1(17)() W WT - -b)) 0)) 02 W
33

---

For convenience, we drop the t's on the learning rates 7i()) = % and (t) A. Applying Wick's
rule the RHS, we get
di (1) = -2 BB)0)=33BBB))0000 -0
+ - Trtpah/we...a..aanaaaa//aaaaaaa0
+ 9rr -
+ 2
+ (I + BB
3 functions: Using (19) applied to [ Yi) (t) yields,
(t) = E [[uof? ((1 ((t)) - (1()) E2 W" - -b))
E [ [w]
2E [w Y2 8- -((1))) W| + 2BE ( T. - -b)))
+ +E [ (A())),) 8 w]
b))) 2BE [w A(I)), 8 (1)) 2.,7 (WO, b))) W]
(t) [w]] (W 2 - 6) 03) W]
+ (B) - I)E (w0 (w ( - b)) @ (WT))( - -))) [W].
Applying Wick's rule, we get
de5) (1) = (0) (1) 271 B))xa(() ())x(() - 22 (1) (0m
d
+ B2 + Ωं BX [|] b)] W| + B(B
- 1))(())3)3()))
Xi functions: Lastly, using (19) applied to E[ Yi) X₁ yields,
E W| E [Mx 0, 72 Σm W =T=;(=)) b)
- 7ss(( A)), m Em= W' 2/(2)) (WO, - -b)) -b
((1 71 Σt WT W" (1)=:/)(ee)))) -6)) [W]
- ()).))))))))) -E[ (0, & Yi) [w]
Applying Wick's rule, we get the following
E W] (t, 7a(1
- - Ax;
+
+ 70( - 71 (1) B1
(2) - (wG_)))000000,,,)00((((///-- b]]l
(TMB(BB1)A330(0 (1008/8000000111//1,aaaaaaa -2
- >)
- - (t)
Putting this altogether yields the following system of linear ODEs.
pairs for K DW and decompose b = wb + such that WT Db = 0. Here the (v x v)-matrix D Djag(j-20 20 1 j S v) and (t) del (e,) = | (we, - b)ll2. Additionally, we drop the time in the learning rate and momentum schedules to simplify notation. The functions
Exact ODEs for p}. E3 Xj- Let (y),wy) be the eigenvalue
p0(1) = E ( (O, )
{)(() = =E(( y,
(21)
and (t) E ((u)22 (0, b) @ W]
34

---

form a closed system of linear ODEs:
d
ह( r(; "(tt; (1) = 22(t A2) x r(t; A1) + g(t; X1);
(22)
where (t(t: X,) = S2(t: 21) + Q(t: 24) and g(t A2)
g(t: X,) + g(t: >1) such that
B22bBB
-2%
(1(1)1)) der
0
22
2n BX,
BB
-73
-
//Bb,/
-2 + (2:1 - (B + 1)43
18 3)-
2%. 2 + AJBA,
(11;2,) mn
BR 1)A9
A
m 13)A3 B 1)A7
7A(22 A)
(7d 2(1
TB
BE (t) W]
(2717273 + +,,73)2, (P())
x(t; /
A (t) (t)) w 3(t:A)) de
0
(-nom- (
(t))
and Az) del
(1))
(1)
(23)
The initial conditions are such that
v(0;Ag)
(0)
Remark C.1. We can write the functions (03 53, Xi) in terms of the non-zero eigenvalues and eigenvectors of the conjugate matrix K def D D1/2 10 We show this equivalence below.
Proposition C.1 (Equivalence of (p}: 23. Suppose K- = D1/2WWT D1/2 where =
Diag(j 2n j = v). Then the following holds
()(()= E (ug, (DV/3WW0, - b))}²
(0 =E (uj, 11/22)2
W
W
A (ug²
(D4/ b)) @ D1/2WY,)
and xx(t)=E
W
where (Aj, u1)= are the nonzero eigemmaluesseeggeneeooss of K.
Proof. Let D1/2W VEUT be the singular value decomposition where V [ug, Ug, Un] Rexe, U = w2,,w,,... wa] E Rdxd and Σ E Rexd rectangular diagonal matrix with non-zero singular values o, where j = 1....,d. d. We prove the result for (1 noting that the results for E3 and Xj follow a similar argument.
We note that WT Di = 0 implies that WT (D!1((2)))) = 0. In particular, this means that
DI/2 0. By hitting both sides by EUT we get that
EU" = 0
D'
Thus for all nonzero singular values oy (or equivalently for all nonzero eigenvalues X, of K WTDW), u² D1/26 O. Additionally, we have that
WT UYT U' -UE = = بزWر.
where ej is 0 except for 1 in the jth position. Note this holds for all non-zero oj.
1hTese are precisely the statistics described in (4).
35

---

We have that = wbb + b where M Db = 0. Therefore,
(DI/ - b))) {uz (W) -D1/28)2
Aj
(ug, D1/2W((, 81}2 (uz. D1/22)2
21 (aylws, (0: (O)-
The last equality follows since of = Ag. Taking expectations, finishes the proof.
Remark C.2 (Initialization). 1f0n = Yo 00, then 13 = F [(u,. 0 for all j = ,,...,d.
(0)
Since the nonzero eigenvalues of K the same K, and in light of Prop. C.I, these {Gp.. E2. x)))s satisfy the same exact ODEs as in (22). Throughout the remaining sections, we use {(o3. <3. x1) defined by the (nonzero) cigenvaluesddigeeveettos of K DD/2
C.1.2 Derivation of the exact Volterra equation
In this section, we derive Volterra equation that deseribes the expected loss using the exact ODEs. To do so, requires abstractly solving the inhomogeneous ODEs.
To solve (22), we employ Duhamel's principle. Let Φ 8) for / 2 be the solution to the IVP
de 2r s) = S11; X3) Ps, (4, s) such that (8, Ida for all 1 S j S d.
(24)
Then by Duhamel's principle, we have for every j
((t;A)) 2x (t))))(0;)))) ''' (t,(t))))((; ds.
As (t) vtt; A1)1, we have that
d
E (()((t) W| =
lim E (t))
W + Σ
A3 x (t)
t1000
j=1
lim (t)( W + Σ t++00
X₂ (t))0())))))
(25)
j=1
J. (( W
B Σ X3 x Φ, (6)))((0 ds.
Here the vector ((x) E IR3 is given by
72 + (2 + +f72)) ?
h(s) def
(-702- - 7PTs)
where the learning rate schedules 7% i 1,,,,3, are time-dependent. We denote the components of hs)) as hs(s) where i 1, 1,2, Therefore, the expected loss satisfies Volterra equation (not necessarily convolution-type"
E [ = #(1+ S' X((s) x [ [ W| ds,
(26)
d
where f(t Σ X₂ x (t., ()) x 23) + lim (t
and Xx(t) del > 23 (s) (t, s)) + h2(ss (t, 8)) + hs)ss) (t, 8)) 1a]-
36

---

Let K VAVT. Define Dæ (t,s) def ( (t,s)) 1 1 S v) and (t, def VD (1,8)VT. Here define for the Where might this come from? By allowing some leniency in the definition of the ODE (24), we can view as parameter (in the case we looked at above, we set À to be an eigenvalue plugs À 0 into (24), then due to the initial conditions, we get that Φ (f, 8) = 1 for À = 0. Using the (t) as defined in Prop. C.I,
zero eigenvalues of K to have (Φx, (1,8))11 1.
K). If one
W]+ + (t) lim (()()))) (-000
Σ (t, (o) (t. (D) - b)) j=1
Similarly defining D0 (1,
Diag((0s,(1,))))
S j s v) and (1, der
VD (t, for k: = 2,3,
X.(1) B Σ X3
Mh PA, ()) h2(2)) (1, 8)) + (n(s)( (1, 13
(27)
B Tr (h
+ (2 (t, + ha(s) (tt
Therefore write Volterra equation for K. From now on, we consider specific initialization setting. Assumption 2 (Initialization). We assume initialization, 0, O and y-1 Yo- 0.
We thus can represent the expected loss function as a Volterra equation
forcing func.
E [ (+:) W]
X((s) ((,) [W] ds.
(28)
deterministic alg.
The forcing function (1) and kernel function (xs) are explicit functions of the eigenvalues the matrix K = D1/2WWTDl 1 where = Djag(j=20 1 J S v). In particular, we have that
(t) (0) (t. 0), (2)
and X,(t) def B Tr R2 (s) / 8) + (2) (t. 8) + ( s)))
with Px (t, s) ltt X1) X,) (t, such that t (s. Ids for all s I and $ 1/ (t, s)
Diag(( (2 5))) j S V)TT for k = 1,2,3.
(Exact ODEs) Volterra equation for K = D1/2WW'''1' = y''Vv Set the (u x v)-
matrix Diag(j-22 -2 1 S j S u) and (t) def (0,)) | b) Let
Φx, (t, s) for t 2 s be the solution to the IVP
d
di Px, (t.s) ((t; st, such that 2x (s, = Id, for all S S d. (29)
Then by Duhamel's principle, we have for every j
r(t; Ag) 2x, (t,0))(0;),) /'" (1, ds.
37

---

The expected loss under the iterates (ere-Mom-SGG) satisfies a Volterra equation
E [ = ((t) / Xi(s) xs W ds,
where $ (1. 0) (1W - b))
(t) def B Tr (k2 (1) (t, (2()) + h2 ()) (t.))
(30)
12 + 12 r
and 1(6)) -
(-7122 7173)
Here (t, = V Diag( (xx, (t, s))ık i V)V for k = 1, and h, are the components of the vector h.
While these representations are easy to see from the derivation of the Volterra equation, a more useful representation of the forcing function and kernel function is through contour integrals over the spectrum K. Let I be contour containing the spectrum of K (recall, K is normalized so that the largest eigenvalue of K is 1; thus I is contour containing [0, 1)). Then the forcing function takes the form
forcing function for algorithm
der
$ ((Ks =2- (2 (t, 0)) da
(31)
2Ti
and the kernel function takes the form
kernel function Kt(t)
$ (s) (
for algorithm
B x Tr
2T Jr
(32)
(t))s)( (K -2)- dz
C.1.3 Deterministic Equivalent of the Expected Loss
The forcing function (t) and kernel function (t) are random as they depend on the random matrix W. Moreover the expressions via contour integration show that both of these functions can be described in terms of the random matrix K D'/2WWT D1/2, Indeed it is the resolvent of K,
2)
which plays significant role in F and X. To analyze the power law behavior of the expected loss, we remove the randomness in K, i.e., the matrix W. We do this by finding deterministic equivalent for the resolvent of K, A (K, 2). using techniques from random matrix theory. Intuitively, want to take the expectation over the random matrix W; though not formally true.
Formally, we define the deterministic equivalent for the resolvent A (K, 2), denoted by R(z) implicitly via a fixed point equation
m(2) def
where R(2) def Diag mm)) .15i5e)
1 1
E 1j=1 fo
j
(33)
As mentioned earlier, this deterministic equivalent R22) can be viewed, roughly as,
E Ew[(k - E w E k, x)] vo R(z);
though it is not formally the expectation over W.
Using this deterministic expression for the resolvent of Y. we define deterministic expressions for the forcing function via the contour representation of F (t) in (31)
F(t) def -1 o (R(z), ( (D)/21) 62) ( ($_ (t, 0)) de
2τ f
(34)
38

---

and the kernel function in (32)
x((() def B x Tr
2 / (2)m(s)(00 + (2, (2)
(35)
(tt,) x R(2) dz).
Additionally, the deterministic equivalent defines two measure on the real line, which come from Stieltjes inversion,
Per (z)) def
lim =10
Im (RR( Im( + iz), (DD/26)88) 22, dæ
(36)
and
pxx (dr)
lim Im (Tr(R(z + i=)=²) dz.
=1o
Using these measures, we can define the forcing and kernel functions on the real line with the forcing function defined
forcing function deterministic equivalent/
g(t) / ((2x(t, 0))1 pr
(37)
and the kernel function defined as
kernel function
deterministic equivalent
)
(t(t)= = B / (In(s)((a(tss)) (e
=B
(38)
ha(s) (1,, 8)) (a)) px(dr).
Using the deterministic expressions for the forcing function F and kernel function X,, we define the deterministic function P R - R as the solution to the Volterra equation:
P(() 5(t++ (t) x P(s) ds.
(39)
Moreover, know quite bit about these two measure fox and jer from [80]. See Section D for details. In Section C.3, we provide some background information on solving Volterra equations.
C.2 Simplified ODEs and simplified Volterra
In general, quite difficult to analyze the exact ODEs in (22). Instead, as way to analyze the scaling laws of (Gen-Mom-SGin we derive system of ODEs using an SDE. These SDEs were studied [78] on a high-dimensionall least squares problem and were shown numerically to reproduce the learning dynamics of the stochastic algorithms (Gen-Moms Mom--MMm-G for variety of learning rate and momentum schedules. We will use these SDEs (and more importantly, the ODEs, denoted by simplified ODEs), to analyze SGD-M, DANA-constani, and DANA- in Section G, Section H, Section I, respectively.
Letting K W DW, we consider the following
dY; = -((()) - =00 + (3) (D
(40)
de; -e
(B (
where the initial conditions given by O 0n. Yo and (B.
(1) (2) :l> 0) are two independent
d-dimensional standard Brownian motions. We note here that V (0.) W D(W@, b).
C.2.1 Derivation of the simplified ODEs.
In this section, we derive a system of ODEs that describe the behavior of the expected loss under the SDEs (40). We denote this system of ODEs as simplified ODEs. This system will be used to analyze SGD-M, Sec. G, DANA- Sec. H and DANA-decaying, See I. To this end,
39

---

suppose (Ag),,))) are the cigenvalucceeigenvecor pairs of K and we write = Wb + b where WTD b b=0 0.
As before, we will consider the follow special statistics/functions for each eigenvector of K (t) = [(u (O; - -) W], (t = E | > and Xy(t) { (0, - b) Y.)
(41)
Instead of using Wick's formula, we now apply Itô Calculus. First, consider the functions IT def (wo, Or b), and 3 def (w), Yı). A simple computation yields for j .....,d.
B.) JB d(w), O: b) = (wj, (-78 - B2 (d (B
(ww, Yi)
Here (B (2)) 0) are two 1-dimensional Brownian motions. Applying Itô Calculus,
d((w), y = x (
(...), dt
d((w,, 0 - b))² aلل 2 x - -1
,2 dt.
As for the cross term, have
d((w), Yi) θ b)) d(,,11,) = -E,
-.E ( 2))
+ (I xy (B
Now taking expectations conditioned on W, we have that E [11]] W| o.. W|
E] and
E |[1,5,|W]
Xj. Thus,
03 222B, 2Tax; +
E) 25E + 21 + + BE W
Xj -383 B2 5x, + BX
Putting this altogether yields the following system of linear ODEs, denoted by Simplified ODEs.
Simplified ODEs for 23 Xj- X3- Let (A,, be the rigenvalue-eigenvectoo pairs for K = W DW and decompose b Wb + b such that T Db 0. Here the (e - D Diag(j-2e 2a
v) and (1) del ((.)) [ - b)1la. Additionally, we drop the time in the learning rate and momentum schedules to simplify notation. The functions
(1 = E (uag,, (0, - 8) 02 [w]
()) = E( ,y.,,,,,")1 y,
(42)
and (1) E ((w (0. -00-b) b) Y]
form closed system of linear ODEs:
d
Te u1 Ay) = 2((11;p) x x((1),) + }gb))))
(43)
40

---

where
(22yBA
-273
A(t;;,) de
0
-2A
2n
?12 BA,
73
- B2
(44)
T BE
(L) W]
gt: Xs)
def
7 7 BE (P) [[]
and
r(t; 2U) der
())
(t)
Under Assumption 2, the initial conditions are such that
p2 (0)
Kug 8002 W]
((:;A)) = (0
(1)
C.2.2 Simplified ODE as the large time limit of a coin-flip algorithm
In the previous paragraph Section C.2.1 we have seen that the simplified ODE (43) models the risk under particular system of ODEs (40). Although precise, this characterization has obvious limitations. Indeed, the SDE formulation requires the learning rates to vanish, with a correct scaling with the time and dimension. This hence does not model a practical algorithm which makes any discussion about compute in the scaling laws.
Instead, in the following show that we can also view the simplified ODEs in (43) as dropping terms that arise due to the higher-order moments of the Gaussian data x in the ODEs of a coin-flip algorithm.
The coin-flip algorithm a simple two-staged version of (13) where each iteration we flip coin and update either the momentum iterate Y or the gradient update ê. Let NJ¹ No be iid Poisson processes (i.e, the coin flips) and let the initialization be such that O, = An and Yo- y-1. At jump of NT, Nř, we generate (x', ...5)) _ new data points whose coordinates follow a power law with parameter a and update by
y= (1 -1-111111 A(t)) +1(()) Σ WT x', (fx,)) b) def Y 0(-Y)_,0i-)
B
9 = 0, - 72(t) x Σ W (xx, (WO,, -7s(t)
- x Y, def O, + 2! θ.-).
(45)
We naturally extended time I to be a continuous parameter with Y, = Y, This is standard way embed discrete processes into a continuous process.
For any (9, 1) + ((,,)) E R, the mean behavior is given by E [f 0.)] Elf(Y - AY 0.), 4, f(Y, 0,)]
(46)
dt
In this paragraph, we derive the system of ODEs that describe the behavior of the expected loss under the coin-ffh -lip algorithm. For this, need to introduce statisties that are closed under differentiation as defined in (46).
Derivation of the coin-flip ODEs.
As previously, we begin by writing R" Im (W) W Thus, there exists a ò E Rd and i R" such that one can write b = wh + b, that is, we can decompose as an element in the image of W and an element in the co-ker of W. Formally, we have that
b = wb + where T Dか = 0.
Indeed, denote E Im) 1. with b b- -bbe Im(W). Then we are looking for some E Im(W) + (b) satisfying w = 0 + Db E Im( + Db E Im(W)= + D Dm(W), Hence such R This is clearly the case since dim Im(W)) = dim (Tm (T)())) = # dim dim(Im(W)t) and because D never sends vector x on a perpendicular one, ie Vz E R", Dr) + 0 since it is diagonal with strictly positive eigenvalues.
exists in particular if D Im [m(W) Im(W)
41

---

Letting K def W DW E Raxd and ()), by the eigenvalue-eigenvvetor pairs for K, we consider the same special functions as in (20).
(3()) = El (0, b W1, () = E ERu , W],
(47)
and X²(1) E [(w] 2 (o, - b) k.)WW.
As before, knowing the pJ's suffices recover the expected loss under the coin-flip algorithm. We proceed like in Section C.1 using Wick's rule to get a closed formula for (0}. 53 x;).
P² functions: Using (46) applied to E [f(O,, Y)) = (t) yields.
B
(1) = 2F [ (0) x
(1)
Σ
WT (())( b) - (a(V)) @ (0, - -8) w
+ +E(W (12(t) x Σ
W W x (We,, -)) 0)) |w] + (t) [wo y, W]
( 1 ( (- 1/2
W " (WO,, - b)) @ (-Y)))
--0 22
+N( W| + B(B 1)) (t) +33(0)3(( (21
functions: Using (46) applied to E [f(,,, Yi) = ()) yields,
B d ()) = E [(w ((1 - A()) + (() Σ WT
-100.--11 0) ww]
[[ Y y [w]
2E 1Wg2.
V. 8 A(t))) W| +2E 22 Y, ( 2, (we, b))) [W]
E [ (A(t)))) W]
- -BE (Wg2 A(I) @ (f)
2.-, (We, (WW 6))) w
+E [W] (W
- b)) )2) [
(WR b)) @ (W)) -))) [w]
= -2208800-200000022---888000000-- -2000000200000000000000000100100 (1())))))
( b)11 W| B(B - 1))
Xi. functions: Lastly, using (19) applied to E [f(,,, YJ] x, yields,
42

---

B
# E w| = E [w (o, - b) 8 ((1 (U)) + (t)
Σ W (we: [w]
- - (0, - b) 0 Yi) W|
+ E |wg . 1/2 W [W -)))((())) - 1)) [
- E ER (0, - b) @ Yi) W
-A(t)E [W/ (0) b) b) @ Yi) W
+ BE [w08? (0,
- []]][W]]]]]]]]]]]]]]
-61)
-12(t)BE [wu) (WTEF" (W' - b)), Yi) W] ra((EE [Wo)] .x2) W
= (t)()x(( +
(1)
Putting this altogether yields the following system of linear ODEs.
Coin-ftip ODEs for pz. Xj. Let wy) be the eigenvalue-eigenvector pairs for K =
WT DW and decompose b = wh + b such that W Di = 0. Here the (w x ---matrix
D Diag(j-20 2a,
j j 5 v) and (1) def ((O)) = -
(1 = E [ (o. b) ) |W], {)(() - (, [W],
(48)
and (1) E [(m) (0. -00-b)
b) @ W]
form closed system of linear ODEs:
d ((t: by: = Sit:: x x(tttt))))))))))))))))))(()))x g(t; g(t
(49)
dt
where
BA,, B(B + 1))(())A3 B + 1)15
T5(t)
2ta 21
S2(1; 21)
(1) + A(1)
2% (1) A(1)
(x)
(1)
-A(1) (L)
(t)(t)
g A,)
and rtt: A1)
(t
0
(l)
(50)
The initial conditions are such that
,(0;A,)
We see that dropping some high-orders terms in the eigenvalue À def o2 and learning rates//moentum schedules from the ODEs (22) yields the simplified ODEs in (43). This is the reason why we believe that working with (43) should not affect our main results. Indeed, when studying the asymptoties for large time the contribution to the risk of the different eigenvalues will come mainly from small cigenvalues o² (and especially o² Bi ), as for larger a's, solutions to the ODE have exponential decay.
Simplified ODE as a high-dimensional limit and re-deriving the SDE system (40) without Itô calculus The goal of this section is to relate the solutions of the simplified ODE Equation (43) to limits of solutions of the ODE eq. (49). This allows to re-derive without Itô calculus, some positiveness result on the sign of the simplified ODE solutions since we know the sign of the coin-flip ODE solutions.
In the following suppose that 71 (t) and that 72 (t), (t),,, A(t) are three continuous functions.
43

---

(t)
(t)= def
m Bo (t)
raff)
Do the change of variable
on Equation (43) and get the new ODE
Pa(())
vn Bo
Bo²
0 -2( + da(t)
-2
( dt
2 2/7na
(t)..
(51)
(a(e)
Viiaa
1371
-22 - A(1)+ 2Ta(0]
Lemma C.I. Let 71 (1) = 1, B, > 0, and (,()))na!),,(() be three continuous functions with
rr()) O. Let X def 0², 8 > -1 and M, E (R)x Denote (-1, oc) - ~Maxa (R) the
solution of the simplified ODE (51) with initial condition for - 1, (s) = M,. Let d >> 0 and
denote (() the solution of the ODE (49) where we used Á(1) def A(t/d) (1) relt/d) for E [3].
4 (dtt
Additionally denote (t) def
10 (dt)
(t(())
and write the initial condition ((s) Mg. Then
(l (ddl)
VT > 1,
((()) )(() = 0
sup (E(a.T]
Proof. Denote K
B+1 B
E]1, 2]. We obtain
-22
(000
(-1+ 2=(()
2
db(1 d d2
(t())ra(t)
2RO B
(()) A(t)
ts(t)
2 (1)1)
0 A(t)
(+))
d
de
Ts(()
d
(1)
v 7500m1008
Ait
fx(t)
d
d
Bo²
2 (t)t)
-2((1) +
Sul
2 (1)
+ +(1)
(Tolt)
(1)
200
(1))=2(()BB'
5.0
2
which implies the result.
Remark C.3. Lemma C.I shows that the simplified ODE (51) is "high- niigh-dimessional" limit of SGD, where the learning rates are scaled inversely proportionally to the dimension with the correct scaling. was done in [78]. Hence one will note in the following the proximity of our results results from 178]. However, we underline that our argument in favor of using the simplified ODE (51) is because this gets rid of higher order terms in the eigenvalue a, whose for large time I. Indeed, we avoid taking the learning rates to vanish for large dimension, as this would make the conpute-optimat scaling laws useless. Hence it is surprising that dropping higher order terms in the eigenvalues a (looking at the dynamics for large t) is equivalent to dropping higher order terms in the learning rates 7/,2,3 2.3 (looking at a deneeegee,,,te, high-dimensional limit of the algorithm for fixed time t). We leave the study of the links between both for future work.
believe the contribution to vanish
Corollary C.1. Under the same conditions as in Lemma C. denote V : (-1, o0) + R the solution of the simplified ODE (51) with initial condition for 1, (s) = Diag(a (a1, 0z 0z) with a, > for i E [2]. Then, Vt > 1, Vij 2 for j.j E [2].
Proof. For any d > 0, consider the solution (t) of the ODE (49) where ((t) def Att/d)
(t)
doBB do /B ). We know that x(t/a) for i E [3] and with initial condition (ss) Diag(a1, 02 71th ag /
(t), (t), follows the same ODE as db. Hence, if the initial condition is admissible, we can deduce that the first and second coefficients, being squares, are positive. We know that (s)(s, (s) are admissible which directly brings that VI -1, (4). (1), 221 (1), &22 (t) (t) 2 0. Then we can pass to the limit-inf when d + by using lemma C. to get the same result V.
D
44

---

C.2.3 Derivation of the simplified Volterra equation
We follow the exact same ideas in Section C.1.2 replacing the matrix ( X) in the exact ODEs (23) with (tt A,) from (44). First, we Φ2, (t,s) s) for 2 s be the solution to the IVP
d
dt P, (t, (tt X,) (t, such that (s,(5,5) Id3 for all 1 S d.
(52)
where (tt;; )2) defined in (44).
Following Section C.1.2, we represent the expected loss function under the SDE (40) Volterra equation following an application of Duhamel's principle, that is,
v A1) = 2x, (1,0))(0; A,) + S'" (x,((,)))(1; 22) ds.
(53)
Note that in this case the function (s)) simplifies,
h(s) +
(Simplified) Volterra equation for K 11/2 D1/2 VTAV. Set the (u v)-
matrix D = Djag(j 2a 1 ك S v) and (1) del ((O)) (D)/2(W,, b)] Let 2x (1,, 8) for 2 s be the solution the IVP
dt $ (t.s) O(t: st such that &x (s, s) = Idg for all و S d, (54)
where (2tt X1) is defined in (44). The expected loss under the iterates of the SDE (40) satisfies Volterra equation
+ / x((s) (({() W| ds,
E
where (1) del (0) 0), (D) (D)/ - b))82).
(55)
and X,l) del B x Tr( K2 (1)(() K (1, + (1) 8)).
Here 1 (t, V Diag(( x, (t. 9) 1 - T for 1: 1,2.
Deterministic equivalent for the SDEs. As the (simplified) forcing and kernel function only depend on K, like the (exact) forcing function (31) and kernel function (32), we use the same deterministic equivalent of the resolvent in (33) and the measures, μer and μoe in (36). Therefore define the deterministic equivalences for the forcing and kernel functions of the simplified ODEs on real line as
(simplified) forcing function deterministic equivalent
F(1) / (1,((, (dx())
())
(56)
and the kernel function defined as
(simplified)
kernel function )
(t) = Bx (1)()) s)) ( s))2 ux(2x).
(57)
deterministic
equivalent
Using the deterministic expressions for the forcing function F and kernel function K, we define the deterministic function + as the solution to the Volterra equation:
P(t) = F(() + S (t) x P(s) ds.
(58)
45

---

We know quite a bit about these two measures μx and jor defined in the forcing and kernel function from [80]. This knowledge allows us to derive scaling properties for the algorithms in Section B. See Section D for details. In the next section, we provide some background information on solving (general) Volterra equations.
In the rest of the appendix, we will use this simplified system of ODEs to analyze SGD-M, DANA- constant, DANA-decaying Section G, H, I.
Assumption 3 (Siimplified Forcing and Kernel Functions and Simplified Volterra Equation)). From now on, we will only work with the (simplified) forcing function and (simplified) kernel function as defined in (56) and (57). The same holds for the corresponding (simplified) Volterra equation (55). We will drop the reference to simplified in the forcing and kernel functtons and Volterra equation from now on.
We finish this section with remark.
Remark C.4 (Thm. 4.1 (({))) vs. F(t) this section.). We note the following connection with Theorem 4.1 in the introduction, which we hope does not add too much confusion. We have
(Theorem. 4.1) ((())) corresponds to F(t) (Section C, (56) Appendix).
In particular, the function F(t) (56) derived in this section and considered in the rest of the appendix,
in an asymptotic sense, includes the (t) defined in Theorem 4.1. For K(() def Ko(t), it is a slightly more complicated correspondence, here
(Theorem. 4.1) {()(t)) corresponds to
1 X(t) (Section C, (57) & Appendix),
7²B
where y is defined in Theorem 4.1. This holds true also for (t) x (pp(t, s) which we will define properly in the next few sections.
C.3 Background on Volterra equations
While Volterra equations such as (58) are quite nice and well-studied in the literature (see, e.g., [45]), we need an approximation of the solution to them in order to have a better understanding of the scaling laws. In particular, we need the (deterministic equivalent) loss function P(t) to be constant multiple of the forcing function T and kernel function K. We state this idea more precisely at the end of the section.
To do this, we need some background on general Volterra equations of the form:
P(t) = F(t) + (K * Pt),, where (K * *P))t) - - K ((t,)))(s)))))))) ds.
(59)
and where F()) is non-negative forcing function and ((1, s) is non-negative and monotonically increasing in second variable and monotonically decreasing in the first variable kernel function. In general, we define (K (t(t,,) s) = K(t, K()r))(rr) s) dr.
Let us define (tt s) def (K = K = * K * K) Kt, s), the n-fold convolution of K where K
times
K(t,s). s). Under mild assumptions such as |K| def suPre(o, So K (t, ds < 1 11 (we call this quantity the kernel norm) and the forcing function F is bounded, there exists a unique (bounded) solution P(() to (59) and the solution is given by repeatedly convolving the forcing function with K (see, e.g., [45]),
P(() = F() + ΣK' * F)(1 = F(1) (K * F)(1 + (K * K F)() +
+ (K K K F)()++
1=1
This representation of the solution (59) enables us to get good bounds on P()). More precisely, [45, Theorem 1] shows that if K : R, x R2+ + IR is continuous, ((1, 8) = 0 for L > a and F bounded, then solution to the Volterra equation is written above. [45, Theorem shows that if additionally lim supr fo K (t, s) ds then the above solution is bounded.
"In the case that K is of convolution type. that is K (t, s) can be expressed as K(t - s), this definition of [[]] simplifies. Particularly, define ||]]] Jo (tt) dt.
46

---

C.4 Reducing the complexity in the Volterra equation
First, we state and prove lemma attributed to Kesten's 's Lemma [8, Lemma IV.4.7).
Lemma C.2 (Generalized Kesten's Lemma). Suppose K((t, 8) is positive kernel. Suppose that for some constant C(K), for all t 2 r 2 0:
[ ((s)s)s S C(K)K(tr))
(60)
Then, for all n 2
(t,tt))ttr) K(L,r
sup 12r20
S <C(K)*
Proof. We follow the proof in [80, Lemma C.1]. Define
sup K
def
an
2rr2 K(t, (((K)))(
Then it is immediate that a1 = 1. If we can show that
On S 1,
then we are done. By definition of the operator we have that: (t,rt
C(K)a L. (t, ()t,s)
(t, ((K))) ds S on [ "((t,))K(s,) ds.
C(K)
C(K)
By the hypothesis (60),
(x,at C(K) S On ((,T))
fort 2 r 22. 0,
Hence the result is shown.
Often we will not be able to analyze directly the forcing and kernel function, F and K resp., in the Volterra equation. Instead, we have access upper and lower bounds on F and K. Using these upper and lower bounds (together with Kesten's 's Lemma). we can still give non-asymptotic bound on the solution to the Volterra equation.
Lemma C.3 (Non NNn-assyppoote Volterra bound). Suppose the same conditions as Lemma C.2 hold. Suppose additionally that C(K) and F(t) < F(t) S F(t), K(t, < K(t,s) S K(t,s) with non-negative F. K. Then,
E(I) + IK E)) < P(t) < F(I) +Cx F)().
where C. C-C(
Proof. This proof parallels the proof in Lemma C.2 in [80]. We include the proof for completeness. We consider the upper and lower bound separately.
Lower bound: Since K and F are negative. Σ (K*) Et) (K*) Et)() (K * (t). Recall the solution to the Volterra equation takes the form.
P(t) F(t) + Σ (K F)()).
1=1
It immediately follows from E (K*) (t) (K (t) the lower bound.
Upper bound: By Lemma C.2, there exists C(K) < such that
(t, < (
47

---

Hence, we have that
0c Σ (K F)()) - Σ (t*(t,))F(r)) dr - Σ C(KY (K F)())
80
;=1
(1 (I-C(K) (K (()))
The upper bound is thus shown.
A main tool for analyzing the deterministic equivalent. We now state of the main tools used to analyze the deterministic equivalent loss function (58).
P(() = 5(() + ! (() x P(s) ds.
For each algorithm, we will show that the forcing function F and kernel function (() have upper and lower bounds, that is,
F(t) l(t) (t) and X,(t) (t) (t).
(61)
Moreover, we will show that the scaling laws for F and F (same for (1)) are the same. Therefore, for each algorithm separately (with the expection of DANA-dec,ying), we will show the following
Informal Theorem. [Reduction of the Volterra Equation| Let K(t)) def Ko(() and define simi- larly JC(t) and K(t). Suppose 2a + 233 > 1, a + > B. and a > _ For various assump- tions specific the individual algorithms, there exists an M > large enough and constants C(a, B, M, alg), č(a, 3, M,, alg) such that if Bt M then:
ēx x [(t)
BB x(t) S P()) S (3()+
x(1) + BB
(62)
def We prove versions of this informal theorem, see Theorem G.1 for SGD-M, see Theorem H.3 for DANA-constant, see Proposition F.9 for SGD with decaying learning rate schedules.
We denote 722 + 73 for classic momentum and def for both DANA algorithms and SGD.
8
For the upper/lower bound on the kernel and forcing functions (61) for each algorithm are given in Section F for SGD, Section for SGD-M, Section H for DANA-constant, and Section I for DANA-decaying. The general asymptotics for the forcing and kernel function can be found in Section D.3 and Section D.4 with specifics for each algorithm in the individual algorithm sections.
D Measure of the Deterministic Equivalent
In light of the Informal Theorem above, to understand scaling laws, need to analyze the (up- per/lower bounds) forcing and kernel functions under the deterministic equivalent for K. We recall the (simplified) forcing function, F(() defined in (56) by
F(1) = /
( ((=()))))
and the (simplified) kernel function, (L) in (57)
(t, = B x
(1) s))11 + (2)) +)) pxx)))
where the measures (dxr and pax (dr) are defined in (36) via the Stieltjes inversion,
Mo (d(dx)) lim Im (m(((( - ie), (D) dr
#10 म
(63)
1 Tr(22Tr(R(z + iE))) dr.
and Pxx (dx)
lim e00
W'hhile we need a > and a + 1> 3 formally, we believe they are only the result of the proof technique and not necessary. We believe the result would hold without these conditions.
48

---

1.5,
v- 2000
ef d5-555 y-2,000
22:500 vs2,500
101
101
124
107
density. A22
dersity.
20
dersity.
107
121
10
8pe -2. 10-
1831
cigenvalues
eigenvalues
(a) Deterministic equivalent vs. em- pirical
(b) jer as / changes
(c) per as changes
Figure 11: Deterministic equivalent of the forcing function measure μig. Numerical set-up: 100 randomly generated K = DWWTD D and the pj computed; for empirical density μg,, 500 bins equal spaced on log scale from 10-8 and counted the number of > that fall into each bin weighted by Pj and then averaged over the 500; For deterministic equivalent μs, solved fixed pointed equation (33) using Newton Method on grid of æ-values. Fig. la: deterministic density of My (63) under the deterministic equivalent for K (red) matches the empirical density of pay. Fig lb: slopes of the densities Mı for all / agree the left edge, but the slopes diverge the right edge as / changes. Number of 'point masses' on the right edge are the same for all B. Fig. llc Left edge cutoff evolves like d-2a illustrated here. Number of "point masses' and slope of the density change as a changes. These slopes effect the scaling laws.
Here R(x) is the deterministic equivalent of the resolvent of K, given by the fixed point equation
m(2)
where
R(2) = Diag 5-20
1 2j=1 m=)==
For verification that the deterministic measures in (63) match the behavior of the random matrix K, see Figure 11.
We begin with first technical lemma stating that will allow us show that μrg px put mass on intervals [1 + oo) for d large enough.
Lemma D.1. Let > a + Y/ and 0. There exists d 0 such that Vd d admissible has real solution and any z E + e, oo). the fixed point equation m(z) 1+ Ei=ı 3- m-
20
m(2) [1 -
Proof. The proof boils down to showing that for any 2 E [1 + e, oo), the function G(m; =) def
1+ Eÿ= ---m has fixed point in c/2, e/2] for d large enough. Clearly for all dl >
G(m; 2) is continuous m [1-e/2, c/2]. We write Vm [1-c/2, 1 1 [1+ +00)
2a
20
S
j
2a
m
c/2
==1
(C((a)
(0.
The above be made arbitrarily small for large d. Hence, by definition of G(m; 2), this brings the existence of d(a,, e) such that Vd d, Vz + e, oo), G(m; is continuous self-map on - e/2, + e/2). Hence admits a fixed point in [1 - e/2, e/2].
49

---

D.1
Estimating Arr
In the following section, we provide upper and lower bounds on the deterministic equivalent measure of the forcing function μor to compare it to reference measures Poo Pgact py so that integrating it against any function f can simplify as
[EE. f(o) d =0 I JERR
f(o) d(ps, + Ms ps (o²)
JERR
This is formalized in Proposition D.13 below.
We begin by noting for all E [0] that we can define 13 ( (dz) by Stieljes inversion, i.e. Va E H,
(dx)
del
j -22m(2) -
I
(64)
or equivalently ( (dx)
def
lim +10 π
Im
dr.
j-20m(x + +ic) - (x ic)
We therefore rewrite Per (dx)
ビ 20 1ig (dx)..
Now let's notice an important identity, which will be very useful to get upper and lower bounds on Mgr. That is for any + in)xx) E IH can rewrite
Mig (ds)
1
Im j m + in(2)) (x in(x))
Im
π
(x in(x))
(ds)m(a
(s x) + 11)
(Pois,(2)) *py (z).
Hence, estimating the left-hand term allows to evaluate the result of the convolution between ply with Poisson kernels. We can hence use the bounds from [80] on specific contour in the complex plane to estimate μr.. For that, we will need technical lemma comparing indicators on segments and specific combinations of Poisson kernels.
Lemma D.2. Let C,C, a
> 0. (log(1/c)/0 maxful-1/220 x a a'-1/(m)) and denote P(u) def
Consider for u E [d-2a, 11. rπ(u)
def
Poism(e) x) dr.
2a d
There exists a constant Ö(C, c, a) independent of d, v, such that VC E ld- 20 (i)
we have
2
Vs E RR (c) SO ( Pois (s - u) du
(ii)
(i [C, P(u)
and if [0,20] P(u) E Emin (1,
xLC) ) with d(u) def d(u, [ d(u)
4C BC)).
Proof.
We can define def inf inf n(u). supuelC.201 n(u) and notice that there exists constant Cle, c, ar) such that
First note that Yu [C,2C], n7u) S C.
def
(f 20 Pois u) du
AC > SYS E R, (1):00(()) so
From the definition of v7(4) which follows power law, we know that & 21+1/(2a), Hence this brings that Vs E [-C, CJ.
sup.elu.6 Pois, (s) e.c.M.a
Int fe (o) S
1. This proves the result.
50

---

iii Similarly, we know that for potentially different constants C(s,,, we have
(LC
20
AA > OVSER,
Poiso(s - u) du SC.
This brings using infgelam) (8) notice that for potentially larger C
( Poisy(s)
s.M.a.
S
1 that [C, P(u) S Ć. Additionally,
4C 5C 3
(o SC ((C)
Vu & [C, CC). Vo E
[CC 50 d(u, 3 D)
we obtain that still for a potentially larger Č
Vu (C, P(u) Cmii (1,
((CCC d(u)?
40 50 with d(u) def d(u, 3 ))-
D.1.1 Upper bound on pa
We begin by proving the upper bound by decomposing the real line into different ranges of 'singular values', a's: large ss a = 0, small o's, and intermediate a's.
(1) A first observation is that Poj is probability measure.
Proposition D.1. Let a > 0. For any E [u], Pig is a positive measure and U (R-) = 1.
( Proof. By definition, PT (dx) lim Im ( dæ. Since for all E IH, m(=))
defined as the fixed point of (33) with negative imaginary part (80, Prop. E.1], is clear that 1o positive measure. Additionally, since the support of plz is bounded and ==m()))= ii
we see that M ( has mass 1.
This implies a first bound on the mass μay for large a. Proposition D.2 (Upper bound for large a's). Let 0. Suppose 2a + 23 1. Then there is constant C(a, 3) such that VM > O, we have
y ([M, S C(a,, 3).
Additionally, for any > 0, there exists some d > 0 such that Vd 2 2, pe ([1 + c. +oo)) =
Proof. We know from Proposition D.1 that
Ms Ms(R) (R) = Σ5
2 2a 28
1 M3
(R) = Σ5
2a-28 23 < 00
==1
1=1
since 2a 23 > 1. Finally, the last claim is direct consequence from Lemma D.I.
Proposition D.3 (Point mass at 0). Let a > 0. Suppose 2a + 2, > 1. Then has an atom at and there exists some C such that
2a 28
(t((())
$ Sdd Cd-22-+(281))
1
2a Ks (v/d)
where k(v/d) solves fotd 6tuz du = 1. In particular, we know ps (o({)) x 22
51

---

Proof. This is direct consequence of [80, Prop. E.3, F.1, H.3] and the weak convergence of the Poisson kernels the Dirac
Proposition D.4 (Upper-bound for small d's). Let > O. There exists some constant c(a) independent of d such that
Vx (0, coodd2e 2a), pr(dx)
Proof. We know that jeur (dx) lim Im (m(((( + ie), (D1/2)) s2) dx. Using [80, Prop. E.3), know the existence of c(a) 0 such that m(=) is analytic in ((0, c(a)d and
Vr (0,ed 20), m(x) <0
which implies the result.
Proposition D.5 (Upper-bound for intermediary d's). Let a > 0 with a + 3. a a 1/2 and B # 2 with 2a + 23 > 1, 2a + 3 + and 2a + 1 > B. There exists M, C > 0 depending only on a, B
such that for any C' [d M, M
NJ pJ (C)) SO (cu+ +
Cg ct) d
Proof. Using lemma D.2, only need estimate ブ 20-221111111100
23 Im m=)-) where z
in(u). Applying [80, Prop. E.5], we get the existence of c, M
0 such that for any
u E [d- M,, 1r and with r7(4)= = (log(1/c)/a max max(ul+1/(2))
a u 1/122))
we have
2n
1 1
m(z(u)) - (1+ 8,) + i(1 52) u a/(2a)dd with (2)
2a
3 3'3
Applying the summation lemma [80, Prop. E.4] we obtain for E [d-20 20 M, M and some constants Č, C of E,
Im Σ 2,5 1
( 2(u)
)(u))
20(1
20
C3 u ia + com(u)u
π
d
$ δ
20_1
CB
1)
u
+ u
d
Here, we used 2a + 1> 3 to get (og(
S u
28-1
We conclude by writing
/
v=([C,2C1)
JR.
((c:2c(())pxdds)
50
Pois (8 - u) du
per (ds)
R+
20
(Poisn(w) *Hr) (du)
c
σē 2c
28-1
+ C3
2 du
d
s Ô
(0
22
Cg c²-
There some intermediate a's that still need to be bounded, x² E (c(a)d-20 20 Md-2aj. This done in the next proposition.
52

---

Proposition D.6. Let a > 0 with 2a + 23 > 1. Then for any > > 0 there exists constant C( C1,1)))))))))) 02) such that
po(ledd 20 , s C(a,,, 02) 2a+(1-28)+
Proof. We use the estimate in [80, Prop. E.3] using points 2 ((c1, 02] +i)d and the upper- upper-bound bound oer I Č Set Pois (x s) ds for some C(c,, oz) and all R4.
D.1.2 Lower bound on jo
To obtain lower-bound, we will the lower-bound the indicator [[ using Poisson kernels proved in Lemma D.2. This introduces some error terms due to the tail of the Poisson kernels that need be controlled using the previous upper-bounds. In particular, we will need r7(u) u, for this bound to be non-vacuous.
Proposition D.7 (Lower bound for intermediary o's). Let a,, > 0, a 1. + 2. 2a 2a > 1, a-
1 > B there exists some M > large enough and C(a, M) such that for any C E [Md-26 M
1 c1 |=((C,2C]) 2
28-1 20
Cp C² 2/
+
Ö
d
def 55033 Poisn(z) (u x) dr where x E [d-2a M, 1
Proof. We consider as in Lemma D.2 P(u)
JC/3C3
M
and r7(x) (log(1/c)/a maxxf=++1/22) x² 1/(2w)
with e, c, M given by [80, Prop. E.5]. We
write using some constant Č' from Lemma D.2
5C/3
Cpr((C,2C)) 2
Im((s + --v (1 m(x + in))d)) - /ue|c20n (du))
4C/3 π
We hence have using some constant Č that
55/3 Im
j g-2a
20 28
2a
Cpr(1C,20) 2
ml ( + in) - (x + in)
40/3
P(w)pos(u))
(
20
+ a 0 2a
( P(u)ps(du)
S W=20
P(u)po(du)).
We need to upper- bound both integrals on the right-hand side using Proposition D.5. For the first one we decompose
2'c
((u)pa(du)
max
P(u)
12-000 pz(du)
ue(2-1-102-101
log(dRn
(CCC
(2-10) 2 + 1
C
1=1
(C)
c+
c²
28-1
0 2x > 1 )
+
C
d if 2a < 1
53

---

Here, we used Proposition D.6 for the left edge and, in the second line 1 + check when 2a
28-1 2a
> 1. There is left to
that
n(C),
d 20 v 2 2á-1
1
+
C
d
C 2 a and C1/(22) 20 S -0 20-1 (-
C S do and ي d- < C++ + C1-((1)1
The second case is always valid but the first case is true in particular when B < a 1. For the second integral, we similarly expand
SSc ((u))
log(1/C) E
P(u) x f fx00
'c
max
2-1--1 1 uo(du)
wE|2-C,2'+10
07
loRl +)
MCCC
(2 29-1
+ (2)))))-
(2'C)²
(C)
7 7
26-1
+
cI
1
C
Here we used Proposition D.2 for the right edge of the integral, 3 < 1 + a for the first term and reduce the less restrictive condition _ 2e
2
0 for the second term. In fact for the first term,
B < + a: by noticing that ((C) $ (c++ 23-1
under this condition.
C
D.2 Estimating Px
Similarly to the previous sections, we will now upper and lower bound μax prove in Proposi- tion 1.14 upper and lower bounds for the integral of functions with respect to μox.
We rewrite, using [80, Lemma E. and dropping the -v term which vanishes
pio (xx( def 1 lim Im ((x il m(x + ie)) equivalently
fo
Hex (ds)
(x((s)
2 =(1-m(z))d)
<00
(65)
An important equivalent identity is to rewrite for = + in((x) EH,
Tm(=((1 m())dd)) / 00
pxx (x(()))(()) (s - x)² + n1z)2 = (Pois)(2) *pxx (x).
D.2.1 Upper-bound on μo
Again we decompose the real line into various ranges of o'ss Using the upper-bound of the indicator function using Poisson kernels, we similarly get:
Proposition D.8 (Upper bound for intermediary o's σ's). Let a > O, a 1. a + 1 2. There exists M, Č depending only on a such that for any C E [d-2a M, 1.1. we have:
Pxc([C,20 2C1) S Ò C d.
Proof. Using Lemma D.2, we only need to estimate 1 IIm(x + in) m(x + in))d) for x [d-2a M, 1 and nT()) (log(1/e)/c) max(2=+11(2a)) - xt 22122))))
J. Applying [80, Prop E.5]. there exists c, M > 0 and c 0 sufficiently small such that
2a
1
Im((x + (n - m(x + in(x)))d) = 2a
1=11((2a)
+ O [()))((++ +n(x))())()))))))))))))))))))))) +ex
ex. 11/12)))))
π
We use the bound on A from [80, Prop.E.4] (when setting 3 = 0) to conclude.
54

---

2/ VM > 0, u .M) z 1. Additionally, for any > O, there exists some d 0 such that Vd d pxx(11 + c, +o0)) 0.
Proposition D.9 (Upper bound for large o's). af
> 1. *
Proof. Consider + i E H for x E IM M). This defines a compact U of distance at least 1 to [0, 1]. From [80, Prop E.6], we have that
C(a)
-
Im
2
2a
+ +01
1 minii, dd,dda-1)
j-2a
2n (j--20 x]
x
+1x
2a x)² +1
( 2a - x)2 +
C(a)
+ +
1 x min (d, dia-1]]
j-
C(a)
0
20 0)3 +1
minin dta-1)
The above being bounded, it implies using Lemma D.2 that pox (1 M)) 1. Finally, the last claim a direct consequence of Lemma D.1.
Proposition D.10 (Upper bound for a's near zero). 3c(a) > such that (e, ((a))d)) 2a() =
Proof. Using [80, Prop. E.3], we know that m(=) is analytic in B(0,, c(a))c(2a))) 20). Hence, from (65) follows that poe is null on [0, c(a))d22))
now move to the last range for a's: 0 with 2a + 233 Then Vez
Proposition D.11. Let a pxc(ccdd220 ,2d (22d-220 201) S d da-
C1 > 0, we have
Proof. We use the estimate in [80, Prop. E.3] using points = E ((e1, (2]+ +i)d and the upper-bound (x) < Ò Sa Poiss, (x) ds. This brings that
20
ux((c 2a c2d 201) x
(Pois) (s)
d Im ((s + i) - m(s + i))d)
d--4a
Im ((s + i) f(s + ())) + Old
< dl-4a
D.2.2 Lower bound on μox
In the next proposition, we prove lower bound on
Proposition D.12 (Lower bound for intermediary a's). Let a > 2. # there exists some M > 0 large enough and Ö(a, M) such that for any C |Md-2x 1.
2 (2-1/(2n)
px((C,2C1) >
55

---

Proof. We consider as in Lemma D.2 P(u) def M, W7 and (xx) = (log(1/e)/e) max(x++1/(2a) Prop.E.5]
Ja=tC/3 Poisn(u)(u 2) dr. Again for
I x² 1/(2a) with e, C, M given by [80, 20
d
We write with some constant C from Lemma D.2 and potentially larger C
50/3 1
+mm((x + in)( (1 mm(x + in)))d) dx /
C)x(((C,C)) 2
P(u)
Jugic
#0// π
(2-1/(2a)
(1 ((w)pxc(du) + S J
P(u)ux((du)
need upper-bound both integrals on the right-hand side. For that we use Proposition D.8. For the first one we decompose
2 px
P(u) x
P(u)poc (du) &
max wE/2
Cc2 ic]
2-11-1
(C)C
(C/3)2 x (2-10)²-2
(C) C2- +o
3 M C- +,
Here we used that 2
for the sum to converge. For the second one we write
2n
log(1/C)
2
((u))xx(du)
max we(2:C,2i-1C)
P(i)
px(du)
2C
2C
(CCC
(C2:)2 x (2
-2-#
5 M 11(2ca) 2 +
All this finally brings that for M large enough, le (1C, 2 - for some Č depending
6, C, Oz.
D.3 Forcing function
In this section, we decompose the measure jur based on the upper and lower bounds previously derived. This decomposition allows us to break the forcing function into three components, (t).). Jpp (t), and F (t), that is
F(I) = Fpp (1) + F F1c Fo(().
We will discuss this in detail.
Proposition D.3 gives an explicit formula for the behavior of the measure frr at = 0. Proposition D.4 gives the "gap" between the 0 eigenvalues of K and the next smallest eigenvalue which occurs c(a) d 2a where c(a) some explicit constant. Moreover Proposition D.2 says the largest eigenvalue of K approximately This is all to say, the support jay is (((a))( 2a ++c(a)), Lastly Propositions D.5 and D.7 can be interpreted as saying something about the density for pg.
56

---

((t) x 5 pp + (o((t +
x pm (1) + Jao (1) + + K (1)
Model Capacity: ((0{)))))
Fo(t) = d 2at max (0,1-2)
Population Bias: Spikes in jay
pt) x t 220
(t) - [Cx Fo))
if 1, 1,20 <
Embedding Bias: Bulk of pg
0,
2f <
and if 23 > 1, 2a > 1, J (t) x d
Variance: Spikes of pex
K Rpp (t) dmax (0,1 12n) 2+1/
Learning rate
Algorithm
Time scale, {())
7
SGD(72)
72
0(1) 1 + 72Bl
SGD-M(12: 73. 6)
72
73
O(t = 1 + (72 + 23 Bt
DANA-constamt(22: 78)
72
((t) + Y/2
DANA-constant(12, 7a)1, .1 1
((1) + 73 rB
72
DANA-dccaying 72. 78)"
8(t) = 1 + B(1 + +1)2-11(22)) 1/(2a)
7/2
7a(t) x (1+
(ANA(12,73)) (t(t; d) general schedule
8(1) = + B2 + (Jo' Vra(s ds)²
7/2
DANA-constant with 73 72 x 1/d. DANAAdecayiig only when 2a > and x 1.
Table 4:
Asymptotics for the forcing and kernel functions for all algorithms. See Sec- tion G/Section F for details/proofs of the derivations for these asymptotics of SGID-/SSDD. See Section H for details/proofs of the asymptoties for DANA-constant. See Section for details about the heuristics used to derive these asymptotics for DANA -AeA-inc Here the constant C' is independent of dimension and
order 1 independent of d. For the DANA class, we do not have a full proof. We believe this is true for 2a > 1 and 7 =s < 72, which we believe is true for stability reasons and would suggest that y 722 this is the most uncertain part. Note that KKp
B where Kpp is defined in Thm. 4.1 and (t) def (t, 0) is defined this section.
Therefore, informally, we have
#s Rd MT + MTT. + PTo0
(66)
where we define the three measures as (=(()))do(dw)) Me (du) def Locugiu( (23-1)/((o)) du,
Peo (du) def
and M3. (du) def 16-20 Czu 1/(2a)
du.
Here c3 = -
if 3 2/2 and 0 otherwise and 8y is Dirac delta function, that is,
50 (x(x) def 0, + 0
where [ 50(x))
( dx 1.
OC, x
With this interpretation, we decompose the forcing function based on integrating against the three different measures uyr pp Pyre, that is,
Fo(t) def [. (t,(t, (5. = ((s.)) 0))
Fpp (1) del / 0 (1,(1) x u3 (du) = 1 (#,(1,))), (28- 12 du,
(67)
pc and Fte(t) def / (#.(t.))) x po,.. (du) = Cg (t)(t,0)) x " 1/(2a) du.
d d-20
57

---

In particular, for any "nice" function (t) R2o Rzo + R, with some regularity parameter A we know the existence of a constant C(a, A) such that v4 0:
C (Fo(1) 3 J.)()) S F()) C (Fo(1) + F ml Fac(t))
We formalize this idea in the next proposition.
Proposition D.13. Let a > 2.3 > 0 with 2a + 23 1, a+1> a, # 2. A1 > 0. There exists
M, M,, M2 > 0 and C(a, A1) such that for any f [0, 1] + R, satisfying vu E (0, 7
min f(u) 2 A1 max f(u). [2,2u]
[2u
we have for any 2 1,
C(a, A1) Mid-on ((u) (UT, + (Fre)) (du) S f Ma--2 ((u))z((uu)
1
J ih
S S(a, J f(u) (ueD, + (Fac)) (du).
Proof. The proof is entirely similar to the one of Proposition D.14 Remark D.1. An alternative (although not identical) definition for po is as a sum of masse-s
Mo (du) Σ
(23+20) 8 (du).
D.4 Kernel function
We can simplify the kernel function in similar way as for the forcing function, that is
K(t, s) def X, pp (t,
where Kyp that is simpler analyze than directly the kernel function. Two differences that there is no mass at 0 and the only the pure-point contribution survives. We define accordingly
part is negligible. Hence
Pxpp (du) def
1ocugin"-1/C0 111(2a))) du where Hxx Re pucpp?
(68)
and the kernel function integration against this measure PXpp by
Ktp (t,
def B
J ( +
pxpp (du)
(69)
B / (72(s) 8))1 ()()) 8))12) u² 1- 11((2)) du.
In particular, for a > 1. * 1/ and for any "nice" functions & (t, s)12, 4 (t, s)12 Ryo x Rzo x Rz0 R, with some regularity parameter know the existence constant Cla, A) such that Vt 2 0:
1 (t, s(t,s) s C x (t,
We formalize this idea in the next proposition. Proposition D.14. Let A There exists M, M1, M2 > 0 and C(a, such that for any f [0, 1] + R4 satisfying Vu E (0, })
min f(u) 2 A max f(u). [u,2a]
[2u,4u]
we have for any 1,
f(u)uec, (du) S f #
f(u) S C(a,Ai) f
f(u) (du).
Cla, A1) 2id-za
Ma--2a
Mad-2a
58

---

Proof. We will instead show the following claim, which directly implies the result.
Claim D.1. There exists k(a) E N and C(a,, A1) such that for any f [0, 1] + RR satisfying
Vu E (0, 2) the following holds
f(u) 2 A1 max f(u). min
[0.2u]
[2u,Lu]
then for any d and any k1, 22 E N4 with 2a 20 log(d) - -k> ki k2 k we have
1 kkkttt C(a, A1)
p² 2- Je (( (du) s L f(w))x((du) S C((a,A)) - f(u))xx,, (du).
Proof of the claim: First, using Proposition D.8 and Proposition D.12 we know the existence of k(a) and C(a) such that for any 2a log(d) k z k1 2 k2 2 k, and & [K2 11 + 11.
Hu ([2 *, 1I) (xx > 1])+ I CuCpp (2-*, 1)-
For the left side we write
-2+2
1-
Jamn f( (du)
u 2-4 f(u) x pox (2 (2-kt1,2² 2-h1)
max
k=k2
1
min k-] f(w)) x x(2+) 2-
A 2
k=
((a)
A1 22 f(u)
Similarly for the right side:
J Su)) (du)
min
S()) x poxp0 (2 2-h1)
wE(2-k-1 2-
k=
h 1
((w))
1 -ox 2 2-k+1
A
max
x ((a)
uel2 2 R+2j
k=kz
2 2-
A1
(u))ux(duu)
(a)
12-
Taking C(a,, A,) def C(a) Ai
yields the result.
Remark D.2. An alternative (although not identical) definition for picpp is as a sum of point-masse
pxX (dw) = Ey-da 25-da88-20 2--20 (du).
We can finally provide the following heuristics on the different loss terms:
"Population bias" (Fpp): This loss term corresponds to the loss dynamics when running full- batch gradient descent on the problem (hence following the population gradient), without the embedding matrix W.
"Model capacity" (F)): This loss term (which only d-dependent) represents the limit of the loss, the number of iterations reaches infinity. It arises from the partial expressivity of our model class, since the learned parameters e Rd cannot encode the whole target vector b IR when > d.
59

---

"Embedding bias" (Fac): This loss term comes from the random embedding matrix W which deforms the spectrum of the data covariance matrix and misaligns it with the target vector b.
"Variance" pppp): This loss term comes from the stochasticity the algorithm which at each step samples new i.i.d. random datapoint compute a stochastic gradient. This stochastic gradient, whose average recovers the population gradient, a non-exact estimate of the gradient and therefore introduces this additional loss term.
E
Compute-optimal curves General
We summarize the results of this Section E in Figure 12 (Phase Diagrams) and for scaling laws and compute-optimal tradeoffs see Figure 13 and 14 (above the high-dimensional line) and Figure 15 (below the high-dimensional line). In the Phase Diagrams, we indicate which phases acceleration occurs.
In this section for finding the compute-optimal curves, we note that DANA-decaying refers to DANA with K3 2n and R2 possible learning rate for stability. DANA-constant refers to DANA with ki = max (0, 2ax,
max{O, 1 2a). This is the same as saying that 72 is the largest
R2 = + K1, and Ki O.
In light of (62) and the simplifications of the forcing function (Section D.3) and kernel function (Section D.4), we have that
P(t) x 3 (t() + Fa + Fo(t) + BB Kpp(t),
(70)
where 7 = 72 for SGD/DANA and 72 I for SGD-M, and K (t) (t,
Since for each algorithm we consider (SGD, SGD-M, DANA-constant, DANA-decaying) these terms are asymptotically equal to d' -t we can now derive the compute-optimal curves and exponents. See Table for the asymptotics of the forcing and kernel functions for each of these algorithms. For derivations, see Section G (SGD-M), Section H (DANA-come) DDANA-constant). and Section I (DANA- (DANA-decaying).
To simplify the computations for compute-optimal curves, we introduce the following curve
P(t) max (tpp(t),
1 Fec (t). Fo(t), ,BB Kpp(tt)))
(71)
The function Ptt, d) achieves the same power law behavior as the original compute-optimal curve P(t, d) (i.e., the slope of the compute-optimal curve is correct) and deviates from the true curve by an absolute constant (independent of d and f). Note that some of the terms in the max function (71) should be taken to be O when not defined for the different phases. Therefore, we derive the compute-optimal curves by solving the problem
min P : and d' (f) def
arg min,
P( d), d.B
(72)
then the compute-optine curve is P (f)
P ( d' (f)).
Using this alternative loss function, It,, d), the compute-optimal line must occur at one of the corner points, i.e., where any pair of functions equal each other. The following lemma gives a useful characterization of these points.
Lemma E.1 (See Lemma D. in [80]). Suppose Co. e₁ > are constants and 70. Po. P e exponents and such that a function Ptt, d) equals
P(1, d) = max (Cot- tod -
Suppose that there exists E {0, 1), j def 1 - such that Pi % > 0 > Pj 7/. Then replacing 1 the minimizer of in d satisfies d' def arg min, {5( d)} = -E:) 1/(m po x 11-000t))0(331-00000000
and the associated value is
min Plf, Co 70 x
60

---

Loss (())
Trade Compute-optimal Curves
Ia Jor Phase (f) x 2x+ -1+TT-1444/0-1//200
prase la x fz/ (2a+1)
lb
Fo
Phase lb (f) x x33 a-
1
Phase
(pp(t) (t)
Ppp = Fo
Phase lb f2
a(2a+28-1)
Ic
Pot
Phase lo (f) x (a(28-3)-
1-2(0+
dpaase lc x 2(0 (2 3 3) 20+1)
20-22-1
(t) + Fec(t) +1)((l)
Ju Phase u (f) x 2(a+A)
Phase II
Fpp = Fac
dp x ((0/a)/(1+B/a)
Fac(t) Fo(() + (pp(t)
T
Phuse In (f) 2(2-1a)/(1a)
Phase III
K Jac
BB
Phase x f2/2
Jom
Phase IVa. a
IVa
75 Kpp
= 3o
dehase IV f2/22
pt) Fo(t) + (pp
Phase IV
Jo FoD
Phase Ivb (f) x (2(208-0-22)
201(2a
IVb
K
dlesse IVb x f(a- (200000220
Table 5: SGD-M/SGD: Loss description Ptt) for SGD-M (7 def 12 + 7 and SGD (- def compute-optimal curves for 5 7. d) across the 4 phases.
72) and
Proof. The proof a straightforward computation. The minimizer of P(j. d) in d must occur where the two terms in the maximum equal, i.c.,
Co (o()) 0 Po -1) dd
Solving for this d gives d², Plugging in the value of dt into Pf d) gives the optimal value.
Remark E.1. The possible minimal values of (72), i.e, where pairs of functions in the max are equal, can be reduced further. For instance, Fac d) exist for the phase, then for some O < ro < rı < r2
F Spp(t,d)
0 < t to
1 B (t.
to t
P(t, d) x
do, Folt,d),
t < t t2
< t.
Thus, there are only maximum of three points to check order to find the optimal compute curve.
Remark F.2. In view of Lemma F, 1, to find the optimal compute curves, we first find the potential curves (i.e., all the possible combinations of two functions in the loss curve are equal while still lying on the loss curve). Then the curve which has the smallest exponent on the flops, f, the optimal compute curve.
E.1 Stochastic momentum (SGD-M), compute-optimal curves
For constant momentum, the loss curve well as the forcing terms Fo, pppp acc and kernel term KMp are entirely similar up to constants the for SGD [80]. The compute-optimal curves are hence identical, see Table when replacing 72 by 72 + 2 (see Remark G.2). See also Figure 12 for description of the phases in the (a, -laane...
61

---

E.2 DANA-constane compute-optimal curves
In all this section, will use the hyperparameters in Lem. H.5 and Cor. H.2 (see Section 3) with B = 1. We discuss the effect of learning rate and batch after. The asymptoties of the forcing and kernel terms for DANA- CANA-constant, valid only in some regions of the (72, (3; B, L, 5) space, are below
3 pp d) = min (12 Bt) B1) 21= ((V ((173B1) 20-1 }.
Fao d) x d ( 1+ 2o (V 2+
72 kpplt, x 72 min (12 2 (V -+ Folt, d) x d max
Derivations for these forcing function and kernel function asymptotics can be found in Section H. For a summary of the compute-optimal curves for DANA-constant, see Table and Figure 12b for description of the phases in the (a, (o,,))planne
Below the high-dimensional line, (Phases Ib, Ic, IVa, IVb). In that case, the limit level Fo is reached for < da < We have, vt d,
I ppt d) (t² 2u
28-1
Ftc x d d (12 1+
and
Kpp 7B
Kp, (t, x 12(9)ttt)
Hence the forcing and kernel terms are similar to SGD and the compute-optimal choices for d", t* are the same as in Table 5.
Above the high-dimensional line (Phases Ia, Ia, Ilb, Ia, IIIb). A first observation is that the limit on the min in pp.. Tapr Kpp can never be computc- -ptpimll by itself, since decreasing d would yield strictly better performance. Hence we only need check the risk between the different terms of the forcing and kernel. Additionally, from [80]. the compute-optimal for SGD happens in all phases for 2 d. Hence it has be similar for DANA-constant since. We can therefore simplify in the following for i 2 d.
I (t, x (VygBt)-2 28-1 a
Fee(t, x d (VrrBt) 2+
1 72 B kml(, 12 Bl)
1
and
Phase Ia. In this phase, the approximate loss curve satisfies
p( d) x maxtFpp Fo( d)} max( max((V/alll) (VrgB1) 2- 28-1 2a 20-28-11.
(73)
Proposition E.1 (Phase I, DANA-constant), Suppose we are in Phase la, i.e. 2a 2, <
Then the compute-optimal curve ( using (73) occurs with d' x f3rt-a and P pff.d d') x
2a 28
3/240
Proof. We apply Lemma E.1 with
% = 2+ 29
23 1
Po =
71 = 0, and
P1 = 2a + 28 1.
2x
Phases Il and Ilb. In this phase, the approximate loss curve satisfies,
(Fpp(ppp))) P(},d) = max ppl d), Fao ( ( d), Fo(
F
(74)
x maamma 20-1 ,d (VBBBt - 2 max
62

---

Loss (())
Trade Compute-optimal Curves
la Ph
Phase (f) x
20-2261
3/21a
paase la (f) x f3/2+0
lb Prr Phase (f) x f3 -a-
Phase
plp(1 (1)
ppp = Fo
Plase 1b (f) =
2222+22-11 1)
Ic Po
Phase lc (f)
1-2(a++8) (1)a(2833-228))
Pnaso (f)
2a (2
Fr Phase Ha (f)
4s
Ila
Fac Fo
4a
(t) (t) (T)(t)
Phase Ha (f)
In
2444
Phase II
28-1
IIb
Ppp Fac
Phase IIb (f) x
33
28
Phuse Ilb (f) x 38
Phase (f) x 4440-- 4a
IIaa
Jac = To
(t) + Fo(t)
Plase llla (f) x 4da2 4a
Phase III
+ 1
(1)(()
fo
Phuse 111 (f) f- 1
raB
IIIb
m/ T Kpp
F
dhase I11b (f) x f1/2
Fo Phase IVs. (f) f-a
IVa
B2B Kpp
= Fu
dpnase Iaa x f1/2
(1) Fol!)
Phase IV
+ BBB (tp(t)
IVb
Kpp
Fpp
fom
Phase IVb (f)
11-201(2022--1
(A:208-0-283)
dphase lVb ( ()/
Table 6: Loss description P(t) for DANA-constant and compute-optimal curves for P(g.p d)
across the 4 phases (and subphases) defined in Figure 12b. We consider DANA-constant with the hyperparameters in Lem. H.5 and Cor. H.2 (see Section 3) and batch size B
Proposition E.2 (Phase II, DANA-constant), Suppose are in Phase II, i.e. 2a > 1, 23 > 1.
B. Then the compute-optimal curve P( ( d²) using (74) occurs
a 3 (Phase Ia), with d' f and P( ( d')
3/21
28 ifa < 2 (Phase llb), with d" निोनल and P( f d d²)
(2a 28-28-11
1,44
Proof. We have two potential cases to check, whether the compute-optimal is attained for Fpp (t) = Fno(t) or F (t) (t).).
(t) Fa (t)( We apply Lemma E.1 with
23 1
23 -
1
70 = 2
a Po -1
71 2
and P1
2a
a
If < 3. we have P1
m 0 Po 7o and it yields optimal dī
f and P 2 d;)
1/3 On the other hand. if 3. we have p() Fae(1) for L. $3/2+0 but the optimal is to take 1* the largest. This brings us to the second case.
(20428-1) 30
T (to(t) Fo(t): We apply Lemma E.1 with
63

---

1
=
2 O²'
Po 2x." 0, and P1 = 2a.
a > 3 j. we have p1 71 > 0 > Po 70- Hence the minimum is attained for d2 = f and
P( P(2. d.) = fo.4-- 2a 3,/ For a < however the optimal is to take t, the smallest. We conclude that for a > 7. in Phase Ila, d² = d and for a < ] in Phase IIb, d" = dz-
Phases Illa and IIIb. In this phase, the approximate loss curve satisfies,
1 PC I d d) x max
12 xp d), F Fac d), Fol2
(75)
x max max((22 (72(( 73 BI) -+2 at (v Bt) -+ ,dd2aa)
Proposition E.3 (Phase III, DANA-constant,, Suppose are in Phase III. i.e. 2a 1, 23 > 1. B. Then the compute-optinge curve P( das d²) using (75) occurs
24
if > å (Phase Illa), with d² x fau and P (P d²) x
ifa < 1 (Phase Illb), with d' f1/2 and PK d' f 1
Proof. We have two potential cases to check, whether the compute-optimal is attained for 22 K pp 1, d) F atc d) or for Ftc d) Fo(t, d).
Kpp d) Fto(t, d): In that case, directly know that di // d²) f-1 t. We
also see that by defining
Tu = 4 a
Po = 2 2a²
"1 =
and P1 =
a
2a:
If a > 3
I: we have
> and we apply Lemma E. to obtain di :2//2, 112, P .(//dd di) =
f-1+ t. If a < 3.8 the optimal is to choose 1* the largest which brings us to the other case.
Fao d) Ft, d): In that case, define
1
= 2 -
Po
0, and P1 = 2a.
2x.'
For a < 3
7. we have P1 - 9 0 > To and hence applying Lemma E.1, it brings an optimal
dž = fo+s.. P( d5) - -45/ For a > 3. the optimal to choose t* the smallest going back
3 in Phase IIIbb the first case. We conclude that for a > 4 in Phase IIIa, d' = dř and for < 4 d' di
E.3 DANA-decaying, compute-optimal curves
In all this section, we will use the hyperparameters in Remark B.3 with B We discuss the effect of learning rate and batch after. We remind below the asymptotics of the forcing and kernel terms, valid only in some regions of the (72, Fa. B, 8) space (r(1) x (())).
(t, d) x min ((t))) 2 Fao d) d min S Bt, (t))² -11 - 2a+
1
T2 (m, d) 72 Bi ((99)) Fo(t, x d 100+maxx0011--00000 10)1
Derivations for these forcing function and kernel function asymptotics can be found in Section I. For a summary the compute-optimal curves for DANA-constant, see Table and Figure 12d for description of the phases in the (a, 3pllnane.
64

---

Loss (())
Trade Compute-optimal Curves
2xr 201(d 1)
la Tre
Phase (1) x
Pruse La (f) ~ 440-1
lb
P Phuse (f) x 12 a-- a-
Phase 1
ptp(t) (t)
ppp = Fo
Prase 1b (f) x f2
a(2n+2881
Ic
To+
Phase To (f) x aa(22-3) -3)---3
1- 1-2(a++8)
denuse
2(((228 -3) 28-11
20(da 1)
P Phase Ha (f) 4a 4a-14 1442
Iaa
Jac Fo
4a-1
(1) Fnc(t) +t)(tt
Phase Iaa (f)
in 1a-1+402
Phase II
(2--28-
Jom
Phase IIb (f) x
2(20 1408
IIb
Jpp = Fac
(da-1)d
Phase 1hh (f)
20
4108-8
2n (1n-1)
P Phase Ma (f)
1n 10-14
Illa
Fac Fo
4n-1
T (t) Fo(L)
Phase x
110-1410
Phase III
(t) BB
To
Phase II1b (f) x 20(6a-1)
(10-1)2
23 Kpp = Jac
d d'mase IBl (f) x f
Ja-1
D Phase IVa (f) f-
IVa
1 72 Kpp = 500
dphase IVa (f) x x11/2
5 (1) + Fo(()
Phase IV
BB (pr(1)
IVb
1 T2 Kpp
Fm
fo
Phuse IVb (f) x (2t2n0
20l22a
dphase lvb x f (2
Table 7: DANA-decaying, Loss description for P(L) when solved with the DANA-decaying algorithm and compute-optimal curves for ( d) across the phases (and sub-phases) defined in Figure 12d. We consider DANA-decaying with the hyperparameters in Remark B.3 and batch size 1.
Below the -ime-dimall line, (Phases Ib, Ic, IVa, IVb). In that case, Bt S 1/(2a) z r(t)?. Hence, write
F (t, = (72 (t2 28
(to (72
1 1 Kpp (t, (t, x
and
B2
Hence the forcing and kernel terms are similar to SGD and the compute compute-optimal choices for d", t are the same as in Table 5.
Above the high-dimensional line (Phases Ia, Ha, II,, Ma, IIlb). In that case, we can check that Bt 2 7s 2 2 r(t)2. Therefore we simplify
(t, x (5s BB. 1/22111(2)---)1
28-1 20
(t, x d' " 17 -120))1-) 20
and
(pplt, d) = 2e Bt 1//22)) -2 to.
B
Phase Ia. In this phase, the approximate loss curve satisfies
P(G. d) x max{Tpp ( (d) Fo d)} x max (0xxx((7g Bt ( ) 20
-1 28-1
2a 20-22+11
(76)
65

---

Proposition E.4 (Phase Ia, DANA-decaying), Suppose we are in Phase Ia, i.e. 2a > 1, 2/ < 1.
Then the compute-optimal curve P f. dd)) using (76) occurs with d" x 4 and P( d')
O- -28](1o-1)
(4n 1)
Proof. We apply Lemma E.1 with
23 - 1 ) 1
70
Ph 0,71 =0,
and pi 2a + 233
2xx
2xx
Phases Ila and Ilb, In this case, the approximate loss curve satisfies
p d) x max acc ( d)}
(77)
max (55g 26-1
d d (73B(211/220) -11-22))))1110_0)101111111111111 ddd
Proposition E.5 (Phase II, DANA-decaying), Suppose we are in Phase H. i.e. 2a > 1, 2B >
B. Then the compute-optinal curve ( ,x,d)) using (77) occurs
if 3+55 (Phase Ila), with d' x fio-1440 and P (2,, d²)
4n-1
2utda 1)
4a-144a"
ifa < 3t/5 (Phase llb), with d x 220 and P dr: d²)) x
(1--1)8
12 11f4a 1)
2(2u 1/2u+1440 1400 D)
The compute-optimal choice can be either attained for 3 (2 d) J Jac( d) or (.c(( Fo(2.d).
Proof.
F Fmp (p.(d) = Facl 2, d): In that case, we introduce
1
2,3
70
2
1
Po = O,
2
1
and Pı 1.
2a
2a
2a
2a
If a < 3+/5 then P1 -
> 0
Pu 7o and applying Lemma E. we obtain an optimal However, if
3+ then P1 - and Po 70 < 0. Hence the optimal is to choose t" the largest which brings us to the second case.
di = faa -
and P 2.
d1) =
(2w
120 194 11 2(2 4a0
Faol d d) = F. d): that case we define
1 2a
7o =
2 2n
Po 1, 71 =0 O, and P1 20x.
If > 3+/5 then 70 0 <
7 and we apply Lemma E. to obtain that d; x f 114a
4n-1
and ( d3) x
2arde
4 On the other hand, if a v+ then Po - To > 0, Pi 7 > 0 and 4
the compute optimal is to take t* the smallest, i.c going back to the first case. We conclude that for 3+53 in Phase IlIa, d² di and for 3+y5 in Phase IIIb, d" d
Phases IIl and IIIb. In this case, the approximate loss curve satisfies
P(G.0) d) x max ( (1 d), For d), Fo d)} -2 tox d (1 - to dd
(78)
x max - 11/(20)))
66

---

Proposition E.6 (Phase III, DANA -deA--deay.ng).. Suppose we are in Phase III. i.e. 2a > 1, 23 > 1,
< B. Then the compute-optimal curve P( (z, d') using (78) occurs
ifa > 1
34y5 (Phase Illa), with d² x flo-1440 and P d- dt) 10-1-104
20(4a-1)
da-1 ifa < 3t/5 (Phase Illb), with d' x foo-i and P (0, dd") = = 20(60-1).
(4a-1)²
The compute-optimal choice can be either attained for F ppp(d: d) = Facl or Fac Fo(2.
Proof.
7B Kp (2. d) = Faol aa. d): In that case, we introduce
70 22-0 2a/ )
1 1
Po 0, 71 (2
1
1
2
and
P1 =
2a
2a
2a
If a < 3+y5 then m
To and we apply Lemma E.1 to obtain d, = and
(1)-
P( dij x 1)) However, when a > 3+y5 then pu 7o < 0, P1 "1 < 0 and the optimal
choice is to take ts the largest, leading to the second case.
Fadl (ao(2.d) d) = Fo d): In that case we define
1
211 2a
70
2
1
Po 1, 71 =0 0, and P1 2a.
2a
4--1 a 34/5 then Po 70 0 P1 71 and we apply Lemma E. obtain that di x f 110-1+10 1
20)4e-1)
and (( d5) f -0013) On the other hand, if < 3+y5 then To > 0, 7 > 0 4
and the compute optimal is to take t" the smallest, i.e. going back to the first case. We conclude that for a > 3- in phase IIIa d' di and for 3+v5 in Phase IIIb, d' = d3.
E.4 Comparison of samples needed at compute optimality
Independently of compute, a bottleneck in the training of large models is the amount of data available. While the size of a model can be arbitrarily increased, data comes with hard limits: the size of internet when it comes to language models [97] or when dealing with resource constrained tasks such as medical imaging [13]. Hence natural question is how do the previous algorithms compare in terms of samples used at compute-optimaligg In what follows, we denote DANA-c DANA-constant and DANA-d as DANA-decaying.
We know that for given number of samples/iterations t 2 0, DANA-d always achieve smaller or equal loss than DANA-c which in turn achieves smaller or equal loss than SGD, Henee for fixed given loss level, DANA-decay needs less samples to achieve this loss than DANA-constant which needs less than SGD; strictly less when the scaling laws exponents are improved the considered regime. However, it is not clear that this ordering remains the same when considering the compute- optimal training regime. Indeed, in some regions of the (a, 3) plane, the corresponding training regimes can be different for two algorithms. For example, for < and B < the compute-optimal regime of DANA-c happens at the frontier Fae/Fo while the compute-optimal regime for SGD is between Kyp/ /F/caacc The compute-optimal regime of DANA-c is shifted later in training (this is general effect of the acceleration, see Remark E.2). Hence, even if for a given loss level DANA-c needs less samples than SGD, DANA-c needs more samples at compute-optimality
In the following denote for E (SGD, DANA-c, DANA-d), p: > such that in given phase, the number of samples required for algorithm is at compute optimality 1* f The smallest Pt. the more data-efficient the algorithm for given compute. A first observation, that even though the loss P (f) is continuous, across all the phases, neither the optimal dimension d' (f) nor the number of
67

---

samples tx (f) are continuous. The points of discontinuity are exactly when the trade-off condition changes, for examples at the borders IIa/Il or IIIa/ IIIb.
Claim DANA-d will always use less samples at compute-optimaliy than DANA-C in all phases.
Additionally; given (a, (9) and for two algorithms in (SGD, DANA-C, DANA-d). if compute- optimality is reached at the same trade-off' between Jac, pp- Kpp then DANA-decaying uses less samples than DANA-constant which in turn uses less samples than SGD at compute- optimality. The improvement is strict if one of the algorithm accelerating with respect to the other. Hence DANA-d, DANA-c and SGD are in that order sample efficient in phases la. Ilb. IlIb. However, for large a, compute-optimalit for DANA-c, DANA-d is reached later in training due to the acceleration, and DANA-c, DANA-d may well use more than less samples than SGD at compute-opimmalit. depending on a; 3..
Proof of the claim. We consider each phase:
Below the high-dimensional line, SGD, DANA-constant and DANA-decaying have the same scaling laws.
In Phase la. PSGD 2a+1
PDANA.
3/2+a
PDANA-d =
2 4a-
4a-1
In Phase IIb (of DANA-c), PSGD = 1
3/a 1+3/a
> PDANA-C = 1
38+0
> PDANA-d
=
1
(4a-1)8
2a 2a-+1ap-3
In Phase IIIb (of DANA-c), PSGD PDANA
PDANA
4a
6a-1
In Phase IIb (of DANA-d), PSGD 1 additionally in phase Ila, then PDANAiC =
B/a T48/a
PDANA =
(4a-1)
I are
2n 2n 4-4
4 > PDANA- = 2a 24084
4a-2
(1a-1)0
In Phase IIIb (of DANA-d), PSGD = =
PDANA-
6a
4a-1
If we are additionally in
Phase IIIa, then PDANA-C = 1 Ia-+1a-3
4a-2
PDANA-d
40-1 6a-I
In Phase Ia (of DANA-d), PDANA = 1
402 2o2-40-00
da-2
PDANA-d =
4a-1 4n 4-2444111
In Phase III (of DANA-d), PDANA-C =
4n-2 +la-3
PDANA-4
=
Ta 1a-+1a- 1a-
4a
In the other cases, there is general rule of whether algorithm uses less samples than the other at compute-optimality. The interested reader may still easily derive them using Tables 5 to 7.
One can explain why DANA-d always use less samples than DANA-c, even when they don't 't share the same trade-off condition (i.e. for 1: 3+
B B> 1/) by the fact that DANA-c shifts the tra (later in training) for smaller a than DANA-d.
E.5 Summary on compute-optimality results
We provide some specific details about compute-optimality for each algorithm in the different phases. Below digh--dimesiooal line (Phases Ib, Ic, IVa, IVb); Limit level To is reached when < DANA/SGD-M M has same scaling laws as SGD compute- compute-optimallty.
SGD-M. Since same scaling law SGD, there is a large portion of the (llaae (Phase III. IVa, Ib) where we have universal scaling, i.e., params exponent /22 or equivalently, the compute- optimal regime is the same as the proportional regime (t. = d). This was observed empirically in [50].
MAMA-conttatDDANA-dccaaaaa To improve the compute-optimal loss exponent for DANA- constant, compute-optimality must occur after iteration d-1 (Thm. H.3). As noted in Thm. 1.2, DANA-decaying improves the loss exponents for all scaling regimes where 2a 1, including the compute-optimalit regime.
68

---

Phase la: Here, ppp accelerates beginning at 2 d until it reaches the limit risk Fo. While the compute-opiine tradeoff occurs at 3 pp and Fo (same as SGD), DANA-constam reaches this point faster than SGD; thus DANA-constant outscales SGD.
Phase II: This phase involves pp. 3 acc Fo and at compute-opimmllty. we always see better loss exponent since the tradeoff point occurs at point where 2 d. Notably the acceleration changes the tradeoff constraints (e.g., where compute-optimal tradeoff occurs). For a S 0.75 (Phase Ilb, DANA-constant)
S + / (Phase IIb, DANA-decaying), tradeoff occurs the same two terms as SGD, i.e., Fpp and Fac. but its get there faster. This means that one uses fewer samples than SGD to achieve compute-optimality For 0.75 (Phase Ila, DANA-constant) (Phase Il, DANA- -ANA-decayigg), the acceleration Jac shifts the compute- ootimall frontier to Fac and To. making DANA model capacity constrained, that is, why the compute-optinat frontier exists changes.
2 (3 + v5) /4
Phase III: This phase involves K Kpp. Facs and Fo and, notably, Xpp and Foe always intersect at t x d. The same as Phase II occurs with Xpp replacing ppp When < 1. though, for DANA-conste the compute-optimal tradeoff occurs Kyp and Fac. so no outscaling SGD.
We summarize the results of Section E below; for Phase Diagrams see Figure 12, for scaling laws and compute-optimen tradeoffs see Figure 13 and 14 (above the high-dimensional line) and Figure 15 (below the figh-dimensiona line).
69

---

Ou tscales
I1a
Ila
la
la
Ila
highel line
I
high-d
11l
111b
high-d
0.5
n.5
line
wta
Line
IVa
0.25
2-v/2
1b
IVa
0.26
2-v/2
Ib
W
IVa
TVb
TVb
Ic
Same
SGD
0.5
0.5
0.5
/
B
(a) Phase Diagram:
(b) Phase Diagram: DANA-constn
(c) Comparison of SGD DANA- constant
SGD & SGD-M
O tscales
Ou tscales
Ila
Ila
Illa
la
Illa
3+
la
0.75
3+v/
111b
high-d
111
111b
high-d line
111b
high-d
n.
line
0.5
0.3
line
lb
IVa
2-v/2
Tb
IVa
2-
1b
IVa
0.25
0.25
0.26
TVb
Same as
SGD
Ic
Same SGD
0.5
B
(d) Phase Diagram: DANA-decaying
(e) Comparison of SGD & DANA- decaying
(f) Comparison of DANA-constant & DANA -ecaying
Figure 12: Phase diagrams for the various momentum algorithms the compute-optimal regime. See Figure 13, Figure 14, and Figure 15 for cartoon pictures of the scaling laws and tradeoffs across all phases and algorithms. The main phases (I.I1, are based on the components of the loss that dominate each time. The phases are further broken down to account for the different tradeoffs in the compute-optimal training regime. We always use K1 max{0. (0, 2a} and we use DANA-constane with K2 1 Kn and R3 0 and DANA-decayinn with K3 1/(2a), K2 Ri. and batch size B 1.
70

---

SGD/SGD-M
DANA-constant
DANA-decaying
compute optimal
Phase loss curve
t (2a)
loss curve
d2/(20) plpa
loss curve
compute optimal
compute
optimal P
P
flops
flops
flops
SGD/SGD-M
DANA conntconttant
DANA-decaying
Phase Ila, o >>> 2
p((a))
compute
t (20)
compute optimal
t (P/(2a)(2-11(2a)
optimal
loss curve
loss curve
do/(2n)
loss curve
compute optimal
d t 141)
d 1/ 1/C2bb & 241/0
(2 /(2a))) 1/(22))
d
d
d
flops
flops
flops
SGD/SGD-M
CANA-constant
DANA-decaying
Phase Ilb, 7.75 < a < 2
t ((aa)
compute
1 (2m))
compute
(=/(22)((-11((22)
optimal
optimal
loss curve
loss eurve
(20() t
loss curve
compute optimal
d t 141/(2o)
1/(22))-220 241/0
t d
flops
flops
flops
SGD/SGD-M
DANA -constnnt
DANA-dggg
Phuse lib, < < 0.75
[ P((a)
" (2aa
compute
t (22(22)(2-11(22)
optimal
loss curve
loss curve
(2at pfa
loss curve
compute
compute
optimal
optimal g
d b 141/(22)
(20))
241/a
y -11/)))))))
-2aa
20
flops
flops
flops
Figure 13: Cartoon plots of the scaling laws for each algorithm above the high-dimensional line (Phase Ia, Ia,, Ilb). Pictures for the scaling laws for all three algorithms in each of the different phases. When < d, DANA-constant behaves like SGD/SGD-M. Observe that the trade off point for compute-optimum changes across phases and algorithms; indicated by (magenta), population bias, F pp (t) = (purple), embedding bias, Fee(t) = (bluc), and model capacity, Fo(t) = (orange). Variance due to the algorithm has impact. Here p 2a 2 1. We always K1 max(0, 1 - 2a} and we use DANA-constalt with Kig = 1 + Ki and Kig = 0 and DANA- CANA-dccaying with <3 = 1/(2a).
K2 K2, and batch size B= = 1.
71

---

Phase III,, a > 3-v5
SGD/SGD-M
DANA-constanl
DANA-decaying
compute optimal
(2 1/(2o))
C (2-11((22))
loss curve
compute optimal
loss curve
(2 1/(2a))
loss curve
compute optimal
d (21)
d 1/ 241/(e)
(2- ty (26))0 g-
di
flops
flops
flops
Phasee IIIb, 0.73 < a < B
SGD/SGD-M
DANA-degaying
compute optimal
t (2-1/(20))
c (2-1/(2a))
loss curve
compute optimal
loss curve
(2 1/(2a))
loss curve
compute optimal
1/(2a)
d 12(d) 241/(a)
(2- 1100010000000000000
20
20
di
flops
flops
flops
Phase I1lb, 0.5 < < 0.75
SGD/SGD-M
DANA-constant
DANA-decaying
(2-1/(20))
t (2- 1/(2)))))
t (2-1/(20))
loss curve
compute optimal
loss curve
compute optimal
loss curve
compute optimal
t (1)(
d 12(e)) 4 241/(a)
(2- -1000
20
20
20
di
flops
flops
flops
Figure 14: Cartoon plots of the scaling laws for each algorithm above the high-dimension line (Phase III,, IIIb). Pictures for the scaling laws for all three algorithms in each of the different phases. When < d, DANA-constant behaves like SGD/SGD-M. Observe that the trade off point for computte-otimmmn changes across phases and algorithms; indicated by (magenta). Variance, (tp(t) = (green), embedding bias, Foc(t) (blue), and model capacity, Fu(t) (orange). Even in the stochastic noise-dominated regime, acceleration occurs for DAAN--conttrn and for DANA- decaying. We always use H₁ = max (0, 1 - 2a} and we use DANA-constant with K2 = 1 + #n and Ki = and decayydy with Ki 1/(2n). N2 = k1, and batch size
72

---

SGD/SGD-M
DANA-constanl
DANA- DANA-decaying
compute optimal
compute optimal
compute optimal
Phase Ib/lc
loss curve
loss curve
loss curve
10
2G
d (lb), d
(Ic)
di (Tb), d
(Ic)
d (Tb), d (le)
flops
flops
flops
SGD/SGD-M
DANA- MANA-decaying compute
compute optimal
compute optimal
optimal
Phase I'a loss curve
26 2 11
loss curve
1 2n 1 20
loss curve
2 kp
d
2a
20
Z
2
t d
(2a-11 1)a
2a
12- 15 13))
d
(20 1)
2a
d
8
(2)) d
flops
flops
flops
SGD/SGD-M
DANA-constant
DANA-decayann
compute optimal
compute optimal
compute
optimal
Phase IVb loss curve
2a1e
loss curve
2
loss curve
2u3p
12aa)
2ox
1/(2))
20x
1/(2a))
flops
flops
flops
Figure 15: Cartoon plots of the scaling laws for each algorithm below the high-dimensional line (Phase Th/c, IVa/b), Pictures for the scaling laws for all three algorithms in each of the different phases. Observe that the scaling laws are all the same for every algorithm; Tradeoff in compute- optimum indicated by (magenta). Variance, Kpp (t) (green), population bias, F Fpp (t) (purple), and model capacity, Fn(() = (orange). SGD/SGD-M and
have the
same scaling behavior. Here p 2a 2,8 1. We always use Kı max (0, 1 2ax} and use
DANA- Dtnn--cntttatt with K2 = 1 + K and R3 = 0 and DANA-decaying with 63 1/(2a), K2 Ki. and batch size B
73

---

0 -2A(1) + (2 Stochastic gradient descent (SGD)
The ODEs (22) precisely yield the same scaling laws and compute-optimal curves as previous works have shown for the stochastic gradient descent (SGD) algorithm with constant learning rate. Of particular interest reproducing the results of [80] and extending them to include learning rate schedules.
E.1 Volterra equation for SGD
When "/3 = O, then the updates in (Gen-Mom-SGD) are precisely SGD with learning rate schedule that the iterates generated by the "coin-flipping" algorithm when ys
(t).). Moreover, give the exact same iterates as SGD. That is, SGD is the "coin-flipping" algorithm with 73 = 0.
When 73 = 0, the ODE in (22) becomes
dv
dt ((t; x ((t; + gt;
BA + B(B + 1))(
where 2,);
B + 1)A5
(1) A(0)
B1)
0
-A(1) - Bx)
We observe that the v(t; Xy)1 = (t) for = I.... d completely decouple from E} and Xy- Therefore, the ODE reduces to solving the following linear ODEs for 07. 1, I,...,d
d de 03(1) = [- 22) B(B+ B (20
This is the same ODE that appeared [80] with (2) constant. Moreover when X, = then d
di pt) and so ()) = (0) for t.
For XÚ > 0, using Duhamel's principle on these 1-dimensional non-homogeneous linear ODEs, we can explicitly solve for the (t):
pī(t) exp (A(t)) + [' exp (A(t, A(s, BE (s)) W] ds,
where A(t, def 2BL B(B ())A, (2(() der I re()) dr, [2(t) def L 12(7) dr.
Alternatively, in terms of the matrix (1, s), we have for SGD
( 8) = exp(A(t, - A(s, A)).
(79)
provided that À > O and when x l =1
Remark F.1. We do not need to compute { (t, 8)]12 since 71 (t) 0 so it does not effect the kernel function.
F.2 SGD with learning rate schedule
Using the results above, we have the Volterra equation for SGD.
Volterra equation for SGD with learning rate schedule (t)... We can now give explicit representation for the expected loss of SGD, that is, following the arguments in Section C, with 0g = O, 0,
forcing func.
{ X((s) x (( W| ds,
E L (t) W]
(t)
(80)
grad. descent
74

---

where, for any contour I' containing K, we have
F def -
2.i /
(1) exp 2B00(()2 (0000//0080000002-222000-------0 (1225)))))=+++00-00000000 dz
-1
Xxtt)
2 (B)
$ 3² exp (eep(-22zz(A0 -
x Tr
(s)))
L(2, R) d2)
and (2(() def I rolt) de and (t) def [ 7t())) dt'.
The Volterra equation using the deterministic equivalent for K immediately follows.
Volterra equation for SGD with learning rate schedule (t(() under the deterministic equivalent. We can now give an explicit representation for the deterministic equivalent of SGD, that is, following the arguments in Section C, with 0g O, O,
forcing
P(L)=
F(L)
(1) x P(s) ds,
(81)
gra.. descent
SGD noise
where, for any contour r containing [0, 1]. we have
F(t) del 1
2τi f. (D) exp ( (t) B(B (t)²() da
[ exp 2BT (t)u B(B+ ())2))) pd((d)
(t)()) def
2 (s)
2τi
x Tr
2² exp ( --2-:110000-166))-------------1- -
x R(2) dz
(2) B
exp ( - [2(x)) + B(B + (f2)(((() - f2()))
X pxx (du),
and (t( del S 12(r) dr and T2(1) / (2()) dr.
Remark F2. The results above also hold when (t()) is a constant and they exactly reproduce the results in [80].
F.21
Simplifying the Volterra equation for SGD with learning rate schedule
We now aim to simplify the Volterra equation for SGD and do so, we will assume that the SGD learning rate schedule is decaying.
Assumption 4 (SGD Learning Rate). Suppose the learning rate schedule 72 R20 R is nonincreasing function. Moreover, we let 7/2
def (0). that (2(()) S for all t.
Throughout this section, we assume that Assumption 4 holds. Prior to this section, i.e., the derivation of the Volterra equation (81), holds regardless of Assumption 4.
75

---

Next, we need to introduce upper and lower bounds on the SGD kernel. For this, we introduce two kernel functions,
K
SGD
def
{ exp( + B(B + 1)72u²t) (d((d)
(82)
SGD (t) def ( exp(- (2But) pxx
and
To simplify the notation, we write K(t) del x (t) and (t) del x.GD (t). Next, we know that (2(() (2()) s 72 (T) - (2))). Using this observation, we have, for any 0 S East
(s) (t) (2())) S (t) S (5) ((T - (2())).
(83)
Similarly, we define an upper and lower bound on the forcing function 'f
F()) exp(-2 B(B + t) per (du) and F(t) S 1xp(-2Brug pe (du).
(84)
It is clear that
g(T((0) S 5(1) S F(T(()))
the following lemma, we show that K and F are monotonically decreasing in L. Lemma F.1 (K and are decreasing functions). Suppose Assumption 4 holds and the learning rate satisfies
2
72-
++1
Then I and K are decreasing functions in t. particular, the function 3 is bounded.
Proof. First, if F is decreasing, then it is clear that F is bounded. Therefore it only remains show that F and K are decreasing. This amounts to showing that
u( w + B(B + 1)52u) < 0, for all E [0, 1].
This holds provided that B+1 72-
Lastly, we introduce Kesten constant for the kernel of SGD the following
K(t)) dt.
(85)
We now proceed to simplify the Volterra equation for SGD in (81) using a result similar to Lemma D.3 in [27]. Lemma F.2. Suppose Assumption 4 holds, that is, 72(t) is a nonincreasing learning rate schedule and 2
B+1 72. Let (t) (r) dr. Then for all 2 0.
(r())
P(() 2 5(1) /'' 72) x K(T(()) - (2())){(ss ds.
If, in addition, there exists a E > 0 and T(:) > 0,
S K(s)K((t s) ds < 2(1 + and 27.B| + + < 1,
(86)
then for all
P(t) s 5(1) +(1)+ +C L :)())B x x(T(() - (2(s)){((s)
for
K(00)
C= (C(T) +1) +1) 1 - BI +
76

---

Proof. We recall that
p(t) - J(t+ + L.¹ (t) x (P(s) ds.
The lower bound holds trivially after noting that (t) 2 72( x - and P(s) 2 (ss). For the upper bound, we start with the following. Let us define the convolution map
(1)) def J. (1)((())) ds,
with the composition of the convolution mapping by
(T) = J. K. (1))(()))) ds ds= S (x)(,,,()))) dr
This naturally extends to (t). Next, let define two functions h Fo and g(u) def B x ( where T2(()= = Jo (s ds. Moreover, we introduce another convolution mapping given by
(()) del [ (((l-)g(u)hat du,
where the composition is given by
('() / x(I - ())()) ds ( 1 ""x(- 8)g())( u)g(u)h(u) du.
Again extend this j compositions, (t))
As the kernel function (t) and forcing function F(t) are non-negative, we have that
P(t) = 3(t) + Σ ())
(87)
Next, we prove the following claim. Claim: 9 (F) < 2 (h) for all j 2
Proof of Claim: To see this, we will do the case when j 3 in detail, but the idea will extend to all j 2 For this, we have that
3" (t) =
T2(())
J. 10
J' TT(2(()) - w)K(w w) ()()))))((()))(()(()))) de du dw.
consider the change of variables v = T(p) where du 72(p) dp. Then
T2(t)
('())(T2())
X(T2(()) w K(w - (2(())
gl du dw.
Now consider the change of variables where a (2( Then,
[2(t)
((())(2(())
K(Tz(t) w)[(( 1 T2(p))
x ( dp du
T2(t) r- (w)
Jo
K(T2(()) - w) w
- Tz(P))
x xg(w))30)) dp dw
Jo TT(T((() - [2(8)) - T2(r)) - T2(p))
B) dp
77

---

The last equality follows from the change of variables w = T2(s). The result immediately follows since (t)) - (s) 2 (t) and the kernels are nonnegative. This proves the claim. Therefore, have from (87)
P(() S 5(1+
Σ s (1)(22())
Next, we show that the map ((h) is contracting and in particular,
(b) = L'' Kit s)g(s) ds = 1//, J'" K[t (5)_(()))){(s
u)g(u)h(u) uu)(((() du ds
.h (L. x(( s)X(s= -u)g(s) ds g(u)h(u) du
L. (t))((u))) ()) du,
where the third equality is since < s < The last transition is by change of variables and the assumption that 72 (t) is a nonincreasing function. Consecutive application of the convolution map will then yield by induction,
((()((1) (1x''(--)l(or (((-w))())h(()) du.
Therefore, expanding the loss and using the upper bound, and denote q B11 + e)|]]]|| such that
< 1,
P(t) < F(t) +
Σ 9
(T)((T2(t)
j=1
r2(()
5((()
K (T'((t) - (l))))(()))) du
1=1
S5(1+
Σ
20
(2% (| + +))0-1
C K(Ta(1)) w)g(u)h(u) du
T2((t)
(1-1) x ( BY²(s) ((T T (s)) ds.
I +
The third transition follows from Lemma F3 with C₁ (TT and the last transition follows from a change of variables (s))..
x(0)
Lemma F.3 (Lemma D.4 in [27] and Lemma IV.4.7 [8]. Suppose f2 < B41 iee. K is monotoni-
2
cally decreasing. Suppose additionally that | and for some > 0, there exists T(s) such that
Lf K(s)K((L- 5) s 2(1 + {)((()) for all 1 2 T.
Then for all n 2
K (t) K(t)
K(O)) + +1) (2() + 1)
sup
We immediately get a corollary of Lemma F.2. Corollary F.1. Under the assumptions of Lemma F22, the following holds
F(I) + (
x x ( ds P(()
< FF(1) +C (
- (ss)))
where C. is the constant in Lemma F.2.
78

---

The goal will be to show that the lower and upper bound on P(t) from Corollary F.1 have the same
asymptotics (see Section F.2.2 for forcing function and Section F.2.3). Moreover, using the results from [80], know that K satisfies the Kesten's 's lemma condition (86).
Proposition F.1 (SGD and Kesten's condition, Proposition G.2 [80]). Suppose > 1.4 For any > 0, there is an M sufficiently large that for Bt E [M, d /M].
S. K(s)((( 5) s (2+ [2+
Moreover, have estimates on the Kesten constant |
Lemma F.4 (Boundedness of ||]]], Corollary G.1 [80]. When 2a > 1 and 72 1)
1
j-2a
|
(1 (1)).
2B
(T
When 2a < 1,
1 x1-2a
2B 1 - 2a (1 + o()))
We remark that | is approximately 28 result.
1 Tr(D); here Tr(D) = E}=13 2a, This leads to the main
Proposition F.2. Let v, be admissible. When 2a suppose
72(B 1) 1)<<2 <2 and T2 <
2=j=1
0c
j-
j 2a
Or when 2a < 1, suppose 9/2 x d where is any constant and
(B + 11 and 1 2a <1
Then the expected loss P is bounded and
I(T2(1)) + [ B)) K(T2(() (2())))())))))))(((()))))))))) ds P())
< (((2())) (( 3)())Bx((T2()) -2()))))))))
where C' is the constant in Lemma F22
Proof. The proof combines the previous results with the idea that the conditions the learning rates ensure that F is bounded and [|3]] < 1.
F.2.2 Forcing function for SGD with decaying learning rate schedule
In light of Corollary F.I, we can now define the different components of the forcing function. Following the outline Section D.I, we have three components for the upper forcing function F and the lower forcing function F.
T
(t) pp
def
exp(-2But) (28-1)/((2) du
2x
F ppp
(t)
def
exp(-2Bud c2pttBBut + B(B + 1)7zu²t) x (20-1)((2a) du
2a
F Fac(t)
def
C3
exp( (xp(-2But) x u 1/(2a) de
2ad
F (t)
def
C8
exp(-2Biie -2But + B(B + ut) x (200-1) du,
2ad
79

---

where cg def ビド=13 j-2 if 3 2
and 0 otherwise, and lastly,
j 20-22 23
/n/ R dr
Fo(t),Fo(t) def
-+j-200 2a d2a k(v/d) where & > is the unique solution
1.
1 1
1 + x2a
For the pure point terms, we get the following proposition.
Proposition F3 (Pure point forcing term, Proposition H.2 [80]). Suppose 2o + 2B > 1. For any E > 0, there is an M > so that for Bt 2 M.
3 (l) 2() N E x g(() and 3 Epp(1) g(() s x g())
where
g(t) def (2x) (2B)
1 + +11 x x x T( 2n
Furthermore, for any M > 0, there exists some constants C, C,C,c > 0 independent of so that
c c F (t), (t) SC if Bt < M
and ift > 2
(t) (t) < Ō Fo(t).
As for Fo. we have the following proposition. Proposition E44 (Asymptotic for Fo. Proposition H.3 [80]). Suppose and d are admissible such
that the ratio v/ d > 1 and suppose 2a + 23 > 1. Let < k(v/d) < be the unique solution
f1/d
du.
K. + 2a
Then as d + oc
d (Ejn
if >
Ft (t)
d d fo fo/d
u 2 Kty du, 2 1.
Lastly, we have a proposition for the absolutely continuous part of T (t). Proposition E.5 (Absolutely continuous forcing function, Proposition H.4 [80]). There exists
constant C(a,, () >> 0 such that
[C x Fo(), i > 2a < 1 23 1.
I (t) <
10,
Suppose now 2a > 1 and 23 > 1. For any E > 0, there an M > 0 so that for Bt E [M, d /M)
Fac(() (()] - & g( and
Fac (1) g S g())
where g(t) def J 23 (2B
T(1
1 )t 1+ 1+2/(22) x d'
2n
j=1
Furthermore, for any M > 0, there exists some constants
C, > 0 independent of so that
[Cxd', (t). s c Fo(),
if Bt < M
Bt 2 Md
Combining all these propositions, we have the following conclusion. Proposition F.6 (Forcing function for SGD, Corollary F.1 [80]). For any a, B with a, # and
1 there is a function C(t) bounded above for all so that
2
1 C(t)
(T
F vp
< (1)) 9(()
<F(T((0) & C() + F.c
Moreover, for any > O, there is a (s) large enough that C(t) < 1 1'2(t) E [M, da /M] and I2(t)B Md Lastly, the upper and lower bounds satisfy
Epp (1) ~ (t)( (ac(t) ~ Fac(t), and In(() ~ Fo()).
80

---

SGD-M
SGD
Fo(t) x d 20 max
Fo(t) = d d=
(t) x BBt
28-1
Fyp(t) x (12Bt) (143/a)+1/122))
C' x Fo(():
2 > 1, 2a <
C Fo((),
if > 1,2x > 2 < 1.
Fac(t) 5
Fac(1) =
0,
2 < 1
0,
21 1, 1, x d Bl) 11)
2f > 1,2x (. x d (12 (12BB)-1+1// 141/(20)
K p0) x 1 BI) B1) 2+1/
Kpp(1) x B 2+1/(2a)
Table 8: Asymptotics for the forcing and kernel functions for constant learning rate SGD/SGD- M. See Section G/Section for details/proofss of the derivations for these asymptotics of SGD-M/SGD. Here the constant C' is independent of dimension and Yeff 72 + 23.
F.2.3 Kernel function for SGD with decaying learning rate schedule
We perform similar analysis for the kernel function for SGD, To do so, introduce two (pure point) kernel functions:
K (t)
(t) def
2a / exp(-2Bul + B(B + 1) x u² 22-1/(2a)
du
and
K del
exp(-2Bul) x u 11(2))))) du.
2x
Jo
The first result states that the upper and lower Kpp are asymptotically the same.
Proposition E77 (Kpp asymptotic, Proposition H.5 [80]). Suppose a 1/4. For any 0, there is an M > 0 so that for Bt 2 M,
K (1) g()] x g(l) and [(p(() - g({) < x g(()
where
g(t) def (2a) (2B) 12(0(2 T(2 2n) x t 2+1 120)
Moreover, for any M 0, there exists constants c, C, Ō 0, such that when 2a 1,
< SC, es (5 SC C, &
and when 2a < 1, Furthermore, for any M. > 0, there exists constant Č > 0, such that
(tp < Ô x d Km Ò x da-1 Bt < M.
(1) S Ô Fo(l). (yp(1) S Ò [o( if Bt Md.
This directly leads the main result for the kernel function. Proposition E8 (Kernel estimation, Proposition G.1 [80]). Suppose > 1/4. There is positive function C(t) so that
Ctt) x (T, [2(s)) < xt(t) C(t) x 72(() - I2(s)),
and C(t) is bounded independent of d by function of M for all [2 d M. Moreover for any E > there is an M sufficiently large so that for [t) E [M, (M.d"/M) /M], C(t) <1+ Lastly, the upper and lower bounds satisfy
Kpp(t) (tp(t)..
Propositions F.4, F.5, F., and F.6 for the foreing function and their companions, Propositions F.7 and F,8 for the kernel function, allow us to simplify the Volterra equation.
81

---

Proposition F9 (Simplification of Volterra equation for SGD with decreasing learning rates), Suppose Ta(t) is a decreasing learning rate schedule (i.e., Assumption 4 holds) and suppose the assumptions of Proposition F.2 hold. Let a >> 1/4. First suppose that the forcing function with the decaying learning rate is integrable, that is,
e def
C(72 B) def [ 00
M(T2(s))72(s)
ds < 0c,
where h(11)
def (2a)
TV 2a + 1) x (2B1) (105))((1+/0)11/221 + + x (2B1)-1+1/00)) (+11/(2)) x d
(88)
Here 0g def (Ej=r3 (2a) (I(1 2a 2f > 1 and 23 > 1 and otherwise it is 0. Then there exists absolute constants M, > such that the following holds
=(Tx(0)+exx((<(() - P(t) S Cx (TT22(1)2)
for all T2(t) e [M, d M)
where e(v2, B) def fo ((T2(s)) (5 (s) ds and C > O is an absolute constants independent of and d.
On the other hand, if the forcing function with learning rate schedule is non-integrable, that is,
h(T ((T(())))())) da oc,
(89)
then there exists absolute constant M, M O such that the following holds
(()2(()) S P()) S C x F(T (1)/ for all T2(1) E [M, d M). where an absolute constant independent of d, 722 and B.
Proof. Let us first consider the case when (88) holds, that is,
e e= F(T2())))2(s) ds < oc.
Now we consider the upper and lower bound separately. Upper bound: From Proposition F.2, we have that
P(() < F(T2(()) + S. (((r(() (2())))))(( ds.
By a change of variables (v l22(s)) and setting u = l'2(t), we have that the RHS of the above inequality equals
F(T2(1)) + / x((Ta(() - (2()())))((()))))))))))
ds =s F(u) I(u) 1/ (((o)))) de where ( = 7T (v)).
Now let us decompose the convolution into two terms
1 K(u - u)F()g((()))) du = /"'
n//
K(u - ()(())((() de + I" K(u - ()u-v)((()g(())) du
Jn/2
(100 F(e)g(b) de + +g(u/2)3(uu/)
ra/2
1a/2 K(() de
s {((u22)
( F(v)g(e) du
11/2
s K(u/2) x
+ 7 x F(u/2).
Here we can bound 72 [| < since we are assuming sufficient conditions on the learning rate for bounded solutions. Now we need to consider the first term. For this, we know by Prop. F.6, Prop. F.3,
82

---

Prop. F.4, and Prop. F.5 that there exists M, M > 0 such that for all 2Bu E [M, we have
F(u/2) ~ h(u/2) + d 1o+xxx(1-2200 and Flu/2) E 1 for all 2Bu S M. It follows that 2/2 h F(v)g(v) do x L M
(2/ F(v)g(e) du + S h(w)g(v) do
(n/2 [M/ d (-2a-max(001-20) dv.
M
(90)
Here X means an absolute constant independent of d, 72. and B. The first and third integrals are bounded by absolute constants and the middle integral upper bounded by e. This proves the upper bound.
Lower bound: From Proposition F.2, we have the following lower bound on the loss curve which after a change of variables gives
F(u) 1- X(u (((u-o)((ugg() de S P(t).
We immediately have the following bound
X(u- du [()) [ I(v)a()) de x ((a())
/
H())((v) de
as X is a decreasing function and (()) are non-negative.
Next we suppose that
L h(T2(s))3(()) ds
(91)
Upper bound: To prove the upper bound, we recall that
12/ (1 s(w)3(v) de +
du S xK(u/2)
As before, we know by Prop. F.6, Prop. F.3, Prop. F.4, and Prop. F55 as well as Prop. F.7 and Prop. F.8 that there exists M, M > such that for all 2Bu E [M, M]. we have F(u/2) ~ h(u/2) + d mox (1220122 and F(u/2) S for all 2Bu I M. Moreover K(u/2) ~ u 12 and K(u/2) $ for all 2Bu & M.
It follows from (90) that
n/2 F(U)g(v) do S [." M F(o)g(e) du du+ l h(v)g(t) do [ d (=20--ax(001-22) dv.
12/2
(92)
Here S means an absolute constant independent of d, 72, and B. We also know that the first and third integrals are bounded by absolute constants. As (91) holds, we know there are two cases to consider.
First suppose 2o be that 2/3 < 1. Now we see that if 2a > 1, then
2/ 1 that h(()) (124))))+1((a) In order for (91) to hold, it must
1 a2/2 h(u)g(v) de ~ u 2+1 2/1/(2)), -145)))+11/200)) <u (1+88//))+1/00 ~ F(u/2). M
K(u/2)
Here Proposition F.3. Moreover always have in this regime that K(u/2) < Ff(u/2) (here used Prop. F.3 and Prop. F.7). Thus, all three integrals in (92) are bounded by F(u/2), proving the result in this setting.
Now we suppose that 2a < As in the prior case K(u/2) S Fu/22), showing the result for the first and third integrals in (92). It remains to show for the 2nd integral in (92). In this case, we note that the second integral
/ M ~ fa-1 2+1 - 2+1/(2a) (1+8/a)+11220411
4/2 M()g(()) do ~ 722 2+1/(2a), -1+t//))-
K(u/2)
83

---

Here we used that when 2a < 1, we need 12 2 d Now we see that the following hold
La-zu 2+ 2+1/(20) (1+8/a)+1/222))) x a -1(+///)++1/2a)) F(u/2)
+ 1 2
Since [M,deaM and -1 1/(22) 0, that the left hand side is maximized dea which shows that for all E [M,d
da-tu-111/000 2 (d SL.
Hence the result holds in this case.
Now we need to consider the case when 2a > and 2 > As we have already shown the Fpp (u/2) asymptotic is not the reason that (91) holds. Therefore if (91) is true, it must be because the Fac asymptotic is causing the problem. As before, we always have that K (u//)) S F(u/2). Thus, the first and third integral of (92) are done. For the middle integral, see that
d''u 2+ 2+1/2222/00 1/2m) 1+1/ S 1+1/(20) S F(u/2).
This holds precisely because 2a 1. Therefore the upper bound in this case is shown.
Lower bound: From Proposition F.2, we have the following lower bound on the loss curve which after a change of variables gives
F(u) S F(u)+ + / x(u- - () du S P(t).
This proves the result.
Remark F.3. While we do not have asymptotics for the kernel function when a < 1/4, we believe that the kernel function is not power law but rather exponentially decaying. This exponential decay would mean that the forcing function (which is still power law when < 1/4) dominates (above and below) the loss function Pt))-
Corollary F.2. Suppose 72 (t) is decreasing learning rate schedule (i.e., Assumption 4 holds) and suppose the assumptions of Proposition F.2 hold. Let a > 1/4. First suppose that the forcing,uunction with the decaying learning rate is integrable, that is,
e def
C(72,
B) def h(T2(s))22(s) ds xc,
where hitt) (2a)-IrrB
def
20 + 1) (L+8/))44)) c x (2Bt)-1+1/(20) d-1
(93)
Here eg del (Σ (2x) - 1n) 2f > and 23 > 1 and otherwise it is 0. Then there
exists absolute constants M, MM such that the following holds
[(Tx(t)+ < P(t) C'x (F(Tz(t)) (xX((r(())) for all I2(t) [M, MM
where C(72, B) def ((T2(s)) ds and C > O is absolute constants independent of 1/2 and
On the other hand, the forcing function with learning rate schedule is non-integrable, that is.
I(T2(s))2(()) ds 0o,
(94)
then there exists absolute constant M. M O Ouch that the following holds F(Iz(t)) & P(t) S C' ((T) for all I2(t) e [M, df where C is an absolute constant independent of d, "/2. and B.
Proof. The result immediately follows from F(u/2) ~ H()) by Propositions F.4, F.5, F.,, and F.6 and K(u/2) K(u) by Propositions F.7 and F.8 for the kernel function.
84

---

A 0.1, 1,000
10-
P2 -tr()
Figure 16: Comparison between SGD and SGD-M. Numerical set-up: Run SGD with constant SGD (14) and SGD-M with A 0.1 and constant 73 SGD-M (15) on the PLRF model with batch size 1. 72
SGD ranges (rl to 1.0/Tr((KK) (100 equally spaced points where selected) SGD learning rate 72
where Tr(K) Di= j- Learning rate 7/ T = A x SGD, SGD and SGD-M were run for 10* iterations and the last iterate was recorded. As you can see the solid (dot) lines, SGD, nearly match the same value as the 'x' marked faded lines, SGD-M. This agrees with our theoretical finding that have the exact same scaling law in the
show SGD and SGD-M under the equivalence 7/2 small batch regime.
SGD
G Classic Stochastic Momentum (SGD-M)
In this section, consider the classic (constant) stochastic momentum algorithm (15) under the simplified ODEs (43). In this case, 8, 72, 73 are constants independent of and ? = 1. In particular, derive the asymptoties for the foreing function, F, and kernel function K (see Table 8 for summary). See also Fig. 17 for a summary of the different loss curve behaviors depending on phases in the (a, -l-lae (Fig 12a for SGD/SGD-M M phase diagram).
To begin with, we need to give an expression for the forcing and kernel functions; specifically, we need to solve the ODEs in (43).
G.1 Solution to the ODE for classic stochastic momentum
Given the Volterra equation (55), we need to understand ( s) where x s) solves the ODE
dt Φ (t. s) = 2tt (t, such that x s) = Id3
(95)
where for this constant momentum algorithm, the matrix 2itt, A) reduces to
Bx
0
-273
Ω def A(t:A)=
0
22
2BA
(96)
BX
73
-5
B>
Since 22d is constant (independent of time), a solution to (43), (x s) = exp(S x (t s)), see [25, Chapter Section2] for systems of constant coefficient linear ODEs. As a result, we just need an expression for exp(12" t).
Computing the exponential of matrix can easily be done if one knows the eigenvalues and eigen- vectors/generalized eigenvectors of the matrix. For example, the eigenvalues of Ω are distinct the eigenvectors of Ω form basis of R3, then we can write exp(S2" x () V((V)' where D(t) = Diag(exp(r; x () : i = 1, 2, and r; eigenvalue of Ω2*) and the columns of V are the eigenvectors of Ω-, the case that this does not occur, then we can use the Jordan form of Ω to find exp(I?= x t). It would involve computing the eigenvalues and generalized eigenvalues of Q
G.2 Asymptotics of the kernel and forcing functions, SGD-M
We can now give an explicit solution to the ODE (95) and use it to obtain an explicit formulation for the loss curve S under the SDE in (40) (see also simplified ODEs (43)).
85

---

Phase Ia, lb, Ic
Phase II
Phase III
compute optimal
compute optimal
compute optimal
loss curve
F pp
loss curve
F pp
loss curve
Ppp
Fac
F
Fo
Fo
flops
flops
flops
Phase IVb
Phase IVa
compute
compute optimal
Ia
II
loss curve
F pp
optimal
loss curve
F pp
111
high-d line
lb
IVa
z-yl
K pp
K
0.25
TVb
3o
Fo
0.5
flops
flops
Figure 17: SGD-M (and SGD) cartoon plots and phase diagram. The loss curve the sum of P(t) pt) T oe + 1 (yp + Fo)),, where 7
def 722 for SGD and y
def 72 + 28 for 8
SGD-M. Each of these terms are explicit and take the form -- (See Table 8 for summary of the asymptotics of the terms and Section G.2 for derivations). Depending on the strength of a relative to 3, some of the terms drop out (Sec. E.1). Compute-optimal derivations can be found in Sec. E.
8 + 72 Bo² μ def Define def (8- 72 4/3 Bo², P def
72 Bo²
Then the eigenvalues of 9 are p, p + Vw, p p Vo and we write
(t- -)p (- Bo² + - (1--) 2Bo2 u/W
2w ((=-))/a[( + 2Bo 73) + μv/w))
4Bo2
e -t- (t- ( ((1 s) + V) - 1)
2w
273
(t, e (t- (t- -(--s)/w)) 1)
Here V E C is defined by e²
[w+2Bo2ns)t
(u+2Bo
We can now give precise bounds on क (t, $11, क (t,
In the following, we will restrain the study to some values of o, 72, 7s B. Assumption 5 (Hyperparameter domain). Suppose that > 0, 12 > 0, 73 0, a (0.2). BEN and that additionally
Bo < 8 73 8 max
4'
Proposition G.1. There exists C' > 0 such that under Assumption 5 we have
82
S W
C
V E R,, and more precisely
w
< <
<
Cw
Cr²Ba
22
86

---

Proof. We clearly have 8² 2² w 2 (g5) & c. Additionally, is clear that (w + 23)2
NVW > O. This implies that V/ E R4 and more precisely
(w + 2B)2 + uV
(w + 2Bo2gs)) - M//
82
- 12022-23- 273 Bo
(6 - 7/ Bo 16 (6
72 2) 4:g Bo²
82
6 2202/3 - 273 Bo² (6 72
Bo
Bo²
[6- (6-7218'11
(6 12 B22)2
4rgs Boa
(o2)?
8a
(73B02)2
w
B2
Proposition G.2. There exists C >> 0 such that under Assumption for any t 2 2 0 we have
1 doll- S (t, s)11 S (C(t--)
1
3
(t-s)( (t, 82 S c 73 clt-s)((55-))
if addiiioonally - 8 2
8'
C
8
8²
1
2
((t - s)o)² - (t, s) I C 2
while ift 8 <
((t s)8)
8'
C 82
82
In particular one can also choose C > 0 such that
1
(t-s)(+2
273 )Bo S (t,)s) S Ce C(t-s)((mz Poo)
1 11 d 270 |Bo < x (t, s) s c 23 -(tts)((rat (20) BB
additionally t
C 82
Proof. Since V >~ aB
> 1 know that
4Bo²
Φ (1,
e -1-- (coss -) + V) 1)
2w
4Bo² 73 -(--x)p exp (11 + +V)
2w
x exp((1 s) - -)).
For the same reason, if (t - s) (t ss) 2 1, we have that
2 (V(=-))
(t, 8)12
while if (1 ([=8)5 < we use that ( - 8) ,)))}² < cosh((1 - 8)6) - 1 s s))82
Using Assumption 5 we obtain that there exists C > 0 such that
87

---

8 + C
72
+
1
B
2
Ta 2aBo²
Bo²
V
6
2
0
0²
1/2 Bo²
12 BB?
- 73 BB
8/ 1 -
+
8
82
12
/
12
+ 2 23 2B
Bo 82
S 8 0
2
where used that 1/2 222
Bo Bo?
4.22 Br
Ba
< 0 is bounded under Assumption 5.
Additionally, since 2= o < t. we can further reduce to
273 8 )=²
2n3 Bo² I vw -p + B
22y +
8
This implies the last result.
Now that we know the behavior of 2n s), (t, s) we can reduce the study to simplified functions that will be used to bound the forcing and kernel functions. This is done in the next definition. Definition G.1. Let E > 0 be a large universal constant corresponding to Proposition G.2. Then, under Assumption 5 we define
(002, 1) def st (a2, def
e - Br2)
e( Bo²t
35 1) del Ce + Bo
(x( 63
C (
1,26 + ((t ((t-s])f (()=
ox(o2 c
2 +m) Bo 128-1 + ((1 ((1- ((1=8)6)21{6-1)
e
Then it is clear from the above that VI 2 2 0,
4 (o t) x 0) S (o)
2 (02 5) 5) < (1, - 00 s).
We hence define:
F(t) _ 2i0² t) (d(2)
xo
F(t) Jo vi (a t) 2z(do²)
F(t) Jo Wr (a 2 t) 2
(() = B x p (7 + ( Po ()) uxx(d=
(t) B x Jo (12 (0 t) by (o (o²tt) pex vx((do)
(I) = B x Jo (1 + ox (2)) uxx(d=))
From the above, we see that F, F F are non-negative functions and that (C,()) = 5c( - s). (,) = K(t - 8), (1) = K(L 8) are non-negative convolution kernels. The following proposition shows that understand the behavior Pf only need understand the solutions of the Volterra equations using the bounds on the forcing and kernel functions.
Proposition G.3. We have the bounds for all S S
0 ( < F(t) $F)),
0 < kt < (t) s (().
88

---

In particular, denoting P.P.P the solutions of the Volterra equations with respectively forcing function J, I.9 and kernel function K.K... K, we have the inequality for all time 2 0:
- 2(t - (()) - P(t).
Proof. This is consequence of proposition G.2 and the fact that μs,, Mu positive measures. The last inequality is a consequence of the Volterra equation definition.
Denote now, similarly to was done in Section D:
Fo(t) pe (00))) = ne (10)),
F (t) del 1
// -2 01t² 20_2 (00 do,
20
(1) def
CB
0² 000 do,
(97)
2nd
12 + (28)
B
(t) def
-3 202 t)
2a
We can compute the following results on Fo. F ac: F pp, Kpp which are summarized in Table 5. Proposition G.A. Suppose a, 33 + 3 and 2a + 23 > 1 and Assumption 5 to hold. For any M > 0, then, there exists a constant C(a, 8) such that we have for all 2 0:
j d 20 (0,1 S FF(()) S Cd. -dt-datttat 2a
min(1. ((12 + 73 -1- 20-1 (12 + (7)) )BB
8 ))tB)
2a
} I 5 (t) S Cmii(11 (112 )Bt)
73
-1-
28-1 2a
C
8
C min (12 + 23 )t)
-2+
} S (t)
S <0(1 + (25)))) min (112 + 33 )Bt)
-2 Ic
5
Fott
2a 1,28
Fe) -
0
if < 1.
If additionally 2a 1,23 > 1 and (12 + )) < 2d we have
dd dd min (12 + 23 ))B) g
112/
} J 5f (t) S Cdd1 (min(1, (12 + 13 g
)
-1+1/(2a)
Proof. The proofs are entirely similar to the one of [80, Proposition H.2, H.3, H.4, H.5] and just
follow from the change of variable (72 + 2 Bo in the integral definitions (97). For Fo. we used Proposition D.3.
Suppose a > 0, a. 3 # 2/ and 2a 23 > a +> 3 and let M 0. Then, there exists a constant C(a,, 3, M) > O such that we have for all t 2 s 0:
Proposition G.5.
1 C (oo(t)+ Foo + F (t) < (1) C (F((t) F. + 5 po (t)
If additionally a > } and (72 + Btt 5) < Md00
1
pp 8) s K(t 8) s C'ppp(t - s). C
89

---

Proof. The proof is a direct consequence of Propositions D.13 and D.14 and the estimates derived in Section D. We apply Proposition D.13 which implies the existence of L, C >> O such that
_/ /1t t
1
+ (2oo) -
C Ld-20
Ld-2a
$cf
t
sc
Óio², (u + (doo))
10-20
Additionally, we know using Propositions D.4 and that
od 2a
/
oo 2x
max (t)ss) + Pr... (do"))
310² ((ostt)ps (do ((o=))))p,p((d))) 2 Fo(t).
Finally, using Proposition D.2, it clear that
maxt S "602, ,us + (doo), I (do))))) 53 Fp)()).
1/M
1/M
For the kernel, we can similarly write for L > 0 large enough
K(( - 5) = (72 + (28)))
L' $602,t +(o2,t-s) - -s)xx(ddo)
20
- Ld
$102t (oo + fra- ( + (' 1602, -
1/4
For the second term, we use Proposition 1.14 and that (72 + 3 B(( -) 8) < Md to compare it to the value in Proposition G.4. For the first term, using Proposition D.6 and 2 (t 8) < Md it is absorbed by the second term. Finally, the third term is absorbed by the second term due to the exponential decay of v.
Now, we can easily bound the forcing and kernel functions using F,K. Proposition G.6. Suppose a, 3 + } 2a + 23 > 1 aa + > B and let M > 0. Then, there exists
constant C(a, 3, M) > such that we have for all 2 0:
C ())()) + F (t) + +F (()(1) S I(1) s F())= s (Fo(l) + Fac(t) + Fpp
If additionally > 1. (12 + m3 Bt < Ma and t 2 8-1
C K (pp(t) S K(t) s K(t)= s (t).
Finally a > 4. (12 + Bt < Md and S 8-1,
?/ 12 B min [1, ((12 + 2 Bt) 2+1/(2a), + BB
(
82 ((t -)))}]
<x(() S {((()
)1) (B))-2+1/(20)) + BB 82
2 min(11 ((72 + 73 BI) SC
((1 -3)8)]
Proof. We know that for any t. i 0,
g(t) F(Ct), F(t) (.
if additionally t5 2621. 1, K(t) = e ((((t), K(t)= - ex t).
1
Additionally, we know that (5o(1) + (() + - 3p)(() and (1) follow power laws. Hence, we only have to prove the bounds on F and K which was done in Proposition G.5. Finally the case t6 < 1 can be handled casily since (t, s) ((t s)0) and we obtain the result.
D
90

---

Proposition G.7 (Kernel norm bound). Let a > 1.
4' a + 1.
2' There exists some constants C(a) >
O, d(a)= > such that for > 0, B E IN" 92 > 0, 73 > 0, BB - 16. 12 $ 16 and d 2 d(a).
5a
8²
we have the bounds (note that for 2a < 1, use C d)
d(1-2a) 20)1
(72 + 23 $ |ل||| = |
8
S Cdll -2))+ (72 +
73).
8
Proof. We will instead show
d 20)
7/2 + (20)²
< | - |
C
72- 8
(12 + ( 1B
S Cdd(1-ai)- -2a)+
(72 + +
B since a² a+b = a and 12B م 28
which directly brings the result Va, b > O we have
82 2 1.
is clear that ||||| - | The bounds are then direct consequences from Definition G.1 and the lower and upper-bounds on jox derived in Section D.
We first use Lemma D.I to get that there exists some d > 1 large enough such that Vd 2 d. (2, oo)) = 0. Hence we reduce to a S 2 and under the assumption of the proposition we see that we are under Assumption 5.
For the lower bound, we write for some M > 0 large enough to apply Proposition D.12,
K
2V
(12 + ( 13 ) (rra F BB''t (yxx(o)) de
1=
Ja Md
100
(12 ( (23 ) Be (ru+ B)a't dus, (o²) dt
2
1=
Md-
12
(2)) 2
3- do
72
Md=
+ (2
2 [1=2a)+
12
72 + T3
For the upper-bound, we proceed similarly and write form some M > large enough to apply Propositions D.5 and D.9 to D.11
13 22) Be (Ta+ (1a+ 2
B)o' duy (o²) dt
(12
5-
0=0
(12 + (25 )) Be - (ra+ 7 Bo² djip (0²) dt
2
1=5-1
8
12 + )2)
3- do
72 + 73
0=
S (1 (2))) 12 (2
72 24
For (t,s,, s small is not a problem as we have directly
Be (tat +
duxx dt
1=0 6=0
7/
S (1)))
7/2 + 73
91

---

For ()2(t, s) we need to be more careful and we write for small times
(
B
2
=)))²
d dt
1=0
2
m du ) dt
1=0
o=
$ (8 2 B8-1
-3- do
8
0=
x d 20)+ 13 B88
Hence we see that ||]|| S (12)) 22))
which brings the claim and the + (2()))) (2)) which
72
result.
Using the bound Proposition G.7, we obtain the following sufficient condition for the algorithm to be stable.
Corollary G.1 ((Sufficient) Stability condition for hyperparameters of SGD-M). Suppose the iterates
{31-11 0 are generated by (15). Let > O, 1/2 0, 7/3 > 0, B e N Let a > 4. a
# 2' B B
3. 2a 233 > a + 1>>>>> B. There exists some constant > O such that if max 8 } S 16 and 72 + 73 remains bounded, i.e. ||P co.
82
a I cd c1-21- then the solution P to the simplified Volterra equation Equation (55)
Proof. Notice from Propositions G.5 and G.6 that F(t) is bounded. Additionally, applying Proposi- tion G.7 yields the existence of a corresponding c such that ||]]| < 1. We saw in Section C.3 that this implies that the loss remains bounded.
Remark G.1 (Necessary stability condition). Similarly, it is clear using the lower bound on |X] from Proposition G.7 that 12 + 23 2 Cd-(1-20)), for some C(a.) > 0 then the risk is unbounded iee. || 00.
Proposition G.8. Let 8 > 0, 72 0, 73 O, BEN" M 0, >> 0. Let 11 #
There 8 S 8 and 72 + 2 S cd 20) then for
exists some constant e>> 0 such that if max (72 B, B any 2 0 with (7/ + Bt < 2d we have
To
K + (t) S {(t).
Proof. We apply Propositions G.4 G.6 obtain that
if:25 5- K(t)= (72 73 )) min(1, ((72 + 73 ) -241/(20))
8
ift < 8- K(I) 2/ B min 1, ((12 + 7
x BB) - 2+1/(2a) } + 5 ) s)0)3.
A crucial observation is that K behaves as power law, i.c. there exists some constant C >>>>>>>>>> 0 such < C. Indeed, if t < 5-1 2 or t 2 8-1 it is clear. This is still true for
that Vt 2 0, 1/4 < X())
T(()))
e 8 by just noticing lim4ys- K(t)) x K(6-1) ) since (72 + $
It is additionally clear from Proposition G.7 that for any > 0, for e > 0 small enough we have | <.
Using these two previous facts, one just writes for any 2 0
92

---

[ * (t) = !. K(t s)](s)
1.²0 K(t -s)
112
K(t - s)){(s) ds
112
2 x(1) '
K(s) ds
2 (t).
By decreasing as needed we obtain the result.
We finally state result on the foreing function norm (see related result [80, Proposition H.6]).
Proposition G.9. Let a > 0, + 2, B * 2.2 a+1> > B, 2a + 23 > 1. There exists such that
if max (12 Ba } I 16
6² and 72 + T & cd od-(1-2a)= then for any 2 0 with 1 I (72 + }) ))Bt <
Mf we have
1 1, 50 (() ds (x
1/23 < 1, (72+ Tu+ty18 K ml(() S (1) So (s) $ F(1)
(12 3 JB K Rm)().
Proof. First case: 23 > 1
First the contribution of Fo gives So M 1(va 77) Fo(s) ds S (1x 3J because Fo(() = 2-
It is clear that Fyp E L.).). Hence since < (12+ ( Bt we know that Jo J (s) ds x (72 4 )B
For F ac: if 2a < 1, we handle it like Fo. If on the other hand 2a > then we write
/ 1(1+ B)
d
Md
ds) S
x 1 dx
(72 BB
(02 7 y
(72 + 2B
Since for (12 + Bt < we know that Fao(t) x d we obtain that
Jo Ma /((x2 2 )B} (sc(s) ds S (72t 5)B This concludes for the case 28 > 1.
2nd case: 28 In this phase, we do not need worry about Fao since it is zero
It is also clear that (Ta+ (na+'y)B ppl(() x Kyp(1) So (s) ds since So (s) ds 2
Jo J (s) ds 2 F)B
We first consider the F contribution.
(t) L'' (s) S (t)
(L (
(s) ds S" (s)
1
Kpp (t) + KC (pp S" F (s) ds
2
Ben
Tes
B K (pp(t) + ven B Bt) zu
1
Bs)
20_2__1 ds
xpp (pp(t) + Yea Bt) 22 du (Yerr Bt) 20
1
Ban
93

---

We only need to show that
(enn(Btn 22+ 2u (ten
2
2 3 pp((t
Temr(Tenr Bt) 2t da
(TerBt)
28-1
1- S (t) B1)
28-1 2G
X Jeff Z (1amBt) 1-
Since by assumption Joft Bt Mda that only need Yett E d2a-1 This is true since by assumption "jer X d=)1-20) da-
We now consider the To contribution that we will bound using F pp (!).
S Fo(s) ds S pp (1)d-2a-22)) 2a-22+11
X kpp(1)
Ten Brea Bt) Bt) 2+ z0 d 2a 28+1
I (en( -1 da d da 2--22224 23+1
We only need to show
(1) x J. Fo(s) ds S F(1)
rem Bt) 1+ t d 2a 2 (Yea -1 -1-28 29
Yell S (Yetr B1) _ 2
Bl BI 2 we only need to show that Yeff S dd2a-+229 which is true since Yeeff / (1--222))) by assumption and (1 2xx++ 2 (1 2a) 23.
Since Yefr
Theorem G.I. Suppose > 1, 2a + 23 > 1, a. a+1> B. Let M > 0,6 > 0,
0, 73 > 0, B E N". There exists some constant c. >> such that if (max(12 B, BB
5 S 16 and
72 2 I cd (1-211+,) then for any I 2 0 with (72 + Bt S Md and I 2 8-1, we huve
(To(t)+ 0c (t) pp (t)+ (12 + 3 )
1
(t)
(t)) S P(t) <
1 (F. Fec (t) pp (t)+ (t)
(72 + 2 )B (t) (t))
c
Proof. Using Proposition G.3, it suffices to prove the result on P(() and P((). We already know that V1 0,
06
I(I) J P()) P(() F(1)
F K
(1).
Let 3. Using Proposition G.8 know that for > 0 small enough, Vt 2 0if (12 3) S 2
π + K (t)
and ||]|| 1.
Hence, applying Lemma C.3 we can bound for some C' > 0
(t) F* (t).
We only have left to respectively lower and upper bound [ K. (t) and F K (t).
94

---

We write using Proposition G.6, and for 2 8-1 and (12 + +B) 2 1,
6 1/2
F(s))((,)) S Ft))
3 K (1)
S kml(t) x 1"² ((s) + (Fo(t) + (t) + pp)(t) x |
Fol( Fao(t) + (t) (72 B) x (t)
Similarly for the lower-bound we write
4/2
(1) 2 ((t) ss
2 (1) /'² 1/2
F())
(72 + 2 (t)(t).
Here we used from Proposition G.9 that if S (12 + )t S Md then F(t) + K (t) Fo
F()) x
F(t) + (rat ()B pt)
Finally, using again Proposition G.5, we obtain the result.
Remark G.2. Theorem G.I together with the asymptotics in Proposition G.4 shows that SGD-M with learning rates 1, 72, Trs and momentum parameter has the exact same scaling laws than SGD where 722 (72 )))
H DANA-constant
In this section, analyze the DANA-constant algorithm introduced in Section B.3. We interested in deriving the forcing and kernel function as well as their asymptotics (See Table 12 for summary). As in SGD-M/SGD, depending (a, -laane, the loss simplifies. We summarize the results in Figure 18 which shows the different phases and a qualitative description of the loss curves.
We begin with description of change of variables in the hyperparameters and then go into details about solving the simplied ODEs in (43) with hyperparameters associated with DANA-constant. This ultimately leads to the asymptotic description of the forcing and kernel functions.
First, since (1), (l) are constants, DANA-constant with hyperparameters (71, 72, 73, I47) yields the same algorithm DANA-constant with parameters (7i, 72, 73. T+z) where chose % def def 73 73 X 71- Indeed, one can check that the updates of the two algorithms on the 0 variable are identical, and also that the forcing and kernel functions are identical. Hence across this section, we will freely use 71 = without loss of generality. Additionally, throughout this section, we use o² to denote the eigenvalue À of K, that def x²
Below we introduce the main parametrization for DANA-constant that will often be used throughout this section. It essentially amounts to consider (DANA) with K3
0.
Parametrization H.1. Let vector of hyperparameters H def (52, Ta. cb, K1, K2, Kb. 8) with
ÿ2, ña, Cb > O, K1, K2, Kb 2 O. We add the restriction H& S min (ii,(x1, K2), 2a S - -K2 -k2 Kbb S O. We parametrize
1,
3/d Ni. Iss "2 B coddt.
(98)
Remark H.1. The reason require Hi S min(RI, (1I, K2) and Kn, 62 2 is to ensure that 722 B, 733 B remain bounded as d + co. Otherwise, eigenvalues of (23) do generally longer have negative real
95

---

part and the algorithm would trivially diverge Additionnally we require -2a S -k2 2kg Rb S 0 B S This condition is in fact not very restrictive on most scalings of interests. We mostly make it to ensure that no edge case for particularly small schedules or very large batch creates problems the scaling laws. particular is satisfied for any B S d, Ys3 x 1/2 3. d= 2 722 E d1-12an) which includes DANA-constant with batch B = in Section 3. In particular note that Parametrization H.1 allows for B z d which the reader check will allow for outscaling when 2a < when reported in Theorem 4.1.
ensure d 20 - 78
H.1 Simplification of the ODE
In theory, the Frobenius method to get asymptotic solutions of an ODE can be applied to Equation (22) (see [25, Chapters 5] (with some care near zero for the 1++1
term). However, the computations
are cumbersome. Instead we will work with the simplified ODE
0
0
0
0
28
0
2o2
0
2ya
(t) dl
0
0
-8)
+
0
27 Bo
(99)
1 1+0 I
Bo²
(79(t)
Bo?
(t)
To simplify the computations, we will additionally do the change of variable
(()) def
31 Bo (t)
Pt(
on Equation (99) and get the ODE
( 0 0
0
26
0
(t)(t)
0
0
Ba²
0
0
2 nan Bo
+
0
0
2 71 Y/
(()).
(100)
dt
1+t
Vo
-vn Bo
Bo²
The following technical lemma shows the decreasing behavior of the norm of the solutions of ODE (100).
Lemma H.1. Consider the ODE (100) on + def (X,Y, : (-1,00) - R3. Then we have the identity for any > -1
(t) (t) + (t) = Bo - 5 (2) - ( [1t + 72 Z(t)
dt
+t
Proof. This comes from the fact that
(x)
-/ BoX())
(v (t) -444 (2) (4
d2 de -4,4 +4173 (() 4orssn (r2Bo22?(1))
H.2 Getting asymptotic solutions of the ODE through Frobenius method
In the following, we state two strenghtened results from [25, Chap. 4, Thm 4.1; Chap. 5, Thms 2.1,
4.1] I] to get uniform estimates in parameter space of asymptotic solutions of (100) for small and large
t.
96

---

Theorem H.1 (Singularity of the 1st kind around 0). Ler Z C R" for some n 2 1 and define for any
Ç Z the ODE:
dec de
(t) = ( R + +AS) &
(t) where Φ R + Ms,3(R) and RS, A E (M).
Suppose the existence of 6 > 0, C. > O such that JÇ Z. R Diag(ri r2,rg) is diagonal with
min d((r; -r51. N) 2 8 and that ||RS loc ||A5 loc S C. Then we have the following:
1. J4 E Z, (t)
(Is [I++ Exz P& th6 is fundamental solution where noting Po = 1a. the matrices Ps are uniquely defined by the recurrence relation Vk 2
0, k+1 [RC + (k + 1)1a} = R P + AS PS Especially, wk 2 0, k,, A E Z, ||P S Ck.
2. As a consequence, AK Z O. (8, K) > O, C) > 0 such that JÇ E Z, the fundamental solution Ho verifies VI S T, Vi E [3), | $C.K S DR where
$C def (Is + EiSK p(() ERS
Theorem H.2 (Singularity of the 2nd kind around infinity). Let Z C. RD for some n 2 1 and define dæc de (t) = At + AS) (t) where Φf : R2 + Ms,a(R) and R AS E (R). Suppose the existence of 8 0,C > such that VC Z, A Diag( p22 16) is diagonal, minisi/isa -u5 1 2 6 and that || ||A5 S C. Then we have the following;
for any Ç E Z the ODE:
1. JC E Z, (1) = 11 + Ex21 +t) ( es is a formal fundamental solution where Ri
)
8i R and noting Po 13, the matrices Ps are uniquely defined by the recurrence relation
Vk 2 0, Pf AS A ps R$ P P fts + kPp. Especially, Vk 2 0, 0,3CkrvVC VC
2, 1 S Ck.
2. AK 2 O, 36, > 0,3D > 0 such that JÇ E Z, there exists a fundamental solution
Φ such that Vt 2 T, Vi E [3). | - $C.K S -K1) where & def
Iss + EeSK 2 f ceas.
H.3 Fundamental solutions around zero and infinity
In this section, we show asymptotic solutions & of the ODE Equation (100) for small and large times. We remind 71 1. To get bounded coefficients in the ODE, we will differentiate the two cases Bo < B and 72 Ba 2 V4 The first case corresponds to small and hence t large with respect to the dimension, where there is acceleration. The second case corresponds to large, hence small and dynamics similar to SGD. Additionally, getting asymptotic estimates using Frobenius method around xo requires the eigenvalues the leading order matrix to be distinct, since the singularity is of the second kind. Hence we will introduce E (0, 1) and first restrict ourselves to the case where 7/2 Ba & (1 /4
1st case: Bo S B. We apply the time change 7(t) o + t) and defining (()) def (t) we obtain the new ODE
(
0
0
0
( def dr
R OV +T
0
22
0
8
-2 Bo VTBB
-2
A (((T)
0
0
((T).
1BB +
1
-1
72 Ba
13
(101)
Denoting again R, A respectively the left and right matrices, we rewrite A = TDT-1 with D del
Diag(P1, Ha, H3) with #1 20 Ba
VTaB
4
3
73
Roy
42
22 Ba
B B
+
-4 +
23
73
B
p3 =
97

---

Ba
In the limit Ba Ba 7B + 0, the eigenvalues are distinct. Hence the matrices T,T* can be
VTBBB
chosen analytic in na Bo v
This implies
72 Bo
T2Bo
T i
i
T
+ +
(102)
V13B
vaaB
1 0/
Notice that the matrix has bounded coeffiiients and distinct eigenvalues bounded away of each other. Additionally define 81 def
(1-
(1 V
B0
82 def
1+
Tre Bo
83 def
B222 4B
In the next proposition, we apply Theorem H.2 on solution & of Equation (101) to obtain an asymptotic solution for large time (uniformly in parameter space).
Proposition H.1 (Asymptotic solutions around infinity). Let e, E (0, 1), C > 0 such that ō E (0, C). There exists a constant M (e, E, C) > 0 such that V 12: 73, > 12 (1 (1-e)y4raB. there exists a fundamental solution h x0 of Equation (100) such that for any > satisfying + (1
M
P +
P₂ (ov +
T (t) (Is
DD(o, B)
ov B( +t)
(oV7aB(1+1)) -5 e
(( +1))
where D(o, 73. B) Diag((o
CaeellooaBBBBa B11 ( i 1,2,3). Moreover the matrices P1. are uniquely determined from Theorem H.2.
Note that above could directly bound the matrix norm of the difference since the module of the decay is the same on all eigen-vectors (Re(ô,) = Re(d)) = Re(0g) and Re(H)) = Re(u2) = Re(ua). Additionally, if 6 is not an integer, we similarly apply Theorem H.l to obtain asymptotic solution for small time (uniformly in parameter space).
Proposition H.2 (Asymptotic solutions around zero). E (0, 1), C> such that (0, C). There exists constant M (e, é, C)) such that V72, "3, B > if 72 Bo S /4raBB and d(6, N) e, there exists a fundamental solution T0 of Equation (100) such for all t > 1, B M
(t) (Is + 13 +1+) (oV +1 P2)
0
((o (lo
B(1 + t)) -26
((o B1 + 1)) 1))
0
0
Se E (oVTaB((+1)) 13x3
0
((oVraB(1+t))
0
0
(10 B1 +t))
Pt. P2 are uniquely determined from Theorem H.l as
%
0
v V
I-
P
0
0
I+4
6+1
I-6
V
1
2m
dig
2
x²
6+1
(2 2-1
2
and
P2 =
6+11126+22
I- 3
(6+2)x
6+2
[1 (1-8)2-88 8)m
2
3+1
VrsB where
98

---

2nd Case B2 2 V4yB. Then we apply the time change 7(t) def
?2 and define ((7) def
((). We obtain the new ODE
0
0
26
0
2
0
y 12
(T(T) dr
0
-
222 Bo
0
0
Ô(T).
"2Bo²
+
Ba
1 Ba
Bo
Again the matrix A has bounded coefficients and distinct eigenvalues bounded away from each other. We rewrite A TDT with some different matrix T and eigenvalues jni -1
B 6
+ L, 42 = - + 72
B
+ 1,
μ3 -1 Denoting x
/TaB "2
we have
1 x² +0
(1))
x +x()
22x +0 x
T
x² +0
1- 1
2x +0 (z)) 1
-x
-0+0
1 3x +0 (7)
+0(22
2x + + 0 (x3)
T-1
x² + ##400
1 1 3x +0(g)
2x + r
x +
x 0+0
1 4x +0 (x3)
Remind 61 def -8 (1
To Bo
82 -6 -1+
Ba
83
dtf
5.
We can
1/ 122
122
again apply Theorem H.1 to obtain that (uniformly in parameter space):
Proposition H.3 (Asymptotic solutions around infinity). Let o, (0,1), C such that 8, (0, C). There exists constant ((e, é, C)) such that V12, 73, a, B. > o 1) > M and Bo > (1 + +))4/4r there exists fundamental solution o0 of Equation (100) such that
P2 T- (t) - (I3 + 12 (12Bo?(1 t)) (1=))2(1+
+
)
S (12 t)) 1axa x D,
where
(12 ( maRiin
D def
(ra + (0) eo
0
(m2 (1+ goDottus
+())
(12 + 1) ga ust
0
D def
Bo + 2))
0
0
Bo (1 + t))
bo
Moreover the matrices Pt, P2 are uniquely determined from Theorem H.2
28 222
2(6-3)x
-2
0
8-1
2x
(6-1)2
36-36 +2
P1
0
2x
P3
2
2
6+1
16+
(26+-))
4- 2_
8+1
[6+1](0+2]
31-3642
2(1--1)
Additionally, if 28 is not an integer, we can again apply Theorem H.2 to obtain that (uniformly in parameter space):
Proposition H.4 (Asymptotic solutions around zero). Let c, E 1), C> such that E (0, C). There exists constant ((,, ē, C) such that V-z, Y3. a,t. B > 0 o +t) < M, Bo 2 /4ra B, and d(8, N) > there exists fundamental solution 80 of Equation (100) such that
(b(()) (Is +
(72B0² +1) PayD.
< E (2 + 1))² 1xxx
99

---

where
def (
0 ((12 (1 + ()) =
0 ((12 (1 + +))++))) f))
0
Moreover the matrices P1, P2 are uniquely determined from Theorem H.l.
H.4 Behavior of (t, and (t)
The goal this section is to derive bouds on & (t, and & I12 s) where {(t, 8) (0), 2 s) - (R) denotes the solution of the IVP (99) with initialization o(s, Iz.
This implies the initialization condition on Equation (100) that is ((s, s) Diag(1, T vBe consider several cases.
v ) We
1st Case: 22 < (1 <) AraB. Depending the time and we get different estimates for the values of 4 (1, s) and 8)).
For ov BB (1+ + t) S 1, we have the following estimate:
Proposition H.5. Ler E (0,1), 1), C> O. Let d E (1, C) with d(6, N) and let 93, B, ler < $ Σ t. Then, there exists a constant M(c, é, C) such that if a BBB +t) < M, and
B2 < M B. the function & (t, = 1 = and, if additionally & < M, the function 112 s) Bo (3B s))² (1 E). If S T3 + t) S 1 or M 4 S <yBo s B4 or M 1 S 1, then (t, 8) (1) and P12 O Bo (B +))²
M
Proof. We apply Proposition H.2 to write that solution d of Equation (99) is:
s(t,s) ('()) Diag(1,
73
73
B
v
Bo
= 1) (o / B( (P2 +c) (ov B( + 1)) 24
Ba + +())
Diag
(, 5))
(oV7gB(1 + +))) + + +8)
11,
73
V 73
x Diag
Bo
Bo
|t
[+( [+s
(o + +0)
( (oV73 + +))) +c)
1 1tc
Be
(ov% (1 s)} (((Pp P2)12
x (oV a)) (((Ff Pa)a te)
11 8))² ((1P)ss =)
-o0773B1 8))² ((P)ss
(a
BgBc +0 Vhae
Notice that p P2
Additionally, 8 > 1, we have that
Lis 8-1 1++)
< 1 for M small. Thus there exists some M small enough such that if a BB < M and 7/2 Bo < M AT then (t)(t, = + c. If, additionally, S M, then $ Piz(t, s) x Ta Boz
(1 s)) (1.
If one of the conditions is fails, still have that the coefficients of the ODE are bounded for o + t) < 1, 72 < V47BB, and & < Hence we get under this condition that (n, s) (1) and (t, s) O+ ( ( (o (o B(1+ s))
For the case o / + and a BL + s) 1, we have the following estimate:
100

---

Proposition H.6. E (0,1), C. > 0 such that å E (1, C). Let 12 73. B. > 0 and suppose
0 S t. Moreover suppose that d(6, N) E and Ya c1 / Then there exists
constant M(e,5, E, C) such that if o 73 (1 + t) > M and (1 + 1) < M there exists
C₁ C2 for E [3] with C1 C1 > 0,
(11 - (oV + 11) -8
8y2Bo
x Ci + C2 cos (10g(0v BBB(1 + 1))
V 4 B 7 B2o2 + 72Boy / 4y2B 7 B2021
87
+ +C sin 1og(oo
+ +72 473B - 132380221
V43B B
B2o
12 Bo
+ +
V73
and
112(1,8) (a 73 + 8)) TrBB''
73
0V (oV 73 B(1 + 1)) -
Bo²
x [6 + E2 (oogoo B(1 +t))
Bo
+ 122
V
44aB 1/ B²02
B72Bo
+ Cgsin
+1))
+ 72 Ba A3 - 13
V 32
+ + 0 B2
V73B
(ovrrB((++s) 1/ or oV T3B(1 + t) S M holds, the following is true
+(t,5)11= =0+ B( + -
and
=-+ ( -72 73
(0 vaB(1 + s))²).
B
Proof. We apply Proposition H.2 and Proposition H.1 to decompose
( = (t)
(t)
T (1)-
(0s)) -Diag(1, 1, 7 73
Bo VBo
i
1/ Bo
+ O(E)) (a (a + 1)) 1))--)))-) 8 Ba
i
+
V1
1
x Diag
( -io B
13 + t)) i8
V438 23
Br
-1
(a (o
(r()) x Diag(1, (oV s)) (av 13 + =)}}))
x S + 0 B
73 Va 13
x
73
Diag (1, Bo VBa
1,
Here, we used that
(10
S + oy no
+2 (1)
/ V3
101

---

where S' is constant matrix. To see that, we need to compare the ODE Equation (101) to the limit ODE with Viah
Ba
O, / B + 0:
0 0
0 26
0
0
-2
(
0
0
-6
T +
0
0
2
((T).
dr
-1
Around zero fundamental solution j(T) is asymptotic Diag(1, T --2 --)) while around infinity a fundamental solution iT) is asymptotic to
1
Diag(e SaaageT7arr 7207 7207+ Aaa 7207 -
i 1
1
The matrix then defined by S def ((1) We, in particular, know that det( +
Lemma H.2. Vt > 0,j(t) E R3x3 Additionally, we have the bounds on the coefficients Re(S > 0 and (e(S)2) > and we can in fact write for some Ci, C22 C3 with Cj > 0, E [2]:
A() - ( C2 cos (VAr) + Cj sin (V4+) ((F-
Proof. Denote (²()) = )u((T), j.s (T) the columns of the fundamental matrices around 0 and oc, j.j. By definition of S, we have for all 7
( = S1) + ((()S21 + (T)
By using the asymptotic of j. and the fact that ji must be positive by Corollary 1 (to better justify because assumption not completely valid, only the limit), we get that Re(Sn) 0. The same argument shows that Re(S12) 0. The last inequality comes by expressing the solution in a cos, sin basis.
For an estimate of 122 we need the 12 coefficient of l3 + o 7 B(1 + P) + ( + def
(P)
We already did this in the previous proposition. We write
U-1 = I3 a BB + s) (a 1/ + (P} -P2) O(1a VTB(1 + +)))
and hence
1
= -oov/aa1+1
(6 11= +ē
Finally, o we obtain the result.
Next, for a BB + t) and /o +) 1, we have the following.
Proposition H.7. Let e, (0,1),C > such that 8 E (1, Let Y2, Yas B, o > and suppose 0 < s < L. Moreover suppose that Ba < (1 tyaBB. Then there exists constant M é, C) such that /B +1+ M,
(1. to (c-
[+t) (1 coslo 8) log
1if T*a
Bo Viuil
O()))
and
8))(1.8) 2 (t
-8 (1 (1 ( c0-ca(o 5) + log (1t)
[tt [44)
tr Be /oar हo 1a
+ (0)
Te8i vis
+ +0(e))
If, on the other hand, a / B(1 + +) E [1, M. we have
(n = O+ (6 (e-naoo' (o ())- 1)) and P12(t,s) = 0+ (o +1))-))
102

---

Proof. We use Proposition H.1 to write
((t, - (1) Diag
(1,
73
V73
Bo
v Ba
1 -
= +T)
1+8 e e
x Diag
tn (1-))
300
x (T x Diag (1,
2a Ba
/Ba
The same estimate (102) on T we used previously, implies the result.
2nd Case: Y2 > (1 + e) As in Case I, we get different estimates for the values of P (1, 8) and 8))
For 7/2 (1 +) <
we have the following.
Proposition H.8. Let e, (0.1), C> such that 8 E (1,C). Let 12, Ts rs B, > 0, ler s I Suppose that d(ō, N) > 6, "2 > V4raB. Then, there exists a constant M(e, C) such that
if 12 B1 + t) < M, &n)(,,) = + e and if additionally & < M, (12(t,8) = 0 (72)- If on the other hand M S B2 + 1) < or M S S 1, then P11(1,)) = 0 (1) and (1)(1,8) = O+ (12)-
Proof. We use Proposition H.4 to write
=(t,s)= = (1)) Diag 1,
73
V 73
2 /Bo
0
0
(Ta + (o² P) (72 Bo + (P3 + +())
0
I+s
0
0
142
-8
x (I3 (12 P1 +72 5)) + +))
"Diag
Diag
1,
73
V'3
B
Bo
[e (m=Bo -
(m + ( (1)
i+c
:
144
(1ta)
l x (12 + ((P) - Pa)n +e)
(1: - a)) P,) P2) +
Y3
vT3
1=
Diag
1
Bo
VBa
(raBo?(1 + ())(-((P))) +
(raBo s)]( (P1)s()))
Note that (P2) (P))22 = O
Viall Bo
and (P1)28, (Pi) = O
V Bo
Hence we obtain
(t, 1 + and P12(t, s) = O+
Ya Boa
(72 Bo (12 +))
(
v/ 2 72
=c) = 0+(13)
Additionally if 72 Bo +1) [M, 1] we still have the bounds: 11 (1, O+ (1) and 112
O_ (72).
For 7/2 (1 + Proposition H.9. Let (0, C such that 5 E (1, C). Let 72, 73, B, >> O and suppose
> 1 and 7/2 20 + +) 1, we have the following.
0 < 8 < f. Moreover, suppose that ?/2 Bo > (1 + 44gBBB d(s, N) > c), 72 Bo ( + 8) S 1, and B2 + 2 I, then
1111,,) = =0+ -
(e-12BoP(L+)) (12 21 + +t))
103

---

and
12 8) = 0+
Ta Bo
(12l8² 5))² (e + (1)) )))
Proof.
We use Proposition H.3 and Proposition H.4 to decompose
(tt, -t)()) (1) (s) Diag (1, Bo² Bo
v 73
T(Is + (12Bo + (12Bo + +)) (P2 - c))
(1200
(12 (12
e (1 (,) (1+
(D+o V (1))
7250
Diag(1, (12 (12
(Is + (12 (1 + s)) + (12 s)) ())
73
V
Diag(1,
Bo2 VBa
Here we have written (1)
=
Vige D (1)) where D = Diag(d) da) for d, > 72
0, i E [3].
To see that compare with the limit of ODE Equation (100) as Ba
+ 0, 72 Bo² - 0:
0
0 0
0 25
0
(T
0
0
2
0
0
+
0
0
0
((T).
dr
0
0
-1
We know that around there is fundamental solution j asymptotic to Diag(e 27 1,,=T) and around O there is a fundamental solution j asymptotic to Diag(1,7-24 r-2)rr=)) It is clear that in fact both solutions are diagonal, positive and hence we can define the positive diagonal matrix D j(1). This implies directly the bound on 4
2 1n s) B0 (eMit (1 1) (1 +e"st(++ +00
(72
where we used that μn S S12 = -2 Bo2 < nB S and 8n < 83 - S with 8, bounded above by some constant of 8. For the bound on 1 notice additionally that since ó > 1, the term (12Bo²(1 + +)) 2 (r2Bo² + s)) 2 (o2
6+1
25
For Bo +t) > 1 and Bo + +) >
we have the following
Proposition H.10. Let (0, 1), C> 0 such that 8 (1, C). Let 72, 73, B,o < S. Suppose that 72 Bo 8) > 1, Bo > (1 + +) B. then
1+t
2 (1) (110 s) + x))
1+
1+t 82
+ +((+ + (12 2 + (72 + - (1
1
+ +0 10 Bo (r 2) (
1+t
80
1 1
104

---

and
61
73 112(1,8) - Bo²
(0(o² +12 (1
(1+ + 8)) 3))-3))))1(( ()emn((=s)
1 + f 1+
1 ++ t
82
+ +0 + B (1+ (2((
1 8
1+ 68 1+
+ (1 + (12l8² + t)) (
where x=
V Y/ Bo
Proof. We use Proposition H.3 to decompose
((t, () = 5 (t) (s) (s)-1 Diag 1, Bo² /Ba
73
v 73
= T + (12 + 1)) P1 + (72 (1 + ()) +))
(c=)
61
I+s
82
0
(t
1+!
0
I+s
(ds(t_s)
1+t 1+3
(Is + (12 2 (1 + 8)) -1 P1 + (12 (1 + +)) -))= (Pp # c)) -1) 'T- Diag(1, (1,
73
V3
Boz
/Bo
? O + (
1+ O(x 1 B0
O + (oBboTt
(t-))
1+L
0
0
I+s
1++ 82
0
1+8
5y
0
0
edHa((-o)
1++ Its
2
1 O(x + BeBo' )
O + 12B00=
O +
1
1 O(x + T2BO22)
Diag(1, 1,
73
73
(2022
Bo
/Bo
2
O(x +
Bor2)
O + (22
where we used the estimate (102) for T and (H.4) for Pi, P2. We additionally remind μ41
-72 - B2200 - 4y3B, U2 = B22 + +ov
/ B202 4rBB x -2 #3 = B2
and 61 = - 1 V13 47
To Bo
82 - 1 +
m Ba
83 = -. Especially
4
have for 72 Ba > (1 (1+c) that μn x 222 Ba² 112 x
278
and x - 273B
T3 &g x -28.
Since M1, 12: Hz S -1 d= this directly implies that In1 P12(t, s) 0+ -e- - Ultimately, this leads to
1+t
$ (11(1,8)= = (1 O1 8)) + (t)))((ss)
1
1++
+ +0 202
+ (n (1+ 2 -2)
(103)
1 +++
82
+ + +0 (1 s)e -1 + (12 (1+ -2))
+
1+ 8
105

---

and
++ 61
112(t,8) 8) = Bo²
73
O(x² + (12 10 (1+ 8)) ))
1
1 + 62
0(x² (12 2 + 1)) -2 3)
1
(104)
+ 0( + + +
+1)) 5))=
1 1+1 1+
83
(
H.S Summary
We informally summarize all the bounds on $1 (t. 8), (1, 8) that we previously developed. In red, outline the important contributions for the first and second terms of the kernel.
o (t, 8 11, $ $ 22
e eBB't
e (-
(t,
s
75
V
BtsBt
VrsBs
(o (ytBB -6
e Bo'
c (o
(1,8111)11:
22 1
0
V Bi
7B
1
(ov 2 ((-)) (+)-a
e e-
1
a (t, 8)
12, $s25
S B
Bt
VogBs
73 (a Bs
22
(aBeel a of saraii tostmmatiii
1 o (t, 8)
1)12 12: 8
22 <t t
0
Bt
22
Vi2Bs
Ta Ho
(a Bs)
72 (a 2 Bo o basea off coneiin
Bo 2 -2
2 (t, 12:
22
5 t
0
V Bt
VssBs
n 22 222
106

---

H.6 Bounds at the singular point Ba E - (1-c)) (1 V 4
Denote w def La
Vira 1. For w bounded away from zero, we derived in propositions H.1 and H.3 asymptotics solutions for large time of Equation (99). We still need to do the same for z 0. To that end, we will show that the solutions constructed near zero Propositions H.2 and H.4 have exponential decay for large time. We will follow similar steps as in [78, Lemma D.3].
The goal is to bound P (1,, 8), (12 8) where @(,, 8) satisfies the IVP Equation (99) with ((8, 6)
T3. Denote
5
1
(105)
then we have the identities
(
1 Ba
2 73 B
0
Yas
B
S-1
0
2 2/B
12Bls
S
0
0
VaB B
73 B
B B
Y2 Ba
0
0
0
+
4
B
V3 B
0
4 B
0
and
0
0
0
8
S-1
0
28
0
8
-28
6
0
26
8
Then (99) can be rewritten
5
0
0
-8
-8
0
0
0
0)
dd((t)
0
28
0
0
1
4
0
$(t)
(106)
dt
B
+
0
0
0
0
4
0
① (
where: ((1) def s
Ta
P22 TV
BoB
Now define {²(1) def
max ov3B
Boz
va
23
We then have:
-2
(1)
0 E
25
5 (11)((1)
8 25)
o
o +
0
(t)
w
46
)
E (t)
4 0
g (t)
107

---

(1)
Denote (t) =
5 (t)
and A(t) the operator norm for the infinity norm of the right-hand
5
2 1
matrix
8
0
2 8
2
5 0
0
0
A(t) def
0
28c
6
+
+ w
45
V +
0
4
0
Then we have the identity:
N(1)
('(t) & A(t)N(t)
10== " o BB B
From the above, if [w| < we know the existence of a continuous constant M(6) such that Vt 2 0 with JB +t 1, A(t) I Mt(t).
By Gronwall's lemma, implies that for all / 2 with a v3 1,
N(1) - (2
(107)
Proposition H.11. Denote def Y2 1. There exists some and some continuous function V4aB M (8) such that if [w] < c and 2 0,
(t, < Me (1-0))
Be
(1-0),
e(t,s)12 I
73
Me
By (t-s)
Bo²
Proof. The proof is similar the one of [78, Lemma D.3].
ao ToovaaB Ba + 1) 2 Ba + 8) 2 we can apply Equation (107) to each column of ((1, 8). For the first column, we know that (s) S {(s)-2 and f E(s) ds S (t-ss since E(t) follows power law. This brings that Vt 2 s, (t) - (E())//(s)22 (t- 3)(() I (t (t))))))), We now repeat the same procedure on the vector
4 (t)
to obtain similarly [t) - ((t)/((s)e( (t-a))(( S (tt ()(())
< 22))
and finally on the single vector (41 (t)) to obtain ( X (t- Notice that (1 -1- 8)(()) S
max V [. Vav/aBB1 Bi}. This all bring ((1, 8)11 S ve ML-M
(1,,))) S
To Ho
Me
2n For the second column notice the additional factor B
Finally, to
extend to a V B(1 s) < 1, notice that we still have [1 s)ll S 1 and we can apply the same
argument for 0 73 (1 + f) 1. For a /3 + t) < 1 this is also clear since (t s)ll S
H.7 Estimation of the forcing function
In the following section, will estimate the three forcing terms To, Jpp Fac.
H.7.1 Asymptotics of F()
Remind the definition
Fo(t) def M3, ((0))(000 0),
Then we have the result from [80]:
108

---

V B
V 1/2
Z
0
Bt
Bt
22
Bi v Bt
2
VTzBt
<
0
Table 9: Cuts a domains for forcing function. This is simplification of Table using $ = 0.
Proposition H.12. Suppose á 1, > and 2a + 23 1. then the constant function Fo(t) satisfies:
Fo(t) x d (020-22
Proof. From [80] or proposition D.3, we can rewrite
1 20 k k(v/d) (1 O(d-)) where k(v/d) solves [
j-2a-2
fuid du
Fo(1) Σ
K u2a
j=1
For large d, the result proven in [80, Lemma H.3]. For small d just notice that Fo(t) > 0.
H.7.2 Asymptotics of Far and Fpp
We remind the definitions the pure-point term and absolutely continuous term of the forcing function
F (1) del ( (o²) (02) 20
Di (()))))((()== =2 2 2/( o
o+t 20 a 1 (1, 0) do,
Fee (t) def ( (7,00)) =2c3 -a
-²- t (t, 0) do.
To compute the forcing function, we only need Φ7 (t, 0) above and to cut the integral in 4 different parts, of which only three can exist together. Informally. the integral on is decomposed using Table 9.
Two regions correspond to the first Ba V B and the two other regions correspond the second case 7/2 Bo 2 /473 B. The two sub-regions in each group correspond respectively to the sub-conditions o. (1 + t) 2 Bo रे Note that, as noted in Table 9, for t Z I2 the third
7a'
regions disappears (and the lower-bound of the fourth integral becomes
Liv ). On the other hand, T2
for the second integral disappears (and upper-bound of first integral becomes Ary T
Y2
The case Σ Proposition H.13. Let a > 0, 2a 23 > 1 and consider the general Parametrization H.I. Suppose Then, if 8 > max(2 28-1 22- - 2. 1), 28 & N, for any M >> O, there exists a constant C> O such that
73
that t 2 73"
(1) x (V13B(1 + 1)) - 2- if vT (t) x1 if V1 +t))s
additionally 1
Fnc(t) d (VaBB(1 +t)) 2+ F (to(t)=( d-1 if B1 + C
Md vr3B(1+t) C
Proof. We first consider the case
VTsB BB < For t > 2 the third region disappears. Let e 7
and 6 small enough, then we know the existence of some c. small enough such that
(148(14) & (1 61) / < (1 + 61) v < 1 we can decompose
B
109

---

F (1) x
o 29-1
$ 11 (1,0), do
==
(1g"((+4)
28-1
28-1
D1 (t, 0) do +
01 (t, 0) do
0=
(1+1)) 1 un
28-1
$1 (1, 0) do +
28-1
$1 (1, 0) do
72
72
(1)))
+ e) do
2
[1 1
72
1+ 20-1
0 re Bo
(oV73B(1 B(1 + +))) do
(1++))+
(1 (1+(1) 72
1 28-1
0+
more
do
26-1
(12 1))
do
0= 11 [c))
V4rg
1=
-2 23
-2 20
(VraB(1+1) (11(1+))
+ +0+ (va((1++)))
29-1 -5 du)
u
2+ 28-1
+0*
Vn
23
1/2BB 1/2B
-2 28-1
12
+0+ (0+((V12211++ +1+
u
22++1
28-1
du)
-2 29
(V73B(1+0))
In the first integral we have used that 1 28-1 -1 + 2a and that (t, a BB + t) < c.- the second integral we have made the change of variable == 13 and used that have made the change of variable u² = 72 B. We use again that negligible with respect to the other integrals
+ 28-1 for the integral on u to converge when 1a - co. In the last integral, we
28-1 to get that
(V (VT2B(1+ +1)
-2
B(1+)
26+++++ 28-1
du
2 28-1
-2-28-1
a +
11 +t))
7/2
-2
$ (Vrab(++t)
Additionally, note that for //7 +1) ce, clearly have Fpp (1) Su ال+ (1+
Finally, in the case V/sB 2 1, it is still clear that if BB (1 +t) S 1, only the first integral BB
contributes and we obtain similarly Fpp(t) x while if v B (1 + t) 1, the second integral is capped and we still obtain F (t) = (V7s ( + - 21-1
A similar calculation shows that (t) x cad- (VraB(1+ +)) 2+ The only difference is that we require > 0 for the first integral to exist and hence to have > 7. It's not a problem since below we have bounds from [80] (see Proposition H.15). Additionally, for the second integral to converge (and the third integral to be negligible) require 8 -
110

---

The case S 72 79
Proposition H.14. Suppose that S 72 2a + 2B 1, for any M we have
Then, if 6 > max +
23-1
1
1}. 28 & N and
"3'
2a
2a
22-1
3 (t) x (v12B(1+t)
if vr2B(1+t)21
F (1) =1 if (1 1 1/2 and B z.
additionally
Faolt) x d- ( (V12B(1+ +1) - Fao(1) x d if v=2B(1+0) 51
if Md 2 B2)+(1+))2 21
Proof. For 1
7/ the second region disappears. Let 21 61 > small enough and some > 0.
We first consider the case where
S 1. For +e1)
/4 72
1, write:
B
B(1+))
(1-e1)
Vagg
VingB (1te1) 22
F
(t)
0+ 20
c1
2+ 20 0+ 3e
do
Vig 72
0=0
(1-a)
+
(1))(1))))
28-1
(1 + c) do 0+(
1
()2Bo² ()) --) do)
V12 ((++))
x (V12B(1+tt) +t))
-2 28-1
V 13
24 20-1
28
72
B2
-2- 20-1
7o B(1+t)
+
(222(11
25 ++ 28-1
du)
= (V122(1++t) -2- 28-1
Here we have used for both first and third integral that 1+ 28-1 > -1 - 2a 1. the last
integral, we have made the change of variable u² Bo + f) and used that > 1+ 2 +
-28 + 1 28-1 1 for the integral on u converge. Also note that we bounded the contribution near the singular point by noticing that
ViaB B2 B
S
V12Bt
73
On the other hand, if / B1 + () S ce. the last integral disappears, while the third one is cut at 1
and we obtain: (t) x 1/ 1+ 28-1 do x 1.
Finally, in the case
LaB 728 Z we can check that necessarily 6 S
B
1
VTsB
while only the first integral remains and we obtain (t)(1) x 1.
2+
similar calculation shows when & > 1. B > that F Fac(t) x d 12 B1
/12 B( +t) > ce and Fac(t) = d if 12 B(1 + 1) < cc- In that case we need as previously that
1 in the first and third integral and 8 > 1 on и to converge.
28 1 1 < 1 for the last integral
We state two simple bounds on Fac.

---

Proposition H.15. Let a 0, 2a 23 > 1, a 1 $ 2' There exists a constant C(a, 5) such that vt 2 0:
{ 3 (1) C Fol() if 2a 1,23
3 Fac 0 <3<1
Proof. The proof is very similar to the of [80, Proposition H.4]. First, if 23 >> ca 0 and
hence Vt 2 0, 3 Ftc(t) 0.
Now, suppose that 2a < 1, 2, > 1. From Lemma H.1, we know that Vo > 0, vt 2 0, (t, 1. Hence this brings:
J Fac S C8
al-ad
do
d-a
d
S Cp
1 - 2
S 22
Since we know that Fo(t) x d + (0,1 299) we get the result.
Corollary H.1. VM > 0, 3C
0, vt > 0, (, > Md we
have:
9 ppp + S CF)(1)
Proof. This a consequence of Propositions H.12 to H.15.
Finally, relate the forcing function F(L) the terms p(), that we just estimated.
Proposition H.16. Let a, 3 0, 2a + 28 > a+1> 3, a, 3 # 1.2 Use Parametrization H.
with > max (1, + 28-1 } and 28 & N. Then there exists constant C(ax, B, H) such that 1I 0,
1 E ( + Fpp - F(t) < C + (p))
Proof. It is clear that for any t 0 and on the range of a's where we can apply Proposition H.5 and Proposition H.8 that 4f (t, 0) 1. More precisely it corresponds the range X Vis and Bo & 44rBB or the range $ (t, satisfies, up to constant, the hypothesis of Proposition D.13 on this range of o's. The reader can check that the contribution from eigenvalues of this range lower-bounds J (t), (t).
TIe and 72 Z V4nBB. Hence the function (o +
Similarly, by upper-bounding the oscillatory part in cosine/sine appearing in Proposition H.9 and Proposition H.6 by constants, one can upper-bound 4 (t, s) by some function which satisfies the hypothesis in Proposition D.13.
This implies the bound for some M, M1, M2 > and some
1 Sra (t, + Hd..)
f*
( ((t
Mia
Ma 2u
I o The
Φ (t,0))) + us (do2).
Mad-za
There remains to bound the integrals on segments [0, Ma
U 1 oo]
for some con-
stants M 0. We know from Lemma H.l that for any t 0 and any
112

---

DANA-constant with S 2
DANA-constant with / 2
Fo(1) d 2a 2a max(0,1-288)
Fo(l) x d 2a maxx (0.1-23)
pl x (12 -1-
Flp(1) (V7a . -2 28-1
(1) 5 Fo(t), 2 > 1, 2a < 1
Fec(1) C Fo(t), if > 1, >
2 <
0: if28 <
23 > 1,20 > 1, (t) x (Batt) 20 d=1
if 1, > (ac (V t) -2+ d=
K (t, B12 Bt)-211/220 (2a)
(1, x By (ViBB.1 1+11aa
Table 10: Large d behavior of the forcing function and kernel function for DANA-constant for small and large f. Here the constant C is independent of dimension.
O, (t,0t,) I
Hence using Propositions D.4, D.6 and H.12 we
can bound for any M 0, max Mad 2a (do²)) S d 20
(t. () S 2a &f 0)
E Folt). Additionally, using Propositions D.2 and H.9, we bound max [ Di Di ()do))))) f & () + (do2))
(1-2)1-
minfe nanfe-) (V1BB/Mt) "0 < Fo(t) + (t) pp).). This concludes the proof.
H.8 Necessary conditions for stability
In this section, we look for necessary conditions learning rates and batch exponents K1, R22 K& >> O in Parametrization H.1 for the risk to remain bounded. We first state a technical lemma that shows Volterra equation when the forcing function is lower-bounded and the kernel noise is too large.
divergence of the solution
Lemma H.3. Let T IR+ R4, K {(t,s) R22 =2 + R4. Let P R4 - R4 solution to the Volterra equation P(() F(1) Jo P(s) s) ds. Suppose that:
v1 2 0, F(t) 2 Fo > 0,
liminffe f1 K(t, 8) ds > 1.
Then,
lim P(I) = 00 10000
Proof. By assumption, we know the existence of T such that 2 T, Se 8) 1 for some e > O. then clear, by recursion, that
vk MEEUU {-1), VL 2 x P(1) 2 + (1+6) k
This proves that PPt)) L230 00.
Lemma H.4. Let > 1 Let 8 > max - 1114-11,, 21, 25 & N. Under Parametrization H.l. if
N1 < (1-2a) or Kn Ki + Kb, then for d large enough,
['
lim inf 120
K(t, s) ds > 1.
J122
Proof. We first consider the case where N2 2x + 2a
Kh. This ensures that
B
da
B
Since Vt > 0, Yo 2 0, 4f (t, s), 2f s) 0, use Fubini theorem to write for large
enough so that B1+ S VTaB
B
113

---

[ K(t. 5) ds 2 B 1/2
fr
Pf (t, ds
uox (do
1/2
Vrg B
Vig
2
(1 8i (f2(t, ds px (do2).
+ B
Note that we integrate only to and this ensures we remain bounded away from the singular point. For the first term, we write for
Y
as we only look for lower-bound on the kernel
8
r B
V BB
a E
(1+5()
[' D (1, 8) ds
[" 1 (1
1+t -8
(1 4ra 8)
1/2
1/2 2
1
1+1
8y2 Bo
+ (0 (
12 Bo
+ +0(1) ds
+ log
1
v ///BBB - B22
VaB
B2
+ +0(1) x Bo²
viaB
B2
1
72
Here used that / B 2 722 Bo² and f large enough.
For the second term, a similar argument brings that again for o E
v B(*))
[/d 2 (t,sss ds = /'
73
(t-s)
1+t
-5
1-
47gB(tt / 5)
1/2
1/2
2
1+8
1+1
812Bo
(0
12Bo
+ +0(e)) ds
+ log
+
1+
44s
7/2 B2
VralB
73
72 BB
+ 0(s))
x Ba +
B
V3
72 Ba
2V
73
B
B
Here we used that Bo²
< o 44a B to obtain that the oscillations from cos average in the integral.
It is hence clear using that
77 B
2 d= and / large enough that
B
[ K(t,s) ds Z BB
72 Bo Mxx (do?) + B
Y
X
Y/2 Bog MXx (do
Bo2
1/2
v Vra1(143)
(13)((14))
min
2V
12
Bo
1/a
B2 + Bo 3-100 1/a
73
do
?2 B2
d-
72
20)
+
73
x
B
The above shows that supposing K2 2kl + 2a Ks. then if < (1 2x)+ or K2 then
Ki + Kb.
lim inf /' K(t,, s) ds +00.
1+00
J1/2
114

---

We know consider the second case where K2 2 2kg + 2a Kb. It is clear that we only need to con- sider the case where K1 < (1 20.+ since
(1 2a)
R2 K3 +
Kbb
and R2 2 2 + 2a Kb
We use Proposition H.10 and obtain for given M 0. max that
Md and (1+881+ > 1
11+
(1 = (1 + (1 ((1200 +)) 2r)
1+ $
+ 0 + x² (12 + +)) + Bo (1 + ()))-))ed2(000
1+1
1+5
+ 0(2² + x Bo (1 + (r2 + -)) (e
1+t
63
1 8
Especially, the most important term is ((n(()=)) which gives rise to 2r in the kernel norm. To see that, notice that we can take s large enough so that (12Bo² as small as we want by increasing d (or M in the particular case where X2 2nz 2a). This brings
(1++)) s)) is negligible and x²
2 B B
[ (1 O((2Bo2 +)) + (t)e)n((=s) 1/2
1+1 1+
5 ds 2
?2 Bo²
It is additionally clear that the last term brings a negligible contribution since
[ O(1² + (12Bo + 8)) ()
11+1
1
=002² (02 + (1++))+ 8))-
1/2
1+
14
1
O (72 5))
We mostly need to bound the second term which is done as follows
[' Olr + (12Bo² + 8)) -)))))))
1 82
ds
O(12' + (1+ -2)
1+8
112
=0
V73 13
72
Bo
7/3
=0(
V
B2
72 B2
We noticed that μ12 want. Additionally,
and vB T2Bo
Additionally, we noticed that can take & as large
72
V Bo
2
Finally, we obtain that
1r K(t,s)ds
-2B 12 (x
1/2
o=d-a
1V
120 1201-1/00 1/aa) do
o=d-a
12 -2a4+ (12a)t
115

---

The above hence implies that supposing R2 2 2kg + 2a -ib, then if < (2a: - 1)+ we have
lim inf [" K(t, s) +00.
1/2
Corollary H.2 (Necessary condition for stability of DANA-constant), Ler a 7. Let 8 > m }): 26 & N. Under Purametrization H.l, suppose λ₁ < (1-2a)+ or R2 < K1 + Kb.
Then for d large enough, P(t) t oc.
Proof. This is direct consequence of Lemma H.3 and Lemma H.4
H.9 Sufficient condition for stability: upper-bound on the kernel norm
Lemma H.5 (Sufficient condition for stability DANA-constant, Under Parametrization H.l, for
given 4: + 2:6 5 > max 4 2), 25 & N, M > O, we know the existence of C >> O such that,
1
1I 2200, 0, /
X(1, 8) ds So (0,2ain00 +d
73
72
As consequence, there exists some 0 such that for any K1 (1-2a)+ R2 2 722 S C, Xxxc and any d 2 we know that
Kbb 1,
sup P(t) < 00.
120
Remark H.2. We believe the following stronger bounds to be true and could be shown in the same way, although a little bit more technical. We discuss this in more detail in Section J. For given /: 2a 2/3 >> a. B z. using Paramelrization H.l with parametrization vector H, 5 > max [1, }), 28 & IN and for any M there exists C(a,, B, H. M) > O such that for any
t 2 0 with M
S 722 Bt S Md
C (2((2Bt)
1
(1-20)
33
(12 1((a)))
- / K(t,,) s) ds
BB
0 (2((2Btt)
73
(12 (t)Bt)/220))
2a
72
Additionally the kernel norm converges in the sense that 30 3, H, M) > 0 such that for any 1 2 0
with 3/4 S 72 Bt 2 2
(2d((-20))
7/3
K(t, ds s (122d1-220) +
73
+ d
$
C'
1BB
B
Proof. We can bound P11 s). 112 8) as follows
from Propositions H.5 to H.7 if B I (1 e) then P11(t,
Ole 1200 B(t-s)), 11 s) 2a O(e 220 B(t-s)) with a corresponding
(1 (1 + DB Proposition H.11 brings that for t, V 4 (t, 0 B(t B()-s)), 1 s) B2 Ole 2 B(t-s)). ((tss)))((tt))
for Bo [11-) 1B
Propositions H.8 and H.9 bring that for B > (1 + <)VTBB if B2oo3 S
1. then (t, s) 23 Baz 020 B(t-s) + (120² B(t -)) with C. > O,C. > 28 > 2. Fi-
00 B(t-s) + (o20 B(t - s))- (1t(ts)
nally, Proposition H.10 shows that (t, 0
(11202 Bt)
116

---

x:² )e
st-siog
), 11) =
73 Bo
O y2d B(t-o) + (()20²Bt)- + x²)e -2
se-eio
with
T
Ba
We just need to integrate these estimates on s and a. It is clear that for S
(1 +1+00 E)
Y B
we have
f @ni(t,) 2
[ e-amBoo((t ds 2
1
0=0
y2Bo²
and similarly On the other hand, for 2 (1+ B
f @12(t,8)
73
(t ds
73
1
== Bo
B
7/2
Bo
still obtain
S Ln((,s) S f 0=
(= (
V73 B
2
(t-s)
T2
ds
Y2Bo
(r2Bo²1)2
B2
Here used that So - [t-
(1-s)
ds S (in
12, 73
We combined that with
V3
72 <
S Y2 r2 2 1
Bo
73
Bo (r2
12 Bo
where we used that in that range, B S 1. Similarly, we have
112(1, ds ds x Ba²
Bo
Now integrating against pe and using Proposition D.14 brings the desired upper- upper-bound bound for some constants e, O,
[ s( 8) ds
f (ii + 12 s)) ds
83/1/0 1/a
12B
B
do
72 Ba²
Bo²
922 Bo²
2a
< C 12 2a)) + B d)
The last claim is straightforward consequence using the fact that F(t) is bounded from Proposi- tion H.16 and (tt s) is continuous.
H.10 Upper- on the kernel function
In all this section, we assume K1 = (1 2a.+, 0 0 S K1, A1 + 1. The previous sections show that these are the largest stable learning rates. In this section we will estimate an upper-bound on the pure-point term of the kernel function by:
upper-bounding (sn))()))sin(())
in Propositions H.7 and H.10, hence defining
1 11
using the upper bound on μx by poxy from in Section D.
117

---

72
B1+tal <0 1+82
(1-(11))) ()
B1++(++))
0
V1aB(1+1)
1 >
72
278
2%3
>
/laB
72
VnBi
72
2a Bt
VyaBa
VyzfF
< a
4rya
Table 11: Cuts of a domains for kernel function
We hence define
K(t,s) s) del K2 + (t,
क (t, s) dp((o) + t f" dp((
def
12
We additionally define the corresponding pure-point contribution of the upper-bound on the kernel.
K Kpp (t, s)
def x-1
pp
(t,
+
+ (t,
B B
$ (1, d/pp)pppp +
81 8) d)ppp(o²)
0=0
o=0
B
-3 1 27 (1, s) do B 0 ( 8) do.
J
In the rest of this section we will provide upper-bounds on Kyp We will divide the integral on in different parts, explained in Table 11. We will show in the following that Kyp provides an upper-bur bound on the kernel function using our estimates on μх. We think this upper-bound is in fact tight but we cannot (and don't t need to) entirely prove it with our current estimates on fox from Proposition 14 due to the oscillatory nature of @ 11, 212 in Proposition H.7, Instead we will rely on a slightly different argument for the lower-bound.
H.10.1 First term: SGD noise
We will use the asymptotic of Φ (1, developed above show the following bounds KK (1, 8):
Proposition H.17. Suppose a 1. a # 1, tion H.I with (1 22) 0 E Kb & K1, K2 = Kbb + 1 and let any M > O. There exists
max(4 1, 1}, 28 & N. Consider Parametriza-
constant Cla, H, M) such that for any (s, t) E R2 with S t, we have the bound:
-4
JeL (t, s C
mint
VaaB(t - s)
8)) -
1++
-8
-2+ 2a
+ B
B2
1
Proof. The proof is divided in three sub-cases depending the relative size of s, 72 Introduce 73"
some e > 0 small enough.
1st case: 22 +51+ss1+
We will suppose the following v 73 (1+1) 2 B7g /
Z and 72 B(++)) Z B1++s) ž We additionally first assume that B < The other cases can be handled similarly by capping the bounds of the integrals. We hence decompose
1 118

---

B
1
11(, do = =
Vigei
03 0 (1,s) do
o==
o==0
(1--1)
+ 12 B
Vistla
1/2 11 8) do +72
12
$ 11 (4
(1+e1)
+ 1
3- a
do +72
1 (t, 8) do.
(1-t1)
Ving B
(1+e))
A 22
Continuing, we have
B
- L (11 do S T/.B
Vislt g3d x do
0=
0=0
+ 1/2
V1g8s - p3-ae
(OV do
-
BT
(1
1+t -8 [+s) do
+ 1B
g o (to
VTyBe
+ 1/2
(1+61)
T2
3 1--1 0(e
(t-10-
do
(1-e1)
22
x²
+ B
b-
2a (1- + (x ()2025)²
e
(1+c))
Vrrg 22
1+1 82
1
+
)
1+8) do
()202))2
-4+1
1
S B
(17aBt)
11
du
==0
+ +12B
(vaBBt)
-4+1
2g du
22
u=1
11+))
-8
(1-)))))1
+ B
(72 - 8))
-2 nm
u 1-/1/a du
1+8)
22(1-0)
+ +0(
61
VTaB
7/2
2
Be
Se 3773
BBB
Bitt
+ +7gB(12B(( - 5)) 22+1/(2a)
- du
Jum
-4+
5 BB2B (3
4-1/a (1-9)13
if
t - 8 <
+ (1+1)
B12 - (0)(tss))--00 -2
273
C if
72
(1+)
72
(t-e 7302
if
228 <
Br2 2))2-1/(22))
if
B2 - <1
+72B(rBB(t (t)t()))-2200) -2+1//2(22))
if
< t
73
(t-eing
if
12 8.
119

---

We have made the change of variable u = / + t) in the first and second integrals. We have used / for the first integral converge. In the last integral we used u² Y2 -) s) and noticed that the integral on u vanishes as >=((=>) + oc, as 8
D2
We check that this result
738
equivalent up constant to the result in Proposition H.17
We additionally bounded the singular point for t 2 22 and hence 73
Ltt 1+a
X (1-4)73
(1te)) v Srggl
T2
des (1-18
12 B
3- Oie
72
)o
(1-E1)
Ce. (t (t-(t-)73 -b-241/(2a)
Vial 4-1/a
x B
72
Y/
1+t
(12 - -)) -))-112((2a)
1+8)
For 8 S 22 we have I+8 X because + 2 22 which brings the upper-bound directly. both cases, this term is absorbed by the other integrals. 7a
1+f
Ts
We additionally observed that the following term is also absorbed by the others
BB
3- x
x22 (72Bo + (r2)2) )
e -cn (e-a)rs
1 (1+1 62
do.
(1+t))
Ving 12
1 s
Indeed, we know that ++21
1121+8
22 2
This implies that
73
72
B2Bo vrBBB
-2
Y2 Bo 2 32
= x. -2
hence we only need to bound
1+t 82
BB
O xx
3-
ozj 72
a o/a
do
[1+t)]
VingB B 12
1+
B
(t-2)13
1+t
$ 12B
do x
Gey
72
V13B
(1+q)
1+
72
-4+1/a
-2
B
(t-2)73 72
1
$ 1BB
X
VirB
1+s
-5
$ 32 -s) -2+1/(2) -3))tt--)))//00)
1 1+t
1 ++
We used when t - s that for vn > 0,4g 1, e S v
with 17 -22 + 24
For t
<
73
-4+1/a
we still have B
BB
S 8)) -
V
120

---

2nd case: 1+ S 1 32
We adopt the same strategy although we will go faster as some
73
computations are very similar
4 (1+e1
BB //' 0²- (1,,) do BBBBB
72
3 1
1 x do
+ 12
V12ll
σ 3- x x 1 x do
(1+c)
Ving
+ 12
V128i
-- _/2 Boott
do + 1BB
/-- = e e-xzBot(L_s)
Vig8i
V1285
x²
(e-elaa 72
1 1++
62
+
x
(2)Bo +
1 do
(12Bo2t)2
2 1/ BB2 Bt) (t))/BB))2 / w²-2 : u²² du
2+1/(20)
m=
V(2B(t)s)
+ B12B(12B(t - -s) -+
u 3- _ -u² du
(12 2))2-1/(22) if
B2 - 8) < 1
+ B(22 - -++1/(2()))
C if BB < ( <<8
e -- if 8 sct
where in the first 3 integrals we used the change of variable u Bo that e- for Vr and that a for convergence of the integral for small u. In the last integral, we made the change of variable (t2 s).
We additionally noticed that
V13BB
4
1 +
62
do
7/2 Bo
1 1 s
VTzliz
v 1/3
2
1/a do
111 62
B2B
1
V12Ba
7/3
(1e AM
1 + t A2
7/2 (12
1 + s
(72 Bt)
B1)- 2+
for 5 large enough. 3rd case: 1 s 5 12 <1+t
121

---

BB [
0² : sddo do B/
VTsBt
-3-21 1
0
(1 [1-61]
+ 2
B
72
3 "
12 Ba
(0 do
Vigit
+ +0+
72 B
(1+61)
72
3- Bo
(t-s) do
(1-61)
22
+ 1
V12Ba
0 a
B2 do +12
3- e (t-o0 do
(1+e))
Lighl
12
V>28a
1411/a
2 12 (178Bi
3-1/a du
11 u=0
(1 (1-e1)
+ BV23(V73B1 1+1//
72
" 1/a-
6-
du
u=1
VB3B
+0+
252
B
V3
+ Brr2(Btt -tt)/22))2222 2+1/12a)
3 1/1 - u² du
)u=(1+e))
72
12
+ B72B(rBB(( - -s) -))-2+11((2)))
3 3- 3/1/ u² du
S T (V73Bi) -4+1/a
B12 2))2211(22))
if B2 - 8) <1
+ B(2 - (0)((-))))+11111 -241/(20)
C if BB <t-s<s
e - if 8 s<t-s
where we used in the first 2 integrals the change = a Bl and that - for S
H.10.2 Second term: momentum noise
will use the asymptotics on P12(1, 8) to show following:
Proposition H.18. Suppose a
i' a * 2a
max(4 1 1}, 25 & N. Consider Parametriza- a:
tion H.I with (1 22)., 0 E Kb I Ki, K2 = Ks + 1 and let any M > 0. There exists
constant C(a, H, M) such that for any (s, t) E R2 with S t, we have the bounds
R2 (1, $ Rc (1, ift
or 22
R2 (t,s) C 1/ (BB(v8B(1++)))+ (1+
2 + 73
1+1
-8
Bt2
I+8)
f1+ 21 22
Proof. The proof is again divided in 3 sub-cases. Let 61 small enough, > and > 4 a²
1st case: 72 <1++$1++
122

---

B
0
112 8) do B
Vig 0 112((, 8) do B VTsBr
- 12 8) do
0==0
0=0
VigBt
(1 (1-c)
(1+e)) Ving 72
+ B
72"
11, do + B
V 0" 3-
(2,
do
(1-c1
Vig
B
112(11,8 do
4
(1+1))
22
[1+1) -2
B
Vig
73
(a V 13 B(1 1))²
do
B
1 1+
+B
0
3- 3
(o
13 B(1 + t)) -
12 Bo?t (o 73 + s))² do
B
(1-61)
1 -5
+ B
724
0 a
12
1ts) do
Bo²
VT3B>
(1+o) V R
73 0+ (e (e do Ba
+
Ving
(1-1)
72
(1+6 81
+ B
73
0(2² + (12 -)) -e) (1-s)
AygB
Bo
1 1+s
(1+t))
22
1 62
+0(2² + (12 (1 -)
1+8
(1+1) 88
+0(2² + B0 Bo (1+ 5)) -1 + (T2 2 (1 1)) (ess(()s)
do
1+
1+t
$ 13
(V7aB(1-1) (1) 22
u du
1+s)
u==
1+t
n2 12
+ 73
(V>a((++t))-
24
u
T3
du
1t8
+ 73
[+8)
1+t
(VraB(t- -
-241
e u
du
12((xx)
3322
B(e-(t
+
+7o(12B(t - 8)) 5)) -1+1/(2a)
11 1- ae
-2 du
2u²
3 73
1+1 I+s)
-2 (v B( t)) t))-2224
2+
73
1 1+( I+ s
(Vr2B(t - s)
-2+
where we have made the change of variable u= = o + t) in the first and second integral, 1 used that a > / for the first integral on u to converge, that d > for the second integral on converge, and note that 7st & 1. Finally, in the last integral used the change of variable u def Bo
Also note that we bounded the con-exponential term similarly as for K! (t, s), since know that 1 + 12 1 + 2 78
72B. this implies that
Bo -2
BB Z B2Boss
-2
Vr
hence only need to bound
123

---

B
σ
+.4 73
x²
72
1 l
do
(1+c1)
V 72
Bo²
1+s)
1/
1++t
26
5 73
do xe
722
VY3B
(1+g) 72
1+
YeB
-2+1/0
11-0h3 72
1+t
26
$ 73
VY3B
1
1 1+t
-8
5 73t ss) -211/(2a)
1+
2nd case:
1 "2 < 11t 73
In that case,
(1, = B /'
0² a (1, 8) do S B / (139 و3 73(1 8)² do
00
VagB 72
(1
+ B
=2 (0 (o ()) -3 + x)² do
Vighl
+ B
(1+t1)
B
3 0 00+
c - do
(1-61)
B
+ B12
V12BB
0¹ a
(V Bo + s)) e 2re do
(1+E1)
V4gg m2
\/aBB
+0
B
σ 3-1 7a
(e (e-212B0(u-s) + + do
B
B2
Vi2Bi
(1+8)nn
S B12
(V1aB(1+t)) 4+
S x st
3rd case: 8 SSt
22 In that case we have: 79
(1te)) singl
K2 (t, 3) /'' J
73
s1, do S B [
I 2gg + 5)² do
+ B
3-1 (1 + 3) do
(1+9)
V
+B
73
(12 B)zz))e 22, e do
B
v/3B 2
+
3- 73
(e
2n2 Bo? 4 +e- -== do
B
Bo
V1282
(1 (1+
1 1+s 2
TBB Bt 2
< BY2
(12 -2+ +
72
1+t
22 ( ( (U-5))a
2
+ BTg(12B(1 B(1 - 5))-
7/2
-
$ (t,s)
124

---

H.10.3 Summary
What we have shown can be summarized as:
72 (1g 2 if 8 7/ t 1 2 2
73
72 1
if 2
K (pp()))) S K (l) (p,)
def
-
L+t
min (1, (72 B(t - -))-
I+a
v13 B 73 + ())
2 ssst.
73
R pp (t,
if 8 78
t 8-
72 73
N2
S (t, pp
(72 8)) - V V +1))
(1++))a 2
73
I+s if
22, -
12 73
Hence Proposition H.17, Proposition H.18 together lead the following proposition:
Proposition H.19. Let a * with a > 1/ and 6 > max 4 - J.. Denote Kpp (t,s) s) the kernel defined for any t 2 0,
(7p 8))
2+ 2t if s5 2 t 2 + 1)) -
or t S
22
"s
73
BiBa(Ba
if 8 <t
Kpp def
1+t
(72 B()
B(1 V
4+
1+s
1+3
12
13
if
12 - 8 8>
72
Then we know that VM 0, 30 3, M)),, vo S < with ( Ma
s
K(t, s) < K(t,s) s) 2 (t,
(t, 5) + (t, & (t,
Proof. By non-negativity of μx, it is clear that Vt, s > 0, x(t,s) < K(t, Additionally, the estimate we previously derived show that (t, s) S (t, s). Hence only need to show that k(t,s) $ (t, 8).
We constructed क (t, s), (t, s) exactly to ensure that we can apply Proposition D.14 by bounding their oscillatory part. Hence, Proposition D.14 brings the existence of M, M2 > 0 and C > 0 such that
fTL J (Y (t. + B /v (t, (t,s))) (d))
sc f
Ma (22 + ( (t, s))px, (du).
To conclude, we only need to show that the small and large eigenvalues do not contribute too much to the kernel, i.e.
7f (1, ())c((du))) 1 (1,8) + B (1, (d)
So Mdd 20
max
S (ppp(t. s).
The above follows immediately from Propositions D.8 to D.11 for large enough to apply Proposi- tion D.9.
125

---

H.11 Verifying the hypothesis of Kesten's Lemma
The goal of this section is to show that the upper-bound on the kernel X in Proposition H.19 satisfies the hypothesis of Kesten's 's Lemma.
Proposition H.20. Let a > 1. a 2.2 8 > m h 26 & N. For any M > 0, e > 0, there exists some C( 6, c, independent of such that considering Parametrization K1 2 (1- -a) R2 2 K1 - Kb + 1,, NG S min(ki, K))) if (52, 323 S C then V0 S a S with max t), (v 73 } I Md we have
I with
f" xpplt,))Rppll dr s st
Proof. This is consequence of the form of Kpp computed above. We differentiate different cases.
1st case: L 22 or
Then the proof sketch can be found in [80, Prop.G.2). More 79
7s
precisely, we have
[" (mm(t))((mmm1))) 2 f" By2B((t rr)
2+ 2h BTBBB 8)
2+ zu dr
st f Rypll, r) dr
$ B2 B(1 -8) 8)
S {t,ppt,
Here we used Lemma H.5 in the bound on the kernel integral.
2nd case: s S 73
and 3 < t Then we write:
{" 75 B22 5)) 24 20 72B(V7aB(1 + 1)) 1))-4++35 )
2
K pppll r) (r, $
dr
73
22
1 1+t
(r2B((-r))) + B1B + +0)
(1 [1+r)n
73
1+r)
22 22
72
B2 73 B(1 + +r) ++) dr
- r)) -2+ x 7 B( B(v B1 + +)) -1+ + dr
E BB +)) +
(1
Indeed it is clear for the first and last term. For the middle term distribute and bound for 22, 73 small enough for a fixed e,
TH B( (BB(V3BBL B1 14+2
[0-r)a
x B333BB11 -1+ dr
2
72
32
(1 + r17)) 2
x B2 B1 + 1)) 4+
73
32 BB + r)) 4+ dr
r=
72
73
$ S 73 B(1 44
because
126

---

(1 + r)a
2
7/ V B( + r)) -))- # dr
7/2
1s
1/2
( V 73
B3 B( + r)) 2+ dr
2 22 B2B
73
$ d B2
Bdaaa (dt(-1+//a)
And the other term,
-5
L+L I+T
(12 - -1 x 73 BL +)) 4- dr
2 13
5 <112 73
1+t
-
(12B(t - -) 2 x BB + r)) 4 dr
1+7
5 BrgB(V1BBl + ()) ++ x
B( - -)) dr
32
S 133BlL1 +t)) -+
where we used the stability condition
73 B2 < 1
3rd case:
2 222 and - 8 2 7/2
5"
Kyp (t, (.rr)KK (r, dr
8)
12
(1+1) -
1 +18B((/3B(++))+ (1)t)n))) 2)
)
2
7/3
(1+r)
r=
2 B222BB B(r - -s)
2+ 2 dr
1/
1 1+t) -8
(1 r)13 2)
(() + B(
7 B(1 + 1)) 4+!
73
12 r=st
1+r
7/2
73
X
13
+r
-8
(12B(r-s))-
- -)) - B33 -4)
++
(1 (1+ 8yns )
2
dr
1+8
72
2+
22 B2 B - r)
P=1-
12
1 [+r)
-5
(r) + +12 (()))) -4+ / (1 + s)133
73
dr
1+8)
72
1 ++
-
(1 + 8)73
2)
Se 6
73
(12 - - 2/ + + B( 73 B(
+ 1)) 1+
-1
1
7/2
The last term is straightforward because of the stability condition on 72 (which implies that we can de- crease "/2 so that f-= 22 By (r r) 7g
/d /(121) B) 2+ Ia dr S 1
22 Ia dr S
72 B2 (ddaa r]
and the fact that L x in that integral. The first term is also straightforward for the same reason, ie stability condition on 72 which and the fact that r x s in the first integral, For the first term we develop and write
127

---

1 1+1 -8
(1+r)
-8
(12B(1 - r)) (1)) xσ x 73
(12(1 8)) -))= za dr
73
1 1+r)
11
-
- T2
< 73
1 +t
(r2B(1 -
72(12 - (r))) -s))-++1/(20) dr
(1+s)
-6 (12B(1 - 5))-1 20.
(1+t) 1+
- 67a
Here we used the stability condition on 73 which implies that y/ 3 ra(12B(r 9)) -141/(20) dr 2
J. (L2B/) Br 1+1/(2a) dr S E.
We additionnally bound
72
(1 + rha
73
2
B ((1
12
72
72
13
1+r
(r2B(r (0)((r-))---0 - dr
73
1+s)
15 12 B(1+))))+++ / (1 (1+r)ms
r=st 12
72
1 (1r)) -
93
(1+s)
(12B(r $)) -1) 20
dr
(1+ + s)
2 22 B(V 7/ B(1 + 44
4
72
12 13
-6+2
1 1+
73
(12B(r- - -1)
+1
dr
r=st
1
(1 + s)133 2
S B(V 13B(l +
(1++)+))) +t) 4+
7/2
Here we used that 8 > 2 and the stability condition on % to get that Ya 1- 22
Br2
- dr S 73 1 (12 B) B2Br 1+ dr Se.
We also write
(1 (1+
B2 V 7B + -4
22 -=s+
7/2
73
(1 8)13
2
(7BBVBBBL1L11 (1L))(()))+4) ++
dr
72
>/8B(V B (1+)) -4
(1 + 73
2
T3 B B(
73
B( + r))
2+ dr
72
72
B
(1 + 8)Tss
12 B x aa 1
if a<l
(c7B(V11BB1 B(1 + 1)) +11))4+3 4+
d
-24
7/2
B
B
if a>l
VTaB
(1 + s)MY3
S 4+
7/2
Here we used the stability condition
d and that a S
< 1.
B
72
128

---

Finally we only have to bound the following term.
I2
73
Itr
lit
(72B(t - -)) La x 7/ B3 + =
(11+))28) 7
dr
-4+
1+t
(72 -r)) -
73
1+r
(1 + 8)
x 72 ( B1 + +r))-++a
4+
dr
(1 + 3)
2
22
S 13BB(V7BBLL+ +0)) 4+! ())
Ta(12 - -)) -1)-11a da dr
x
r=s
(1 + 8)73
S B BBBB(1
(1 + 1)) 1
++!
7/2
where we used that 8 4 - 1, and the stability condition on "ia for the second integral to be small.
Below we show an intermediary result that gives a lower-bound on the term I * K of the Volterra equation solution.
Proposition H.21. Let az B with > 1 a, B + 2: 2a + 28 > 1, a + 1> B. Then for any M > 0, there exists some C > 0 such that for any t 0. if m2, (VrsBt)²) Md and Bt 1,
[ + x3(1
C B (1.0)
Proof. We remind that under Parametrization H.1
T
S This ensures that 7B S T2 bB
"3"
Using Propositions 13, H.14 and H.16 we have precise asymptotics on the forcing function F(I)
for all times. Using these it is clear that for any T 2 0 with BT > we have 1fd F(s) ds Z T.B
We now differentiate the two cases t S and t Z for which (t, 0) shows two distinct behaviors.
For t 2 22 we know that for any 2 0 with St. 7a
Vigiie
K(1,8) Z
BB (t, duoc
M 1/2d-a
2V
VTgBt
BBB x
73
x djpxp (do²)
Bo
M-1134-0
1+1
A2 A2
73( / Bt)
(t,
(t, O).
Similarly i 22 we write for any 2 0 with 8 56 t
78
V128t
g(1,,) 2
7g B (1, djps (do?)
max 1/2d
V128i
22
B x 1 x d/yp pp (do2)
maxx 111d/ad-0,
$2
2+
2/ (V
(t. (t, 0).
129

---

The above directly brings that [5 * (](()) 2 (1% T(s) ds) Kpp(t, 0) which concludes.
Similarly, we now state an upper-bound result on the convolution between the forcing function and the kernel function.
Proposition H.22. Let a, 3 with a > 1 a, B * 2. 2a + 23 > 1, a 1> B. Thenfor for any M > 0, there exists some C' > 0 such that for any t 2 0. if max (1, (v) I Ma and Bt 2 1,
72 x no (1.0) (10)++
[ * x (1) S C
Proof. Indeed, we distinguish two cases.
1st case t S 22 then we write
Kmn * (5)(1) 1// B2 2+1/(2a) F(s) ds
10" 112
Bl 2+1/(20) J(s) ds + '
2+1/(2a)
2 (12BBrt F(s) ds
1/2
S Kxplt,0) I/' F(s) s(t) /." s))
1 2B (m)))))))
Here to bound the forcing function norm, we used that either 2B > 1 and
22
dea
75
B2
V
d
B
(oss) ds + fot T (s) ds
/0²
2 Ts) ds
(s)) ds
$
0
00
da + BB + 12a21d
1/2a
S
2a 20
2 (1+
72
72
73
x B2B''
When 2/ < 1 we do not need to worry about Fac and we write for the Fo term
(pp(t. 0) ( Fo(s) ds S By
2+ 26
td 2d-24-234 23+1
1 to -2a-2224
S B22Bt
-1-21-1
$ Bt
$ F(t)
where we used 72 dt (1-2aa) and Bt S da Finally, for the pure-point term we write
(t)(t, S (m)ls)ds S BB
2+ 1
B (T2Bt)
1
20-1
12B
2- 2-1-1
2=1
S B
Kyp (t, + 7272BB
BB 5 (t, +(t))
130

---

where we used that 7/2 S (1--2a)) x B2 2a
2nd case: 2 2 we write
73"
4411/a
pp
(1) -
T5 73
F(s) ds
===
79
1 1+1
73
(18)
-
22 73
+ 2 B( BBB B11 +
1)/1aaaaaa
1
F(s) ds
Ta
1/2 Br BBBr2Btt 2+11((2))
2
(ss)
22
75
1
S F(t) +
B2 pp (t, 0).
Indeed, we used for the first term the same arguments than in the first case to directly obtain
BB
ds S
12 K Kop (t Ftt)
For the last term note that
2
2+1/(20)
F(s) df do)3() 5()).
1- 2g
The middle term needs more care For the first part, we used that ó > max 4 - 2.2 2 28-1 This ensures that (1) + 8) diverges as power law. Hence we can write
73
(T2 - s)) (11)) ds (1+t)
73
(1+))
22
2 7a + 1) [F(()1 + 1)8x x x1] (12 (2))
53(1) x 13 (72B1)2/(2a)
722 B Ι
pm (1,0)
B2
1
< F(t)+ +
Kpp (t)
B2
where we used
F(s)(1 + 8)8 ds (()(1 +1)² x t.
We additionally used the stability condition For the other part we write
T2 x d S (12
2
(1+$)
2
(((+)B(++t) 441/a
-411//
F(s) ds
12
73
$ (T(()) X
+ F( x
B(V B( + t)) 1+1/a
2
73
S5(1) x Vral x 11B(
+ 12 73
x Kyp (t, 0)
5 1
pp (t,
BBB
131

---

Here we wrote that (using
72 Va
2 for the pure-poin and absolutely continuous- part terms)
2
+1) ds 5() +t² +3( 22)
Y2
73
73.
We additionally used in the last step for the first term that either S and v + S do and
vTa d-1122 or 1 and Bt 2
ViBB 2 and ys I 1 to obtain that
3 2 F(t)
/ Is (v>sB(1+t)) -41// s(t).
T2 B
B(1 + t)) +
S slo
B
For the second term, we used F(22)) S 1 2
1 B"
We now state the main theorem for bounding the loss of DANA-constant.
Theorem H.3 (Bounds for stable algorithm). Let a, + 2: > 1. 2a + 28 > 1, a > B. Suppose 28 N, ò max 1, 4 1.2 23-1 some с > such that if K1 2 (1 - 2a)+, R2 2 11 Kb +1, Ri I min 122 S c,
a 2)- Under Parametrization H.l, there exists
%/2 S c,
then P(t) bounded and there exists an M large enough and a constant Ö(a B, M) such that if Ma 72Bl, (V7:3 ((1 + 1))² > then:
1
Foll) + Fac F p() + B2
K ppp(1,0) S S((()
?
<Ox x Fo(() (1) + F pm (1) +
7/2 B k (1
Proof. Upper-bound
We first apply Lemma C.2. Using the bounds on the kernel norm and forcing function convergence guarantees for e small enough that if K1 2 (1 2a)+, K2 2 Kb 1 Ko I min 72 s
fa < then,
P(I)= Ō x (5(1) + K 31(1))
We only need to estimate [Kc pp +t( which was done in Proposition H.22, and obtain for potentially different ē, Ć,
KKpp * +)( < <Ōx F(t)+ + BB (t.(0)
Lower bound
We know that V\ 0, P()) [F)]+
+((()) and hence only need lower bound [F (t))). This
was done in Proposition H.21.
I DANA-decaying
Below introduce the main parametrization for DANA-dccaying that will be used throughout this
Section I.
Parametrization L.1. For 7/2, C K constants
71 1, (2())= = 72, 73(t) def 73 +
def
72c(1 (oc[1= + t) "K.
(108)
132

---

Phase la
Phase lb, Ic
Phase Ila
t d
compute optimal
d
compute optimal
loss curve
F pp
compute
loss curve
F pp
loss curve
optimal
pp
t
acc
Fo
Fo
Fo
flops
flops
flops
Phase IIb
Phase IIIa
Phase IIIb
compute
K Kpp
compute
Kpp
compute
optimal
optimal
optimal
loss curve
loss curve
loss curve
F Jec
F ac
F ac
Fo
Fo
30
flops
flops
flops
Phase IVa
Phase IVb
TTa
compute optimal
compute optimal
Ia
Iila
0.75
Iilb
loss curve
9 pp
loss curve
pp
IVa
Kpp
=
d
Kyp
t =
Te
0.25
TVb
Fo
Fo
0.5
i3
flops
flops
Phase Diagram: DANA-constant
Figure 18: Cartoon Plots DANA-constant. Pictures for the scaling laws for DANA-constant each of the different phases. When < d, DANA-constant behaves like SGD/SGD-M. Observe that the trade off point for compute-optimum changes across phases (see Sec. E for derivations and Table 6). See Table 10 for summary of asymptotics of I pps F Jacs Fo, and Kpp. This uses DANA-constant with R2 = and batch size B
Assumption 6. In all this section, we suppose 1 >> K 2a with d-independent constant learning rates "/2 T3 2 1, and B 1. Moreover, we only work above the high-dinensional line 2o and under the technical assumption 3 < a + 1. Finally we suppose 20. + 2B > and B + 2 also suppose large enough (independent of d).
1 We
We will later heuristically extend the results this section under this assumption to the general (DANA) algorithm in Section I.5.
Remark L.I We supposed a > as it will become clear that for a < }. this algorithm equivalent
to SGD. Indeed in that case 0(t) x 1 + Bt. The scaling ?22 x T33 x 1 and 1 2 K > 2a will imply the relation between momentum and SGD times for eigenvalue o²,
Vr(()) +1) (1 +
We also remind the definition + Bt (Jo Vrs(s)Bds) that will be used throughout this section.
133

---

I.1 Solutions to the simplified ODE
As before, we work with the simplified ODEs (43) and the resulting forcing function, kernel function and simplified Volterra equation solution (55).
We now provide explicit estimates for the DANA-decaying solution, (ta(t,)), to the simplified ODEs (43).
Theorem L.1 (Fundamental solutions). We summarize the results of this section below. Suppose that 5 + K > 1 and suppose that m = I S 1. for convenience
s(t) c11 + t)= and Bo where K E (0, 1) is constant.
(109)
There is c > 0 so that for all > 0 there is my sufficiently small so that for all m my the following hold.
We use the time scale T = m(1 + t) and
- -1 -T du vc
-T -1--//2
((T:TO)
m 1/2+
1 k//2
We let E(T) E(T; m) for short.
There are (1) and w2({) continuously differentiable with w/ + w2 bounded away from and above by a constant, and they satisfy
wi((E)
()
wi (5) w2(1))
/coos(E + wn, + O((-1) & 00, V2
E +0 and
w2(E))
sin(E + (,)
for some constant wi (0).
For any 8 2 O, we use To = m(1 + 8) and Eo the corresponding value of &. Then with 6/24k/4 we have the following uniform estimate for Eu <e
1-x/2
111
()())²+
= -1+ (/(6))(1-/00+++00
711 ((-)) of (56()) 012(lis)
(8+K- 1)2
(110)
+ +0 (em (em Sa a(w)du) where Q(u) = minf1, (E())))
The error term (()) also vanishes quadratically in the second entry as E - 0. For all To 1 with s large enough that En > Lthere are other bounded oscillatory W1,
3ga()) 2 P12(t;s) ) = = (x-
(1;(6;))
(1+5]
-2p
Wi (w)(E + 0(e) w2(E:t0) + O
1
0 (cm
(em (r)e f% R(w)du)
If >
@11(15)) 8) (59(s) 112 8)
0 ( + (1P(T)) 17 ((w)da)
We summarize the regimes for this fundamental solution below. In the first training regime, when is small (5 < 1), which is to say
E = m 1/2(1 + 1)1 < 1,
we have that the ODEs have not begun moving. On the time scale from & > but < wc not yet using the curvature, but we have decay like
3 3 (1
Finally once / > m there is exponential decay with speed m.
In the remainder of this section, prove this result.
134

---

Changing variables.
We recall (99)
0
0
0
0
28
0
d
0
0
-8)
22
Bo²
0
- 27((11
+
0
27 Bo
1 1):
dl
1+t
Y1 Bo
-73(t)
B2
(1)
To simplify the computations, we will additionally do the change of variable P(t)) def on Equation (99) and get the new ODE
B B (t) Bo Pt(
272 Bo²
27/ B
dt(t)
0
26
+
2
$(t). (111)
dt
1 +
8
1a
B
-Y1 Bo²
In terms of the variables m and 5 we can rewrite the simplified ODE (111) as
(1(1))
2m
0
- -2ms
26
+
0
(()).
(112)
dt
1+t
8
m115
m
now introduce a time change (() del obtain the new ODE:
m( 1) and defining ((T) del'
((1), terms of which we
d (T) def dr
2
0
2s
- + ((T)== =
- 25
+
0
0
2
((T).
(113)
1
5
m
Lemma I.1. Suppose that X and Y solve the ODE
d
X Y
(114)
dr
Then the vector (x2,Y² Y2, XY) solves (113), Hence, a fundamental matrix for (113) can be given in terms of solutions to (114) with initial data at time given by
(x2, xY, where (X(To), ((o)) = (1,0),
(x2, (22x))) where (((ro),((To)) = (0.1/(Bo²)),
Im(x? y2, XY) where ( (1,i/(Bo²)).
Proof. We just need verify that the equation is satisfied by the vector (X2 y2 (2, XY).
x
2X dx
2X(-1X -
d
2Y dv
dr
22
$
dr
XY
Y2
X dy
+ y
dr
=
X Y Y--1X sY)
2 -)
d7
dr
-20, 2
2 Φ 2 $2
(1 + )
This is precisely the equation (113).
To formalize the estimates, we divide the range of 7 into three regimes. Set
1/2 - x/2
1.
(115)
1 1- k/2
135

---

(Entrance) 7 cm² (5 < 1) Here we will approximate X and Y by simple power expressions.
(Transition) em? < r < m : (5 = 1) Here the fundamental matrix solves a non- degenerate rescaled ODE, which after changing variables is approximate solution of the Bessel equation.
(Bulk) mm² < < - ( > and T < 1) Here the system develops very strong oscillatory behavior which we can describe explicitly.
(Exponential decay) 1 T S cm cml (T > and ('(T) > 1) Here we will give estimates showing uniform exponential decay of the fundamental matrix.
(Slow decay) em - but at a slower stretched-exponentaal rate.
(E'(T) < 1) Here the fundamental matrix continues to decay,
First change of variables: Isotropic coordinates We start by making a change of variables. Introduce y /smY. Then, changing variables from (114) V, get
dy (Vsml)
dr dr
d
(Vsm) Y + /sm dY
dr
ds 1
VsmY + sm
X Y
dr 5
m
Since s(t) = c + +t -* and r m(1 + t), we have
da
= R
Thus:
di
dy
1
X y
+ v v
dr
2T
111
1 X
8 k//2
V
V.
m
Therefore, our system becomes:
d
X
(116)
dr
y
y1
We define fundamental matrices of this matrix equation
(v/
P(T;T))
P(T;TO),
P( m)) = ld.
(117)
dr
Second change of variables: Rotating frame
Define the orthogonal matrix R(T) by
(co((E) - -ii((E) sin(E) cos(E)
E(T;TO) l" s(u) du.
R(T;m)
(118)
We observe that the matrix R. satisfies the following ODE:
-sin(E) - -cos(E)
g' = (0 -1
R
(119)
dr
cos(E) - sin(E)
Thus we set
U
X y
= R-
136

---

then differentiating with respect to т. we get
d
()
X
R
-1
dr
v
dr
41
dR
X
R-1 d X
dr
dr 41)
dR
X
d
x y
R
R
-1
R
dr
y
dr
X
Using equation (119)
and equation (1 (11 for
d dr
we obtain:
dr
y
d
(
-1
= R 1
5'
R
R
1
dr
v
1
y
-1
X
+ R-
-
5
S+r/2
y
0
= -
R
+ R -1
R
As we have chosen such that g' = v T then the off-diagonal terms cancel, and we get
dr v (v)= R (-)
d
5+x/2
R
(120)
We also introduce the fundamental matrix this system
d dr
U(: R R (-'
0
6+x/2 7
RU((:
U( vo) ld.
(121)
We also record for convenience the relation between the fundamental matrices that we will need, (with r = m(1 + t) and To = m(1 + s)):
(u((t,s) ) (1)(t,s) (P12(T:To)122((s)
PII(T:T)
P(r; R(T; To (T) and
(122)
To verify the second matrix entry, we use Lemma I.1 to show that
P12(t, = =t)
where X(s) = 0 and =ss)=
Thus
P12(t, s) = (ia(T:TO) s(s)m
(Bo2)2 (T: To) 725(s)
Lemma 1.2 (Entrance). Suppose 8 > 2. In what follows we E(T) {(T: m). There is a > so that for all > 0 sufficiently small, there is an ma so that if m < my and m < to < < emP (hence E(T) S cxx/22
[P(T:T0)11 - -r- s ce (r-m)g)))) and [P(T: 4(T s (()
where 1T(T)T)) = 1" 6tx/2 E'(w)e- (u-7o) vo) du,
and we note (T;:TO) S cE(T;To). We also have that
IP(T:70)12 - I2(T; S cl((T:(o)E((7) and IP(T; 6--0/2 s (h((T)))(()))
where 12(T;T0)= ["
-5-//2
C'(u)e (r-u)--(--) wT- du.
To
137

---

We have that 1o) & <T(To 0), and that
1-1/2
To) b 1
with the error o(1 tending to 0 as T/TO + 00 T
Consequently, for 4 have the following bounds in this regime:
(t, 8) S and P 8) S or + +) s)2(1-n),
Proof. Under the assumptions of the lemma, we have that for all To < < cm"
1
0 S {(T) - 1 - x/2 wx/2 and 5'(F) Voe--=/2
First column. the 11 entry is
We start with the first column of P(T; To). The natural candidate for the solution to
P(T:To) ( (r-70)
So we introduce
P(T:To)I e (1-To)) (1 + V(T)),
where V(TO)= = O. From (116), we have
Integrating by parts, we have
V(r) = - [ =f" ('(u)221 du
(T:T)) + S" E(u) (1u())
5 - k/2 5-k/2
P2n (w:; to) du.
As for P21, we have
(T; u 6tn/2
(7) (()
(u-to)- V(u) du.
Let U(T) = max [uu)]. Then we have bound
IP21 (P:((:))) (1+ 1 b(T)) ["
8tn/2
{'(u) du
+ D(T)
Vem Vem- 1/22t/// 0+1
1+
7
(1 11+0/0 ((T)
(1+ + v(T)) {((T)..
\Vmm-1/2ta/22--x//
1+8
1+8
So we have, substituting these bounds,
(1 (1+b(7))
((T))((7) + <²(7)
(2 x/22)
s 1+8
(1 (1)))) +
(1 + D(7)),
2
2(1+
and hence
5 (1++
and so by monotonicity of v, we have
C(8,K)cE2-R
b(T) C(ō, x)
This proves the first part of the lemma.
138

---

Second column. We now turn to the second column of P(T; To). We again introduce a function V(T) now defined by
x/2
T0):
(1+V(T)
To
with V(To) 0. Then changing variables, from (117), we have
6+x/2
= <(T))
-
x PI2(T;TO)
(123)
As for P12, we have from (117)
PIa(T: TO) =
'"'
T
u
8-n/2
(((u)e-(r-u)1 + V(u)) du.
To
TO
- -5-6/2
Let D(T) sugt
|V(w))
Define
10
/'
88-6/2
(((u)e((ru) du.
7a
To
Then we have a bound
IP1a(T; To) 12(T: To)1 b(T)E(T).
So have, integrating (123) and using the above bound,
l' u 6+n/2
E(T)²
6+x/2
V(T)
((())(()T)) de
T
b(T).
1 8 - K
TO
Using that is increasing, have that
[(+) < E(T)
T
8tx/2
I(T:70)
+ E(T)²
T
ótn/2
D(T).
T 5 K
TO
1 1 8
This proves the second part of the lemma. This leads the inequality
{(F)
((7) 1 -- 5(7)2 S (((
(using boundedness of {) which turn gives
6+n/2
I(7;70),
T
and which concludes the proof.
Completing the proof. We start with
To) = [ 1
x/2
-8 -8-x/2
('(w)e du
Vou k// m0-mmmmm/222/ -1/2-n/2 (r--w) du.
To
Dropping the exponential, we therefore have
[2(TTO) S
cm
1/24n/22
(8+ K 1)TO
(6 + k 1) 1/2+x/2 -x/2+1
1- x/2 8 + k {(To:
To
Moreover, this inequality becomes an asymptotic when 7 7/TO + x but 7 + O.
Turning to the claims on & in terms of P (122), we have (with 7 = m(1 and
m( s))
L (t, s) PiI(T: $1.
139

---

Using E(TO)= S Vem(1 + s) s//= and s(s) = c(1 + 8) s)
Pia (20(8)
_E(T0) 2}
S cg2(1 s)((-m))
m
m
which concludes the proof.
Lemma 1.3 (Transition). For any > O, there is an ma an M1 > 0 and M2 > 0 so that if m < mo so that
1. Uniformly for em? < TO < T < fmP, the fundamental matrices
(IP TO)il M1.
Moreover, the fundamental matrix uniformly close to an explicit expression involving Bessel functions.
2. For < cm
84m-1 (u((Tm) (1=m/2) Piz(T;m))
2 (2) E-D ((cooo
ws)) ((cos((-aus000 (o))
(( + +(()))
In this regime, |&] is bounded above by constant depending only on and not on m.
Proof. Starting from (116), have
x
-1
(((T)
X
dr
y
=
(((T)
6tr/22
y/
We change time to &, which leads
X
1//((T))
(
X
dg
3
1
6+n/2 t)(())
y
We have by definition that Hence solving for this, we have
7l-x//2
1-x/2
To
{ {(T: Vc-
11/24n/2
-k/2
1--/22
& m 1/2-n/2 - -k/2) + TO
1-x/2
Vc
We also have that
TF'(T) Verl-=/2 22mp2m//m2//2x/// 1/2+k/22
{(1 2) 70 1-x/2 1/22 def Eo) kk/2).
We note that (TE Omm)) and hence, in conclusion.
x
(
0 (mP)
-1
X
3 kt1/2 1- x/2
where a
> 1.
dg y
E+Eo
We will use continuity of the fundamental solution in the limit that [ remains bounded independent of m). Therefore, it suffices to solve the equation where we have taken this error term to 0. As a second order differential equation in &, this is
0 to solve the equation (note
(E) E E
X'(E) X(() 0.
From [37, (10.13.4)), with v (a 1)/2 > we have the solutions
(6) + Eu) + c2(E + "Y Eo)
where J, and Y, are Bessel functions of the first and second kind, respectively. This leads to the claimed estimates on the fundamental matrix P(r; To) and an explicit expression terms of Bessel functions.
140

---

Matching solutions from the entrance regime. From the initial conditions which come from Lemma 1.2, in the case that we start with (X)) (1,0) 0) have with corresponding to cm2 that X (8) = + Oe)) and V(s) = O()). Now on sending εo - O we conclude that c₁ + (
and C2 = ( (using [37, (10.7.3))).
If we instead start with (X (X())) = (0,1) 1) we now get at s (provided (em(em?) < E)
k/2 (() + n E(To: O(e))
(taking the asymptotic of I2) whereas
)((s) 0 ());
which therefore leads to the same ej and C2 as in the first case, up to rescaling the solution by 1-x/2 3+x-I5(701 0).
For V, we have
dy
+
y= X.
dp
& + Eo
Using the integrating factor uuE) = (5 + En)", we multiply both sides:
(5 + En) _ dy
+ a(5 Eo)a-1y² (5 + X.
dg
The left side is the derivative of (5 + 50) " so
((E + Eo) ()) = +5 dg
Therefore
V(E) = (5 + Eo) -a (E))(0) + / (s+ 2 ds
Using [37, (10.22.1)), we have
!. (s + (0) ds = ci + 51 Ju+I + 5o) + cz(5 + t + Eo) + 0(1),
and hence
)(E) = ci (5 + Eo) Joti + Eo) + c2(E +
Eo) + O((E + En)
We recall the following large & asymptotic of the Bessel function:
From [37, (10.17.2)], we have
1 1
w(z) = 2
2 1T
4 n
and from [37, (10.173)), as 0o with v fixed,
2
(=)
co8 W, +1=-11))
πT2
2
(=)2
sin Wp + 03223/))
Hence have (for large {) and with W, w(E + Eo)
X(5) = 2
(E + -v- ((1 (e)) cos wt + O(e) sinw, + O(E-)))
π
(2)
(2) $
V())
(5 Eo) -0-122 ((1 + 0c)) cos w + O #1 + +((-1)).
x
Note that sin(w))1) cos(w>) and cos (W sin(w,).
141

---

Finally, we note that
1 8/2+ x/4 V + 2 (1 (1-k/2)
We also note that Eo = m/2 which tends to
Lemma 1,4 (Bulk), For any 6, O, there is an mo if < mo so that
(P); /02 (T - ((r ro)ll & ce (T-10) -1/11))61/---111011 6/2-n/44
Proof. We recall that the range of T is given by
m s 7 <
We first estimate (((T; ro). recalling that
To: RR-1 =R-1(1)
RU(T: (o)
s+x/2
dr
Using trig identities, there matrices with non-constant trig polynomials WA(E) (having constant terms) so that
)
R (-) 6tn/2 R
Wi(E) Wi(E - Id ld
(124)
We need an integration-by-part estimate for oscillatory integrals.
Lemma I,5 contegration-by-aree by-parts estimate). Let p be a positive integer. There is constant ((s) so that for any absolutely continuous
|rre (os( g(u) du s
C(k) ('(To)
(S (o((()) du + sup
lo(w))
ue(soir
Proof. We recall that ('(7) 7 1/2tm/2 Applying integration by parts, have
{ (r((u) g(u) du = l" ine(u) ipE (u) du
(p("(u)
eine(w)g(u)
(pC((u) - S (ne(w) d g(u) du inp((u) du.
Expanding the derivatives and bounding, brings us the claim.
We now divide the range of time into two parts.
Large 7: > To
Suppose that to as well. By standard approximation arguments, follows that for any bounded continuous function f on [-1/e, 1/c).
1/e
(pp(() f(u) du ++ 0 and S dp((4) (1(()) du + 0
-1/c
-1/e
as m + It follows from (124) that the matrix converges weak-* as m -
R (T:(7:0) (-)
6+n/2
0
R(T: 0)
m-+
-
8/2 x/4
Id.
Solving differential equations continuous with respect to weak convergence, so follows that
U(T; m-+0 -/2--//2-///
And therefore for any (including those which depend e) there sufficiently large that
(2 ////2//// Id]| I and [[[(T;; -/2,--00 6/2 x/ Id| <
142

---

Small Ti m <T<E,
We choose mm I < e. Using (124), we change variables to let V(T:Ta) = U(T - 6/2tx/4
Then we have
dr V = (wi(t)+ ) (V).
Hence if we let
|| Ilmax be the maximum entry norm, and we apply Lemma 1.5 entrywise, we get
S
du
d
v(u; ro)
('(Tu)
C(k)
du
d
V(u; To)
du +
sup We [r(oo
[M ro)
max
(125)
+ (())
d
V(u; Tb)
V(u To)
du + sup ME 70/7..7
(r))
du
max
Now we use that
V(u; ru)
(V(u;
V(u;
du
= du
w
And hence
d
((u; ru)
)
1 d
1
V(u
To)
+
sup
|| To)
du
u
TO
du
u²
uE(To.T]
max
max
We also use that Combining all of these estimates, and rearranging (125), we have
d (u;
sup
[V TO)) max - 1
du.
J0
du
max
(Ln du
(1- 2C(k)
2C(Ks)
5 C(k)
v(u; ro)
du
(T())) {'(To)TO
(T()))
max
Recall that
5'(ro) 50 m -1/2+e/2
-x/2
and hence using that _m² < T0
{'(T)) To
1-x/2 m
1/2 n// 2 1+x/2 1 (1 11/2tn/2 -1+n//2
We conclude that for sufficiently small, This leads the estimate
(Lr d du
V(u mo)
du
2C(()) 6/2 1/2222/222222//22 1/2-221/2
max
[IV(T:T)) - -Id] < 2C(r:) 6/2 ( 1/2-k/2
and hence for 21 we have
[[((T:T)) ( (1 6/2-n/4 Idll < 2C(())
152 1/2 - (T)
Combining the estimates. for any To < in the regime.
Using that P and U only differ by a rotation, conclude that for any c,
P((:)))) -e (r (T/ 6622 "R=(
I (r- (1/0)))) -//--6/4
Finally, for larger v, while sharp asymptoties are possible, it suffices to bound the decay of the solutions. We do this with two separate estimates.
143

---

Lemma 1.6 (Exponential Decay). Suppose 8 > 2. There is constant c >> 0 and én so that for all E < and t > To > 2. and (T( = / m 2 1/E
IP(T: ml & (1 ce)e - 6//2-u/44 -17
Proof. We compute the eigenvalues of the matrix that appears in (117). We have
6+w/22 (1+ 6+x/2 )² 4(5))
2
In particular in the regime in which we operate, have the eigenvalues are complex conjugate pairs (as 5' is large), and in fact we have A
S =ič'
We introduce a change of basis matrix H so that
d p= H (
( x+
0 2)
HP.
dr
1) i | = =O(c). Hence differentiating
We note that this eigenvector matrix is within of particular we have H H ld + O(c) and furthermore
owing to the magnitude of g' In
d = (A'
(p- ( dr ( )))
Thus for any fixed vector
C if ((T) |HPu||
d w S (11
6tN/2 + O(:)) w
dr
and hence
|Pu||2 < (1 + O(e)ee -11 (1-)((--o)) (T/To) (T/Ta)-6-n/2 68-1/2
Lemma 1,7 (Slow Decay). Suppose 8 >>> 2. There is constant >> 0 and so that for all 60 and f > To > 2/2
[P(T S (T/To) 5/2-x/44
We further have an improved estimate when ('(To) S is small
max (IPL..(::T)) S (1- C(s, 8) (E(())) exp (S (((((( du
Proof.
From the equation (120) we have
dr ( (u V)RT (-)
O
(V)
5tn/2
R
By assumption we have > 1, and hence the matrix in the middle is dominated by 5+x/22 Therefore
dr (T)((::T)) (T(TT)))) s 8+ x/2 (U²(T: (T(T))))
d
Integrating this from To get for any unit vector E IR2
[M(T; (T/
and hence the operator norm of 14 decays the same way. Since P a rotation of it decays the same rate.
144

---

Improved estimate.
We start with (117)
X
(3)
dr 3
-5'
6tx/2
We set
+ v 4(6))2
2
where we note the radical is real for the regime chosen and it approaches 1 as & - 0. Define W by
W+ = X A
From a direct computation
K dr (gg(A/14))
T / 4()))
Hence we get
d dr
W+ = \ W + +W+ W_)0(1/r),
where (11/T) is bounded by C(K, 6). Uniformly over the range of considered, we can bound S eg' (u), and so
dr (w2 + W2) S (A 0( (W + W3) s 22 (WF + w2)
Integrating from To with g To at most 1
(W?+ (T) & (W? + w2) (To) exp 2 (2/ -('(=((
We have
=(T)= = V W+ W
d+ - X_
so that for an absolute constant 0>0
D2(T) (g'(t) (W3 + W3) (7)
<CE((T) (w? (w? + w3) (TO) exp ( ((w)(((w)) du
Returning to the differential equation,
dr -
Integrating both sides and bounding, we arrive at
[X(T) < (x(r))ee-rr-o)
+C(E(())) / (W3 + W3) (To) exp ( (())((())))dd) du).
I.2 Computing the forcing function
In this section, we prove the scaling law for DANA-decaying as parametrized in Parametrization I.1 using the estimates in Theorem I.1.
We begin by estimating the forcing function.
145

---

Proposition 1.1 (Forcing function). Let a > 2,2 2a + 23 > 1, a. + 2/2 a+1> B. Suppose
that 2p def
a+
20+20-1 Moreover, denote 0(t) def + Bt + (Ic Vro(ss) ds) =
42
1+ (Io
(s) and suppose Parametrization 1.1 holds. Then there exists some C(x,, 3)
such that for any 0 and d large enough
Fo(t) + rnl(t + +aa(()) 55(() S Fo(t) (1) Foo(t) +
where
Fo x d (1-((1-
Plp = A(1) 0(1--1-22 28-1
x -t) 2 Fo.
f0 1, 28 >}
and Fac(l)
2a < 1, 28
0,
else
Proof. We define 3p. aac as Section and use the estimate on क (f, 0) from Theorem I.1 and more precisely Equation (110). Fo is unchanged, however for example for Fpp pp (t) we compute
def f o 2==1 e (1 (t)
V /2cBo² ((1+ t)+ 1)) 2p
1-
+ O(e(72Bo)) 4 (o22 (1 1)) (1
5-
min
V 1- 11)
11
20-1
do
0(1)
28-1 2m
where we used that -2p + 23-1
<-
20+28-1 A similar computation on Fac - 2p
a
gives the result, with the additional condition 2p- > which is automatically satisfied since 20-23-1 - =.
Finally, to bound using by noting that in the range of interest (integral on second line), Φ (t, 0) satisfies the hypothesis in Proposition D.13 (it is constant).
we proceed with similar proof as Proposition H.16
I.3 Stability condition for DANA-decaying.
We can now prove stability condition for DANA-decaying. Proposition 1.2 (Stability of DANA-dccaying using Parametrization L.I under Assumption 6). Con- sider Parametrization 1.1 under Assumption 6. Suppose that
2p
def 8 +
4a
max
22 + 2 - -
_4-
_ and 8+ 1. 2
K
(126)
40
Then we have for any 0 that there exists g(R,E) > 0 and do large enough, such that for any particular. since F(t) is bounded
do Y2 g and S g. then supezo 16 K(t, s) ds (indep of d), we have supazd, ||P|| < 30.
Proof. We again consider the estimates in Theorem I.1. The goal is to have sufficient conditions so that f 3(t, s) ds < e for any t 0.
146

---

We compute the first term by bounding the oscillatory w/ (6) by constants and that 1 (t, s) S e (())) For M > 0 large.
1+C I+50
S 1 to get
f x (1, ds S 2 !'' J
I 3--de (
do
o== -
(Vro(tBB) e --00(1//72l-0)) do ds
+ + B
Tytt T2
Y2Bo
1=0
2 1/
0 g--a [
12Bort do + 92
720 1- 1/a do
+da
Vagtenn 72
$ 72.
Here we used that S e (0(1)//2(000)) ds x mo and that
V /a(BBB
1
x
v 7B
2
Ba
ma()) 1 72 B2
For the second term we do the same to obtain that 112(4, 8) S ((())))))( ((()) 22(e)
Br
We write for B2 (1)++)) da
1 x2 (ts
-3-1 73
(1+8) " Bo do ds
Bo²
o= i
+
B
(Vas(1)B) )
(1+ + s)= "a do ds Bo2
73
0
V B
32Bo
1
73
min (1 + 1)=
Bo (1 (1+1) 11++)=++1) do
0= 4d
0 1-
$ 73
V128(144)
a + 1) - do + 73
(1 (1+ 1)
B2 do
d
$ ra +1)
B12 + 1)) - + 73
(1 1) "122)
(12 B(1 + +1)) #
B2
Va(
Here we bounded for
7
6 3- 1 (
V (t
4 733
(1 7)(00(0))))/ ds
72 Ba
Bo
minf 72 v (t)
(t))
71(1)
Y2
1
S min
(1+
B2 (1 + +)++)=++}
When 72B(1+ 1) 2 La the first term vanishes and we obtain
L/ (1, S
(1 + 1) x
do
72 B
td
7
<
(1 + t) x x d.
72 B
Evaluating at the worst case 1 de
yields the stability condition
bB
2a
73
x d L 1.
(127)
72 B
B
147

---

1.4
Kernel function
Now that we have shown stability of Parametrization I.1 under Assumption 6, we can proceed with computing the kernel function. For that we define upper on the solutions of the ODE:
+ E(1)
$ 8) def (-(())))(())
(++5(s) E(s)
c-(((()-(+))
(1+-E(t)
2pp
Ta(a) Bo² + =((t)2e- f% 0(u) du
(1+E(t) -2p (+{(s))
(I<(())
((s) >
and (f((1,8) def
c (7(1))((())
[++((t) [-((s)
Ta(a)
B (o (1) + ( 1% a(a) du
[+e(()) I+E(s)
-2p
if E(s) 1.
We define accordingly for all / 2 2220
Ke 8) def /- BB (t, 8) duse(o²)
(t, s) df [
2 BB (t, do2)
(((t,)) def K (t, 8) + (t,
K1 (t, del S B2 (t, 5) diuu (0²)
X2 pp 8) del f 7/ (t. 8) dppx (02)
and K Ktp(t,
def Kpp (t, s) Ktp (t, s).
Proposition L.3 (Upper- bound on Kernel), Consider Parametrization 1.1 under Assumption 6. Then an upper-bound on the kernel function is
K(t, s) S 5cl pp (t, 8) + 522 (t,(tss) (t,
where for some 8 with 5 6 20p0 00
11t -8
(t. 7/2B (1 (a+ -2 20
+)
[++)
1+t -3
R2 (1, = 72c(1 + xx (1+ (12 ()))
- 11
1+ +
(Vr2cB(1 + 1)2- /3( + s))2
Proof. We know since pox is a positive measure and (1, 8) S T5 s), (t, S (1, that
K(t,s) s) S (t,, s). There remains to upper-bound bound (t, (t, s). For that, notice that since क (t, s). T (f, s) are well-behaved in the a variable, then using Proposition D.14 and by a similar proof in Proposition 19 we have Vt 2 s 2
xt 8) $ Kg. (1,8) k3(t, 8) $ x2 (t,(t, (t, s).
148

---

\TS(0)B
We only need to upper-bound (t, (t,,,, K2p (t, s). For the first term we have in the case (the other cases can be handled similarly)
K (1, S B
V 7 -3- -e- - (t) do
aa
+ 1
VyzeB(1-a)= 1 4 g - (t- (o (ov (tt))(ttt)-((5 (6+ do
V12ef((L4)* 1-
Vos 722
-(6+
+ +72
o² 3- (12o
ae
1+t
do
(188)
(12cB((+4)) -8
((t)
-2p
+ +22
1.2 =2((1)2 1o.-5 f7 1(u) du
do
V 4221
1+ {(((s)
-8
S B
(1 (1+12B(1- - [)) -
(1+t)
+ (v T2c (1 + +0)1-3)-4+1) -)++3
(1+8)
Here we made the change of variable a - ov V Bt. Notice that the integral behaves as u² (8++22/) It is integrable for 6 large enough, ie 8 + K/2
J. We also computed
f (s' (u)) du
(1+)) (s(t(u) du
Bo 72 Bo²
(1( + +17
Bo
Bo2
where we used d((u) du
Te since u = 1.
This allowed to bound the non-exponential term as
#
S% ((u)) du
1 1++((t)
-2p
do
{(+)
72
V Va(()
(1+5(t)
-20
(18(0(((-))//1
do
V 22
B2
1+ {(s)
V VTa(t)
4
1 ++t
-8
σ do
72 B
1 +
/ 12"
V /a(t) 4-1
1
e (ta(())(t-)))/
72 B
1+8
-8
1 +
(r2 $)) -++ 2
1+ s
Indeed, when Ts(()((-s) 72
> we know that -0(t)(t-e))
(s)
-2+ +
and when (1)((())))
72
72
4
1 we have
S (12 s))- 1.. Hence the error term is absorbed into the other
B2 terms. We additionally used that
149

---

1+ E(t)
-2p
1 E(t)
1 + E(t)
-2p
22
+
1 1+
E(s)
-6++
1+1
2A
+ +((t)=
1+5
-(++4)
1+t
2A
+t (1-N3)pp
1+5
where ò can be chosen arbitrarily large by choosing arbitrarily large. We proceed similarly for the second term
K2 (t,(t,)) S B
VigeB(140 -3- 1 72C
12 Bo (1-3)))))))0) (1+ ((0 do
Ba²
+ B
Vrzchil4e)
03 _ n Ba (ov (a)) -(++)) 12C
(8+
B 8) (oV do
V1zeB(10011g
-66
B
g 3-bb (1=
1+1
B (1 + 8) dd
72C
1+$)
V12eB(1+0)
4
=2(t)2e 1% R(w)
1+
-2p (s())
+
B
do
1+ E(s)
Bo²
7/44
-8
11+1 1
$ 12c(1 8)
(1 (1+ - --)- 20
1 1+
(Vv7(s))(11+ 5)) ()
Here the second integral we made the change of variable u Btt We used that ++Ni K// > 4 1 for the integral to converge.
Proposition L.4 (Lower bound on Kernel), Consider Parametrizatio 1.1 under Assumption Then lower bound the kernel function is Vt 2 2 0
KC(,, 8) 2 2 B( mcB(1ccl++ 1))
Proof.
Note that (=)()=)))))))))())) x around 0. This gives a lower-bound for o S
V 7t
1+ E(t)
-2p
$ (1,8)) (F(()-(()) (((0--))
1 E(s)
(0 ((t) (1+
Since (T,(t,s) s) is positive, that gives lower bound on the kernel by integrating
x(1,8) 2 B j
VigeE(111) - 1 d
2 B/ + 1 ++ -4
150

---

Now just integrate as we did for the forcing function and apply Proposition D.14 since o H e + E(t)) -2p is approximately constant in the region d
VisliBBi'
Corollary L1. Let l(t) def 2 + (Io' Bs x 1 (I'Vis(oBBds) ? For
Parumetrization 1.1 under Assumption 6, 12 B1 2 and (1) S daa, we have
[ * (1) 2 B(O(() 4
(128)
Proof. Just apply Proposition 1.4 and note that Sd f(s) ds 2 So J pz (6) 21.
Proposition I.5 (Kesten Lemma). For Parametrization 1.1 under Assumption 6, and for ō large enough, define
(1, def (T((5)s
12 (
(I())BI) -44 14
7/2
-8 (1+t)
+ (1t s) s)= 1+1/(20)
12
+ +7((s)
(+s)
72 B2 s)
Then under the assumptions in Proposition 1.2 (stability conditions), we have Vt 2 s 2 0 with ((t x max (tt)) $ dda
[ (t 8) dr S (t,
Proof. To prove this upper-bound, the idea is to use that Rpp (t,, s) behaves as power law and that the stability condition Proposition I.2 ensures that its integral sums at most to one S Kyp (t. ds S 1. For-example we compute for any 2 s 20
(r(r)r 2 72 B2 (B (Vss(t)) Bt)
73(8)8) 2
4+ x
4+ dr
92
(8(8)8
1"
(>a(r)r)
2
M
BB v Bt()
1+
BL (r) -1- dr
Jrm
722
72
22
(a(s)s]
(BB(vast))Bt)++r
72
Here we used that (
2 BBB((3a())Br
75
4+ -2-2x-44 +26 is integrable for
2a 1
In the same way we have
1+1
1
1
7/2 B(t -r) + 7a(r))
1+r
-3
x (r2 - 8) (+11((2n)
1+r
2
72 B(r -s) + *n((s) dr
1+
1 +
(72B(1 r) -1+11(20)) (72B(r- -) (1(1/
1+8)
BB 7/ B(1 r)
+ Ta(r))
T 72 B - 8) + +7(()) dr.
151

---

There are four possibilities.
We write
(12 - 5)) 1141/(2a) x rs())a((s)d dr (12B(t - r)) -1+1/(2))))
Brr (12
8)) -))-1+1/(20) x (3( dr
+ (re - r)) 1 B(r -]) - 1+ 1+1) x (s(r)s(ss)
1= X 7a(s)( - $)) - 121/
Here we used for the first integral that y (r) - -1)111/((22))) is integrable when K > for the second integral that (r2 r)) -))--1+((((a))) is integrable. Additionally,
1 and
2a
-1) -
-1111222 (s))
<
12 B(r - -2) (12 - -))--1-1/(2)) (s)) dr
B(r2 - r)) 2+1/(2a) (12 - 8)) 1+1/(2a) (38)
S ra(o)(12B(( - $)) 1+1/(20)
Here we used for the first integral that - 18 B1g r)) r)) dr S 12(12B(t 8)) -1+1/(2)) and for the second integral that ن:= 12 B B(1 r))- 2+1 2+1/(20) dr S 12(12B(1
8)) -1)11(2((2)) S1. 1.
Similarly,
f (2 -))) 2
- (12B(t-r)) - s)) 2+1/(22) / dr
l 1 (12B(t - r)) -)) -11112 (1 8)) 2+ 2+2/ 73 B dr
S 1Ts(s) 8)) s)) -+11 + B32 s))- -
We used in the first integral that (r) X (x) and that Br -))- 2 integrable. We used in the second integral that (12B(t- - r)) 1+1 1201 7a (r) is integrable.
Finally it is clear since (12B(1 - r)) - - is integrable that
8)) -2)1212()))) dr 12 B(2 (f))))))))))))))))))))) r)) -112))
B22 B(t - 2 B(r b))- -2-1/2a)) dr
r= 222 B(22 r)) 2+1/ 73 -2 BBr2BBBrrr - })) - dr
2 2 B(2 B(t 22+1/(22))
152

---

There remains to bound
S. (r(r)rr) 2
12 Be (3t((ttt +
72
x Br2B(r 5)) -))-14)1(20)
1 1++
1 Br - 5) + +ns(s) dr
1
1 11s)
(3(5))
12 B(v (a(1))l)) 4
72
-++((++x)
(12B(r- 5)) 1+1/(20)
1+r
12 B2 - 3) +na(s) dr
x
1+s
(3()))
2
T/ BL (CB(V33(1)) (1())BI) 4 -4
B(2 (] -)) 2+1/(20) dr
72
-842(1+
(1+r)
+
13(5)(12B(r - s)) -))=-)1+((-00 1+1/(20)
dr
1+88
rms
(3())) (BB())))))B)) 443.
2
7/
do is integrable and for the second term that Here we used for the first term that (r - -))-2++ -3 + 2(1 + x) - so that Brr s) -1+ 20 (1 +r) is integrable since k > 2a and that ra(s) + 5)* L.
Finally, we need to bound the symmetric term i.e.
>3())8)
2 B(
Brr 1+3 -1
72
1+t -5
x (12B(t r)) 1+1/(2a)
B2B(1 r) + ra(1)) dr
1+r
72
(3(())
Bl V Bt)
-+
72
-5+4-1
/'
(12B(1 - r)) -1221/222
1+1
12
1 72 B(t - r) + 70(11) dr
1+r)
(Ta())8)
22 ( (LBB(178) Btt 1
72
Here we used that
-(6 + 5) + 4 1 < 0 and that
f B12 r)) - 1+1/(2a)
T2 Btt - r) + ra(r)) dr $1.
Proposition I.6 (Upper-bounl on kernel contribution). Denote for 3 large enough
Kpp (t,, s) def (18(8)5) ) (²B((7s(t)
4+
92
11t -8
72 B(1 8) + 77(6)) )
+ (r2B(11 s) 1+11(220
1
1+8)
153

---

Then if x 1 andd 73 x 1 and under stability in Proposition 1.2 we have Vt 2 2 0 with O(t) 2 d
F * X p1) (() + (t).
Proof. (t, 8) has two main terms we treat separately. For the first contribution of the first term, there are two cases, whether f(t) r 2
is integrable not.
722
Not integrable Then we check that
1 F(s)
Ta($)8
(CBBvra(t)B((++))) ds
72
(
$ 5(()) x
(1+ (1+1)3B((2(()81100)))
72
S F(t).
Here used that since
Ta(t){
(1+ B( +1) 4+)
7/2
(1 + t)
11 - 2i.
$
Integrable Then we write
1 F(s)
13(5)s
++ ds
72
S f" F(s) x
Eepp B (() (7a($)8)
72
$ (t,(t.)).
For the contribution of the second term, just notice that since 5 is large enough, F(s) + s) is not integrable. By cutting the integral in two pieces and noticing that fo (4, ds < we obtain that the contribution is smaller than (t). This concludes the proof.
Theorem I.2 (Scaling Law for DANA deca-yecay,ngg.. Let > 2 and B x 1. Consider Parametriza-
tion 1.1 under Assumption 6. Suppose that 8 is large enough, ie 8 > 8(K, a. (3), Then for any 0
there exists g(n, c) > 0, C > and dg large enough, such that for any d 2 do if = g and c S
we have
Kpp) (Fo + Jac + (t ((1) P(t)
C +J (1) Kpp (1)
where Fo, Fpp. Foe have the asymptotics given in Proposition I.1
(tp = B(1 + 71((t) Bt) 4+ 1_
We believe to behave at least such that
5+ T-
> max 20+28-1 4 ). and (ō + x/2) > 2 3x
154

---

Specific learning rate schedule, ra(() (1+ 1/(2o)
General learning rate schedule, 7t(t)
Fo(t) = 2r 4max (0.1 288)
Fo(t) x =d m-x
(r2Bt) 26-1
if < 1
(t) = (v(()(
Fpp(t) x
(1(1)))-1--2 -1-1 if >
Fo((), ir > 1,
Fo((), 2f 1,2a >
(tc(t) x
0. if2 <1
0, if
28 > 1,203 >1
2 > 1, > 1, Fac x (T())) 2+1/(()) d=
Foe(t) (O(()
Bya(12Bt) Bt) 2411(222)
if <
Kn (1,0))
(pp(1.0)
BT2(7(113) -211/222))
if >
Table 12: DANA-decaying: Large d behavior of the forcing function and kernel function for general and specific 7/3 schedules. See Section B.4 for details about the algorithm. The constant C independent of dimension and the function T(4, 8) remind from Theorem 4.1 that (t) x +++ Bt - (t))². In the case where (t) ~ + t) 1/(aa) T(,,0) 0) ~ B(1 + ()1-1/(1a) The definitions 9, and Xpp can be found in the introduction.
def 1 ru) B du with r(l) def T(1, 0) and
Proof. We know that the solution of the Volterra equation can be written by repeated convolution of the forcing with kernel function, i.e. vt 2 0
P(t) = F(t) + Σ F (t).
For the lower-bound, we apply Proposition I.1 and Corollary I.1. For the upper-bound we apply Lemma C.3 with Kesten's 's Lemma Proposition 1.5 on the upper-bound Kpp the kernel (Proposi- tion I.3) with the upper-bound on F * from Proposition I.6.
1.5 Extended heuristics for general algorithm
In the following we discuss heuristics to justify the scaling laws of the general (DANA) algorithm that are formulated in Theorem 4.1 and more precisely (10).
Claim I.1. Denote ? (1) = 1, (1( ni 73(() = a2 + 1) A 8(1
-) and B cgd Denote p def =+
1- T
For a 1/ take 2p > max 20
4 } and 2p k2/22 >
2 1, 72 TJs = 0. Same scaling laws SGD.
M Ki
722 = 173 = 0. Scaling laws in-between DANA-decaying and SGD.
2x
0 S
< -/2 1/2 = 0, 173 = 2a 2x - -Ka). Scaling laws in-between DANA-constant and
DANA-decaying.
The scaling laws are given by (10).
If 2p max 20+25-1 2) then the exponent in Fac, is replaced by the minimum between its exponent and plus SGD exponent after started accelerating.
155

---

15,1 Forcing function
Claim 1.2 (Forcing function), Let a > 0, 2a + 2 > 1, ox + 2: a x+1> 3. Suppose that
2p def
54 2a+28-1 There exists some C( 3 > 0 such that under Parametrization I.1, denote
(1) def
+ 2y2 Bl (Io vra())Bd)) we have for any i 2 and d large enough
Fo(t) + ppp(t) + Fac(t)
S F(() S Folt) + Fpp(t) +
F Fac(t)
C
where
F0 x d 2add-att1-22))
0(t)))((t)) 28-1
= 0(1)-1+2 ff 1, 28 >
F (t)
Fo if2 <1, 23>1
0 else.
We define Fpp F ac as in Section D and use the estimate on Φ (t, from Theorem I.1. To is unchanged, however for example for I pp((t we compute
Idea:
F (t) def /¹¹
ol+ 2011
V2 (
(1++))-- 1)) --))
1 g
2
+ D(e(122 6 f e
min
V12 V 2- 11
28-1
0 do
o=0
x 0(t)
20-1
2a
Where we used that 2p 1- 2011-2-112 + 2p 2a+28-1 Similar computation on Fac
a
brings the result, with the additional condition 2p > 2 which is automatically satisfied since 204229 > 2-1 1.
Finally, to bound F using Fo()) p() + FI) we proceed with similar proof Proposi- tion H.16 by noting that in the range of interest (last integral), Φ (t, 0) satisfies the hypothesis in Proposition D.13 (it is constant).
I.6 Stability conditions
Claim L33 (Necessary and sufficient conditions for stability above the high dine-iiimnsional line a > and with batch B 1). Let > 1/2 Consider Paranetrization I.1 with B 1. > 0, 72 > 0.
Suppose that 2p > max 22+2281 4 3. We have
(Sufficient condition) For any > 0 there exists g(K, E) > O and do large enough, such that for any d 2 do if (N then sup,20 50 K(t, s) ds < e. In particular, since F(t) (is bounded indep of d), we have supazda |Pll < o0.
2 72 and c: S g) or (N < 2a 72 g and S gufaax--tt )))
1
(Necessary condition) For any c > O, any g > and any do E large enough, for any d 2 do and if (R < Taa g and gd with 0 > k > 2a(x 12),, then there exists
σ1 > 0, 02 such that for any daa > t N da-oi we have Sip2 K((, s) ds d In particular under this scaling supd2da ||
156

---

For any e >> 0 and >> 0, there exists some g > O such that 12 S then for do large
enough and any d 2 do. sup So K(t, s) ds S and consequence, || co.
Sufficient condition We again consider the estimates in Theorem L.. The goal is to have sufficient conditions so that Jg X(t, s) ds < for any 2 0.
Idea:
We compute the first term by bounding the oscillatory w (E. Eo), (E, by constants and 1+C For M >> 0 large.
S 1.
T+
S (((),)) S B
[ [
I3 _ B2
do ds
J 3d
x 72
12Bo do
x 12d
For the second term do the same and that (725( Bo (1 + s)-* to write for B1 + 1) daa
0 K2 (1, SB
-3d 73
(1 + +) - da ds
t Bo
S 73
o min(((+1))
B2 (1 + 1) xx+1x+1111 do
1
73
V128(144)
o² a (1 + n
x+1 do + 73
(1 (1+ t) t)=
B do
V-28(140)
B12 + - $ 7311 + +1)
73
(1
B( + f)) _.2
B2
When B2 + 1) Z La the first term vanishes and we obtain
I (t,
x
73
(1 + t)
-- do
B
4d-0
2 73
(1 -* d.
7/2 B
d=
Evaluating at the worst case f
yields the stability condition
73
d S
(129)
72
Y2 B
Necessary condition To obtain the necessary condition, just observe that the previous upper-bounds can be treated constant. Since we
are in fact tights. Indeed, for Eo x E, the term
1+E
additionally know lower-bounds on w (E), (E),w?(() we deduce the corresponding lower bound on (1, (1, s). This brings that for S dea
L" K(t, 9) ds 2 ((2))) 1/2
(1+ kt 2..
Since by assumption, R + 2a( + ) we obtain the existence of 01 0, 02 > 0 such that
1A 2 Bo--0 we have f K(t, 8) ds > d Now we know that 1t 2 0, Va 0. F(I) Z do
157

---

Aftert t 2 d we hence obtain by recursion that the loss grows as d'3 ("2))ooz(()) This brings that P( 2 d02 (dg, We see that this diverges as d + oo which concludes.
Third point The last point in the proposition states that broken, the loss remains bounded (even though can increase arbitrarily with dimension for times
when the conditions for stability
t I ddn To see it just observe before that Jo XC (t, s) ds S and that for 72 2 d we have
d + O. So (1, S -)=* T B
1.7 Kernel function
Now that we have the stability conditions we can compute the kernel function. We will focus on the contribution from & (t, which we remind the behavior from Theorem I.1
{(t) 1+ {(s)
20
(t, s) x (r(()-(()))
Then we can give the behavior of (t, Claim 1.4. Oft x da have
x (t, 0) = 1 (t2 2t + +VV7a(1)B +
+ 4+ -+
Idea; As for the forcing function, there are two cases depending on which contribution from B2 or (v (())B)) is dominant in 0(). Hence we decompose
minl m221 Vitcoe (7(1) (-((1))(()) E(t)) 2p do
x (t, 5) x 12
d-a
x B
Compute-optimality beyond stability and motivation for DANA-decaying schedule
In the previous sections, we restricted the learning rates domain to describe stable algorithm, ie. requiring that the risk P(t) stays bounded for all time (0, oo). In this section, ask and heuristically answer the following question:
Does there exist for classical momentum (equiv. SGD) or DANA-constant a scaling in of the learning rates 72, 73 which yields a better compute-optimal frontier without requiring the algorithm to be stable for any time E R4?
J.1 Strategy
In the previous study, we imposed stability of the algorithm by controlling the kernel norm Vt 2 0, f' [(t,, 8) ds < 1. This, combined with corresponding Kesten's 's lemma (see [80, Lem. C.1] for SGD, Lemma F.3 for classical momentum, Lemma C.2 for DANA-constant) ensures that the resolvent s), def E X** (t, s) of the Volterra equation (55) is bounded for all time and leveraging [45] implies that the solution P(t) = F(t) + [9 (t))) is bounded.
Instead it is clear by following the same steps, that for some 7 > 0, T, Jr K(t, 8) < 1, we can recover that Ptt) stays bounded for t < T. More precisely, retracing the same steps in the proofs that give rise to Eq. (62) (see Theorem G.1 for SGD-M,
Theorem H.3 for DANA-constant), would obtain Vt < T, P(f) ((t) 72 K(t, 0). On the other hand, if for T, S ((t, s) ds 1 then, starting 2 T, P(t) will start diverging exponentially. Taking T = +oo (or equivalently La (122 B) for SGD and min (d²//(1BB) /v for DANA-constant) recovers the previous study.
158

---

Largest )/2 stable at compute-optimality
K
Compute-optinal 72
Phases
Any
72 x
7/2 x 1
la, II, III
Phase lb
K
20
72 x (
Y2 x d2a-
Phase Ic
4a²
2a(2a 1) 92 1
722 20+20=1
2n -)
20-+28-1
2
Phase IVa
2x
1/2 x da-1
12 x
(cf [80])
Phase IVb
K o220
X
1
(a-1
72 x
20
(cf [80])
Table 13: and optimal 7/2 across the 4 phases for SGD.
J.2
Stochastic Gradient Descent
We will focus on SGD, although classical momentum can be handled entirely similarly. From [80, Sections 1.1, 1.2] batch has no effect on the compute-optimal frontier. We choose WLOG B = 1. We additionally know from [80, Prop. G.I, H.5] and Proposition G.5 that for 7/2 Bt S d2x we have K(t,,) s) = 7 B min (12
s)) -2+11(20)) It is hence clear that
[ KT
VI < 2a
sup
K(t, 8) ds x 72 min BT (1BBTT))-+1((2a))) 1+1/22a)3.
Taking BT x 2 we recover the stability for all time condition 72 S (0,ax00.2a_)) Additionally, if a 1. 2. we still get the stability condition 722 S 1 independent of T. However, for < 1 we now
obtain for S 12 BT S d2a the condition 7/2 S (T2
2 d2a-1
The question is now: for a < 2- can such a larger "/2 improve the compute-optimal frontier? To that end, denote T the time at which compute optimality is reached. Further introduce x such that BT = d" For a < 2. we hence obtain the condition 7/2 3 d d=x(-1++)))
1
We can now proceed to the same compute-optmal study as in Section E with the new k variable which gives results summarized in Table 13. Above the high-dimensional line, we obtain the same results as in Table 5. In phase I.b, since compute optimality is reached for BT, we do not see
improvement over Table 5. In phase I.c, we obtain 7/2 X d d (20-211)
>> d2a-1 which yields faster compute-optimal curve. However in this phase, we do not have proofs about the kernel asymptotics which stops being power law. In fact it may be possible that a much larger learning rate can in fact be used. finally phase IV, although the learning rate could be chosen larger, [80] showed that compute-opiimalit was reached by a smaller learning rate.
JJ3 DANA constant
The kernel is slightly more complicated in DANA constant and for simplicity we will use here simplified form which becomes valid for s ス I2 and 1+t Its 2 1 This will not affect the main results as this is the dominant term in the kernel norm. To see that, the reader can either notice that the upper-bound K derived in Section 10 is in fact tight, or directly integrate the estimates 111, 4 first on time and then on a for a complete proof.
K(1, = (1, 8) + (1, 8) = 73 B(2 x)) 2+ + (a(((1 +)) -1-
Again, denote for T 2 0. BT dd the compute-optimal time for some 0. As for SGD, we obtain the stability conditions
{
122 + - (12 $1 if 2a > 1, 12(72BT) -1(1/2))+++1(20)) (21/((2))
+ - (T) 1/20g) $
2a 1.
159

---

Optimal 72
Optimal Y3
2a
Phase la
Y/2 x 1
73 x d
2n x d
Phase IIa
x 1
Phase IIb
72 x
73 x d 3[10-1)
2a
Phase IIa
Y/2 x 1
x d
Phase IIIb
12 x
Ya x [dn-1)2
Table 14: Optimal 72, 73 for DANA constant for a > 2. The results are valid up to a constant independent of the dimension
Again, if we use the upper-bound 72 d=, we find the stability conditions "22 S d (0,20-1)
and 7a B S d. However, these bounds are overly conservative, and we can use our study of DANA-
constant to improve them by computing s. For that, use the write P(L) Fo(() (()+))+++(((((aaa)a) Fao))) pp 0) with F() 3 (72 Bt (v73 (1))), Bi))), pp, 0) KSGD (72 Bi V VT3BB)) We believe this can be made formal with an entirely similar proof 8) although this is technical and not very enlightening. Again we will focus on B 1 and more precisely a 1 which has slightly easier stability conditions. For 2- we necessarily have the bound "/2 S 1. Therefore, we can easily compute as a function of Y3. B at compute optimality, and hence deduce
pp
did in Section B.3 by using our
estimates of 4 ( s), 12,24,
722 Using this bound, we can also check that it is feasible which means T > 22 which used for the derivations, and hence find the results for phase Ia, II, III Table 14.
"
Finally, one will note that (V1gBT)² = r())² the corresponding time of DANA-decaying around compute-optimaliyy time T. This implies the following:
The previously computed 23 for DANA-constane induces the same dynamics as DANA AANA-dccaying around the compute- compute-optimal time T after which DANA-constant starts diverging exponentially.
J.4 DANA-decaying
Exactly as we did originally for DANA-decaying, we can try to construct a better learning rate than DANA-constant, by making sure that the value 73 wanted is attained compute-optimality 72r = d= In other words, suppose that in some phase, DANA constant asks for 733 72 Bd- na and that compute optimality is attained BT d Then DANA-decaying would use a schedule as (t))) Bt) -*a/N.
Computing these values from Table 14 recovers exactly in all cases (at least above the high-dim) (1+ (1++) 1)(22)), This is because this step-size is already optimal by ensuring that f K2 s) ds x B (12 BT) (/2a) < is on the verge of diverging as t - co. Note that this was also the criterion chosen for DANA-constant in the previous section at the point of compute optimality which explains the correspondence.
7a(t)
Above the high-ding dige-dimensional line, we can write the three equivalent characterizations of the def DANA-decaying learning rate (1) 72 ((12 ( + ()) (2(());
the largest power law decay that satisfies the DANA-constant condition (x(t) < d B
for í 2a the time where the problem starts being strongly convex and the solutions stop behaving as power law and decay-exponentialyy
160

---

the largest power law such that the stochastic noise induced by momentum Jo K2(t, 8) ds Ta(0) B (t) remains bounded for all time & 0.
the power law which recovers ?a(T" " the constant DANA-constan learning rate computed in section Section J.3 for which the kernel norm remains bounded up to the compute-optimal time T.
161

---

K Power-Law Random Features Experiments & Numerical Simulations
In this section, we describe the PLRF experiments, measurements of the empirical scaling law exponents, and the numerical simulations of the ODEs shown in Figure 8, Figures 25-33 and Figures 34-42.
K.1 Power-Law Random Features Experiment Details
For each (a, 3) pair, run the stochastic algorithms for SGD, DANA-constant, and DANA- on PLRF with parameters
d E [200, 300, 400, 600, 800, 1200, 1600, 2400, 3200, 4800, 6400, 9600, 12800].
and set 4 x d. We selected 84 pairs of (a., 3)'s that provide good representation of all the different phases spanning a E [0.2, 2] and 3 E [-0.15, 1.4]. We use batch size = 1 for all PLRF experiments and compute the mean population risk across a of random seeds. We 10 random seeds for most experiments but use 100 random seeds when 3 < 0.3 due to the increased gradient noise when B is small.
We calculate flops as the number of training steps x We train all models for 1e12 flops or convergence, whichever occurs sooner.
The loss curves in Figures 25-33 and Figures 34-42 show nearly perfect agreement with the numerical simulations of the simplified ODEs (23) across all three algorithms and all values of (a,, () and d. In particular, this numerical agreement validates that the simplification in the ODEs between (6) and (7) has trivial numerical effect.
We define Tr(D) def Ei11
K.2 Measuring the Empirical Scaling Law Exponents: Chinchilla Approach
To measure the empirical scaling law exponents for the loss and parameter count, we follow Chinchilla Approach [50]. We note that for SGD on PLRF, similar experiments [80] found close agreement for exponent measurements between Chinchilla Approach and Approach 2. We do not expect perfect agreement between theoretical predictions for the scaling exponents due to noise the loss curves, discretization in the parameter counts, and finite-size effects noted in Approach 0 in [80] that showed the 'instantaneous' power-law exponents vary as a function of flops well beyond practical scales for experimentation.
We follow Chinchilla Approach I as described in Section 1 of [50]. First, we choose a flops window Ihinn fuax] and construct fy's fs's using geometric spacing between fi = fmin and f. fmax- For each flops slice fj. we find the minimum loss across all d. We denote this minimum value of the loss, 3o* (fy), and its associated parameter count, d" (fz). This creates 'envelope' of optimal losses vs flops as shown in Figures 25-33 or a 'staircase' of optimal parameter count vs flops as shown in Figures 34-42.
The coppull loss exponent found by plotting
(ff T0 (. (fa)))ssisn)
(130)
and fitting a power law curve the form (P(()) = c f-, For the 84 different pairs of (a, B). plot the fitted power law exponents and list the predicted theoretical loss exponents (see Ta- ble (SGD/SGD-M), Table (DANA-constant), and Table (DANA-decaying)) in the legends in Figures 25, 26, 27 for SGD, Figures 28, 29, 30 for DANA-constant, and Figures 31, 32, 33 for DANA-decaying. The (absolute) maximum difference between theory vs. empirical is 0.09. Figure (middle) used the loss exponents from Figures 25-33.
The compute-optimal parameter exponent found by plotting,
{(tz, d' (f,))sissn
(131)
NNot that for consistency with the theoretical setup. for the PLRF experiments we omit the the 6 x num parameters x tokens heuristic [50. 54J often used for calculating flops. We include this 6 in the LSTM experiments for accuracy and consistency with the empirical scaling laws literature.
162

---

and fitting the function d' = fd where c, b are constants. The empirical parameter exponent measurement is given by the exponent For the 84 different pairs of (a, 3). we plot the fitted power law exponents and list the predicted theoretical parameter exponents in the legends in Figures 34, 35, 36 for SGD, Figures 37, 38, 39 for DANA-constant, and Figures 40, 41, 42. The predicted theoretical parameter exponents can be found in Table 5 (SGD/SGD-M), Table 6 (DANA-constanl) and Table 7 (DANA-dccaying). The absolute maximum difference between theory and Approach is 0.132. Figure 8 (right) used the parameter exponents from Figures 34-42.
important note that the fit of the exponents is sensitive the choice of the flops window; changing the window can have dramatic effect on the predictions. For Approach 1, we have discrete set of parameter counts that approximate the 'true' compute-optimal frontier that would result from an arbitrarily dense sweep over parameter counts. We therefore want to fit the power laws for Approach within the range of flops where the parameter counts in our experiments act as a good approximation of the compute-optimal frontier. More specifically, as a heuristic, we want to fit Approach I within the flops window whose minimum the flops where the two smallest models crossover, and whose maximum is the flops where the two largest models crossover. In particular, for larger values of alpha, the loss curves from the larger values of d do not intersect because the experiment could not run long enough, so we want the flops window maximum to be the crossover between the two largest models that did reach a crossover.
As first pass, we used this heuristic to determine the flops window for each algorithm and (a, 3) pair using smoothed loss curves to locate the crossover points between model sizes. However, due to noise in the loss curves, it is nontrivial to systematically distinguish which crossovers in the noisy curves are due to noise and which represent the location of the 'true' crossover that would result averaging over an arbitrarily large number of random seeds. We therefore manually adjusted the flops windows using visual inspection to adhere to the intent of this heuristic.
K.3
Computing the deterministic equivalent for K
In order to derive fully deterministic curve for the loss, we need get precise estimate for the deterministic equivalent of K = D+/2WW/2 D1/2 where D = Diag(j 2a particular, we need to solve the fixed point equation in (33),
S S v}. In
1
ml m(2)= = 1 + 2 1j=1
where R(2) = Diag (-20 - : 1 $1sjs0
m=))=
(132)
In order to implement and solve such a fixed point equation, given that || llop Re 1. we uniformly discretize the z's from c x j-2d to 1 + for some small and add small imaginary component. For each of these 2's values, we solve the fixed point equation in (33) with Newton's 's method constrained so that the Newton update always remains in the upper half plane.
K.4 Implementation of the ODE
To implement the ODEs in (23), we use implicit Euler with step h and use the deterministic equivalent as the initialization. While other ODEs can be used, implicit Euler is known to work well on problems whose solution is exponentially decreasing, but ill-conditioned exactly our setting. To improve the speed of implicit Euler, we perform a time change mapping t - e'. This allows for log log-spaced time and an exponential speed up in solving the ODEs. We note that one does need solve d ODEs simultaneously, but the ODEs form nearly decoupled system of 3 ODEs as the coupling only occurs through the forcing function (I)). Numerical experiments in this section (Sec. K) and Fig. use a step length of h = 10-3 (after changing into log spacing). All other numerical experiments in the paper use a step length of h 10-2
163

---

L LSTM Language Model Experiments
In this section we present experiments for SGD and DANA-decaying on LSTMs trained on text data. We constructed a language model setup using cross-entropy loss on next-token prediction on the C4 text dataset [84]. We use the LSTM [49] architecture from [53] with standard parameterization.
To sweep model sizes, we co-scale the embedding and hidden dimensions while keeping the depth (number of layers), sequence length, batch size and vocabulary size fixed. The embedding sizes used are [16, 24, 32, 48, 64, 96, 128, 192, 256, 512, 1024, 2048, 4096] and the hidden dim is set to 2 x emb dim. All models use depth = layers, sequence length 20 and batch size 32.
L.1 LSTM Results: DANA-decaying across fig
In this section, we train models with SGD with learning rate = 0.5 and DANA-decaying with 12 0.5 (t) 72 (1 t) -)- where sweep K3 E (0.0....,, 1.4].
We first consider single model size for the LSTMs (embedding dim 128). In Figure 19, we repeat Figure 3a on the left and Figure 2a on the right to show side-by-side that the behavior of DANA-decaying across Kig qualitatively similar on PLRF (left) and LSTM (right). On PLRE, we see that DANA-decaying diverges when K3 < 1/2a. On LSTMs, we see that DANA-decaying similarly diverges for K3 < 0.6. Moreover, we see that for moderate values of K3 DANA-decaying outperforms SGD in both model settings, and that K3 approaches 1.0, DANA-decaying smoothly deteriorates back to SGD.
1.7, P 1.0, 800, vilt (t+ 1
theary, dana-decaying theory. sqd thesry, dana-constart
10
1o
08
Validation Lobes
0.55
67
risk
0.5
09
10-4
0.4
055
10a4
10-11
10a
10a
flops
10a
10a
10t2
DANA-decaying
1d
10a
104
Traiiing
Figure 19: DANA DANA-decaying shows similar stability and divergence behavior across K3 on PLRF and LSTMs. (left) We repeat Figure 3a PLRF where (t((t) x (1 (1+1) +) - This shows DANA- decaying diverges when f3 < 1/2a and degrades to SGD as K3 goes towards 1.0. (right) We repeat Figure 2a on LSTMs with embedding size 128 showing DANA-decaying diverging when K3 0.6 and degrading to SGD when Kg exceeds 1.0.
We next repeat this experiment across model sizes, with embedding sizes ranging from 16 to 4,096. In Figure 20 we see that the behavior of DANA-decayin is qualitatively similar across model sizes: we continue to see divergence on small #3 values, outperforming SGD on moderate K3 values, and degrading back to SGD on #3 values above 1.0.
One interesting question is whether the optimal value of K3 consistent across model sizes. To investigate this, we plot the final loss versus K3 in Figure 21. We see that the optimal K3 is fairly consistent across model sizes, with the optimal value of K3 falling within 0.55 0.75 on all model sizes with a trend towards 0.75 for the large models. We note that the largest models may slightly overestimate the optimal Ki they did not run far enough to reach divergence. Using the PLRF as an analogue where Kig 2a this optimal 83 near 0.75 would correspond to 0.67, but we note that the precise meaning a clear for STMs.. We leave this question about how to measure and interpret real-world problems for future work.
164

---

Embied dim:
Embed dim:
Embed dim:
10
16-
Validation Loss
220
563
550
d4li-dccyyng
DM decayycyg
Embed dim:
Embed dim: 64
Embed dim:: 96
10
Validation Loss
SGD
SGD
167
10"
321
10
300
187
16
107
17'
10
10
170
27
16"
107
17'
10
10
Embed dim: 192
Embed dlim: 256
Embed dim:: 512
Validation Loos
DANA-ckca/1g
BollA-bccap
12
167
18
10
10
10
10
15"
10
17
17
10
10
100
10
32"
10P
10
18
Embed dim: 1024
Embed dim: 2048
Embed dim: 4096
10
Validation Loss
525
5GD
550
10 Training Takens
Friining Tokens
Tra Training Takens
0.6
DANA-decaylng Na
Figure 20: Sweeps of Kg for DANA decaying show qualitatively similar behavior across model sizes. Small values of K3 diverge, moderate values of K3 outperform SGD, and large values of Kgg trend back towards SGD, Each panel shows one model size with SGD in black and DANA-decaylin in bright colors indicating the value of K3.
165

---

6.0
Emb Dm
5.5
4096
2048
5.0
512
4.5
256
4.0
12
3.5
0.4
0.5
0.6
1.0
1.2
deca-ig R3
Figure 21: Final loss versus DANA-decaying K3 showing the R3 which optimizes final loss is fairly consistent across model sizes. The optimal value of R3 for the final loss is between 0.55 to 0.75 on all models with trend towards 0.75 for the largest model sizes. Color indicates model size. Black dotted lines indicate the final SGD loss for the model size indicated by the color of the highlight.
L.2 LSTM Results: Loss Exponents
We next measure the LSTM loss exponents for SGD and for each individual value of Ki for DANA- decaying, by plotting the loss curves across model sizes and considering the compute-optimal frontier defined by the lower envelope of these curves. We observe that the compute-optima frontier appears to follow power-law trend for an intermediate set of models, but the power law breaks on the largest model sizes similar to what was observed in Figure 7 of [54] on a similar LSTTM scaling setup.
We perform a version of Chinchilla Approach [50] using an intermediate set of model sizes ranging from embedding size 16 192. We use an abbreviated version of this method where we fit power laws through the crossover points in the loss curves between adjacent model sizes. Due to the smoothness in the loss curves, the locations of the crossover points between model sizes can be precisely determined. The power law fits have very high R values, exceeding 0.99 on all settings except the most unstable Ki 0.0 which has R 0.984.
In Figure 22, we show the loss curves and compute-optimal power law fits for each algorithm and value of Kg, and report the loss exponents and R2 values in the legends. Finally, in Figure 23a plot the loss exponents as function of K3 and show the SGD loss exponent as a baseline. The magnitude of the DANA-dccaying loss exponent maximized when K3 0.7 and exceeds the SGD loss exponent magnitude. We note that the measurement of K3 = 0.7 that maximizes the loss exponent magnitude is very close the K3 near 0.75 that optimizes the final loss, which is interesting because the loss exponent measurement is determined using earlier phases of training whereas the final loss values naturally describe the behavior from late training. While the numerical differences between the SGD and DANA exponents are small, the high R values and the smoothness of the DANA-decaying loss exponents as function of K3 indicates that these measurements might be robust and imply true improvement in the loss exponent over SGD.
More notably, the DANA-decaying loss exponents traverse exactly the regimes seen in PLRF: diverging for small values K3, outscaling SGD for intermediate values of Ki, and reverting back SGD-like scaling for K3 1.0. We note that A3 1.0 corresponds Schedule-Free SGD and shows almost exactly the same loss exponent as SGD for LSTMs.
166

---

5-0.0
5--2
Valiortton
0030.
Kem0.4
Valiaaoonn
0001 4.54 02-1
44 81-0
Lces 00 -330 87-3392
Lc0000 -030 P--05
Tonttonn
0..5
50-
201 122
127
127 12"
5,70.
1-10 4-12
10
Eanton
1003 6.59 -1-63
LCH0 +.30 55-95
Lcsl 0.39 (81-65990)
Figure 22: Loss exponents for SGD and DANA-decaying across K3. Loss curves across model sizes (black light blue) and power law fits (orange) through the compute-optimal frontiers showing the loss exponent measurements (with R² values) for SGD and DANA-decaying, The first panel shows SGD and remaining panels each represent one value of fg for DANA-decaying.
167

---

-0,050
22
DANA
SGD
SGD
2.0
Divergence Stabi boundary SGD boundary
211
1000
0.054
Loce Exponent
-0.056
1.9
-0.058
1.0
DANA-constant
177
0.060
-0066
0.5
DAMM-decaying
1.5
-.0664
Schedule-Free
0.4
0.5
0.8
1.0
1.2
0.0
0.2
0.4 Mi2000.6
DANA-decaying
Ks
Figure 23: DANA-decaying loss exponents show improvement over SGD, and traverse the regimes seen PLRE (left) Loss exponent versus K3 for DANA-decaying. The loss exponent for SGD is shown with black dotted line. The loss exponents for DANA-decaying have higher magnitude indicating an improvement in the loss exponent for 0.6 S Ki I 0.9. (right) We repeat Figure 6 showing the regimes for PLRF (left) traverse the same regimes as the æ-axis of the right figure: divergence for small values of Kg. outscaling SGD for intermediate values of K3, and reverting back SGD-like scaling for из 2 1.0. Note K3 = 1.0 corresponds to Schedule-Free SGD, and shows almost exactly the same loss exponent SGD for LSTMs.
function of K3- We note that the LSTM loss exponents
168

---

10
SGD Equivalent Learning Rate
Validation loss
0.5
0.2
0.1
0.05
SGD
0.02
SGD + momentum 0.9
SGD + momentum 0.95
0.01
SGD
momentum 0.99
5
10a
104
10
10
107
10
10
Training tokens
Figure 24: The equivalence in risk dynamics between SGD and SGD-M holds approximately in LSTMs. We sweep different values of momentum (different line styles) and different effective learning rates where effective learning rate learning rate / (1 momentum). Color corresponds to the effective SGD learning rate. Loss curves of the same color have similar dynamics showing that the equivalence in PLRF dynamics extends fairly closely to the LSTM setting.
L.3 Equivalent risk dynamics for SGD-M and SGD
As noted in Section 5 and Remark G.2, there is equivalence in risk dynamics on PLRF when NOD SGD-M SGD-M/ /6.
In Figure 24, show this equivalence holds closely in the LSTM setting especially after the early phase in training. We perform a two-dimensional sweep across learning rates and momentum values. We train SGD (corresponding to m 1.0) and SGD-M with three different momentum values of m E [0.9, 0.95, 0.99]. Following the equivalence, we sweep the SGD "effective learning rate" denoted by 77 and set the learning rate to n/(11 m) for each value of m. We sweep r) [0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 4.0).
In Figure 24, learning rate (same color), particularly after the early phase of training. While we do not measure the loss exponent for SGD-M directly, this equivalence suggests that SGD-M should have nearly identical scaling behavior to SGD and result in the same loss exponent.
that all values of momentum have similar risk dynamics for the same effective
L.4 Experiment Details
Initialization. We use standard parameterization to initialize the LSTM parameters. All matrix parameters use Gaussian initialization with zero mean. The embedding parameters are initialized with standard deviation 0.1 and the hidden (LSTM cell kernel and recurrent) parameters and readout parameters are initialized with standard deviation 1/vfan-in. Bias parameters are initialized to zero.
Dataset. All models are trained on the C4 language dataset [84] encoded with the T5 Senten- cePiece [57] tokenizer. with an additional beginning-of -oaiinng-o--uuueeee sequence (BOS) token, resulting in the vocabulary size of V = 32, 001 (32, 000 original vocabulary + BOS). The effective vocabulary dimension in experiments is 32, 101 due 100 unused tokens. Both training inputs and evaluation inputs are padded rather than sequence-per aequence-packed. We use the 'train' split for training and the 'validation' split for evaluation from the TFDS [1] version of the dataset.
We train in the single epoch regime: the C4 dataset contains 156 billion tokens, but our models never train past the first few billion tokens and we never repeat training examples.
Implementation details. Our LSTT implementation started from a fork of the code released with [53], with changes made for performance, initialization, and switching the dataset. The LSTM
169

---

implementation uses JAX [32] with Optax optimizers [32]. We use a modified version of the data pipeline from NanoDO [65].
Hyperparameters. We first swept the SGD learning rate and selected stable learning rate of 0.5 which we use for all SGD experiments except the learning rate sweeps in Fig. 24.
For DANA-decaying, we set 72 0.5 to correspond to the SGD learning rate, and 7/((t) = 722 x (1 f)= We then sweep over K3; the minimum K3 where DANA-decaying converges appears optimal
(Fig. 2c 20). Since just needs to be large, we set
Flops calculation. We compute flops as 6 non-embedding parameters x training tokens using the '6ND' heuristic [50, 54]. We follow the common practice of parameters as this has been shown in [50, 54] to induce better power law fits.
counting only the emmedddnnn
170

---

M Additional Information on the Experimental Set-Ups for Figures
In this section, we provide additional information regarding the experimental set-ups used to generate some of the figures. For most figures, we have included these in the captions. Due to space limitations, some figure experimental set-ups are included here.
Figure 1a. Experimental setup: d = 500, 2500, batch size 1, 10 iterations, with learning rates 72 = (D) for all methods, " 0.1 0.1/(Tr(D) x d) for DANA-constant, and 7a(()= 0.1 + t) 1/(aa) for DANA-decay.
Figure 1b. Numerical setup: SGD (blue curves) learning rate 722 0.5/Tr(()) DANA- (green) has (t) = 1, (2(() = 0.5/Tr((), 7t(() = 0.1/d, (t) = 6/(1 + t) where 8 = max
1/a,
DANAAA-eAaiN (orange) same y1: 72, and A as DANA-constant, 0.1/(1 + f)2/(2a); Schedule-free SGD (red) 3 0.9 and 7 (D) in [33]. Algorithms run using the ODEs (43) with hyperparameters given in Table 3 for Schedule-Free SGD; 10 iterations of algorithm, d {100 x 2 ]. i I...., 10 and v 10 d. Schedule-free SGD (red) scales very closely with SGD (blue) for this (a, 3). We see both DANA-decaa and DANA-constant accelerate.
Figure 1c. PLRF setup: number of iterations for all algorithms is 10², 1.2, = 0.7, batch size with (250, 500, 1000, 2000} and t = d x 2. Adam [56] algorithm was run from Optax [32] where the default parameters (i.e., 81 0.9, {3 0.999, the Adam learning rate, a cosane-decay schedule, also from Optax, was used. Initial value for the cosine learning rate is 0.001. For each d-value, multiple runs of adam with the decay steps in the cosine-decay learning rate given by 2 x 1000 X mi, m = 0, l,... until decay steps is greater than total iterations, 10". For each d, (solid green) lower envelope of the Adam runs; (faded green) one single run of Adam with a fixed decay step. SGD (solid orange) and DANA-decaying (sold blue) plotted using ODEs. DANA-decaying hyperparameters: 5 max (2 + (2.3 1)/a:), 2 1/a) + with (t) 6/(1 + f): (t) 0.1 O.) (1 1) 1/(2)); (2(() ()/22a))) 0.3 (.3/Tr(D):niln
10-8) were applied. For
1. SGD
hyperparameters: (f) = (.))
Figure 3a and Figure 3b. Optimal learning rate schedule, varying K3 and X2 in learning rate
na((). Numerical set-up: d = 800 and v = 5 x d; batch size is 1, number of iterations is 10 SGD: learning rate is 72 = 0.5/Tr((D) DD where Tr(D)
ンド=1 j 22a j-2a; DANA-constant: 72 same as SGD,
(t()) /.1/ x d where K3, if not specified is (t) = 6/ + t) where d max
(2.8 - 1)/a, - 1/a) + 1; DANA-decaying: 72 same as SGD, (t) = 0.1 +t x I/Tr(D) where K3, if not specified, is 1/(2a), ((t) same as DANA-constant. Curves generated by using the ODEs in (7) with the deterministic equivalent for K; In Fig. 3a, if K3 < 1/(2a) in 73, DANA- decaying will eventually diverge, but until it does it exactly follows the curve for DANA-dccaying with K3- = 1/(2a). Moreover when K3- = 1/(2x) in Fig. 3a, DANA-decaying does not diverge and it appears optimal. When K3 = 1, i.c., momentum is a pure average, the loss curve is similar to SGD and does not accelerate. Similarly, in Fig. 3b, for DANA-constant with 7t(t) > d= R2 = 1 is the only power in d for which the algorithm does not diverge, however the algorithm is slower. For 82 < 1, initially the curves decrease before diverging and, up until divergence, follow the trajectory of DANA- DANA-dccaying with (t(()) x (1+t) 1/(2o), -))(2)))) This justifies the learning rate schedule
(a( x t 12(aa) in DANA-decaying.
Figure 3c. Numerical set-up: 100 randomly generated K DWWTD and the pj 's computed; for empirical density Ngy, 500 bins equal spaced on log scale from 10-8 to 1 and counted the number of xy that fall into each bin weighted by Pj's 's and then averaged over the 500: For deterministic cquivalent μg. solved fixed pointed equation (33) using Newton Method on a grid of x-values.
0 and we are above the high-dimensionts line, 2a > 1. Using Theorem 4.1, you can find the number of iterations needed to reach the optimum, Fo. when 2a > it independent of which phase (Phase Ia, II, III). In particular, we have that
Figure 6. Suppose 11
((t) x max d K24 kst2
171

---

For 2a > and B = 1, quick calculation shows that
min
(2rt02 2-n3
2ar,
20
>
time to reach irreducible loss, t x d where
2-ka
2a,
else.
When x d2a, this is the same amount of time that SGD requires to reach the irreducible loss level. Equality holds in the minimum exactly
(magenta) line, #2 = 2a Ka).
In terms of divergence, we can look at the stability condition (stability) in Theorem 4.1. From this, we have divergence if R2 -2a x Ki This precisely below the (red) line. We plotted the values of for a = 1.1.
Figure Numerical set-up: 100 - 10 for Phase Ia, IlIa, d 100 2', 2 - 1/ with = 5x 5xd; SGD: )/2 same SGD, %33 (t) 0.15/(Tr(D)) x d=! DANA-decaying: "/2 same SGD, (I(() 0.25(1 + 1) 1--1)-14 (t) same as DANA- constant;; colored lines computed using deterministic ODEs (7) over different d and black lines the predicted compute-optimal curves from Table 5,6,7. Our predictions match the scaling laws of the deterministic ODEs and show that DANA-decaying yields better scaling laws. Phase diagram, (bottom, right). Green region: both DANA-constant and DANA- decaying accelerate; red dashed region: DANA-ccontall doesn't accelerate but DANA-decaying does; red region: neither DANA-constant nor DANA-decaying accelerate.
10...., 15 included for Phase Ilb, IIIb; 5 2 + - 1)/a.
learning rate 7/2 = 0.5 D.); DANA-
A() = 6/(1 () where = max (2 + (28 1)/a,2 1/oc 11
Figure 8. (left) Simplified ODEs (43) and exact ODEs (22) match the stochastic algorithms across multiple d-values. Prediction for the loss exponents match empirical (Approach I used from [50]). (middle) Plotted is the loss exponent 7,, Pff/d"; d²) = f against the our predicted 17's given in theory (see Table 6, 5, 7) by fixing B 0.7 and going through different values; as 1 go through Phase > IVa + IVb - + II. Solid dots computed from multiple runs of the stochastic algorithm using Approach 1 from [50] (see Section K for details). Predictions match well estimated values and show that as SGD, DANA-constant, and DANA- DANA-decaying all coincide when in IV. Exponent of DANA-constant and DANA-decay Ic, there a small discrepancy with theory versus estimated due to dimensonality, d, effects (see [80] for details). (right) Same plot for params exponents (top). & on d' j and data exponents (bottom), C. i.c., exponent at compute-opimmum for number of samples; Fixed a = 1.0 and varying B.
significantly larger than SGD III and II. In
172

---

5GD: la, - 0.72.077
50D: Io. 61 -00000000 0.71
500 In. (t.6)- (0.26.07)
5CD (s. (1 -20227.220 02.0.251
Fnrpikal Sired DBC
Thears may 4.22
Thealy mgp 6.00
Theryy Hap
828
1ha0 PX0 0
la.
(a 6t (8.27.055
(a.
Enrpiai
Impatilt
Emprica Hap
Criliial spp
20
37
18
SGD 122--0020 (0.27,871
SCD:: 81- (027 2.1
5GP: (.2..0.9
SGP: B1 (0.27.10)
18
15.m
Sinpl CD2
Err Crrptaated eg 624
EmprRd 105
222.
CYpiiaal KOS -427
Oripccal 4.20
5C22: Lc. - (0-27, 1:33
5422: In. 6) - (0-28 D.71
500b (s. 2) (0.5 0.71
SG2:: La. 01 10.4 0.7)
32
Naak
Emgiiill
Strpt CDC
Thebry tog -.30
187
1140y O9 +-6.2
ISD
Therry 0g -530
Orgakel L22
36
SGD: (15.05.0.0
5GD: F1- (0.6.00.000 0.71
-0151
107 107.0-00111
10.k
Enraaaal
164
Empatitleo 655
Emprik 9.5 ThERy Ei
Empliaa
$.20
Theary taz
438
4.30
Thory op
Thusay KRS
18
18
20
18
120
SCD
SGD: (07.000)
SGD (07.0.0))
SGD: (07.033
187
31"0
Kak
Siredee DnC
Ergantalegg..oo --13
Expredl a 4.3)
Crpitall 502 -033
01p85el -44L
SGD: 10, B1 (0.7. (0.7. 0.4)
SCD: 10. 00.700
5GD: (2, f) 0.7.
SGD: (2, B1 00.70.71
127
354
Mick
Emriical
Siept bed CDC
Orokkca
0.50
Theey mg 6.30
TINY mg)
PR0 0.
Hops
Plops
Payss
Hapa
Figure 25: SGD loss curves and compute-optimal loss vs. theory on PLRE For various (a, B). we plot the mean population risk for stochastic algorithm runs over 10-100 seeds (see individual figures) for SGD (solid lines) with 72 0.375/Tr(D), Colors indicate dimensionality d ranging from 200 to 12.800. The stochastic nuns of SGD match the solutions of the simplified ODEs in (43) (dotted lines) nearly perfectly. The empirical compute-optimal power-law (dashed black line) generated using Chinchilla Approach [50] where the solid red highlighted section shows the region fit. The empirical compute-optimal loss exponents nearly match theory for all (a:, B) tested, with absolute maximum difference of 0.08.
173

---

61
13-
Enpirial
SGD: la. 122)
5tAD: La. (..7
5GDD: (2.7.14)
3GB: 81
Ceplrcal los m-111
Orgatal CO -
Enpiaaico CO
+-034
Explica 122 +-0
SGD: Io. 61 10.0..0.71 0.0.)
Se F1- 116.0.71
3G0 (s. (2.00.
50D: le. 6 (1..0.9)
16
Ceparnol Smmp Exd 002
18
Thooy cq = 8.60
Theap LS2
+-A.1
18
Trepy LE2
+AA22
Theay op # M
51
Eapelou
(1.4.14)
112.
12.
42
Cipifial 0op 3.7
orgiisl RQ 6.7
ImpRaasg STA
Expliic 43
5CD: 10.00 10.00112..000 (12.09)
SED: SED: (0.20-121.11 (12.10)
SGD SGD (1...11 (12.11)
5G5 68 (19.1 (12.12)
12-
R12
Eigitial SPP End 005
Thaory
6.70
Theop 62 2-077
Teey 85 a -.7
Thay 8-8
la. 61
51 (12.141
(24.0.))
22
22
Depariul Theary ag
Expiticl Thesy
Thaorry -.13
TheEy
Floss
Figure 26: SGD loss and compute-optimal loss theory PLRE See Figure 25 for details.
174

---

Enaiiial
014,
SGD (1.4.14)
(16.000 0.7.1
(1.6.
R15
Ceplrcal los - --111
Exprica 622 --AR
Expred 822 0
50D: In. (1 (1.6. 0.91
3G0 le. A1 (2.6.10)
3G0 Ls. (2.6..1
SGD (e. (16.1 121
Creeroo Smmp Exd 002
Thory cq 880
Thpoy w:8.
Tee9y op 4.8.
Theay 6x7 0.0
8t
Eaperou
12
227
(..8.0.9)
1..8.1.0
cillica aop 3.22
Englice AUS -.0
Emp-Ed P03 0
Ens 0g
SGD: lo, (3.
SED (2.0..07)
5:D: 65, (2.0....1
50D 65, 2.0.D.
21096
13
R12
Enditii
Smpp End 002
Thaory 664
T1a0 up 6-2.8L
Tpeay.op 2-4.0
Theyy 083
Le 225.
L8 81
(8, 51
Lo 2.0.0
Deparou Theary ag
Thao'y V ag
Thaary op
Thaar
17
18
Flops
Figure 27: SGD loss and compute-optimal loss theory PLRE See Figure 25 for details.
175

---

DAMA-coaatar:: B1- 10-102..2
DoNd-consiart: le 51-1- (0.25 87)
DAhsh-orssaant (s. A1- (0.25. 07)
DAWA-cordaant d1 (0.27.2.5)
Nak
Emrpekal
d0:
Imppitatio
Emoric Theyy Ex -4200
Grppical
022
Thesry uas
+-0.2
Thear ep -030
Thasey 652
13
30
10-1
301
35
DANA constant:
(0.2.0..4
DANA 02.22.
DANA corstant: 27.
R1k
Sindl CD2
tmpitatee ug 666
Emprid 10p 9.3.
-.1 DVNA condtant (2. B1 -1.20.77..00
Crotiial 602
4.2
DANA constant: lo. B1- (0.770.7)
DANA constant: (a. 51- (0.7.0.8)
DAN constant: (or B1 10.27.100
Nakk
Empiial
Sinpited CDC
30
16-
Thebry sog
0.29
Theay Ag +6..0
Theory ug
Orokkello B22
CANAAconstant:: ln. 10 -(227,
34N44constant:: la f) (0.28, 07)
DeetA-constant 10.3. 0.7)
Daskk-constant: 10.4 0.2)
101
Empataal
CD:
Therry tap +-333
T5E8y cp +-6.0
Thony ep -230
Thusky 6X2
DANA (0.505..77
(0660007
DARIA
DANA (07)
10117
80ak
16-
30-
Irrpiial
14-
Irypakateo em 250
Exprid 953
A31
Eripiiaa 5
DANA www.onntttttt..o Co. (07. 0.05)
DANA.conata cons-constartt RE le (3..000)
DANA constant Le. =1 (37.02) 0.)
DAN0 constant: (o. () - (07.03)
10"
Hak
100
Enriical Sinelieed DBC
414
Empris KC AMp
TINNy 9gp
ANS
6111
350
Thar99 0.45
DMMil-constaot:: (n. (0. 84
cohil-cneetat (2. (0.
Do0A-rmmetnn: B0 =0.7. 0.0)
DeNb--omsaant: to. =.77.
1440
16:
Sinel CDC
Enriiccal
42.0
224
Theery A0p
Thenry 453
Figure 28: DANA-constant loss curves and compute-optimal loss vs. theory on PLRE For various (a, 8). we plot the mean population risk for stochastic algorithm runs over 10-100 seeds (see individual figures)
= 5(1 + .8 m 2a+28-1, 2- } + 1. Colors indicate dimensionality d ranging from 200 to 12,800. The stochastic runs DANA-constail match the solutions of the simplified ODES in (43) (dotted lines) nearly perfectly. The empirical compute-optimal power-law (dashed black line) is generated using Chinchilla Approach [50] where the solid red highlighted section shows the region fit. The empirical compute-optimal loss exponents nearly match theory for all (o, 3) tested, with absolute maximum difference of 0.09.
for DANA-constant (solid lines) with ?1 = 1, = P) 73 0.1 x .E A
0.375
176

---

DMW/ecommdant: lo. 8i (07.0.8)
Delibweonetant: is 8i (87, 6.9)
DeW.ccnddent: lo. (i 11-000111
DMMMccemtent: (o. -1-107 1.11
10
Engirial
Thaor 6664
Theop +-0
Treay * --444
Thaay +-2 54
DMMdeconstant:: 18. 61 (0.7.
DANik-otolann (8. (077
DeWA-corstant (a. (0.7
DWMMM-crsten:: 8)
Orgrical
Explical
227
121
DAMb--onsttan:: 80 (079
DdMGccomshnn: B1 11..1.,,000
Dohd-ranetent (100000
18
CeaFiial 1og --2
gripiaa EO -
Emp DO + 4T%
Exglica 652 + 4T%
DANAecoteeant: (.D. 101
DWW.enndent ty 11.00.0111 (.0.1.11
DWW.enddam: lo -1.1111111
D6hd-cantttt Io 0 (10.11) 1.3
R12
Engiiial
12.
22-
Smpp Ehd 201
Thaory 8.80
Thecry 6-4
EQ
T146 682
a -48
T12y op
4-2 20
DANA-comotant
DANo-comoiant (a (11
orottnt (a -11112.77
-1
Zepericl Theory oo
Thaop
DANA
DANA
ceeinol es A11
orgekal og
Impraalop ER
Explick AD.
D4NA cwnstconstart.. (a E1- (12.1.31
constnt (8, 21.2 2.41
DAN condtant (y Bi (1.40...
DANA sonstat: to E (14.00
15-
101--331 Thavy 630
Sinp Rid 0ps
Oriaicel cg ros
12-
og
Tenvy +.5..
Tieey Pes AX7
Hlps
Plopa
Flags
Faps
Figure 29: DANA-constant loss curves and compute-optinal loss vs. theory on PLRE See Figure 28 for details.
177

---

DAVAvommtnn:: 15 -1.44..666
DWV.coraten:: (o. B-11..110
DMlleomdant: (D. (14.111
DAblecnndant -1. (1-4.22
Thacy 622
+.007
Treay 6C2
-007
Teeyy 0p
8-3
Therry op
--287
(1.44....
DeMeecrattad: (1..4.11.....
DWM6-oratart: (8.
D4hi-comtirt (1.6
124
Ins
1240
1071
1270
Expliial
22
Expital
27
DoNAvmodhen (1.6.0.9)
DMMM-cratant: 11.0.1.0)
DMMe-ronetnt:: (1.6,113
DMMd-ooeetnt: (26
14444
OFS real
1.00
LeiS (531
1EL
Expleed 82 LE.
Expiial 60
1OL
DANAvommeent: Lo -....11... 1.31
D0MMeerataan: 1....1.4 1.41
DAM0ecommkant In. -1.8 0.7)
02 0Ablecmmian 11 3.81
Scppiie beg ODC
Tlacey 682 4.0
Ensrral T1ayy 662
400
Engetial 6-.00
Theory op $-
1,01
Orgikel
300
Smp
622
T022y 09
Thabey 0p
(..8.144
(2.0.077)
1244
1941
Imgrraain 9D
Expetidd c3 1.00
Expiial 101
DANA condtant ta #1 12.0.1.01
DANA: condtant lo (2.0..1.1
DANA: constant: (a. (2.01.12)
DANA constant La B) (20.144
Sir DIC
Ensrral ppe-1
Empetial
Tharey 1.14
Tveey CE:
114
Tevvy map
Flops
Flops
flops.
Plpp
Figure 30: DANA-constant loss curves and compute-optinal loss vs. theory on PLRE See Figure 28 for details.
178

---

DAtAA-dwapping (e. (02
DAN DAW.drcayiigg ic, 8)- 10.25. 0.71
DAN decceiigg ie. 01 10226.077
Sond-decayigg (0 0255
Nak
iipptall
Imrpital
Emprica Thery Ex
Grppiical
022
Thesry unp
+0.22
-4200
Theary exp -030
Thabey 6X2
12
30
12
301 26
25
DANA 0..27..33
DANA decosing: (..27.0.4
DANA 10.27,0.55
DANA decaving 51- (.22..06)
Ri5k
Sinol CD2
tmpitaleg eg 6.6
Emprid 105 9.3.
Crptiial 602
Gripieal ug 4-22
DANA DANA decaying le Ri- (027.0.))
DAN docaying (G. B1 (0.270.8)
D6NN decaying: (e. 8 - 10.27.0.91
DANA decaying (a. 51- (.227..))
10
Empiiaal
Sinpiled CDC
trraitalao
Thebry
-.02
IDD Theay Ag +6..0
Therry ug -3300
decc-dccayng fr Al - (0.27
DAN decA-decaying (c. A)-
DAbA-dacaying f0) 0.7)
dacayidc (s A1 (8.4.
101
Empaaaal
CD:
Therry tag -036
T5EEy op +-220
Thoay ep -230
Thusky 6R2
10.50.5.00
(0.7
DANI
10117
16-
80ak
30-
Irrpiial
36m
Irrpakatton emp R6A
Empriia 2AR
Eripiial 5:20
DANA: decaying lo. B1 -0.7. 0.051
DANU- Sec.ying Bi 10.70000
DAN-A Secoying - 10.
DANA decaying Ca fi (6..0.3)
Hak
Enriical Sinelded DB5
4.4
Empric Aup mg 6.18
AUS
35.
1ha099
Agp 5.25
DeWi-derappng lo, 83 20..0.0
deMayyinn Bl (.77.
DAbl-decay wng (8. 51 10.7. 7.3
DAbi-decay PIg (8 5) (.7
350
Enrpitial Simol CBC
Thebry A0p
2700
Thenry 453
Figure 31: DANA-decaying loss curves and compute-optimal loss vs. theory on PLRE For various (a, 8).
we plot the mean population risk for stochastic algorithm runs over 10-100 seeds (see individual figures) for DANA-decaying (solid lines) with ?1 = 1,
D) 7t(t)
0.375
0.1 ((f) 8(1 + +)
(144)1/133
max{ 20-22-1
} + 1. Colors indicate dimensionality ranging from 200 to 12,800. The stochastic
runs of DANA DANA--eccying match the solutions of the simplified ODEs in (43) (dotted lines) nearly perfectly. The empirical compute-optinn optimal power-law (dashed black line) is generated using Chinchilla Approach [50] where the solid red highlighted section shows the region fit. The empirical compute-optimal loss exponents nearly match theory for all (a. .9) tested, with absolute maximum difference of 0.085.
179

---

D.M.edderaiing to, 10.2
Bf (a 21 107, 0.9)
Dehibdaccyying (2 (0.-0077711 1.00
De0.edecapping to B) 00. 1.8)
Stpekeenne Thece 602
+-4.2
14.9
Treay 652 --772
Theay 0p 9-272
Theery 0p -2222
DwN.odecaaiigg B1 1.2)
D4hsk-iccyping (t (2.7)
d40A-drcayynng (0.
Fial
164
Eng
Expiiial
Expiial
14-
12
12
DAMeedecaytgg (0.9.0.))
b4 1..0,0.1
DAMbdderaying (1.0.00
Se00-irryynng (1..0.9)
11-
OFS real -.570
LeiS (53)
Explted 622 &.
Expltal 60 683
DMi0edecayingg Lo 81- 11.0
Ddhubodeceyingg -- 1.33
DANbudecayingg 12 (1.00 1.1
DWW.ederagigg 11. 1.51
02-1
Scpp bes ODC
Tlacyy 682
16:0
Ensrraa T1ary 622
422
Empeiial EED 2-2
Thaory op -.30
decaving lo -1
11-
Orgakell
1565y 623
17027 02
Thadey op
DANA
meprcalo 2.97
Emprral mg 28)
Exgetid c33 0.5
DANA decoytng: fy Bi 112.1.1
D4NA DANA decaying. (e -12 1.4)
DANA decalig 0
DANI decaying: In. BI 11..0.8)
124
1240
Ensrialon ope-13
Empltial
Tharey ear
0.8
Tvery ex: 637
Teeey map
Flops
Flops
flops
Plps
Figure 32: DANA-decaying loss curves and compute-optimal loss vs. theory on PLRE See Figure 31 for details.
180

---

12M00-deccaying: -0 (14.0.9)
Dehibdeccrgingg 2- (14.18))
DAW.eerapping 15 -l 1.1
www.oeer to -14.11111
RICE
124
Thece 602 1.10
Treyy 6C2
T1eyy 0p 6-1
Therry op 0-244
04toh-decayin:: 1.444
D4ho-dicaping (e (1.66.130
decapinn ld.
144
12.0
real
100
1941
Expliial
37
D490Addeccying:
Dihobdeccyiing: 1.4).14)
DeMeedr (1.6
18*
OFS rell
LeiS F63T 2.1
Explted 822 LE.
Expiial 60
LLL
DM00edecayin B1- 11.6
Delobeleceying. 51 (1.8 0.71
DwW.edecepingg Lo. -1.11-111111 0.81
1444
Stpphie bes ODC
1941
Tlacey E02
Ensrral Trayy 6C2
Engetial
Tharry
decoving:
0.91
D4Nohdersgiing
DANA-dacaying
37
144
12:0
Orgskel Smp
Thasey
13027 02
Thatey 02
DANA (2.000.))
DAN6
2.6
EmS PEREE11000
Expeticd E03 1.11
Expet ca) 124
DANA DANAA dec.yleg. to 8i 12.0.1.0
D4NA 34NA-de:aying le -2.00 (2)
DANA decayiig la (2.0..12
DAN decaying ln Bi 1.4...4)
10100
Shr Tharey eao
Empeiial 1.22
1.2
Every eeo
127
Thevry
Plops
Flops
Plops
Figure 33: DANA-decaying loss curves and compute-optimal loss vs. theory on PLRE See Figure 31 for details.
181

---

SGD:
0.25)
Thatry
Paramveter
SGD: 1o. =1 (6.2).03)
SGD: (2.
SGD: l 1027..0.55
SGD: 09, (0.27.04)
Parmmeer 3
Ostee leae Dali
Ugriel laantoun
3poer Peri
Gilimel Ieae Cou
5GD (8,81
SGD: (00:(e.5)0 F) (022,
5GD: fo (027.0.9)
SGD:15, (6.22
poviles 241
Brppssiien 5.4
Parsmeter
(1.22..1.1)
S60 (0.28....0
(0..0.0.)
SGD: (0, 10.3.0.7)
SGD (2. 53 (2.66,0 6,0.7
SGD 80 (0.1)))13)
SCD: (0,8) 10.7-0.11
Sator/Peae Gevi
Oytoall Crant
Daliivl feam Cool
Perameter
500: (1 50-1001 10.7,
Uatraer (z (0.7.00..
Lor ( (0.7. 0.2)
1200 L8, (0.7, 0.3)
Seud
pooulee E41
5.3
Parsmeter
Permmeteer 0
Ther 572
Treary
Figure 34: SGD Chinchilla Approach 1. Gray stars plot the parameter count of the model size that is optimal for each value of flops using the loss curves from Figure 25. Power laws (red line) through these points give the empirical parameter exponent, which matches theoretical predictions within 0.09.
182

---

Parameter
50D: 10, -0 (0.7112)
500 (0..1.3
SCD: (a. (0.7.144
5GD: (a, B1 0.7)
0steri rear Oniy
Fatre Inan
Scer Peice
101
Permmee
In, (0.9.0.)
(1.8,07)
SGD A 0.0.8
11
eu 2/0
Siae Maleeeeee 1.00
Brain de
Parseeete
11.0.1.0
SGD (1.00.1.
10/
Parsmeter
SC0: -0 (1...44)
SCD 80 (1.1.07)
SCD 8 (12.0.7)
SGD: Øl (1.2008)
Gatrr/ Gevd
QuieeiaarrCo Cown
Suae Cord
Perameter
12
101
10.8)-11 1110111...00
(1.2.1.0)
(g, (1.2.111
SGD (8, (1.222211
Sveriidd ne 200
10"
Suler Pffice
Plecto
Parsmeter
ca E18
Srance
Heero de
Persieteer
19
Flags
Figure 35: SGD Chinchilla Approach 1. Gray stars plot the parameter count of the model size that is optimal for each value of flops using the loss curves from Figure 26. Power laws (red line) through these points give the empirical parameter exponent, which matches theoretical predictions within 0.09.
183

---

301
Buary
Paramveter
SCD: (8, a) (14,
SCD: (a, R -1.4111.4
SGD, (2, a (1.6..))
SCD: - (L.6,08)
Datee lever Daui
Uetine hame
Detreiiteer Coco
Parmeee 15
5GD: Lo, - 0237
SED: 11.61,1111
[e,
SGD: B (1.A111A 1221
Bramnttee EX1
Parsmete
(1.0070)
(1...00.
Vieenii da 1.1
Porsimeter
12
SGD: (a. Fi (1.8.09))
SCD: Bl (1.0..0))
5CD: T0, B1 (1.A.1.))
SGD: (0 (1..122
Satee/teer Gevi
Quiieel Dran
Detinal Connt
Perameter
5GD (8,
(2.5.00.)
502: (1.06-12.0...))
Inerraatie mel
Seud
Braenaleen 2.2
Thei cp
Taaro
Persintter count
Figure 36: SGD Chinchilla Approach 1. Gray stars plot the parameter count of the model size that is optimal for each value of flops using the loss curves from Figure 27. Power laws (red line) through these points give the empirical parameter exponent, which matches theoretical predictions within 0.09.
184

---

(0.20.07
DANA 10.25.077)
DANA (0.26.
10.27.0.9)
Paramveter
14
DANA-consaat (o, B1 (0.27.0 0.3)
DAMM-ondtant: (o, 80 10.27, 9.4)
DAN constant (a, (5.27.05)
DAMM-coeptatt: (o. g 10.27.0.0)
Pormmeer 3
Oatee lear Ouui
Sytee harscum
Bcinn Pese Ovd
Optral kaantuur
W
34hl-ronttnt: (s, (3.22,27
DAMMccomtant:: (0. 8 -0 10.27, 0.8)
Dehhrrmmtnnt: (5 (0.20) 0.3)
DAM-ccmmtat:1110 5)- 10.27
Ieared ke 2.44
aecuees 8.3
Hery por 6i
Parsmene
DANA consant (0.27.1.11
DANA constant. -l..8..1...00
DANA constant: (0.30.3.00
DANA constant 1.74.071
Porsmeter
DANA-constant: lo. = 00.5, (.75
DANN-constant: (s. = 0.7)
DAM--constan lo. g = IDJ, -0.151
DANA: constant la. d 0077 -0.3-
-atarl here devi
Sytoel Stanviunn
Salre Gend
optnal Druil
Perameter
19
DWNlecordent:: fo (0.00..-3.00)
D0000-ometet: (g, 5) [87, 8.0)
D6Ndermeent:: [a, (0.00-10.7.02)
DMNA-comdeet (o. Ri- [0.7 8.31
Parsmeter 1
feoer Soud
Sptrer mariicuat plaiiee
Sakr Seit Periuree
8.23
12 lf
vericcee
273
bf paice as
Hinnoree 2.40
eq- 323
Parmmetter Count
Flags
1lpps
Figure 37: Chinchilla Approach 1. Gray stars plot the parameter count of the model size
that optimal for each value of flops using the loss curves from Figure 28. Power laws (red line) fit through these points give the empirical parameter exponent, which matches theoretical predictions within 0.10..
185

---

DANAA (07.0.8))
DANA cp
DANA
DANA
Parameter
DMMll-contannt lo, a (0..1.2)
DANA-constont: f (.7,
DAN conotant (2, B1 (0.7.14
DMMM-cnntan:: BI (08,0.7)
Oster lear Oniy
acer Phee Cony
Perimetr
DWNll-comttent le, (0.9 0.7)
Dohiwccoeteentco (o, Bi (1.0. 0.7)
DMM-cromtant: Bi 0.8)
B6bi-conntan (e. (1.60 (1.,0.3)
Inarred eu 2/0
Brain
de C.49 Heroee
Parsmetee
10
DANA constant
DANA (1.00...
DANA (1.0 (1.0..1.2)
DANA (1.011113
Porsmeter
DANA-constont.co f (3.0.144
DANi--conctont f0 (L..0.
DANA condtant (a, dl (12.12
MOMM-constant: =1 (12,0.8
Gaterl Paier Gevi
Spimel Saaviuun
Suaen Cold
Perameter
104
D0NAlconstarf:cco (e. (12.0.9)
DANAecteeentt (g, (1.2.1.0)
DMM- L0, (1.2111..1.
DAMMcnreeann: (o. (122,12
beeiee Seud
Sole Parce Sard Plero
Parsmeter
vepiccie E27
Serance -RA7 ne
Teero me 232
Persieeter CoUnT
Figure 38: HANNA-oosstant Chinchilla Approach 1. Gray stars plot the parameter count of the model size
that optimal for each value of flops using the loss curves from Figure 29. Power laws (red line) fit through these points give empirical parameter exponent, which matches theoretical predictions within 0.103.
186

---

DANA
DANA
op
Parameter g
DANA-constont: 8 (1.4.13)
MoM--onnttnt G8, B1 (1.4.1.4)
DAN constant (2, B1 (1.6..7)
DAMA-constantt consta nt: (e.R)= (14,0.8)
Oster lear Our
Scer Pars Coll
Parmmee
g
DAMAecmeeeent: 5) (2.5 0.9)
DoMoemneeent: (2.6.10)
DAMeccomsamt: -1 3.33 Braeue
DMM (0. (1.5,12)
Herrp HD -032
Parsmeter
15"
101
DANA constant (3.6.13)
DANA (1.66...
DANA (1.8
DANA (1..8...
Year Meriiaiee 1.4
Miuridl dng 2314
Porsiettrr
DANA-constart 50 (1.8.09))
DANA coNNtonnetttttco fl (1.8.1.0)
DAMM-contant: Bl (1.,1.3)
DAMM- wwW--cootant:: to (1.8 -uee Coui
Saborl Peae Gevi
Sammel Serni
Perameter
20N0l-ccnetent: (o. 51- (1.8.1.4)
DoNalectteeenttco (o, 81- 0.7)
D6hib-oontann: (2.0)- (2.0,,0)
Raae
Cameer daa 2.1
Sole Prie Card Plero
Seeieed 200
es
Parsmeter
3
T
Liga 40-1- ap
prance ne
Herr de
Parsietter Count
Figure 39: DANN-constant Chinchilla Approach 1. Gray stars plot the parameter count of the model size
that optimal for each value of flops using the loss curves from Figure 30. Power laws (red line) fit through these points give the empirical parameter exponent, which matches theoretical predictions within 0.103.
187

---

1.0200077
00.26.077
DANA 0.5)
Paramveter
DoNA-decaying: [2, a (3.27.0.))
DAMA-dectylng (o (027.0.44
DANA decoyfng: 10, RI (0.27.0.5) 0.5)
DANA-decaying (a. = (027.0.0
Ostee lear Caly
splrel cann
oytnel leae
Parmmetr
B4hhheeesyyiggcco (o. 8 (8.27,077)
DAMMedeceying: lo, fi 1027.0.8)
deNbb-ecayyng: 15, (6.27,
DANAdersying (c.5) (827,1.0)
Ieeriat 2.44
me 2211 500
Parsmetee
DANA decaying. (1.27.2.11))
DANA
DANIA (0.3.077)
DANA (0.0077
Virrie 2MM
aadd as
Vianra 440 124
Porsmeter
DANA decayeea a (0.5.077
DANA-deckylng le & (0.6,07) 0.711
DAMAAdecayyg (a. B) (0.7-0..15)
DAN decaytng: lo BI (077 -0.11
Gaterieeae devi
Cannel Cuan
calnel taln Dooal
Perameter
DoNiedecastrg: (g.81 10.7. 22200.0000
DANAdde sy.ng Ie. (07.0.0)
D&Woodecaying Lg, Bl (0.2.02)
DANArdecaying (a, (3...00)
Ouber Pate Soud
Salner 2.4
Parsmeter
(s.
Snanne ap
Persmeterr Count
Flag
Figure 40: DANA-decaying Chinchilla Approach 1. Gray stars plot the parameter count of the model size
that optimal for each value of flops using the loss curves from Figure 31. Power laws (red line) fit through these points give the empirical parameter exponent, which matches theoretical predictions within 132.
188

---

DANA
DANIA
Parameter
DAN ng: 68, 4 (07, 1.2)
DANA decoyyng: lo, #1 (0.7.1.3)
DoNA-deccying (o (07.1.4)
DANA-decaying (e. ( 0.71
Fatrel Inanicue
Permmeer
W
D6No-leanynng 15, B1 (0.9.6.)
De 18 - (1.0
DAMledecaying (e (1.6
DeModdcayyng 11.081111000000 (1.0.0.9)
Be Brannd
Heros -337
Parsmetee
DANA (1.0..00
DANA
(1..0.1..
DANA decaytng: lo #1 (1.0.1.1
DANU decaytng loi a1 (L.)
DAWA-decayinn: to, Bi 112,071
DANA-decaying 8) (1.20.08
twae Gevd
Suee
Perameter
101
DANo-decaying Lo, (1.2001.)))
DAMM by.ig Lo. (1.2.18)
DANiddecaying (0 (12.111..111 (12.1.11
DANAddeayynng (2.8)- (1.2.1.7)
Soid
Su00 Pirce Plert
Seered ne
17
Parsmete
pranc
Herr de
Persietter Count
Figure 41: DANA-decaying Chinchilla Approach 1. Gray stars plot the parameter count of the model size
that optimal for each value of flops using the loss curves from Figure 32. Power laws (red line) fit through these points give empirical parameter exponent, which matches theoretical predictions within 0.132.
189

---

DANIA
Paramvtter
DANA decaytng: Io, B1 (1,4,1.3
DoMA-dccyying lo, (114.44
DONA-decaying 88 (16,0.7
DANN-decaying B) 0.81
Oster levrr Deui
Pormmee
DeNo-decmytng: Ie, (16.6.9)
DMMMtderaying (e. (166611)
DANAdeeaynng (o, E (15,1.3)
DAWAderaying (5, 8)- (1.8
CC.
Parsmetee
DANAA
DANA
DANA
Portimtter
DAN- decaying (6, $1 (LB 5.91
DeMMAdecaylnn lo, B1 (1.8.1.)
DANA-decaying le, f (1..
D/NA-decayynng () (1.8.22)
Sameel Saan vouui
Perameter
DANo-decaying Lo, (1-0)
D4N0o-erayye Jo, (2.0.0.7))
DANivdecaying
Calner (0.00-(220 0.3)
DMMoodeaayhg (g, (1........)
Racer Seid
Parsmeter
prance
ne
Herr ea
Persietter Count
Figure 42: DANA-decayion Chinchilla Approach 1. Gray stars plot the parameter count of the model size
that optimal for each value of flops using the loss curves from Figure 33. Power laws (red line) fit through these points give empirical parameter exponent, which matches theoretical predictions within 0.132.
190

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The statements in the abstract and introduction accurately reflect the paper's contributions. The introductionnaasstract make clear which of the algorithms we have proven (e.g., only provide a heuristic for DANA-decaying). We also make it explicit in the statements of the theorems, propositions, etc, the assumptions that we are making. We also verify these assumptions in figures.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: At the end of the main paper, we provide a 'Limitations' section. our statements of propositions, theorems, etc, make very clear assumptions. For example, we are explicit that we do create the deterministic system of ODEs directly from the stochastic algorithms but rather SDE (see Section 3). Another example is in Section H, we are clear that we can not solve the system of ODEs in (22) directly instead we solve simplified system of ODEs and show empirically that the solution this simplified system of ODEs matches .oNNtacnntnnt.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they seale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
191

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: All theorems, lemmas, etc are proven in the Supplementary Materials. We provide a short outline of the arguments in the main document. In some cases, c.g., DANA- decaying, only prove for 2a > 1, but make this explicitly clear in the main document as well as Scction I. All assumptions are clearly stated in the statements the theorems. For example, we make clear that of main theorems only holds for 1/4 and a + 1 B. We do provide numerical experiments which show that our results hold beyond this setting.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: In each of the figures, we provide an explicit description of how the image was generated including the numerical set-up. In some cases, due to space limitations, the experimental set-ups are discussed in detail Section M. We provide a description of the power law random features model (PLRF) in Section 2 and experimental details in Section K which allows for reproducibility description of the set-up for the LSTM experiments in Section L;; we include a citation to the datasets that we are using as well.
our model on the synthetic data. We also provide explicit
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (c.g., in the case
192

---

of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
The paper does not include experiments that require significant code. The model we analyze a simple power law random features model (PLRF) applied to synthetic data. As such, the code can be readily produced by following the set-up seen in the captions and/or our numerical simulations section, Section K. The model has been used before in other papers. Additionally for the LSTM experiments, we are using a set-up similar to other papers [54] and reproducing Fig. 2 with DANA-decaying and the C4 text dataset [84].
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (https://nips.. cc/ httlcc/ggg/s////ee htte///giideee///ouuiiiieeeeee
conPoliccy oolicyyyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. cctppccpppp.ccccccccccccccccccct wwwpc//ggg///eeeeeeeeeeeee /ttggdd/d////duunmiiiiieeeeeeeee (CodeSubmissii ssionpolicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question:
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
193

---

The experimental design, including the power law exponents, hyperparameters, fixed stepsizes, choices of d and v, and the numerical simulations for solving the ODEs are all written in the captions of the figures and/or Section K, Section I. or Section M. We also intend to release the code for numerically computing the ODEs.
Justification:
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We record how we generate the empirical compute-optimal exponents using statistical tools that were first deployed in other papers such as [50]. We report the number of random seeds used for each PLRF experiment in Section K. We report R2 values for all LSTM loss exponents. We are careful to explain when and why the theory deviates from the numerical simulations. Often this is due to finite d and v effects and the slow behavior of the theory to the asymptotics.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: As this paper is about compute-optimal curves, we provide details on the exact number of flops required perform the experiments in Section K and Section L. The compute resources are also well known in the community using the standard 6N heuristic (50, 54] for flops.
Guidelines:
The answer NA means that the paper does not include experiments.
194

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /ttps://///u....../// wwweeris..cppuuaeeeeeeeee httee/eiiiiiiieeeeeeeeeeeeeeeeee
Answer: [Yes]
Justification: We have reviewed the code of ethics. We have made our utmost attempt to adhere the guidelines provided by NeurIPS. We do not use any human subjects nor any datasets. We did our best to cite all the relevent related work. Given that our work is in the foundational research, it is difficult mitigate all the risks as the downstream effects of theory are long, but we have done our best. The model is completely synthetic using the standard stochastic momentum- anticipate any risks. We have included "Broader Impact" statement at the beginning of the "Supplemental Materials."
algorithms; thus we don't, t, the best of our knowledge,
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The work presented is purely foundational research and is not directly tied to any particular application. We study simple random features model with power law data and target and we solve the model using a common algorithm SGD. Given the theoretical nature of this work, we do not anticipate any direct ethical and societal issues. See our Broader Impact Statement in the appendix.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
195

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks. The majority of the work is theoretical and focuses the power-law random features using only synthetic data generated from normal distribution. This model is a standard statistical model (e.g.. least squares) which is textbook learning problem. The LSTM experiments serve only illustrate the theoretical results act as proxy in more realistic setting, and are trained on a public language dataset (C4,[84)). We do not release any data pretrained models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We acknowledge via citations that the model we study was introduced before by others (e.g.. Maloney, Roberts, and Sully paper). We cite the datasets we are using, e.g., C4 language dataset [84] and software use [18]. We include an extensive relate work section in Section where we provide tables comparing sample complexity across multiple algorithms.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi www.whhthcdd..co com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
196

---

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: We do not intend to release any new assets. The work is theoretical analysis random features model along with limited language model experiments that illustrate the theoretical behavior. We will release some code for solving the ODEs.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects. The work is purely theoretical on a simple model.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The work does not involve crowdsourcing nor research with human subjects. The data used in this work is generated synthetically.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
197

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification:
The core development of the research in this paper does not involve LLM as any important, original, or non-standard components. LLMs were used to clean up some of the grammar aspects well as make our code more efficient, but the ideas and proofs (core of this work) did not involve any LLMss
Guidelines:
The answer NA means that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: http://eeuppp.eeeeeee cc//onnferences/2005// for what should or should not be described.
198


Reviews:
[{"review_id": "3aTEc0WeZJ", "reviewer": "Reviewer_fBHr", "summary": "The paper introduces DANA, a momentum algorithm that adapts its momentum and learning rate hyperparameters based on model size and data complexity. The core claim is that DANA achieves a superior scaling law exponent compared to SGD and classical SGD-M in several regimes, particularly when $2\\alpha > 1$. The authors leverage a continuous-time analysis of the risk dynamics, combined with RMT, to theoretically prove these advantages. These claims are supported by extensive experiments on the synthetic PLRF model and small LSTM language models.", "questions": "1. Although I have not thoroughly verified the correspondence between the ODE and the discrete-time algorithm, I noticed that the SDE model (Equation (7)) is not a straightforward continuous-time counterpart of DANA, i.e., DANA does not appear to be a direct discretization of the SDE. How should I interpret their relationship? What are the errors between the continuous-time models and DANA?\n2. Given that DANA shares some similarities with the algorithms listed in Table 3, I’m curious about its empirical performance on real-world tasks, particularly when hyperparameter tuning is permitted.", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "My evaluation of this paper is based on limited expertise in the area. Some parts of the paper may be beyond my full understanding, and my review reflects this limitation.\n\nStrengths:\n1.  The paper is well-written and clearly structured.\n2.  The primary contribution, a theoretical proof that DANA can improve scaling law exponents is novel. The theoretical framework appears solid, and the claims are supported by extensive empirical validation.\n\nWeaknesses:\n1.  The theoretical analysis is limited to the PLRF model, which has a quadratic objective. While the LSTM experiments suggest it is a useful proxy, the extent to which these results generalize to the non-convex landscapes of more complex architectures remains unknown.\n2.  The practical guidance is limited by the fact that the optimal DANA schedule requires knowing the exponents $\\alpha$ and $\\beta$, which are unknown for real-world tasks. The proposed heuristic of sweeping a new hyperparameter, $\\kappa_3$, while practical, makes it less clear how DANA's tuning process compares in difficulty and outcome to other advanced optimizers that also require tuning.\n3.  The theoretical proof of outscaling is established for the $2\\alpha > 1$ regime. There is no direct evidence as whether this applies across modern machine learning applications.", "was_revised": true, "final_justification": "I believe the overall quality of the paper is appropriate for NeurIPS. The authors’ rebuttal addressed my questions regarding hyperparameter tuning and the outscaling regime of $2 \\alpha > 0$. Therefore, I am inclined to increase my score.", "rebuttal": "We thank the reviewer for their time and valuable feedback. We are heartened to hear that the reviewer finds our theoretical contribution “novel” and “solid”, and that they appreciated the “extensive empirical validation”. We thank the reviewer for considering our manuscript to be “well-written” and “clearly structured”. We now address below the key clarification points raised by the reviewer.\n\n**Limitation of the theoretical analysis to the PLRF model**\n\nWe acknowledge the reviewer’s concern regarding the limitation of our analysis to the PLRF model which uses a quadratic objective. We agree that this is a clear limitation, and we put it in the “Conclusion, limitations and future work” Section lines 259-260. Going beyond this simplified model constitutes an interesting and non-trivial research question.\n\nWe believe that although the PLRF model might not reflect the entire training trajectory of deep neural network’s training, different regions in the $\\alpha, \\beta$ plane may reflect parts of the training dynamics where for example, optimization may shift from a population dominated region (phase 1) to a more noise-dominated behavior later in training (phase 3). Hence, we hope that the PLRF may help us understand various training behaviors. \n\nFinally, we note that recent works are pointing towards very good agreement of behaviors between quadratic/convex optimization and the dynamics of complex large neural networks, see e.g. [a, b]. This encourages the study of simplified quadratic objectives as a proxy for more complex training settings. \n\n**Dependency of DANA schedules on $\\alpha, \\beta$ and applicability to real-world scenarios**\n\nWe thank the reviewer for giving us the opportunity to strengthen the discussion regarding the dependency of our schedules to the $\\alpha, \\beta$ hyperparameters. We agree with the reviewer that hyperparameter tuning is a major difficulty in the training of large models and answer below in three main points:\n\n* DANA-decaying schedule **depends only on the $\\alpha$** parameter, not on $\\beta$. This characteristic is both surprising, and promising since $\\beta$ is less tractable in real world problems. $\\alpha$ on the other hand may still be measurable as the power-law decay exponent of the data covariance spectrum. Even better, DANA-constant schedule **depends neither on $\\alpha$ nor on $\\beta$,** but only on $d$. Note that we believe this may also have drawbacks in the training of multi-layer networks where $d$ becomes ambiguous. Characterizing the precise meaning of $d$ and $\\alpha$ in the training of deep networks is a very interesting direction for future research.\n* We suspect that $\\alpha$ may depend more on the data than the precise model architecture. Indeed in both LSTMs and GPT2 preliminary experiments (which we cannot share due to NeurIPS policy), the optimal $\\kappa_3=1/(2\\alpha)$ value is around 0.6-0.7. This bears promise for **transferability of hyperparameters**. Additionally, from our PLRF theoretical results and underscored by LSTMs results in Figure 2c, $\\alpha$ does not depend on model size. We only need to measure it once and transfer to higher scales. This again needs more empirical evidence and is a very interesting direction for future research.\n* Finally even if characterizing the optimal $\\alpha$ allows to have maximal acceleration on the edge of stability, note that using $\\kappa_3=1/(2\\alpha)$ with any $\\alpha>0.5$ above the high-dimensional line will outscale SGD provided the algorithm is stable (this result for which we provide heuristics in Section I.5 was not proved but we believe it to be true using very similar proof techniques). Hence although HP tuning allows for optimal performance, DANA-decaying should bear huge advantages at scale even without optimal tuning.\n\nFinally, we agree with the reviewer that making DANA a competitive optimizer on real world problems is a highly non-trivial but very interesting direction for future research. While the main goal of our work was to provably show that outscaling is possible using a different algorithm, we suspect that to make it a competitive algorithm, DANA should be mixed with adaptive updates as in RMSProp or ADAM. For example, Schedule-Free [32] with AdamW does well on many tasks, but not Schedule-Free with SGD which is what we compared with. Preliminary experiments on GPT2 models (which we cannot share due to NeurIPS policy) show similar acceleration of DANA+RMSProp against vanilla RMSProp. \n\n**Outscaling is established for the $2\\alpha>1$ regime**\n\nWe acknowledge the reviewer’s healthy skepticism regarding the applicability of our outscaling results below the high-dimensional line. We mentioned it in our limitation section lines 260-261 (“It is an open question if outscaling is possible for either the 2α<1 case“) and lines 266-267 (“Defining and measuring this $\\alpha$ on real-world problems, particularly determining whether $2\\alpha>1$, is an important direction for future work“ but are happy to provide a more detailed discussion on it as it is a natural and interesting future research direction.\n\nFirst, we kindly remark to the reviewer that our analysis also shows outscaling of DANA-constant for $2\\alpha>1$ for large batch $B\\asymp d$. We didn’t explicitly mention this result in the main paper because of lack of space and that we believe a general theory for large batches requires more work but our analysis for DANA-constant also applies in the scaling regime $B\\asymp d$ (see parametrization H.1 and remark H.1). When batch $B=1$, the reviewer is indeed right that DANA-constant only outscales SGD for $\\alpha>1/2$ and reduces to SGD for $\\alpha<1/2$. However, when $B\\asymp d$, the stability condition (Lemma H.5) imposes $\\gamma_3/(\\gamma_2B)\\lesssim d \\iff \\gamma_3 \\lesssim \\gamma_2$. Hence, taking $\\gamma_3 \\asymp \\gamma_2$ and applying Theorem H.3, it becomes clear that **DANA-constant outscales SGD below the high-dimensional line** ($¼ < \\alpha < ½$). We thank the reviewer for pointing out this ambiguity to us and will update the manuscript to add a remark about this. Finally, we believe that a general theoretical framework to determine the effect of large batches on all algorithms in (Gen-Mom-SGD) is a very interesting direction for future research.\n\nExperiments on LSTMs and GPT2 show a similar optimal $\\kappa_3= 1/(2\\alpha) \\approx 0.7$ which is situated above the high-dimensional line $2\\alpha>1$ and corresponds to DANA-decaying outscaling SGD. Although this requires further investigation, we believe that this exponent does not depend on the architecture but mostly on the data modality (here language). Characterizing $\\alpha$ on real-world problems beyond LSTMs and GPT2, and whether acceleration is possible constitutes a very interesting research direction.\n\n**Correspondence between the ODE and the discrete-time algorithm**\n\nInterpret relationship between SDE and ODE: You should interpret the two independent Brownian motions in the SDE as using two independent gradient estimates, one for the momentum buffer $Y_t$ and one for the direct state update of $\\Theta_t$. The second change is a simplification of the diffusion coefficient (i.e. the covariance matrix of the gradient update), where we effectively drop the dependence of $\\langle X, \\Theta_t - \\Theta_*\\rangle$ and the data covariance $\\mathbb{E}(X \\otimes X)$ – this is a correct approximation in the ‘high-dimensional case’ [references 75-77]; see also Remark C.3 where we discuss this using the ODEs.  We’re happy to expand on how this approximation relates to Gen-Mom-SGD in the paper.\n\nAs for the ODEs, the ‘simplified ODEs’ exactly describe the mean risk behavior of the SDEs.  We also have a family of ‘exact ODEs’ which give the exact risk behavior of DANA when you ‘continuize it’, meaning you compute the exact risk of DANA but at a random time (see eq (17) and (18)).  See also [13]; and just to be clear, for the ‘exact ODEs’ we are using exactly the update rule of Gen-Mom-SGD – no simplifications. We judged the ‘exact ODEs’ too complicated to analyze directly, and so we created the ‘simplified ODEs’ by removing higher-order terms. \n\n**Empirical performance comparison with Algorithms in Table 3 on real datasets**\n\nIt’s a great question.  We suspect that for DANA to be competitive on real problems, one needs an adaptive version (e.g., Schedule-Free with AdamW does well on many tasks, but not Schedule-Free with SGD). We have hope that such an adaptive result may be possible especially given the performance boost over SGD that we see in the LSTM experiments. An adaptive version of DANA is a future research direction we are pursuing, and we think that doing a rigorous comparison of adaptive-DANA on real-world problems is worth it.  (We have been running adaptive-DANA on GPT experiments, which unfortunately due to current NeurIPS rules, we cannot share).\n\nMoreover, almost all of the algorithms in Table 3 require knowledge of ‘model size’ and the momentum parameters have to scale with this quantity (otherwise you do not get any acceleration). Because of this dependency on model size, it might be harder to transfer knowledge of a small model’s hyperparameter to a larger model since ‘model size’ is not clearly defined for (non-PLRF models) – model size is probably not the full number of parameters on more complex architectures. We do think there’s some chance that the $\\kappa_3=1/(2\\alpha)$ in DANA-decaying remains stable across scales (see Figure 2).\n\n[a] Schaipp, Fabian, et al. \"The surprising agreement between convex optimization theory and learning-rate scheduling for large model training.\" arXiv preprint arXiv:2501.18965 (2025).\n\n[b] Qiu, Shikai, et al. \"Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks.\" arXiv preprint arXiv:2507.02119 (2025).\n\n[c] Kaiyue Wen, Fantastic Optimizers and Where to Find Them (2025)"}, {"review_id": "FTqjpzp5Sj", "reviewer": "Reviewer_P7JW", "summary": "This paper provides a comprehensive study of the scaling laws of SGD with momentum in linear regression.", "questions": "1. How does the variance error occur in the problem setting? From my understanding, the variance error should stem from the label noise, but there is no label noise in this setting.\n2. How does DANA-decay compare with the hyperparameter selection in Li et al. (2023) (Ref. [59])?", "limitations": "This paper does not have any outstanding limitations or negative societal impacts.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n1. This paper extends the scaling law in linear regression optimized with SGD to SGD with momentum acceleration.\n2. The findings that SGD with heavy-ball momentum does not outscale vanilla SGD but Nesterov momentum with dimension adaptation can are impressive.\n3. The authors perform extensive experiments.\n4. The authors discuss the comparison of DANA against the most important variants of SGD-M extensively in the appendices.\n\n## Weaknesses\n\n1. The technical terms \"population bias\", \"model capacity\", \"embedding bias\" and \"variance\" are not clearly defined, although they already appear in the literature.\n2. It is understandable for such an extensive study, but a little unpleasant to see that the paper has only 4 sections and no subsections at all. Section 3 is also very long, which makes the understanding of the main idea a bit more difficult.", "was_revised": true, "final_justification": "This is a paper with strong techniques, solid findings, and comprehensive discussions. After reading the authors' reply to all reviewers, I find no more confusions. The only drawback of the paper seems to be the dilemma of extensive length and limited space, but could be significantly improved given the additional page if the paper is accepted. I would strongly recommend the acceptance of this paper.", "rebuttal": "We would like to thank the reviewer for their time, feedback, and positive evaluation of our work which helped us strengthen the updated manuscript. We are happy to hear that the reviewer feels that our results are “impressive” and that our empirical evaluation was compelling. We now address the key clarification points raised in the review:\n\n**Clarification of technical terms**\n\nWe acknowledge the reviewer’s concern regarding the definition of some technical terms which would help clarify our results. We mainly ran out of space.  We have added the following definitions in the paper:\n* *\"Population bias\"* ($\\mathcal{F}_{pp}$): This loss term corresponds to the loss dynamics when running full-batch gradient descent on the problem (hence following the population gradient), without the embedding matrix $W$. \n* *\"Model capacity\"* ($\\mathcal{F}_0$): This loss term (which is only $d$-dependent) represents the limit of the loss, as the number of iterations reaches infinity. It arises from the partial expressivity of our model class, since the learned parameters $\\theta \\in \\mathbb{R}^d$ cannot encode the whole target vector $b\\in \\mathbb{R}^v$ when $v > d$\n* *\"Embedding bias\"* ($\\mathcal{F}_{ac}$): This loss term comes from the random embedding matrix $W$ which deforms the spectrum of the data covariance matrix and misaligns it with the target vector $b$. \n* *\"Variance\"* ($\\mathcal{K}_{pp}$): This loss term comes from the stochasticity of the algorithm which at each step samples a new i.i.d. random datapoint $x_t$ to compute a stochastic gradient. This stochastic gradient, whose average recovers the population gradient, is a non-exact estimate of the gradient and therefore introduces this additional loss term.\n\n**Sections unclear**\n\nWe acknowledge the reviewer’s legitimate concern about the division of our sections, and more precisely Section 3. We ran out of space but will break up the sections in more precise and delimited subsections in the updated manuscript to improve clarity of the paper. In particular, we updated the manuscript to end Section 3 after the SDE part, and created a new Section 4 titled “Deriving the scaling laws of (Gen-Mom-SGD) on the PLRF model” to present our main results and Theorem 3.1.\n\n**Origin of the variance term**\n\nThe reviewer raises a very interesting point regarding the origin of the variance term in the scaling laws eq.9. The reviewer is indeed correct that the variance term $\\mathcal{K}_{pp}$ does not arise from label noise in our setting. The variance here comes from the stochasticity of the gradients which samples a new random point $x_t$ from the data distribution at each iteration to compute a stochastic estimate of the gradient. The gradient is a stochastic sample whose average equals the population gradient, but whose variance gives rise to an additional loss term in the scaling laws.\n\nOn a different note, we kindly remark that due to model capacity limitations ($\\theta \\in \\mathbb{R}^d$ while $b\\in \\mathbb{R}^v$ with $v>d$) the target label $\\langle b, x\\rangle$ cannot be recovered from the model $\\langle W^Tx,\\theta\\rangle$. This creates some intrinsic label noise which in turn gives rise to the $\\mathcal{F}_0$ term in the loss. This term can equivalently be seen as the model capacity limitation.\n\nFinally, we believe that introducing some additional additive noise in this model would not change the study substantially. It would mainly add a term in $\\mathcal{F}_0$ and consequently would shift the loss curves. We thank the reviewer for raising this valuable question and added a remark about it in the updated manuscript.\n\n**Comparison with hyperparameter selection in Li et al. (2023)**\n\nWe thank the reviewer for this interesting question. We reported extensive hyperparameters comparison with other algorithms including Li et al. (2023) in Table 3. We note that neither DANA-constant, nor DANA-decaying can be straightforwardly compared to Li et al. (2023). Indeed, as derived in the first column of Table 3, the corresponding $\\Delta(t)$ in Li et al. (2023) is constant but $d$-dependent (as are other hyperparameters).  As we took $\\Delta(t)$ non-$d$-dependent, we believe that the algorithm in Li et al. (2023) does not straightforwardly enter in the category of SGD-M or DANA. The case where $\\Delta$ is d-dependent and learning rates d-dependent was also studied in [*] under DAHB. It was shown (not under power-law) that DAHB did not accelerate, but did improve the constant of SGD. We suspect that a similar (non-acceleration) result holds for Li et al. parameters under the power-law covariance setting.  \n\n[*] C. Paquette & E. Paquette, Dynamics of Stochastic Momentum Methods on Large-scale, Quadratic Models, NeurIPS 2021\n\n**Presentation improvement page 5**\n\nWe thank the reviewer for pointing this presentation issue to our attention. We originally ran out of space and reported the description of the phase space on its right. As suggested by the reviewer, we will improve the updated manuscript with clear captions to strengthen the clarity of the flow.\n\nWe are grateful to the reviewer for their great questions and hope that our answers clarified them. We believe we have addressed to the best of our ability the weaknesses and questions raised by the reviewers and that they strengthened the manuscript. We politely encourage the reviewer to ask any further questions they may have and, if they are presently satisfied with our responses, to consider a fresher evaluation of our paper."}, {"review_id": "hlEXZ2WZ2N", "reviewer": "Reviewer_3zYo", "summary": "This paper offers a theoretical analysis of momentum-based stochastic optimization in high-dimensional settings, particularly when only a single pass over the data is permitted. The authors study random feature models with power-law spectra (PLRF), employing tools from random matrix theory, continuous-time dynamics, and Volterra integral equations to characterize generalization error. Their main contribution is a dimension-adaptive momentum schedule (DANA), given by \\$\\mu(d) = 1 - c/d^{1/3}\\$, which yields provable improvements over standard SGD.", "questions": "1. The paper compares DANA to Adam, but why is the momentum coefficient scalar (\\$\\gamma\\_3\\$)? Would a coordinate-wise or adaptive scheme be more effective in anisotropic, high-dimensional settings?\n\n2. The authors model gradient noise via deterministic quantities (e.g., $K\\_{pp}(t)$, Appendix C, Eq. 7) and reduce the analysis to ODEs. How might an SDE-based approach, in the spirit of \\[a], \\[b], or \\[c], alter the results or insights?\n\n3. LSTM experiments in Section L are interesting, but serve only as indirect evidence. Could DANA be tested on architectures like CNNs or Transformers to assess its relevance in more realistic settings?\n\n4. The paper is too long. Consider trimming the main text and moving technical appendices to a companion or journal version. The current length exceeds what is practical for peer review.\n\n5. The Volterra formulation is elegant but dense. Would it be possible to include plots of bias terms (\\$F\\_{pp}\\$, \\$F\\_{ac}\\$, \\$F\\_0\\$) and variance (\\$K\\_{pp}\\$) over time, for varying \\$\\alpha\\$, \\$\\beta\\$, and \\$d\\$? Visualizing these could aid interpretation.\n\n**References**\n\n[a] Mandt, Stephan, Matthew D. Hoffman, and David M. Blei. \"Stochastic gradient descent as approximate bayesian inference.\" Journal of Machine Learning Research 18.134 (2017): 1-35.\n\n[b] Li, Qianxiao, and Cheng Tai. \"Stochastic modified equations and adaptive stochastic gradient algorithms.\" International Conference on Machine Learning. PMLR, 2017.\n\n[c] Yang, Ning, Chao Tang, and Yuhai Tu. \"Stochastic gradient descent introduces an effective landscape-dependent regularization favoring flat solutions.\" Physical Review Letters 130.23 (2023): 237101.", "limitations": "The paper does not fully acknowledge its limitations. In particular, the assumptions of fixed features, a global momentum parameter, and deterministic dynamics constrain the analysis. The authors are encouraged to expand the limitations section and clarify to what extent DANA generalizes to modern, feature-learning models.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n\n1. **Originality**: The proposed dimension-dependent momentum schedule is novel and well-justified, especially when contrasted with heuristic or fixed hyperparameter choices commonly used in SGD.\n\n2. **Theoretical Depth**: The analysis draws on a sophisticated toolkit (Volterra integral equations, high-dimensional asymptotics, and random matrix theory) to derive interpretable bounds on generalization error.\n\n3. **Clarity of Exposition**: Despite the technical complexity, the paper maintains a coherent high-level narrative. Key ideas are well-motivated, and derivations are presented in a readable form.\n\n4. **Empirical Support**: Though primarily theoretical, the work includes experiments on PLRF models and linear settings that reinforce its claims.\n\n**Weaknesses:**\n\n1. **Excessive Length**: The paper reads more like a dissertation than a conference submission. Its length far exceeds the typical scope of a NeurIPS paper, and the sheer volume of technical material renders comprehensive verification by reviewers infeasible. While key lemmas appear sound, I did not attempt to check every derivation.\n\n2. **Restrictive Momentum Design**: The analysis assumes a scalar momentum coefficient (\\$\\gamma\\_3\\$) applied uniformly across all dimensions. This simplification may limit DANA’s relevance in anisotropic, high-dimensional settings. The paper compares DANA to Adam, but does not explore vector-valued or coordinate-adaptive alternatives. A discussion of their feasibility could broaden DANA’s appeal.\n\n3. **Limited Noise Modeling**: The authors model gradient noise through deterministic equivalents (e.g., \\$\\mathcal{K}\\_{pp}(t)\\$) and rely on ODE approximations. This sidesteps the stochastic dynamics that drive implicit regularization. A more direct SDE-based approach, as explored in \\[a] and \\[b], could better capture generalization phenomena like the preference for flatter minima \\[c].\n\n4. **Restricted Applicability**: The analysis centers on fixed-feature PLRF models, which limits insights into DANA’s performance in deep networks with learned representations. The LSTM experiments (Section L) offer some evidence, but serve mainly as a proxy. Experiments on models like CNNs or Transformers would better demonstrate practical utility.\n\n**References**\n\n[a] Mandt, Stephan, Matthew D. Hoffman, and David M. Blei. \"Stochastic gradient descent as approximate bayesian inference.\" Journal of Machine Learning Research 18.134 (2017): 1-35.\n\n[b] Li, Qianxiao, and Cheng Tai. \"Stochastic modified equations and adaptive stochastic gradient algorithms.\" International Conference on Machine Learning. PMLR, 2017.\n\n[c] Yang, Ning, Chao Tang, and Yuhai Tu. \"Stochastic gradient descent introduces an effective landscape-dependent regularization favoring flat solutions.\" Physical Review Letters 130.23 (2023): 237101.", "was_revised": true, "final_justification": "I thank the authors for a thoughtful and comprehensive rebuttal. The authors have successfully addressed my primary concerns. The clarification in Appendix C.2.1, which links the simplified ODEs to the expected risk of an SDE model, was particularly helpful and resolved my question about the deterministic approach. Furthermore, the authors' justification for using a scalar momentum coefficient for the isotropic PLRF model is sound, and I appreciate the authors' candid discussion of this as a limitation for more general, anisotropic settings. In light of these clarifications, I have raised my score to 5.", "rebuttal": "​​We thank the reviewer for their thoughtful comments and feedback on our work. We are especially appreciative that the reviewer views our contribution as “novel” and “well-justified”. We are also happy that the reviewer feels that our theoretical results are “sophisticated”, and that they were appreciative of the clarity “despite the technical complexity”. Building on the reviewer’s summary of our paper, we would like to take this opportunity to reiterate the core contributions of our work:\n1. We provide **the first theoretical proof that one can change the scaling law exponent over SGD**, which requires new proof techniques in particular by including model-size dependent behaviors. \n2. **New algorithmic development with the use of data-dependent learning rate.** Our work shows that a data-dependent learning rate, $\\gamma_3(t) = (1+t)^{-1/(2\\alpha)}$ can greatly improve algorithmic performance.\n3. **Hyperparameter transferability.** The data exponent, $\\alpha$ is independent of model size. Therefore, one could find the $\\alpha$ on a small model and it potentially transfers to larger models without additional hyperparameter tuning (see Fig. 2, right). \n4. **Potential applications to practical algorithms.** Many of the outscaling phenomena that we proved on the PLRF, such as $\\gamma_3(t) = (1+t)^{-\\kappa_3}$ living near the edge of stability persisted to LSTMs (see Fig. 2,  Fig. 18-23). We believe that this bears promise to build a competitive optimizer at scale, backed by theoretical guarantees. \n\nWe now address the key clarification points raised by the reviewer:\n\n**Excessive length**\n\nWe acknowledge the reviewer’s concern about the length of the manuscript. We wanted to cover many related materials to provide as much of a complete picture. In particular, we provided proofs of outscaling for many different algorithms as well as extensive experiment. These are taking a lot of space but we believe them to be core contributions of our work. We tried a lot but would be more than happy to answer any question or suggestion the reviewer may have to improve the clarity of our results. In particular, as suggested by Reviewer NJzg02, we will add a detailed proof sketch (see our answer to their review) to clarify the role of the different sections of the manuscript.\n\n**Comparison to Adam and would coordinate-wise or adaptive scheme be more effective?**\n\nWe acknowledge the reviewer’s concern regarding our comparison to Adam while using a coordinate independent scalar learning rate. The features of the PLRF are isotropic, and so coordinate-wise step-sizes are not very interesting.  We mostly put the comparison (Fig. 1) because Adam is so widely used.\n\nThe reviewer is correct: in an anisotropic model, a coordinate-wise optimizer would be more effective and would be an interesting subject of study.  To our knowledge there is no proven scaling law for Adam on either an anisotropic model or the PLRF.\n\nWe also think a coordinate-wise optimizer, something which combines RMS-Prop and DANA, is important for practical usability of DANA and is a super direction for future research. We have some promising preliminary experimentation on transformers language models (GPT2 architectures) on the fineweb dataset that shows that mixing DANA with adaptive schemes from RMSProp shows a similar acceleration against the vanilla adaptive baseline RMSProp. We unfortunately cannot share this picture due to NeurIPS policy.\n\n**Deterministic vs SDE-based approach**\n\nThe reviewer raises a very interesting question regarding the link between SDE and ODE based approaches. We agree with the reviewer’s comment that our analysis is deterministic since we study ODE on expected quantities (see sec and derivation of ODEs). We show in Appendix C.2.1 that the ‘simplified ODEs’ are the expected risk of the SDEs. There could be potentially other interesting directions of inquiry (like in [a],[b],[c]), such as how concentrated the risk curves are around the ODEs, or other properties of the distribution of the iterates under the SDE (e.g. tail behavior). \n\n**Practical relevance and test on other architectures**\n\nThe reviewer raises the interesting question of the practical performance of DANA-decaying on real-world problems. We kindly highlight that our work's main contribution is a theoretical proof of the outscaling of a stochastic momentum algorithm against SGD. We additionally performed extensive experiments on the PLRF demonstrating this outscaling phenomenon, and large scale experiments on an LSTM model showing promising results of outscaling in practice with an improved scaling exponent compared to SGD. Additional promising experiments on a transformer GPT2 model suggests outscaling could occur with $\\kappa_3\\approx 0.7$  very similar to the one used for LSTMs. This highlights the benefit of DANA-decaying which depends on a data-dependent hyper-parameter $\\alpha$ that may be easily transferable across architectures and scale (see Fig. 2 right). We suspect together with the reviewer that an adaptive version may be needed to make it competitive and believe this is an open and crucial avenue for future research.\n\n**Visualization of Volterra components throughout training**\n\nThis is a great question and we can indeed simulate these. The way one would do this is that we would write down the ODEs that describe the behavior of F and K (similar to what is currently in the paper) and we can simulate these ODEs.\n\n**Limitations**\n\nWe thank the reviewer for their detailed feedback and suggestions. We believe to have already addressed some of the limitations raised by the reviewer, especially on the use of **non-adaptive learning rate** (lines 262-264 “While we compare DANA against other non-adaptive momentum methods, most real-world problems benefit from adaptive methods such as Adam. Hence it would be interesting to have an analysis of DANA combined with Adam or other preconditioned methods.“) and of **fixed features, quadratic model** (line 259 “A full convergence argument beyond quadratics for DANA-decaying would be desirable”). We do not want to overclaim any of our results and will **extend the limitation section** using the reviewer’s suggestions as follows:\n* **fixed feature** Our analysis is restricted to the setting of linear regression with quadratic loss. Hence this does not incorporate feature learning which has been shown to be a critical factor of success of recent machine learning systems. Incorporating such behavior in momentum analysis, for example using multi-index models as pursued in [a] is a very interesting and non-trivial direction for future work.\n* **global momentum parameter** Our analysis uses a global learning rate. While our core contribution focuses on how to optimally scale this learning rate using model size or data dependent quantities, this has to be compared with most state-of-the-art optimizers which use per-parameter learning rate by incorporating adaptive updates. We believe that extending our analysis to adaptive learning rates is a rich and non-trivial future research direction. No work has yet derived scaling laws for adaptive SGD (Adam or RMSProp), even without momentum.\n* **deterministic dynamics** Our analysis is based on ODE description of the risk by taking expectation over the random sampling of data points. Hence our results do not capture stochastic variance of sampling particular datapoints. Extensive PLRF experiments suggest general agreement between deterministic dynamics and actual evolution but any theoretical guarantee would be a very valuable future contribution. In particular, in light of recent work [77, Theorem 1.1] we believe that at least below the high-dimensional line $2\\alpha<1$, we could similarly prove concentration of the training dynamics around deterministic quantities.\n\nWe thank the reviewer for their valuable feedback and great questions. We believe we have answered to the best of our ability all the great questions raised by the reviewer and we kindly ask the reviewer to potentially upgrade their score if they are satisfied with our responses. We are also more than happy to answer any further questions that arise.\n\n[a] Ren, Yunwei, et al. \"Emergence and scaling laws in sgd learning of shallow neural networks.\" arXiv preprint arXiv:2504.19983 (2025)."}, {"review_id": "4iNAhYljP6", "reviewer": "Reviewer_NJzg", "summary": "The authors study the scaling laws for stochastic momentum algorithms for random feature models with the input \ndistribution and ground truth direction satisfying power laws. They show via both theoretical analysis (under certain \nsimplifying assumptions) and experiments that (a) SGD with momentum yields the same scaling law exponent as SGD, and \n(b) the scaling law exponent can be improved if we allow the momentum learning rate to depend on the dimension \n(or the time-dependent effective dimension).", "questions": "* In Assumption 1, you implicitly assume that the input distribution and the ground truth direction share the same \n  eigenvectors. Is this essential? What will happen if they are not completely aligned?\n* Is your rates optimal in any sense? (Is there a way to define optimal in this setting, that is more fine-grained than the minimax rate?)", "limitations": "The authors suggest in the introduction that the theoretical analysis is rigorous (line 30-32), while they actually \nneed to ignore some higher-order terms and use the deterministic equivalent (which is only mentioned in one sentence \naround line 131) at certain point of the analysis. To be fair, the authors do provide experimental results to justify \nthese approximations, but they should be more clear about this non-rigorousness.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "### Strengths\n\nThis seems to be an important paper, as it shows that at least in certain cases, changing the algorithm, in particular, \nletting the momentum parameter depend on the (effective) dimension, can potentially change the exponent of the scaling law,\ninstead of just the constant. Moreover, the authors conduct extensive theoretical analysis and experiments, covering a\nwide range of possible power law random feature parameters, algorithms, and training regimes, to support their claims. \n\nThe framework (Gen-Mom-SGD) the authors use covers many existing algorithms as special cases and might be of \nindependent interest. In particular, this framework/strategy allows continuous-analysis of momentum-based algorithms, \nwhich is often quite tricky.\n\n### Weakness\n\nOne issue of this paper is the lack of a proof sketch section. The main text is mostly about the results, experiments,\nand interpretation. Without a proof sketch section (which can be put in the appendix), it is extremely hard to check the \nsoundness of the proof of a 194 pages long paper.", "was_revised": true, "final_justification": "I believe this is an important paper as it proves that, at least in certain cases, the scaling law exponent can be changed by changing the algorithm and contains extensive experiments to support their theoretical results. \nThe authors provided a proof sketch in the rebuttal. It looks reasonable to me and this resolves my main concern.", "rebuttal": "We thank the reviewer for their time, constructive feedback, and positive evaluation of our work. We are heartened to hear that the reviewer feels that our results on changing the exponent make “an important paper”, and that our theoretical and empirical evaluation was “extensive”.  We address below the key clarification points raised by the reviewer.\n\n**Lack of a proof sketch section**\n\nWe acknowledge the reviewer’s legitimate concern about the lack of a proof sketch section. We provide below a detailed proof sketch which we will add to the updated manuscript.\n\n*Detailed proof sketch:*\n\nI. The first step is to derive an evolution equation for the quadratic risk using the raw algorithm updates Gen-Mom-SGD. \n1. We adopt in section C.1 a continuized analysis in (Gen-Mom-SGD) by replacing the update times by the jumps of a Poisson process of rate $1$ (eq. 17-18). This allows us to reduce the analysis to ODE systems without sending the learning rates to $0$.\n2. In Sec. C.1.1 we derive ODEs for the projected risk on each eigenvalue direction of the covariance matrix for two algorithms: the original  one in Sec. C.1.1 and the coin-flip algorithm in Sec. C.2.2. The coin-flip algorithm uses two independent Poisson processes for the momentum and parameter updates in (Gen-Mom-SGD) which gives rise to much simpler ODEs.\n3. The ODEs obtained are theoretically solvable but a lot of small high-order terms make the analysis difficult. Hence we drop high-order terms in $\\lambda$ in the exact and coin-flip ODEs to obtain the ‘simplified ODEs’ (Sec. C.2.2). We prove in Sec. C.2.1 that the ‘simplified ODE’ exactly represents the evolution of the SDEs presented in Sec. 3.\n4. Finally, using these ODEs, we can derive a Volterra equation for the full quadratic risk $\\mathcal{P}(t) = \\mathcal{F}(t) + (\\mathcal{K} * \\mathcal{P})(t)$ where $\\mathcal{P}(t)$ is the risk, $\\mathcal{F}(t)$ is the forcing function, $\\mathcal{K}(t,s)$ is the kernel function. This is done in Sec C.1.2 and C.1.3.\n\nThe rest of the proof focuses on computing asymptotics for $\\mathcal{F}$ and $\\mathcal{K}$ (eq. 54) by a) quantifying the empirical distribution of the covariance matrix spectrum and b) solving the ODEs.\n\nII. Quantifying the spectrum of the data covariance matrix $\\hat{K}$. \n1. We first rewrite $\\mu_{F}, \\mu_{K}$ in eq. 54 as integrals against two measures $\\mu_\\mathcal{F},\\mu_\\mathcal{K}$ which are defined through the deterministic equivalent in eq. 62. \n2. Using Random Matrix Theory and previous results from [78], we provide upper and lower bounds on the deterministic equivalent measures $\\mu_{F}, \\mu_{K}$ in Sec. D.1 and D.2.\n3. In Sec. D.3 and D.4 we use these bounds to define diverse components of the forcing function and kernel function (eq. 66, 68), formalized in Prop. D.13, D.14.\n\nIII. We then solve the ‘simplified ODE’ eq. 53 for each class of algorithm in (Gen-Mom-SGD) separately. \n1. For SGD-M we solve the ODE in Sec. G.1. These are the simplest as the ODE has constant coefficients, and the solutions are exponentials. \n2. For DANA-constant, the ODE has time-varying coefficients. We employ Frobenius theory to obtain asymptotic solutions of the ODE in Sec. H.2 to H.6. \n3. For DANA-decaying, we employ orthogonal polynomial theory to solve the ODEs in Sec. I.1.\n\nIV. Finally, we need to solve the Volterra equation for the risk.\n1. For each algorithm to remain stable, the kernel norm has to be smaller than $1$. This leads to stability conditions on the learning rates for SGD-M (Cor. G.1), DANA-constant (lemmas H.4, H.5) and DANA-decaying (Prop. I.5).\n2. Using this stability condition we can simplify the Volterra equation as $\\mathcal{P} \\asymp \\mathcal{F} + \\mathcal{K}$ using Kesten’s lemma (Prop. G.8, H.20, I.5).\n3. We then compute $\\mathcal{F}, \\mathcal{K}$ for each algorithm by integrating the previous solutions against $\\mu_{F}, \\mu_{K}$ and obtain explicit scaling laws for the loss $\\mathcal{P}$. This is done in Prop. G.4 for SGD-M, Prop. H.13, H.14, H.19 for DANA-constant and Prop. I.1, I.3, I.4 for DANA-decaying.\n4. Finally, we determine the compute-optimal regime in Sec. E by a straightforward comparison of the different terms in the loss.\n\n**Alignment between data and target**\n\nWe understand the reviewer’s skepticism regarding this particular assumption and asking the question whether allowing for a different structure would modify the results. First, we acknowledge that Assumption 1 and more generally the PLRF model are crude approximations of real-world scenarios. The alignment between data and target was already used in numerous previous works under the source-capacity conditions. This is the simplest model we decided to study as it is already giving rise to deep and rich behaviors as exemplified by DANA-constant and DANA-decaying stability conditions and observable in more complex systems such as LSTM where even the notion of target vector becomes ambiguous. \n\nRegarding removing this assumption towards a more general scenario of misalignment, we believe that it would indeed affect the loss curves, although it is not essential for the analysis. Indeed, the analysis only depends on the projection of the target vector in the data-covariance eigenbasis. As an example, if we multiplied the target vector by some random matrix, the expansion in the kernel eigenbasis would no longer be a pure power law as we have but saturate after some point; this could potentially result in a different scaling-law phase plane.\n\nFinally, we highlight that due to the random matrix $W$, the data covariance eigenbasis and target vector are not completely aligned. This was the main focus of the random matrix part on the deterministic equivalent, which adds the embedding bias term $\\mathcal{F}_{ac}$ to the loss. This important term has a dominant role in phases 2 and 3 and was the main reason behind recovering the Chinchilla scaling in phase 3 for SGD (see [78]).\n\n**Optimality of our rates**\n\nWe thank the reviewers for raising a very interesting question regarding the optimality of our rates.\n\nWithin the algorithm class of Gen-mom-sgd, our theoretical analysis suggests that DANA-decaying is the stable algorithm that reaches the irreducible loss level with the least number of iterations. We kindly refer the reviewer to Section I.5 where we propose extended heuristics to derive scaling laws of all stable algorithms with the DANA class, extending results from Theorem 1. These heuristics show a continuous evolution of scaling laws from SGD to DANA-decaying and DANA-decaying to DANA-constant. This is consistent with Fig. 5 where we plotted the time to reach irreducible loss for algorithms within Gen-Mom-SGD for a specific choice of $\\alpha, \\beta$. We indeed observe continuous evolution of scaling laws between the three remarkable points SGD, DANA-constant and DANA-decaying, with DANA-decaying being the optimal one.\n\nThe reviewer additionally raised the non-trivial question of how to generally define optimality. In [91], the authors define a class of first-order stochastic algorithms by allowing access to gradients at random data points in the same way as [73] defined for full-gradient. [91] then showed a minimax lower bound on the convergence rate by showing existence of a data distribution for which the rate for any first-order stochastic algorithm after $d$ steps is at least $\\Omega(1/d)$ compared to the initialization risk, ie $\\mathcal{P}(d) \\gtrsim \\frac{||\\theta_0-\\theta_*||^2}{d}$. We reported these in Table 2. As we noted in the caption of Fig. 2, “our results improve over the bounds in [91]. In particular, the two results agree as one approaches stationarity.“ Around $\\alpha\\approx \\frac{1}{2}$, we especially recover the lower bound $\\mathcal{P}(d) \\gtrsim \\frac{||\\theta_0-\\theta_*||^2}{d}$.\n\nFinally note that Gen-Mom-SGD assumes sampling a new iid data point at each iteration (as in [91] definition 4). Allowing the reuse of samples may improve sample complexity. In particular, we refer to [79] which derives improved rates by allowing multiple passes over the data. Moreover in a concurrent work [a], the authors show an improved scaling law via data reuse when the data satisfies a power-law covariance (no embedding matrix $W$) only for some choices of $\\alpha, \\beta$. It is a very interesting question to see if one can derive improved rates using a  DANA-decaying style algorithm with data reuse on the PLRF model. One should be able to modify the results of [77] to apply to the PLRF and compute the scaling laws similar to this work. To the best of our knowledge, this has not been done.   \n\n**Limitations**\n\nWe acknowledge the reviewer’s healthy skepticism regarding our assumptions and approximations and are happy that they found our empirical justification compelling. We will make it clear on lines 30-32 that we are studying a model of the problem.  \n\nWe tried to make every approximation clear and to provide clear definitions of the mathematical objects we use, eg ‘simplified ODEs’ (line 117 “We formulate all results going forward for the simplified ODEs“) and deterministic equivalent measures (line 132: “all the scaling law statements that follow are proven for the deterministic equivalent“). Note that the gap in the deterministic equivalent study can be removed and made entirely formal at the cost of removing the random embedding matrix $W$ which will remove the $\\mathcal{F}_{ac}$ loss term. Finally, note that although approximated from the exact ODEs of Gen-Mom-SGD, the ‘simplified ODEs’ track **exactly** the evolution equation of the SDE system in Section 3, without approximation.\n\nWe thank the reviewer again for their review and detailed comments that helped strengthen the paper. We hope our answer allows the reviewer to consider potentially upgrading their score if they see fit. We are also more than happy to answer any further questions.\n\n[a] Lin, Licong, Wu, Jingfeng, and Bartlett, Peter. “Improved Scaling Laws in Linear Regression via Data Reuse.”"}]

Author rebuttal:

Meta-review: This paper investigates the scaling law of SGD with momentum under power-law inputs. Both theoretical analysis and empirical results demonstrate that SGD with momentum yields the same scaling law exponent as plain SGD. The authors also introduce DANA, a method in which the momentum parameter adapts to the model dimension, leading to improved scaling behavior.

Most reviewers agreed that the paper deserves acceptance. After carefully reading the rebuttal and subsequent discussion, I also share this view. However, I recommend that the authors address the reviewers’ feedback in the final version of the paper, with particular emphasis on:

1-Providing more detailed explanations of the theoretical analysis.

2- Including experiments on models beyond large language models (LLMs).
Decision: Accept (spotlight)