4+3 Phases of Compute-Optimal Neural Scaling Laws

Introduction

The advent of large language models (LLMs) has changed our perceptions of the landscape of optimization and is resulting in the emergence of new interesting questions related to scaling. Prior to LLMs and other large models, we often viewed the large-scale optimization problems as being limited by the amount of data. In training language models, in contrast, data can be effectively infinite. Thus, compute budgets can be the limitation. This leads to the following natural question: given an architecture, given a fixed compute budget, and having unlimited data, how should one select the model size to minimize loss?

To formally address this question, let us consider the general learning problem,

the number of parameters $d$ is large, and the data vector $x$ is drawn from an unknown distribution. We solve (1) using stochastic algorithms, such as stochastic gradient descent (SGD) with batch size $B$, under various parameter sizes $d$, that produce a sequence of iterates $\{\theta_{r}\}$. A standard formula used in practice to measure compute is the "6ND" formula , that is,

Therefore, we can plot the loss curve $\mathscr{P}(\theta_{r};d)=\mathscr{P}(r;d)=\mathscr{P}({\mathfrak{f}}/(d\cdot B);d)$ as a function of flops (see Fig. 1). The question now is: given a fixed number of flops ${\mathfrak{f}}$ and given batch size $B$, how should we choose the parameters $d$ so that we get the best loss, i.e. $d^{\star}$ solves the constrained problem

In this work, we analyze a three parameter simple model, which we call power-law random features (PLRF) . The three parameters in the PLRF are the data complexity ($\alpha$), target complexity ($\beta$) and model-parameter count $d$. Using this model, we derive a deterministic equivalent for the expected loss, as a function of $\alpha$, $\beta$, and $d$, that captures the training dynamics of one-pass SGD. This can be used to derive numerical predictions for the scaling laws. We also extract exact expressions for the compute-optimal scaling laws and the optimal parameter $d^{\star}({\mathfrak{f}})\in\text{arg\,min}_{d}\mathscr{P}(\tfrac{{\mathfrak{f}}}{d\cdot B};d)$ for largeWe discuss how large is large, but the truth is somewhat complicated and also quite dependent on the desired precision. If $\pm 0.05$ on the achieved scaling laws is tolerable, a flat $d>1000$ seems to suffice across all phases. $d$, and give some estimates on the order of $d$ necessary for these scaling laws to take hold.

We also observe for a large portion of the $(\alpha,\beta)$-phase plane, the optimal parameter is $d^{\star}={\mathfrak{f}}^{1/2}$, suggesting a regime of universal scaling behavior (see Fig. 3(b) and Table LABEL:table:phases_intro).

The PLRF is not only analyzable, but also exhibits a rich behavior of compute-optimal curves/loss curves, which are qualitatively and quantitatively different depending on the strengths of the data $(\alpha)$ vs. target $(\beta)$ complexity. Particularly, we show that there are 4 distinct (+3 sub phases) compute-optimal curve/loss curve behaviors.

Model constrained compute-optimal curves. In two of the phases (Phase Ia,b,c and Phase II), it is the underlying model that dictates the curves. The algorithm has little/no impact. This appears in two forms. The first behavior are compute-optimal curves controlled by the capacity of the model (Phase Ia,b,c). Here once the algorithm reaches the limiting risk value possible (capacity), it is better to increase the model-parameter $d$. Another type of loss dynamics is due to poor model feature embedding (Phase II). Here the features are embedded in a way which is difficult to train. After an initial large decrease in the loss value, this feature embedding distortion frustrates the algorithm and training slows, but it continues to solve. However, solving to capacity wastes compute, in that it is compute-favored to increase the model parameter count $d$.

Algorithm constrained compute-optimal curves. For some choices of $(\alpha,\beta)$ (Phase III and IV), it is the noise produced by the SGD algorithm that ultimately controls the tradeoff. Here the algorithm matters. Indeed, another algorithm could change the compute-optimal curves for these phases.

The key source of inspiration for this work are , which identified compute optimality as a fundamental concept in scaling large language models and made a substantial empirical exploration of it. The problem setup was formulated by , where additionally data-limited scalings were considered, but compute optimality was not (nor indeed any algorithmic considerations); see also where gradient flow was considered in the same setting.

There is a substantial body of work considering scaling laws of losses (trained to minimum-loss) of dataset size vs parameter count, in a variety of settings (linear, random features, deep networks). See especially: , where in more complex models a “hidden-manifold” model is often adopted for the data. We note that as we consider one-pass SGD, some dataset/parameter-count scaling laws are implicit from the results here; however, the training method (one-pass SGD) is, in some regimes, suboptimal given unlimited compute.

1 Problem setup: SGD on Power-law Random Features

In this work, we analyze the three parameter power-law random features (PLRF) model, that is,

The dimensions we consider throughout are always such that $v\geq Cd$ for $C>1$. Throughout both $v$ and $d$ need to be large, but for some choices of $\alpha$ and $\beta$, the $v$ will need to be comparable to $d$.

We assume that $v\geq Cd$ with $C>1$ and $v,d\to\infty$. Above the high-dimensional line, which is when $2\alpha>1$, we suppose $v/d\to r\in(1,\infty)\cup\{\infty\}$.In fact, we may take $v=\infty$ for $2\alpha>1$. On the other hand, below the high-dimensional line $(2\alpha<1)$ we limit $v$ to be $v/d\to r\in(1,\infty)$.Indeed one can, in the former case, take $d\leq v\leq d^{1/(1-2\alpha)}$, but for simplicity of presentation we focus on the proportional regime when $2\alpha<1$.

One can rewrite the expression in (4) using the convenient form:

To solve the minimization problem in (5), we use one-pass SGD with mini-batches of size $B$ (independent of $d$)One can study batch size $B$ growing with $d$, but for $2\alpha>1$, $B$ must be indep. of $d$ (see Prop. 2.1). Thus we only consider $B$ independent of $d$ setting. and constant learning rate $\gamma>0$: letting $\theta_{0}=0$, we iterate

The learning rate and batch size will need to satisfy a condition to ensure convergence (Prop. 2.1).

Under this setup, our main goal is to characterize the compute-optimal frontier. Precisely, we want to find the parameter count exponent $\xi$ and scaling law exponent $\eta$, such that,

We use $\mathscr{P}(\theta_{r})=\mathscr{P}(r)$ when we want to emphasize the iteration counter $r$. We say $\mathscr{A}(r,v,d)\sim\mathscr{A}(r,v,d)$ for functions $\mathscr{A}(r,v,d),\mathscr{A}(r,v,d)>0$ if for every $\varepsilon>0$ and for all admissible $v$ and $d$, there exists an $r_{0},d_{0}$ such that for all $d>d_{0}$ and $r\geq r_{0}$

We write $\asymp$ if the upper and lower bounds hold with some constants $c,C$ in place of $1\mp\varepsilon$ respectively and $\lesssim,\gtrsim$ if only one inequality holds.

[notespar, caption = Large $d$ behavior of the forcing function and kernel function. See Sec. 10 for proofs. ,label = table:forcing function, captionskip=2ex, pos =!t ]l Function ${}^{*}\Gamma(x)$ is the Gamma function $\mathscr{F}_{0}(r)\asymp d^{-2\alpha+\max\{0,1-2\beta\}}$ \mathscr{F}_{pp}(r)\sim(2\alpha)^{-1}\times\Gamma\big{(}\tfrac{\beta}{\alpha}-\tfrac{1}{2\alpha}+1\big{)}\times(2\gamma B\times r)^{-(1+\beta/\alpha)+1/(2\alpha)} \mathscr{F}_{ac}(r)\leq\begin{cases}C\times\mathscr{F}_{0}(r),&\text{if2\beta>1$,$2\alpha<1}\\ 0,&\text{if2\beta<1}\end{cases} for $C>0$, independent of $d$ If $2\beta>1,2\alpha>1$, \mathscr{F}_{ac}(r)\sim\big{(}\sum_{j=1}^{V}j^{-2\beta}\big{)}(2\alpha)^{-1}\Gamma\big{(}1-\tfrac{1}{2\alpha}\big{)}\times(2\gamma B\times r)^{-1+1/(2\alpha)}\times d^{-1} \mathscr{K}_{pp}(r)\sim(2\alpha)^{-1}\times\Gamma\big{(}2-\tfrac{1}{2\alpha}\big{)}\times(2\gamma B\times r)^{-2+1/(2\alpha)}

Learning dynamics of SGD

Compute-optimal curves (3) for the random features model (4) rely on accurate predictions for the learning trajectory of SGD. Similar to the works of , we show that the expected loss under SGD satisfies a convolution-type Volterra equation (for background on Volterra equations, see Section 5.3)

Intuitively, the forcing function is gradient descent on the random features model and the kernel function is the excess risk due to 1 unit of SGD noise.

The forcing function $\mathscr{F}(r)$ and kernel function $\mathscr{K}(r)$ are random functions depending on the random matrix $\hat{K}$. Indeed, it is the resolvent of $\hat{K}$, $(\hat{K}-z)^{-1}$, which plays a significant role in $\mathscr{F}$ and $\mathscr{K}$. We remove this randomness from the expression by using a deterministic equivalent – a technique from random matrix theory.

Formally, we define the deterministic equivalent for the resolvent of $\hat{K}$, denoted by $\mathscr{R}(z)$, implicitly via a fixed point equation

The solution to the Volterra equation with deterministic equivalent (10) numerically exactly matches the training dynamics of SGD, see Fig. 2. A discussion of the deterministic equivalent for $(\hat{K}-z)^{-1}$ can be found in Sec. 7. All our mathematical analysis will be for the deterministic equivalents, going forward.There is good numerical evidence that the deterministic equivalent captures all interesting features of the PLRF. There is a vast random matrix theory literature on making precise comparisons between resolvents and their deterministic equivalents. It seems a custom analysis will be needed for this problem, given the relatively high precision required, and we do not attempt to resolve this mathematically here. The derivation of the Volterra equation for the expected loss can be found in Sec. 4.

An immediate consequence of (10) is that for convolution Volterra equations bounded solutions occur if and only if the forcing function is bounded and the kernel norm $\|\mathscr{K}\|\stackrel{{\scriptstyle\text{def}}}{{=}}\sum_{s=0}^{\infty}\mathscr{K}(s)<1$. This directly translates into a sufficient condition on the batch size and learning rate of SGD.

Suppose learning rate $\gamma$ and batch $B$ satisfy $\|\mathscr{K}\|<1\text{ and }\gamma B<1.$ Then $\mathscr{P}(r)$ is bounded.

Below the line $2\alpha=1$, the kernel norm diverges with $v$ for fixed constant $\gamma$, and so we must take $\gamma\to 0$ to ensure bounded solutions. Thus, provided $\gamma\sim v^{2\alpha-1}$, then

Thus, the kernel norm, $\|\mathscr{K}\|$, is always constant order for all $\alpha$.

For $2\alpha>1$, the restriction on $\|\mathscr{K}\|$ and $\gamma B$ imply that $\gamma$ and $B$ must be independent of $d$. For $2\alpha<1$, $B$ can grow with $d$. We only consider $B$ order 1 in this work. For a proof and necessary and sufficient conditions on $\gamma$ and $B$, see Prop. 5.2, and see Cor. 9.1 for the asymptotic on $\|\mathscr{K}\|$.

The Volterra equation in (10), while useful as it removes the randomness, is not sufficient to derive the compute-optimal curves (3). We need a more explicit formula for $\mathscr{P}$ (see Section 5.3.2 for proof).

Suppose $\gamma$ and $B$ are at most half the convergence threshold and $2\alpha+2\beta>1$, $\alpha>\tfrac{1}{4}$.In spite of Theorem 2.1 holding only for $\alpha>\tfrac{1}{4}$, we expect this to hold for all $2\alpha+2\beta>1$ as supported numerically. When $\alpha<\tfrac{1}{4}$, the kernel function stops being power law and, as a result, requires a different set of tools to prove the result. There exists an $M>0$ large and a constant $C=C(\alpha,\beta,M)$, independent of $d$, so that for all admissible $v$ and $d$, for all $\gamma Br>M$,

The convolution $\mathscr{K}*\mathscr{F}$ further simplifies

If we were to run gradient descent instead of SGD (i.e., $\gamma$ small), then we would only have the forcing term, that is, $\mathscr{P}(r)=\mathscr{F}(r).$ The measurable effect of SGD comes from the second term that contains the kernel function. For this reason, we refer to SGD noise as $\frac{1}{\gamma B}\cdot\mathscr{K}(r).$

In light of (13) and (14), we have trapped the training loss between the sum of $\mathscr{F}$ and $\mathscr{K}$, so it suffices now to understand the forcing and kernel functions.

1 Forcing function and kernel function

We decompose the forcing function (11), $\mathscr{F}$, and the kernel function, (12), $\mathscr{K}$, into

Each term is explicit and has an asymptotic equivalence (when $1\lesssim\gamma Br\lesssim d^{2\alpha}$) given by

The two error terms are such that for large $d$ with $1\lesssim\gamma Br\lesssim d^{2\alpha}$,

for some constant $C>0$. For $\gamma Br\gtrsim d^{2\alpha}$, the forcing function $\mathscr{F}(r)\asymp\mathscr{F}_{0}(r)$, the limiting risk value. The terms arise from different parts of the spectrum of the deterministic equivalent for $\hat{K}$ (see Fig. 6).

Point mass at : $\mathscr{F}_{0}(0)=\mathscr{F}_{0}(r)$ is the limiting value of $\mathscr{P}(r)\asymp d^{-2\alpha+\max\{0,1-2\beta\}}$ as $r\to\infty$. It occurs because the loss is irreducible $(d<v)$, that is a component of the target is not in the image of the RF model (or equivalently that $\hat{K}$ has a kernel).

Aligned features: The function $\mathscr{F}_{pp}(r)$ represents gradient descent on the components of features which are aligned to the underlying population features. Indeed, if we ran gradient descent on the population loss without a random features map (or a diagonal $W$), this would be the loss curve.

Distorted features: The function $\mathscr{F}_{ac}(r)$ is the result of feature distortion, where the matrix $W$ leads to an embedding where a small component of the leading features is distributed across many different eigenmodes. These are still solvable, and given enough compute these will eventually be used, but they are much slower to solve.

Aligned kernel: $\mathscr{K}_{pp}(r)$ is the excess risk due to $1$ unit of SGD noise, which is then solved according to population gradient descent.

Out of brevity, we relegate the exact definitions of $\mathscr{F}_{0}$, $\mathscr{F}_{pp}$, $\mathscr{F}_{ac}$, and $\mathscr{K}_{pp}$ and all proofs of the asymptotics in Table LABEL:table:forcing_function and analyses of the functions to Section 8, 9, and 10.

The 4 Phases

We now put together a coherent picture of the effect of different choices of $\alpha$ (data complexity) and $\beta$ (target complexity) and their impact on the compute-optimal frontier. By Theorem 2.1, we estimate

Fig. 4a. shows empirically that this equivalence of $\mathscr{P}(r)$ is quite good.

The 4 distinct phases (see Fig. 3(a)) decompose the $(\alpha,\beta)$-plane based on the shape of the loss curve $\mathscr{P}(r)$, that is, which of the distinct components of the forcing function (i.e., $\mathscr{F}_{0},\mathscr{F}_{pp},\mathscr{F}_{ac},$) and/or kernel function (i.e., $\mathscr{K}_{pp}$) dominate the loss curve at a given iteration $r$. See Table LABEL:table:phases_intro for loss description in each phase. Cartoon pictures of the different features of the loss curves are shown in Fig. 5. For each phase, we derive a compute-optimal curve in Section 3.1.

The high-dimensional line, which occurs where $2\alpha=1$, distinguishes the phases where the $v$-dimension can be big and independent of $d$ (Phase Ia, II, III, $2\alpha>1$) and the phases where $d$ and $v$ must be related to each other (Phase Ib, Ic, IVa, IVb, $2\alpha<1$). When $2\alpha+2\beta<1$, the loss does not exhibit any power-law decay as the limit level stops going to as $d\to\infty$ (purely as a consequence of having selected the regime $v>d$). Moreover, there exists an interesting critical point $\alpha=\beta=\tfrac{1}{2}$ where all the parts of the forcing function and kernel mix and interact with each other. The behavior of the loss at the pentuple point (see Fig 3(a)) we leave for future research. Across each of the phase boundaries the compute-optimal curves are continuous, but not necessarily differentiable; in contrast, $d^{\star}$ is discontinuous across some phase boundaries.

To simplify the computations for compute-optimal curves, we introduce the following curve

We describe the qualitative and quantitative properties of compute-optimal curves for each phase. These are broken down into model constrained (Phase I, II) vs. algorithm constrained (Phase III, IV), i.e., whether the PLRF model or SGD is the constraining feature.

Phase Ia ($2\alpha>1,2\beta<1)$, Ib ($2\alpha<1,2\beta<1,2(\alpha+\beta)>1$), Ic are characterized by having the simplest loss description, $\mathscr{P}(r)\asymp\mathscr{F}_{pp}(r)+\mathscr{F}_{0}(r)$. Here the SGD noise is irrelevant and one would have the same loss (and thus compute-optimal curve) as gradient descent on the population loss. Compute optimality is characterized by training the model completely (to its limit loss) and choosing the model parameter count large enough so that at the end of training, the smallest loss is attained. The main distinctions between Phase Ia, Ib, Ic are the model capacities (i.e., $\mathscr{F}_{0}(r,d)=d^{-2\alpha+1-2\beta}$ in Ia, Ib, and $\mathscr{F}_{0}(r,d)=d^{-2\alpha}$ in Ic) and the dependence of dimension in the learning rate due to Ib,Ic being below the high-dimensional line. Consequently, while the qualitative features of the loss curve are the same for Ia, Ib, and Ic, the actual values of the compute-optimal curve vary across the different regions. Notably, in Phase Ib, the compute-optimal parameter is $d^{\star}={\mathfrak{f}}^{1/2}$ and it is independent of $\alpha$ and $\beta$.

Phase II ($2\alpha>1$, $2\beta>1$, $\beta<\alpha$) has a loss curve where the $\mathscr{F}_{ac}$ is important, that is, $\mathscr{P}(r)\asymp\mathscr{F}_{pp}(r)+\mathscr{F}_{ac}(r)+\mathscr{F}_{0}(r)$. The $\mathscr{F}_{ac}$ term becomes the dominant term after running for some intermediate amount of time $d^{c}$; in fact it is compute-optimal to stop at this point, and then select the number of model parameters so to minimize the loss with this early stopping criterion. It transpires that across all phases, it never pays to solve through the $\mathscr{F}_{ac}$ part of the loss curve – it is always better to just increase the number of model parameters.

In this phase $(2\alpha>1,2\beta>1,\beta>\alpha)$, SGD noise is important. The loss curve is $\mathscr{P}(r)\asymp\mathscr{F}_{ac}(r)+\mathscr{F}_{0}(r)+\frac{1}{\gamma B}\mathscr{K}_{pp}(r)$. Notably, in this phase, the compute-optimal parameter is $d^{\star}({\mathfrak{f}})={\mathfrak{f}}^{1/2}$, which is independent of $\alpha$ and $\beta$. PLRF that fall within this phase have the same scaling law regardless of data complexity and target complexity. Moreover, the tradeoff occurs, like in Phase II, once the optimizer reaches the $\mathscr{F}_{ac}$-dominated part of the loss curve. Unlike in Phase II, the optimization is slowed by SGD noise ($\mathscr{K}_{pp}$) leading up to that point. We note that there is a dimension-independent burn-in period required for SGD noise to dominate, and for small numerical simulations, one may actually observe an $(\mathscr{F}_{pp},\mathscr{F}_{ac})$ tradeoff.

Like Phase III, SGD noise is important. The SGD algorithm in Phase IV will be distinguished from gradient descent. As one approaches the high-dimensional line ($2\alpha=1$) in Phase III, the $\mathscr{F}_{ac}(r)$ disappears. It becomes too small relative to $\mathscr{F}_{pp}$ and $\mathscr{K}_{pp}$. Moreover at the high-dimensional line, $\mathscr{F}_{pp}$ becomes important again. Thus, the loss curve in Phase IV (a and b) look like \mathscr{P}\big{(}r,d\big{)}\asymp\mathscr{F}_{pp}(r,d)+\mathscr{F}_{0}(r,d)+\tfrac{1}{\gamma B}\mathscr{K}_{pp}(r,d). The distinction between Phase IVa ($1-\tfrac{1}{\sqrt{2}}<\alpha<0.5,2\beta>1$) and Phase IVb ($\tfrac{1}{4}<\alpha<1-\tfrac{1}{\sqrt{2}},2\beta>1$) is where the compute-optimal tradeoff occurs. It changes from $\mathscr{K}_{pp}=\mathscr{F}_{0}$ (Phase IVa) to $\mathscr{F}_{pp}=\mathscr{K}_{pp}$ (Phase IVb). In particular it can be (Phase IVb) the SGD noise is so large that increasing the model parameter count is compute-optimal. We note that in this phase $d$ must be taken very large (in particular larger than we could numerically attain) to get quantitative agreement between the exponents and theory.

In Phase III, Ib, and IVa, the optimal parameter $d^{\star}={\mathfrak{f}}^{1/2}$ (see dashed lines in Fig. 3(b)). These phases, taken together, encompass a large section of the $(\alpha,\beta)$-phase plane. This suggests that there is a potential universal scaling law. Moreover using 1 GPU-day of compute, one reaches scales of $d$ where the observed exponents in the scaling laws – SGD, the theoretically-derived Volterra equation eq. (10), and the equivalence of $\mathscr{P}(r)$ eq. (17) – are still changing (see Fig. 4b and c). This serves as a potential warning for empirically derived scaling laws. Additionally, although we have identified the lower-left of the phase diagram ($\alpha+\beta<1/2$) as "no power-law", this designation relies on the assumption $v>d$, which could be relaxed to interesting effect in more realistic (e.g. non-linear) models.

In summary, we analyze a simple three parameter model, PLRF, and derive deterministic expressions for the training dynamics (see Volterra equation (10)). We then extract compute-optimal scaling laws for large $d$. We identify 4 phases (+3 subphases) in the $(\alpha,\beta)$-phase plane, corresponding to different compute-optimal curve/loss behaviors. These phase boundaries are determined by the relative importance of model capacity (Phase I, IV), poor embedding of the features (Phase II, III), and the noise produced by the SGD algorithm (Phase III, IV). The latter suggesting that another stochastic algorithm might change the compute-optimal curve; we leave this interesting direction to future research. We also show evidence of a universal scaling law which we also leave for future research to explore. Lastly, we did not fully explore the effects of batch and learning rate on the compute-optimal curves; but we note that only in Phase IV is there a potential gain to be had from increasing the batch. This is again an interesting direction to pursue.

The remainder of the article is structured as follows: in Section 4, we derive the convolution-type Volterra equation for the expected risk under SGD, (7). In Section 5, we analyze the Volterra equation under the deterministic equivalent. A discussion on the convergence threshold for $\mathscr{P}(r)$ including a necessary and sufficient condition for bounded solutions of (10) (Proposition 5.2) and a proof of Proposition 2.1 are provided in Section 5.2. Some background on Volterra equations and their solutions are provided in Section 5.3.1 followed by the proof of Theorem 2.1 in Section 5.3.2. We finish this section with a detailed description and proofs for the risk curves in all phases, Section 5.4. Section 6 is devoted to deriving and proving the compute-optimal curves in Table LABEL:table:phases_intro. We follow this by Section 7 which analyzes the deterministic equivalent for the resolvent of $\hat{K}$. Here we examine the spectrum of $\hat{K}$ from a random matrix point of view. In particular, in this section, we prove estimates on the fixed point equation, $m$, see eq. (9). We then give explicit descriptions of the components of the forcing function, $\mathscr{F}_{0},\mathscr{F}_{pp},\mathscr{F}_{ac}$, as contour integrals and show that the error terms $\text{error}_{\mathscr{F}}$ are small, see Section 8. We do the same with the kernel function $\mathscr{K}$ and kernel norm in Section 9. In Section 10, we derive the asymptotic formulas for the components of the forcing and kernel functions (see Table LABEL:table:forcing_function) used in the compute-optimal curve derivations. Finally, we end with some additional numerical experiments (and their experimental setups) as well as detailed descriptions of the different approaches to estimating the exponents in the scaling law and optimal model-parameter, Section 11.

Derivation of Volterra equation for expected risk under SGD

Here we recall that the $(v\times v)$-matrix $D\stackrel{{\scriptstyle\text{def}}}{{=}}\text{Diag}(j^{-2\alpha}\,:\,1\leq j\leq v)$. Using these moment calculations, we can compute explicitly each of the terms in (20).

First, we consider the gradient term in (20). A simple computation yields

We now turn to the quadratic term in (20). For this, we have the following

This, after taking expectations, is in the form for us to apply the moment computations in (21). Using these moments, we get the following expression:

Now we simplify the 2nd term in the summand

As a result, we deduce by combining (24), (25), and (26) with (23) gives the following representation for the expected quadratic term

We will write our Volterra equation in terms of $\rho_{j}$’s. Note we can express the loss $\mathscr{P}(\theta_{r})=\|D^{1/2}(W\theta_{r}-b)\|^{2}$ by

We can now plug $\rho_{j}^{2}$ into (30). For this, we need to compute $\nabla\rho_{j}^{2}$ and $\nabla^{2}\rho_{j}^{2}$:

Using an integrating factor, we can implicitly solve this expression

and thus, we have a discrete Volterra equation

Let us define $\check{K}\stackrel{{\scriptstyle\text{def}}}{{=}}W^{T}DW$. Using the expression in (32),

Analysis of Volterra equation under the deterministic equivalent

From now on, we consider the setting where the initialization of SGD is $\theta_{0}=0$. Let us introduce the forcing function:

and recall the kernel function $\mathscr{K}(s)$:

While these representations are easy to see from the derivation of the Volterra equation, a more useful representation of the forcing function and the kernel function is through contour integrals over the spectrum of $\hat{K}$. With this in mind, let $\Gamma$ be a contour containing $$. Note that by the assumptions on$\hat{K}$, the largest eigenvalue is normalized to be$1$; hence$\Gamma$contains the spectrum of$\hat{K}$. Then the forcing function takes the form

Then one can write the Volterra equation (37) as the forcing function plus a convolution with the kernel and the expected loss, i.e.,

The forcing functions $\mathscr{F}(r)$ and kernel function $\mathscr{K}(r)$ are random functions as they depend on the random matrix $W$. Moreover the expressions via contour integration show that both of these functions can be described in terms of the random matrix $\hat{K}=D^{1/2}WW^{T}D^{1/2}$. Indeed it is the resolvent of $\hat{K}$,

Formally, we define the deterministic equivalent for the resolvent $\mathscr{R}(\hat{K},z)$, denoted by $\mathscr{R}(z)$ implicitly via a fixed point equation

As mentioned early, this deterministic equivalent, $\mathscr{R}(z)$ can be viewed, roughly as,

though it is not formally the expectation over $W$.

Using this deterministic expression for the resolvent of $\hat{K}$, we defined deterministic expressions for the forcing function via the contour representation of $\mathscr{F}(r)$ in (39)

We note the similarity with the Volterra equation for SGD. We conjecture that the two processes are close: for $\{\theta_{r}\}$ the sequence of iterates generated by SGD with $\theta_{0}=0$ and any $\varepsilon>0$,

for all admissible $V$, $d$ with probability going to 1 as $d\to\infty$.

We leave this for future research and suspect it is true because of deterministic equivalence for random matrices and our numerical simulations.

2 Convergence threshold.

A natural question is: for what choices of batch $B$ and learning rate $\gamma$ does $\mathscr{P}$ converge? To answer this, we introduce an additional quantity, the kernel norm defined as

If $2\alpha>1$, then $v$ be taken to equal $\infty$, that is,

In all cases, we choose $\gamma$ so that the kernel norm is asymptotic to a strictly positive constant.

A well-known result about convolution-type Volterra such as (45) is that the solution of convolution-type Volterra equation is bounded if and only the forcing function $\mathscr{F}(r)$ is bounded and the kernel norm $\|\mathscr{K}\|<1$. This naturally leads to conditions for our specific forcing function and kernel function.

The forcing function $\mathscr{F}$ is bounded and the kernel norm $\|\mathscr{K}\|<1$ for (43) and (44), respectively, if and only if

The first term ensures that the forcing function of the Volterra equation in (45) goes to (i.e., bounded) and the second condition is the same kernel norm bound. Moreover, we can think of condition (i). as the same condition needed for gradient descent to converge while the kernel norm is the effect of noise from SGD.

We also note that in light of Proposition 5.1 the kernel norm does not involve the batch size $B$. Therefore the condition $\|\mathscr{K}\|<1$ only places a condition on the learning rate (see below).

We now state necessary/sufficient conditions on the batch size and learning rate (47).

The learning rate, $\gamma>0$ and batch size, $B>0$, satisfy

if and only if the solution $\mathscr{P}(r)$ to the convolution-type Volterra equation (10) is bounded.

From (47), we need that |1-2\gamma B\lambda_{j}+2\gamma^{2}B\lambda_{j}|<1,\,\,\text{for all\lambda_{j}\in}. For this, we consider two cases.

First, suppose that $1-2\gamma Bx+2\gamma^{2}Bx^{2}<1$ for all $x\in$. We have that

The roots are precisely $x=0$ and $x=\frac{1}{\gamma}$. If $1/\gamma>1$, then the inequality in (49) always holds. Therefore, we need that $\gamma<1$.

Now suppose $-1+2\gamma Bx-2\gamma^{2}Bx^{2}<1$. Then we have

The roots of the left-hand side are precisely

For $B=1,2,3$, we have complex roots and so (50) automatically holds.

Below the high-dimensional line, $2\alpha<1$, the kernel norm diverges with $v$ for fixed constant $\gamma$, and so we must take $\gamma\to 0$ to ensure bounded solutions. Furthermore, with $\gamma\to 0$ (at any rate depending on $V$) we have the asymptotic equivalence

For a proof of the asymptotic for $\|\mathscr{K}\|$, see Corollary 9.1.

Following the proof of Prop. 5.2 , we need \frac{1}{2\gamma}\big{(}1\pm\sqrt{1-\tfrac{4}{B}}\big{)}>1 and $\gamma<1$. For $B=1,2,3$, we have complex roots and so (50) automatically holds. For $B\geq 4$, we need that

and this will imply (50). Suppose $B=4$, then we get that $\frac{1}{2}>\gamma$. For $B>4$, Taylor expanding (51) at $x=\infty$

Hence in this case, we need $\gamma B<1$. ∎

Similar results as for Prop. 2.1 and 5.2 hold for the expected SGD loss (via the Volterra equation (41)) by replacing $\|\mathscr{K}\|$ with $\|\mathscr{K}\|$.

3 Simplification of the Volterra Equation

While convolution-type Volterra equation such as (45) are quite nice and well studied in the literature (e.g., ), we need an approximation of the solution to it to have better understanding of compute-optimal curves. In this section, we show that we can bound (above and below) $\mathscr{P}(r)$ by a constant multiple of the forcing function $\mathscr{F}$ and kernel function $\mathscr{K}$.

To do this, we need some background on general convolution-type Volterra equations of the form:

where $f(t)$ is a non-negative forcing function and $K(t)$ is a monotonically decreasing non-negative kernel function.

Let us define K^{*n}\stackrel{{\scriptstyle\text{def}}}{{=}}(\underbrace{K*K*\ldots*K*K}_{\text{ntimes}})(t), the $n$-fold convolution of $K$ where $K^{*1}=K(t)$.

Under mild assumptions such as $\|K\|=\sum_{t=0}^{\infty}K(t)<1$ and the forcing function $f$ is bounded, then there exists a unique (bounded) solution $P(t)$ to (52) and the solution is given by repeatedly convolving the forcing function with $K$ (see, e.g., [6, Theorem 3.5]),

This representation of the solution to (52) enables us to get good bounds on $P(t)$. First, we state and prove a lemma attributed to Kesten’s Lemma [1, Lemma IV.4.7].

Suppose the kernel function $K$ is positive and monotonically decreasing and $\|K\|<\infty$. Moreover suppose for some $\varepsilon>0$, there exists a $T(\varepsilon)>0$ such that

Then $a_{1}=1$, and we are trying to prove

By definition of the convolution, we have that

where the last equality follows by the equality $\|K^{*n}\|=\|K\|^{n}$, [6, Theorem 2.2(i)].

In conclusion, by monotonicity, we have that

If the assumption (53) holds only for $\hat{T}>t>T$, then the statement of Lemma 5.1 still holds with

We now give a non-asymptotic bound for the general convolution-type Volterra equation.

Let $K$ and $f$ be non-negative functions. Suppose $K$ is monotonically decreasing and for some $\varepsilon>0$, there exists a $T(\varepsilon)>0$ such that

Moreover, suppose the convergence threshold condition $2(1+\varepsilon)\|K\|<1$ holds. Then

where $C=\left(\frac{K(0)}{K(T)}+1\right)\left(\frac{1}{1-2\|K\|(1+\varepsilon)}\right)$.

We consider the upper and lower bound separately.

Lower bound: Since $K$ and $f$ is are non-negative, then $\sum_{j=1}^{\infty}(K^{*j}*f)(t)\geq(K^{*1}*f)(t)\geq(K*f)(t)$. Recall the solution to the convolution-type Volterra equation takes the form,

It immediately follows from $\sum_{j=1}^{\infty}(K^{*j}*f)(t)\geq(K*f)(t)$ the lower bound.

Upper bound: The solution to a Volterra equation (in $L^{1}$) is

By Lemma 5.1 and the hypothesis, there exists a $T>0$ and $\varepsilon>0$ such that

and $(2\|K\|(1+\varepsilon))^{j-1}<1$. Hence, we have that

3.2 Proof of Theorem 2.1.

We are now ready to show one of the main tools used to analyze the loss function, Theorem 2.1. The result relies on approximations for the kernel and forcing functions found in Section 8 and Section 9. We restate the theorem statement to remind the reader of the result.

Suppose $\gamma$ and $B$ is at most half the convergence threshold and $\alpha>\tfrac{1}{4}$. There exists an $M>0$ large and a constant $C=C(\alpha,\beta,M)$, independent of $d$, so that for all admissible $v$ and $d$, for all $M<\gamma Br$,

The convolution further simplifies. For any $\epsilon>0$, there exists an $M>0$ and a constant $C=C(\alpha,\beta,M)$ independent of $d$ so that for all $M<\gamma Br$,

Note for all $\gamma Br>1/Md^{2\alpha}$, we have that $c\mathscr{F}_{0}\leq\mathscr{F}(r),\mathscr{K}(r)\leq C\mathscr{F}_{0}(r)$ for some $C,c>0$. This is where the limiting level starts to dominate. We begin by showing (55). Fix $\varepsilon>0$. From Proposition 9.2, we have that there exists an $M>0$ sufficiently large so that the hypothesis for Kesten’s Lemma, i.e.,

where we used monotonicity of $\mathscr{F}$ and $\mathscr{K}$.

Using Proposition 10.2 and Proposition 10.4, for large $d^{2\alpha}/M\geq\gamma Br\geq M$, we have that $\mathscr{F}(\tfrac{r}{2})\asymp\mathscr{F}(r)$ since $\mathscr{F}$ is power law for large $r$ (see also Corollary 8.1). The same holds for $\mathscr{K}$, using Proposition 10.5 and Proposition 9.2, $\mathscr{K}(\tfrac{r}{2})\asymp\mathscr{K}(r)$ for $d^{2\alpha}/M\geq\gamma Br\geq M$ for some $M>0$.

For small $\gamma Br\leq M$, we have that $\mathscr{F}(r/2)\leq C$ and $\mathscr{K}(r/2)\leq C$ for some $C>0$. Since $\mathscr{F}$ and $\mathscr{K}$ are monotonic, we can choose a constant so that $\mathscr{F}(r/2)\lesssim\mathscr{F}(r)$ and $\mathscr{K}(r/2)\lesssim\mathscr{K}(r)$ for $\gamma Br\leq M$.

Now using Proposition 5.1 and Proposition 10.6, we have that

where we used monotonicity of $\mathscr{K}$ and $\mathscr{F}$.

We note that $\mathscr{F}(s)\asymp C$ for $\gamma Bs\leq M$ for all $M>0$. Therefore,

On the other hand, by Proposition 5.1, for any $\epsilon>0$, there is an $M$ so that for any $\gamma Br\geq M$,

4 Details of risk curves for the phases

We can now put together a coherent picture of the effect of different choices of $\alpha$ and $\beta$ and their impact on the Pareto frontier. We will have 4 distinct phases where the expected loss will exhibit a power law decay and 1 region $(\alpha+\beta\leq 0.5)$ for which the expected loss has no power law decay (see Figure 3(a)). We will describe each of the 4 power law phases below marked by their boundaries.

First, we recall the forcing function $\mathscr{F}$ and kernel function $\mathscr{K}$ introduced in Section 2.1.

The function $\mathscr{F}_{0}(r)$ is the component of the forcing function corresponding to the point mass at , $\mathscr{F}_{pp}(r)$ is the component of the forcing function corresponding to the pure point part of the spectrum, and lastly, the most complicated part of the spectrum, the forcing function corresponding to the distorted features. In particular, we will show in Section 8 the exact definitions of $\mathscr{F}_{0},\mathscr{F}_{pp},\mathscr{F}_{ac}$ and $|\text{error}_{\mathscr{F}}|$ are small, and, in Section 10, we derive asymptotic-like behaviors for these functions. See Table LABEL:table:forcing_function_appendix for definitions and asymptotics.

Similarly, the kernel function $\mathscr{K}$ is

Note here that the kernel function has a multiplication by the eigenvalue of $\hat{K}$ and so the point mass at will not contribute. In Section 9, we will give an explicit definition of $\mathscr{K}_{pp}$ and show error terms are small and, in Proposition 10.5, we give the asymptotic-like behavior of $\mathscr{K}_{pp}$.

Now we describe in detail the different risk curves that arise for the different phases.

4.1 Above the high-dimensional line (Phases Ia, II, III).

This setting is commonly known as the trace class. It is characterized by four components:

learning rate $\gamma$ can be picked independent of dimension;

loss curve does not self average, that is, the loss curve does not concentrate around a deterministic function;

$v\geq d$, but $v$ has no upper bound and so we can take $v\to\infty$;

batch, $B$, is constrained to be small (see Proposition 5.2).

When $2\alpha>1$, or the trace class phase, the loss will exhibit 3 different phases. We described these phases in detail below.

In this phase, it notable for three characteristics:

absolutely continuous part of the forcing function does not participate;

level at which SGD saturates is affected by $\beta$;

In this case, the loss curve is just a constant multiple of gradient flow. Hence, we have that

Suppose $2\beta<1$ and $2\alpha>1$. Suppose the learning rate $\gamma$ and batch $B>0$ satisfy at most half the convergence threshold in Proposition 5.2. Then there exists an $M>0$ large and constants $C=C(\alpha,\beta,M)$ and $c=c(\alpha,\beta,M)$, independent of $d$, so that for all admissible $v$ and $d$, for all $\gamma Br>M$

By Theorem 5.1, we know that it suffices to look at the forcing function $\mathscr{F}$ and kernel function $\mathscr{K}$. Moreover, in this regime, we have that $\gamma$ and $B$ are constant (see Proposition 5.2).

The rest of the argument relies on the bounds found in Proposition 10.2, ($\mathscr{F}_{pp})$, Proposition 10.4 ($\mathscr{F}_{ac}$), Proposition 10.3, ($\mathscr{F}_{0})$, and Proposition 10.5 ($\mathscr{K}_{pp}$).

For the forcing function, $\mathscr{F}_{ac}(r)=0$ as $2\beta<1$ (Proposition 10.4). Therefore the forcing function is composed of $\mathscr{F}_{pp}(r)$ and $\mathscr{F}_{0}(r)$.

First, we have that $(\gamma Br)^{-2+1/(2\alpha)}<(\gamma Br)^{-(1+\beta/\alpha)+1/(2\alpha)}$ as $\beta<\alpha$ in this phase. Thus, using Proposition 10.2 and Proposition 10.5, for $\gamma Br>M$, where $M$ is some constant, we have that $\frac{1}{\gamma B}\mathscr{K}_{pp}(r)\leq C\times\mathscr{F}_{pp}(r)$ for some $C>0$. Hence the result is shown. ∎

As a consequence of the argument above, we know that

absolutely continuous spectrum takes over for $r\in(M,d^{2\alpha}/M)$ for some $M$;

Suppose $2\beta>1$, $2\alpha>1$, and $\beta<\alpha$. Suppose the learning rate $\gamma$ and batch $B>0$ satisfy at most half the convergence threshold in Proposition 5.2. Then there exists an $M>0$ large and constants $C=C(\alpha,\beta,M)$ and $c=c(\alpha,\beta,M)$, independent of $d$, so that for all admissible $v$ and $d$, for all $\gamma Br>M$

By Theorem 5.1, we know that it suffices to look at the forcing function $\mathscr{F}$ and kernel function $\mathscr{K}$. Moreover, in this regime, we have that $\gamma$ and $B$ are constant (see Proposition 5.2).

The rest of the argument relies on the bounds found in Proposition 10.2, ($\mathscr{F}_{pp})$, Proposition 10.4 ($\mathscr{F}_{ac}$), Proposition 10.3, ($\mathscr{F}_{0})$, and Proposition 10.5 ($\mathscr{K}_{pp}$).

$\gamma Br\leq M_{0}$, for some $M_{0}$: First, we have that $\tfrac{1}{\gamma B}\mathscr{K}_{pp}(r)\leq C_{0}\times\mathscr{F}_{pp}(r)$ (Proposition 10.5) and $\mathscr{F}_{ac}(r)\leq C_{0}\times\mathscr{F}_{pp}(r)$ (Proposition 10.4) for some constant $C_{0}>0$. The constant $M_{0}$ is where the asymptotic of $\mathscr{F}_{pp}$ starts to apply.

$M_{0}\leq\gamma Br\leq M_{1}$, for some $M_{0}$ and for all $M_{1}>M_{0}$: We see that $(\gamma Br)^{-2+1/(2\alpha)}<(\gamma Br)^{-(1+\beta/\alpha)+1/(2\alpha)}$ as $\beta<\alpha$ in this phase. Thus, using Proposition 10.2 and Proposition 10.5, we have that $\frac{1}{\gamma B}\mathscr{K}_{pp}(r)\leq C_{1}\times\mathscr{F}_{pp}(r)$ for some $C_{1}>0$. A quick computation shows that $\mathscr{F}_{ac}(r)\leq\mathscr{F}_{pp}(r)$.

$M_{1}\leq\gamma Br\leq M_{2}d^{2\alpha}$, for any $M_{1}$ and some $M_{2}$: The $M_{2}$ is the smallest of the two endpoints for the asymptotics of $\mathscr{F}_{pp}$ and $\mathscr{F}_{ac}$. As in the previous regime, we have that $\tfrac{1}{\gamma B}\mathscr{K}_{pp}(r)\lesssim\mathscr{F}_{pp}(r)$. In this region, $\mathscr{F}_{ac}(r)\asymp d^{-1}(\gamma Br)^{-1+1/(2\alpha)}$ and $\mathscr{F}_{pp}(r)\asymp(\gamma Br)^{-(1+\beta/\alpha)+1/(2\alpha)}$. We see at $(\gamma Br)=d^{2\alpha}$ that $(\gamma Br)^{-\beta/\alpha}\leq(d^{2\alpha})^{-\beta/\alpha}=d^{-2\beta}\leq d^{-1}$ as $2\beta>1$. Therefore, at $\gamma Br=d^{2\alpha}$, $\mathscr{F}_{pp}(r)\lesssim\mathscr{F}_{ac}(r)$ and we started, i.e., when $r=M_{1}$ with $\mathscr{F}_{ac}(r)\lesssim\mathscr{F}_{pp}(r)$. Therefore, we must change in this regime to being $\mathscr{F}_{ac}$ dominate.

$M_{2}d^{2\alpha}\leq\gamma Br$ for all $M_{2}$: In this case, all terms are bounded above by $\mathscr{F}_{0}(r)$. ∎

As a consequence of the argument above, we know that

In this case, we see that SGD changes the dynamics over gradient flow. In particular,

absolutely continuous forcing function takes over for iterations $r\in(M,d^{2\alpha}/M)$ for some $M$;

Suppose $2\beta>1$, $2\alpha>1$, and $\beta>\alpha$. Suppose the learning rate $\gamma$ and batch $B>0$ satisfy at most half the convergence threshold in Proposition 5.2. Then there exists an $M>0$ large and constants $C=C(\alpha,\beta,M)$ and $c=c(\alpha,\beta,M)$, independent of $d$, so that for all admissible $v$ and $d$, for all $\gamma Br>M$

By Theorem 5.1, we know that it suffices to look at the forcing function $\mathscr{F}$ and kernel function $\mathscr{K}$. Moreover, in this regime, we have that $\gamma$ and $B$ are constant (see Proposition 5.2). The rest of the argument relies on the bounds found in Proposition 10.2, ($\mathscr{F}_{pp})$, Proposition 10.4 ($\mathscr{F}_{ac}$), Proposition 10.3, ($\mathscr{F}_{0})$, and Proposition 10.5 ($\mathscr{K}_{pp}$).

$\gamma Br\leq M_{0}$, for some $M_{0}$: First, we have that $\mathscr{F}_{pp}(r)\leq C_{0}\times\tfrac{1}{\gamma B}\mathscr{K}_{pp}(r)$ (Proposition 10.5) and $\mathscr{F}_{ac}(r)\leq C_{0}\times\tfrac{1}{\gamma B}\mathscr{K}_{pp}(r)$ (Proposition 10.4) for some constant $C_{0}>0$. The constant $M_{0}$ is where the asymptotic of $\mathscr{K}_{pp}$ starts to apply.

$M_{0}\leq\gamma Br\leq M_{1}$, for some $M_{0}$ and for all $M_{1}>M_{0}$: We see that $(\gamma Br)^{-2+1/(2\alpha)}>(\gamma Br)^{-(1+\beta/\alpha)+1/(2\alpha)}$ as $\beta>\alpha$ in this phase. Thus, using Proposition 10.2 and Proposition 10.5, we have that $\mathscr{F}_{pp}\leq C_{1}\times\frac{1}{\gamma B}\mathscr{K}_{pp}(r)$ for some $C_{1}>0$. A quick computation shows that $\mathscr{F}_{ac}(r)\leq\mathscr{K}_{pp}(r)$.

$M_{1}\leq\gamma Br\leq M_{2}d^{2\alpha}$, for any $M_{1}$ and some $M_{2}$: The $M_{2}$ is the smallest of the two endpoints for the asymptotics of $\mathscr{K}_{pp}$ and $\mathscr{F}_{ac}$. As in the previous regime, we have that $\tfrac{1}{\gamma B}\mathscr{K}_{pp}(r)\lesssim\mathscr{F}_{pp}(r)$. In this region, $\mathscr{F}_{ac}(r)\asymp d^{-1}r^{-1+1/(2\alpha)}$ and $\mathscr{K}_{pp}(r)\asymp r^{-2+1/(2\alpha)}$. We see at $\gamma Br=d^{2\alpha}$ that $(\gamma Br)^{-1}\leq(d^{2\alpha})^{-1}=d^{-2\alpha}\leq d^{-1}$ as $2\alpha>1$. Therefore, at $(\gamma Br)\asymp d^{2\alpha}$, $\mathscr{K}_{pp}(r)\lesssim\mathscr{F}_{ac}(r)$ and we started, i.e., when $r=M_{1}$ with $\mathscr{F}_{ac}(r)\lesssim\mathscr{K}_{pp}(r)$. Therefore, we must change in this regime to being $\mathscr{F}_{ac}$ dominate.

$M_{2}d^{2\alpha}\leq r\gamma B$ for all $M_{2}$: In this case, all terms are bounded above by $\mathscr{F}_{0}(r)$. ∎

As a consequence of the argument above, we know that

4.2 Below the high-dimensional line (Phases IVa, IVb, Ib, Ic).

One of the main differences between the previous regime and this regime is that $V$ can not be taken to $\infty$ independent of $d$. As a result, we call this below the high-dimensional line and it is precisely bounded by whether $2\alpha$ is summable or not.

The four main characteristics of this regime are:

learning rate $\gamma$ scales like $v^{-1+2\alpha}$;

$v$ can not be too large, i.e., $d$ and $v$ are proportional;

batch can be large (i.e., $\gamma B\leq 1)$ since the learning rate is small ($\gamma\sim v^{-1+2\alpha}$).

In Phases IV, Ib, and Ic, because $j^{-2\alpha}$ is not summable, the summation of the $j$ depends on the dimension $v$. Thus, the kernel norm is

limiting value of the loss that SGD converges to is unaffected by $\beta$;

pure point forcing function plays a role;

absolutely continuous part of the spectrum does not contribute to the forcing function;

The following gives the precise statement.

Suppose $2\beta>1$ and $\tfrac{1}{4}<\alpha<\tfrac{1}{2}$. Suppose the learning rate $\gamma$ and batch $B>0$ satisfy at most half the convergence threshold in Proposition 5.2. Then there exists an $M>0$ large and constants $C=C(\alpha,\beta,M)$ and $c=c(\alpha,\beta,M)$, independent of $d$, so that for all admissible $v$ and $d$, for all $\gamma Br>M$

By Theorem 5.1, we know that it suffices to look at the forcing function $\mathscr{F}$ and kernel function $\mathscr{K}$. Moreover, in this regime, we have that $\gamma$ decreases like $d^{2\alpha-1}$ (see Proposition 5.2). The rest of the argument relies on the bounds found in Proposition 10.2, ($\mathscr{F}_{pp})$, Proposition 10.4 ($\mathscr{F}_{ac}$), Proposition 10.3, ($\mathscr{F}_{0})$, and Proposition 10.5 ($\mathscr{K}_{pp}$). We first note there is no $\mathscr{F}_{ac}(r)\lesssim\mathscr{F}_{0}$ and therefore it is too small to contribute.

$\gamma Br\leq M_{0}$, for some $M_{0}$: First, we have that $\tfrac{1}{\gamma B}\mathscr{K}_{pp}(r)\leq C_{0}\times\mathscr{F}_{pp}(r)$ for some constant $C_{0}>0$. The constant $M_{0}$ is where the asymptotic of $\mathscr{F}_{pp}$ starts to apply.

$M_{0}\leq\gamma Br\leq M_{1}$, for some $M_{0}$ and for all $M_{1}>M_{0}$: We see that $\gamma(\gamma Br)^{-2+1/(2\alpha)}<(\gamma Br)^{-(1+\beta/\alpha)+1/(2\alpha)}$ since $\gamma\asymp d^{2\alpha-1}$ and $2\alpha<1$ in this phase. Thus, using Proposition 10.2 and Proposition 10.5, we have that $\frac{1}{\gamma B}\mathscr{K}_{pp}(r)\leq C_{1}\times\mathscr{F}_{pp}$ for some $C_{1}>0$.

$M_{1}\leq\gamma Br\leq M_{2}d^{2\alpha}$, for any $M_{1}$ and some $M_{2}$: The $M_{2}$ is the smallest of the two endpoints for the asymptotics of $\mathscr{K}_{pp}$ and $\mathscr{F}_{pp}$. In this region, $\mathscr{F}_{pp}(r)\asymp(\gamma Br)^{-(1+\beta/\alpha)+1/(2\alpha)}$ and $\gamma\times\frac{1}{\gamma B}\mathscr{K}_{pp}(r)\asymp\gamma\times(\gamma Br)^{-2+1/(2\alpha)}\asymp d^{2\alpha-1}\times(\gamma Br)^{-2+1/(2\alpha)}$. We see at $r=d^{2\alpha}$ that $d^{2\alpha-1}(\gamma Br)^{-1}=d^{-1}\geq d^{-2\beta}=(d^{2\alpha})^{-\beta/\alpha}$. Thus $\mathscr{F}_{pp}(r)\lesssim\tfrac{1}{\gamma B}\mathscr{K}_{pp}(r)$ and we started, i.e., when $r=M_{1}$ with $\mathscr{K}_{pp}(r)\lesssim\mathscr{F}_{pp}(r)$. Therefore, we must change in this regime to being $\mathscr{K}_{pp}$ dominate.

$M_{2}d^{2\alpha}\leq\gamma Br$ for all $M_{2}$: In this case, all terms are bounded above by $\mathscr{F}_{0}(r)$. ∎

As a consequence of the argument above, we know that

𝛼𝛽1(2\beta<1,0.25<\alpha<0.5,2(\alpha+\beta)>1). Phase Ia, Ib, and Ic are quite similar as the dynamics of SGD only depend on the forcing function pure point and limiting value. In this phase, the learning rate $\gamma$ is dimension dependent, unlike Phase Ia, and the following hold

limiting value of the loss that SGD converges to is $d^{-2\alpha+1-2\beta}$;

absolutely continuous part of the spectrum does not contribute to the forcing function;

SGD noise not does affect the loss curves.

Although we did not prove the statement for $\alpha<\tfrac{1}{4}$ as we do not have estimates for the kernel function, we believe that statement still holds. We believe that the kernel function stops becoming power law when $\alpha<\tfrac{1}{4}$, but the forcing function is still power law. The following gives the precise statement.

𝛼𝛽12(\alpha+\beta)>1). Suppose $2\beta<1$, $2(\alpha+\beta)>1$, and $\tfrac{1}{4}<\alpha<\tfrac{1}{2}$. Suppose the learning rate $\gamma$ and batch $B>0$ satisfy at most half the convergence threshold in Proposition 5.2. Then there exists an $M>0$ large and constants $C=C(\alpha,\beta,M)$ and $c=c(\alpha,\beta,M)$, independent of $d$, so that for all admissible $v$ and $d$, for all $\gamma Br>M$

By Theorem 5.1, we know that it suffices to look at the forcing function $\mathscr{F}$ and kernel function $\mathscr{K}$. Moreover, in this regime, we have that $\gamma$ decreases like $d^{2\alpha-1}$ (see Proposition 5.2). The rest of the argument relies on the bounds found in Proposition 10.2, ($\mathscr{F}_{pp})$, Proposition 10.4 ($\mathscr{F}_{ac}$), Proposition 10.3, ($\mathscr{F}_{0})$, and Proposition 10.5 ($\mathscr{K}_{pp}$). We first note there is no $\mathscr{F}_{ac}(r)$.