Title: The High Line: Exact Risk and Learning Rate Curves of Stochastic Adaptive Learning Rate Algorithms
Venue: NeurIPS 2024
Accepted: yes
Keywords: Stochastic optimization

high-dimensional probability/statistics

adaptive stepsizes
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.

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

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.

1 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

10

2

10

1

10 0

10 1

SGD Iterations/

d

10 0

Learning rate

10 0

6×10 1

Risk

AdaGrad-Norm Least Squares

d

=256

d

=1024

d

=4096

d

=16384

Theory, learn. rate

Theory, risk

10

1

10 0

10 1

SGD Iterations/

d

10 0

9.2 × 10

1

9.3 × 10

1

9.4 × 10

1

9.5 × 10

1

9.6 × 10

1

9.7 × 10

1

9.8 × 10

1

9.9 × 10

1

Learning rate

10

4

10

3

10

2

10

1

10 0

Risk

AdaGrad-Norm Logistic Regression

d

=16

d

=32

d

=64

d

=128

Theory, learn. rate

Theory, risk

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 : R 3 → R composed with the linear functions ⟨ X,a ⟩ , ⟨ X ⋆ , a ⟩ . That is, we seek to solve

We 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 better 2 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).

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 ∈ A , we have

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 than E [ ∥ 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.

1.1 Model Set-up

Wesuppose 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 composite 4 , of the form

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 .

Assumption 3 (Risk representation) . There exists a function h : R 2 × 2 → R such that h ( B ) = R ( X ) is differentiable and satisfies

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).

1.2 Algorithmic set-up

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 update

where 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 γ : R + × D ([0 , ∞ )) × D ([0 , ∞ )) × D ([0 , ∞ )) → R + by 6

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.

We also define a conditional expectation version of G k where the filtration F k = σ ( X ⋆ , X 0 , . . . , X k ) :

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 ,

Finally, γ is bounded, i.e., there exists a constant ˆ C = ˆ C ( γ ) > 0 independent of d so that

The 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 Table 2, Sec. A, and Sec. C.3).

2 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 projections V 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

The pair of functions h (Assumption 3) and I (Assumption 4) can be evaluated simply:

The deterministic dynamics for the risk R ( t ) in this case can be simplified to:

This is a convolution Volterra equation with a convergence threshold of γ t < 2 d Tr K [14, 44, 46, 47].

In the noiseless label case (i.e., ϵ = 0 ), the risk is given by R ( t ) = 1 2 d ∑ 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:

We 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. 7

We now introduce a coupled system of differential equations, which will allow us to model the behaviour of our learning algorithms. For the i th ( λ i , ω i ) -eigenvalue/eigenvector of K , set

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

We now introduce a system of coupled ODEs for each ( λ i , ω i ) -eigenvalue/eigenvector pair of K

with 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 any T > 0

The 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 are

where D 2 i ( t ) corresponds D 2 i ( X k ) def = d × ( ⟨ X k -X ⋆ , ω i ⟩ ) 2 .

3 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.

SGD (line search)

10-2

theory (line search)

10-8

SGD (Polyak)

theory (Polyak)

risk

10-14

Tr(K²) ≈ 2.25

10-20

Tr(K²) ≈ 1.77

Tr(K2) ≈ 1.27

0

20

40

60

time (t)

Tr(K²) ≈ 1.12

4 × 100

3×100

quotientt/

2 × 100

100

20

40

60

time (t)

Polyak Stepsize. A natural threshold to consider is the largest learning rate such that d D ( t ) < 0 , which we denote by ¯ γ D t . Using the least squares ODE (8), this is precisely

Without label noise, (12) simplifies to ¯ γ D t = ¯ g D k = 2 Tr ( K ) /d , the exact threshold for convergence of least squares.

Agreedy stepsize strategy would maximize the decrease in the distance to optimality at each iteration, denoted by us as Polyak stepsize , γ Polyak t ∈ arg min γ d D ( t ) . In the case of least squares, this is

The latter yields the optimal fixed learning rate (up to absolute constant factors) for a noiseless target on a least squares problem [35, 43]). 8

Exact Line Search. In the context of risk, using (8) and noting that R ( t ) = 1 2 d ∑ d i =1 λ i D 2 i ( t ) , we can find γ line t ∈ arg min d R ( 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

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 . Then

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).

10 0

10 2

10 4

10 6

SGD iterations

5×10 1

6×10 1

7×10 1

8×10 1

Risk

(Noisy) Least Squares AdaGrad,

=1.0,

b

= 1.0,

= 2.5, d = 500

10 0

8.8 × 10

1

9×10 1

9.2 × 10

1

9.4 × 10

1

9.6 × 10

1

9.8 × 10

1

Learning rate / t^(-1/2)

SGD

Predicted

10 0

10 1

10 2

10 3

10 4

10 5

SGD iterations

10 0

2×10 0

3×10 0

4×10 0

learning rate

SGD,  =719.69

SGD,  =100.00

SGD,  =51.79

SGD,  =31.62

SGD,  =12.33

10 0

10 1

10 2

10 3

10 4

10 5

SGD iterations

10 0

learning rate

SGD, ||

X

0

X

*

|| 2 =1.0

SGD, ||

X

0

X

*

|| 2 =2.0

SGD, ||

X

0

X

*

|| 2 =4.0

SGD, ||

X

0

X

*

|| 2 =8.0

SGD, ||

X

0

X

*

|| 2 =16.0

4 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 AdaGradNorm 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 equation

The solution ∫ t 0 R ( s ) d s 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 d s < ∞ with γ s as in Table 2 (AdaGrad-Norm for least squares), then γ t ≍ 1 b η + η 2 4 d 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 of K 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 .

10

1

10 0

10 1

10 2

10 3

10 4

10 5

10 6

10 7

time (t)

10

4

10

2

10 0

10 2

10 4

10 6

risk

t^(beta + delta - 2)

delta + beta = 0.2

delta + beta = 0.4

delta + beta = 0.6

delta + beta = 0.8

delta + beta = 1.0

delta + beta = 1.2

delta + beta = 1.4

delta + beta = 1.6

delta + beta = 1.8

t^(1-2/(beta + delta))

delta + beta = 2.0

10

5

10

3

10

1

10 1

10 3

10 5

10 7

time (t)

10

6

10

5

10

4

10

3

10

2

10

1

10 0

learning rate

constant

delta + beta = 0.2

delta + beta = 0.4

delta + beta = 0.6

delta + beta = 0.8

delta + beta = 1.0

delta + beta = 1.2

delta + beta = 1.4

delta + beta = 1.6

delta + beta = 1.8

delta + beta = 2.0

t^(-1 + 1/(beta+delta))

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.

5 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).

Acknowledgments and Disclosure of Funding

E. Collins-Woodfin was supported by Fonds de recherche du Québec - Nature et technologies (FRQNT) postdoctoral training scholarship and Centre de recherches mathématiques (CRM) Applied math postdoctoral fellowship. Research of B. García Malaxechebarría was in part funded by NSF DMS 2023166 (NSF TRIPODS II). Research by E. Paquette was supported by a Discovery Grant from the Natural Science and Engineering Council (NSERC). C. Paquette is a Canadian Institute for Advanced Research (CIFAR) AI chair, Quebec AI Institute (MILA) and a Sloan Research Fellow in Computer Science (2024). C. Paquette was supported by a Discovery Grant from the Natural Science and Engineering Research Council (NSERC) of Canada, NSERC CREATE grant Interdisciplinary Math and Artificial Intelligence Program (INTER-MATH-AI), Google research grant, and Fonds de recherche du Québec - Nature et technologies (FRQNT) New University Researcher's Start-Up Program. Additional revenues related to this work: C. Paquette has 20% part-time employment at Google DeepMind.

References

Luca Arnaboldi, Florent Krzakala, Bruno Loureiro, and Ludovic Stephan. Escaping mediocrity: how two-layer networks learn hard single-index models with SGD. arXiv preprint arXiv:2305.18502 , 2023.

Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From highdimensional and mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks. arXiv preprint arXiv:2302.05882 , 2023.

Krishna B Athreya, Peter E Ney, and PE Ney. Branching processes . Courier Corporation, 2004.

Amit Attia and Tomer Koren. SGD with adagrad stepsizes: full adaptivity with high probability to unknown parameters, unbounded gradients and affine variance. In International Conference on Machine Learning , pages 1147-1171. PMLR, 2023.

Krishnakumar Balasubramanian, Promit Ghosal, and Ye He. High-dimensional scaling limits and fluctuations of online least-squares SGD with smooth covariance. arXiv preprint arXiv:2304.00707 , 2023.

Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for SGD: Effective dynamics and critical scaling. In Advances in Neural Information Processing Systems , volume 35, pages 25349-25362, New York, 2022. Curran Associates, Inc.

Leonard Berrada, Andrew Zisserman, and M Pawan Kumar. Training neural networks for and by interpolation. In International conference on machine learning , pages 799-809. PMLR, 2020.

R. Berthier, F. Bach, and P. Gaillard. Tight Nonparametric Convergence Rates for Stochastic Gradient Descent under the Noiseless Linear Model. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Michael Biehl and Peter Riegler. On-line learning with a perceptron. Europhysics Letters , 28 (7):525, 1994.

Michael Biehl and Holm Schwarze. Learning by on-line gradient descent. Journal of Physics A: Mathematical and general , 28(3):643, 1995.

B. Bordelon and C. Pehlevan. Learning Curves for SGD on Structured Features. In International Conference on Learning Representations (ICLR) , 2022.

Michael Celentano, Chen Cheng, and Andrea Montanari. The high-dimensional asymptotics of first order methods with random data. arXiv preprint arXiv:2112.07572 , 2021.

Kabir Aladin Chandrasekher, Ashwin Pananjady, and Christos Thrampoulidis. Sharp global convergence guarantees for iterative nonconvex optimization with random data. Ann. Statist. , 51(1):179-210, 2023. ISSN 0090-5364,2168-8966. doi: 10.1214/22-aos2246. URL https: //doi.org/10.1214/22-aos2246 .

Elizabeth Collins-Woodfin and Elliot Paquette. High-dimensional limit of one-pass SGD on least squares. Electronic Communications in Probability , 29:1-15, 2024. doi: 10.1214/23-ECP571.

Elizabeth Collins-Woodfin, Courtney Paquette, Elliot Paquette, and Inbar Seroussi. Hitting the high-dimensional notes: An ODE for SGD learning dynamics on GLMs and multi-index models. arXiv preprint arXiv:2308.08977 , 2023.

Alex Damian, Eshaan Nichani, Rong Ge, and Jason D Lee. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. In Advances in Neural Information Processing Systems , volume 36, pages 752784, 2023. URL https://proceedings.neurips.cc/paper_files/paper/2023/file/ 02763667a5761ff92bb15d8751bcd223-Paper-Conference.pdf .

Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborová, and Florent Krzakala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. arXiv preprint arXiv:2402.03220 , 2024.

Aaron Defazio and Konstantin Mishchenko. Learning-rate-free learning by d-adaptation. In International Conference on Machine Learning , pages 7449-7479. PMLR, 2023.

John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research , 12(7), 2011.

Darina Dvinskikh, Aleksandr Ogaltsov, Alexander Gasnikov, Pavel Dvurechensky, Alexander Tyurin, and Vladimir Spokoiny. Adaptive gradient descent for convex and non-convex stochastic optimization. arXiv preprint arXiv:1911.08380 , 2019.

Matthew Faw, Litu Rout, Constantine Caramanis, and Sanjay Shakkottai. Beyond uniform smoothness: A stopped analysis of adaptive SGD. In The Thirty Sixth Annual Conference on Learning Theory , pages 89-160. PMLR, 2023.

Cédric Gerbelot, Emanuele Troiani, Francesca Mignacco, Florent Krzakala, and Lenka Zdeborová. Rigorous dynamical mean-field theory for stochastic gradient descent methods. SIAM Journal on Mathematics of Data Science , 6(2):400-427, 2024. doi: 10.1137/23M1594388. URL https://doi.org/10.1137/23M1594388 .

Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems , 32, 2019.

Sebastian Goldt, Marc Mézard, Florent Krzakala, and Lenka Zdeborová. Modeling the influence of data structure on learning in neural networks: The hidden manifold model. Physical Review X , 10(4):041044, 2020.

Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Marc Mézard, and Lenka Zdeborová. The gaussian equivalence of generative models for learning with shallow neural networks. In Mathematical and Scientific Machine Learning , pages 426-471, New York, New York, USA, 2022. PMLR.

Robert M Gower, Aaron Defazio, and Michael Rabbat. Stochastic polyak stepsize with a moving target. arXiv preprint arXiv:2106.11851 , 2021.

Elad Hazan and Sham Kakade. Revisiting the polyak step size. arXiv preprint arXiv:1905.00313 , 2019.

Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Cited on , 14(8):2, 2012.

Maor Ivgi, Oliver Hinder, and Yair Carmon. DoG is SGD's best friend: A parameter-free dynamic step size schedule. arXiv preprint arXiv:2302.12022 , 2023.

Xiaowen Jiang and Sebastian U Stich. Adaptive SGD with polyak stepsize and line-search: Robust convergence and variance reduction. Advances in Neural Information Processing Systems , 36, 2024.

Kiwon Lee, Andrew N. Cheng, Courtney Paquette, and Elliot Paquette. Trajectory of MiniBatch Momentum: Batch Size Saturation and Convergence in High Dimensions. To Appear in NeurIPS 2022 , art. arXiv:2206.01029, June 2022.

Kfir Levy. Online to offline conversions, universality and adaptive minibatch sizes. Advances in Neural Information Processing Systems , 30, 2017.

Kfir Y Levy, Alp Yurtsever, and Volkan Cevher. Online adaptive methods, universality and acceleration. Advances in neural information processing systems , 31, 2018.

Xiaoyu Li and Francesco Orabona. On the convergence of stochastic gradient descent with adaptive stepsizes. In Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics , volume 89 of Proceedings of Machine Learning Research , pages 983-992, 2019.

Nicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacoste-Julien. Stochastic polyak step-size for SGD: An adaptive learning rate for fast convergence. In International Conference on Artificial Intelligence and Statistics , pages 1306-1314. PMLR, 2021.

H Brendan McMahan and Matthew Streeter. Adaptive bound optimization for online convex optimization. arXiv preprint arXiv:1002.4908 , 2010.

F. Mignacco, F. Krzakala, P. Urbani, and L. Zdeborová. Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification. In Advances in Neural Information Processing Systems , volume 33, pages 9540-9550, 2020.

P. Nakkiran, B. Neyshabur, and H. Sedghi. The Deep Bootstrap Framework: Good Online Learners are Good Offline Generalizers. In International Conference on Learning Representations (ICLR) , 2021.

Angelia Nedi´ c and Dimitri Bertsekas. Convergence rate of incremental subgradient algorithms. Stochastic optimization: algorithms and applications , pages 223-264, 2001.

J. Nocedal and S. J. Wright. Numerical Optimization . Springer New York, 2 edition, 2006.

Antonio Orvieto, Simon Lacoste-Julien, and Nicolas Loizou. Dynamics of SGD with stochastic polyak stepsizes: Truly adaptive variants and convergence to exact solution. Advances in Neural Information Processing Systems , 35:26943-26954, 2022.

C. Paquette and K. Scheinberg. A stochastic line search method with expected complexity analysis. SIAM J. Optim. , 30(1):349-376, 2020. ISSN 1052-6234. doi: 10.1137/18M1216250. URL https://doi.org/10.1137/18M1216250 .

C. Paquette, K. Lee, F. Pedregosa, and E. Paquette. SGD in the Large: Average-case Analysis, Asymptotics, and Stepsize Criticality. In Proceedings of Thirty Fourth Conference on Learning Theory (COLT) , volume 134, pages 3548-3626, 2021.

Courtney Paquette and Elliot Paquette. Dynamics of stochastic momentum methods on largescale, quadratic models. In Advances in Neural Information Processing Systems , volume 34, pages 9229-9240, 2021. URL https://proceedings.neurips.cc/paper/2021/file/ 4cf0ed8641cfcbbf46784e620a0316fb-Paper.pdf .

Courtney Paquette and Elliot Paquette. High-dimensional optimization. SIAM Views and News , 20:16pp, December 2022.

Courtney Paquette, Elliot Paquette, Ben Adlam, and Jeffrey Pennington. Homogenization of SGD in high-dimensions: Exact dynamics and generalization properties. arXiv e-prints , art. arXiv:2205.07069, May 2022.

Courtney Paquette, Elliot Paquette, Ben Adlam, and Jeffrey Pennington. Implicit regularization or implicit conditioning? exact risk trajectories of sgd in high dimensions. In Advances in Neural Information Processing Systems , volume 35, pages 35984-35999, New York, 2022. Curran Associates, Inc.

Boris T Polyak. Introduction to optimization. 1987.

Michal Rolinek and Georg Martius. L4: Practical loss-based stepsize adaptation for deep learning. Advances in neural information processing systems , 31, 2018.

David Saad and Sara Solla. Dynamics of on-line gradient descent learning for multilayer neural networks. In Advances in Neural Information Processing Systems , volume 8. MIT Press, 1995.

David Saad and Sara A Solla. Exact solution for on-line learning in multilayer neural networks. Physical Review Letters , 74(21):4337, 1995.

A. Varre, L. Pillaud-Vivien, and N. Flammarion. Last iterate convergence of SGD for LeastSquares in the Interpolation regime. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

S. Vaswani, A. Mishkin, I. Laradji, M. Schmidt, G. Gidel, and S. Lacoste-Julien. Painless Stochastic Gradient: Interpolation, Line-Search, and Convergence Rates. In Advances in Neural Information Processing Systems (NeurIPS) , volume 32, pages 3732-3745, 2019.

Sharan Vaswani, Issam Laradji, Frederik Kunstner, Si Yi Meng, Mark Schmidt, and Simon Lacoste-Julien. Adaptive gradient methods converge faster with over-parameterization (but you should do a line-search). arXiv preprint arXiv:2006.06835 , 2020.

M. Velikanov, D. Kuznedelev, and D. Yarotsky. A view of mini-batch SGD via generating functions: conditions of convergence, phase transitions, benefit from negative momenta. In International Conference on Learning Representations (ICLR) , 2023.

R. Vershynin. High-dimensional probability: An introduction with applications in data science . Cambridge University Press, Cambridge, UK, 2018. doi: 10.1017/9781108231596. URL https://doi.org/10.1017/9781108231596 .

Bohan Wang, Huishuai Zhang, Zhiming Ma, and Wei Chen. Convergence of adagrad for non-convex objectives: Simple proofs and relaxed assumptions. In The Thirty Sixth Annual Conference on Learning Theory , pages 161-190. PMLR, 2023.

Chuang Wang, Hong Hu, and Yue Lu. A solvable high-dimensional model of GAN. In Advances in Neural Information Processing Systems , volume 32, New York, 2019. Curran Associates, Inc.

Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. The Journal of Machine Learning Research , 21(1):9047-9076, 2020.

Alexander Wei, Wei Hu, and Jacob Steinhardt. More than a toy: Random matrix models predict how real-world neural representations generalize. In International Conference on Machine Learning , pages 23549-23588. PMLR, 2022.

Xiaoxia Wu, Rachel Ward, and Léon Bottou. Wngrad: Learn the learning rate in gradient descent. arXiv preprint arXiv:1803.02865 , 2018.

Yuege Xie, Xiaoxia Wu, and Rachel Ward. Linear convergence of adaptive stochastic gradient descent. In Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics , volume 108 of Proceedings of Machine Learning Research , pages 1475-1485, 2020.

Junchi Yang, Xiang Li, Ilyas Fatkhullin, and Niao He. Two sides of one coin: the limits of untuned SGD and the power of adaptive methods. Advances in Neural Information Processing Systems , 36, 2024.

Yuki Yoshida and Masato Okada. Data-dependence of plateau phenomenon in learning with neural network-statistical mechanical analysis. In Advances in Neural Information Processing Systems , volume 32, New York, 2019. Curran Associates, Inc.

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

Supplementary material

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.

A SGD adaptive learning rate algorithms and stepsizes

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 that

with the deterministic equivalent (see Section 2 and also C.3) for this learning rate being

In the case of the least squares problem, the quantity I ( B ( t )) is explicit and

Algorithm 1 AdaGrad-Norm

Require: Initialize η > 0 , X 0 ∈ R d , b ∈ R and set b 0 = b × d

Generate new sample

a

∼ N

(0

, K

)

,

ϵ

k

∼ N

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) is

In the case of the least squares problem, the quantity I ( B ( t )) is explicit and

(0

, ω

2

)

;

▷ updating learning rate

▷ updating step with stochastic gradient

Algorithm 2 RMSprop-Norm, α Exponential Moving Average

Require: Initialize η > 0 , X 0 ∈ R d , b ∈ R and set b 0 = d × b , α > 0 exponential moving avg. g -1 = d × η b 0 ; 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 for

B The Dynamical nexus

In 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 use W + = [ 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 .

B.1 Discussion of the assumptions on 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:

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

Furthermore, h is continuously differentiable on { B : det B = 0 } and its derivative ∇ h satisfies an estimate

̸

Proof. To derive the existence of h , note that

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 = √ Bz , and hence we have

As the function f is absolutely continuous with a Lipschitz gradient, we can differentiate under the integral sign and conclude

For the differentiability of h , suppose for the moment that f is C 2 with bounded second derivatives. 10 Setting Q = √ B the positive semi-definite square root of B , we have

Then using the chain rule, and setting ∂ i f to be the i -th partial derivative of f ,

̸

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

As a matrix equation, this can be written as

This is a linear equation in Dh . When Q ≻ 0 , we can define

and note

Moreover, the mapping M ↦→ ∫ ∞ 0 e -tQ Me -tQ d t defines a two-sided inverse for M ↦→ MQ + QM , and so Dh = A. Note that by symmetry of J , Q , and Dh

and therefore and so taking inverses on both sides, Dh = J.

Undoing Stein's Lemma, we have Q ( Dh ) = ( Dh ) Q = M, where M ij = E ( E ( z j ∂ i f ( Qz,ϵ ) | ϵ )) . From L -smoothness of f

Hence

Note Q -1 2 M ( Q 2 ) = ( Dh )( Q 2 2 ) is bounded by L ( f ) , and so we arrive at

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 with Q = √ B ,

Then

Applying Cauchy-Schwarz and using the L -smoothness of 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

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.

B.2 Integro-differential equation for S ( t, z )

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

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 4

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,

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 : d S ( t, · ) = F ( z, S ( t, · )) d t (20) where F ( z, S ( t, · )) def = -2 γ t (( -1 2 πi ∮ Ω S ( t, z ) d z ) H ( B ( t )) + H T ( B ( t )) ( -1 2 πi ∮ Ω S ( t, z ) d z )) + γ 2 t d [ Tr ( KR ( z ; K )) I ( B ( t )) 0 0 0 ] (21) -γ t ( S ( t, z )(2 zH ( B ( t ))) + (2 zH T ( B ( t ))) S ( t, z )) . Here B ( t ) = -1 2 πi ∮ Ω z S ( t, z ) d z, H ( B ) = [ ∇ h 11 ( B ) 0 ∇ h 21 ( B ) 0 ] , γ t is defined in (9) , and the initialization is S (0 , z ) = W ⊤ 0 R ( z ; K ) W 0 . (22)

The 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 [15, Lemma 4.1].

Lemma B.2 (Equivalence to coupled ODEs.) . The unique solution of (21) with initial condition (22) is given by

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 function A : C → R as

Definition B.1 ( ( ε, M, T ) -approximate solution to the integro-differential equation) . For constants M,T,ε > 0 , we call a continuous function S : [0 , ∞ ) × C → R 2 × 2 an ( ε, M, T ) -approximate solution of (20) if with

then 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 probability

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 that

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

Then there is an ε > 0 such that for d sufficiently large, with overwhelming probability,

Proof. This is basically equivalent to [15, Corollary 4.2]. The only difference is that [15, Corollary 4.2] requires the boundedness of N ; however, since our function f is α -pseudo-Lipschitz with α ≤ 1 , this boundedness follows from [15, Proposition 1.2], 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.

C SGD-AL is an approximate solution

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 Ω ξ having cardinality, | Ω ξ | = C ( | Ω | ) d ξ .

Lemma C.1 (Net argument, [15], Lemma 5.1) . Fix T, M > 0 and let ξ > 0 . Suppose Ω ξ is a d -ξ mesh of Ω with | Ω ξ | = C · d ξ and positive C > 0 . Let the function S ( t, z ) = S ( W td , z ) satisfy

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,

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 have

Remark C.1. Fix M > 0 and define the stopping time on ∥ W td ∥ , ϑ = ϑ M , by

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 .

C.1 SGD-AL is an approximated solution

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,

Again, 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].

C.1.1 Shorthand notation

In the following sections, we will be using various versions of the stepsize γ . In order to simplify notation, we set

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.

C.1.2 SGD-AL under the statistic

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 obtain

We will now relate this equation to its expectation by performing a Doob decomposition, involving the following martingale increments and error terms:

We can then write

Extending X k into continuous time by defining X t = X ⌊ t ⌋ , we sum up (integrate). For this, we introduce the forward difference

giving us

where | ξ td | = ∣ ∣ ∣ ∣ ∫ t ( ⌊ td ⌋-1) /d d · ∆ φ ( X sd ) d s ∣ ∣ ∣ ∣ ≤ max 0 ≤ j ≤⌈ td ⌉ {| ∆ φ ( X j ) |} . With this, we obtain the Doob decomposition for SGD-AL:

From here, we can proceed as in [15, Section 5.3] to show that SGD-AL is an ( ε, M, T ) -approximated solution.

C.1.3 S ( W td , z ) is an approximate solution

Proof of Proposition C.1. The appropriate stepsize, as a function of W td , is

(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 , · ) satisfies

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 ∈ Ω ,

Next, fix a constant ξ > 0 . Let Ω ξ ⊂ Ω such that there exists a ¯ z ∈ Ω ξ such that | z -¯ z | ≤ d -ξ and the cardinality of Ω ξ , | Ω ξ | = Cd ξ where C > 0 can depend on ∥ K ∥ op. For all z ∈ Ω , we note that ˆ τ M ≤ ϑ cM (see Lemma C.2). Consequently, we evaluate the error with the stopped process W ϑ td def = W d ( t ∧ ϑ ) instead of using ˆ τ M . By Proposition C.2, the proof of which we have deferred to Section C.2, we have, for any ˆ δ > 0

We deduce that

An application of the net argument, Lemma C.1, finishes the proof after setting ˆ δ = δ/ 8 and ξ = δ/ 8 .

C.2 Error bounds

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 as

where W = [ X | X ⋆ ] . Then, for any z ∈ Ω and T, M, ζ > 0 , with overwhelming probability,

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.3, C.4, C.5, C.6, and C.7.

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.

C.2.1 Bounds on the lower order terms in the gradient and hessian

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,

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 ) . Then

We look at |E 1 | first.

For the first term, we use (6). We have

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 conclude

For the second term, we use (5). Again, since ∥ N k ∥ ∞ and ∥ Q k ∥ ∞ are bounded due to our stopping time, we have

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.

Now let us consider |E 2 | :

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.

Finally, consider |E 3 | :

By (6), the first term is bounded by ˆ C ( γ ) 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,

Summing up to k = Td and dividing through by 2 d , 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,

Proof. We have

We then have

Just as in the Hessian argument, (6) lets us bound E [ | γ ( G k ) + γ ( G k ) | 3 | F k ] 1 4 by d ε w.o.p., (5) lets us bound E [ | γ ( G k ) -γ ( G k ) | | F k ] 1 4 by d -δ/ 4 w.o.p., and Lemma C.5 lets us bound E [ | I 3 ( k ) | 2 | F k ] 1 2 by 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,

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,

Proof. For notational convenience, set ∆ M k = ∆ M grad d ( k/d ∧ ϑ ) , and F k = -γ ( G k ) I 3 ( k ) /d , so that

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 that

Set β = d -1+ ζ/ 2 and t = d -1 / 2+ ζ ; this becomes

However, M β k is not quite the martingale we started with: there is still an error term,

We bound this term in overwhelming probability. We have

̸

The second term is superpolynomially small by Lemma C.5; the first term is superpolynomially small by (6) and (C.6).

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

(any power of d would have worked). Combining the error term and the projected martingale, we find that, with overwhelming probability,

We can now take the maximum over k from 0 to Td 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,

Proof. The proof here is basically identical to the previous one. Again, set F 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 us

The error term is also quite similar:

̸

We have

both of which are superpolynomially small by (6) and Lemma C.5. For the expectation, we have

this product is superpolynomially small by (6), Lemma C.6, and Lemma C.5. Overall, we have, with overwhelming probability,

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 ∈ Ω ,

which is bounded by d -1 / 2 w.o.p. by the boundedness of I 1 , I 2 , I 3 , and γ ( B k ) .

Proof. We have, as above,

C.2.2 General bounds

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 as

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

Lemma C.4 (Preliminary bounds) . With f and ∆ k defined as above, for ε > 0 and λ ≥ 0 , we have

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 , then

Then

Now, since a k +1 is Gaussian, we can write a 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 ε , and

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 have

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

which is bounded by d 2 ε w.o.p. Now for the expectation:

For the first term, we have

where we are using the notation v i k to refer to the i th 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.

for any ε > 0 and λ ≥ 0 , we have

Proof of (43) in Lemma C.5. Using the fact that ∥∇ 2 φ ( X k ) ∥ op ≤ ∥ S ( W k , · ) ∥ Ω ,

Now, ∥ W k ∥ is bounded by the stopping time. From Lemma C.4, ∥ ∆ k ∥ 2 d is bounded by d ε w.o.p., and every moment of this expression is bounded independent of d , as desired.

Proof of (44) in Lemma C.5. We have

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

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:

The 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,

Proof. The first line, (46), follows from (40). For the first inequality, ∥ G k ∥ ∞ = max 0 ≤ j ≤ k ∥ ∆ j ∥ 2 d , 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,

as desired. The second line is more straightforward:

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 ∥ ∞ :

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:

both of which are bounded independent of d . Finally, a similar argument applies to G k :

We now prove a concentration result that closely follows [15, Proposition 5.6].

Lemma C.7 ([15], Lemma 5.2) . Suppose v ∈ R d is distributed N (0 , I d ) and U ∈ R d × 2 has orthonormal columns. Then

where v -U ( U T v ) ∼ N ( 0 , I d -UU T ) and UU T v ∼ N ( 0 , UU T ) with v -U ( U T v ) independent of UU 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

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 have F k = σ ( { W i } k i =0 ) ; set ˆ F k = σ ( { W i } k i =0 , { r i } k i =0 ) . A simple calculation shows that

To compute the conditional mean E [ a k +1 | ˆ F k ] and covariance ( a k +1 -E [ a k +1 | ˆ F k ])( a k +1 -E [ a k +1 | ˆ F k ]) ⊤ , we use Lemma C.7. By Assumption 1, we can write a k +1 = √ Kv k , for v k ∼ N (0 , I d ) .

Now we perform a QR-decomposition on √ KW k def = Q k R k where Q k ∈ R d × 2 with orthonormal columns and R k ∈ R 2 × 2 is upper triangular (and invertible). Set Π k def = Q k Q T k . In distribution,

As R k is invertible, by Lemma C.7,

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 that

Moreover the conditional covariance of a k +1 is precisely

Next, using that E [ H k | F k ] = H k , we expand (49) to get the leading order behavior

Taking the expectation with respect to F k , we obtain

where the error E k is defined as

The proof now turns to bounding the expectation of this error quantity.

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:

Taking expectations and using Cauchy Schwarz, we obtain

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

Proof. Using the Lipschitz condition on the stepsize, we have

as desired.

C.3 Specific learning rates

In this section, we confirm that AdaGrad-Norm satisfies Assumption 6. In the notation of Assumption 6, we have, for AdaGrad-Norm,

Note that this reduces to the discrete stepsize if we plug in g = G k :

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,

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. Setting F ( x ) = η √ b 2 + x , we have

for x ∈ [0 , ∞ ) . We thus have

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 2 b 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:

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:

Proof. Looking to remove the square roots, we have

For AdaGrad-Norm, we have

We now bound the sum above. Set F 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 us

where we set β = d ε/ 2 . This is close to the bound we want: the error is

̸

We have

which superpolynomially small by (40). The expectation is similar:

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

with overwhelming probability (note that this would be true for any power of d , by the definition of superpolynomially small.) We thus conclude that

with overwhelming probability. Multiplying by d , we find that

Plugging this back into (57), we find that

with overwhelming probability, and so, taking the square root,

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 ,

as desired.

D Proofs for AdaGrad-Norm analysis

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.

D.1 Strongly convex setting

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 that

exists and has norm bounded independent of d. Then one has,

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.

Definition D.1 ( ˆ L -smoothness of outer function f ) . A function f : R 3 → R that is C 1 -smooth (in the first variable) is called ˆ L ( f ) -smooth if the following quadratic upper bound holds for any x, ˆ x, y, z ∈ R

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 inequality (58) holds with constant ˆ L . Suppose x ⋆ ∈ arg min x { f ( x, y, z ) } exists. An immediate consequence of (58) is that

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 and x ⋆ ∈ arg min x { f ( x ) } ,

If 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 d s < ∞ with γ s as in Table 2 (AdaGrad-Norm, general formula), then

Proof. Given the Eq. (87) for the distance to optimality, with ( x, x ⋆ ) ∼ N (0 , B ) ,

By the RSI (with constant ˆ µ ( f ) ) condition on f , we have that

where x = ⟨ X,a ⟩ and x ⋆ = ⟨ X ⋆ , a ⟩ and we note that x has t -dependence due to the t -dependence in B . By ˆ L ( f ) -smoothness,

This implies that

Thus by (60) and (61), we have that

Which then yield:

Changing variables u = Γ( t ) = ∫ t 0 γ s d s , we have that ∫ ∞ 0 R ( t ) γ t d t = ∫ ∞ 0 r ( u ) d u = ∥ r ∥ 1 . Rearranging the term in the above equation and taking t → ∞ . We obtain: ∥ r ∥ 1 ≤ D 2 (0) ( 2ˆ µ ( f ) -( ˆ L ( f )) 2 1 d Tr( K ) γ 0 ) , given that 2 ˆ µ ( f ) ( ˆ L ( f )) 2 1 d Tr( K ) > γ 0 . Using Lemma D.1, with i ( v ) = I ( B (Γ -1 ( v ))) = E a,ϵ [ ( f ′ ( x, x ⋆ ) 2 ] instead of the risk

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 d s , then g ( u ) = η 2 b η + 1 2 d Tr( K ) ∫ u 0 i ( v ) d v 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:

such that i ( v ) = I ( B (Γ -1 ( v ))) . Taking derivative with respect to u , rearranging terms and integrating leads to the desired result.

D.2 Least squares setting

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,

where Γ( t ) := ∫ t 0 γ s d s with,

In 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).

D.2.1 Proofs for case of fixed d

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:

Taking t →∞ , we get

Using ∥K∥ 1 = ∫ ∞ 0 K ( x ) d x < γ -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 2 d Tr( K ) ,

On the hand following Lemma D.3, 1 4 D 2 (0)(1 + γ 0 2 d 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 obtain

Taking u →∞ , and using ∥ r ∥ 1 ≍ 1 4 D 2 (0) ,

This 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 ) d s :

For K = I d , this equation simplifies,

Plugging 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.

D.2.2 Vanishingly few eigenvalues near 0 as d →∞

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 ,

Next, 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 risk R ( t ) = 1 2 d ∑ d i =1 λ i D 2 i ( t ) ,

with R 2 ( s ) = 1 2 d ∑ i : λ i ≤ α s λ i D 2 i ( s ) . Next, choosing α s = γ s 1 d Tr( K 2 ) , we show that the last term is of order o d (1) . By Lemma D.2 ∀ i , D 2 i ( t ) ≤ max ( γ t 1 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,

where N 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 ) d s = o d (1) , provided that d is taken to be large before t .

We then have that up to o d (1) constant,

Using Gronwall's inequality,

where 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 that

for all t > 0 . Finally, we plug this into the formula for γ t and conclude that, for all t > 0 ,

Lemma 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( γ t 1 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), d d t D 2 i (0) ≥ 0 . However, since D 2 i ( t ) , R ( t ) are continuous, this equation implies that D 2 i ( t ) ≤ γ t R ( t ) for all t and thus D 2 i ( t ) ≤ γ t 1 R ( t 1 ) for all t .

Case 2: Suppose that D 2 i (0) > γ 0 R (0) . Then, by equation (8), d d t D 2 i (0) < 0 . If d d t D 2 i ( t ) < 0 for all t , then D 2 i ( t ) ≤ D 2 i (0) for all t . If at some point d d t D 2 i ( t ) > 0 , this implies D 2 i ( t ) ≤ γ t R ( t ) and we are in Case 1.

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.

D.2.3 Power law behavior at d →∞

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 d s and let

where γ t , K are monotonically decreasing, with ∥K∥ 1 < ∞ . Then all t ,

If in addition, there exist ϵ > 0 and T > 0 such that, for all t > T ,

then for all t for

Proof. The lower bound holds trivially, using R ( s ) ≥ F (Γ( s )) . For the upper bound, we start with the following change of variables:

with g ( u ) = γ Γ -1 ( u ) . Let us define the convolution map

Next we show that this map is contracting and in particular,

where 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,

Therefore, expanding the loss and using the above upper bound, and denote by q = 2(1 + ε ) ∥K∥ 1 γ 0 such that q < 1 ,

where 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 ,

Then,

Proof. Define α n = sup t ≥ 0 K ∗ n ( t ) K ( t )(2 ∥K∥ 1 ) n -1 , trivially α 1 = 1 . Consider the n +1 convolution,

By the assumption of the Lemma, we know that there exists some T > 0 such that for ∀ t ≥ T

Therefore, if t ≥ T , we have

On the other hand, if t < T ,

Taking supremum in Eq. (82), and combining the results of Eq. (85), and Eq. (84), we obtain that,

Solving the above recursion equation,

rearranging the terms we arrived at the required result.

Asymptotic analysis of the risk Here, we consider a family of models with d →∞ , for which the following power law asymptotics assumption is satisfied:

Corollary D.1 apply Lemma D.3 in the setting for which F , and K has a power law behavior asymptotically. It shows that the risk will then be dominated by F only. Corollary D.2 shows the behavior of the learning rate in this setting. Finally, Lemma D.5 shows that Assumption 8 is a consequence of a power law spectrum near zero on the eigenvalues of the covariance matrix and a power law assumption on the projected discrepancy at initialization.

Corollary D.1. Suppose Assumption 8 is satisfied, then R ( t ) ≍ F (Γ( t )) .

Proof. Define g ( u ) = γ Γ -1 ( u ) and r ( u ) = R (Γ -1 ( u )) and observe that g ( u ) is a decreasing function. Then, from the upper bound in Lemma D.3, we have

Combining this upper bound with the lower bound from Lemma D.3 and that κ 2 > 1 , we conclude that r ( u ) ≍ F ( u ) and R ( t ) ≍ F (Γ( t )) .

Next, we derive the asymptotics of γ t . There are three different cases, depending on whether the risk is integrable, which translates to a threshold with respect to the parameter κ 1 .

Corollary D.2. Suppose Assumption 8 then the following asymptotics for the learning rate hold:

For κ 1 > 1 , there exists ˜ γ such that γ t ≥ ˜ γ and R ( t ) ≍ t -κ 1 for all t ≥ 0 .

Proof. Using the notations g ( u ) and r ( u ) defined above along with the change of variable u = Γ( t ) , we get ∫ t 0 R ( s ) d s = ∫ u 0 r ( v ) g ( v ) d v . Combining this with Corollary D.1 and the formula for γ t we get

Let I ( u ) = b 2 + 2 d Tr ( K ) ∫ u 0 (1+ v ) -κ 1 g ( v ) d v and observe that g ( u ) ≍ 1 √ I ( u ) and I ′ ( u ) =

2 d Tr ( K ) (1+ u ) -κ 1 g ( u ) . Thus, I ( u ) satisfies I ′ ( u ) √ I ( u ) ≍ (1 + u ) -κ 1 so we have

In the case of κ 1 > 1 , this implies √ I ( u ) ≤ √ I (0) + C ∫ (1 + v ) -κ 1 d v . This upper bound on I ( u ) gives a corresponding lower bound on g ( u ) and thus a lower bound on γ t .

In the case of κ 1 < 1 , we have √ I ( u ) -√ I (0) ≍ (1 + v ) 1 -κ 1 so, for u sufficiently large, g ( u ) ≍ (1 + u ) κ 1 -1 . To recover the asymptotic for γ t , we observe that d d u Γ -1 ( u ) = 1 g ( u ) ≍ (1 + u ) 1 -κ 1 . Integrating both sides and changing back to t variables, we get t ≍ (1 + Γ( t )) 2 -κ 1 (or equivalently

1 + Γ( t ) ≍ t 1 / (2 -κ 1 ) ). Finally, plugging this into the formula for γ t and applying Corollary D.1, we get

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)

Similarly for K ,

with γ ( s, z ) = ∫ z 0 x s -1 e -x d x 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),

by 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 ] .

E Polyak Stepsize

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 d D 2 ( t ) < 0 with the maximum decrease being min γ t d D 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))

101

SGD

103

()6YS

10-1

limit

theory

quotient x(t)

102

rate

100

10-3

入2 = 1/21

learning

入2 = 1/22

10-5

入2 = 1/23

100

10-7

入2 = 1/24

0

20

40

60

入2 = 1/25

0

20

40

60

0

20

40

60

time (t)

time (t)

time (t)

The distance to optimality threshold, ¯ γ D t , occurs precisely when d D 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

A greedy learning rate that maximizes the decrease at each iteration is simply given by g Polyak t ∈ arg min d D 2 ( t ) . This has a closed form and we call this Polyak stepsize 11 . Again translating this back to SGD, we have

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.

The quantities A ( B ) and I ( B ) in (88) and (89) only depend on the low-dimensional function f and thus do not carry any covariance K or d dependence. Moreover, under additional assumptions on the function such as (strong) convexity, we can bound from below A ( B ) /I ( B ) . Thus, in terms covariance K and d , the Polyak stepsize g Polyak k ≍ 1 Tr( K ) / ( d ) = 1 avg. eig of K .

In the case of least squares (see (7)), we get

The latter gives the best fixed learning rate for a noiseless target on a LS problem (as noted in [35, 43]).

F Line Search

F.1 General Line Search

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),

Therefore 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 ODE

From this, we get an immediate learning rate (stability) threshold for the risk, that is, ¯ g R k is the largest learning rate for which SGD is guaranteed to decrease at each iteration, i.e., when the deterministic equivalent of R satisfies d R < 0 or equivalently after translating relevant terms into SGD quantities

The greediest approach, which we call exact line search , would choose the learning rate such that γ line t ∈ arg min γ d R . In this case, we get

F.2 Line Search on least squares

In this section, we provide a proof of Proposition 3.1, but, we show more than this including the exact limiting value for γ t .

Proposition F.1. Consider the noiseless ( ω = 0 ) least squares problem (7) . Then the learning rate is always lower bounded by

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 . In this context, the exact limiting value of γ line t is given by

where x is the positive real root of the second-degree polynomial

This leads to

Proof. We establish the inequality

by observing

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

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 of K that |{ λ = λ i } d i =1 | = d 2 . Thus, we establish the following system of ODEs

where R ( t ) = 1 2 d ( λ 1 D λ 1 ( t ) + λ 2 D λ 2 ( t )) and γ line t = 2 ( λ 2 1 D λ 1 ( t )+ λ 2 2 D λ 2 ( t ) ) ( λ 1 D λ 1 ( t )+ λ 2 D λ 2 ( t ) )( λ 2 1 + λ 2 2 ) .

Since D λ 2 ( t ) ≥ 0 and λ 1 > λ 2 > 0 , we infer that R ( t ) = 1 2 d ( λ 1 D λ 1 ( t ) + λ 2 D λ 2 ( t )) ≥ 1 2 d λ 1 D λ 1 ( t ) ≥ 0 . The structure of the exact line search algorithm ensures lim t →∞ R ( t ) = 0 , hence lim t →∞ D λ 1 ( t ) = 0 . Similarly, we deduce lim t →∞ D λ 2 ( t ) = 0 .

By applying L'Hôpital's rule and substituting the expressions for γ line t and R ( t ) in terms of D λ 1 ( t ) and D λ 2 ( t ) , we derive

Therefore, lim t →∞ D λ 2 ( t ) D λ ( t )

Solving for x > 0 , we derive the explicit formula

Given

we have

By substituting (98), we get

A direct calculation reveals that

λ

> λ

>

0

implies lim

γ

line

2

λ

1

d

t

min

Tr

(

(

K

K

)

)

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 by

where x 1 , . . . , x n -1 > 0 satisfy a more intricate coupled system of n -1 equations.

2

.

1

2

t

→∞

G Examples

Any single index model with α -pseudo Lipschitz ( α ≤ 1 ) activation function is covered by our SGD+AL theory. In this section, we provide key learning problems within this family of models.

G.1 Binary logistic regression

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,

To study the ODE dynamics of SGD in Eq. (9) one needs the deterministic risk h ( B ) , and I ( B ) = E a [ f ′ ( ⟨ X,a ⟩ , ⟨ X ⋆ , a ⟩ ) 2 ] , with B = W T KW . Following the computation in Appendix D example D.4 in [15] we obtain that

where z, w ∼ N (0 , 1) . The I function can also be computed explicitly by solving the following Gaussian integral, where we define g ( x ) def = exp( x )

1+exp( y )

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].

G.2 CIFAR 5m

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., find

where f i is the i th row of F . The SGD we now consider is

where γ k is the usual AdaGrad-Norm stepsize, as in (15). Our empirical covariance matrix K (remembering that f i is a row vector) is then

We 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.

10 0

10 1

10 2

10 3

10

1

2×10 2

3×10 2

4×10 2

6×10 2

Empirical Risk

CIFAR AdaGrad-Norm Least Squares

n = 2048

n = 4096

n = 8192

n = 16384

SGD Iterations/d

H Numerical simulation details

Here we provide more details for the figures that appear in the main paper.

Figure 1: Concentration learning rate and risk for AdaGrad-Norm on a least squares problem with label noise ω = 1 (left) and on a logistic regression problem with no label noise (right). For logistic, see Section G. 30 runs of AdaGrad-Norm with parameters b = 1 and η = 1 for each d ; X ⋆ ∼ N (0 , I d /d ) , X 0 = 0 , and K = I d . 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.

Figure 2: Comparison for Exact Line Search and Polyak Stepsize on a noiseless least squares problem. The left plot illustrates the convergence of the risk function, while the right plot depicts the convergence of the quotient γ t / λ min ( K ) 1 d Tr( K 2 ) for Polyak stepsize and exact line search. Both ODE The covariance matrix K is generated such that the eigenvalues follow the expression λ i ( K ) =

theory and SGD results are presented, showing a close agreement between the two approaches. , the spectrum becomes more spread out, resulting in larger values of 1 d Tr( K 2 ) . 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.

Figure 3: Quantities effecting AdaGrad-Norm learning rate. (left): The effect of adding noise to the targets ( ω = 1 . 0 ) to the risk (left axis) and learning rate (right axis). Ran AdaGradNorm( b = 1 . 0 , η = 2 . 5 ) on least squares problem with d = 500 . X 0 , X ⋆ ∼ N (0 , I d /d ) . A single run of the SGD (solid line purple) matches exactly the prediction (ODE, teal). The shaded region represents 10 runs of SGD with 90% confidence interval. The learning rate decays at the exact predicted rate of η √ b 2 + Tr Kω 2 d t . 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

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 .

Figure 4: Power law covariance in AdaGrad Norm on a least squares problem. Generated covariance K such that the density of eigenvalues are (1 -β ) λ -β where β = 0 . 2 and set X 0 = 0 . Choose ( X ⋆ i ) d i =1 = ( λ -δ/ 2 i ) d i =1 where λ i is the i -th eigenvalue of K and we vary δ ∈ (0 , 1 . 8) so that 0 < δ + β ≤ 2 . Setting of Prop. 4.4.

Figure 5: Convergence in Exact Line Search on a noiseless least squares problem. The plot on the left illustrates the convergence of the risk function, while the center and right plots depict the convergence of the quotient D λ 2 ( t ) 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.

Figure 6 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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction introduce the class of problems we consider in this work. We do not exaggerate our claims and state precisely the limitations of our work, namely 'high-dimensional linear composite setting' with normally distributed data in Section 1.1.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We clearly state all the assumptions on the model and algorithm in Section 1.1. When the theory switches to the least squares problem, we indicate clearly in the theorem statements. For example, see AdaGrad (Section 4) and Line Search (Section 3). We are very careful in the statements of our theorems/propositions/lemmas/corollaries to include all the assumptions needed to prove the result.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The paper contains 2 sections detailing the exact set-up of the problem class (Section 1.1) as well as the algorithmic set-up (Section 1.2). All the proofs for the main results are provided in the appendix and we are careful to include all the necessary assumptions to prove these results.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Either in the captions or in Section H, we state all the parameters, covariances of the data used, initialization, etc that are needed to reproduce the results. We also intend to release the code and make it available via github. We also remark that the numerical simulations are relatively simple to reproduce as they are generated using synthetic data and the algorithms are already in use.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: In the captions of the figures and/or in Section H, we have provided all the details to reproduce the experiments. The data is synthetic. We also intend to release the code.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Either in the captions of the figures or in Section H, we state all the parameters, covariance of the data, initialization, etc that are needed to reproduce the results. We also intend to release the code and make it available via github. We remark that the numerical simulations are relatively simple to reproduce as they are generated using synthetic data and the algorithms are already in use.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: In the figures (especially Figure 1), we provided statistical error bars (shaded region). Part of our theory is that the training dynamics of SGD with adaptive learning rate in high-dimensions concentrate. We then predict these training dynamics exactly. In this sense, we have used statistical tests to illustrate our theory.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The computing resources are minimal in our case as the models are just least squares and logistic regression with synthetic data. We explicitly state all the relevant details for another person to run the experiments. We remark that due to the simplicity of our simulations, no experiment required extensive compute capabilities.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed the code of ethics. We have made our utmost attempt to adhere to the guidelines provided by NeurIPS. We do not use any human subjects nor any datasets (privacy is not a concern). We did our best to cite all the relevent related work. Given that our work is in the foundational research and, in addition, the theory side, it is difficult to mitigate all the risks as the downstream effects of theory are long. The model is completely synthetic using the standard SGD algorithm; thus we don't, to the best of our knowledge, anticipate any risks.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: We do not anticipate any negative societal impact of the work. The work is a purely theoretical result on foundational research and it is not tied to any particular application. The work uses only synthetic data, standard algorithms (e.g., SGD, AdaGradNorm) that have already been used in ML/AI, and models such as least squares and logistics that are textbook. We have included a paragraph at the beginning of our appendix justifying this answer to the broader impacts of our work.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks. The work is purely theoretical and uses only synthetic data generated from a normal distribution. The models employed are standard statistical model (e.g., least squares and logistic regression) which are textbook problems.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We acknowledge via citations that the algorithms we study have been introduced before by others. We are not using any datasets or other assets beyond numpy and thus do not need to name any license or cite any dataset.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We do not intend to release any new assets. Neither the model nor the algorithms are new. We are simply analyzing known algorithms and models in this paper. We do intend to release code for reproducibility. In particular, we want readers to be able to generate the deterministic ODEs.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. The work is purely theoretical on a simple model.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The work does not involve crowdsourcing nor research with human subjects. The data used in this work is generated synthetically.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe motivation in lines 17-31 is very well stated. As a reader, I would be inclined to continue parsing the paper. \n- Proposition 4.3 states that with a null initialization $X_0 = 0$ and subject to a condition on the covariance (which roughly says that $\\lambda_{\\min}$ can be small but that there is a very small number of vanishing eigenvalues), if the ground truth signal is spread out, then the learning rate is never zero. For symmetric activations (say, phase retrieval), and constant learning rate, I think you cannot really start at zero because you have a symmetry, you need at least a tiny bit of overlap. From this result, it appears that you can start from the most useless guess and still get away with it, provably. Interesting. \n- The higher arching aim of the work is very important, as knowing that the behavior of SGD is predictable is pivotal for understanding how our models work. \n- The techniques are of independent interest and are a continuation of a line of work that is systematically, step by step, developing a global answer to these questions. \n- The result that Polyak's step-size converges to the inverse of the average eigenvalue is yet another validation of a method originated from convex optimization holding in more generality. This is reassuring.\n\nWEAKNESSES:\nWeaknesses 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\".  \n\n- 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. \n- 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. \n- 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. \n- 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.\n1. 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. \n2. Can you clarify why it makes sense to choose your scaling? Apart from it allowing the proofs. \n3. 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? \n4. 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?\n\n\n#### Typos\n__NOTE__: I am including for completeness typos in \"weaknesses\". Please do not count them as such. \n- (line 560) In the reminder of $\\Omega$, you are missing a $\\max$. \n- (equation 20) there is no $\\Gamma(t)$ in the expression, right? \n- (lines 1140-1141) the way of quoting \"noise\" and \"variance\" is not the correct one, TeX-wise.\n\nQUESTIONS:\nPlease note that some of the questions below are very ingenuous. \n\n- (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\". \n- 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. \n- 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. \n- 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). \n- 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. \n- 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? \n- In Figure 1, how many runs do you make? It would be nice to know.\n- The same question above but for any other plot where you show error bars. \n- (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)$. \n- Proposition 4.4 is nice. What happens when $\\beta +\\delta > 2$?\n\nMy soundness score is due to the many questions. I hope I will raise it after the discussion. The overall score follows the same principle.\n\nLIMITATIONS:\nNone.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-03"}

comments: SUMMARY:
Analyzes the behavior of AdaGrad-Norm, Line-search and Polyak step size using ODEs.

STRENGTHS:
This paper provides an interesting characterization of the behavior of several optimization methods using ODE tools, which are not as heavily used in optimization theory as they should be. This work seems technically sound, although I didn't review the (long) Appendix. 

I am in favor of acceptance, although with reservations. The methods studied are interesting to theoreticians, but they are rarely if ever used in practice, which limits the practical appeal of this work. Especially the idealized special cases for the line search and Polyak step size. The focus on quadratic problems is also a major limitation, especially since the analyzed methods are not typically used on quadratic problems.

There is definitely potential for followup work to build on this work.

WEAKNESSES:
- 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:
See above.

LIMITATIONS:
See above
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n### Update after rebuttal\nAfter reading the feedback carefully, I updated my score for the paper.\n\n### Original review\n\nIn 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.\n\nI 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.\n\nSTRENGTHS:\n1. The theory reveals that line search leads to slower convergence that Polyak stepsize.\n2. In some simplified cases, there exist exact expressions that describe the dynamic of the studied methods.\n3. The theory is also supported by numerical simulations.\n\nWEAKNESSES:\n1. The theory for line search and Polyak stepsize is only for the idealized version.\n2. 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\"\n3. 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.\n\n### Minor issues\nAbbreviation \"SGD\" is used before (line 23) it is introduced (line 30)  \n\"We fully expect\"  -> \"We expect\"  \n\"worst case convergence guarantees\" -> \"worst-case convergence guarantees\"  \n\"strongly-convex\" -> \"strongly convex\"\n\nQUESTIONS:\n1. 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. \n2. Why do you say in Appendix B that $R(X)$ is \"an expectation of a Gaussian vector\"? \n3. The authors wrote \"We shall remove this condition by approximation in what follows\", where is it done?\n4. The authors say they need \"mild assumptions on the learning rate\", but why are they mild?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nThis 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.\n\nSTRENGTHS:\nThis paper provides a fine-grained analysis of practical optimization algorithms, and also proves a general theorem on convergence to deterministic dynamics that I believe could be used to analyze dynamics for different adaptive-learning-rate algorithms in the future. There is also a strong match between theory and experiments. I believe it is an important step in the direction of understanding these practical optimization algorithms.\n\nWEAKNESSES:\nIn 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.\n\nThe data is also assumed to be Gaussian, although anisotropy is allowed.\n\nQUESTIONS:\nI am satisfied with the presentation in the paper, and do not have questions.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nFirst, 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\n- Plolyak stepsize. Obtaining learning rate. \n- Exact line search. Obtaining learning rate for two cases of data covariance $K$: isotropic and with two distinct eigenvalues\n- 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$.\n\nThe theoretical results are validated with numerical experiments on synthetic data.\n\nSTRENGTHS:\n- The main result (convergence to deterministic continuous-time ODE) is valid for general loss functions. The precise characterization of the optimization trajectory, as opposed to upper/lower bounds, is typically harder to obtain for non-quadratic loss functions\n- For AdaGrad Norm algorithm, the authors characterized the dynamics of risk and learning rate for several distinct types of covariance matrix spectrum. \n\nHowever, both these points are achieved only partially (see *weaknesses* for details). Completing them would significantly improve the paper.\n\nWEAKNESSES:\nThe 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: \n  - Five pages are reserved for introduction and setup. \n  - 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\n  - 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.\n  - Any kind of concluding section is absent\n\nRegarding 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. \n\n**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.  \n- 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. \n- 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.\n\n**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.\n\nQUESTIONS:\n**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 \n- rank 1 with non-trivial eigenvalue $\\lambda=1$ (extreme case).\n- with power-law spectrum $\\lambda_i=i^{-\\alpha}, \\alpha>1$  (e.g. classical *capacity* and *source* conditions).\n\nIndeed, 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. \n- Could the techniques developed in this paper handle such problems?\n- 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.\n\n**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))? \n\n**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$?\n\nLIMITATIONS:\nThe paper surely doesn't have any negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-30"}