High-dimensional limit theorems for SGD: Momentum and Adaptive Step-sizes
Authors: Aukosh Jagannath
Taj Jones-McCormick
Varnan Sarangian
Venue: ICLR 2026 Poster (2026)
Keywords: stochastic gradient descent
momentum
adaptive step-sizes
scaling limits
high dimensional probability
spiked tensor PCA
single index model
Primary area: learning theory

Abstract: We develop a high-dimensional scaling limit for Stochastic Gradient Descent with Polyak Momentum (SGD-M) and adaptive step-sizes. This provides a framework to rigourously compare online SGD with some of its popular variants. We show that the scaling limits of SGD-M coincide with those of online SGD after an appropriate time rescaling and a specific choice of step-size. However, if the step-size is kept the same between the two algorithms, SGD-M will amplify high-dimensional effects, potentially degrading performance relative to online SGD. We demonstrate our framework on two popular learning problems: Spiked Tensor PCA and Single Index Models. In both cases, we also examine online SGD with an adaptive step-size based on normalized gradients. In the high-dimensional regime, this algorithm yields multiple benefits: its dynamics admit fixed points closer to the population minimum and widens the range of admissible step-sizes for which the iterates converge to such solutions. These examples provide a rigorous account, aligning with empirical motivation, of how early preconditioners can stabilize and improve dynamics in settings where online SGD fails.

Published as a conference paper at ICLR 2026
HIGH-DIMENSIONAL LIMIT THEOREMS FOR SGD: MOMENTUM AND ADAPTIVE STEP-SIZES
Aukosh Jagannath
Department of Statistics and Actuarial Science, Department of Applied Mathematics Cheriton School of Computer Science University of Waterloo, Canada
a. bosgannath@watterlloo
Taj Jones-MeCormick, Varnan Sarangian Department of Statisties and Actuarial Science University of Waterloo, Canada (tejonesm, &
ABSTRACT
We develop high-dimensional scaling limit for Stochastic Gradient Descent with Polyak Momentum (SGD-M) and adaptive step-sizes. This provides framework to rigourously compare online SGD with some its popular variants. We show that the scaling limits of SGD-M coincide with those of online SGD after appro- priate time rescaling and a specific choice of step-size. However, the step-size is kept the same between the algorithms, SGD-M will amplify high-dimensional effects, potentially degrading performance relative to online SGD. We demon- strate our framework on two popular learning problems: Spiked Tensor PCA and Single Index Models. In both cases, we also examine online SGD with an adap- tive step-size based on normalized gradients. the high-dimensional regime, this algorithm yields multiple benefits: its dynamics admit fixed points closer to the population minimum and widens the range admissible step-sizes which the iterates converge to such solutions. These examples provide rigorous account, aligning with empirical motivation, of how early preconditioners can stabilize and improve dynamics in settings where online SGD fails.
INTRODUCTION
Stochastic gradient descent (SGD) and its variants are central to large-scale optimization modern machine learning. A significant theoretical interest lies in characterizing and comparing the per- formance of these algorithms in the high-dimensional setting, where data and compute are limited relative to the data dimension and model complexity.
The fixed-dimensional asymptotic theory of SGD classical, dating back to Robbins & Monro (1951): McLeish (1976): Kushner (1984). In this regime, the small step-size limit of online SGD converges to the gradient flow on the population loss. More recently, in the high dimensional regime, there have been numerous works on scaling limits where the dimension tends to infinity for specific
problems, including linear regression (Wang et al. 2017 Paquette et al. 2021 2024), online PCA (Li al. 2016 Wang al.) (0177), single-index models (Saad & Solla 1995a Goldt al. 2019
Veiga et al. 2022), and multi-index models (Collins- Woodffn et | 2024). ) A unifying framework for proving such scaling limits was introduced by Ben Arous al. (2022 2024), including an extension to study diffusive dynamics. A key insight from the figh-dimensional perspective is that when the dimension scales up as the step-size tends to zero, there exists critical scaling regime of the step- size in which high-dimensional effects yield distinct dynamics. The fixed-dimensional dynamics recovered with sub-critical step-size, but at the critical scaling, there are additional corrections, such as the "population corrector," in addition to the gradient flow drift.
(2022]2024),
For SGD with momentum (Polyak) 1964 Nesterov 1983 Sutskever et al.
2013), recent works the
fixed dimensional setting consider similar continuous-tima ballistic limits (Su al. 2016 Kovachki

---

Published as a conference paper at ICLR 2026
& Stuart 2021: Wilson al. 2021 Feng et al. 2023 There is also well-established body of literature characterizing the benefits of SGD-M over online SGD, such as improved convergence
rates, reduced escape times from saddle points, and implicit regularization effects (Liu et al. [021]
Wang et al. 2021 Cowsik et al. 2022; Gess & Kassing 2023; Ghosh et al. 2023 Dang et al. 2025).
Conversely; in the high-dimensional setting, Paqqeette & Paquette (2021) Ferbach al. (2025) study
the limiting dynamics of several comentuum-assed stochastic algorithms in the setting where they
scale the momentum parameter accordingly.
Preconditioned SGD methods incorporating adaptive step-sizes (Duchi et al. 2011 Hinton 2012 Zeiler, 2012 Kingma& & Ba 2014,) Yu et al. [017 [Shazeer & Stem (2018) are also theoretical
terest, particularly for explannigg their empinical success in high- nigh-dimensional, non-convex optimiza- tion (Zhang al. 2020b). fxed-dimensional settings, notable works addressed similar questions
regarding convergence behaviour and continuous-time dynamics (da Silva & Gazeau 2018: Reddi
al., 2019 Barakat & Bianchi 2021 Ma et al., 2022] Malladi et al. [022]] Ghosh al., 2023). The
preconditioner we consider in Section 3 gradient normalization, was first introduced to mitigate ex-
ploding gradients in deep sequence models (Mikolov) 2012 Pascanu et al. 2013). Since then, there
have been several works looking at the effectiveness of "clipped" gradient control (Li et al.) 2017;
Chen et al.] 2020 Zhang et al., [0200a] Koloskova et al.] 2023) as well as high- nigh-dimenssonal scaling
limits for random least squares regression (Marshall al. 2024).
1.1 CONTRIBUTIONS
We extend the effective dynamics framework of Ben Arous et al. (022)) to derive the high- dimensional limiting dynamics of SGD with Polyak momentum (SGD-M) and of SGD with adaptive step-sizes. We find that the critical step-size scalings are the same between these dynamics. Further- more, we find that in the critical step-size regime, SGD-M amplifies the high-dimensional effects as compared to SGD, potentially steering the dynamics even further from the population gradient. More precisely, for any sequence of step-sizes in the critical scaling regime, the effective dynam- ics of SGD-M accentuates the aforementioned emergent "population corrector" observed in online SGD. However, for any given step-size of SGD-M, there exists a corresponding step-size for online SGD that yields the same effective dynamics, up to a time rescaling, establishing an equivalence between the algorithms.
In contrast, we find that training can be improved by incorporating pre-conditioning. To this end, make of our framework to compare SGD-M with version of online SGD equipped with a simple adaptive step-size that restricts the gradients to have unit norm (Mikolov) 2012] Pascanu al. 2013); we refer this variant as "SGD-U" in subsequent sections. We prove, in two popular classes of learning tasks, that the SGD-U can exhibit dynamics with not only superior fixed points (i.e. closer the population minimum) but can significantly broaden the range of admissible step- sizes that ensure convergence to these solutions. The admissibility of larger step-sizes that maintain stable dynamics and yield high-quality solutions provides a rigourous demonstration of how pre- conditioners can mitigate high-dimensional effects and empirical phenomena such as exploding or vanishing gradients. Thus, our framework offers rigorous justification for the empirical motivations underlying the design of early preconditioners.
While our framework is broadly applicable, we illustrate our results on two problems in high dimen-
sional inference: spiked tensor PCA (Johnstone) 2001 Péché 2006; Montanari & Richard 2014) and single index models (Bishop & Nasrabadi] 2006;
Hastie et al. 2009 Barbier et al.) 2019
MAIN RESULT
In this section, we present our main results, starting with the scaling limits of SGD-M, followed by discussion of how the same methodology be used to study effective dynamics for scalar preconditioners. All proofs are deferred to the appendix.
We consider the following online learning setup. Suppose y1, Yy.... taking values in y C Rd with law P, E Mi (R**) and loss L X, x Vn, + IR where X, C RP is the parameter space. Consider online SGD with learning rate 8, and momentum
given i.idd. data sequence
2

---

Published as a conference paper at ICLR 2026
rate B E [0, 1), given by
Pr = 3pr-1 8n V (2
Tt - pt
with possibly random initialization tion of a finite collection of summary statistics un (x) = both pv and dn may grow with and 6n, - 0 as nl oc.
Mn E (M)) Our interest is understanding the evolu-
( (z)) the regime where
We define the functions
(x Y) (x, Y) (x) where (+(x)= = - Y)
In the following, suppress the dependence of H Y and let V(x) (T) VH(x)) be the covariance matrix for VH at z.
To proceed, we state two assumptions: 8, 8,-localizability and asymptotic closability. The first, which we present next, provides an upper bound on the learning rate in terms the regularity of the sum- mary statistics and data distribution. These bounds ensure tightness of trajectories of the summary statistics.
Definition 2.1. A tuple (un, Ln,, Pn, B) is o, -localizablle with localizing sequence (Ek) if there is an exhaustion by compacts (Eạ) of R and constants Cк such that for any S k,
1. supreu, (Exx
V u ol CK on 1/2 and sup supzeu.. (Ex)
| lop S CRi
2. sup
V S Ck. and supseur. (Ex) E ||| CkO
3. sup suPreuu,*(ER VH, V 4 S C 2 and
suPnen,' E V u (x). V H (z, y) @ (H( - V(x)) = o(5 and for
SUP(r))Eun. F y2 2V (z1), (I v1) & (2)))
The second assumption ensures that the limiting coefficients for the evolution of the statistics close. To this end, we define the first and second-order differential operators,
A, ( V) and L. 2 (V, V
Definition 2.2. A family of summary statistics (u,) are asymptotically closable for learning and momentum rates ([6n), 3) if (un, Pa.,P 3) are l..-localizable with localizing sequence (Ex)x, and furthermore there exist locally Lipschitz functions h IR IR - R* and Σ R2* k such that
11-1-1
+ (1 (1-30 B)2 5n5n (x)(x)) - (3, un (z)) +
sup .uu, (FEx)
sup s
V -((u(()))
with Jn the Jacobian of un. We call h the effective drift and Σ the effective volatility of u respec- tively.
We refer the reader to Ben Arous et al. (2024) for detailed discussion regarding the interpretations of Definition 2.1 In particular, these scaling assumptions are more general than the regularity as- sumptions commonly imposed in the literature, e.g.. L-smoothness and convexity of the loss. These assumptions can be checked in broad class of models, such as those studied here or in (Ben Arous jet al., 2024) and related works. For example, one can verify that Definition 2. holds for com- mon models such as binary classification and Gaussian XOR classification with two-layer neural networks with ReLU activations (Ben Arous et al. 2024)..
Let briefly compare the notions of localizability and asymptotically closability to the related notions in Ben Arous et al.] (2024) Compared to online SGD, we now require an additional item in Definition (3) to control the correlation between the random fluctuations H (x1, Y)), V V Hx,, YY, when * j. Additionally, for Definition 2.2 we note that the two operators seale differently in B3. We revisit this point in Remark 2.4mar22.4afffe presenting main result below.

---

Published as a conference paper at ICLR 2026
Theorem 2.3. Let (Xx")e e be SGD initialized from Xu Mn for Mn 11 (RPt) with learn- ing rate on and fixed momentum parameter 3 E [0. 1) for the loss Ln (:s.) and data distribution Pa. For a family of summary statistics u, (u.] Jim1 let (u), be the linear interpolation of (un (X [(m-"]]e--
Suppose that un are asymptotically closable with learning rate on, effective drift h, and effective volatility E. and that the pushforward of the initial data has (un), , fn + weakly for some M1 (R*). Then (t)( - (u+), weakly as - oc, where u, solves:
du = h(3, u.) dt + 1 3
Σ U1 dB,
initialized from where B, is standard Brownian motion in R
Taking 3 recovers the original result due to Ben Arous et al. (2024) for online SGD. The proof is presented in Appendix B
Remark 2.4. It is beneficial to compare SGD-M with online SGD as well as compare the fixed and
high dimensional regimes in light of Theorem 2.3 To do so, suppose that we further impose the individual limits
sup (x) - (x)]l + O, a (Ek)
sup sER, (Bk)
[1 - g(un (x))l +
for some functions f.g. IR + I2* such that the revised dynamics are spelt out as
1 du. - - f(u.) (1- (1-)))(u))) dt + 1- 3 E(u,) dB,
One can loosely interpret the population drift I as the "learning" or "signal" term which pushes the dynamics along the descent direction of the population loss ④. The population corrector g is then an emergent "variance" term in the critical scaling regime.
To compare the behaviour of SGD-M with online SGD, consider the dynamics of the latter with step-size dn ōn.(1 3). The limiting dynamics then:
du: = [f(ue) + 1- g(( dt +
E(u,) dB,
t/(1 B). In particular, the relative reweighting between f and g are modified so that the corrector poses a stronger influence on the dynamies, potentially overwhelming the underlying signal. Evidently, for the case of SGD-M, we see that becomes more prominent as 3 tends to one unless we adjust the step-size accordingly.
which coincide with equation 2 after additional time time-reccalingt
As consequence of Remark 2.4 any trajectory produced by SGD-M some time interval [0, T] can be replicated by online SGD after a time change where both the step-size and time horizon are rescaled by 33- The effective number of iterations 7/8 remains unchanged, confirming similar observation to both Kovachki & Stuart (202I), who considers the fixed-dimensional deter- ministic regime, and PPqquette & Paquette (2021). who considers the high-dimensional regime for the case of linear regression.
Finally, we note that our results subsume the classical ODE theory for diceddiimesional SGD. That is, the high--iigenn -igh-dimensional dynamics reduce to the ballistic phase
du, =
(u,) dt
by taking 5, below the critical scaling threshold. In this subcritical regime, the relative scaling of the population drift and population corrector becomes moot g is negligible.
2.1 SCALING LIMITS OF ADAPTIVE STEP-SIZES
Following the setup of Theorem 2.3 consider the modified online SGD update step with a scalar-
valued preconditioner 7)n X, yn RR+, "iivvee by
Et Tt-1 8 17(2 1, 4t V VL(
(3)
4

---

Published as a conference paper at ICLR 2026
Under the same approach of online SGD, can develop scaling limit by decoupling from the data-dependent component of 77, i.c. we define
(z, Y) - Th (2, V (z, Y)= (x(x)
where
< - (x Y) V (2, Y)]
which serve the analogue to V and V H defined before. Provided VH and V satisfy Definition 2.1 and 2.2 then Theorem 2.3 with 3 = 0 extends naturally to the preconditioned case rigorously derive scaling limits for this family of pre-conditioners under this same framework, which we demonstrate next.
(see Appendix A for extended discussion). Hence,
EXAMPLES
In the remaining subsections, we demonstrate Theorem 2.3 and the followup discussion in Sec- tion 2. for the spiked tensor model and single index model. In particular, consider the dynamics of SGD-M as well as the dynamics for version of online SGD with,
11 Y) (1, Y)I
(4)
We impose the additional dimension factor /n (see notation below) when normalizing the gradient to preserve the scaling relationships presented Definition 2.1 Since | |L O(vn) in both of our examples, we require the rescaling to ensure that the limiting dynamics are non-trivial. We abbreviate this version of SGD equipped with 77 as "SGD-U" in subsequent sections.
3.1 SPIKED MATRIX AND TENSOR PCA
Suppose given i.i.d. samples of data of the form y Ax We where W are i.i.d. k-tensors with Gaussian entries, E Rn is a unit vector, and = An > 0 is the signal-to-noise ratio. Our goal is infer the direction from the noisy sample.
We take as loss the (negative) log- liglliiel,hood,
Lxx, Y) IY - 121
to optimize parameter vector r. We consider the pair of summary statistics UA = (m, with
m mxx = (x, v) and = (z)
(x) == || mu| = || m
We also define R22 = m + 72 | 2 For both cases we take the critical scaling of 8, = cs/n where 0 is constant. Since goal is estimate the direction v, wish to optimize the ratio R representing the magnitude of the cosine ratio between x and v.
Proposition 3.1. Fix & 2, A 0, ca > 0, let 8n cs/n. For 3 [0, 1), u. (I) following dynamics induced by SGD-M converges + oo to the solution of the following ODE initialized from 1im (ux) pn:
dm 1- 3 2m (akm* - kRR((-11) dt, dr 2kR2(k-1)
1
(1² 3) cs) dt (5)
By²
Likewise for SGD-U, (L) converges weakly to the solution of the system
dm = / Vi
m
12 122-1m
de, dr = -2vk
R 1
dt
(6)
R
2
We verify localizability and closability for both SGD-M and SGD-U in Sections C.I.J C.I.2 as part of the proof of Proposition 3.1 We analyze the fixed points for both systems in the presented order. We find that the dynamics yield fixed points away from the axis m = O only if À exceeds a critical value, say Aeni (k, 3, cs). that depends on the tensor order, step-size, са and 3. For both SGD-M and SGD-U, we can show (see Appendix C.2 that the critical X for the respective methods are
M MM,, (k: B, ca)
C5
4/2
12
1))*-1
k(1
B)
kk
2)/2 21
5

---

Published as a conference paper at ICLR 2026
A - 0.8
À 1 1.2
2
[m(E)R(
SCD 44 5
500
02
0.
1000
25
10.0 22.5
25
10.0 125
20.
Figure 1: Matrix PCA in dimension 10000 for X 0.8 (left figure), A 1.2 (middle figure), and À 2.2 (right figure) with с = 1. Depicted is the evolution the statistic (m(t)/ for 20n steps with random initialization. We note that in the left figure, only SGD-U is supercritical. In the middle figure, both SGD-U and online SGD are supercritical, however the former attains better alignment with the direction vector The rightmost figure shows only SGD-M with 0.9 as subcritical, however the alignment order is consistent.
20 1/(k+1)
( 1
((k- 1)(k +2)
ca)=
(x-2)1++1)
(2 ((k: 2)( 1))
We fix in the following to determine the stability of the fixed points.
Proposition 3.2 (Fixed Points SGD-MM. Let 2, A > 0, cs 0, 3 (0, and denote oj by
cg cs/(1 (3) with AM (2, B, cs) cg- We then have the following fixed points:
1. An unstable fixed point (0, 0) and fixed point at (g) is stable if A AM (2, and unstable if X X M (2, cs)
2. If X M (2, 3, es): two stable fixed points at (+:m., (g) where m, V X 0, cs > 0, and denote of by cf cs/ (2/(2 with A (2, cs) (c/) We then have the following fixed points:
Proposition 3.3 (Fixed Points SGD-U). Let
1. An unstable fixed point at (0, and fixed point at (0, (c])2/3) that is stable if X v (2, c6) and unstable if X X (2,
2. If A (2, cg): two stable fixed points (#m., c// V with m. VA (G//V)
In Figure T we provide simulations for Matrix PCA (i.e. 2) with fixed cs to demonstrate how SGD-U and SGD-M with different values for 3 and behave in the ballistic phase. We see that as U (2, 0g) < XM (2, 3, cg) ) for 3 2 0, there are values of for which SGD-U is supercritical when SGD-M and online SGD are subcritical. This means that SGD-U able to converge towards one of the fixed points away from m 0 for values of À where the other methods fail to do so. A careful analysis reveals different settings of с and À for which SGD-M may converge desirable critical points instead. For example, fixing Cs, the fixed points suggest that SGD-U is preferred whenever À (1 3)8/8. Similarly, the emergence of fixed points away from m 0 depends on the choice of relative to X. Indeed, one can verify that the largest for which non-trivial fixed points emerge larger under SGD-U than SGD-M whenever
(1 - 3)2/8. This suggests that there is problem-specife critical X 1/8 dictating which of online SGD SGD-U is preferred.
In practice, the choice of cs is specified by the user. Thus one may be inclined to simply take cs small enough such that the non-trivial fixed points emerge. However, the number of iterations (and thus the requisite samples and algorithm run-time) scales as Tn/cs. i.e., inversely with cs. which naturally constrains the step-size from being arbitrarily small. Consequently, taking cs to be too small may prevent one from running the algorithm long enough for the dynamics to even reach the basin of stability for the emergent fixed point.

---

Published as a conference paper at ICLR 2026
À
A
2=22
0.5
0.15
010
0.10
0.05
0.00
SOO bett-O
SGD betanoo
500
Figure 2: Matrix PCA dimension n 10000 for À = 0.8 (left figure), = 1.2 (middle figure), and À 2.2 (right figure) with с 1. Depicted is the evolution of the rescaled statistic ūn (t) for 6n steps around a fixed window about m = We note the increased volatility of the diffusive limits for SGD-M as 3 increases. We that SGD-U becomes mean repellent for smaller values of X, similar to Figure
3.1.1 DIFFUSIVE LIMITS AT THE EQUATOR (m 0)
Here, we demonstrate an example of diffusive dynamics for Tensor PCA around m for SGD-U. We consider the rescaled observables ü. (f) = (ύ1, (22) = (Vrm, r²). Proposition 3.4. Fix 2 2, À > and 6, = ns/ Then for SGD-U, ú, (t) converges as n - oo
the solution of the following SDE initialized from lim,, (ūn (ūn)))/. (0, @ 61:
m
h-1
1==2 - dt + Vcs (k- m +1
1/2
dm
dB.
(7)
r
dr -2 pptt]
dt
(8)
We see that 7 now solves an autonomous ODE which converges exponentially to [cs/2]2/((+1)) In particular, when & 2 and t tends co, the equation for m behaves like
21/3m
2/3 2
22/3mm2 dt +
1/2
dr
1/3
A
1/3
C8
dB.
CB
C6
Interestingly, in contrast to the dynamics for online SGD (Ben Arous et al. 2024 Proposition 3.4), equation 7 no longer a standard OU process as the diffusion coefficient explicitly depends on Wi. However it is still reminiscent of one since, for example, the drift is mean-reverting or mean- repellent depending on the choice of сs relative to Figure presents several simulations of the diffusive limits for új under different setups of SGD-M and SGD-U.
3.2
SINGLE INDEX MODELS
Let E gn-1 be fixed direction and suppose we are given i.i.d. samples of data (y/, a De) under the model
o' = S(a v) +
where f R + R is possibly non-line non-linear link function (which is assumed to be known) and (c") are zero-mean additive emrors with E[e2] = σ² < co. Our goal is to infer from (y/, a²) where a are i.i.d. standard Gaussian feature vectors by optimizing the quadratic loss,
L(, (V- f(a' x))²
(f(a' v) f(a" .)) +
Similar to the previous example, we determine the evolution of the pair of summary statistics u, (m.r3) r)). The scaling limit for this problem was analyzed for online SGD by Saad & Solla [1995a b); Goldt et al. (2019): Veiga et al. (2022): Rangriz (002)). The critical scaling is determined by the
step-size 0, ca/n. n.

---

Published as a conference paper at ICLR 2026
Proposition 3.5. Fix C5 O, non-constant link function f R R, let 8n, cs/n. For B [0. 1), (t) following the dynamics induced by SGD-M converges as n + oo to the solution of the following ODE initialized from lim, (u.), Mni
-2 dm (1- 3)
E fa1 (a1m+ a2r) (f (aim + agr) - f(a1))) dt,
(9)
-4
dr
E [aarrf (a1m + a2r) (f (arm + a2r) - f(a1))]
(1- - 5)
C8
+ (a1m+ o2r) ((f (aim + agr) - f(a1))² + o²) dt
1- 3 (a1m
(10)
where d1, dg N0 1) are i.i.d, standard Gaussian variables. Likewise for SGD-U, (1) con- verges weakly to the solution of the system,
dm -Bla sga(J"(mai + (a2) + ra2) s(ai) + +)]] dl,
(11)
dr² 2 rEl sgn(f (mar + ra2) raz) f f(11 ) + c))l C8
dt (12)
2
where the sign function satisfies sgn(x) = I .>0 1x00-
We verify localizability and closability for both SGD-M and SGD-U as part of the proof of Propo- sition 3.5 in Section D what follows, we consider f(z) = x² for k 2 This family of link functions includes well-studied examples such phase retrieval (k 2). In the small noise regime o² + 0, the population optimum (m, r²) ) = (1, 0) is a fixed point for SGD-M (see Proposition 3.8), but not for SGD-U. Despite this, show that SGD-U is able to converge to fixed points that still close to the optimum in settings where SGD-M diverges (i.e. dr² > 0, vt).
For concreteness, present some specified choices of gauge the exact limiting dynamics and fixed points for SGD-M. We obtain closed-form dynamics and fixed points for SGD-U in the next section where we consider the small noise regime.
Proposition 3.6. If we take fz) = x, then the continuous-tim time dynamics for SGD-M in Proposi-
tion 3.5 with B [0, 1) reduce to
2
(m- 1) dt, dr 4
Co ((m +)) + 0²) dt
dm
1 1- r
with a² > 0. In particular, the dynamics admit unique fixed point (m, (1, /11 B cs))
and only if cs < 1 -.
Proposition 3.7. If we take f(x) x then the continuous-tima dynamics for SGD-M in Proposi- tion 3.5 with / E [0, 1) reduce to
12
4cs
dm
1-3 ml m( R² dt, dr 1
23 - -1) 1- - ()
+
with の² > and where we define the polynomial
P(m,r) 15m + 45m²,2 + 45m2r4 + 15. 30m 36m2 2_ 6r 15m² + 3
In particular, the dynamics admit the following fixed points:
1. For ca > (1 /2) we have a fixed point at (m,+²) = (0, and additional fixed points
(m, r2) = (0, - for r > O where 7e are roots (assuming they exist) to the equation,
4cs Gr - 1- 3 (15r* - 6r + +3+0²)
2. For C9 < (1 B)/12, we obtain the pair of fixed points (m. 1?) for ,2 where 2 = 2/022(1 3 - 12cs) and m. = + + -
The proof for both results is provided in Appendix D.2 These results inform of the inherent difficulty for the SGD-M (and consequently online SGD) dynamics to converge to fixed points the population minimum. To attain these fixed points, require taking smaller cs. thus increasing number of iterations to run the algorithm as per our discussion in Section 3.I We capture this phenomena more generally in the small noise limit next.
8

---

Published as a conference paper at ICLR 2026
fx) x2
f(x)
xx3
fxx==x x2
4x*
1.0
1.0
1.0
cole dgmaa))
95D-0 ced sigram)
300-1/ood) sayraud
sigmam0.22
0.6-
0.4
0.4
5G0 o-d digma-01
001 dgma-d
550 codd digrtbbi-4
001
dggaaae 01
100 cedit 6gmapa
0.0
0.0
50D- x1gga-11
0.0
20 60
100
20
40
60
80
100
20 40 60
80
100
Figure 3: We plot the value of mm/ft over the course of training for independent runs of SGD (full lines) and SGD-U (dashed lines) for various functions f with different amounts of additive noise. We set n 10, 000, 8 = ca/n and the total number of steps is taken as one million. We consider f(x) = x², f(x) = x²3 and f(x) x² + 4x each case we choose cs = 10- for the smallest integer k such that the dynamics are not effected by exploding gradients.
3.2. 1 SMALL NOISE 0² + LIMIT
In this section, consider the regime where we take x² - after co obtain closed-form dynamics and fixed points for SGD-U. We also generalize our observations from Propositions 3.6 3.7 to the family of monomial link functions (x) x² 1. To do so, we determine the existence of stable fixed points along the axis m 1. The following result considers the fixed point along for even-powered monomials by symmetry.
m 1, however the same holds for m
Proposition 3.8. Let k 1, B E 1) and take x² + 0. For the link function f(x) x the dynamics of SGD-M, as given in Proposition 3.5 admit a locally stable fixed point at (m,, r..) (1, if and only if we take c5 such that
1- (2k -3)!!
0< <
(13)
122 (4k 5)!!
For SGD-M, increasing the degree of the monomial link restricts the admissible values of cs for which the dynamics converge to fixed points away from m = In particular, one can show (e.g. Appendix D.2) that the rate for which the upper bound decreases is of the order A k---2 In particular, we observe that even though (1, 0) is fixed point for any cs >> O when o² = 0, is not stable unless cs satisfies equation 13 We next determine the dynamics for large collection of link functions for SGD-U in the small-noise regime.
Proposition 3.9. If a - 0, then for every strictly increasing link function f. SGD-U admits the following dynamics,
(m 1)2 dt, dr [- V(m - 1)2 + p2
m
2 2/
dm
Cg
dt
These dynamics yield a unique fixed point at (m, r2) (1, 1
This can be viewed as a type of "universality" result in the link function. For odd-powered mono- mials, equation 13 illustrated how сN must be supor-exponentially small in for the dynamics of SGD-M to converge to fixed point about m = 1. However, in Proposition 3.9. we see that regard- less of the degree, the dynamics under SGD-U always converge to fixed point along m = with ?² > 0. Setting сş arbitrarily small thus recovers fixed points near the population optimum. We can view this as a rigorous justification for how using normalized gradients, as in SGD-U, allows for choosing larger step-sizes for which the learning dynamics remain reasonable. This example validates some the heuristics and intuition behind using mnit-norm gradients to control exploding and vanishing gradients, causing instability for SGD (Mikolov) 2012 Pascanu et al., (203)).
To complement Proposition 3.9 we provide additional analysis for even-powered monomials in Appendix D.5 We show similar universality result where the dynamics again invariant with

---

Published as a conference paper at ICLR 2026
respect to the monomial's degree. To summarize results, we show that there also exists a thresh- old for O5 that determines the existence of fixed points away from m = 0. In particular, we find this threshold to be 4/ / 10w s 0.71, which is larger than, for example, the 1/12 threshold observed
for phase retrieval in Proposition 3.7 polynomials.
We suspect similar phenomenon to occur for more general
In Figure of >0 O where e' ~ (0,, o²) are i.idd. Gaussian variables. For each algorithm, we choose cæ such that the dynamics converge a fixed point. We see in the first recover the solution, although the step-size required to do so is smaller by one and two orders of magnitude respectively for SGD compared to SGD-U. For the third figure, we take an arbitrary large degree polynomial and see that we require G5 I 10-10 10 for SGD not to be effected by exploding gradients. As a result, the dynamics fail to progress within the specified number of iterations. In contrast, SGD-U is able to recover the true direction with the same step-size in all our examples within the same timescale.
provide simulations for online SGD and SGD-U with varying choices of and
figures that both algorithms
3.2.2 DIFFUSIVE LIMITS AT FIXED POINTS
Using our results, one also obtain diffusive limits around appropriate fixed points for rescaled observables, similar to Section 3.1.I We illustrate this by considering the diffusive limit for phase retrieval at m = 0 (i.e. initialization).
Proposition 3.10. For the link function fxx) = x²² The dynamics for the rescaled statistics ü. (1) (m,r²) = (Vrm, r²) converges to the solution of,
12(r2 1)
4r
dm
md +
co(15r - -182 + 15
dd)
1- 3
1- 3
8
20a(15r* -6r2 +3+ o2-] dt
2r
37 1+
(14)
1 1
3
1-
The proof is provided in Appendix D.3 We note that the evolution of x² is independent of m. Furthermore, the equator m 0 unstable when y < 1, providing similar insight to the ballistic phase where we observed the emergence of stable fixed points away from the equator for sufficiently small es (which indirectly forces 2 < at such point(s)). One could also obtain a similar diffusive limit for the SGD-U dynamics obtained in Proposition 3.9 however as the fixed point is unique and stable we find that these dynamics may not be as informative.
4 CONCLUDING REMARKS
We develop a general theory for the high-dimensional scaling limits of online SGD with Polyak momentum and scalar Markovian pre-conditioning under the effective dynamics framework of Ben Arous et al. (2024). We illustrate our results on two canonical models in nigh-dimensional inference: Spiked Tensor PCA and Single Index models. For both problems, consider both the ballistic and diffusive limits SGD-M and SGD-U, with the latter setting the preconditioner as the inverse gradient norm. In doing so, we provide rigourous examples of how SGI--U is able to stabi- lize training dynamics and high-dimensional effects leading to more desirable training outcomes as observed empirically. In particular, we illustrate how the effective dynamics framework can sharply capture empirical phenomenon such as apploding/vanishing gradients in the limiting dynamics.
Given these results, there are several interesting directions explore for future work. Of particular interest is the problem of developing limits of more general scalar preconditioners and develop a systematic analysis comparing and contrasting preconditoning and momentum based methods on broader classes of statistical problems. It would also be interesting to determine the limiting dynamics of popular matrix-based preconditioners (Duchi et al., 2011 Hinton 2012 Kingma & Ba 2014) and whether they mitigate the high digh-diimnsional corrections.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
A.J. acknowledges the support the Natural Sciences and Engineering Research Council of Canada (NSERC), the Canada Research Chairs program,the Canadian Foundation for Innovation- John Evan's Leaders fund, and the Ontario Research Fund. TJ and V.S. acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC). Cette recherche a été enterprise grâce, en partie, au soutien financier du Conseil de Recherches en Sciences Naturelles et en Génie du Canada (CRSNG), du Programme des chaires de recherche du Canada, et du Foundation canadienne pour I' l'innovation FLJE, et les Fonds pour la recherce en Ontario. 1RGP1N-2020-04597, 1BECC-220000019,9CC-000000000000
Project 43994]
REFERENCES
Theodore W Anderson. The integral of symmetric unimodal function over a symmetric convex set and some probability inequalities. Proceedings of the American Mathematical Society, 6(2): 170-176, 1955.
Anas Barakat and Pascal Bianchi. Convergence and dynamical behavior of the adam algorithm for nonconvex stochastic optimization. SIAM Journal Optimization, 31 31(1):244-274, 2021.
Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions in high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, 116(12):5451-4400 2019.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. Advances in neural information processing systems, 25:25349-25362 2022.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for SGD: Effective dynamics and critical scaling. Communications on Pure and Applied Mathemat- ics, 77(3):2030-20000 2024.
Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, vol- ume 4. Springer, 2006.
Xiangyi Chen, Steven Z Wu, and Mingyi Hong. Understanding gradient clipping in private sgd: A geometric perspective. Advances in Neural Information Processing Systems, 33: 33:13773-13782, 2020.
Elizabeth Collins-Woodfin, Courtney Paquette, Elliot Paquette, and Inbar Seroussi. Hitting the high-dimensional notes: An ode for sgd learning dynamics on glms and multi-index models. Information and Inference: Journal of the IMA, 13(4)iaae0228 2024.
Aditya Cowsik, Tankut Can, and Paolo Glorioso. Flatter, faster: scaling momentum for optimal speedup of sgd. arXiv preprint arXiv:2210.164000 2022.
André Belotto da Silva and Maxime Gazeau. A general system of differential equations model first order adaptive algorithms (2018). URL https://aax org/abs/ 1otiis////11..333 13 1(3):9 2018.
Thanh Dang, Melih Barsbey, AKM Sonet, Mert Gurbuzbalaban, Umut Simsckli, and Lingjiong Zhu. Algorithmic stability of stochastic gradient descent with momentum under heavy-tailed noise. arXiv preprint w:Xi::2502.00885. 2025.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Yasong Feng, Yifan Jiang, Tianyu Wang, and Zhiliang Ying. The anytime convergence of stochas- tic gradient descent with momentum: From a continuous-time perspective. arXiv preprint urXiv:2310.19590 2023.
Damien Ferbach, Katie Everett, Gauthier Gidel, Elliot Paquette, and Courtney Paquette. Dimension- adapted momentum outscales sgd. arXiv preprint arXiv:2505. 16098, 2025.

---

Published as a conference paper at ICLR 2026
Benjamin Gess and Sebastian Kassing. Exponential convergence rates for momentum stochastic gradient descent in the overparametrized setting. arXiv e-prints, pp. arXiv-2302, 2023.
Avrajit Ghosh, He Lyu, Xitong Zhang, and Rongrong Wang. Implicit regularization in heavy-ball momentum accelerated stochastic gradient descent. arXiv preprini arXiv:2302.0088) 2023.
Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dy- namics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019.
Trevor Hastie, Robert Tibshirani, Jerome Friedman, et al. The elements of statistical learning, 2009.
Geoffrey Hinton. imsprop: Divide the gradient by running average of its recent mag- nitude, 2012. URL https:// http///ww.ccc www.cs. toronto.eda/ ttt lecture_s s1 ides secc.pcf
wttp://tt.rrrttoeeeeee/ee//e///a idds
lain M Johnstone. On the distribution of the largest eigenvalue in principal components analysis.
The Amnals of statistics, 29(2):295-327, 2001.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Anastasia Koloskova, Hadrien Hendrikx, and Sebastian U Stich. Revisiting gradient clipping: Stochastic bias and tight convergence guarantees. In International Conference on Machine Learn- ing. pp. 17343-17363. PMLR, 2023.
Nikola B Kovachki and Andrew M Stuart. Continuous time analysis of momentum methods. Journal of Machine Learning Research, 22(17):1-40, 2021.
Harold J Kushner. Asymptotic behavior of stochastic approximation and large deviations. transactions on automatic control, 29(1 1):984 990, 1984.
1EEE
Chris Junchi Li, Zhaoran Wang, and Han Liu. Online ica: Understanding global dynamics of non- convex optimization via diffusion processes. Advances in Neural Information Processing Systems, 29, 2016.
Qianxiao Li, Cheng Tai, et al. Stochastic modified equations and adaptive stochastic gradient algo- rithms. In International Conference Machine Learning, pp. 2101-2110. PMLR, 2017.
Qiang Liu, Jason Lee, and Michael Jordan. A kernelized stein discrepancy for goodness-of-fe tests. In International conference machine learning, pp. 276 276-284. PMLR, 2016.
Tianyi Liu, Zhehui Chen, Enlu Zhou, and Tuo Zhao. A diffusion approximation theory momen- tum stochastic gradient descent in nonconvex optimization. Stochastic Systems, 11(4):307-323, 2021.
Chao Ma, Lei Wu, al. A qualitative study of the dynamic behavior for adaptive gradient algo- rithms. In Mathematical and scientific machine learning, pp. 671 PMLR, 2022.
Jan R Magnus. The exact moments of ratio of quadratic forms normal variables. Annales d' E'conomii et de Statistique, pp. 95-109, 1986.
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms. Advances in Neural Information Processing Systems, 35: 7697-7711, 2022.
Noah Marshall, Ke Liang Xiao, Atish Agarwala, and Elliot Paquette. To clip or not to clip: the dynamics sgd with gradient clipping high-dimensions. arXiv preprint arXiv:2406. 11733, 2024.
DL McLeish. Functional and random central limit theorems for the robbins munro process. Journal of Applied Probability, 13(1):148-15) 1976.
12

---

Published as a conference paper at ICLR 2026
Tomáš Mikolov. STATISTICAL LANGUAGE MODELS BASED ON NEURAL NETWORKS. Ph.d. thesis, Brno University of Technology. Faculty of Information Technology, 2012. URL https:
www. fit vut. cz / tattuddypppddthsiissese
Andrea Montanari and Emile Richard. statistical model for tensor pea. Advances neural information processing systems, 27, 2014.
Y Nesterov. A method of solving convex programming problem with convergence rate O(1/k²). Proceedings of the USSR Academy of Sciences, 269:3, 1983.
Courtney Paquette and Elliot Paquette. Dynamics of stochastic momentum methods on large-scale, quadratic models. Advances in Neural Information Processing Systems, 34:9229-9240, 2021.
Courtney Paquette, Kiwon Lee, Fabian Pedregosa, and Elliot Paquette. Sgd in the large: Average- analysis, asymptotics, and stepsize criticality, In Conference on Learning Theory, pp. 3548- 3626. PMLR, 2021.
Courtney Paquette, Elliot Paquette, Ben Adlam, and Jeffrey Pennington. Homogenization of sgd in high-dimensions: Exact dynamics and generalization properties. Mathematical Programming, pp. 1-90, 2024.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pp. 1310-1318. Pmlr, 2013.
Sandrine Péché. The largest eigenvalue of small rank perturbations of hermitian random matrices. Probability Theory and Related Fields, 134(1):127-173, 2006.
Boris Polyak. Some methods speeding up the convergence of iteration methods. Ussr compu- tational mathematics and mathematical physics, 4(5):1-17, 1964.
Parsa Rangriz. High-dimensional scaling limits of online stochastic gradient descent single-index models. University of Waterloo Thesis, 2025.
Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. arXiv preprint arXiv:1904 100.009790 2019.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathemati- cal statistics, pp. 400-407, 1951.
David Saad and Sara Solla. Dynamies on-line gradient descent learning for multilayer neural networks. Advances in neural information processing systems, 8, 1995a.
David Saad and Sara A Solla. On-line learning in soft committee machines. Physical Review E, 52 (4):4225, 1995b.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604. PMLR, 2018.
Joram Soch, TJ Faulkenberry, K Petrykowski, and Allefeld. Shaprrofbook/saappoooooook github. io: Statproofbook 2023, 2023.
Karl R Stromberg. An introduction to classical real analysis, volume 376. American Mathematical Soc., 2015.
Daniel W. Stroock and S. R. Srinivasa Varadhan. Multidimensionall Diffiusion Processes. Classics in Mathematics. Springer- Verlag, Berlin, 2006. Reprint of the 1997 edition.
Weijie Su, Stephen Boyd, and Emmanuel J Candes. A differential equation for modeling nesterov's 's accelerated gradient method: Theory and insights. Journal of Machine Learning Research, 17 (153): 2016
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International conference on machine learning. pp- 1139-1147. pmlr, 2013.
13

---

Published as a conference paper at ICLR 2026
Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase diagram of stochastic gradient descent in nigh-dimensional two-layer neural networks. Advances in Neural Information Processing Systems, 02-22224422255 2022.
Roman Vershynin. Hligh-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
Chuang Wang, Jonathan Mattingly, and Yue M Lu. Scaling limit: Exact and tractable analysis of online learning algorithms with applications to regularized regression and pea. arXiv preprint arXiv:1712.04332, 2017.
Jun-Kun Wang, Chi-Heng Lin, and Jacob Abemnethy. Escaping saddle points faster with stochastic momentum. arXiv preprint wrXiv:2106.0230 2021.
Ashia C Wilson, Ben Recht, and Michael Jordan. A Iyapunov analysis of accelerated methods in optimization. Journal of Machine Learning Research, 22(113):1-34 2021.
Adams Wei Yu, Lei Huang, Qihang Lin, Ruslan Salakhutdinov, and Jaime Carbonell, Block- normalized gradient method: An empirical study for training deep neural network. arXiv preprint arXiv: (707,0007 2017.
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv:1212.5701, 2012.
Bohang Zhang, Jikai Jin, Cong Fang, and Liwei Wang. Improved analysis of clipping algorithms for non-convex optimization. Advances in Neural Information Processing Systems, 33:1551 1-15521, 2020a.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 13:15333-1539,, 2020b.
14

---

Published as a conference paper at ICLR 2026
A GENERALIZATION OF THEOREM 2.3 FOR PRE-CONDITIONING
Here, we provide a generalized theorem (with adequate conditions) that allows one to prove similar scaling limits in the setting where we adopt scalar Markovian preconditioners for SGD.
Consider the following online learning setup. Suppose we are given an i.i.d. data sequence yı, 9/,..........................9 ... taking values in dR with law P. MM (R2)) and loss L : Xr 3n, + I where Xr C R is the parameter space. Consider online SGD with learning rate on > 0. momen- tum rate [0, 1) and the scalar and Markov preconditioner (I. X, x Vn + R given by:
Pe = 8p:-1 - 8n Th (2 ye)
æt -4-1 + Pt
with possibly random initialization tion of finite collection of summary statistics un (x) = (u)..)).... (x)) the regime where both Ps and dn may grow with and 6n - 0 as n oc.
pn E Mi (Xn). Our interest is understanding the evolu-
We define the functions
(a Y) nn Y) LI, (2, (②(x) where VZ(x) EEn, ()VLn(,,)) (15)
Let (xx) = (x) & v=)( be the covariance matrix for V at æ.
We state two assumptions: (on. -ocalliiizayil and asymptotic-olooabbiiyy
Definition A.1. A tuple (un, Ln. Pn, Tn. 3) is (6n, nh Hocalizaze with localizing sequence (Ex) K there is an exhaustion by compacts (Ex) K of R and constants Cк (independent of n) such that
1. sup SuPnenn (Fx) / llop S CK 5/2 and sup zEM (EE) 3
u. llop S Ck:
2. suppeu (Ex)
V < Ck, and supzeu, (Ex) B B < Ckò (Ex) VI, Vuf) Ck and
3. sup (Ex)
supreu, (Ex)
E
V u (x), V y) @ (H()))) v(z)) = ( and
SuP(aa)feu77 Eg va (V (m)), v ( @ v (=2, | = o(o ) for , > 0
The second assumption ensures that the limiting coefficients for the evolution of the statistics close. To this end, we define the first and second-order differential operators,
Å = (V and É. 2 (V, V2
Definition A.2. A family of summary statistics (u,) are asymptotically closable for learning and momentum rates ((6m), 3), and preconditioners (n,), if (u,, V Pn, mn) (6 nm)- localizable with localizing sequence (EK),, and furthermore there exist locally Lipschitz functions h R x R + I2 and Σ I + R kxk such that
8n Ln ( An + (1 (1-9p_n
un (n(x)) - (3, (x)) +
sup #Eu, (Ex)
[1 (x))()))]
sup
zEu, (Ex)
with Jn the Jacobian of u,. We call h the effective drift and the effective volatility of u respec- tively.
Theorem A.3. Let (Xxx). )r be SGD initialized from Xo ~ yh for un E M, (IPP)) with learning rate 5n, momentum rate / E [0, 1) and pre-conditioner 7n (x. y) for the loss Ln(: -) and data distribution Pa. For a family of summary statisties u. (uf))=== let (un(()); be the linear interpolation of
(un, 8x [16,]] (()=))))e.
15

---

Published as a conference paper at ICLR 2026
Suppose that un asymptotically closable with learning rate 8,, momentum rate 3. precondition- ers (Nn),. effective drift h. and effective volatility Σ, and that the pushforward of the initial data has (un), Mr - weakly for some v E A11 (R*)) Then (u, (t)), - (ut), weakly as - co,
where u, solves:
du, h(3 u.) dt +
Σ uu dB,
B
initialized from V, where B, is standard Brownian motion in R
This result is essentially identical to Theorem 2.3 under the revised constructions of V VH, to 4 and V.. That is, by augmenting the loss gradient with the preconditioner to obtain the coupled gradient 1/,, LI, we can leverage the same mechanics as those for SGD-M to immediately obtain scaling limits for SGD-M equipped with Markovian scalar preconditioners. To understand why this holds, we note that V L, is treated as an arbitrary vector field satisfying specific regularity conditions in the proof of Theorem 2.3 The fact that L,, is the gradient is irrelevant, and principle can be replaced with any arbitrary estimator satisfying Definitions A.1 A.2 In this case, we replace L with the "estimator" v
B PROOF OF THEOREM 2.3
Our goal is establish the weak convergence u. + random variables the space of continu- ous functions C((0, oo)), where u solves the limiting SDE. It suffices to show the same on C[O,, 7)), endowed with the uniform (i.e. sup-norm) topology, for every T > 0. The proof proceeds by structing an auxiliary process Zn, defined in equation I6 such that
sup Blllu, z, II] 0 EEER
Let T denote the exit time of the interpolated process (t) from the compact set Ex. Further- more, let Ex,,, u (ERR denote the corresponding pre-image, and LR (EE the associated space of localized bounded functions. The subsequent analysis relies on several technical lemmas. which hold under the sysmptotic-closabee
and 5, -o-looaaizability assumptions detailed in We assume Po = unless stated otherwise.
Definition 2.1 2.2
To simplify notation, we use the subscript notation fy := f(x,) for quantities that depend on æ,. In particular, we write 9j := g(x,,, Y5+1) for functions that further depend on the random data Y Finally, we say that x S y if there some constant > such that f S Cg and that f Sa g is there is some constant C(a)) 0 depending only such that f S C(a)g.
Lemma B.1. For any integers 0
(-1, the momentum variable satisfies the recurrence
j-i-1
Pi-1-i = 3 P:ppppp--0 Pi-1-5 8
3th
LLVLL-2-i-m-
Lemma B.2. For any 0 Usise- 1,, 1, the position variable admits the representation
i-1
1 3m+1
g 2-2
Ct- = 81-1-i
2-m 1-3
Σ
g
Le-2-ii--
Lemma B.3. For any 0 i 1, the Hessian difference is bounded uniformly in time in L by
\V ( V=u((ze-1-i)lop O(ivo).
Lemma B.4. For any Sis 1, the gradient admits the expansion
g
- -u(-V 8
1 8 E-2-i
gm+ 3ot
VH (0 +O(P3),
3
where the trailing term holds uniformly in time in 1,1
16

---

Published as a conference paper at ICLR 2026
We proceed by analyzing the process coordinate-wise Let v denote the j-th coordinate function of u, and we again adopt the notation ue == u(xt). For any I TK/ /ö, a second-order Taylor expansion yields
1
u(æz) u( + pi) ue-1 + (pe, Vuc-1) + (pr @ pe, Viug-) + Re.
2
where by 6-1ocalizability, the remainder term is bounded in L by
sup xE
E [[R]]] S
sup E ZEER,
|
||pel
)
3
S 100
Σ $20
B'IV0lL + +IV 1x
SK 8312.
This bound holds uniformly for é S TR/6. Substituting the definition of pa and rearranging the expansion provides the explicit form of the increment:
u1 14-1
8
8 V
(V. V
8²
8-1
+
guts (4 8 Pulee1
4
+ +
Σ
gits V 1 VH Pun)e-
i==
82 1-1
+ 2
Σ
8
8 VVL.
Σ
Here, we employ the compact subscript notation (a, b)g (az, bz), {a- bie (uz-i, be) and (a_. = (ac-s @ ct-s, bt), assuming that the indexed quantities at - i, - are all
well-defined.
The next several lemmas summarize how we construct the auxillary process n, by identifying the terms in the expansion contributing the drift and quadratic variation of the limiting SDE. We show that all other terms in the expansion vanish in L
Lemma B.5. The evolution of the j-th coordinate of u, which denote u, and Vul for the spective gradient, admits the decomposition:
8²
uy(ar)) uz(w 5
3
(Vor
Vul
U.5
estice
ΣΣ PESICE
where the remainder term o(1) T.1, uniformly for l S / In particular, we have that the first two series contribute to the limiting drift and the final term contributes to the limiting quadratic variation.
Lemma B.6. The martingale component, i.c. the third series in the expansion due to Lemma B.5 admits the asymptotic expansion
[8/8| é-1
8/8
Σ
Σ
H Vuje-1
1
E=
17

---

Published as a conference paper at ICLR 2026
1t/6
g 1 B 1- (V, 2
8²
where remainder is o(1) L', uniformly for t S T
This final lemma obtains a similar asymptotic representation for the remaining two series in
Lemma B.5
Lemma B.7. The following asymptotic equivalences hold in Lt, uniformly for & & T 12
1
('V Vuje-1
V4, 11)e-++000))
1-
1 1-3 (V.-i, v 2(1 - 3)2 (V,
84
1
g
2
2
By consolidating Lemmas
B.5
B.6 and B.7 construct the auxiliary process nn follows.
Consider the j-th component z", the discrete-ting process is defined by the Doob decomposition:
0
(A 1 +
(M -M
(16)
k=1
where the previsible increments (AA 1 and martingale differences (M. ME-1 are defined as
1 An (1 3)2
8
AV AW =
Lan u1
1
M - Mi- = 1- (V Vuje-1
we write 2n as the corresponding continuous time interpolated process of equation [6 then the difference |/u,, 2n || is bounded by terms which vanish in L uniformly in time. In particular, have that this construction satisfies the requisite limit:
sup ((nn(s) zn (s)l + in L.
0SASTR
The remainder of the proof closely follows (Ben Arous et al. 2024 Theorem 2.3); we provide the details for completeness. Let (s) be the linear interpolation of the discrete process (z[s/a) given in equation 16 We decompose this process
(s) zn (0) + a.))) + (s),
where (s) and (s) are the continuous-time interpolations the cumulative predictable and martingale parts, respectively. More precisely, define for E [),
ds(s) = A A and 8,(s) = M7ro Mi0/8|-1
and write the j-th coordinate of a, (s) by
ag(s) [ ds (s') ds' = (5(5)//)) 5(z/5]) (A///8 A/618.-1)
and similarly for by(s) 1o ds')
We now establish tightness for the sequence of stopped processes (zn(s A Th (esooT in C((0, TT) for each K. By Kolmogorov's continuity criterion, it suffices to verify the moment condition: for all 0 < T.
E Zz, (t ((ATK) Zn (s A 1 SKT (t s)
This implics that the limit points are (1/4)-Hölder continuous. We control the drift and martingale components separately. Fix j and let Uj a ag, b by to simplify subsequent notation.
18

---

Published as a conference paper at ICLR 2026
For the previsible (drift) term, we have
6
82
E -
SE
Σ
An
(1 3)2 Ln ur-1
1
4
S E Σ
hn(B,uunee- un)e-
+ +ol(1 - -8)²)
S (t s)* (11llit + +0(1) EK (t - 5)*
The summation ranges over l from [8/5] + to [4/8], stopped at / The bound follows from the continuity h as stipulated by asymptotic closability, i.e. Definition 2.2
For the martingale term, applying the Burkholder- -aaiiisde---viis---ie Davii-Guudy inequality yields
( (E 1(M--MEL)) S E
2 Σ(M Mi-1)²
(( R = (b((AT))) -
)2
Substituting the martingale increment and applying Cauchy-Schwarz yields,
27
82
Σ (VHe-,, (VHooI.Wu.-1)() 3)2
SB
8t
Vu Vu ( V1r
Σ
SK (t s)- s)²
where the final inequality follows by condition (3) of o,.-localizability.
Since both terms are O ((t - ()²) for 0 s,t S T. we apply Kolmogorov's continuity theorem to conclude that the processes (z, ( })s are uniformly (1/4)-Hölder continuous, thus establishing tightness. Consequently, the sequence martingale components (b,, A ()) also tight, and its limit points are continuous martingales.
We conclude the proof by showing that the limit points are solutions to the limiting SDE by the martingale problem. Let zKK (t) (1 A TR). and define a (l) and b (t) analogously. Denote their respective limits by z K (t), a (t), and (t)..
We first compute the limiting quadratic variation of bK (t), and thus = xx(t) (t). Let AM; := M M and notice that for 1 < i.j < k, the limiting predictable quadratic covariation is given by the integral
[. SE AM
1' (Vu, VVus) ds AM [6/8)ATK] ds (1 31-
(/60)AK
By asymptoticc-closabiityy we ensure that 0C.
[
8 ((Vu, VVuj) - Enj (s)) ds
sup EST
ST T sup (V, VVu)) (x) Eig (z, (x)) aEFk
goes zero as - oc. Thus, the angle bracket of the limiting continuous martingale bk ((()) is by definition,
(b*), - f Σ [z (s)) ds
By ItO's formula for continuous semimartingales (Stroock & Varadhan 2006 Theorem 6.3.4), for any f E Og (R*), the process f(z - So ( Jds is a martingale, where L is the infinitesimal generator
1 k Σ
L
Ei E.y + Σ h
W=
19

---

Published as a conference paper at ICLR 2026
Given the assumption that h and v V are locally Lipschitz (and thus Lipschitz continuous Ex), the martingale problem associated with is well-posed. This uniquely characterizes the limit xx as the solution to the SDE
E(u(t)) dB, du(t) = h(3 u(t) dt + I- 3
stopped at the exit time тк- By employing a standard localization argument (see, e.g. (Stroock & Varadhan 2006. Chapter 11)), conclude that every limit point =(t) of zt) is solution to the limiting SDE.
B. PROOF LEMMAS B. B.3
In this section, prove the preparatory lemmas that characterize the behaviour of the momentum variable ps as well as the summary statistic un. In subsequent subsections we treat each of the lemmas dedicated constructing the auxillary process nn separately.
Proof of Lemma B.
The momentum variable evolves according to the recurrence pt = 8p:-1 8V Lt-1 Iterating this relation from time t down time < yields
P: = B p
VLj VLj.
8=8
Reindexing the summation via m = - - j. we obtain the equivalent representation
Pt g - Σ
3 V L
1-
To recover the first claim, we set -1-1-1 and 1-11-1 j. which satisfies t s and t t- = for S
Proof of Lemma B.2
We expand the position variable recursively: æz-1 x0 + e-1 ing the expression for P& derived from Lemma B.
Pa. Utiliz-
we have
30
ΣΣ
Σ
3"h
m==0
Interchanging the order of summation and evaluating the finite geometric series yields,
f-2
8-2
1 -
3mm+
1:-1 20
8
Σ
g
V 2-m
-
Σ
V Li 2-m
m=0
m=0
To establish the relation between 07-1 and F:-1-i we partition the summation:
1
gm+
6-2
1 - gm++1
300
6
V Vee
Σ
LB 1--2-m
(17)
1- 3
m==
1=
We reindex the final summation by setting m
gm++1
2-2-1
1 - gh+i+1
V
L
B
1 1-3 k=0
Decomposing the numerator gh+it1 -1- (3 gh+itl), recognize that the first component, combined with Vo. constitutes Øz-1-i- Substituting this identity back into the expression for #2-1 and consolidating the remaining terms yields the desired result.
n
20

---

Published as a conference paper at ICLR 2026
Proof of Lemma B.3 By the integral form of the Mean Value Theorem, have
L. V (1 + 1)) [ - 24-1-1] dt.
Viu(ae-1) (
Taking the operator norm (with respect to the tensor V and bounding the integral on the right- hand side, obtain
[[Ven((x=1) Viu(mo-1-iillo x sup
| (tx0-1 (1- 1))x)-)))) |00-1 26-1-il
(E[0.1]
< sup
V | x2-1-ill
=E
By 8, o.-llocaliaabiitty, U/oppp Cx and by Lemma B.2 the displacement vector (xe-1 1e-1-i) satisfies
1-2-i Sivo
xt-1 112--IT-- 1 So 8[[V + [H
i ++2
Σ m=0
gm+1
which is quantity of order O(iv 0) in I uniformly in time. Consequently, combining both bounds gives the desired I,1 norm bound for the Hessian difference.
Proof of Lemma B.4 Take the first-order Taylor expansion of Vu(x;(1) around 04-1-1
Vu((x-1) Vu( + V=u(z(-1-))= x6-1-1] Rr
where the remainder RG satisfies T
|| Rilll,, S sup =EER,
| (V')ll) | =
=
using the bounds derived for the displacement vector in Lemma B.3 and 6-localizability.
Substituting the expression for the displacement vector using Lemma B.2
gm++
Vu(ax-1) Vu(
V
1 3
2-
+
3 gm+
2
0(20).
m=
We proceed by aligning the arguments of the Hessian with those of the gradients with respect to æ. For the first summation, Lemma B.3 to bound the error introduced by realigning indices get
gm+
Pu(BA-1-i 2
1
YT 0
i-1
gm++1
8
Vollaaa-mmmmlaaammm-- //1111/111,,... .iv8)
m=0
1
gm+1
8
m=0
where the additional error term in the last line follows by expanding VL V VH and bounding the series involving VΦ by ò-localizability The remainder term can be further bounded O V8) in L! similar analysis applies to the second summation which we may write as
2-2-
1-2-1
gm+ 3m --
m=0
Substituting the expansions and noting that the error terms are dominated by (6) completes the proof.
21

---

Published as a conference paper at ICLR 2026
B.2 PROOF OF LEMMA B.5
Proof of Lemma B.5 The proof proceeds by demonstrating that the extraneous second-order terms in the expansion of u (zt), when accumulated over the trajectory, vanish usymptotically in L'.
First, by -o-aliaiizability. the contribution at step & for the deterministic quadratic term is bounded in L by
8
Σ
3 (V & Voi-1-j, Vue--1 SK 83/2
2
Summing over t 1,..., [1/0] yields total contribution of 0(vō) o(1)
For terms involving the fluctuations VH, use Lemma B.3 to align the indices of the Hessian V with the fluctuation term. For example, we write the cross term as
8
Σ
OV. 8²
Σ
&
R2
1,==
1.j=0
The error R( arises from the Hessian difference over j steps and is of the order
R RRI | O (00/2 1, -
O(6²)
uniformly in 1 Similarly, for the centered quadratic fluctuation term:
8
git
& VH VLL
V|H,), V H &
2
+
82
8-1
VH Vuu.- + o(8)
(18)
2 +,=0
where the error follows by item (3) of 6-localizability. Accumulating the emrors over the trajectory yields o(1) remannder
remains demonstrate that the accumulated aligned stochastic series vanish in L', We show this by proving the stronger condition that these terms vanish in L We first introduce the notation:
Pij 44 & VH,, Puy Hig = WH & v H1 V. @ VH2J, V
Now consider equation
18 which is further accumulated over & = 1,... L1/ō] steps and set T = [t/8]. In particular, we want bound the term
T
SP. : 82
g P
82
Pk.m
Σ
g 1)
0
k.m=0
(=
where the last equality follows by reindexing with f-1 and f- 1-j and interchanging the order of summation. The inner geometric series is uniformly bounded by
32 max(k
3lk-ml
32(t-1)-k--1
B2
32
By Lemma B.8 follows that the 1-2norm vanishes: E[(S{)]] = o(1), uniformly in An identical argument applies the series involving Ha.m since write
Le/81
10/8
[
8²
Σ
git H
82
Σ
Has
Σ
B2f
(=max(+,))+1
thus completing the proof.
22

---

Published as a conference paper at ICLR 2026
B.3 PROOF OF LEMMA B.6
Proof. We first analyze the inner series by applying the gradient expansion provided by Lemma B.4 Vug-1, centered l i.
Σ B ( Vuje- Σ Vu) S. + 0(8
where the terms s and s
are the second-order components:
- gm+1
s(1)
VH.. & VH. Vuum-mee-2
1-
t=1
m=0
s((t
8²
8-1
+
B
8-2-i
Σ
gm+l
VV. & VH. V vu.m)
We decompose these terms into centered fluctuations and deterministic correctors. In Sp the diagonal contribution occurs when - we note that E H, @ VH,] for + j. This yields the decomposition:
m. i.e., m = -- 1. For off- off-diagonal terms,
8 (V. 2 u---1-i
5 0 = 82
Σ
1 1-3
1=1
1 3m+1
8
Σ R' Σ
1-8 VVH VH -H---EEVH.... & VH. Vu..mie
From here, all of the remaining terms constitute fluctuations, which we show vanish in I after accumulating over e = 1, . [t/8] as in Lemma B.5 Adopting the same notation H Ha,m, reindex with n - i, m = j. and reorganize the first summation to obtain
1-6--1i- ΣΣΣ
1 1- 30+1 3-+1
2 [t/8]-2
(t/6)-1 Σ (=m+1
ge-1
g
H -
Σ
Σ
Hn,m
8
1 B
3
f=1
For m n, the inner geometric series is bounded by
[1/6-1
8r 86-1-m 1-m
gTh (1- 3)2
8 8m-n+1
Σ
3)(1 32) S m-
1-
where C3 0 is some constant that only depends B. Setting w(n, m) 0 for apply
Lemma B.8 to conclude that the -nom vanishes as d + 0. An identical argument applies to s as we can rewrite the series as
[6/6]
1 + 8 2-2-1
8
Σ
8-He-1-1--2----
[t/6] - [t/8]
6/8
1 14-18-111 1-n
H
Σ
8
3
m==0
1=
(=
To conclude the proof, we rewrite the leading first-order term by first reindexing the summation by = -(-1
j
10/8| t-1
[0/8|
BL1/61-5
(VH,
Σ
VH,Vu)
3
t=1
1=0
[4/8)
8
H, - 0
1- 3
t=1
Putting all the terms together shows that the series converges to a rescaled martingale sum as desired.
23

---

Published as a conference paper at ICLR 2026
B.4 PROOF OF LEMMA B.7
Proof. We analyze the asymptotic behavior of the first-order (drift) and second-order (diffusion) components separately. Let 7 L4/5). Starting with the second-order drift, align the Hessian with the covariance matrix by Lemma B.3 to write
82
L8/6J
é-1
1 8r
821
+
28
(Vi-1-1 v
2
1-
t=1
1=0
82
[0/61 t-1
821 + 28¹ - -387
Σ
Σ
(V. W..V2)/---ita
/ (1/5)
2
1 1-3
1=1
1=0
8²
11
[0/6] 1--
(1 Bitty2 (1 Bi)2
Σ
Σ
(V. (V. + +(1
2
(1 (1- -)2
1=0
1=0
[r/6]
82
Σ
(V. v -
2(1
2(1-3y
where the simplification the 3-weights follows since
(1 (1- = 23¹² (1 3) -1 (1-
3) [23 -1 3)]
For the first-order (drift) component, the analysis mirrors that of Lemma B.6 As reminder, apply Lemma B.4 to expand Vuz-1 around Vuy- yields
1- Vun-i)
1=1
6-1
8-1
i-1
gm++1
B
( Vu)e- 82
Pe-1-i.8-2-
8
1=1
1=1
m=0
3
2-2-1
8 1=1
P P 0
3
==
where the two series involving P,y vanish in 2 by Lemma B.8 following identical arguments as Lemma [.6
It remains to analyze the leading first-order term which mirrors how we handled the first-order martingale sum in Lemma B.6 Using the same arguments, we find that
(V4, Vul-1-1 1-3 V4, -
t=1 ==0
which completes the proof.
B.5 AUXILIAARY BOUND
We establish a technical lemma essential for controlling the accumulation of recond-order stochastic
emrors in Lemma B.8 B.6 B.7 We recall the following random variables representing the stochastic interactions:
Pay b, @ VH, V²u,), :y H & VH, T E @ VH,), V..).
We observe the following orthogonality properties derived from the martingale difference structure of V with respect to the filtration generated by the process.
If j + I, then E(P., PM] 0 by towering the expectation with XjvI. Furthermore, by ô-localizability we have the correlation bound
E Pml < sup V"
Sx 8-1 -0000))000000
- 0
Fx,n
24

---

Published as a conference paper at ICLR 2026
If max(i, ) # maxfk, 1), then HjH,,HHi)) Hal 0 by towering the expectation with Xivyvevi- Furthermore, by 5-localizability we have the correlation bound
Hal - (E[H)) ((Hl)) 1/2
S sup K.yEEk
(v @ VH(v), Fu(y)) S o(6
With this observation, we construct the following lemma that controls the correlation bound over an accumulated memory.
Lemma B.8. Let (X)., be a collection of r.v.s where take either X P X Hi. as defined above. Then for the series constructed by
S. = w(i,j)Xi.i
5 0
where w(i.j) are deterministic weights such that (wci,))l I Cg.9l-vil for any i.j with Cp an absolute constant that only depends on B, then the 1,2 norm of Ss is bounded by
E E Sa.K o(ŌN)
In partieular, if we take N O(8-1) then B[$3) o(1) as 5 when take X as either P or H.
Proof. We first note that the decay condition on the weights w((, imply uniform bound over single-index sums, That is,
sup
wiiij)) S
sup Σ Cagli-il
1+ 3 <0 1- 3
and similarly for the uniform bound over index i. We now determine the -o-mo,
B[ 81
Σ
u(i,))u(t,()))x,,
A.
we set P. then the summation restricted to j = which yields bound
2
E E
84
((: 8 sup k.
E Pkr]
Σ
Σ (w(i,))) ()
(19)
Ei.
On the other hand, for X = H the summation is restricted to (i, j) = max(k,1) = so that,
E
8 Σ
w(i.j)w(k: (k, Ha.1
(a, =m max (.!) =m
< 8 sup [E, Har] E
Σ
Σ
[w
(20)
max(i.j) =m
where the inner series involving the geometric coefficients is bounded by
20
Σ [w(i, w(m, m)
[w) +
w(i,j)) Cp
1-3
max(FJ)
i=0
which is an absolute constant that only depends on 3. In either case, we see that both Equation equa- tion 19 and equation areare uniformly bounded by
B[S Ss 8''0(-3) o(8N)
where the o 3) error attained from the previous correlation bounds. Thus the L -norm vanishes for N < 0(1-1) steps desired.
n
25

---

Published as a conference paper at ICLR 2026
C PROOFS FOR TENSOR PCA
C.1 PROOF OF PROPOSITION 3.1
The first part of the proof pertaining SGD-M follows from Theorem 2.3 and Section of Ben Arous et al.] (2024) along with the verification of the final item of (3) in 8, -ocloiiiizayilityy (which verified by the same bounds the previous item). For example,
sup L [(Vu(()) (H, y1) & V v2)) zEER,
sup DE
E E & n²
for (m, 2). For our diffusive limits about O, we have V = so that the bound is trivial. Moving on to SGD-U, we split the proof into two parts, The proof follows from the preconditioning theorem provided in appendix A.3 Thus we simply must verify asymptotic closability and (6n, 7n)- localizability. We prove these two assumptions separately below starting with asymptotic closability.
C.I.1 PROOF OF ASYMPTOTIC CLOSABILITY
We check asymptotic closability starting with the pre-limits for V and V under the normalized gradient pre-conditioner. Recall that from (Ben Arous et al. 2024 Section 7), VL is Gaussian vector with mean Vh and covariance V given by
V 22km k- + 2k R2 mw + 2k (xx - me) = -2kmm * 2kR22-2
V = 4k I + 4k(k 1) 2k -Ax
prior to the action of the preconditioner. Furthermore, the distribution of V only depends on the summary statistics (m, r²).
Before proceeding, we compute the expected norm gradient vector,
|| V - +
[|| 2E VH) | tt
Substituting the values yields
D -| || tr(()
(4X² - 2 Rdk- 4h1 m R22--2) + (4n + 4k(A: 12R22-2) 2)
( 2k + RAK-2 Am R2 + 4(n (n + k - 12
We compute the first and second-order generators separately.
Lemma C.1. In the asymptotic limit where asymptotic expansion for the first-order drift for SGD-U,
and VL is defined as above, we have the
VL
vn V
V E
O(n
V
D
Proof. For h(x)) = x 1/2, we take the Taylor series centered at D
h(z)) h/ h(D2) h'(D2 (x D2) 1 "(
D222 1 11" D233 -D233
1
3
15
2D (-- D2) + 8D (-- D2
(r D213
D
6457
where & E (x, D attains the remainder in Lagrange form. Using 1[[ 2||² we write
VL
E
4Φ 1
EVL D2) + 8D - D2>
D
2D
T(1V)/)))
26

---

Published as a conference paper at ICLR 2026
15
E
E 12
64
L(1))
follows that the remainder term satisfies,
15
EIVL(IV
-D3y
SB
D
64
|
which is O(n 3) 3) by Lemma E.4. To conclude, we analyze the remaining polynomial
4Φ 1
[[ - D
203 EIVLLL
EI
DD
8D$
To do so, Stein's lemma: EI(Z F()) EE[VF(Z) for Gaussian vector Z N Ni, Σ). We bound each term individually by considering the respective numerators.
(( - D²)], set fz) || and center the variable:
E D$] E[V 4 (I|VI - D²)1 + 2 D
VE 2
which is of the normed order (with respect to the Frobenius norm) VIV VΦ S
[Vllop | 0(1).
For V L|| - D2)a), use f(=) = (1|=|2 c)2 so V 4[|==1² -c)z.
E VV = v D2) (I D2
= -VE
-V + v& V L11 SK n
with respect to the Frobenius norm since Var(I 2.133 n.
Putting the full expansion together, we conclude that
V
VΦ 1 0(1) +0)+-((1))=
V$ + +
E
VL
D
2D3
as desired.
Lemma C.2. In the asymptotic limit where n + 0o and VL is defined as above, we have the asymptotic expansion for the second-order drift for SGD-U,
Vn V D2 + O(n
Proof. The arguments for the proof follow those of Lemma
C.I
First, we write
V
V H
E
V
vi.
V
60 V
V
VL
V
VΦ
0 (n-
|
where the last term follows by the trailing O(n Subtracting both sides by the quantity V D/2
3/2)
error for V acting on / DD (n12)/2).
VL & VL
4 ④ VΦ
E E &
VL
VΦ &
VΦ
=
E
+ +0
D2
D2
D2
VL &
VL
V & V
= E
+ +O(n
|
D
27

---

Published as a conference paper at ICLR 2026
using the expansion V V X VL VΦ @ VΦ, To control the leading term, conduct
similar analysis as before using the expansion of the function h(y) = y around = D which yields,
1 1 (v D²)
1 (v -V- D6 (v DD)2
h(y)
D
D4
some & E [y, D2 to express remainder Lagrangian form. to the leading term,
Substituting this expansion back
VL VL
V VL
E
V
VIVL @ VL
E VL VL M([V]]]3)
D
V VL (1V D2)
E V
&
VL
(| D2y2]
+ R(E. VL)
D6
For the trailing remainder function R, one can show that
D
| - E
which is (n-² again by Lemma E.4
For the leading two terms, we use Stein's lemma for vector-
valued functions (e.g. (Liu al. 2016 Lemma 2.2)),
E[Xf(X) [ + VE[V
for Gaussian vector X ((, V) and Vf(X) is the Jacobian of the vector- function f.
For E V & VL (|V/||2 ) apply Stein's 's Lemma to write
V VL 2 Dd>1 V & V Dd) + VE[[|L]] D2)1
2VE 8 VL]
= 2V - 2 v & 4 + 0(1) 2V +0(1)
where obtained the quantity B[V D2) in the previous derivation 0(1)
and the other O(1) term is of the same order by symmetry. We note that the entire quantity is of the Frobenius norm order || O( following the spectrum of V.
For V L & VL (|V - D2 we again apply Stein's lemma to write E
E
VL (| D2>
V & [[V D2) VI D2
+ 8VE & VL (1/ D²)]
V. Var 8V O(vr) 0(m)
V. Var(I O(n)
where in the penultimate line we substitute the quantities from the previous step (and derivation VΦ). We again note that the entire quantity is of the order 3/2 with respect to the Frobenius norm on the variance quantity and spectrum of V.
Putting the full expansion together, we conclude that
V 1
V
12n + D6
-((3/2)+ - 0( -2)
+ O(n -3/2)
D
D
D
28

---

Published as a conference paper at ICLR 2026
C.1.2 PROOF OF (6n: 7-) LOCALIZABILITY
Proof. We verify (ōn, 7m Hocalizabbily for tensor pea with ū(x) vn(mm()) (x); r (x)). We note
that verification will then follow immediately for the cases u(x) = (v (z) r (x)) and u(x) = (m(x),r² (z)) or any centering of these statistics. We also note item (1) of (6n, 7n localizabiiity has already been verified for the no-preconditioning case (e.g. see (Ben Arous et al. 2024 Section
7)). We move on to (2):
(2)x)) +(LLvv))
sup mEN (FE)
| = sup
Vn
NEN (FR)
(2, m)
sup (Ba)
0(1/vn)
v / 2( + 1x=111
We now note the following bound [1
v = 2Amm k 1 + 2% |2 k||==||a~>20 22xx - mv)l
S |||=||26~2 + ||=|||-||
Applying this to the above we get:
CA | 2k- +
| Rー
sup V
vn
SC
ZEu (Ba)
sEu (E&)
n |1|100 2hh + (x))
vT
We now turn to the martingale component:
VL f)1l 1 + +1 +000))))+++000))))))))
VI v)
sup V S nd
sup (Ex)
E
xBt)
< C&n
We now tur to item (3) of 8, localizzzzbl
LVL(1, v) v((x), vū(
sup (VV Vū.()))
sup xEu (EE)
E
VVL w)1
#Eu- (EEN)
V ( y) VV v)1 vū.(z))
n²
sup
E(
zea- xEL)
< n
sup sEa ((E))
E(VL(z, v). (ū,()))
VL(x, y)l
Where in the second line, we use (a + $ +1 and Jensen's Inequality to remove the V term at the cost of constant factor. We now recall that VL(x, y) is n-dimensional Gaussian with mean O(1) and covariance D La + Arx' ,|>illop = =1),, |2]| = (n)). For both yrim and yūr we have that (V, v), (i,())) is Gaussian with variance bounded by C&n and by technical lemma for Gaussians we have that E v((x)))1l = O(1/n²) yielding:
E(vnVil( v), =0)))))==)(m)) -Ex
sup
As required. Lastly we consider:
sup E(VLu,()), (z & VH - 4-2 =Eu-
sup Etr((T - ( & vH(,,y) 5))²
(Eu
sup Etr(VH((, & V.(..)) V) Eur(ve (V & VM(m.y) VD)² (Ex)
We now bound the first term above noting that the second term above is bounded by (e.g. using the Von Neumann trace inequality) the first term multiplied by ||ev" = O(1):
tr(vil( v) & vil(z,)) v) - -V = ir (1
VVVL(x 4) - vQ(T) VL
82
LVL(( 8)
82
- E
VV
[lll
29

---

Published as a conference paper at ICLR 2026
VnVI(a, 0)02
VnVL(a, v)
& vo(r)) =
tr
2tr
( (Z, fi112
(x, x)ll
VnV L1x, 1)
VVV ((x,))
2tr
+ +
( w)1l
L w)ll
VnVI(x; v)
VnVL(z,)) [[V y)ll
2r
@ vQ(xz)
- 2EErr
@ VQ(x))
VL y)|l
Noting that tr( Vlzzg) [[
deterministically. We thus get the following bound:
VnVL(a, v)
sup (Ex)
Et (z, & (H(..)) V)
sup (Bx)
ERt
[VL 8 vo(x))
S 2 sup
[|V((1)]1
(BE)
$ C
where in the first line we use the fact (a - 6)2 S a² + 122 and Jensen's inequality. In the second line use Cauchy Schwartz to bound the trace (which is here equal to an inner product) and in the last line we use item (2) of (ōn: 7m loc-laiaiizability above. Which yields the desired:
max sup ((V"u(s), (Bx)
Vil =
We conclude this subsection by proving the limiting dynamics of Proposition 3.1 for SGD-M and SGD-U.
Proof of Proposition 3.7
Following Theorem 2.3 and the discussion following it, we find that the
dynamics of SGD-M follow those of online SGD up to the additional 3-dependent constants. With this, our analysis immediately follows from (Ben Arous et al. 2024 Section 7) where we take the mappings.
f,
(1 (1- (1-3)29
for 3 E (0, 1). For example, by (Ben Arous et al. 2024 Proposition 3.1), the limiting generators for online SGD are given by
fm 2 + 22 2m,
f,= 4r 2
and 9m = 0, g.- = R2k for the correctors. The dynamics of SGD-M immediately following from the rescaled mappings discussed.
We now turn our attention to SGD-U and recall that Vm = v and V, = 2(x - me) and note that
|ml] 1 and || < drifts
2||==||- Substituting all of these into Lemma C.1 yields the pre-limiting
Vn VΦ, Vm) O(n Onnn'I|Iml]]
fm
V tr(V)) + / fm + o(1)
Ank + 4k(k 11/22-11 (- 22km + 2 + o(1)
11
VT
(R20-2m - 11m--mm--1- o(1)
R o(1
f.r vn (V4, Vr 0(n 1|n- 2V o(1)
1
fr + o(1)
D
4, 2r2k222-2
2 V
+o(1)
1+ o(1)
2v o(()
30

---

Published as a conference paper at ICLR 2026
Under the scaling that -> oc, have the limiting functions
k-1
fm = V Avk
===22= =2 2
R
For the second-order generators, recall that 4 m
and 22p² 2(1 or and note that
|| vn. Substituting all of these into Lemma C.2 yields the limiting corrector term for r²,
Ca g,- = n 2D
Vr + O (n 3/2 112
4cgk R2k
71
2 + o(1)
+ +0(1) " 4k o(1)
4nk 4k 1)R2
To conclude, we determine the volatility matrix using established approximations
8JVJ ca nJ
DE + +n 3/2)) S c D2 O O(n (1-
WJJ
The error term is thus negligible as || = O(v by ò-localizability (i.c. the norms of the summary statistics satisfy |/uu;]| O(Vn)). Recalling the volatility matrix for Tensor PCA under online SGD, we have the following:
D2 D2 JVJ
T1
1
4k(k - 1)m² R2k-4 + 4kR2k-2
4k(k 1)m((R - m R2K- 4k(k 1) - -m)) R2 22)22k
4kR + o(1)
4h:(k: 1) 1)ml R? m)) R2
(k 1) +1 1)E(R
(k: 1) - - (k: 1)(R 5)3
where we set the direction cosine & = m/R. The volatility matrix evidently vanishes for e8/n.
C.2 PROOF OF PROPOSITIONS 3.2 AND 3.3
The proof of Proposition 3.2 follows directly from the of online SGD in Ben Arous et al. (2024) and Theorem 2.3 In particular the equivalence of SGD and SGD-M after appropriate time and step-size rescaling. Below prove Proposition 3.3
Proof. We first determine the fixed points for the case when h = 2. In particular, we solve the system
mvkR(A - R22) 0, 2vhr² R = Cs
When = 0, it follows that R² = 12 so that we obtain
2/3
2/3
3 C5
2
C6
c
2/2
2/2
On the other hand when m + 0, we have that Ri > which implies that
22
C8
m
C8
2V2X
2V2X
for fixed À > 0. In particular, the solutions exist if and only if
2/3 = Xe
8
which is the critical À for the case of preconditioning and k = 2. To treat the general k 2 3, the system under consideration reads
k-1
R² m, cs = 2R - VRr²
where rearranging the equation involving 2 yields,
2
C5
2VTRR111
31

---

Published as a conference paper at ICLR 2026
Introduce the direction cosine := m/R R 1] so that {R and r² = R² {2). This in tur allows us to isolate the equation above in terms R² by writing,
-
n 2v
In particular, from the first equation we see that
x R2
= 0 or R2th xg#=2
implying that all equilibria are obtained by solving the second equation together with either of the two conditions. Hereafter, we A = cs /s/O (2vk).
First, f = (i.e. m = 0), then we obtain
2/(k+1)
R² =2
C3
2//
On the other hand, for & + 0 we have
R
A
R = Ag 2gk-2
(1- = X
A
g-
This yields the equivalent condition with respect to & as,
A
F(g) = ())-))))))(11111 <²)* A
g [0, 1).
To conclude our analysis, we record the following facts about F:
1. F(O) = F(1) = and F({) for & (0,
2. has a unique critical point on (0, 1), which is global maximum at
(k- 2) + 1)
E (0.
(k: - 1) + 2)
This can be by considering the log F(E) that
(E)
(k - ()
2kg
F(E)
E
1
Solving for (g) = O will recover the critical point above.
3. The maximum value is
(6-2)(++1)
(k)= - 2)( (1: + 1))
(2h:)²
F Frax = F(g.)
(6- (6-2)(++1) +
((k - 1)( + 2))
Therefore, the solutions for F({) only exist so long as
((k: 1)( + 2)) (4-2)-+1) 1/(k+1)
A F(E) = X Fmax Akrn c6) =
2V
(k 2106+11)
(2k)* ((k: 2) + 1))
C.3 DIFFUSIVE LIMITS FOR TENSOR PCA: PROOF PROPOSITION
Proof. Applying Theorem 2.3 in light of 2.1 simply check asymptotic closability, noting that (ōn, Th Jooaaliizabbb was already verified for this setting in the proof of Proposition 3.1 We check asymptotic closability which also follows easily from the same proof. The drifts under the new coordinates are written as
V Vūn) =
n
(V Vm) + O(n 0(n-1/3) = VnR" 'm m A/n
m k-1
D
R
32

---

Published as a conference paper at ICLR 2026
k-1
= - + (ū
(27 (ū//n2)
k-1
11 An 7
i1
Vr2 + (ig /m)
(v,, Vr = 2r (r² + [ū] /n/
Taking limits as n + 00, as long as À is fixed in n, we obtain the rescaled drift
k-1
21
fu
k 2
für 2ūg 11/2
U2
U2
k 3 3
The volatility matrix (under the new coordinates) is given by the limit
JV.J = (A: () + 0
where the only surviving entry after the limit is E11 cs - 1) /īz) co.
D PROOFS FOR SINGLE-INDES MODELS
We start with SGD-M below. The limiting generators were first determined (Rangriz) 2025 The- orem 3.1, Corollary 3.2). We provide a slightly different derivation of the result for completeness. Refer to (Rangriz 2025 Lemma 3.5) to show that the problem satisfies the 6-localizability and closability assumptions they also apply to SGD-M.
D.1 PROOF OF PROPOSITION 3.5
Proof for SGD-M Dynamics. We first establish some preliminaries to simplify subsequent argu- ments and notation. For the direction vector v, we write x according to the decomposition
T (z, u)w (T vv )ar me bu
where u /r3 (x - m/)///=> -mvl is an unit vector orthogonal to v, i.e. (u, 0
and | mw/ V
By rotational invariance, any standard Gaussian vector N 7) admits the decompasition
01 + 0z + w, = (v. (u,9)
where a1 a2 are independent Gaussian variables from (0, 1) and w g 1a-ayu-a20 02 is centered Gaussian vector independent from a1, 02 and orthogonal Span (u, v) with covariance E[[w"" uu vv
We define the scalar projections := (x, a) ma1 raz (as (z, w) by construction) and 1 = (u, a) dj Let A == f(s) -[(() + denote the residual error.
To determine the expected loss V | T.), note that by the chain rule
VL (s) (t) + + (s)a
Substituting the original decomposition a 0 d2u + w and rearranging terms yields,
V,L (s) azu + w) (2 (s)) (2azAf 22 In particular, noting that is independent α1, 02, we find that
VΦ 2E(a, (+) (f(a)A
Now determine the covariance matrix Cov( V VH), first show that the component of V lying on span u) and the remainder are uncorrelated to simplify our analysis. This follows by noting
Cov((2a (s)) (2azAf((s) u, 22 (s) 4 (s)) + azu)w]
33

---

Published as a conference paper at ICLR 2026
due independence of from a2). Thus, have the orthogonal decomposition for the covariance matrix
V Cov(VL) 4Cov (AI' [ azu) wl)
4Cov (Af (A)) azu)) + 4Cov (Af
4Var|Af"(saa [Var|A] T (f()s)))))))))) + +vu" vu]]
4E 4f (s) -
We now compute each of the generators using the quantities obtained above. In particular, the first order generators given by
fm = 2 Vm) = 2
2B(ai (mma + raz) (f(ma) + ra2) - f(a1))
fA2 = ( -ソン) = 4r ff f()] = , f (ma, + ra2) (f(ma) + ra2) - f(ai)))
where also simplify the expectation using the fact that are independent zero-mean errors. The second-order generators are 9m = 0 and
C8 (V. V2
C8 (4Var(Af' 4E|
2n
n2
30 4cBE
[ [f'(ma] + raa)) raz) f(a1)) + +o2] = =
where we again use the independence of and fact that E[e?] = o² by assumption. Finally, we determine the pre-limit for the volatility matrix as
(ar(AA/(saa)]]
8r (s)() {
(+)
JVJ = Sr (fv(Af(s)a
16r Af (s)a)
which evidently vanishes as co. Recall that we set == (5,9) = + and A f(s) f(t) +
We now prove the remainder of the proposition for SGD-U. We note that verifying (6n, nm))
localizability is a similar exercise as to the case of Tensor PCA in Section C.1.2 We thus cover the key observations below. Again, consider here ū(x) vni((z) Vn(m(r) (2), ( (z:) which implies all cases of interest (i.c. m and y² up to a rescaling of /n with appropriate centering).
Proof for SGD-U (6n; T/o) Tho) -localizabiiity.. Item (1) is the same as the SGD-M case above. For item (2), the preconditioned gradient is:
VL
2 Ao (s)a
vn sgn (s)
[1v + 02
W
|a|
|
Which implies that [[)())] = (1) given the expectation of each fraction above is O(I//T). The above also trivially gives |/ H| O(6,²), concluding item (2). For item (3) we again (as in the tensor PCA example) have:
sup E(Vi((z,)) v), Vū(())) < n² NEw ((Ex)
sup *Eu (Fx) E(VL(x, (nVu( VL(a (x, !
We remind the reader that VL ( and Vu(x) (v, mv)), hence
(VL(x,v), (z), (( Additionally, V x|| 0(1/VR) which yields:
(x, Vu,(z)) -
= (E.)
The last part of item (3) follows identically to the Tensor PCA example.
34

---

Published as a conference paper at ICLR 2026
Proof for SGD-U asymptotic closability. Recall that the gradient of the loss function is given by (( Y) 2 (s)a follow the established notation in the online SGD case. The norm of the gradient is |L = 2)
The preconditioned gradient G(z,, Y) is:
G(x,Y) 1/ 2|Ao'(s) [s]l||||| (2 vn sgn (s)
To simplify notation, we write S(x, Y) sgn( (s)) and à a/ a/||a|| that
G(æ, Y) - vnSa
The expected gradient vΦ(x) (x, Y)] - - // SGi] can be further decomposed as
VQ(+) VnB
5 aiv agu +
||||
vni E
|al|
v
E
[s
a2
aall
1
To see why the last component vanishes in expectation, recall that the law of w is spherical and symmetric about zero. Furthermore, for fixed 01 02.
||a||
a
a3
an odd function of w. By conditioning on d1, 02, we conclude that
SIE
||a|
01 02.
0
Now consider the covariance matrix ((x) Cov(G Y)) - moment matrix E[GG"] ] can be deduced as
The second
E[GG EBnS ad [āa J.
where 2 almost surely. Since a is an isotropic Gaussian vector, á is uniformly distributed on the sphere. It follows that E[āà" 1 I/n (see, for instance (Vershynin, 2018 Proposition 3.3.8) for construction). It follows then that the covariance matrix is:
V(x) 1 (Φ(x) X (Φ(z))
We now turn our attention to the limiting generators. First, we see that V admits asymptotic simplification due to the behaviour of ||a]]. particular, we see that for the coefficient of the basis vector
E San
Saz
E (SPaa)) E
E
1
|al|
||a|
||a|
by Cauchy-Schward where S and Jensen's 's inequality, we have
almost surely and B[aa] 1. By Lemma E.3 (assuming
9)
72
n
n - 2
and
E
+
a
2
2
which shows that v ||Sai/||al]] E[Sai] + Similarly, we can show that v E|Sa//laa]]
E[Sag] + o(). useful observation is that | = O(1) with respect to n as each of the cocffi- cients for the unit vectors v, bounded.
With this, it follows that the first-order generators for Vm = e and 2 = 2ru are
(, Vm) = [[ + o[] "=00 aan sgn (o'(s)(((s) - ((1)) c))] = fm
V -ソ²) 2r Sa2] + o[[lu]]) "3 Br sgn (o (s)(((s) a(a1) + c)) = f,-
For the second-order generator, again we have that Jm and
C5
Co (nn 1) (V Vे, vo
"=000
Cs -
2n
35

---

Published as a conference paper at ICLR 2026
since (Vă & VΦ, vv | - | vv ol V = O(1). It remains to compute the volatility matrix for the preconditioned covariance matrix.
JUJ = - (B[Saa])²
-2rB[Sa [[Sa2] + o(1)
[r[ o(1) 1- 4r (B[S22)) + o(1)
which evidently vanishes as n + oo once we incorporate the additional factor of ó = cs/n. Recall that we define
S =
+ + - o(01) +c))
and ma1 r02-
D.2 FIXED POINT ANALYSIS FOR SGD-M (MONOMIAL)
In this section we provide the proof for Propositions 3.7 and 3.8 Before doing so, we prove a more general result and obtain the dynamics for any monomial link (x) = x² with k 2 1. To simplify notation, recall that we are given the dynamics for (m, r1) by:
dm -2 Ems
dr -4
(E,,1 C'Er,)),
dt
dt
where C² c6 and the expectations are defined as:
Em = [a (())
E = E [a (2)
E,2 = E [f' (Z) (Z) - -f(a)))
Here Z = aim + 02f1, and a1, (2 ~N(0,) 1) are independent. By Stein's Lemma may further rewrite the expectations as
F=1=rEloof(
((a)))==r11
(((()(()))()))))))
- s(an))
002
- ((a)))rr((())) = 2
=
Ja -00)) f(a)))
Em = Blai
= = =2)m/"(2)(((2) ((2)) (((a)))] mP
which yields us the reduced dynamics
dm dt
dr -4
1 - (mP - Q),
1- -
dt
We set f(x) = 2t, We define the joint moments Ma (m, 12) 21 m and rT given by the Binomial expansion:
These are polynomials
p
Mp.9 E
P i
Σ 5=0
p
m 1italloーj
1=0
where Pn is the n-th moment of a standard Gaussian. Substituting f(x) = x into P. P.Q, E,,2:
P Elk(k - at)] Elk k(2k 1)M2x-2. k(k 1)
Q E|kZ*-1 kag = 2 M -
E,2 22 )² + ا2]] 122 (M 2M M22-2,24 + + 02
Following this notation, the fixed points to the system satisfy m.P Q = 0 and P- C²E,, = 0. We necessarily have that any fixed point must satisfy r? > 0 under the construction that 0² > 0, k > and 0 since they imply that E, 0 and C' > 0. The global optimum is (1, however it is not necessarily fixed point in the presence of noise. We now prove the two results of interest.
36

---

Published as a conference paper at ICLR 2026
Proof of Proposition 3.7 For f(x) x², we have that the quantities in the dynamics reduce
P = 2(4 - [a + a2r) - 2(1 ((1) = 6R² 22
Q 41E[(am d2r)a1] 4m
E 4 (Bl(a + - a2r) >aaaaa/ aÿ] I[amamm + a + o + azr))
=1 30m 36m- - 6r 15m + 3r o² R²)
where we use the fact that a m + a2? ~ (N R2) and apply Wick's lemma for computing the higher-order joint moments. To determine the fixed points, note that they must satisfy 6m [t2 - = 0 which implies m = 0 or R² = 1.
In the case of m O, we have R2 r² and (m, r²) (0, 0) is a trivial fixed point. For y2 > we find that fixed points exist only if there are roots to the quadratic equation with respect to r²
4ca
6 -
1- (15r Gr +3
For R² = 1, we solve the equation using m²
+ +² =1 1,
0²
CBO²
3 (12r* + +
1-E 12
12cs
with the corresponding m = - r². Finally, under the constraint that p22 > 0, it follows that this fixed point exists only if cs - 3)/12 which completes the proof.
Proof of Proposition 3.8 We analyze the local stability of the dynamics near (m, (1, in the regime where o² > O. We do so by examining the Jacobian J. Let Fm = dm//dt and F, = dr /dl denote the time derivatives of both statistics. One can show that J(1, 0) is upper-triangular since On Fr ((1.0) depends on 0m Er which has terms depending on products of f(Z) - f(a1). J. At (1, 0), = 41 so the derivative is zero. Thus is sufficient to look the diagonal elements (eigenvalues).
For Am = Om Fml(1.10))) we have
-2
-2
Am 1-
0m om Em((1.0) = 1- Ela - ((a)
(1.0)
At (1, 0), the expectation is strictly positive that Am < 0 and we that the dynamics are stable along the m direction.
For X 0,3 Frlc,,), we have
4 (P(1,0)- C²a,3 Fr.2(
X = 1-3
(f = 122 p2x-2 for f(x) = xt. Evaluating each term independently, we see that P(1.0) = (a1)
On the other hand, analyze the expansion of E,,(11, rf) near T1 0. In particular, using f(Z) f(11) = f'(a1)a2r1 + O(r2):
E,2(1,12) = ([( ((a)))
E [(f( (f Pa3r1 + 0(+)))
0
so that 0,3 E,21(1,0) = (lf = ka 4= for f(x) *
Stability requires X < O which yields the condition
k2 1212-22 1
C8
yak
cs (1 3) 222 (4k 5)!!
1 (2k: 3)
by substituting the closed-form moments of the standard Gaussian. (-1)!! = 1 for convenience.
We adopt the convention
37

---

Published as a conference paper at ICLR 2026
If we wish to obtain the decay rate for the upper bound, we first that
(2k- 3) (2k: - 2) 22k 22k-2(2k - 2) = 2h-1 (k- - 1)!(4k: 4)! (4k - 4)!
((2k 2)1)
(4k: - 5)!! 2 1)!
where we use the fact that 2k 3 3 and 4 5 are odd numbers allowing us to substitute the closed- form expression for the double factorials. Apply the Stirling approximation m!! ~ V2Tm(m/e)) with = k: yo get,
((2(A) 1))1)2 1)((4 - 1))!
k-1
2
8(k: -
which implies that the upper bound of is of the order h k for sufficiently large degree k.
D.3 DIFFUSIVE LIMITS FOR SGD-M
We present a more general version of Proposition 3.10 below.
Theorem D.1. Suppose the link function f E C² satisfies (z)l 0 for Z ~ (0, 1). Then for SGD-M, the rescaled statisties (1) = (m, r²) = (Vnm, r²) converges as n + oo to the solution of the following SDE,
2
dm 1-
mB aa (frr))((f(rr ((a))) fraz))) dt
3
f(a1)) [af ((f(ra2) f())))))))(((a)))))) o o dB,
3
4
dr
rE [a2 (raz) - f(a1)))-
((a)))
1-3 E (ra)) ((f(ra)) f(a))) 2 +o]]]] df
(21)
where 01,00 02 (O,0 independently and B, is standard one-dimensional Brownian motion.
Proof.
From Proposition 3.5 we rewrite the drift components for the dynamics associated with
U2 2 as
( Vr²) = 4r [a2l' (al(ui/ + a2r) (f (ui(/// + a2r) - f(a:))]
C6 VV. 22,) = 4cg ff (ai(ūn/ + a2r)² ((f (ai(61/ a
+ azr) - f(a1)) + o²)
2n
Taking the limit as n + oo recovers the drift components for 2 For the dynamics of ū₁ Vnm, we linearize the drift about m m= 00 0,
Fm (m, = E aa) agr) f(a1))]
=
- f(ai))) + m0,, ((0,)) + o(m)
By independence of a1, 02, the leading term reduces to
Blai (aa
By Stein's lemma, have that [ar) E[f'(a))) which is zero by assumption. For the first-order term, the partial with respect to m is given by
in F(m, r) = fla (aim + (f + agr) - f(a1)) + +
Thus under the rescaling of ã Vum, we have that
4 Vūn)
인 vn
fa)) + f'(a2)) f( o(n oon-onnn// 1/2)
which recovers the limiting drift for म as n volatility matrix from the proof of Proposition 3.5
oc. For the diffusion term, recall the limiting
(s)r
8r (f V
JVJT = 8r (f
38

---

Published as a conference paper at ICLR 2026
where we set 8 (x. maj ra2 and A f(s) f((a)) + E. Evidently, under the rescaling only En1 survives the n + oo limit as reduces to
(ū1, = 4Var ((f(raz) f(a1) + +) )) (raz)a1
To simplify the variance, note that
E [[((ra2)) f(ai) + +) [r(a))ai )B
where the reduction follows by independence of (1, 02, followed by our assumption that B[aif(ai)] = [a)) = by Stein's 'ii'ss lemma. Thus the limiting volatility coefficient co is
En1(01.) 4E (raz) ((f(raz) f(a1))2++
which completes the proof.
The condition (z) = 0 is to ensure that the fixed point for which we are considering the diffusive limit exists along 0. This holds, for example, if is function such as in phase retrieval where (() x.
Proof of Proposition 3.10
With Theorem D.1 we compute each of the coefficients using f(x)
x², f((x) = 2x, and f"(x) 2,
2E
[a (f"
f(az)) (raz)²)] 2B 2a = 12(r² - 1)
(reloaf(ro))(((rra)))) f(a.)))))))))))))))))) 8 " at) = 8r2 - 1)
and with + 0,
4E (raz) ((f(raz) f(ar))++
16
= 16r 16r² - -6r +3+
4E (ro))) ((f(raz) f(a1)) 0²) 16r'E|ajo o1) 16r
=16r 18 +5+ 0
Substituting the quantities recovers the diffusive dynamics in equation 3.10
D.4 DYNAMICS FOR SGD-U (STRICTLY INCREASING LINK FUNCTIONS)
If f is strictly monotonically increasing, then (s) 0 almost everywhere (a.c.) and stn(f(t) f(s)) = st - 8).
Lemma D.2. Let (X, Y) be jointly centered Gaussian variables. Then
E E[Xsgn (V)) = / π Var(Y)
Cov(X, Y)
Proof. Consider the random variable W = X - aY with constant a = Cov(X, Coo(XXY)/Var(()) /Yr.)). Since (X, Y) is centered, have EW 0. Moreover, Cov(W, Y) which implies that W is inde-
pendent of as they are jointly Gaussian. By the towering property and independence, it follows that
Y ElaY + Y = =Y Cov(XX
Var(Y)
Returning to the original quantity [E))snn(Y)), we have again by towering that
E[X sgn (Y) E [ sgn()) Y| Cov(X,Y) E Var(Y)
To complete the proof, write Y Var(Y) Z where Z ~ NNO,, 1) and use the fact that E| 22/m.
39

---

Published as a conference paper at ICLR 2026
From here, collect the following quantities
Cov (a1, (1 m)a1 ra2) 1 m, Cov (a2, (1 m)a1 -ra2) -r,
Var ((1 m)a1 ra2) = m)² +
so that by Lemma D.2 we conclude that
2 1 m
lim lim (V Vm) = Ela, sgn ((1 - m)a1 - ro2)]
=V v(1 mm² + +
a n-00
2
2r²
lim lim (vé, V 2r Ela2 sgn ((1 mmmam - raa)) =
n-
v (1 m
D.5 DYNAMICS FOR SGD-U (EVEN POWERED MONOMIALS)
To complement Proposition 3.9 we provide additional analysis for even-powered monomials (i.e. f() = 2.2h for k 2 1) next. We first define the auxillary functions,
F(x,y)=
T
V(r + 1)2 V 12+
V yy²
G( = - 113 +x+ y²
Vx² + y²²
We show similar universality result where the dynamics are again invariant with respect to the monomial's degree.
Proposition D.3. If f is an even power polynomial (e.g. f(x) = 2k for & 2 1), then with
v /n/l| obtain the following universal dynamics,
2
4r
dm
F(m, r) dt, dr = 22T
G(m,r) cs dt
2m
Of interest to is a more refined analysis of the fixed points that approach the population opti-
mum (1.00 0). As cannot determine closed-form solutions for the system of equations in Proposi-
tion D.3 we instead obtain the following asymptotic expansion fixed points in the regime where
? approaches zero.
Proposition D.4.
The asymmetric fixed points (m, r2) of the dynamics in Proposition D.3 satisfy
F(m,r) and
4r G(m, r) = c5, where 2r2. More precisely, we write
1. For 0 < cs I 4/44/11/// 100, these fixed points satisfy |ml] E G2. 1) and , E (0, 1].
2. Fix 0 < < 0.2. Extend c by c(0) == O and define
2r
c(r)
G(m(r),r),
Cy max c(r)- osršy
V2T
Then C. 0, and for every 0 cs C, there exists asymmetric fixed point with E (0, and m = +m(r). In particular, - 33 < m(r) < 1. Moreover, as C9 0 (equivalently A: = csV2 /4 0), the solutions obey
1 r k+ k 2
_ k +O(K),
m= + (1 (1-3"++067)
The first item in Proposition D.4 provides the existence of threshold for ca such that for cs smaller than this threshold, we obtain fixed points away from m = 0. In particular, we find this threshold to be 4/ 10 z 0.71, which is larger than, for example, the 1/12 threshold observed for phase re- trieval in Proposition 3.7 We suspect similar phenomenon occur for more general polynomials.
40

---

Published as a conference paper at ICLR 2026
The second item states that for сş satisfying another threshold constraint with upper bound Cy, the system converges to a fixed point within some neighbourhood of the optimum (1.0). In particular, have that for y S 0.2, we can compute Cy 0.286 for the precise asymptotic expressions for (m,r²) to hold. Even if we did not take y in this range, still have larger range of admissible step-sizes for SGD-U for which we attain a desirable fixed point whereas online SGD still at the equator (i.e. 1/12 < os < 4/v 10w in the case of phase retrieval).
D.5.1 PROOF OF PROPOSITION D.3
Let f(x) = 2k for k 2 Applying the identity at bb (a E/= a² 1---", we have that
k-1
2 2h = (t- (t + 8)
Σ (()
where the series is polynomial with strictly positive components for s o Thus we can reduce the quantity inside the expectation again to get
sgn(t - 22h) sgn((ggntt () + s))
s)
for k 2 1, thus the following result universal for any even-powered monomial since we have
S(s, t) sgn (f((s)))())) - f())) sgn (s(t s) + s))
Recall that f = 01 = (a. v) and s ma1 + +raz (a. me + ru) so that s, s. s are linear in (a1, (2). we write as a function (a1, 02), we see that S(a U2) = S(pa, pu2) for any p p> 0, thus it is a homogeneous function of degree zero. In particular, if we write the joiny vector (4), a2) in terms of polar coordinates (a1, a2) = (p cos 0, p psin 0), we see that S(a1, 02) (p cos psin 0) S(0)
so that S only depends on the angle 6
Lemma D.5. The inner expectations in the drift terms simplfy to
/ 2n
/
2x
E E|a
cos 0 S(0) d0,
Ela2S]
sin 0 S(0) d0.
V8
Proof. We consider E(a, S] and the other integral follows by symmetry. Substituting the joint den- sity of (a), 02) in polar coordinates (recall the Jacobian is the radius p) and noting that S bounded (to invoke Fubini-Tonelli), we have that
S]
1
an a2) c (af700))/2 da daz
2T g22
1
/
25
(pcost))(()) - 02 dp
2m
1
(fee
pe ap) (1e
2x
cos 0 S(A) d8
2m
The first integral simplifies by the standard T-formula to / Thus the integrals reduce
π 1/a
:2n
{"
2x
E[aiS]- = 2π 2 2
cos 0 5(0) d0
cos 0 cos8 dB
desired.
To further simplify the integrals, we consider the behaviour of S(0). In particular we see when the sign changes as one of the three linear factors vanishes. These boundaries are solutions to the equations
0
cos 0 + sin 0 0, 0 m) cos 0 sin (1 + m) cos + r sin
[ t+s 0
41

---

Published as a conference paper at ICLR 2026
and parameterize the auxillary tangent quantities with a1, 02. ag (0, 2) by,
m Lan Or lan 02 = m
1 + m
tan 0kg =
The solutions to the three systems gives two antipodal angles in [0, 2x) (by also accounting for the quadrant they lie in to ensure the signs align with the system):
8 = 0
0 E { 01, 2x a1)
t -- = 0 = 0 E {az, π + a2).
+ 8 = 0
0 E { - 0g 2n 033
For simplicity assume 0 < < 1 (we have previously established > 0), however the same arguments hold for -1 < m < 0 in the case of even-powered monomials as only the ordering of angles will change, not their values. With this, ordering the angles yields
80
2
=02
0s
0g
using the observation that a2 I a1 ag g since (by monotonicity of tan for x E (0, 2/2.
a2
1
01
03
To determine the boundary behaviour,
that 0 0+ have di >
0, 02
0 hence
> 0, - s >> 0, t + s > 0, so S = +1 for small 0 > 0. Crossing a single boundary changes the sign of exactly one factor, hence flipping the overall sign we have an odd number of quantities to consider. This exhibits a behaviour where the signs alternate across the intervals
(0, 0o), (6u. 01), (91,02) (02,0s), (0s, 01), 0s), (0s.27). 2x).
The boundary sets themselves have Lebesgue zeasure-zero, and thus are not considered when com- puting the integral.
To compute the integrals using the alternating sign pattern (and simplifying the notation to illustrate the argument), we first see that
| cos 0 S(0) do
|
00
01
82
.03
0.
0,
2x
+
cos 0 d0
0
0p
01
02
00
0a
0s
80
0
2x
= sin
sin
+
sin
0s
= 2 (sin 0o - sin 0 + sin 02
sin 0g + sin 04 sin 0s)
Substituting the values for with i 1...., yields the concrete integrals
2
cos 0 S(8) dl (sin + sin 0g sin 03)
0 (()) d0 (cos a COS 02 - cos (3))
Writing the trigonometric ratioes in terms of (m, r2), conclude that
lim lim VVē, Vm)
2
m - 1
m +
m
22
(mm - 1)2 + p
Vm² + 2 Vmm + 1)2 +
2
1
lim lim V
2
8-00
( 1)2
Vm² + , V(m 1)2
which recovers the first-order drifts as desired.
42

---

Published as a conference paper at ICLR 2026
D.5.2 PROOF OF PROPOSITION D.4 D.4
We prove this in several steps, starting with the existence of fixed points and respective global bounds. The proof only requires elementary arguments from calculus.
Lemma D.6. For any fixed r E (0, 11, there exists a unique m(r) E 1 such that F(m(rr, = 0.
Proof. We analyze the function m + F(m, r) that
on the interval
[ 1j. Starting with 6m F, we find
,2
y2
OmF(mm, r) ((m - 1)² + r2)3/2 ((m + 1)2 + r2)3/2 (m² + r2)3/2
For m 2 1/2, we have (m 1)² I m² which implies that 0m F ((m, r) 0. At the endpoints, we find that F(1, r) 0 and F(2. r)
for any ' 27/20. By assumption that r (0, 1], the
existence unique root m(p) ($ 1) follows by the Intermediate Value Theorem (IVT). 0
Lemma D.7. The function c(r) given by
c(r) = 4r
G(m(r),r),
E (0.
V2n
is continuous on (0, 1]. satisfies c(1) 2 4/ 1/00, and lim, c(r) =
Proof. By Lemma D.6 we have that for any fixed ro E (0, 1], there exists some mo = m(ro) such that F(m(ro)) ro)
with 2,, F(m(rr) ro)
By the Implicit Function Theorem (IFT),
follows that m(r) is C on (0, 1], thus c(r) is continuous.
On the other hand, for m(r) (1/2,3) 1) we have (m 1)² m² which implies that
4
G(m(r),r) v (m- 1)2 + + c(r)
2
+
At the boundary r = 1, we have c(1) > 4/ 10T. Finally at the limit r - 0, it follows by the Squeeze Theorem and non-negativity of e(r) that
4
0 ~ c(r) <
+ 0
2
as desired.
Proof of Part (1). Let 0< co < 4/v10m. By Lemma D.7 c(1) 2 and lim c(r) = 0. By
IVT, there exists r. (0, 1] such that c(r-) cs where the corresponding m, m(r)) satisfying m, E (2, 1) follows by Lemma D.6
To prove (2), we first analyze the behaviour when is small, focusing on the regime E rg) where ro = 0.2. Then we connect this asymptotic behaviour with the specific choice of cs as stated. The arguments remain elementary, albeit tedious.
Lemma D.8. For r E (0,0.2), the unique root m(r) satisfies - 3 < m(r) < 1.
73, 1]. We first note that the function Proof. We analyze F(m,r) r) on the interval [11 -
h(s) (1 $) -)-1/2(1++))) -(1+
concave on the interval s E (0, 1] since (s(s)) < 0. particular we have that h(s)) < (h))+((0))
where h())) = 0 and h1((0= = 3/8. Thus, we obtain the upper bound F(1rr) = h((r) < 3/2
On the interval m E I, we have mm 1/5 S 3 so that &m F is bounded below by
Om Fm, 2
r(1 +1
43

---

Published as a conference paper at ICLR 2026
by bounding the first term in F, dropping the second term, and noting that for the third term that
m + r² 2 (1 - r²) + = r²
2r + r²) 2 1
for r I 0.2. Thus for r I 0.2 we can establish that 3. Fm,,r) 2 1/(2)) r². To conclude, by the Mean Value Theorem for some & I we have
0m F(1 - ,r) = - F(,,r) ,
where we substitute the previous bounds to get that
5 ( 2r - p2 x²
---
33 - < 0 for r S 0.2
+ 1
Thus since F(1 (11-r") ((1-r9,)) < 0 and F(1,r) > 0, the conclusion follows by IVT.
Lemma D.9. As r +r000 0.
m(r) = 1
3
3 0(rr), dr) = 2m
(r
1 n² + 0(rt))
8
2
Consequently, with k: == c(r) / /4, the inverse admits
1 k.a + O(K*).
r(k) k + 2 2
k2
Proof. We write 8 == mm((r) where F(m(r),r) r E (0, 0.2) one has
by Lemma D.7 By Lemma D.8 for
1 1- r² < (m) <1 = =
so that on the interval ((r) == [1 - م3م 1]. we obtain Fm Fn (m,r) 2 1/(2r) for r S 0.2. Using a similar MVT argument the previous lemma, write
F(1,r) - F (1- - 8,r) = F(1r)) F = 5 == F(1,))
for some E (1- (1-8,1) c. I(r). Taking the expansion of F(1,r) as - 0, we have
F(1,r) 1+
p2
-1++22 -1/2
,2 + 0(1*)
Next, we similarly expand F., (E, noting that
1 s 5 Orr²). For the first term, we have 2
22
3 (1+ + 0 +
((E 1)2 22)3/2
while the remaining terms in F,, (E, are O(r²) uniformly in E I(r):
2
,2 (ga + r
0 < (EE + 2233/2
2
(since e² + 2 1).
Putting all of the terms together yields the asymptotic expansion
(E. r) 1
O(+²)
uniformly for [(r).
(22)
where substituting it back the quotient representation for d yields.
3 (3 ( (r (r+0(r4)
3 3
+(r3))
m(r)= 11 8 8=1 , 0(55)
8
Next obtain the expansion of c(r), first expand G m = 1:
1
1
1
G(I,r) = -
- 0
V4+7
VI+
44

---

Published as a conference paper at ICLR 2026
By MVT, for some Ç (1 C (rr with ó as defined before gives
G(m(r),r) G(1,r) = G(C,r) (m)()) 1)
We can show that 8n G(, - 0(1) uniformly on (r) by observing that
m -
m + 1
(m 1)2 S 1
((m 1)2 + r2)3/2
p
(m + 1) + r2)3/2
m
m 2
(m² + r2 2233/2
where we also recall that m 2 21-r
2
Hence 10 GI < 4 uniformly and thus
1 +0(72) = c(r)
4
(r
1 2-- + 0(74)
r
22T
To show the final claim about inversion, let h == dr)V22/4 and note that d(()) > 0 as - 0. Following our expansion for c(r), we write
k=r 2 + 0
and revert the series up to the leading terms, say r= = r(k:) = 1: + a2k? + agh3 + O(kt). To do so, simply substitute and match powers to obtain that
1 (k: + azk3)² + 0(K4) k
1
k: = (Is + agk² azhh3)
a2
222 (as 02) k O(k*),
2
which implies that a2 = a3 = which gives
rkk) + 5 + 1 k3 + O(K)). 7
k2
2
Proof (2). Let < < 0.2. By Lemma D.6 D.7 and the extension of c(0) 0, c continuous on the compact interval [0, 7]. hence attains a maximum C.2 max(0,7) c(r).
By Lemma
D.9 we have the asymptotic expansion as r + 0,
4
<(r)
O(r4))
2n
which again verifies that c(r) 0 for sufficiently small
0, in particular C, > 0.
Fix cs E (0, C.] and choose T. E [0, 7] with c(r,) = C., Since is continuous and c(0) = 0, by IVT there exists E r0,r,] C (0, 7] such that c(r.) cs. Set m, == m(rr), then by construction
4r G(m,, r.) c(r.) C5:
F(m., r.) O,
22T
so (+m,, r2) are fixed points. Furthermore, since r. S S 0.2, Lemma D88 implies - -3 m. < 1.
Finally, because e(r) - 0 as r + it follows that cs 0 asymptotically as well. By Lemma D.9 we obtain the expansions,
3 m(r) =1 -1 3,3 +06),,
c(r)
.
(r r- 5 1
22T
which the local inversion
2T
rkk = k:
2
k3 + O(K4), k =
Cs
2
4
which together yield the claimed asymptotics for (m. r.) C8 - 0.
45

---

Published as a conference paper at ICLR 2026
E TECHNICAL PROOFS OF AUXILLARY RESULTS
E. NORMS OF GAUSSIAN VECTORS
Lemma E.1. Let X ~ (u, )) be an n-dimennional Gaussian vector with > 0 and ||>]| < oo (i.e. X non-degenerate). If || = O(vn). then the cumulants Kp of || are also of the order ((n) for each fixed 2 1.
Proof. Noting that | X X is quadratic form, the cumulants can be determined exactly, for example see (Magnus 1986 Lemma 2):
Ky = 2n - 1) (tr((t(()))))))))))))(")) + p(u, 100,,2~11p) uu}
Using the assumption that || = O(vn) and tr(2*) (n) for any k 1 completes the proof by assumptions on the spectrum.
Lemma E.2. Let X ~ (M, M) be an non-degenerate n-dimensional Gaussian vector. If n >> 2k, then TE| X|| 2h: < 00.
Proof. Denote 7x as the corresponding density of X. suffices to demonstrate integrability near the origin since,
E| X|| 2h [1())21)]
||||| xx dx S ,22 < oo, for r >> 0
To this end, consider the restricted expectation over the open ball (0) for r > O. By continuity and positivity of yx on (0),
|||| 7x (x) (x) dz < |/ loo
||||||
2k dr
B.)
B,((0)
As |||||| a radial function, by the coarea formula (e.g. see Stromberg 2015 pg. 370) the remain- ing integral becomes
Sa.00
|0|||
2h dx
|sn-1|
p²
1=-1-22 dp
B.(())
where [S"-2| is the surface measure of the (n dimenssioa unit sphere. Noting that the integral finite only if > 2k completes the proof.
Lemma F.3. Let (N, Σ))
IR" with E > Then for any integer k 2 with n. > 2
1
E X| < Amiu (E)=
n2 2k
In particular, if Amin (E) 2 > 0 uniformly in n, then E||X| -28 = O(n
||
Proof. Let Z (0, 1,) be standard Gaussian vector. We first determine comparison inequality involving ||| 2k for μ = 0, and show that the upper bound is still applicable μ * O by Ander- son's 's inequality Anderson (955).. In particular, we have that for the centered Gaussian density y2 and origin-symmetric, convex set K C R
f T5(a - 14) dx = f( (E(x) do S / 7x(a) dr
(23)
Proceeding with the argument, if we write X 1 and obtain the comparison inequality
Amin [ 2 < |X|| < Amax (E)
Raising both sides to the power 2k for 1: 2 1 and applying expectations yields
Amaxe (E) ~E El X|| 24 Amin (E) -
(24)
46

---

Published as a conference paper at ICLR 2026
On the other hand if u # let Y ~ (0,, Σ) be centered version of X and denote (2( as the Gaussian density of Y. Note that the density of is (2( p)). Then we write
E|
X 2k
2k yy |||||
pz
1 [[:]]=|]<<)-v1a))) 22(x - μu) dt dr R2 J
R2
by "layer-cake" representation. Note that the inner integral is over the centered Euclidean ball
B(0,t-1/22)) = {a: ||=||| < <()
which is symmetric and convex. Furthermore, by verified integrability in Lemma E.2 for n >> 2k, we invoke Mubini-Tonelll to write
[ B 1/(a)) 7x(x (1) dx = 1.0 0
x 2
JB(0,-1//00) xx (1) dr di
By equation 23 the inner integral can be bounded (and further simplified by reversing the layer-cake arguments above),
B(0,r-1/(20) x - p1) d. dt & 1 00
|||||| 24
(2(2) dæ df E||||| 2k
Thus the upper bound in equation 24 still holds as the analysis applies to the centered Y.
remains to determine the asymptotic rate for Z|||
22] Since 2 follows a Xa distribbbttion
and n > 2k with & 2 1, the inverse moments are given by Soch et al. (023)
r
k)
1
k
B 2
2k
2-
n 2-
n2
2k
Combining this bound with equation 24 completes the proof.
Lemma E.4. Consider an n-dimensional non-degenerate Gaussian vector N14, E) such that
tr(2)- = (n)) and |||||| o(n). Then for 9,
1 X
E
O(n
with ||||X||]] = tr())) + |A1
Proof. To simplify notation, put Q = |X 2 and a = EQ = =t(t) + ||a||2 and write 5 = (Q a//a, so (1 a/)) (5/(1 + 5))3 To control the expectation around the singularity d = 11, split the
expectation on the events
A (6 2 1/2),
B= < < = {Q a/2).
where we slightly abuse notation to incorporate the null event {Q = a/2) For 2 -1/2, we consider the bound
3/z
| x1) x)-
S |xd
(1-1 1/2) 48|x/
1+
by the mean-value theorem applied to + (1 (L+t) -3 expectation over as
Thus using the expansion above, we write the
8
3
E
1
=E
(1T16)111])
1+8)
Now 1883 (Q) and E = 44 (Q)/a² By Lemma
E.1 we can determine the order of the
central moments through the cumulants,
na (Q) xa(Q) O(n),
[a(Q) mu(Q) ( 0(n²).
47

---

Published as a conference paper at ICLR 2026
Since 2 for some expectation that
by assumption, conclude for the first term in the partitioned
8
O(n) (on)a
(On2t)
E
1.
+0
Onn
(cn)
remains control the tail event B. When Q a/2, it follows that Q and Q S that
(1-0) 3
3 Q
3
Q
where the last inequality follows as Q 2 0 by definition. By Lemma E.2 with 2 9 and Hölder's
inequality with exponents p = 4, 9- 4/3,
a (@))
3
1B S a (EQ (EQ-1)//
We bound each factor individually. First, note that a < (tr Σ + || = O(n²). Second, we have by Lemma E.3 that EQ = OO so that (EQ +)3/4 + O(n 3)
To bound the final term P(B) IPQ $ a/2). we use a left-tail Chernoff argument
P(Q (1 =)a) S exp (t(1 - =)a KQ(-t)), for any t (0, where the cumulant generating function K2(-t) is given by,
1 KQ(- = log det + 21E) - (I+ (I+22E)
2
Expanding the log determinant and noting that log(1 + x) - x2/2 for 0, obtain the upper bound on the egf
1
Ko(-t) log log ded 21E)
log(1 2(Ax) -ttr(E) t tr(E²)
the trailing term is quadratic form positive definite matrix. Applying the cgf bound on the tail probability, we have that
exp (t(1 E)a + Ko(-t)) s exp (-et tr(E) + l tr(E + ((1 - =)]1
recalling that a = ||a1|2 +r To obtain the desired decay rate, set / = E (r(EE) (r((²))) hence,
-el tr(E) 12 (r(==)= (tr Σ)²
e² (r(E) S -c'n
trr(E²)
4C
for an absolute constant c > ( where the penultimate inequality follows by the von-Neumann trace inequality
Lr tr(E)) < ||E|op tr(E) S C tr(E)
using ||D|op C and (r(E)= 2 c'n. Finally, noting that the remaining term satisfies ((1 - [1-=)]|m||² 2 O(1), obtain the desired exponential decay P(B) & Ke-enn for constant K 00 inde-
pendent of n. Combining, [[((a/Q) 1) S 50(e m'n) so that altogether,
E
(1-1)
[[(1
a
3
[
3
1
+
8
3
1B
O(n
Q
1+8
1
1
claimed.
48


Reviews:
[{"review_id": "l6yugF9c5R", "reviewer": "Reviewer_17tv", "summary": "The paper extends the main ideas of high-dimensional scaling limit for Stochastic Gradient Descent presented in Ben Arous et al. (2024) to two important variants: SGD with momentum and SGD with adaptive step-sizes. The paper provides a solid theory, and extensive experiments verify the results and justify the need for the exploration of these ideas.", "questions": "See the Weaknesses part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and well-organized with clear contributions. I enjoy reading the motivation of the main ideas, and in particular, i find the writing of section 2 on the main result of the paper very informative. Theorem 3.2 is the main theoretical contribution of the work, with the examples in Section 3 to justify the need for the theoretical result.\n\nThe proofs for the major steps I checked seem correct, and the theoretical statements are reasonable.", "weaknesses": "I do not have to point out a specific weakness. I like this work and I believe it is worth being accepted to ICLR.\n\nSome questions:\nTypically, the HIGH-DIMENSIONAL LIMIT THEOREMS of an algorithm are purely theoretical results. In this work, the authors show that SGD-M will amplify high-dimensional effects, potentially degrading performance relative to online SGD. Can they comment on what else can be an interesting follow-up to these ideas? What other algorithms might have a similar outcome?\n\nI believe having such remarks and paragraphs at the end of the paper would be particularly valuable, in a section, let's say, named Conclusion, “Remarks and Future impact”.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7iufPIIgdP", "reviewer": "Reviewer_L1BR", "summary": "This paper extends the high-dimensional effective dynamics of summary statistics framework of online SGD in [1] to online SGD with Polyak momentum (SGD-M) and online SGD with adaptive step-sizes. Based on these results, the authors looked at two applications: the spiked tensor PCA and single index models. In these applications, the authors \n\n(1) consider a specific online SGD with adaptive step-sizes defined by normalizing gradient, referred as SGD-U;\n\n(2) derive the ballistic limits of the summary statistics of SGD-M and SGD-U and study the fixed points in different regimes;\n\n(3) derive the diffusive limits of the summary statistics of SGD-U at the equator in spiked tensor PCA.\n\n[1] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. *High-dimensional limit theorems for\nSGD: Effective dynamics and critical scaling*. Communications on Pure and Applied Mathemat-\nics, 77(3):2030–2080, 2024.", "questions": "1. How does $\\beta$ affect the $\\delta_n$-localizable condition defined in Definition 2.1?\n\n2. in line 322-323, the coefficient of the diffusion term is incorrect;\n\n3. Why do you scaling by $\\sqrt{n}$ in SGD-U? As explained in the paper, if $\\lVert \\nabla L_n \\rVert=O(\\sqrt{n})$, wouldn't the adaptive step-size of constant order, which contradicts to that we want the step-size tend to zero?  \n\n4. When is there a diffusive limit in the single index model example?\n\n5. Please double check for typos. Some that I noticed are listed below:\n\n    (1) $\\mu$ in line095 should be $P_n$;\n   \n    (2) $\\nabla L$ in equation (1) should be $\\nabla L_n$;\n\n    (3) in line 230: $R^2$ should be $m^2+r_\\perp^2$;\n\n    (4) $r$ and $r_\\perp$ are the same. Please be consistent in notations.", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "1. The paper provides a **valuable generalization** of high-dimensional diffusion limits to include both momentum and adaptive step-size variants of SGD. This extension broadens the understanding of optimization dynamics beyond standard online SGD, offering a unified framework for several widely used algorithms.\n\n2. The main theorems and assumptions are rigorously stated and internally consistent.", "weaknesses": "1. The manuscript does not clearly highlight its key novelty—the extension from online SGD to momentum and adaptive-step methods. The relation to existing online SGD results is underdeveloped theoretically, which makes the contribution appear incremental despite being substantial.\n\n2. The effective limits of SGD-U is difficult to follow. Neither a general limiting theorem or a comparison to SGD/SGD-M is provided.\n\n3. The paper’s structure could better emphasize the main theoretical message. Sections on motivation, scaling intuition, and comparison to prior results are brief, while technical proofs dominate early sections, making it hard to identify the central takeaway.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bL0f4OIkwm", "reviewer": "Reviewer_WtAq", "summary": "This paper extends the high-dimensional limit theory of Ben Arous et al. (2022, 2024) to analyze Stochastic Gradient Descent (SGD) with Polyak momentum (SGD-M) and adaptive step-size variants (including a normalized-gradient variant, denoted SGD-U). The authors derive the asymptotic stochastic differential equations governing the evolution of summary statistics of SGD-M and adaptive SGD, under suitable regularity assumptions (delta-localizability and delta-closability). The theory establishes an equivalence between SGD-M and online SGD up to a time rescaling, and provides rigorous justification for the stabilizing role of gradient normalization in high-dimensional regimes. These results are illustrated on two canonical problems: spiked tensor PCA and single-index models, supported by both analytical fixed-point analyses and simulations.", "questions": "1. Is it possible to verify, even heuristically, that the paper's assumption are satisfied in one of the examples, or provide sufficient conditions under which they hold?", "limitations": "", "rating": 6, "confidence": 1, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "I would like to note that I am not a specialist in stochastic differential equations or high-dimensional diffusion limits. While I have done my best to assess the work carefully, my understanding of some of the deeper technical aspects is limited.\n\n1. The paper presents a non-trivial and technically rigorous generalization of existing high-dimensional diffusion limit results to momentum and adaptive-step-size algorithms. Extending the practical dynamics framework of Ben Arous et al. to handle additional algorithmic components is a clear and meaningful step forward. The identification of the equivalence between SGD-M and SGD via time rescaling (Theorem 2.3 and Remark 2.4) is elegant and provides theoretical insight on known empirical observations.\n\n2. The analysis of normalized-gradient SGD (SGD-U) offers a rigorous interpretation of empirical stabilization mechanisms (e.g., gradient clipping, normalization), bridging stochastic process theory and optimization practice.", "weaknesses": "1. The paper is extremely technical and assumes strong familiarity with high-dimensional probability, weak convergence, and diffusion processes. While this is common in this line of work, some readers may find it challenging to connect the formal definitions. For instance, the concepts of delta-localizability and delta-closability are poorly described\n\n2. I am not entirely sure about this claim, but it seems he localizability and closability conditions are mathematically strong (requiring smoothness up to third derivatives and bounded high-order moments). It seems to me that this condition is basically assuming the algorithm’s stochastic trajectories converge to an SDE, which compresses both randomness and nonlinearity into a self-contained stochastic differential equation. These may hold for the idealized models considered but might limit its applicability to more realistic scenarios (where similar behavior of SGDM and SGD have been observed, like in deep learning).\n\n3. I might have missed this part, but it seems the authors does not provide examples of problem that satisfies delta-localizability and delta-closability. Moreover, it seems these are never explicitly verified for either of the two worked examples (spiked tensor PCA or single-index models).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)