The High Line: Exact Risk and Learning Rate Curves of Stochastic Adaptive Learning Rate Algorithms

Authors: elizabeth  collins-woodfin
inbar  seroussi
begoña  garcía malaxechebarría
andrew w mackenzie
elliot  paquette
courtney  paquette
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-11 00:18:57

Abstract: We develop a framework for analyzing the training and learning rate dynamics on a large class of high-dimensional optimization problems, which we call the high line, trained using one-pass stochastic gradient descent (SGD) with adaptive learning rates. We give exact expressions for the risk and learning rate curves in terms of a deterministic solution to a system of ODEs. We then investigate in detail two adaptive learning rates  -- an idealized exact line search and AdaGrad-Norm -- on the least squares problem. When the data covariance matrix has strictly positive eigenvalues, this idealized exact line search strategy can exhibit arbitrarily slower convergence when compared to the optimal fixed learning rate with SGD. Moreover we exactly characterize the limiting learning rate (as time goes to infinity) for line search in the setting where the data covariance has only two distinct eigenvalues. For noiseless targets, we further demonstrate that the AdaGrad-Norm learning rate converges to a deterministic constant inversely proportional to the average eigenvalue of the data covariance matrix, and identify a phase transition when the covariance density of eigenvalues follows a power law distribution. We provide
our code for evaluation at https://github.com/amackenzie1/highline2024.

Main paper: 
 introduction: 
 In deterministic optimization, adaptive stepsize strategies, such as line search (see [40] , therein), AdaGrad-Norm [59] , Polyak stepsize [48] , and others were developed to provide stability and improve efficiency and adaptivity to unknown parameters. While the practical benefits for deterministic optimization problems are well-documented, much of our understanding of adaptive learning rate strategies for stochastic algorithms are still in their infancy.There are many adaptive learning rate strategies used in machine learning with many design goals. Some are known to adapt to stochastic gradient descent (SGD) gradient noise while others are robust to hyper-parameters (e.g., [4, 63] ). Theoretical results for adaptive algorithms tend to focus on guaranteeing minimax-optimal rates, but this theory is not engineered to provide realistic performance Concentration of learning rate and risk for AdaGrad-Norm on least squares with label noise ω = 1 (left) and logistic regression with no noise (right). As dimension increases, both risk and learning rate concentrate around a deterministic limit (red) described by our ODE in Theorem 2.1. The initial risk increase (left) suggests the learning rate started too high, but AdaGrad-Norm adapts. Our ODEs predict this behavior. See Sec. H for simulation details. comparisons; indeed many adaptive algorithms are minimax-optimal, and so more precise statements are needed to distinguish them. For instance, the exact learning rates (or rate schedules) to which these strategies converge are unknown, nor their dependence on the geometry of the problem. Moreover, we often do not know how these adaptive stepsizes compare with well-tuned constant or decaying fixed learning rate SGD, which can be viewed as a cost associated with selecting the adaptive strategy in comparison to tuning by hand.In this work, we develop a framework for analyzing the exact dynamics of the risk and adaptive learning rate strategies for a wide class of optimization problems that we call high-dimensional linear (high line) composite functions. In this class, the objective function takes the form of an expected risk R : R d → R over high-dimensional data (a, ϵ) ∼ D ⊂ R d × R of a function f : R3 → R composed with the linear functions ⟨X, a⟩, ⟨X ⋆ , a⟩. That is, we seek to solve min X∈R ddef = E a,ϵ [f (⟨a, X⟩, ⟨a, X ⋆ ⟩, ϵ)] for (a, ϵ) ∼ D, X ⋆ ∈ R d .(EQUATIONWe suppose a ∼ N (0, K) where K ∈ R d×d is the covariance matrix. We train (1) using (one-pass) stochastic gradient descent with adaptive learning rates, g k (SGD+AL). Our main goal is to give a framework for better2 performance analysis of these adaptive methods. We then illustrate this framework by considering two adaptive learning rate algorithms on the least squares problem 3 , the results of which appear in Table 1 : exact line-search (idealistic) (Sec. 3) and AdaGrad-Norm (Sec. 4). We expect other losses and adaptive learning rates can be studied using this approach.Main contributions. Performance analysis framework. We provide an equivalence of R(X k ) and learning rate g k under SGD+AL to deterministic functions R(t) and γ t via solving a deterministic system of ODEs (see Section 2), which we then analyze to show how the covariance spectrum influences the optimization. See Figure 1 . As the dimension d of the problem grows, the learning curves of R(X k ) become closer to R(t) and the curves concentrate around R(t) with probability better than any inverse power of d (See Theorem 2.1).Greed can be arbitrarily bad in the presence of strong anisotropy (that is, Tr(K)/d ≪ Tr(K 2 )/d).Our analysis reveals that exact line search, which is to say optimally decreasing the risk at each step, can run arbitrarily slower than the best fixed learning rate for SGD on a least squares problem when λ min def = λ min (K) > C > 0. The best fixed stepsize (least squares problem) is (Tr(K)/d) -1 or the inverse of the average eigenvalue, see Polyak stepsize [48] . Line search, on the other hand, converges to a fixed stepsize of order λ min /(Tr(K 2 )/d). It can be that λ min /(Tr(K 2 )/d) ≪ (Tr(K)/d) -1 making exact line search substantially underperform Polyak stepsize. We further explore this and, in the case where d-eigenvalues of K take only two values λ 1 > λ 2 > 0, we give an exact expression as a function of λ 1 and λ 2 for the limiting behavior of γ t as t → ∞ (See Fig. 5 ).Table 1 : Summary of adaptive learning rates results on the least squares problem. We summarize our results for line search and AdaGrad-Norm under various assumptions on the covariance matrix K. We denote λ min the smallest non-zero eigenvalue of K and Tr(K) d the average eigenvalue. Power law(δ, β) assumes the eigenvalues of K, {λ i } d i=1 , follow a power law distribution, that is, for 0 < β < 1, λ i ∼ (1 -β)λ -β 1 (0,1) for all 1 ≤ i ≤ d and ⟨X 0 -X ⋆ , ω i ⟩ 2 ∼ λ -δ i where {ω i } d i=1 are eigenvectors of K (see Prop 4.4) . For * (see Prop. 4.2), requires a good initialization on b, η.K assumption Limiting γ ∞ Convergence rateAdaGrad-Norm(b, η)(see Sec. 4) λmin > C γt ≍ η 2 b η + 1 4d Tr(K)∥X 0 -X ⋆ ∥ 2 log(R) * ≍ -λminγ∞tAdaGrad-Norm(b, η) Power law (see Sec. 4)β + δ < 1 γ t ≍ δ,β 1 R(t) ≍ δ,β t β+δ-2 β + δ = 1 γ t ≍ δ,β 1 log(t+1) R(t) ≍ δ,β t log(t+1) -1 1 < β + δ < 2 γ t ≍ δ,β t -1+ 1 β+δ R(t) ≍ δ,β t -2 β+δ +1Exact line search, idealized (see Sec. AdaGrad-Norm selects the optimal step-size, provided it has a warm start. In the absence of label noise and when the smallest eigenvalue of K satisfies λ min > C > 0, the learning rate converges to a deterministic constant that depends on the average condition number (like in Polyak) and scales inversely with Tr(K) d ∥X 0 -X ⋆ ∥ 2 . Therefore it attains automatically the optimal fixed stepsize in terms of the covariance without knowledge of Tr(K), but pays a penalty in the constant, namely ∥X 0 -X ⋆ ∥ 2 . If one knew ∥X 0 -X ⋆ ∥ 2 then by tuning the parameters of AdaGrad-Norm one might achieve performance consistent with Polyak; this also motivates more sophisticated adaptive algorithms such as DoG [29] and D-Adaptation [18] , which adaptively compensate and/or estimate ∥X 0 -X ⋆ ∥ 2 .AdaGrad-Norm can use overly pessimistic decaying schedules on hard problems. Consider power law behavior for the spectrum of K and the signal X ⋆ . This is a natural setting as power law distributions have been observed in many datasets [60] . Here the learning rate and asymptotic convergence of K undergo a phase transition. For power laws corresponding to easier optimization problems, the learning rate goes to a constant and the risk decays at t -α1 . For harder problems, the learning rate decays like t -η1 and the risk decays at a different sublinear rate t -α2 . See Table 1 and Sec. 4 for details.Notation. Define R + = [0, ∞). We say an event holds with overwhelming probability, w.o.p., if there is a function ω : N → R with ω(d)/ log d → ∞ so that the event holds with probability at least 1 -e -ω (d) . We let 1 A (x) be the indicator function of the set A where it is 1 if x ∈ A and 0 otherwise. For a matrix A ∈ R m×d , we use ∥A∥ to denote the Frobenius norm and ∥A∥ op to denote the operator-2 norm. If unspecified, we assume that the norm is the Frobenius norm. For normed vector spaces A, B with norms ∥ • ∥ A and ∥ • ∥ B , respectively, and for α ≥ 0, we say a function F : A → B is α-pseudo-Lipschitz with constant L if for any A, Â ∈ A, we have∥F (A) -F ( Â)∥ B ≤ L∥A -Â∥ A (1 + ∥A∥ α A + ∥ Â∥ α A ). We write f (t) ≍ g(t) if there exist absolute constants C, c > 0 such that c • g(t) ≤ f (t) ≤ C • g(t) for all t. If the constants depend on parameters, e.g., α, then we write ≍ α .Related work. Some notable adaptive learning rates in the literature are AdaGrad-Norm [32, 59, 61] , RMSprop [28] , stochastic line search, stochastic Polyak stepsize [35] , and more recently DoG [29] and D-Adaptation [18] . In this work, we introduce a framework for analyzing these algorithms, and we strongly believe it can be used to analyze many more of these adaptive algorithms. We highlight below a nonexhaustive list of related work.AdaGrad-Norm. AdaGrad, introduced by [19, 36] , updates the learning rate at each iteration using the stochastic gradient information. The single stepsize version [32, 59, 61] , that depends on the norm of the gradient, (see Table 2 for the updates), has been shown to be robust to input parameters [34] . Several works have shown worst-case convergence guarantees [21, 33, 57, 59] . A linear rate of O(exp(-κT )) is possible for µ-strongly convex, L-smooth functions (κ is the condition number µ/L). In [62] (similar idea in [61] ), the authors show for strongly convex, smooth stochastic objectives (with additional assumptions) that the AdaGrad-Norm learning rate exhibits a two stage behavior -a burn in phase and then when it reaches the smoothness constant it self-stablizes.Stochastic line search and Polyak stepsizes. Recently there has been renewed interest in studying stochastic line search [20, 42, 54] and stochastic Polyak stepsize (and their variants) [7, 26, 27, 30, 35, 39, 41, 49] . Much of this research focuses on worst-case convergence guarantees for strongly convex and smooth functions (see e.g., [35] ) and designing practical algorithms. In [53] , the authors provide a bound on the learning rate for Armijo line search in the finite sum setting with a rate of L max /avg. µ where avg. µ is the avg. strong convexity and L max is the max. Lipschitz constant of the individual functions. In this work, we consider a slightly different problem. We work with the population loss and we note that the analogue to L max for us would require that the samples a satisfy ∥aa T ∥ op ≤ L max for all a; this fails to hold for a ∼ N (0, K). Moreover, L max could be much worse thanE [∥aa T ∥ op ].Deterministic dynamics of stochastic algorithms in high-dimensions. The literature on deterministic dynamics for isotropic Gaussian data has a long history [9, 10, 50, 51] . These results have been rigorously proven and extended to other models under the isotropic Gaussian assumption [1, 2, 6, 16, 17, 23, 58] . Extensions to multi-pass SGD with small mini-batches [46] as well as momentum [31] have also been studied. Other high-dimensional limits leading to a different class of dynamics also exist [11-13, 22, 37] . Recently, significant contributions have been made in understanding the effects of a non-identity data covariance matrix on the training dynamics [5, 14, 15, 24, 25, 64] . The non-identity covariance modifies the optimization landscape and affects convergence properties, as discussed in [15] . This work extends the findings of [15] to stochastic adaptive algorithms, exploring the effect of non-identity covariance within these algorithms. Notably, Theorem 1.1 from [15] is restricted to deterministic learning rate schedules, limiting its applicability in many practical scenarios. In contrast, our Theorem 2.1 accommodates stochastic adaptive learning rates, aligning with widely used algorithms in practice.We suppose that a sequence of independent samples {(a k , y k )} drawn from a distribution D ⊂ R d ×R is provided where y k is the target. The target y k is a function of some random label noise ϵ k ∈ R and the input feature a k dotted with a ground truth signal X ⋆ ∈ R d , ⟨a k , X ⋆ ⟩. Therefore, the distribution of the data is only determined by the input feature and the noise, i.e., the pair (a, ϵ). In particular, we assume (a, ϵ) follows a distributional assumption.Assumption 1 (Data and label noise). The samples (a, ϵ) ∼ D are normally distributed: ϵ ∼ N (0, ω 2 ) where ω ∈ R, and a ∼ N (0, K), with a covariance matrix K ∈ R d×d that is bounded in operator norm independent of d; i.e., ∥K∥ op ≤ C. Furthermore, a and ϵ are independent.For a, X, X ⋆ ∈ R d , ϵ ∈ R, and a function f : R 3 → R, we seek to minimize an expected risk function R : R d → R, which we refer to as the high-dimensional linear composite4 , of the formR(X) def = E a,ϵ [Ψ(X; a, ϵ)] for (a, ϵ) ∼ D, and Ψ(X; a, ϵ) = f (⟨a, X⟩, ⟨a, X ⋆ ⟩, ϵ). (2)In what follows, we use the matrix W = [X|X ⋆ ] ∈ R d×2 that concatenates X and X ⋆ , and we shall let B = B(W ) = W T KW be the covariance matrix of the Gaussian vector (⟨a, X⟩, ⟨a, X ⋆ ⟩).Assumption 2 (Pseudo-lipschitz f ). The function f : R 3 → R is α-pseudo-Lipschitz with α ≤ 1.By assumption, R(X) involves an expectation over the correlated Gaussians ⟨a, X⟩ and ⟨a, X ⋆ ⟩. We can express this as R(X) def = h(B) for some well-behaved function h : R 2×2 → R.Table 2 : Two adaptive learning rates considered in detail. The stochastic adaptive learning rate, g k , is the learning rate directly used in the update for SGD whereas the deterministic, γ t , is the deterministic equivalent of g k after scaling.General update Least squaresAdaGrad- Norm(b, η) b0 = b × d g k b 2 k = b 2 k-1 + ∥∇Ψ(X k-1 )∥ 2 ; g k-1 = d × η |b k | same γ t η b 2 + Tr(K) d t 0 I(B(s)) ds η b 2 + 2Tr(K) d t 0 R(s) ds Exact line search (idealized) g k ∥∇R(X k )∥ 2 Tr(∇ 2 R(X k )K) d Ea,ϵ[(f ′ (⟨a,X⟩;⟨a,X ⋆ ⟩,ϵ)) 2 ] ∥∇R(X k )∥ 2 2Tr(K 2 ) d R(X k ) γ t arg min γ dR(t) d i=1 λ 2 i D 2 i (t) 2Tr(K 2 )R(t)h : R 2×2 → R such that h(B) = R(X) is differentiable and satisfies ∇ X R(X) = E a,ϵ ∇ X Ψ(X; a, ϵ). Furthermore, h is continuously differentiable and its derivative ∇h is α-pseudo-Lipschitz for some 0 ≤ α ≤ 1, with constant L(∇h).The final assumption is the well-behavior of the Fisher information matrix of the gradients. The first coordinate of f is special, as the optimizer must be able to differentiate it. Thus, we treat f (x, x ⋆ , ϵ) as a function of a single variable with two parameters: f (x, x ⋆ , ϵ) = f (x; x ⋆ , ϵ) and denote the (almost everywhere) derivative with respect to the first variable as f ′ .Assumption 4 (Fisher matrix). Define I(B)def = E a,ϵ [(f ′ (⟨a, X⟩; ⟨a, X ⋆ ⟩, ϵ)) 2 ]where the function I : R 2×2 → R. Furthermore, I is α-pseudo-Lipschitz with constant L(I) for some α ≤ 1.A large class of natural regression problems fit within this framework, such as logistic regression and least squares (see [15, Appendix B] ). We also note that Assumptions 3 and 4 are nearly satisfied for L-smooth objectives f (see Lemma B.1), and a version of the main theorem holds under just this assumption (albeit with a weaker conclusion).We apply one-pass or streaming SGD with an adaptive learning rate g k (SGD+AL) to solve min X∈R d R(X), (2) . Let X 0 ∈ R d be an initial vector (random or non-random). Then SGD+AL iterates by selecting a new data point (a k+1 , ϵ k+1 ) such that a k+1 ∼ N (0, K) and ϵ k+1 ∼ N (0, ω 2 ) and makes the updateEQUATIONwhere g k > 0 is a learning rate (see assumptions below) 5 . To perform our analysis, we place the following assumption on the initialization X 0 and signal X ⋆ . Assumption 5 (Initialization and signal). The initialization point X 0 and the signal X ⋆ are bounded independent of d, that is, max{∥X 0 ∥, ∥X ⋆ ∥} ≤ C for some C independent of d.Adaptive learning rate. Our analysis requires some mild assumptions on the learning rate. To this end, we define a learning rate function γ : In this definition, for functions taking integer arguments, we extend them to real-valued inputs by first taking the floor function of its argument. Note that the adaptive learning rates can depend on the whole history of stochastic iterates (N k ), gradients (G k ), and risk (Q k ) via this definition.EQUATIONWe also define a conditional expectation version of G k where the filtration F k = σ(X ⋆ , X 0 , . . . , X k ):G k (t) def = 1 {t<k} (•) 1 d E[∥∇ X Ψ(X t ; a t+1 , ϵ t+1 )∥ 2 |F t ] for t ≥ 0.With this, we impose the following learning rate condition. Assumption 6 (Learning rate). The learning rate function γ : R+ × D([0, ∞)) × D([0, ∞)) × D([0, ∞)) → R is α-pseudo-Lipschitz with constant L(γ) (independent of d) in D([0, ∞)) × D([0, ∞)) × D([0, ∞)).Moreover, for some constant C = C(γ) > 0 independent of d and δ > 0,E [|γ(k, f, G k (d ×•), q) -γ(k, f, G k (d ×•), q)| | F k ] ≤ Cd -δ (1 + ∥f ∥ α ∞ + ∥q∥ α ∞ ) w.o.p. (5)Finally, γ is bounded, i.e., there exists a constant Ĉ = Ĉ(γ) > 0 independent of d so thatEQUATIONThe inequality (5) ensures that the learning rate concentrates around the mean behavior of the stochastic gradients. Many well-known adaptive stepsizes satisfy (4) and Assumption 6 including AdaGrad-Norm, DoG, D-Adaptation, and RMSProp (see  
deterministic dynamics for sgd with adaptive learning rates: 
 Intuition for deriving dynamics: The risk R(X) and Fisher matrix can be evaluated solely in terms of the covariance matrix B. Thus, to know the evolution of the risk over time, it would suffice to know the evolution of B. Alas, except in the isotropic case where K is a multiple of the identity, the evolution of B is not autonomous (i.e., its time evolution depends on other unknown variables). However, if we let (λ i , ω i ) be the eigenvalues and corresponding orthonormal eigenvectors of K, we can consider projectionsV i (X k ) = d • W T k ω i ω T i W k, and it turns out that these behave autonomously.Example: Least Squares. One canonical example of (2) is least squares, where we aim to recover the target X ⋆ given noisy observations ⟨a, X ⋆ ⟩ + ϵ. In this case, the least squares problem is minEQUATIONThe pair of functions h (Assumption 3) and I (Assumption 4) can be evaluated simply:h(B(W )) = 1 2 I(B(W )) = 1 2 (X -X ⋆ ) T K(X -X ⋆ ) + 1 2 ω 2 .The deterministic dynamics for the risk R(t) in this case can be simplified to:R(t) = 1 2 (X 0 -X ⋆ ) T Ke -2K t 0 γs ds (X 0 -X ⋆ ) + 1 2 ω 2 + 1 d t 0 γ 2 s Tr(K 2 e -2K t s γτ dτ )R(s) ds.This is a convolution Volterra equation with a convergence threshold of γ t < 2d TrK [14, 44, 46, 47] . In the noiseless label case (i.e., ϵ = 0), the risk is given by R(t) = 1 2d d i=1 λ i D 2 i (t).Using the ODEs in (9), we get the following deterministic equivalent ODE for the D 2 i 's:EQUATIONWe will perform a deep analysis of the dynamics of the learning rate on least squares (7) , which will generalize to settings where the outer function f is strongly convex (see D.1).Deterministic dynamics. To derive deterministic dynamics, we make the following change to continuous time by setting k iterations of SGD = ⌊td⌋, where t ∈ R is the continuous time parameter.This time change is necessary, as when we scale the size of the problem, more time is needed to solve the underlying problem. This scaling law scales SGD so all training dynamics live on the same space. One can solve a smaller d problem and scale it to recover the training dynamics of the larger problem. 7We now introduce a coupled system of differential equations, which will allow us to model the behaviour of our learning algorithms. For the ith (λ i , ω i )-eigenvalue/eigenvector of K, setV i (t) def = V 11,i (t) V 12,i (t) V 12,i (t) V 22,i (t)and averaging over i, B(t) def = 1 d d i=1 λ i V i (t).The V i (t) and B(t) are deterministic continuous analogues of V i (X td ) and B(X td ) respectively. Define the following continuous analogues∇h(B(t)) def = H 1,t H 2,t H 2,t H 3,t , N (t) def = 1 d d i=1 V i (t), R(t) def = h(B(t)), I (t) def = I(B(t)),and finallyγ t def = γ(t, 1 {•≤t} N (•), Tr(K) d 1 {•≤t} I (•), 1 {•≤t} R(•)). We now introduce a system of coupled ODEs for each (λ i , ω i )-eigenvalue/eigenvector pair of KEQUATIONwith the initialization of V i (0) given by V i (X 0 ). We finally state the deterministic dynamics for the risk and learning rate. Theorem 2.1. Under Assumptions 1, 2, 3, 4, 5, 6, then for any ε ∈ (0, 1 2 ) and anyEQUATIONThe same statements hold comparing W T td W td to N (t) and W T td KW td to B(t).In fact, we can derive deterministic dynamics for a large class of statistics which are linear combinations of V (t) and functions thereof (See Theorem B.1, and Corollary B.1).One important corollary is a deterministic limit for the distance to optimality, D 2 (X k ) = ∥X k -X ⋆ ∥ 2 , which is a quadratic form of W T k W k and hence covered by Thm. 2.1. The equivalent deterministic dynamics areEQUATIONwhereD 2 i (t) corresponds D 2 i (X k ) def = d × (⟨X k -X ⋆ , ω i ⟩) 2 .I(B) = E a,ϵ [(f ′ (⟨a, X⟩; ⟨a, X ⋆ ⟩, ϵ)) 2 ] satisfies an estimate |I(B 1 ) -I(B 2 )| ≤ L(f ) I(B 1 ) + I(B 2 ) min{∥B -1 1 ∥ op , ∥B -1 2 ∥ op }∥B 1 -B 2 ∥ F . Proof. To derive the existence of h, note that R(X) = E(E(f (⟨a, X⟩, ⟨a, X ⋆ ⟩, ϵ)|ϵ)) is an expectation of a Gaussian vector r = (⟨a, X⟩, ⟨a, X ⋆ ⟩). This vector can be expressed as an image of an iid Gaussian vector z by representing r = 
idealized exact line search and polyak stepsize: 
 In this section, we consider two classical idealized algorithms -exact line search and Polyak stepsize.In deterministic optimization, these learning rate strategies are chosen so that the function value (exact line search) or distance to optimality (Polyak) produces the largest decrease in function value (resp. distance to optimality) at the next iteration. For stochastic algorithms, we can ask this to hold for the deterministic equivalent to the risk R(t) (resp. distance to optimality, D(t)) since we know that SGD is close to these deterministic equivalents. Thus, the question is: what choice of learning rate decreases the R(t) (exact line search) and/or D(t) (Polyak stepsize)? We will restrict to least squares in this section -see Appendix F.1 and F.2 for general functions as well as proofs for least squares. These are idealized algorithms because we can not implement them as they require distributional knowledge of a or X ⋆ . Despite this, they provide a basis for more practical algorithms. 13) in high-dimensional settings, where a broader spectrum of K results in λmin(K)1 d Tr(K 2 ) ≪ 1 1 d Tr(K), indicating slower risk convergence and poorer performance of exact line search (unmarked) as it deviates from the Polyak stepsize (circle markers) . The gray shaded region demonstrates that equation ( 13) is satisfied. See Appendix H for simulation details.Polyak Stepsize. A natural threshold to consider is the largest learning rate such that dD(t) < 0, which we denote by γD t . Using the least squares ODE (8) , this is preciselyγD t = (2R(t)-ω 2 ) Tr(K) d R(t)and ḡD k = (2R(X k )-ω 2 )EQUATIONWithout label noise, (12) EQUATIONTr(K)/d , the exact threshold for convergence of least squares.A greedy stepsize strategy would maximize the decrease in the distance to optimality at each iteration, denoted by us as Polyak stepsize, γ Polyak t ∈ arg min γ dD(t). In the case of least squares, this isγ Polyak t = 1 2 γD t and g Polyak k = 1 2 ḡDk . The latter yields the optimal fixed learning rate (up to absolute constant factors) for a noiseless target on a least squares problem [35, 43] ). 8Exact Line Search. In the context of risk, using (8) and noting that R(t) = 1 2d d i=1 λ i D 2 i (t), we can find γ line t ∈ arg min dR(t); i.e., the greedy learning rate that decreases the risk the most in the next iteration. We call this exact line search. Expressions for the learning rates are given in Table 2 , (c.f. Appendix F.1 for general losses). Because these come from ODEs, we can use ODE theory to give exact limiting values for the deterministic equivalent of g line k . Proposition 3.1. [Limiting learning rate; line search on noiseless least squares] Consider the noiseless (ω = 0) least squares problem (7) . Then the learning rate is always lower bounded byλmin(K) 1 d Tr(K 2 ) ≤ γ line t for all t ≥ 0.Moreover, suppose K has only two distinct eigenvalues λ 1 > λ 2 > 0, i.e., K has d/2 eigenvalues equal to λ 1 eigenvalues and d/2 eigenvalues equal to λ 2 . ThenEQUATION)For a proof and explicit formula for lim t→∞ γ line t , see Section F.2. Hence, being greedy for the risk in a sufficiently anisotropic setting will badly underperform Polyak stepsize (see Fig. 2 ).  
adagrad-norm analysis: 
 In this section, we analyze the behavior of AdaGrad-Norm learning rate in the least squares setting (see Sec. D for general strongly convex functions). In the presence of additive noise, the AdaGrad-Norm learning rate decays like t -1/2 , regardless of the data covariance K. In contrast, the model with no noise exhibits a learning rate that depends on the spectrum of K, as illustrated in Figure 3 . The learning rate is bounded below by a constant when λ min (K) > 0 is fixed as d → ∞, and we quantify this lower bound. If the limiting spectral measure of K has unbounded density near 0 (e.g. power law spectrum), then the learning rate can approach zero and we quantify the rate of this convergence in the least squares setting as a function of spectral parameters.For least squares with additive noise, the learning rate asymptoticγ t ≍ η/(b 2 + ω 2 d Tr(K)t) (1/2)is the fastest decay that AdaGrad-Norm can exhibit. In contrast, the propositions below concern the noiseless case where, for various covariance examples, the decay rate of γ t changes. This is tightly connected to whether the risk is integrable or not. In the simple case of identity covariance, we obtain a closed formula for the trajectory of the integral of the risk and therefore also the learning rate. Proposition 4.1. In the case of identity covariance (K = I d ), the risk solves the differential equationEQUATIONThe solution t 0 R(s) ds approaches (from below) a positive constant which yields a computable lower bound to which γ t will converge. Generalizing this to a broader class of covariance matrices, we get the next proposition, which captures the dependence of γ t on Tr(K). Proposition 4.2. Suppose 1 d Tr(K) ≤ b/η, and that∞ 0 R(s)γ s ds < ∞ with γ s as in Table 2 (AdaGrad-Norm for least squares), then γ t ≍ 1 b η + η 2 4d Tr(K)D 2 (0).An analog of Proposition 4.2 for the strongly convex setting appears in Sec. D (see Prop. D.1). We now consider two cases in which, as d → ∞, there are eigenvalues of K arbitrarily close to 0. Proposition 4.3. Assume that, for some C > 0, the number of eigenvalues ofK below C is o(d), and that ⟨X ⋆ , ω i ⟩ = O(d -1/2) for all i, (i.e. X ⋆ is not concentrated in any eigenvector direction). Then, with the initialization X 0 = 0, there exists some γ > 0 such that γ t > γ for all t > 0.Proposition 4.4. Let K have a spectrum that converges as d → ∞ to the power law measure ρ(λ) = (1 -β)λ -β 1 (0,1) , for some β < 1 9 , and suppose that D 2 i (0) ∼ λ -δ i for δ ≥ 0. Then:• For 1 > β + δ, there exists γ such that γ t ≥ γ, and R(t) ≍ δ,β t β+δ-2 for all t ≥ 1.• For 1 < β + δ < 2, γ t ≍ δ,β t -1+ 1 β+δ δ,β, and R(t) ≍ δ,β t -2 β+δ +1 for all t ≥ 1. •For 1 = β + δ, γ t ≍ δ,β 1 log(t+1) , and R(t) ≍ δ,β ( t log(t+1) ) -1 for all, t ≥ 1.This proposition shows non-trivial decay of the learning rate is dictated by the residuals (distance to optimality at initialization) and the spectrum of K. We note that δ = 0 corresponds to uniform contribution of each mode (e.g. X 0 normally distributed). As the eigenmodes of the residuals become more localized, the decay of the learning rate is closer to the behaviour in the presence of additive noise. Furthermore, the scaling behaviour of the loss is affected by the structure of the AdaGrad-Norm algorithm (see Fig. 4 ). Lastly, constant stepsize SGD yields R(t) ≍ t β+δ-2 , with no transition occurring at β + δ = 1.Proofs of the above propositions, in a slightly more general setting, are deferred to Sec. D. 
conclusions and limitations: 
 This work studies stochastic adaptive optimization algorithms when data size and parameter size are large, allowing for nonconvex and nonlinear risk functions, as well as data with general covariance structure. The theory shows a concentration of the risk, the learning rate and other key functions to a deterministic limit, which is described by a set of ODEs. The theory is then used to derive the asymptotic behavior of the AdaGrad-Norm and idealized exact line search on strongly convex and least square problems, revealing the influence of the covariance matrix structure on the optimization. A potential extension of this work would be to study other adaptive algorithms such as D-adaptation, DOG, and RMSprop which are covered by the theory. Studying the asymptotic behavior of the risk and the learning rate may improve our understanding of the performance and scalability of these algorithms on more realistic data. Another important application of the theory would be to analyze the ODEs presented here on nonconvex problems.The current form of the theory is limited to Gaussian data, though many parts of the proof can be extended easily beyond Gaussian data. The main ODE comparison theorem is also only tuned for analyzing problem setups where the trace of the covariance is on the order of the ambient dimension; when the trace of the covariance is much smaller than ambient dimension, other stepsize scalings of SGD are needed. In addition, the analysis is limited to the streaming stochastic adaptive methods. We conjecture that a similar deterministic equivalent holds also for multi-pass algorithms at least for convex problems. This has already been shown in the least square problem for SGD with a fixed deterministic learning rate [43, 45] . Lastly, numerical simulations on real datasets (e.g., CIFAR-5m) suggests that the predicted risk derived by our theory matches the empirical risk of multipass SGD beyond Gaussian data (see for example Figure 6 ). 
 
 Appendix: Broader Impact Statement. The work presented in this paper is foundational research and it is not tied to any particular application. The set-up is on a simple well-studied high-dimensional linear composites (e.g., least squares, logistic regression, phase retrieval) with synthetic data and solved using known algorithms, e.g., AdaGrad-Norm. We present deterministic dynamics for the training loss and adaptive stepsizes. The results are theoretical and we do not anticipate any direct ethical and societal issues. We believe the results will be used by machine learning practitioners and we encourage them to use it to build a more just, prosperous world.In this section, we write down the explicit update rules for 2 different adaptive stochastic gradient descent algorithms.Example: AdaGrad-Norm. We begin with AdaGrad-Norm (see Algorithm 1). Note by unraveling the recursion, we have thatEQUATION)with the deterministic equivalent (see Section 2 and also C.3) for this learning rate beingEQUATIONIn the case of the least squares problem, the quantity I(B(t)) is explicit andEQUATION) Algorithm 1 AdaGrad-Norm Require: Initialize η > 0, X 0 ∈ R d , b ∈ R and set b 0 = b × d for k = 1, 2, . . . , do Generate new sample a k ∼ N (0, K), ϵ k ∼ N (0, ω 2 ); b 2 k ← b 2 k-1 + ∥∇ X Ψ(X k-1 ; a k , ϵ k )∥ 2 ; g k-1 = d × η |b k | ; ▷ updating learning rate X k ← X k-1 -g k-1 d ∇ X Ψ(X k-1 ; a k , ϵ k );▷ updating step with stochastic gradient end for Example: RMSprop-Norm We consider the "normed" version of RMSprop, that is, where there is only one learning rate parameter.We consider Algorithm 2 where we put a factor of the learning into the exponential moving average for RMSprop. The deterministic equivalent for g k for Alg. 2 (see Section 2) isEQUATION)In the case of the least squares problem, the quantity I(B(t)) is explicit andEQUATION)Algorithm 2 RMSprop-Norm, α Exponential Moving AverageRequire: Initialize η > 0, X 0 ∈ R d , b ∈ R and set b 0 = d × b, α > 0 exponential moving avg. g -1 = d × η b0 ; for k = 1, 2, . . . , do Generate new sample a k ∼ N (0, K), ϵ k ∼ N (0, ω 2 ); b 2 k ← α • b 2 k-1 + (1 -α)∥∇ X Ψ(X k-1 ; a k , ϵ k )∥ 2 ; g k-1 = d × η |b k | ; ▷ updating learning rate X k ← X k-1 -g k-1 d ∇ X Ψ(X k-1 ; a k , ϵ k );▷ updating step with stochastic gradient end forIn this section, we prove the main theorem on concentration of the risk curves and learning rates. We shall set some notation. In what follows, we again use W = [X|X ⋆ ] ∈ R d×2 . We also useW + = [W |X 0 ] = [X|X ⋆ |X 0 ].We shall also use the shorthand r = ⟨a, W ⟩, and x = ⟨a, X⟩ so that f (⟨a, X⟩, ⟨a,X ⋆ ⟩; ϵ) = f (⟨a, W ⟩; ϵ) = f (r; ϵ).We shall let B = B(X) = W T KW be the covariance matrix of the Gaussian vector r. We also write f ′ for the ∂ x f .In this section we show how the assumptions we put on h and I are almost satisfied for L-smooth f . We say that f is L-smooth if:∥∇f (r 1 , ϵ 1 ) -∇f (r 2 , ϵ 2 )∥ ≤ L (∥r 1 -r 2 ∥ 2 + ∥ϵ 1 -ϵ 2 ∥ 2 ), which we note implies f is α-pseudo Lipschitz with α = 1. Lemma B.1.1. There exists a function h : R 2×2 → R such that h(B(X)) = R(X) is differentiable and satisfies ∇ X R(X) = E a,ϵ ∇ X Ψ(X; a, ϵ). Furthermore, h is continuously differentiable on {B : det B ̸ = 0} and its derivative ∇h satisfies an estimate∥∇h(B 1 ) -∇h(B 2 )∥ ≤ ( √ 2 + 1)L(f ) min{∥B -1 1 ∥ op , ∥B -1 2 ∥ op }∥B 1 -B 2 ∥ F .Bz, and hence we haveh(B) def = E(E(f ( √ Bz, ϵ)|ϵ)).As the function f is absolutely continuous with a Lipschitz gradient, we can differentiate under the integral sign and conclude∇ X R(X) = ∇ X E f (⟨a, X⟩, ⟨a, X ⋆ ⟩, ϵ) = E ∇ X f (⟨a, X⟩, ⟨a, X ⋆ ⟩, ϵ).For the differentiability of h, suppose for the moment that f is C 2 with bounded second derivatives. 10Setting Q =B the positive semi-definite square root of B, we have∂ Qij h(Q 2 ) = E(E(∂ Qij f (Qz, ϵ)|ϵ)).Then using the chain rule, and setting ∂ i f to be the i-th partial derivative of f ,∂ Qij h(Q 2 ) = E(E(z j ∂ i f (Qz, ϵ)|ϵ)) = E(E([Q ij ∂ i + Q jj ∂ j ]∂ i f (Qz, ϵ)|ϵ)), where we have applied Stein's Lemma. We conclude when det Q ̸ = 0 by the implicit function theorem that h is differentiable and we have∂ Qij h(Q 2 ) = ∂ kl h∂ Qij (Q 2 ) kl = l (∂ il h)Q jl + k (∂ kj h)Q ik .As a matrix equation, this can be written as (Dh)Q + Q(Dh) = JQ where J kl = E(E((∂ k ∂ l f )(Qz, ϵ)|ϵ)).(Dh)Q + Q(Dh) = 1 2 (JQ + QJ),and so taking inverses on both sides, Dh = J.Undoing Stein's Lemma, we haveQ(Dh) = (Dh)Q = M, where M ij = E(E(z j ∂ i f (Qz, ϵ)|ϵ)). From L-smoothness of f ∥M (Q 1 ) -M (Q 2 )∥ ≤ L E(∥z∥∥Q 1 z -Q 2 z∥) ≤ √ 2L∥Q 1 -Q 2 ∥ F . Hence ∥Dh(Q 2 1 ) -Dh(Q 2 2 )∥ = ∥Q -1 1 M (Q 1 ) -Q -1 2 M (Q 2 )∥ ≤ ∥Q -1 1 ∥ op ∥M (Q 1 ) -Q 1 Q -1 2 M (Q 2 )∥ ≤ ∥Q -1 1 ∥ op ∥M (Q 1 ) -M (Q 2 )∥ + ∥(Q 2 -Q 1 )Q -1 2 M (Q 2 )∥ . Note Q -1 2 M (Q 2 ) = (Dh)(Q 2 2) is bounded by L(f ), and so we arrive at∥Dh(Q 2 1 ) -Dh(Q 2 2 )∥ ≤ ( √ 2 + 1)L(f )∥Q -1 1 ∥ op ∥Q 1 -Q 2 ∥ F ≤ ( √ 2 + 1)L(f )∥Q -2 1 ∥ op ∥Q 2 1 -Q 2 2 ∥ F .We note the bound is symmetric in Q 1 and Q 2 , and by density of C 2 in space of C 1,lip , this holds for L-smooth f . This concludes the estimates for the derivative of h.For the Fisher matrix, I(B), from L-smoothness, we have again withQ = √ B, I(Q 2 ) = E(E((∂ 1 f (Qz, ϵ)) 2 |ϵ)). Then |I(Q 2 1 ) -I(Q 2 2 )| ≤ E(E((∂ 1 f (Q 1 z, ϵ)) 2 -(∂ 1 f (Q 2 z, ϵ)) 2 |ϵ)). Applying Cauchy-Schwarz and using the L-smoothness of f ,|I(Q 2 1 ) -I(Q 2 2 )| ≤ I(Q 2 1 ) + I(Q 2 2 ) × L(f )∥Q 1 -Q 2 ∥ F .This lemma shows that an L-smooth function nearly satisfies Assumption 3 and 4 provided that ∥B -1 ∥ op is bounded. Therefore, our concentration result Theorem B.1 and its Corollaries will hold provided we add a stopping time. Fix M > 0 and letℏ M (B) def = inf{t > 0 : ∥B -1 ∥ op > M }.Then the concentration of the risk under SGD to a deterministic function, Theorem B.1, holds with t replaced with t ∧ ℏ M (B) ∧ ℏ M (B). The corollaries of Theorem B.1 also follow under this added stopping time.In the next section, we prove this concentration theorem, Theorem B.1.A goal of this paper is to show that quadratic statistics φ : R d → R applied to SGD converge to a deterministic function. This argument hinges on understanding the deterministic dynamics of one important statistic, defined as S(W, z) = W ⊤ R(z; K)W, applied to W ⌊td⌋ (SGD updates). Here W = [X|X ⋆ ] and R(z; K) = (K -zI d ) -1 for z ∈ C is the resolvent of the matrix K. The statistic S(W, z) is valuable because it encodes many other important quantities including W ⊤ q(K)W for all polynomials q. We show that S(W ⌊td⌋ , z), is close to a deterministic function (t, z) → S (t, z) which satisfies an integro-differential equation.To introduce the integro-differential equation, recall by Assumptions 3 and 4R(X) = h • B(W ) and E a,ϵ [f ′ (a ⊤ W ) 2 ] = I • B(W ) with B(W ) = W ⊤ KW,and α-pseudo-Lipschitz functions h : R 2×2 → R differentiable and I : R 2×2 → R. It will be useful, throughout the remaining paper, to express ∇h explicitly as a 2 × 2 matrix, that is,∇h ∼ = ∇h 11 ∇h 12 ∇h 21 ∇h 22 .With these recollections, the integro-differential equation is defined below.Integro-Differential Equation for S (t, z). For any contour Ω ⊂ C enclosing the eigenvalues of K, we have an expression for the derivative of S :EQUATIONwhereEQUATION-γ t (S (t, z)(2zH(B(t))) + (2zH T (B(t)))S (t, z)).HereB(t) = -1 2πi Ω zS (t, z) dz, H(B) = ∇h 11 (B) 0 ∇h 21 (B) 0 , γ t is defined in (9), and the initialization isEQUATIONThe functions h : R 2×2 → R and I : R 2×2 → R are defined in Assumption 3 and Assumption 4, respectively.We first note that there is an actual solution to the integro-differential equation. This solution is the same as the ODEs defined in the introduction (see (9) ) and proved in [ S (t, z) = 1 d d i=1 1 λ i -z V i (t).In this section, we will be working with approximate solutions to the integro-differential equation (20) (see below for specifics). For working with these solutions, we introduce some notation. We shall always work on a fixed contour Ω surrounding the spectrum of K, given by Ω def = {z : |z| = max {1, 2∥K∥ op }}. We note that this contour is always distance at least 1 2 from the spectrum of K. We define a norm, ∥ • ∥ Ω , on a continuous functionEQUATION)Definition B.1 ((ε, M, T )-approximate solution to the integro-differential equation). For constants M, T, ε > 0, we call a continuous functionS : [0, ∞) × C → R 2×2 an (ε, M, T )-approximate solution of (20) if with τM (S) def = inf t ≥ 0 : ∥S(t, •)∥ Ω > M , then sup 0≤t≤(τ M ∧T ) S(t, •) -S(0, •) - t 0 F(•, S(s, •)) ds Ω ≤ ε and S(0, •) = W ⊤ 0 R(•, K)W 0 , where W 0 = [X 0 |X ⋆ ]is the initialization of SGD. We suppress the S in the notation for τM , that is, τM = τM (S), when the function S is clear from the context.We are now ready to state and prove one of our main results. Theorem B.1 (Concentration of SGD and deterministic function S (t, z)). Suppose the risk function R(X) (2) satisfies Assumptions 2, 3, and 4. Suppose the learning rate satisfies Assumption 6, and the initialization X 0 and hidden parameters X ⋆ satisfy Assumption 5. Moreover the data a ∼ N (0, K) and label noise ϵ satisfy Assumption 1. Let {W ⌊td⌋ } be generated from the iterates of SGD. Then there is an ε > 0 so that for any T, M > 0 and d sufficiently large, with overwhelming probabilityEQUATION)where the deterministic function S (t, z) solves the integro-differential equation (20) .Proof. By Proposition C.1, for any M and T , we can find a ε > 0 such that the function S(W td , z) is an (d -ε , M, T )-approximate solution. (For the deterministic function S , it is an (0, M, T )approximate solution by definition.) We now apply the stability result, [15, Prop. 4 .1], to conclude that there exists a ε > 0 such thatEQUATION)where τM is shorthand for τM (S(W, •)) ∧ τM (S ). The result immediately follows.Corollary B.1. Suppose the assumptions of Theorem B.1 hold. Let f be an α-pseudo-Lipschitz function with α ≤ 1 and let q be a polynomial. Setφ(X) def = f (W T q(K)W ), ϕ(t) def = f -1 2πi Ω q(z)S (t, z) dz , where S (t, z) solves (20).Then there is an ε > 0 such that for d sufficiently large, with overwhelming probability,sup 0≤t≤T |φ (X td ) -ϕ(t)| ≤ d -ε .Proof. This is basically equivalent to [ , and the rest of the proof is identical to the one in [15] .Remark B.1. The learning rate g k , technically, is not a function of W T q(K)W . However, Assumption 6 ensures that the learning rate concentrates around a function W T q(K)W . Therefore, Corollary B.1 applies to the learning rate.We introduce a rescaling of time to relate the k-th iteration of SGD to the continuous time parameter t in the differential equation through the relationship k = ⌊td⌋. Thus, when t = 1, SGD has done exactly d updates. Since the parameter t is continuous and the iteration counter k (integer) discrete, to simplify the discussion below, we extend k to continuous values through the floor operation, X k def = X ⌊k⌋ . Using the continuous parameter t, the iterates are related by X td = X ⌊td⌋ .The paper [15] provides a net argument showing that we do not need to work with every z on the contour Ω defining the integro-differential equation, but only polynomially many in d. Recall that Ω = {z : |z| = max{2∥K∥ op , 1}}. For a fixed ξ > 0, we say that Ω ξ is a d -ξ -mesh of Ω if Ω ξ ⊂ Ω and for every z ∈ Ω there exists a z ∈ Ω ξ such that |z -z| < d -ξ . We can achieve this with EQUATION)with τM = inf{t ≥ 0 : ∥S(t, •)∥ Ω > M }. Then S is a (ε + C(M, T, ∥K∥ op )d -ξ , M, T )- approximate solution to the integro-differential equation, that is, sup 0≤t≤(τ M ∧T ) ∥S(t, •) -S(0, •) - t 0 F(•, S(s, •)) ds∥ Ω ≤ ε + C • d -ξ ,where C = C(M, T, ∥K∥ op , L(I), L(h)) is a positive constant.(We prove in Section C.1 that S(t, z) does indeed satisfy inequality (26) .) We also cite the following lemma, which relates two stopping times used throughout this paper. Lemma C.2 (Stopping time, [15] , Lemma 4.2). For a constant C depending on ∥K∥ op , we haveC ≤ ∥S(W td , •)∥ Ω ∥W td ∥ 2 ≤ 2.Remark C.1. Fix M > 0 and define the stopping time on ∥W td ∥, ϑ = ϑ M , byϑ M (W td ) def = inf t ≥ 0 : ∥W td ∥ 2 > M .Due to the previous lemma, any stopping time τM defined on ∥S(t, •)∥ Ω corresponds to a stopping time ϑ on ∥W td ∥, that is, for c = C -1 , τM ≤ ϑ cM .Proposition C.1 (SGD-AL is an approximate solution). Fix a T, M > 0 and 0 < ε < δ/8, where δ is defined in Assumption 6. Then S(W td , z) is a (d -ε , M, T )-approximate solution w.o.p., that is,EQUATIONAgain, the proof is very similar to [15, Prop. 5.2] . The one difference is that the martingales and error terms are slightly more involved, because of the non-deterministic stepsize we are using. The remainder of this section, along with section C.2, fills in the details of bounding these lower-order terms, so that the proof can proceed as in [15] .In the following sections, we will be using various versions of the stepsize γ. In order to simplify notation, we setγ(G k ) = γ(k, N k (d × •), G k (d × •), Q k (d × •)), γ(G k ) = γ(k, N k (d × •), G k (d × •), Q k (d × •)), γ(B k ) = γ(k, N k (d × •), Tr(K)I(B k (d × •))/d, Q k (d × •)). Further, setting ∆ k def = f ′ (r k )a k+1 , define I 1 (k) def = ∆ ⊤ k ∇ 2 φ(X k )∆ k /d, I 2 (k) def = Tr(∇ 2 φ(X k )K)E [f ′ (r k ) 2 | F k ]/d, I 3 (k) def = ∇φ(X k ) ⊤ ∆ k .The normalization here (dividing by d) is chosen so that the I terms are all O(1); this is formally shown in Lemma C.5.We follow the approach in [15, Section 5.3] to rewrite the SGD adaptive learning rate update rule as an integral equation. Considering a quadratic function φ : R d → R and performing Taylor expansion, we obtainEQUATIONWe will now relate this equation to its expectation by performing a Doob decomposition, involving the following martingale increments and error terms:EQUATIONEQUATIONEQUATIONEQUATIONWe can then writeφ(X k+1 ) = φ(X k ) - γ(B k ) d ∇φ(X k ) ⊤ ∇R(X k ) + γ(B k ) 2 2d 2 Tr(∇ 2 φ(X k )K)E [f ′ (r k ) 2 | F k ] + ∆M grad k (φ) + ∆M Hess k (φ) + E [E Hess k (φ) | F k ] + E [E grad k (φ) | F k ].Extending X k into continuous time by defining X t = X ⌊t⌋ , we sum up (integrate). For this, we introduce the forward difference (∆φ)(X j ) def = φ(X j+1 ) -φ(X j ), giving us φ(X td ) = φ(X 0 ) + ⌊td⌋-1 j=0 (∆φ)(X j ) def = φ(X 0 ) + t 0 d • (∆φ)(X sd ) ds + ξ td ,EQUATIONEQUATION+ E [E Hess j (φ) | F j ] + E [E grad j (φ) | F j ] + ξ td (φ).From here, we can proceed as in [15, Section 5.3] to show that SGD-AL is an (ε, M, T )-approximated solution.Proof of Proposition C.1. The appropriate stepsize, as a function of W td , isγ t = γ(td, N td , Tr(K)I(B td )/d, Q td ).(Note that N , I and Q can all be found as functions of S(W td , •) using contour integration.) It is shown in the proof of [15, Proposition 5.2] that given the analogue of (33) for deterministic stepsize, S(W td , •) satisfiesS (W td , z) = S (W 0 , z) + t 0 F (z, S (W sd , z)) ds + ⌊td⌋-1 i=0 E all j (S).The only terms of (33) that differ in our case are the martingale and error terms. Thus to show that S(W td , •) is an approximate solution of the integro-differential equation (20) all we need is to bound the martingales and error terms contained in E all j . Let Ω = {z : |z| = max{1, 2∥K∥ op }}, as previously. We thus have that for all z ∈ Ω, sup EQUATIONEQUATION)We deduce thatsup 0≤t≤T ∧τ M ∥S(W td , z) -S(W 0 , z) - t 0 F(z, S(W sd , z)) ds∥ Ω ξ ≤ d δ-δ/4 w.o.p.An application of the net argument, Lemma C.1, finishes the proof after setting δ = δ/8 and ξ = δ/8.All the martingale and error terms (34) go to 0 as d grows. Formally, Proposition C.2. Let the function f be defined as in Assumption 2. Let the statistic S : [0, ∞) × C → R 2×2 be defined asEQUATIONwhere W = [X|X ⋆ ]. Then, for any z ∈ Ω and T, M, ζ > 0, with overwhelming probability,sup 0≤t≤T ∧ϑ E all dt (S(•, z)) ≤ d -δ/4+ζ ,where to suppress notation we use ϑ as shorthand for ϑ cM , and c is the constant from Lemma C.2.Proof. This follows from combining Propositions C. The remainder of this subsection is devoted to proving these supporting propositions; throughout these proofs we will work with the stopping time ϑ as defined in the proposition above.Proposition C.3 (Hessian error term). Let f and S be defined as in Assumption 2 and (37). Then, for any z ∈ Ω, T > 0 and ζ > 0, with overwhelming probability,sup 0≤t≤T ∧ϑ ⌊td⌋-1 k=0 E E Hess k (S(•, z)) | F k ≤ d -δ/4+ζ .Proof. For arbitrary z ∈ Ω and k ≤ (T ∧ ϑ)d -1, set φ(X) = S ij (W, z) to be the ij-th entry of the matrix S(W, z). Then2d E[E Hess k (φ) | F k ] = E γ(G k ) 2 I 1 (k) | F k -γ(B k ) 2 I 2 (k) = E[(γ(G k ) 2 -γ(G k ) 2 )I 1 (k) | F k ] + (γ(G k ) 2 -γ(B k ) 2 ) E[I 1 (k) | F k ] + γ(B k ) 2 E[(I 1 (k) -I 2 (k) | F k ] = E 1 + E 2 + E 3 . We look at |E 1 | first. |E 1 | = E (γ(G k ) 2 -γ(G k ) 2 )I 1 (k) | F k ≤ E (γ(G k ) 2 -γ(G k ) 2 ) 2 | F k 1 2 • E I 1 (k) | 2 F k 1 2 ≤ E |γ(G k ) + γ(G k )| 7 2 |γ(G k ) -γ(G k )| 1 2 | F k 1 2 • E I 1 (k) | 2 F k 1 2 ≤ E |γ(G k ) + γ(G k )| 7 | F k 1 4 • E [|γ(G k ) -γ(G k )| | F k ] 1 4 • E |I 1 (k)| 2 | F k 1 2 .For the first term, we use (6) . We haveE |γ(G k ) + γ(G k )| 7 | F k ≤ Ĉ(γ) • E |2 + 2∥N k ∥ α ∞ + 2∥Q k ∥ α ∞ + ∥G k ∥ α ∞ + ∥G k ∥ α ∞ | 7 | F k .All the terms inside the expectation, apart from ∥G k ∥ α ∞ , are deterministic with respect to F k and bounded by a constant independent of d (see Lemma C.6). Since we know from Lemma C.6 that for any ε > 0, all moments of ∥G k ∥ ∞ are bounded by d ε w.o.p., we concludeE |γ(G k ) + γ(G k )| 7 | F k ≤ d ε w.o.p.For the second term, we use (5) . Again, since ∥N k ∥ ∞ and ∥Q k ∥ ∞ are bounded due to our stopping time, we haveE [|γ(G k ) -γ(G k )| | F k ] 1 4 ≤ d -δ/4 . The last term, E |I 1 (k)| 2 | F k 1 2, is also bounded by a constant (see Lemma C.5), and all together, we find that |E 1 | ≤ d ε-δ/4 with overwhelming probability.2 |: |E 2 | = |(γ(G k ) 2 -γ(B k ) 2 ) E[I 1 (k) | F k ]| = |γ(G k ) + γ(B k )| • |γ(G k ) -γ(B k )| • | E[I 1 (k) | F k ]|.The first term is bounded by (6) , since G k and Tr(K)I(B k )/d are bounded independent of d; the second term is bounded Cd -1 by Lemma C.9, and the last term is bounded by a constant by Lemma C.5.3 |: |E 3 | = γ(B k ) 2 • |E [(I 1 (k) -I 2 (k) | F k ]| .By (6) , the first term is bounded byĈ(γ) 2 (1 + ∥N k ∥ α ∞ + ∥Q k ∥ α ∞ + ∥ Tr(K)I(B k )/d∥ α ∞ ) 2 .All of these terms are bounded by a constant independent of d (because of the stopping time.) The second term satisfies the assumptions of Lemma C.8 with H = ∇ 2 φ(X k ), and is thus bounded by Cd-1 . All together, 2d E[E Hess k (φ) | F k ] ≤ d -δ/4+ε .Summing up to k = T d and dividing through by 2d, we obtain the desired bound.Proposition C.4 (Gradient error term). Let f and S be defined as in Assumption 2 and (37). Then, for any z ∈ Ω, ζ > 0 and T > 0, with overwhelming probability,sup 0≤t≤T ∧ϑ ⌊td⌋-1 k=0 E E grad k (S(•, z)) | F k ≤ d -δ/4+ζ .Proof. We havedE [E grad k | F k ] = -E [γ(G k )⟨∇φ(X k ), ∆ k ⟩ | F k ] + γ(B k )⟨∇φ(X k ), ∇R(X k )⟩ = -E [(γ(G k ) -γ(G k ))I 3 (k) | F k ] -(γ(G k ) -γ(B k )E [I 3 (k) | F k ] = E 1 + E 2 .We then have|E 1 | ≤ E |γ(G k ) -γ(G k )| 2 | F k 1 2 • E |I 3 (k)| 2 | F k 1 2 ≤ E |γ(G k ) + γ(G k )| 3 | F k 1 4 • E [|γ(G k ) -γ(G k )| | F k ] 1 4 • E |I 3 (k)| 2 | F k 1 2 .Just as in the Hessian argument, (6) lets us bound E |γ(G k ) + γ(G k )| 3 | F k 1 4by d ε w.o.p., (5) lets us bound E [|γ(G k ) -γ(G k )| | F k ]E |I 3 (k)| 2 | F k 1 2by a constant, giving an overall bound of|E 1 | ≤ d -δ/4+ε .By the same argument as in the Hessian case,|E 2 | is bounded by Cd -1 ; in conclusion, dE [E grad k | F k ] ≤ d ε-δ/4 .Summing and dividing through by d, we obtain the desired result with ζ = ε.Proposition C.5 (Gradient martingale). Let f and S be defined as in Assumption 2 and (37). Then, for any z ∈ Ω, ζ > 0 and T > 0, with overwhelming probability,sup 0≤t≤T ∧ϑ M grad ⌊dt⌋ (S(•, z)) ≤ d -1/2+ζ .Proof. For notational convenience, set ∆M k = ∆M grad d(k/d∧ϑ) , andF k = -γ(G k )I 3 (k)/d, so that ∆M k = F k -E[F k | F k ]. Set F β k = Proj β (F k ), that is, ensuring F k stays in [-β, β]. Then F β k -E[F β k | F k ] is in [-2β, 2β], and so for the martingale M β k with increments∆M β k = F β k -E[F β k | F k ],Azuma's inequality tells us thatP |M β k | ≥ t ≤ 2 exp -t 2 2 k i=0 (2β) 2 ≤ 2 exp -t 2 2T d(2β) 2 .Set β = d -1+ζ/2 and t = d -1/2+ζ ; this becomesP |M β k | ≥ d -1/2+ζ ≤ 2 exp -d ζ 8T .However, M β k is not quite the martingale we started with: there is still an error term,|M k -M β k | = k i=0 (F k -E[F k | F k ]) -(F β k -E[F β k | F k ]) ≤ k i=0 F k -F β k + E[F k -F β k | F k ] .We bound this term in overwhelming probability. We haveP F k -F β k ̸ = 0 = P (|F k | > β) = P |γ(G k )I 3 (k)/d| > d -1+ζ/2 ≤ P γ(G k ) ≥ d ζ/4 + P |I 3 (k)| ≥ d ζ/4 .The second term is superpolynomially small by Lemma C.5; the first term is superpolynomially small by ( 6) and (C.6).E[F k -F β k | F k ] = E[(F k -F β k )1 {|F k |>β} | F k ] ≤ E[(F k -F β k ) 2 | F k ] 1 2 • E[1 2 {|F k |>β} | F k ] 1 2 ≤ 4 E[F 2 k | F k ] 1 2 • E[1 {|F k |>β} | F k ] 1 2 ≤ 4d -1 E[γ(G k ) 4 | F k ] 1 4 • E[I 3 (k) 4 | F k ] 1 4 • E[1 {|F k |>β} | F k ] 1 2. As before, the first and second expectations are bounded by constants, and the last expectation is just the probability that |F k | > β, which we have already shown is superpolynomially small. So with overwhelming probability, we have|M k -M β k | = k i=0 (F k -E[F k | F k ]) -(F β k -E[F β k | F k ]) ≤ d -1/2+ζ(any power of d would have worked). Combining the error term and the projected martingale, we find that, with overwhelming probability,|M k | ≤ d -1/2+ζ. We can now take the maximum over k from 0 to T d using a union bound; this does not affect the overwhelming probability statement.Proposition C.6 (Hessian martingale). Let f and S be defined as in Assumption 2 and (37). Then, for any z ∈ Ω, ζ > 0 and T > 0, with overwhelming probability,sup 0≤t≤T ∧ϑ M Hess ⌊td⌋ (S(•, z)) ≤ d -1/2+ζ .Proof. The proof here is basically identical to the previous one. Again, setF k = γ(G k ) 2 I 1 (k)/d and F β k = Proj β (F k ), with their associated martingales being M k = F k -E[F k | F k ] and M β k = F β k -E[F β k | F k ].As before, Azuma's inequality, with β = d -1+ζ/2 , gives usP(M β k ≥ d -1/2+ζ ) ≤ 2 exp - d ζ 8T .The error term is also quite similar:|M k -M β k | ≤ k i=0 |F k -F β k | + | E[F k -F β k | F k ]|.We haveP(F k -F β k ̸ = 0) ≤ P(γ(G k ) 2 ≤ d ζ/4 ) + P(|I 2 (k)| ≤ d ζ/4), both of which are superpolynomially small by ( 6) and Lemma C.5. For the expectation, we haveEQUATION; this product is superpolynomially small by ( 6), Lemma C.6, and Lemma C.5. Overall, we have, with overwhelming probability,|M k | ≤ d -1/2+ζ. Taking the supremum, we obtain the desired result.Proposition C.7 (Integral error term). Let f and S be defined as in Assumption 2 and (37). Then, for z ∈ Ω,|ξ td (S(•, z))| ≤ d -1/2 .Proof. We have, as above,|ξ td | = t (⌊td⌋-1)/d d • ∆φ(X sd ) ds ≤ max 0≤j≤⌈td⌉ {|∆φ(X j )|},which is bounded by d -1/2 w.o.p. by the boundedness of I 1 , I 2 , I 3 , and γ(B k ).In this section, we make use of the subgaussian norm ∥ • ∥ ψ2 of a random variable (see [56] for details.) When it exists, this norm is defined asEQUATION)In particular, Gaussian random variables have a well-defined subgaussian norm. Lemma C.3 ([15], Lemma 5.3). There exist constants c, C > 0 such that c∥W ∥ 2 ≤ ∥S(W, z)∥ Ω ≤ C∥W ∥ 2 , ∥∇ X S(W, z)∥ Ω ≤ C∥W ∥, and ∥∇ 2 X S(W, z)∥ Ω ≤ C. Lemma C.4 (Preliminary bounds). With f and ∆ k defined as above, for ε > 0 and λ ≥ 0, we haveEQUATIONEQUATION)Proof of (39) in Lemma C. 4 . By [15, Lemma 3.4], if function f is α-pseudo-Lipschitz with Lipschitz constant L(f ) (as in ( 2)) and the noise ϵ is independent of a, thenEQUATIONNow, since a k+1 is Gaussian, we can writea k+1 = √ Kv k , for a standard normal v k . Then we see that X ⊤ k a k+1 = X ⊤ k √ Kv k is a single-variable Gaussian, with variance |X ⊤ k KX k | ≤ ∥X k ∥ 2•∥K∥ op (bounded independently of d because of the stopping time on X k ). Similarly, ϵ is Gaussian and independent of a k+1 , so the expression ( 41) is bounded w.o.p. by d ε , andE C(α)(L(f ))(1 + |X ⊤ k a k+1 | + |ϵ|) max{1,α} λ F k ≤ C(λ)for some constant C(λ).Proof of (40) in Lemma C.4. We can write a k+1 = √ Kv k , where v k is a standard d-dimensional normal vector. Then, by Hanson-Wright, we haveP ∥a k+1 ∥ 2 -E[∥a k+1 ∥ 2 | F k ] ≥ d = P v ⊤ k Kv k -E[v ⊤ k Kv k | F k ] ≥ d ≤ 2 exp - cd 2 ∥K∥ 2 F + ∥K∥ op d ≤ 2 exp - cd 2 d(∥K∥ op + ∥K∥ 2 op ≤ 2 exp (-Cd) . Now, note that E[v ⊤ k Kv k | F k ] = Tr(K) ≤ d∥K∥ op .Together, we get that ∥a k+1 ∥ 2 ≤ d 1+ϵ with overwhelming probability. Then∥∆ k ∥ 2 d = ∥f ′ (r k )a k+1 ∥ 2 d = ∥a k+1 ∥ 2 f ′ (r k ) 2 d ,which is bounded by d 2ε w.o.p. Now for the expectation:EQUATIONFor the first term, we haveE ∥K∥ op • ∥v k ∥ 2 d 2λ | F k = ∥K∥ 2λ op • E ∥v k ∥ 2 d 2λ | F k ≤ ∥K∥ 2λ op • 1 d d-1 i=0 E ∥v i k ∥ 2 2λ | F k (Jensen's inequality) = ∥K∥ 2λ op • E ∥v 0 k ∥ 4λ | F k , (i.i.d. assumption)where we are using the notation v i k to refer to the ith component of the vector v k . Now, since v 0 k is just a standard Gaussian, all of its moments are bounded. The second term in (42) is bounded by a constant by (39) , as desired.SettingI 1 (k) def = ∆ ⊤ k ∇ 2 φ(X k )∆ k /d, I 2 (k) def = Tr(∇ 2 φ(X k )K)E [f ′ (r k ) 2 | F k ]/d, I 3 (k) def = ∇φ(X k ) ⊤ ∆ k ,for any ε > 0 and λ ≥ 0, we haveEQUATIONEQUATION)EQUATIONProof of (43) in Lemma C.5. Using the fact that Proof of (44) in Lemma C.5. We haveEQUATIONTr(∇ 2 φ(X k )K)E [f ′ (r k ) 2 | F k ] d ≤ d∥∇ 2 φ(X k )K∥ op • E[f ′ (r k ) 2 | F k ] d ≤ ∥∇ 2 φ(X k )∥ op • ∥K∥ op • E[f ′ (r k ) 2 | F k ] ≤ CM 2 E[f ′ (r k ) 2 | F k ]. (Lemma C.3) From Lemma C.4, E[f ′ (r k ) 2 | F k ] is bounded by a constant independent of d, as desired.Proof of (45) in Lemma C.5. We have|∇φ(X k ) ⊤ ∆ k | ≤ |∇φ(X k ) ⊤ a k+1 | • |f ′ (r k )|.By Lemma C.3, ∥∇φ(X k )∥ ≤ C∥W k ∥ ≤ CM (since we are working under a stopping time), and so ∇φ(X k ) ⊤ a k+1 is subgaussian (and thus bounded by d ε w.o.p.). By (39) , f ′ (r k ) is bounded by d ε w.o.p., and so their product is bounded by d 2ε w.o.p., as desired. Now for the expectation:E |∇φ(X k ) ⊤ ∆ k | | F k ≤ E |∇φ(X k ) ⊤ a k+1 | • |f ′ (r k )| | F k ≤ E |∇φ(X k ) ⊤ a k+1 | 2 | F k 1 2 • E f ′ (r k ) 2 | F k 1 2The first term is bounded by a constant independent of d, since subgaussian moments are bounded.The second term is bounded by Lemma C.4, completing the proof.Lemma C.6 (Infinity norm bounds). For G k , N k , Q k as defined in 1.2, we have, for any ε, λ > 0, there exists C > 0 such that,EQUATION)Proof. The first line, (46) , follows from (40) . For the first inequality,∥G k ∥ ∞ = max 0≤j≤k ∥∆j ∥ 2d , which are all bounded by d ε with overwhelming probability. A union bound tells us that the maximum is also bounded by d ε w.o.p.. For the second inequality,E [|G k | λ ∞ | F k ] ≤ E ∥∆ k ∥ 2 d λ | F k + E max 0≤j≤k-1 ∥∆ j ∥ 2 d λ | F k ≤ E ∥∆ k ∥ 2 d λ | F k + max 0≤j≤k-1 ∥∆ j ∥ 2 d λ ≤ d ε , (w.o.p.)as desired. The second line is more straightforward:∥N k ∥ ∞ = max 0≤j≤k ∥(W + j ) ⊤ W + j ∥.Now, ∥X ⋆ ∥ and ∥X 0 ∥ are bounded independent of d, and ∥X j ∥ is bounded by cM (because of the stopping time we are using.) Thus the maximum over j of their inner products are bounded by a constant. The same thing holds for ∥Q k ∥ ∞ :∥Q k ∥ ∞ = max 0≤j≤k R(X j ) = max 0≤j≤k h(W ⊤ j KW j ).Since the derivative of h is pseudo-Lipschitz, h is continuous, and thus bounded for bounded arguments. And indeed, the argument to h is bounded:∥W ⊤ j KW j ∥ ≤ ∥W j ∥ 2 ∥K∥op , both of which are bounded independent of d. Finally, a similar argument applies to G k :∥G k ∥ ∞ = max 0≤j≤k E ∥∆ j ∥ 2 d | F j ≤ max 0≤j≤k C = C by Lemma C.4.We now prove a concentration result that closely follows [15, Proposition 5.6] .EQUATION)where v-U U T v ∼ N 0, I d -U U T and U U T v ∼ N 0, U U T with v-U U T v independent of U U T v. Lemma C.8. For a matrix H = H k with bounded operator norm, or ∥H∥ op < C and E[H k | F k ] = H k , set q(a) = a ⊤ Ha. Then E[q(a k+1 )f ′ (r k ) 2 | F k ] -Tr(KH) E[f ′ (r k ) 2 | F k ] ≤ C(H).Note that the H used here is not the same as the matrix used in the integro-differential equation.Proof. Many of the computations in this proof are taken directly from [15] , but we repeat them here for completeness. We haveEQUATIONTo compute the conditional mean E[a k+1 | Fk ] and covariance (a k+1 -E[a k+1 | Fk ])(a k+1 - E[a k+1 | Fk ]) ⊤ ,= √ Kv k , for v k ∼ N (0, I d ).EQUATION)We note that(I d -Π k )v k ∼ N (0, I d -Π k ) and Π k v k ∼ N (0, Π k ) with (I d -Π k )v k independent of Π k v k .From this, we have thatEQUATION)Moreover the conditional covariance of a k+1 is precisely (49) to get the leading order behavior(E [(a k+1 -E [a k+1 | Fk ])(a k+1 -E [a k+1 | Fk ]) ⊤ | Fk ]) (52) = √ K(I d -Π k ) √ K, where Π k = Q k Q T k . Next, using that E[H k | F k ] = H k , we expandEQUATIONTaking the expectation with respect to F k , we obtainEQUATIONwhere the error E k is defined asEQUATION)EQUATIONThe proof now turns to bounding the expectation of this error quantity.| Tr(H√ KΠ k √ K) E[f ′ (r k ) 2 | F k ]| = | Tr(H √ KΠ k √ K)| • E[f ′ (r k ) 2 | F k ] ≤ ∥H∥ op ∥K∥ op | Tr(Π k )| • E[f ′ (r k ) 2 | F k ] ≤ ∥H∥ op ∥K∥ op • rank(Q k ) E[f ′ (r k ) 2 | F k ] ≤ 2∥H∥ op ∥K∥ op E[f ′ (r k ) 2 | F k ].By (39) , the expectation is bounded by a constant, so this term is overall bounded by a constant. We move on to the next term in the error:q( √ KΠ k v k )f ′ (r k ) 2 ≤ ∥H∥ op ∥K∥ op ∥Π k v k ∥ 2 f ′ (r k ) 2 .Taking expectations and using Cauchy Schwarz, we obtainE[q( √ KΠ k v k )f ′ (r k ) 2 | F k ] ≤ ∥H∥ op ∥K∥ op • E[∥Π k v k ∥ 4 | F k ] • E[f ′ (r k ) 4 | F k ]. The first expectation is E[∥Π k v k ∥ 2 | F k ] = ∥Π k ∥ 4 F = 8, and the second is bounded by (39) as before. We thus conclude that E[E k | F k ] is bounded by a constant depending on ∥H∥ op , completing the proof.Lemma C.9. There is a constant C such that|γ(G k ) -γ(B k )| ≤ Cd -1 .Proof. Using the Lipschitz condition on the stepsize, we have|γ(G k ) -γ(B k )| ≤ ∥G k -Tr(K)I(B k )/d∥ ∞ × (1 + 2∥N k ∥ α ∞ + ∥G k ∥ α ∞ + ∥ Tr(K)I(B k )/d∥ α ∞ + 2∥Q k ∥ α ∞ ) ≤ C∥G k -Tr(K)I(B k )/d∥ ∞ (Lemma C.6) ≤ Cd -1 max 0≤j≤k E[a ⊤ j+1 a j+1 f ′ (r j ) 2 | F j ] -Tr(K) E[f ′ (r j ) 2 | F j ] ≤ Cd -1 , (Lemma C.8)as desired.In this section, we confirm that AdaGrad-Norm satisfies Assumption 6. In the notation of Assumption 6, we have, for AdaGrad-Norm,γ(td, f, g, q) = η b 2 + ∞ 0 g(s) ds .Note that this reduces to the discrete stepsize if we plug in g = G k :γ(td, f, G k (d × •), q) = η b 2 + ∞ 0 G k (ds) ds = η b 2 + ∞ 0 1 {ds≤k} 1 d k i=0 ∥∇ X Ψ(X i ; a i+1 , ϵ i+1 )∥ 2 1 [i,i+1) (ds) ds = η b 2 + ∞ 0 1 {u≤k} 1 d 2 k i=0 ∥∇ X Ψ(X i ; a i+1 , ϵ i+1 )∥ 2 1 [i,i+1) (u) du = η b 2 + 1 d 2 k i=0 ∥∇ X Ψ(X i ; a i+1 , ϵ i+1 )∥ 2 ,which is exactly the discrete version of the AdaGrad-Norm stepsize.Proposition C.8 (Lipschitz). For functions f, g, q such that f (ds) = g(ds) = q(ds) = 0 for s > t, the AdaGrad stepsize γ is Lipschitz. That is,|γ(td, f (d × •), g(d × •), q(d × •)) -γ(td, f (d × •), ĝ(d × •), q(d × •))| ≤ C(t, γ)(∥g -ĝ∥ ∞ ).Remark C.2. This is a stronger condition than the α-pseudo Lipschitz one in Assumption 6.Proof. To show this, we look at the derivative of the AdaGrad stepsize function. SettingF (x) = η √ b 2 +x , we have |F ′ (x)| = η 2(b 2 + x) 3/2 ≤ η 2b 3 for x ∈ [0, ∞). We thus have |γ(td, f (d × •), g(d × •), q(d × •)) -γ(td, f (d × •), ĝ(d × •), q(d × •))| = η b 2 + ∞ 0 g(ds) ds - η b 2 + ∞ 0 ĝ(ds) ds = F ∞ 0 g(ds) ds -F ∞ 0 ĝ(ds) ds ≤ η 2b 3 ∞ 0 g(ds) ds - ∞ 0 ĝ(ds) ds ≤ η 2b 3 t 0 g(ds) ds - t 0 ĝ(ds) ds ≤ η 2b 3 (t • ∥g -ĝ∥ ∞ ) ≤ ηt 2b 3 • ∥g -ĝ∥ ∞ ,where we were able to replace the ∞ with a t because g(ds) = 0 for s > t. We have thus obtained a Lipschitz constant ηt 2b 3 depending only on t.Next we show that the AdaGrad-Norm is bounded. Proposition C.9 (Boundedness). Suppose γ is AdaGrad-Norm. Then (6), as part of Assumption 6, holds.Proof. This is immediate:γ(td, f, g, q) = η b 2 + t 0 g(s) ds ≤ η b .It remains to show that AdaGrad-Norm satisfies (5) in Assumption 6. Proposition C.10 (Concentration). Suppose γ is AdaGrad-Norm, with G k and G k being defined as before. Then Equation (5), as part of Assumption 6, holds:E[|γ(G k ) -γ(G k )| | F k ] ≤ Cd -δ (1 + ∥f ∥ α ∞ + ∥q∥ α ∞ ).Proof. Looking to remove the square roots, we haveEQUATIONWe now bound the sum above. SetF i = ∥∆ i ∥ 2 /d, F β i = Proj β (F i ), ∆M i = F i -E[F i | F i ],and∆M β i = F β i -E[F β i | F i ]. Then |∆M β i | ∈ [-2β, 2β], so Azuma's inequality gives usP |M β k | ≥ t ≤ 2 exp - -t 2 2 k i=0 (2β) 2 , P |M β k | ≥ d 1/2+ε ≤ 2 exp - -d 1+2ε 2T d(2d ε/2 ) 2 = exp - d ε 8T.where we set β = d ε/2 . This is close to the bound we want: the error is|M k -M β k | ≤ k i=0 |F i -F β i | + | E[F i -F β i | F i ]|.We haveP(F i -F β i ̸ = 0) = P(|F i | > β) = P ∥∆ i ∥ 2 d > d ε/2 ,which superpolynomially small by (40) . The expectation is similar:| E[F i -F β i | F i ]| = | E[(F i -F β i )1 {|Fi|>β} | F i ]| ≤ E[|F i -F β i | 2 | F i ] 1 2 • E[1 {|Fi|>β} | F i ] 1 2 ≤ 4 E[|F i | 2 | F i ] 1 2 • E[1 {|Fi|>β} | F i ] 1 2. The first expectation is bounded by a constant independent of d by (40) , and the second expectation is superpolynomially small by the same argument as above. We then have|M k -M β k | ≤ d 1/2+εwith overwhelming probability (note that this would be true for any power of d, by the definition of superpolynomially small.) We thus conclude that|M k | ≤ d 1/2+εwith overwhelming probability. Multiplying by d, we find thatk j=0 (E [∥∆ j ∥ 2 | F j ] -∥∆ j ∥ 2 ) ≤ d 3/2+ε w.o.p.Plugging this back into (57), we find thatγ(G k ) 2 -γ(G k ) 2 ≤ η 2 d 2 b 4 d 3/2+ε ≤ Cd -1/2+εwith overwhelming probability, and so, taking the square root,|γ(G k ) -γ(G k )| ≤ Cd -1/4+ε/2 w.o.p,which is less than d -1/4+ε as d grows (we replaced the constant with an extra factor of d ε/2 .) Controlling the expectation via the boundedness of γ, we find that with δ = 1/8,E[|γ(G k ) -γ(G k )| | F k ] ≤ d -δ w.o.p., as desired.In this section we provide proofs of the propositions related to AdaGrad-Norm in the least squares setting as well as the more general strongly convex setting. Statements of the propositions for least squares examples are found in Section 4.In order to derive the limiting learning rate in this case, we need the following assumption and some standard definitions of strong convexity. Assumption 7 (Risk and loss minimizer). Suppose thatX ⋆ ∈ arg min X R(X) = E a,ϵ [f (⟨X, a⟩, ⟨X ⋆ , a⟩), ϵ]exists and has norm bounded independent of d. Then one has,⟨X ⋆ , a⟩ ∈ arg min x {f (x, ⟨X ⋆ , a⟩, ϵ)},for almost surely a ∼ N (0, K) and ϵ.While at first, this assumption seems quite strong, in fact, in a typical student-teacher setup when label noise is 0 (i.e., ϵ = 0), where the targets have the same model as the outputs, the assumption is satisfied. Our goal here is not to be exhaustive, but simply to illustrate that our framework admits a nontrivial and useful analysis and which gives nontrivial conclusions for the optimization theory of these problems. (58) holds with constant L. Suppose x ⋆ ∈ arg min x {f (x, y, z)} exists. An immediate consequence of ( 58) is thatEQUATION) Note that if f ′ = ∂ ∂x f (x, y, z) is L(f )-Lipschitz, i.e., |f ′ (x, y, z) -f ′ (x, y, z)| ≤ L(f )|x -x|, then the inequalityEQUATION)Definition D.2 (Restricted Secant Inequality). A function f : R 3 → R that is C 1 -smooth (in the first variable) satisfies the (µ, θ)-restricted secant inequality (RSI) if, for any x ∈ R andEQUATIONIf f satisfies the above for θ = ∞, then we say f satisfies the µ-RSI.Proposition D.1. Let the outer function f : R 3 → R be a L(f )-smooth function satisfying the RSI condition with μ(f ) with respect to x ∈ R. Suppose X ⋆ ∈ argmin X {R(X)} exists bounded, independent of d and Assumption 7 holds and thatγ 0 = η b = 2μ(f ) ( L(f )) 2 1 d Tr(K)ζ, for some ζ ∈ (0, 1), and that ∞ 0 R(s)γ s ds < ∞ with γ s as in Table 2 (AdaGrad-Norm, general formula), thenEQUATION.Proof. Given the Eq. (87) for the distance to optimality, with (x, x ⋆ ) ∼ N (0, B),d dt D 2 (t) = -2γ t E a,ϵ [⟨x -x ⋆ , f ′ (x, x ⋆ )⟩] + γ 2 t d Tr(K)E a,ϵ [(f ′ (x, x ⋆ )) 2 ]By the RSI (with constant μ(f )) condition on f , we have thatEQUATIONwhere x = ⟨X, a⟩ and x ⋆ = ⟨X ⋆ , a⟩ and we note that x has t-dependence due to the t-dependence inB. By L(f )-smoothness, 1 2 L(f ) (f ′ (x)) 2 ≤ L(f ) 2 (x -x ⋆ ) 2 .This implies that 1EQUATIONThus by ( 60) and ( 61), we have thatd dt D 2 (t) ≤ -γ t 4μ(f ) -2( L(f )) 2 1 d Tr(K)γ t R(t)Which then yield:D 2 (t) ≤ D 2 (0) -2 2μ(f ) -( L(f )) 2 1 d Tr(K)γ 0 t 0 R(s)γ s ds. Changing variables u = Γ(t) = t 0 γ s ds, we have that ∞ 0 R(t)γ t dt = ∞ 0 r(u) du = ∥r∥ 1 .Rearranging the term in the above equation and taking t → ∞. We obtain:EQUATION≥ η 2 b η + 1 d Tr(K) ( L(f )) 2 D 2 (0) (2μ(f)-( L(f )) 2 1 d Tr(K)γ0) = η 2 b η + 1 d Tr(K)( L(f )) 2 2μ(f )(1-ζ) D 2 (0).where the first inequality is by Eq. 61, and the last transition is by taking the initial learning rate to beγ 0 = 2μ(f ) ( L(f )) 2 1 d Tr(K) ζ, for ζ ∈ (0, 1). Lemma D.1. Given γ t as in Table 2 (AdaGrad-Norm), defining g(u) = γ(Γ -1 (u)), with Γ(t) = t 0 γ s ds, then g(u) = η 2 b η + 1 2d Tr(K) u 0 i(v) dv with i(v) = I(B(Γ -1 (v))).Proof. Taking the square of both sides of the γ t equation in Table 2 (AdaGrad-Norm), changing variables to u = Γ(t) and rearranging the terms:EQUATIONsuch that i(v) = I(B(Γ -1 (v))).Taking derivative with respect to u, rearranging terms and integrating leads to the desired result.To study the effect of the structured covariance matrix and cases in which the problem is not strongly convex, we will focus on the linear least square problem. In this setting, the continuum limit of the risk for the AdaGrad-Norm algorithm has the form of a convolutional integral Volterra equation,EQUATIONwhere Γ(t) := t 0 γ s ds with,EQUATIONEQUATIONIn the following we consider three cases, a strongly convex risk in which the spectrum of the eigenvalues is bounded from below (section D.2.1). A case in which the spectrum is not bounded from below as d → ∞, but the number of eigenvalues below some fixed threshold is o(d) (section D.2.2). Finally, power law spectrum supported on [0, 1] with d → ∞ (section D.2.3).Proof of Proposition 4.2. Define the composite functions r(u) = R(Γ -1 (u)), and g(u) = γ(Γ -1 (u)). Integrating the formula for the risk:t 0 r(u) du = t 0 F (u) du + t 0 Γ -1 (u) 0 γ 2 s K(u -Γ(s))R(s) ds du = t 0 F (u) du + t 0 u 0 K(u -x)r(x)g(x) dx du ≤ t 0 F (u) du + γ 0 t 0 r(x) t x K(u -x) du dx Taking t → ∞, we get ∥r∥ 1 ≤ ∥F ∥ 1 + γ 0 ∥K∥ 1 ∥r∥ 1 .Using ∥K∥ 1 = ∞ 0 K(x) dx < γ -1 0 , and noting that by Eq. (66), and Eq. ( 65), we have that ∥F ∥ 1 = 1 4 D 2 (0), and ∥K∥ 1 = 1 2d Tr(K),∥r∥ 1 ≤ ∥F ∥ 1 1 -γ 0 ∥K∥ 1 = 1 4 D 2 (0) 1 -γ0 2d Tr(K).On the hand following Lemma D.3,1 4 D 2 (0)(1 + γ0 2d Tr(K)) ≤ ∥r∥ 1 . Therefore, ∥r∥ 1 ≍ 1 4 D 2 (0). Next, rewriting the γ t equation in Table 2 (AdaGrad-Norm for least squares) in terms of g(u) (Lemma D.1), we obtainEQUATIONTaking u → ∞, and usingEQUATIONThis then completes the proof.Remark D.1. We note that, on the Least square problem L(f ) = μ(f ) = 1, therefore, the bound in Proposition D.1 yieldsη 2 b η + 1 2(1-ζ) 1 d Tr(K)D 2 (0) .Proof of Proposition 4.1. Using the equation for the distance to optimality (Eq. 8), we can derive an equation for the integral of the risk (with no target noise) which we denote by g(t) = t 0 R(s) ds:EQUATIONFor K = I d , this equation simplifies,EQUATIONPlugging in the equation for the AdaGrad-Norm learning rate (Table 2 ) leads to the desired result. We note that by using the equation for the learning rate, one can also derive a close equation for the learning rate itself.We now consider the case where, as d → ∞, there are eigenvalues of K arbitrarily close to 0. In Proposition 4.2 we saw a constant lower bound on γ t when d is fixed (and thus there are finitely many eigenvalues within any fixed distance of 0). This can be extended to the case where we have some C > 0 such that the number of eigenvalues of K below C is o(d) (see Proposition 4.3).Proof of Proposition 4.3. Following the structure of the loss, after some time the risk starts to decrease, and therefore R(t) ≤ R 0 for and t ≥ 0. Using these observations, we obtain a preliminary lower bound of γ t > C 1 t -1/2 (for t > 0), which enables us to deduce that R(t) is integrable and finally obtain a constant lower bound for γ t . The details of this are below.For t ≥ 0 and some C 1 > 0,EQUATIONNext, to show that the risk is integrable, we divide the matrix K into two parts K + , and K -, such that the eigenvalues of K + are greater than some α s > 0 and the eigenvalues of K -are smaller than α s where α s is a decreasing function of s to be determined later. We then have that, following Eq. ( 8), and the definition of the riskR(t) = 1 2d d i=1 λ i D 2 i (t), R(t) = R(0) - 1 d d i=1 λ 2 i t 0 γ s D i (s)ds + 1 d t 0 γ 2 s Tr(K 2 ) • R(s)ds (72) ≤ R(0) - t 0 γ s (2α s -γ s 1 d Tr(K 2 )) • R(s)ds + 2 t 0 γ s R 2 (s) ds with R 2 (s) = 1 2d i:λi≤αs λ i D 2 i (s). Next, choosing α s = γ s 1 d Tr(K 2 ), we that the last term is of order o d (1) . By Lemma D.2 ∀i, D 2 i (t) ≤ max γ t1 R(t 1 ), D 2 i (0) = c 0 where the bound c 0 comes from the assumption ⟨X ⋆ , ω i ⟩ = O(d -1/2 ) and the initialization X 0 = 0. Therefore,EQUATION)whereN s = d i=1 1 λi≤γs 1 d Tr(K 2 ) . This implies that, if γ s N s = o(d), then 2 t 0 γ s R 2 (s) ds = o d (1), provided that d is taken to be large before t.We then have that up to o d (1) constant,EQUATIONUsing Gronwall's inequality,EQUATIONwhere in the last transition we used the lower bound on the learning rate derived in Eq. (71). Thus, the risk is integrable, i.e. there is some C 3 such thatt 0 R(s) ds ≤ R(0) 1 d Tr(K 2 )C 2 1for all t > 0. Finally, we plug this into the formula for γ t and conclude that, for all t > 0,EQUATIONLemma D.2. Assume that the risk is bounded and attains its maximum at time t 1 . Then, for each i, we have D 2 i (t) ≤ max(γ t1 R(t 1 ), D 2 i (0)) for all t ≥ 0.Proof. Case 1: Suppose that D 2 i (0) ≤ γ 0 R(0). Then, by equation ( 8 In the next section, we consider cases in which the risk is not integrable, an example of such case is when the spectrum of K is supported on the interval [0, 1] or has power-law behavior near 0.Non-asymptotic bound for the Convolutional Volterra In this section, we use the convolutional Volterra structure of the risk (Eq. ( 64)) to derive non-asymptotic bounds on the risk, which will be useful in Section D.2.3 to derive the asymptotic behavior of the risk and the learning rate under power law assumption on the spectrum of the covariance matrix and the discrepancy from the target at initialization. Lemma D.3. Let Γ(t) := t 0 γ s ds and letR(t) = F (Γ(t)) + t 0 γ 2 s K(Γ(t) -Γ(s))R(s) dswhere γ t , K are monotonically decreasing, with ∥K∥ 1 < ∞. Then all t,R(t) ≥ F (Γ(t)) + t γ 2 s K(Γ(t) -Γ(s))F (Γ(s)) dsIf in addition, there exist ϵ > 0 and T > 0 such that, for all t > T ,t 0 K(s)K(t -s) ds ≤ 2(1 + ϵ)∥K∥ 1 K(t) and 2∥K∥ 1 (1 + ϵ)γ 0 < 1 then for all t R(t) ≤ F (Γ(t)) + C t 0 γ 2 s K(Γ(t) -Γ(s))F (Γ(s)) ds for C = K(0) K(T )(2ϵ + 1) + 2 1 1 -2γ(0)∥K∥ 1 (1 + ϵ) .Proof. The lower bound holds trivially, using R(s) ≥ F (Γ(s)). For the upper bound, we start with the following change of variables:R(t) = F (Γ(t)) + Γ(t) 0 g(u)K(Γ(t) -u))R(u) du,with g(u) = γ Γ -1 (u) . Let us define the convolution mapG(f )(Γ) = K * (gf )(Γ) = Γ 0 K(Γ -u)g(u)f (u) du.Next we show that this map is contracting and in particular,G 2 (f ) = G(G(f ))(t) = t 0 K(t -s)G(f )(s)g(s) ds (77) = t 0 K(t -s) s 0 K(s -u)g(u)f (u) dug(s) ds = t 0 t u K(t -s)K(s -u)g(s) ds g(u)f (u) du ≤ t 0 K * 2 (t -u)g(u) 2 f (u) duwhere the third transition is since u < s < t. The last transition is by change of variables and the assumption that γ t is a monotone decreasing function. Consecutive application of the convolution map will then yield by induction,G j (f )(t) ≤ t 0 K * (j) (t -u)g(u) j f (u) du.Therefore, expanding the loss and using the above upper bound, and denote q = 2(1 + ε)∥K∥ 1 γ 0 such that q < 1,EQUATIONwhere the third transition is by Lemma D.4, with C 1 = K(0) K(T )(2ϵ+1) + 1, which then completes the proof.Lemma D.4 (Lemma IV.4.7 in [3] ). Suppose K is monotonically decreasing, with ∥K∥ 1 < ∞, and that there exists T > 0 such that ∀t ≥ T , and ϵ ≥ 0,EQUATIONThen,EQUATIONProof. Define α n = sup t≥0K * n (t) K(t)(2∥K∥1) n-1 , trivially α 1 = 1. Consider the n + 1 convolution, K * (n+1) (t) K(t)(2∥K∥ 1 ) n = 1 K(t) t 0 K(s)K * n (t -s) (2∥K∥ 1 ) n ds (82)By the assumption of the Lemma, we know that there exists some T > 0 such that for ∀t ≥ TEQUATIONTherefore, if t ≥ T , we have1 K(t) t 0 K(s)K * n (t -s) (2∥K∥ 1 ) n ds (84) = t 0 K(s)K(t -s) 2∥K∥ 1 K * n (t -s) K(t -s)(2∥K∥ 1 ) n-1 ds ≤ α n (1 + ϵ) On the other hand, if t < T , 1 K(t) t 0 K(s)K * n (t -s) (2∥K∥ 1 ) n ds ≤ K(0) K(T ) ∥K * n (t)∥ 1 (2∥K∥ 1 ) n ≤ K(0) K(T )2 n (85)Taking supremum in Eq. ( 82), and combining the results of Eq. ( 85), and Eq. ( 84), we obtain that,α n+1 ≤ K(0) K(T )2 n + α n (1 + ϵ)Solving the above recursion equation,α n ≤ K(0) K(T ) n-2 k=0 1 2 n-k-1 (1 + ϵ) k + (1 + ϵ) n-1 = K(0) K(T )2 n-1 1 -(2(1 + ϵ)) n-1 1 -2(1 + ϵ) + (1 + ϵ) n-1 ≤ (1 + ϵ) n-1 K(0) K(T )(2ϵ + 1) + 1 ,rearranging the terms we arrived at the required result.1 + Γ(t) ≍ t 1/(2-κ1) ). Finally, plugging this into the for γ t and applying Corollary D.1, we getγ t ≍ η b 2 + 2 d Tr(K) t 0 F (Γ(s)) ds ≍ (1 + t) -(1-κ1)/(2-κ1) .In the case of κ 1 = 1, we follow a similar procedure as for κ 1 < 1 to show that t ≍ Γ(t) log(Γ(t)) for sufficiently large t. This implies Γ(t) ≍ t/ log(t) which gives the desired result after integration. The decay rate of the risk is then immediate using Corollary D.1.Lemma D.5. Let K have a spectrum that converges as d → ∞ to the power law measure ρ(λ) = Cλ -β 1 (0,λmax) , with C -1 = λ 1-β max 1-β for some β < 1, and λ max > 0, and suppose that D 2 i (0) ∼ λ -δ i , then F (t) ≍ t -κ1 , and K(t) ≍ t -κ2 , with κ 1 = 2-β -δ, and κ 2 = 3-β. In addition, K(t) ≍ t -κ2 , satisfies Eq. (80).Proof. Following the definition in Eq. ( 66), and Eq. ( 65)F (x) = 1 -β 2λ 1-β max λmax 0 λ 1-β-δ e -2λx dλ = 1 -β 2λ 1-β max (2x) 2-β-δ 2λmaxx 0 y 1-β-δ e -y dy = 1 -β λ 1-β max 2 3-β-δ γ(2 -β -δ, 2λ max x) x 2-β-δ .Similarly for K,K(x) = 1 -β λ 1-β max λmax 0 λ 2-β e -2λx dλ = 1 -β λ 1-β max 2 3-β γ(3 -β, 2λ max x) x 3-β .with γ(s, z) = z 0 x s-1 e -x dx is the incomplete gamma function. For large z, γ(s, z) ≍ Γ(s), the complete gamma function. We therefore obtain κ 1 = 2 -β -δ, and κ 2 = 3 -β. Next, we show that K(x) ≍ x -κ2 satisfies Eq. (80),t 0 K(s)K(t -s) ds ≤ t/2 0 K(t)K(t -s) ds + t t/2 K(t)K(t -s) ds ≤ K(t/2) t/2 0 K(s) ds + t t/2 K(t -s) ds ≤ 2K(t/2)∥K∥ 1by the power-law assumption for t > T , K(t/2) ≍ K(t) which then complete the proof.Proof of Proposition 4.4. The proof is an immediate application of Corollary D.2 with, κ 1 = 2-β-δ as implied by Lemma D.5.Remark D.2. This includes the case β = 0, which is the uniform measure on [0, λ max ].The distance to optimality of SGD is measured say by D 2 (X) = ∥X -X ⋆ ∥ 2 . Let us consider the deterministic equivalent for the distance to optimality D 2 (t) in (11) . Fixing T > 0 and any ε ∈ (0, 1/2), we have by Theorem 2.1 (see also corollary B.1 which show concentration for large class of statistics) that sup 0≤t≤T |∥X ⌊td⌋ -X ⋆ ∥ 2 -D 2 (t)| ≤ d -ε , w.o.p. In this way, if we want to guarantee that the distance to optimality of SGD decreases, we need dD 2 (t) < 0 with the maximum decrease being min γt dD 2 (t).As it turns out, the evolution of D 2 is particular simple, as it solves the differential equation (derived from the ODE in ( 9)) D λ 1 (t) and the learning rate γ t , respectively. Further details and formulas for the limiting behavior can be found in the Appendix F.2. See Appendix H for simulation details.EQUATIONThe distance to optimality threshold, γD t , occurs precisely when dD 2 < 0. This choice of γ makes the ODE for the distance to optimality stable. By translating the relevant deterministic quantities in γD t back to SGD quantities, we getḡD k def = 2⟨X k -X ⋆ , ∇R(X k )⟩ Tr(K) d E a,ϵ [f ′ (⟨X k , a⟩; ⟨X ⋆ , a⟩, ϵ) 2 ]with the deterministic equiv. γD t = 2A(B(t))Tr(K) d I(B(t)).(88) A greedy learning rate that maximizes the decrease at each iteration is simply given by g Polyak In this context, the Polyak learning rate is impractical because we do not known X ⋆ . In spite of this, we can learn some things about this learning rate as it is the natural extension of Polyak learning rate to SGD. 1 avg. eig of K . In the case of least squares (see ( 7)), we getg Polyak k = 2R(X k ) -ω 2 2Tr(K) d R(X k )and on a noiseless least squares, g Polyak k = 1Tr(K) d.The latter gives the best fixed learning rate for a noiseless target on a LS problem (as noted in [35, 43] ).Naturally, one can ask a similar question as in Polyak in the context of line search (i.e., decreasing risk at each iteration of SGD). First, by the structure of the risk (Assumption 3 and 4),EQUATIONTherefore using (9), we have that the deterministic equivalent for ∥∇R(X)∥ 2 is M (t) = 1 2 d i=1 m(V i (t)λ 2 i ).In this case, the deterministic equivalent for the risk R satisfies the following ODEEQUATIONFrom this, we get an immediate learning rate (stability) threshold for the risk, that is, ḡR k is the largest rate for which SGD is guaranteed to decrease at each iteration, i.e., when the deterministic equivalent of R satisfies dR < 0 or equivalently after translating relevant terms into SGD quantities risk threshold ḡRk = ∥∇R(X k )∥ 2 Tr(K∇ 2 R(X k )) d I(W T k KW k )and deterministic equiv γRt = M (t) v(K) d I(B(t)).(92) The greediest approach, which we call exact line search, would choose the learning rate such that γ line t ∈ arg min γ dR. In this case, we getg line k = 1 2 g R k and deterministic equiv γ line t = 1 2 γ R t .In this section, we provide a proof of Proposition 3. EQUATIONwhere x is the positive real root of the second-degree polynomialEQUATIONThis leads toEQUATIONProof. We establish the inequalityλ min (K) 1 d Tr(K 2 ) ≤ γ line t for all t ≥ 0 by observing 1 d d i=1 λ 2 i D 2 i (t) ≥ 2λ min (K) 1 2d d i=1 λ i D 2 i (t) = 2λ min (K)R(t). Now let us consider K ∼ 1 2 λ 1 + 1 2 λ 2 for λ 1 > λ 2 > 0. We define D λ (t) def = d λi=λ D 2 i (t). Utilizing the ODEs in (9), we derive d dt D λ (t) = -2γ t λD λ (t) + 2γ 2 t λ × |{λ = λ i } d i=1 | × R(t)for each distinct eigenvalue λ of K. Here |{λ = λ i } d i=1 | is the number of eigenvalues of K that are equal to λ. It immediately follows by our construction ofK that |{λ = λ i } d i=1 | = d 2 .Thus, we establish the following system of ODEsEQUATIONwhereR(t) = 1 2d (λ 1 D λ1 (t) + λ 2 D λ2 (t)) and γ line t = 2(λ 2 1 D λ 1 (t)+λ 2 2 D λ 2 (t)) (λ1Dλ 1 (t)+λ2D λ 2 (t))(λ 2 1 +λ 2 2 ).Since D λ2 (t) ≥ 0 and λ 1 > λ 2 > 0, we infer that R(t) (t) = lim t→∞ -2γ t λ 2 D λ2 (t) + dγ 2 t λ 2 R(t) -2γ t λ 1 D λ1 (t) + dγ 2 t λ 1 R(t) = lim t→∞ -2λ 2 D λ2 (t) + dγ t λ 2 R(t) -2λ 1 D λ1 (t) + dγ t λ 1 R(t) = lim t→∞ γ t λ1λ2 2 D λ1 (t) + λ 2 D λ2 (t) γ t λ2 2 -2 γ t λ1λ2 2 D λ2 (t) + λ 1 D λ1 (t) γ t λ1 2 -2 = lim t→∞ D λ1 (t) 2 λ 3 1 λ 2 + D λ1 (t)D λ2 (t)(-λ 1 λ 3 2 + λ 2 1 λ 2 2 -2λ 3 1 λ 2 ) + D λ2 (t) 2 (-λ 4 2 -2λ 2 1 λ 2 2 ) D λ1 (t) 2 (-λ 4 1 -2λ 2 1 λ 2 2 ) + D λ1 (t)D λ2 (t)(-λ 3 1 λ 2 + λ 2 1 λ 2 2 -2λ 1 λ 3 2 ) + D λ2 (t) 2 λ 1 λ 3 2 = λ 3 1 λ 2 + lim t→∞ D λ 2 (t) D λ 1 (t) (-λ 1 λ 3 2 + λ 2 1 λ 2 2 -2λ 3 1 λ 2 ) + lim t→∞ D λ 2 (t) D λ 1 (t) 2 (-λ 4 2 -2λ 2 1 λ 2 2 ) (-λ 4 1 -2λ 2 1 λ 2 2 ) + lim t→∞ D λ 2 (t) D λ 1 (t) (-λ 3 1 λ 2 + λ 2 1 λ 2 2 -2λ 1 λ 3 2 ) + lim t→∞ D λ 2 (t) D λ 1 (t) 2 λ 1 λ 3 2 . Therefore, lim t→∞ D λ 2 (t) D λ 1 (t)is the positive real root of the second-degree polynomialEQUATIONSolving for x > 0, we derive the explicit formulaEQUATIONGivenEQUATIONwe haveEQUATIONBy substituting (98), we getlim t→∞ γ line t = λ 3 1 + 2λ 2 1 λ 2 + 2λ 1 λ 2 2 + λ 3 2 -λ 6 1 -4λ 5 1 λ 2 + 8λ 4 1 λ 2 2 -6λ 3 1 λ 3 2 + 8λ 2 1 λ 4 2 -4λ 1 λ 5 2 + λ 6 2 (λ 2 1 + λ 2 2 ) 2 . (101) A direct calculation reveals that λ 1 > λ 2 > 0 implies lim t→∞ γ line t ≤ 2λmin(K) 1 d Tr(K 2 ) .Remark F.1. For the scenario where K has an arbitrary number n of distinct eigenvalues, equation (13) remains valid. The proof parallels the one outlined above. However, in this case, the expression for lim k→∞ g k is given byEQUATIONwhere x 1 , . . . , x n-1 > 0 satisfy a more intricate coupled system of n -1 equations.Any single index model with α-pseudo Lipschitz (α ≤ 1) activation function is covered by our SGD+AL theory. this section, we provide key learning problems within this family of models.We consider a binary logistic regression problem with ϵ = 0 where we are trying to classify two classes. We will follow a Student-Teacher model, in which there exists a true vector X ⋆ to be the true direction such that possible labels are, y = exp(⟨X ⋆ ,a⟩) exp(⟨X ⋆ ,a⟩)+1 . or 1 -y. In order to classify the data we minimize the KL-divergence between the label y and our estimate defined by the above formula, We note that, the logistic regression is (µ, θ)-RSI with µ = 1 ℓe √ 4θ see section 2.2 in [15] . Its Lipschitz constant is L(f ) = 1. Using Proposition D.1 one can derive a lower bound on the limiting learning of AdaGrad Norm.For more details and more examples, see [15] .Finally, we include an example that uses real-world data, that is, the CIFAR 5m dataset [38] . Our theory does not explicitly deal with non-Gaussian distributions, but we find that the theoretical risk curves generalize cleanly to that case.As we are now working with discrete data points rather than a distribution, the learning setup, while closely analogous to what was presented earlier, has some slight differences.We start with a subset of the data consisting of n grayscale images, each of which is 32 × 32 pixels, that is, A ∈ R n×1024 . We fill a vector b ∈ R n with the corresponding labels (0 for an image of a plane, 1 for an image of a car.) We then randomly choose a matrix W ∈ R 1024×d with i.i.d. Gaussian entries to generate the features F = relu(AW ). We want to use least squares to predict the label from the features, i.e., findEQUATION)where f i is the ith row of F . The SGD we now consider isEQUATIONwhere γ k is the usual AdaGrad-Norm stepsize, as in (15) . Our empirical covariance matrix K (remembering that f i is a row vector) is thenEQUATIONWe now use (64) , with the AdaGrad-Norm stepsize, to numerically simulate the SGD loss, which we then compare to the actual loss. Our theory matches empirical results very closely. Here we provide more details for the figures that appear in the main paper. The shaded region represents a 90% confidence interval for the SGD runs. As the dimension increases, the risk and stepsize both concentrate around a deterministic limit (red). The deterministic limit is described by an ODE in Theorem 2.1. The initial loss increase in the least squares problem suggesting that the learning rate was initially too high, but AdaGrad-Norm naturally adapts and still the loss converges. Our ODEs predict this behavior. ). Larger values of s correspond to smaller spreads in the spectrum. Additionally, Tr(K)/d = 1 for all s. Both plots highlight the implication of equation ( 13) in high-dimensional settings, where a broader spectrum of K results in λmin(K)1 d Tr(K 2 ) ≪ 1 1 d Tr(K), indicating slower risk convergence and poorer performance of exact line search (unmarked) as it deviates from the Polyak stepsize (circle markers). The gray shaded region demonstrates that equation ( 13) is satisfied. . Depicted is learning rate asymptotic so it approaches 1. (center, right): Noiseless least squares setting (ω = 0). (center): Prop. 4.2 predicts the avg. eig of K (Tr(K)/d) as compared with λ max effects the lim k→∞ g k . Indeed, this is true. We varied the κ = λ max /λ min while keeping the Tr(K)/d and all other parameters fixed. All the learning rates behave identically verifying our theory about the effect of Tr(K)/d on learning rates. (right): Varying the learning rate of AdaGrad norm by ∥X 0 -X ⋆ ∥ 2 ; our predictions (dashed) match and we see the inverse relationship predicted by Prop. 4.2. See Appendix D for details. Additionally, we did the following.• Center plot: AdaGrad with b = 0.5, η = 2.5 is run on the least squares problem with d = 1000 and X 0 , X ⋆ ∼ 1 √ d N (0, I). The covariance matrix K is generated so that the eigenvalues areλ i (K) = d d i=1 i d+1 -2/s • i d + 1 -1/s , i = 1, . . . , d.The constant s > 2. When s is near 2, the spectrum is more spread out, i.e., κ = λmax λmin is large. Larger values of s mean smaller the spreads. Moreover Tr(K)/d = 1 for all s. In the simulations, we used s ∈ {2.1, 3.0, 3.5, 4.0, 5.5} and recorded the condition number κ.• Right plot: Ran AdaGrad with b = 0.5, η = 2.5 on the least squares problem with d = 1000.X ⋆ = 0 and X 0 ∼ p d N (0, I) where p ∈ {1, 2, 4, 8, 16}. In this way, ∥X 0 -X ⋆ ∥ 2 = p. D λ 1 (t) and the learning rate γ t , respectively. Predictions from ODE theory are compared with results obtained from SGD, demonstrating close agreement between the two approaches. Initialization was performed randomly, with X 0 ∼ N (0, I d /d) and X ⋆ ∼ 1 √ d 1, where d = 400. The covariance matrix K has two distinct eigenvalues λ 1 = 1 > λ 2 > 0, and was constructed by specifying the spectrum, with λ i sampled from a discrete uniform distribution U{1, λ 2 } for i = 1, . . . , d = 400, and setting K = diag(λ i : i = 1, . . . , 400). Further details and formulas for the limiting behavior can be found in the Appendix F.2.Convergence on CIFAR 5m [38] . We train a classifier to distinguish between images of airplanes and cars. Fix d = 2000. Then for multiple values of n, we run AdaGrad-Norm with initialization X 0 = 0, b = 0.1 and η = 5, randomly sampling a datapoint from F at every step. Details of the setup can be found in Appendix G.2.More realistic, in that it deals with high-dimensional anisotropic loss geometries and more precise, in that it can distinguish minimax optimal algorithms as more-or-less performant.We extend some results to the general strongly convex setting.Note that d need not be large to define this, but the structure allows us to consider d as a tunable parameter. Moreover, as we increase d, the analysis we do will be more meaningful.Note that, holding time fixed, we perform O(d) gradient updates for a problem of dimension d. For the problems considered here, this scaling leads to consistent dynamics, but there do exist related problems where a different scaling is more appropriate. For example, under random initialization, to capture the escape of phase retrieval from the high-dimensional saddle, O(d log d) iterations are needed; see for example[56].The Polyak stepsize we analyze in this paper differs slightly from the "classic" stepsize in the literature, that is, R(X k )-R(X * )||∇R(X k )|| 2 . Rather than using this form, we skip an approximation step in the derivation[27] and use the exactly optimal form. Both variations of the Polyak stepsize can be analyzed under our assumptions; the choice was admittedly somewhat arbitrary. (Note that in the case of least squares, the two stepsizes coincide.)Our result can be compared to existing findings for SGD under power-law distributions in[8,52,55]. While these works explore similar assumptions regarding the covariance matrix spectrum, they do not address the high-dimensional regime with diverging Tr(K), focusing primarily on β > 1.This condition can be removed in a standard way: one creates an fϵ which is an approximation to f formed by convolving with an isotropic Gaussian of variance ϵ. This is C 2 and has bounded second derivatives (as f was smooth). One then takes the limit as ϵ → 0.This is the idea of Polyak stepsize when the problem is deterministic. 


Decision: This is an excellent paper that makes a number of strong theoretical and conceptual contributions to our understanding of the dynamics of adaptive stochastic optimization algorithms in high dimensions. Moreover, the mathematical techniques developed are likely to be of use to others. The authors have addressed the major concerns of the reviewers. 

Given the depth of the paper and the additional areas for further exposition and exploration, I encourage the authors to consider the suggestion of one reviewer to write an extended, journal-length version of the paper that could do the material more justice.

Reviews: review: paper_summary: paper_summary: The paper focuses on analytical analysis of dynamics of a class of optimization algorithms with adaptive learning rate applied to linear models with Gaussian data. The class of algorithms includes AdaGrad-Norm, RMSprop-Norm, Polyak stepsize and line search, but doesn't include, for example, Adam, classical RMSprop and AdaGrad. Also, authors consider $\frac{1}{d}$ scaling of learning rate with dimension $d$ of the problem.

First, the authors establish convergence of the optimization trajectory to a deterministic ODE in the limit $d\to\infty$. Then, they focus on analysing the resulting dynamics for quadratic loss functions and specific algorithms
- Plolyak stepsize. Obtaining learning rate. 
- Exact line search. Obtaining learning rate for two cases of data covariance $K$: isotropic and with two distinct eigenvalues
- AdaGrad-Norm. $O(t^{-\frac{1}{2}})$ learning rate asymptotic for noisy observations. Also, learning rate and risk asymptotic for isotropic and power-law data covariance $K$.

The theoretical results are validated with numerical experiments on synthetic data.
strength_weakness: strength_weakness: The main drawback of the manuscript is its writing. The paper mostly focuses on the introduction and setting, leaving too little space for actual results and their discussion. Out of 9 pages: 
  - Five pages are reserved for introduction and setup. 
  - Two pages are occupied by Section **2**, which mainly introduces notations and demonstrates them on the quadratic problem example. Indeed, it contains Theorem **2.1**, but that one is quite short
  - Only the last two pages contain the main results of the paper. This leaves very little space for discussing them. Some of the discussions are present in the (quite lengthy) *main contributions* part, but, for readability purposes, it would be nicer to have this discussion together with or after the results.
  - Any kind of concluding section is absent

Regarding the figures, only figure **1** is referenced in the text while figures **2**,**3**, and **4** are not referenced. This adds an extra effort to interpret them and put into right context. 

**New developments from previous work**. As mentioned by the authors, the current paper follows the technical framework developed in [14], extending it to adaptive algorithms. Then, it would be good if the manuscript mentioned more clearly which results/notions were already obtained in [14] and which are new additions of the current paper.  
- In particular, Theorem **2.1** of the current paper seems to be very close to Theorem 1.1 of [14] which can also handle non-constant learning rates. If two theorems are indeed close, it would be interesting to discuss new aspects introduced by adaptive algorithms. 
- Assumptions **1-6**, and the notations of sections **1.1**, **1.2**, **2** look very close to that of [14]. If that is the case, it would be interesting to discuss more specifically their differences. Also in that case, sections **1.1**, **1.2**, **2** could be compressed by partially moving them in appendix, leaving more space for the results and their discussion.

**Results for non-quadratic losses**. Although the framework, as introduced in section **2** is designed for general non-quadratic problems, the specific results are given only for quadratic problems in the main paper. Some results for non-quadratic problems are given in appendix sections **D.1**, **E**, **F.1**, **G.1** (mainly based on results of [14]). In such form of these results, it is hard to extract interesting conclusions, or determine the contributions compared to the previous literature.
questions: questions: **High dimensional vs. low dimensional problems**. Some covariance matrices $K$ satisfy assumption **1** but have $\frac{\mathrm{Tr} K^2}{d}\to0$ and therefore produce trivial ODE dynamics coinciding with full-batch Gradient Flow (as can be seen from eq. 11), for example 
- rank 1 with non-trivial eigenvalue $\lambda=1$ (extreme case).
- with power-law spectrum $\lambda_i=i^{-\alpha}, \alpha>1$  (e.g. classical *capacity* and *source* conditions).

Indeed, for these two examples, the dynamics converge to full-batch Gradient Flow due to $\propto \frac{1}{d}$ scaling of learning rate in eq. 3, while learning rate independent of $d$ would be more natural. 
- Could the techniques developed in this paper handle such problems?
- If no, is it possible to formulate a criterion on covariance matrix $K$ to distinguish non-trivial ODE dynamics from trivial ones (i.e. coinciding with full-batch gradient flow)? For general non-quadratic problems.

**Results for power-law spectrum** The last section provides variety of results in case of power-law distribution of eigenvalues (with exponent $\beta$), and target (with exponent $\delta$). In the non-vanishing learning rate phase $\beta+\delta<1$, how the obtained rate $\mathcal{R}(t)\sim t^{\beta+\delta-2}$ compares with other results for SGD under power-law spectrum (e.g. [1](https://arxiv.org/abs/2006.08212),[2](https://arxiv.org/abs/2102.03183),[3](https://arxiv.org/abs/2206.11124))? 

**Learning rate phase transition**. For power-law spectrum, learning rate asymptotic exhibits transition at $\beta+\delta=1$, as demonstrated on fig. **4**. Interestingly, on this figure we can see that in both phases learning rate reduces similarly up until it reaches very small value ($\gamma\sim 10^{-4}$) at quite late optimization times ($t\sim10^{3}$). Can this dynamic transition (w.r.t. time $t$) be explained theoretically, for example, with some quantities characterizing transition scale with $d$?
limitations: limitations: The paper surely doesn't have any negative societal impact.
paper_summary: paper_summary: This paper studies Stochastic Gradient Descent (SGD) training with adaptive step sizes. The setting is a generalization of single-index models. There is a ground truth vector X^*, and the model must find a vector X such that a loss L(X) = E_{a,epsilon}[f(<X,a>, <X^*,a>, epsilon)] is minimized. Here epsilon is a random additive error, and a is distributed according to N(0,K) for a potentially anisotropic covariance K.

This paper shows that (1) in the high-dimensional regime, the training dynamics converge to a deterministic limit given by an ODE. It studies these dynamics in two significant cases of adaptive step-size algorithms: (2) exact-line-search (i.e., greedily decreasing the risk optimally at each step), and (3) AdaGrad-Norm. For exact-line search, it is shown that this can be very suboptimal if the data covariance is highly anisotropic. For AdaGrad-Norm, it is shown that if the data covariance has lower-bounded eigenvalues the optimal step size (within a constant factor) is reached. However, for harder problems with worse conditioning of the data covariance, AdaGrad-Norm can be overly pessimistic.
strength_weakness: strength_weakness: In terms of techniques, Theorem 2.1 on convergence to an ODE is based in large part on modifying a previous analysis of [Collins-Woodfin ’23] to the case of adaptive step sizes, so the techniques are not completely novel.

The data is also assumed to be Gaussian, although anisotropy is allowed.
questions: questions: I am satisfied with the presentation in the paper, and do not have questions.
limitations: limitations: Yes
paper_summary: paper_summary: ### Update after rebuttal
After reading the feedback carefully, I updated my score for the paper.

### Original review

In this work, the authors study SGD and its adaptive variants in the setting of noisy generalized linear models. Assuming the covariance matrix of the data is bounded by a dimension-independent constant and a few more similar assumptions, the authors establish that the expected risk of SGD and the expectation of adaptive learning rates both converge to deterministic curves described by an ODE.

I found the paper rather hard to follow and I couldn't understand the significance of the obtained results. The theory requires a lot of simplifications: the prediction model is linear; the learning rates that are studied are "idealized", so they are not the same as those used in practice; Adagrad-norm isn't really a method popular in practice; the ODE solution is not trivial in the case of least squares, as equation (11) is implicit, so it requires further assumptions to make a statement about the solution. I did not understand what insight we can draw from all of this.
strength_weakness: strength_weakness: 1. The theory for line search and Polyak stepsize is only for the idealized version.
2. The insight that line search methods do not work particularly well is not exactly new, see for instance J Bolte, E Pauwels "Curiosities and counterexamples in smooth convex optimization"
3. The complexity of the theory and the discussion is such that it is very difficult to draw any insight, and it seems the theory wouldn't be extended by others because of how convoluted it is.

### Minor issues
Abbreviation "SGD" is used before (line 23) it is introduced (line 30)  
"We fully expect"  -> "We expect"  
"worst case convergence guarantees" -> "worst-case convergence guarantees"  
"strongly-convex" -> "strongly convex"
questions: questions: 1. What's the intuition for why $h$ exists and why would it be well-behaved? I can see a discussion in Appendix B, but it doesn't provide any high-level intuition. 
2. Why do you say in Appendix B that $R(X)$ is "an expectation of a Gaussian vector"? 
3. The authors wrote "We shall remove this condition by approximation in what follows", where is it done?
4. The authors say they need "mild assumptions on the learning rate", but why are they mild?
limitations: limitations: The authors could do a better job when discussing the limitations of their work. There is no explicit limitations section and in the ckecklist, the authors state that they already did everything needed by clearly outlining the required assumptions. However, I think it'd benefit the paper if the authors made an explicit statement on how some aspects of the theory are too restrictive.
paper_summary: paper_summary: Analyzes the behavior of AdaGrad-Norm, Line-search and Polyak step size using ODEs.
strength_weakness: strength_weakness: - Paper spends a lot of time setting up a general framework, which it then applies to specific settings. The amount of notation introduced seems over the top for a conference paper. For instance, the results for least-squares can be presented without establishing the full general framework, which could be left to the Appendix. Section 1.2 and 2 are just .... a lot to get through.
- Paper ends very abruptly.
- Some plot fonts are too small to read when printed.
- Further discussion of the implications and potential applications would be good.
questions: questions: See above.
limitations: limitations: See above
paper_summary: paper_summary: The paper is a continuation of a broad line of work exploring deterministic limits of stochastic gradient descent. In particular, the author(s) follow one specific branch of the many, which derives a solution to a coupled integro-differential equation by passing on the complex space. The novelty lies mainly in the fact that they derive, for least squares, deterministic limits of dynamics that have adaptive step-sizes, for non identity covariance of the observed dataset. Industrially, this is important as the idealized scenario of fixed step-sizes is __very much__ idealized, and realistic datasets have structure. Thanks to the general expressions (under 6 assumptions that are claimed to be not super-restrictive), they can proceed to study the various differences between standard schedules. The limiting learning rate can be derived, as well as the scaling of the risk along time. This is done for Line Search, Polyak step-size, AdaGrad Norm. Interestingly, their bounds suggest that for covariances with largely separated eigenvalues (say, power law spectrum), there should be a change in the phenomenology (i.e. a phase transition). They verify this for a power law covariance, which exhibits three phases depending on the parameters of the problem (notably the starting distance from $X^{\star}$). All is accompanied by experiments corroborating the predictions.
strength_weakness: strength_weakness: Weaknesses are questions and questions are weakenesses in some sense. Once the author(s) engage, I will reorder things that get answered and things that remain as "concerns".  

- Proposition 4.4 supposes that $\mathscr{D}_i^2(0)$ could potentially have a different scaling at each site $i$. I understand that the assumption includes an uninformative initialization $\delta = 0$. However, when is it really that one can think of a setting in which you can have an initialization that scales like $\lambda_i^{-\delta}$ for non trivial $\delta$? The result is very aesthetic, but I did not think through of its realistic value. Anyways, it is very nice so I am __not__ criticizing the math. Maybe I am just interested in understanding if you have in mind something. 
- In the Checklist, you state that you have been "__very careful__ \[...\] to include __all__ the assumptions needed". I agree with this statement. However, as a reader, I do not agree with this being the limitation of your work, as in some sense you claim over the whole text that your assumptions cover standard adaptive step-sizes (hence you hiddenly claim that this theory is all we need). I would be very happy to see a paragraph where you clearly state your limitations. Where does your theory fail? Where does the technique __not__ extend naturally? If you want to sell me your paper, I do not believe its value is in the assumptions. On a first aspect, the ideas and the techniques are of independent interest. Secondly, it would be sad if the bottleneck were the assumptions. 
- Your equation (7) is found across some other works. One of the earliest appearances is (Yoshida and Okada, 2019). However, from a first check, I would guess that you are missing a term on the first object. To be fair, I quickly tried to re-derive it and I would guess that you need a further $H_{2, t}\mathscr{V}_{12, i}(t)$ in the first expression. Am I wrong? If so, please correct me. It looks like the equations are missing some symmetry of the dynamics. 
- I am concerned about the scaling. My understanding of the literature is that scalings are very important. Your assumptions require that $\lVert K\rVert$ is bounded independently of the dimension and that $\lVert X_0\rVert$ is bounded, independently of the dimension. Moreover, your step-size is $\frac{1}{d}\mathfrak{g}_k$. Therefore, at the beginning you move by very little, I think by $\propto \frac{1}{d}$. What is done in some other references is very different. Let me take an example, and correct me if I am wrong.
1. You can derive heuristically deterministic dynamics for step-size $\mathfrak{g}_k$ (not normalized), starting from a standard Gaussian standard gaussian vector $X_0$, with data-points being i.i.d. Gaussians $a\sim \mathcal{N}\left(0,\frac{1}{d}I_d\right)$. Now, the step-size is not divided by dimension, the $X_0\sim \mathcal{N}(0, I_d)$ is not bounded in norm by something independent of $d$, and the $a$ matches your assumptions. In this case however, the signal is far stronger, as we are removing two normalizations you indeed have. Below my questions regarding this. 
2. Can you clarify why it makes sense to choose your scaling? Apart from it allowing the proofs. 
3. In particular, how is it not possible to allow for norms dependent on $d$, and in parallel, how is it that the signal is enough to be not stuck at the initialization? 
4. Does this implicitly mean that you are exploring a slow regime? Say, a Gradient Flow? What happens if in your experiments you just change the scalings to something larger in magnitude?


#### Typos
__NOTE__: I am including for completeness typos in "weaknesses". Please do not count them as such. 
- (line 560) In the reminder of $\Omega$, you are missing a $\max$. 
- (equation 20) there is no $\Gamma(t)$ in the expression, right? 
- (lines 1140-1141) the way of quoting "noise" and "variance" is not the correct one, TeX-wise.
questions: questions: Please note that some of the questions below are very ingenuous. 

- (lines 49-57) This is more of an impression. Why is $\lambda_{\min}(K) > C > 0$ for you a situation of strong anisotropy? The power-law case is anisotropic in that the minimum eigenvalue is not bounded below by a constant as $d\to\infty$. Put another way, I do not understand how the __minumum__ eigenvalue being bounded below from zero would describe the whole behavior of the covariance profile of the data-points $a$. I understand that this might be related to the fact that the optimal strategy uses the trace so you say something in the line of "if the minimum eigenvalue is very well distant from the average then there is anisotropy". At the very least, it appeared to me as a reader that you somehow __define__ strong anisotropy as this eigenvalue condition, while in fact it is a consequence of the algorithmic behavior that you can tune the eigenvalues in this "anisotropic" way and observe that "greed can be arbitrarily bad". 
- The original introduction of the Polyak step-size is that of a step-size that uses the optimal value of the function, and the value of the current iterate. In your case, what you call Polyak step-size has a connection with the original method, which is related to your assumptions, as you find that eventually it gets to a $\gamma_{\infty}$ that takes care of the geometry of the function. From a pedagogical point of view, it would be nice to comment further on this. 
- Theorem 2.1 states that the concentration is for any $\epsilon \in \left(0,\frac{1}{2}\right)$, line 48 says "for any inverse power of d", and Theorem B.1, Corollary B.1  cited below Theorem 3.1 in line 195 give a result for existance of an $\epsilon$ such that the approximate solution of the integro-differential equation is verified. Can you shed light on these differences? Since it is your main Theorem, it would be nice to have an explicit discussion of what goes on where. More in detail, why do you get from existance in the appendix to a whole interval of $\epsilon$ (that by the way gives a rather slow decay wrt to say $\epsilon > 1$), to claim in text "for any inverse power"? I have to be honest here and I did not spend time re-doing all the proofs. 
- Related to the point above, in equation (27), your supremum is over $T\wedge \hat{\tau}_M$. In line 531 you state that you will not specify which stopping time it is when it is "clear from context", but for (27), given that you want to prove (26), my natural question is if stability is for $\hat{\tau}_M(S)$ or $\hat{\tau}_M(\mathscr{S})$, as you only mention one. At the very least, it is not clear from context here which stopping time it is in (27). 
- Lemma C.1 requires as condition that $S(t, z)$ is an approximate solution over the $\xi$-mesh. How do you guarantee that this condition holds in your case? If ever, it is not evident from text that this is the case. If it is a direct consequence of arguments in reference 14, I would appreciate it being made explicit. 
- Can you clarify the sentence of lines 239-240? Ignore the "sufficiently anisotropic" aspect which I already discussed above. The idea that smaller learning rates lead to under-performance is not entirely true. In the simplest setting if you over-shoot with your learning rate you will never converge. I believe here this is tightly linked to your Table 1, but there line search and Polyak have the same convergence rate for the risk. Where is the difference? 
- In Figure 1, how many runs do you make? It would be nice to know.
- The same question above but for any other plot where you show error bars. 
- (line 1125) You choose as scaling for the initialization $X_0\sim\mathcal{N}\left(0,\frac{I_d}{\sqrt{d}}\right)$  and the ground truth $X^{\star} = \frac{1}{\sqrt{d}}\mathbf{1}_d$. The norm of the latter is $1$, but the former has norm concentrating at $\sqrt{\sqrt{d}}$, and this violates your assumptions (not bounded by a constant independent of $d$). Is it a typo? Or is it that the experiment does not match the assumptions? If I look at the scaling of Figure 3, line 1099, you indeed sample $X_0\sim \mathcal{N}\left(0,\frac{I_d}{d}\right)$. 
- Proposition 4.4 is nice. What happens when $\beta +\delta > 2$?

My soundness score is due to the many questions. I hope I will raise it after the discussion. The overall score follows the same principle.
limitations: limitations: None.

Comments: title: Response
Rebuttal update
Regarding gradient flow, different scalings
General Response
(Response to questions)
(Response to questions)
Proposed conclusion/limitations paragraphs
comment: I thank the authors for their feedback.

> Line search

I agree the practical limitation of line search is mostly seen in stochastic optimization. Orvieto et al. [40] had a counter-example for the convergence of stochastic Polyak stepsize due to a bias, and I think the same example should reveal a bias of stochastic line search. Your result, however, is quite different, and I agree your characterization of the rate is new.

> the standard analysis would say that you need to use a decreasing learning rate with SGD, but nobody does this in practice

I beg to disagree. Cosine annealing is commonly used with SGD, not just Adam. I agree, however, that nobody uses the standard $1/\sqrt{t}$ and $1/t$ schedules. 

> The point of this paper is to provide another tool that allows for a finer analysis of these algorithms

I wish there was something in the middle: more precise than the standard theory and less convoluted than the one you offer in this submission. I do agree however that the one presented here has value when we require a finer analysis.

> The related work paragraph (Page 4 Lines 106-110) describes past works which use a high-dimensional framework for analyzing dynamics of stochastic algorithms.

True, I did underestimate the amount of work that has been done in this setting.

> Intuition and existence of $h$

To be honest, I am still struggling with understanding the existence and role of this function.

> Mild assumptions on the learning rate

I wouldn't use the term "mild" in this setting since it's hard to quantify how mild they are. My (optional) suggestion is to explicitly state the the assumption applies to several learning rates that are of interest.

I'll update my score for the paper given the points above.
Thanks, I will keep my score.
Dear authors, thank you for your response, including the general rebuttal. Based on the rebuttal, I am willing to increase the score to 5. Let me clarify the reasoning behind it, so that the authors and AC have a possibility to calibrate the score. The positive things are 
  - proposed changes in writing are good
  - the discussion of relation with [14] adds clarity and transparency 
  - the comments regarding different scaling of learning rate with dimension $d$, as well as additional discussion of power-law rates, help to better identify the positioning and relation of the current paper with other approaches.

What prevents me from raising the score further (e.g. 7) is that the paper still explores just a little bit of the behaviors of adaptive algorithms. The theorem 2.1 is a solid foundation for studying adaptive algorithms. However, it really pays off with its subsequent application to different adaptive optimization scenarios. The applications in the current manuscript (sections 3 and 4) still feel as kind of short and not properly explored. The same goes to the application to non-quadratic problems - authors comment that they are roughly the same as quadratic problems, which does not sound satisfying.
Dear reviewer,

We want to address the question you re-raised, which we didn't completely answer:

> I am concerned about the scaling. My understanding of the literature is that scalings are very important. Your assumptions require that is bounded independently of the dimension and that is bounded, independently of the dimension. Moreover, your step-size is ...
  
> What is done in some other references is very different. Let me take an example, and correct me if I am wrong.

Let's take the example you posed.  The steps of SGD look like

$$ X_{k+1} = X_k - \eta a_{k+1} f'\left( \langle a_{k+1}, X_k \rangle \right) $$

Now you put $a_{k+1}$ to be scaled to have norm $O(1)$ and $X_{0}$ to have norm $O(\sqrt{d}))$.  This is not our setup, but it is possible to change variables to put it in our setup.

Let $Y_{k} = X_{k}/\sqrt{d}$ and let $b = a \sqrt{d}$.  Now in these variables, the equation above becomes 

$$ Y_{k+1} = Y_k - \frac{\eta}{d} b_{k+1}  f'\left( \langle b_{k+1}, Y_k \rangle \right) $$

So indeed, in the setup that you proposed, it would be correct to take $\eta  = O(1)$, and it can be equivalently represented in our setup with stepsize scaling like $O(1/d)$.  What's important in both cases is the scale of the inner product $\langle b_{k+1}, Y_k \rangle$ is $O(1)$.

So our setup is not gradient flow, and the (not-degenerate) high-dimensional limit for the scaling you propose is the same as ours.  

By the way, we'd be happy to look at any literature involving the scalings that you propose, so that we can make a better comparison with existing literature.

Please let us know if this answers all your questions regarding this setup.
Thanks!!
Oops! Yes, we forgot to add the responses to the last questions.  Sorry!

* **Comparison of Line search and Polyak risk convergence rate.** As you have observed, Table 1 indicates that the convergence rates for the risk in line search and Polyak have the same formula as a function of the limiting learning rate, $\\gamma\_{\\infty}$.  The key difference is in the value of $\\gamma\_{\\infty}$ in these two cases.  For data that is strongly anisotropic in the sense described above, the value of $\\gamma\_{\\infty}$ can be much smaller for line search than for Polyak, thus yielding slower convergence.  
    
* **Number of runs for simulations.** The number of runs for the simulations, as well as other implementation details, are provided in Appendix H. Due to space constraints, we could not include these details in the main body of the paper.  
    
* **Scaling of initialization in line 1125\.** This is indeed a typo, and thank you for pointing it out.  *We corrected it.*  
    
* **What happens when $\\beta+\\delta\>2$.** This is a natural question and thank you for asking.  We do not address the case of $\\beta+\\delta\>2$ because it represents a setting in which the high-dimensional problem devolves to a finite-dimensional one where our scalings no longer really make sense.  This is perhaps easiest to see in our formulation of the power law set-up for the covariance, where we take the spectrum of $K$ to approach a continuous density function, supported on $(0,1)$, with unbounded density near 0\.  This density is only well-defined for $0\<\\beta\<1$.  Similarly, one can view the power-law set-up for $\\mathscr{D}^2\_i(0)$ giving the distribution of the projections of $(X\_0-X^{\\star})^2$ in the eigenvector directions.  This distribution approaches a well-defined, continuous density function when $0\<\\delta\<1$.  Thus, the limits that we consider cease to make sense when $\\beta+\\delta\>2$.

  One can also see that the proofs break down when $\\beta+\\delta\>2$.  Our analysis of Adagrad-Norm in the power-law setting (see Section D.2.3) relies upon the fact that the deterministic limit of the risk can be expressed as a convolution Volterra equation (see (67)).  Under the assumption that the forcing function F has power law decay (see Assumption 8), we can derive the asymptotics of the learning rate and the risk for various regimes, as displayed in Table 1\.  Assumption 8 is satisfied by our power-law set-up, provided that $\\beta+\\delta\<2$ (see Lemma D.5).
Dear author(s), 
thank you for your general rebuttal and for the rebuttal concerning my specific comments. I understand all your points, and after the corrections/specifications you put, I will raise my score. Some quick points: 
- strong anisotropy and verifying the approximate solution: you comments in italics are necessary in my opinion. 
- Polyak comment: yes, please. 
- scalings comment: see one of the questions missing (at the very end) in my long list to close this aspect. 
- all other typos: thank you for fixing them. 


 last comment: This is a very deep paper, maybe a journal deserves it.

My score is higher. I hope the other reviewers will raise their very low grades. Some scores, in their very description, are too harsh with respect to the actual review. 

My score reflects my opinion on the exact description provided by Neurips: moderate-to-high impact. I stand with the idea that this paper is definitely not a 3. 

Good luck!!
Acknowledged, see full comment below. I believe you are missing the last three questions (one might be a typo therefore I ping you about this).
**Answers to the reviewers’ questions**

* **Strong anisotropy.** When we say “strong anisotropy” we are not referring to the lower bound on the eigenvalues.  Rather, we refer to context where, for fixed $\\lambda\_{\\min}$, we have $\\lambda\_{\\min}/(\\text{Tr}(K^2)/d)\\ll (\\text{Tr}(K)/d)^{-1}$ or in other words, $\\text{Tr}(K)/d\\ll\\text{Tr}(K^2)/d$.  This is in contrast to the isotropic case where $\\text{Tr}(K)/d=\\text{Tr}(K^2)/d$.  *We can make this more clear by writing $\\text{Tr}(K)/d\\ll\\text{Tr}(K^2)/d$ in parentheses after the phrase “strong anisotropy.”*  
    
* **Polyak stepsize.** Indeed, we are using the derivation of Polyak stepsize (see, e.g.,  \[1\]) as motivation to construct what should be the “idealized” stochastic version and not the $( R(X\_k)-R(X^\*) ) / ||\\nabla R(X\_k)||^2$. We will add commentary about the derivation of the Polyak step size in the paper and how it deviates from the classical Polyak stepsize. We appreciate the reviewer pointing this out.   
    
  \[1\] Elad Hazan and Sham Kakade. “Revisiting the Polyak Step Size”.   
    
* **Concentration as an inverse power of d.** The concentration for the distance from SGD statistics to their deterministic equivalent is indeed $d^{-\\varepsilon}$ for $0\<\\varepsilon\<1/2$, so the best bound we obtain for this concentration is slightly worse than $1/\\sqrt{d}$.  The comment in line 48 is that the concentration occurs with **probability** better than any inverse power of d.  In other words, the probability that our statistic will be less than $d^{-\\varepsilon}$ is at least $1-d^{-C}$ for arbitrarily large C.  In the theorem, we describe this more succinctly in the phrase “with overwhelming probability.”  However, the definition of overwhelming probability is not introduced until line 74, so we do not use it in the comment on line 48\.   
    
* **Stopping times.** This is indeed unclear, and thank you for pointing it out.  It is meant to be the minimum of the two stopping times, as in the preceding equation and in the Proposition that is cited from \[14\].  *We can clarify this in the updated version of the paper.*  
    
* **Verifying $S(t,z)$ is an approximate solution.** Indeed, this is a very important condition to verify, and doing so is the primary purpose of Section C.2.  The way in which this is utilized with regard to the mesh is briefly described in Section C.1.3 where we prove Proposition C.1 (pending the details that are in Section C.2).  *Based on your question, we can add a sentence at the location that you referenced summarizing the flow of the proof and where these details can be found.*
We thank the reviewer for finding typos, which have now been fixed. We answer below the reviewer's additional questions. 

**Intuition and existence of $h$ and why is $R(X)$ is “an expectation of a Gaussian vector”**

The intuition for $h$ comes from the fact that $a^TX$ and $a^TX^{\\star}$ are correlated Gaussians with covariance given by $\[X|X^{\\star}\]^T K \[X|X^{\\star}\]$. Here $[X|X^{\\star}]$ means one concatenates $X$ and $X^{\\star}$.  

 Since $a \\sim N(0,K)$, for any fixed vector $X$, then $y \= a^TX \\sim N(0, X^TKX)$. Both $y \= a^T X$ and $y^{\\star} \= a^TX^{\\star}$ are Gaussians, but they are correlated normal since they both depend on the same $a$. A simple computation shows that the correlation between $y$ and $y^{\\star}$ is given by $\[X |X^{\\star}\]^T K \[X | X^{\\star}\]$ where $\[X|X^{\\star}\]$ is the concatenated vector. This is to say that $(y, y^{\\star})$ is Gaussian and $(y, y^{\\star}) \\sim N(0, \[X|X^{\\star}\]^T K \[X|X^{\\star}\]$. (See Section 1.1). 

Suppose $\epsilon =0$ as otherwise this will just add another integral. From Eq (1), using the density of multivariate Gaussian $(N(0, \[X|X^{\\star}\]^T K \[X|X^{\\star}\])$,  
$$  
R(X) \= \\int f(y,y^{\\star}, \\epsilon) e^{(y, y^{\\star})^T G (y, y^{\\star})} \\, dy dy^{\\star}
$$  
where G is the 2 x 2 matrix, $\[X|X^{\\star}\]^T K \[X|X^{\\star}\]$. Integrating out $y$ and $y^{\\star}$,   
$\\int f(y,y^{\\star}, \\epsilon) e^{(y, y^{\\star})^T G (y, y^{\\star})} \\, dy dy^{\\star}$ is a function of $\[X|X^{\\star}\]^T K \[X|X^{\\star}\]$. 

This function is $h$. Moreover, this also explains why $R(X)$ is an “expectation of a Gaussian vector” as $(y, y^{\\star}) \= (a^TX, a^T X^{\\star})$ is distributed as a multivariate Gaussian with covariance $\[X|X^{\\star}\]^T K \[X|X^{\\star}\]$ and you are taking the expectation over $(y, y^{\\star}, \epsilon)$. 

**“We shall remove this condition by approximation in what follows”**

The reviewer is absolutely right: we did not end up removing the condition, although there is a straightforward way to do so, and we will correct this.  In short, one creates an $f\_\\epsilon$ which is an approximation to $f$ formed by convolving with an isotropic Gaussian of variance $\\epsilon$.  This is $C^2$ and has bounded second derivatives (as $f$ was smooth).  We then will take limits as $\\epsilon \\to 0$.  

We note that Lemma B.1 was intended as discussion, and it is not used in the main theorem.  Theorem B.1 generalizes to the more standard L-smooth f, at the cost of additional complexities in the formulation (discussed below the Lemma), and is intended to show why we didn’t simply use the more standard L-smooth f assumptions.

**“Mild assumptions on the learning rate” Why mild?**

The learning rate assumptions are mild since they encompass many adaptive learning rates used in practice; they do exclude learning rates which are non-concentrating.  Among non-concentrating stepsizes, one can further consider those which bias the gradient and those which do not bias the gradient.  Those which bias the gradient require a totally different theory (e.g. scaling the gradient to be norm-1 or things like gradient clipping).  Those which do not bias the gradient would have a similar theory as those which we consider, but with an additional variance term in the ODEs (e.g. RMS-prop with a very short window of averaging, with respect to dimension).  So we chose this case to keep the theory *as simple as possible* (which we know is aligned with the reviewer’s desires\!)
**Conclusion paragraph (to be added):**  *This work studies stochastic adaptive optimization algorithms when data size and parameter size are large, allowing for nonconvex and nonlinear risk functions, as well as data with general covariance structure. The theory shows a concentration of the risk, the learning rate and other key functions to a deterministic limit, which is described by a set of ODEs. The theory is then used to derive the asymptotic behavior of the AdaGrad-Norm and idealized exact line search on strongly convex and least square problems, revealing the influence of the covariance matrix structure on the optimization.  A potential extension of this work would be to study other adaptive algorithms such as D-adaptation, DOG, and RMSprop which are covered by the theory. Studying the asymptotic behavior of the risk and the learning rate may improve our understanding of the performance and scalability of these algorithms on more realistic data. Another important application of the theory would be to analyze the ODEs presented here on nonconvex problems.*  

**Limitations (to be added):** *The current form of the theory is limited to Gaussian data, though many parts of the proof can be extended easily beyond Gaussian data. The main ODE comparison theorem is also only tuned for analyzing problem setups where the trace of the covariance is on the order of the ambient dimension; when the trace of the covariance is much smaller than ambient dimension, other stepsize scalings of SGD are needed. In addition, the analysis is limited to the streaming stochastic adaptive methods. We conjecture that a similar deterministic equivalent holds also for multi-pass algorithms at least for convex problems. This has already been shown in the least square problem for SGD with a fixed deterministic learning rate \[42,44\]. Lastly, numerical simulations on real datasets (e.g., CIFAR-5m) suggests that the predicted risk derived by our theory matches the empirical risk of multipass SGD beyond Gaussian data (see for example Figure 6).*