Title: Scaling Laws in Linear Regression: Compute, Parameters, and Data
Venue: NeurIPS 2024
Accepted: yes
Keywords: learning theory

scaling law

deep learning

SGD

optimization
Abstract: Empirically, large-scale deep learning models often satisfy a neural scaling law: the test error of the trained model improves polynomially as the model size and data size grow. However, conventional wisdom suggests the test error consists of approximation, bias, and variance errors, where the variance error increases with model size. This disagrees with the general form of neural scaling laws, which predict that increasing model size monotonically improves performance.

We study the theory of scaling laws in an infinite dimensional linear regression setup. Specifically, we consider a model with $M$ parameters as a linear function of sketched covariates. The model is trained by one-pass stochastic gradient descent (SGD) using $N$ data. Assuming the optimal parameter satisfies a Gaussian prior and the data covariance matrix has a power-law spectrum of degree $a>1$, we show that the reducible part of the test error is $\Theta(M^{-(a-1)} + N^{-(a-1)/a})$. The variance error, which increases with $M$, is dominated by the other errors due to the implicit regularization of SGD, thus disappearing from the bound. Our theory is consistent with the empirical neural scaling laws and verified by numerical simulation.

Scaling Laws in Linear Regression: Compute, Parameters, and Data

Abstract

Empirically, large-scale deep learning models often satisfy a neural scaling law: the test error of the trained model improves polynomially as the model size and data size grow. However, conventional wisdom suggests the test error consists of approximation, bias, and variance errors, where the variance error increases with model size. This disagrees with the general form of neural scaling laws, which predict that increasing model size monotonically improves performance.

We study the theory of scaling laws in an infinite dimensional linear regression setup. Specifically, we consider a model with M parameters as a linear function of sketched covariates. The model is trained by one-pass stochastic gradient descent (SGD) using N data. Assuming the optimal parameter satisfies a Gaussian prior and the data covariance matrix has a power-law spectrum of degree a > 1 , we show that the reducible part of the test error is Θ( M -( a -1) + N -( a -1) /a ) . The variance error, which increases with M , is dominated by the other errors due to the implicit regularization of SGD, thus disappearing from the bound. Our theory is consistent with the empirical neural scaling laws and verified by numerical simulation.

1 Introduction

Deep learning models, particularly those on a large scale, are pivotal in advancing the state-of-theart across various fields. Recent empirical studies have shed light on the so-called neural scaling laws [see 26, 21, for example], which suggest that the generalization performance of these models improves polynomially as both model size, denoted by M , and data size, denoted by N , increase. The neural scaling law quantitatively describes the population risk as:

where R ∗ is a positive irreducible risk and c 1 , c 2 , a 1 , a 2 are positive constants independent of M and N . For instance, by fitting the above formula with empirical measurements in standard large-scale language benchmarks, Hoffmann et al. [21] estimated a 1 ≈ 0 . 34 and a 2 ≈ 0 . 28 , while Besiroglu et al. [7] estimated that a 1 ≈ 0 . 35 and a 2 ≈ 0 . 37 . Though the exact exponents depend on the tasks, neural scaling laws in (1) are observed consistently in practice and are used as principled guidance to build state-of-the-art models, especially under a compute budget [21].

From the perspective of statistical learning theory, (1) is rather intriguing. Standard statistical learning bounds [see 30, 41, for example] often decompose the population risk into the sum of irreducible

Jason D. Lee

10 2

10 3

Number of effective samples (

N

eff

)

10 1

10 2

Model size (M)

a=1.5

1.2

1.3

1.4

1.5

1.6

Risk

10 2

10 3

Number of effective samples (

N

eff

)

10 1

10 2

10 3

Model size (M)

a=2.0

1.05

1.10

1.15

1.20

1.25

1.30

1.35

Risk

error, approximation error, bias error, and variance error (some theory replaces bias and variance errors by optimization and generalization errors, respectively) as in the form of

where a 1 , a 2 , a 3 are positive constants and c ( M ) is a measure of model complexity that typically increases with the model size M . In (2), the approximation error is induced by the mismatch of the best-in-class predictor and the best possible predictor, hence decreasing with the model size M . The bias error is induced by the mismatch of the expected algorithm output and the best-in-class predictor, hence decreasing with the data size N . The variance error measures the uncertainty of the algorithm output, which decreases with the data size N but increases with the model size M (since the model complexity c ( M ) increases).

A mystery. The empirical neural scaling law (1) is incompatible with the typical statistical learning theory bound (2). While the two error terms in the neural scaling law (1) can be explained by the approximation and bias errors in the theoretical bound (2) respectively, it is not clear why the variance error is unobservable when fitting the neural scaling law empirically . This difference must be reconciled, otherwise, the statistical learning theory and the empirical scaling law make conflict predictions: as the model size M increases, the theoretical bound (2) predicts an increase of variance error that eventually causes an increase of the population risk, but the neural scaling law (1) predicts a decrease of the population risk. In other words, it remains unclear when to follow the prediction of the empirical scaling law (1) and when to follow that of the statistical learning bound (2).

Certain prior works provided risk upper bounds that do not grow with model size [see for example 36, 12]. Still, their results are insufficient for studying scaling law as those bounds require a large model size such that the approximation error is ignorable. Moreover, they do not provide instancewise matching lower bounds to verify the tightness of the upper bounds. See a detailed discussion in Section 2.

Our explanation. We investigate this issue in an infinite dimensional linear regression setup. We only assume access to M -dimensional sketched covariates given by a fixed Gaussian sketch and their responses. We consider a linear predictor with M trainable parameters, which is trained by one-pass stochastic gradient descent (SGD) with geometrically decaying stepsizes using N sketched data. Assuming that the spectrum of the data covariance matrix satisfies a power-law of degree a > 1 and that the optimal model parameters satisfy a Gaussian prior, we derive matching upper and lower bounds on the population risk achieved by the SGD output (see Theorem 4.1). Specifically, we show

that

where γ = O (1) is the initial stepsize used in SGD and ˜ Θ( · ) hides log( N ) factors. In our bound, the sum of the approximation and bias errors determines the order of the excess risk, while the variance error is of a strictly higher order and is therefore nearly unobservable when fitting R ( M,N ) as a function of M and N empirically. In addition, our analysis reveals that the small variance error is due to the implicit regularization effect of one-pass SGD [47]. Our theory suggests that the empirical neural scaling law (1) is a simplification of the statistical learning bound (2) in a special regime when strong regularization (either implicit or explicit) is employed.

Moreover, we generalize the above scaling law to (1) constant stepsize SGD with iterate average (see Theorem F.6), (2) cases where the optimal model parameter satisfies an anisotropic prior (see Theorem 4.2), and (3) where the spectrum of the data covariance matrix satisfies a logarithmic power law (see Theorem 4.3).

Emprical evidence. Based on our theoretical results, we conjecture that the clean neural scaling law (1) observed in practice is due to the disappearance of variance error caused by strong regularization. Two pieces of empirical evidence to support our understanding. First, large language models that follow the scaling law (1) are often underfitted , as the models are trained over a single pass or a few passes over the data [27, 31, 9, 39]. When models are underfitted, the variance error tends to be smaller. Second, when language models are trained with multiple passes (up to 7 passes), Muennighoff et al. [31] found that the clean scaling law in (1) no longer holds and they proposed a more sophisticated scaling law to explain their data. This can be explained by a relatively large variance error caused by multiple passes.

Notation. For two positive-valued functions f ( x ) and g ( x ) , we write f ( x ) ≲ g ( x ) (and f ( x ) = O ( g ( x )) ) or f ( x ) ≳ g ( x ) (and f ( x ) = Ω( g ( x )) ) if f ( x ) ≤ cg ( x ) or f ( x ) ≥ cg ( x ) holds for some absolute (if not otherwise specified) constant c > 0 respectively. We write f ( x ) ≂ g ( x ) (and f ( x ) = Θ( g ( x )) ) if f ( x ) ≲ g ( x ) ≲ f ( x ) . For two vectors u and v in a Hilbert space, we denote their inner product by ⟨ u , v ⟩ or u ⊤ v . For two matrices A and B of appropriate dimensions, we define their inner product by ⟨ A , B ⟩ := tr( A ⊤ B ) . We use ∥ · ∥ to denote the operator norm for matrices and ℓ 2 -norm for vectors. For a positive semi-definite (PSD) matrix A and a vector v of appropriate dimension, we write ∥ v ∥ 2 A := v ⊤ Av . For a symmetric matrix A , we use µ j ( A ) to refer to the j -th eigenvalue of A and r ( A ) to refer to its rank. Finally, log( · ) refers to logarithm base 2 .

2 Related work

Empirical scaling laws. In recent years, the scaling laws of deep neural networks in compute, sample size, and model size have been widely studied across different models and domains [20, 35, 26, 19, 21, 46, 31]. The early work by Kaplan et al. [26] first proposed the neural scaling laws of transformer-based models. They observed that the test loss exhibits a power-law decay in quantities including the amount of compute, sample size, and model size, and provided joint formulas in these quantities to predict the test loss. The proposed formulas were later generalized and refined in subsequent works [19, 21, 1, 10, 31]. Notably, Hoffmann et al. [21] proposed the Chinchilla law, that is, (1) with a 1 ≈ 0 . 34 and a 2 ≈ 0 . 28 . The empirical observation guided them to allocate data and model size under a given compute budget. The Chinchilla law is further revised by Besiroglu et al. [7]. Motivated by the Chinchilla law, Muennighoff et al. [31] considered the effect of multiple passes over training data and empirically fitted a more sophisticated scaling law that takes account of the effect of data reusing.

Theory of scaling laws. Although neural scaling laws have been empirically observed over a broad spectrum of problems, there is a relatively limited literature on understanding these scaling laws from a theoretical perspective [37, 4, 28, 22, 42, 29, 23, 8, 2, 32, 17]. Among these works, [37] showed that the test loss scales as N 4 /d for regression on data with intrinsic dimension d . Hutter [22] studied a toy problem under which a non-trivial power of N arises in the test loss. Jain et al. [23]

considered scaling laws in data selection. Bahri et al. [4] considered a linear teacher-student model under a power-law spectrum assumption on the covariates, and they showed that the test loss of the ordinary least square estimator decreases following a power law in sample size N (resp. model size M ) when the model size M (resp. sample size N ) is infinite. Bordelon et al. [8] considered a linear random feature model and analyzed the test loss of the solution found by (batch) gradient flow. They focused on the bottleneck regimes where two of the quantities N,M , T (training steps) are infinite and showed that the risk has a power-law decay in the remaining quantity. The problem in Bahri et al. [4], Bordelon et al. [8] can be viewed as a sketched linear regression model similar to ours. It should be noted that both Bahri et al. [4] and Bordelon et al. [8] only derived the dependence of population risk on one of the data size, model size, or training steps in the asymptotic regime where the remaining quantities go to infinity, and their derivations are based on statistical physics heuristics. In comparison, we prove matching (ignoring constant factors) upper and lower risk bounds jointly depending on the finite model size M and data size N .

Implicit regularization of SGD. One-pass SGD in linear regression has been extensively studied in both the classical finite-dimensional setting [34, 3, 14, 16, 25, 24, 18] and the modern highdimensional setting [15, 6, 48, 47, 44, 45, 40]. In particular, Zou et al. [47] showed that SGD induces an implicit regularization effect that is comparable to, and in certain cases even more preferable than, the explicit regularization effect induced by ridge regression. This is one of the key motivations of our scaling law interpretation. From a technical perspective, we utilize the sharp finite-sample and dimension-free analysis of SGD developed by Zou et al. [48], Wu et al. [44, 45]. Different from them, we consider a sequence of linear regression models with an increasing number of trainable parameters given by data sketch. Our main technical innovation is to sharply control the effect of data sketch. Some of our intermediate results, for example, tight bounds on the spectrum of the sketched data covariance under the power law (see Lemma 6.2), might be of independent interest.

Prior works investigated linear regression with random features [36, 12], which can be viewed as a kind of sketched features via random coordinate selection. They mainly focused on the small approximation error regime, where the model size (or the number of features) is much larger than the data size. In comparison, we treat both model size and data size as free variables. Moreover, we provide matching upper and lower bounds while prior works mainly focused on upper bounds. These two innovations are crucial for studying scaling laws that predict test error as a function of both model size and data size. Finally, in the comparable regimes with small or zero approximation error, our excess risk bounds recover the bounds in prior works [36, 12, 33, 15, 13].

3 Setup

We use x ∈ H to denote a feature vector, where H is a finite d -dimensional or countably infinite dimensional Hilbert space, and y ∈ R to denote its label. In linear regression, we measure the population risk of a parameter w ∈ H by the mean squared error,

where the expectation is over ( x , y ) ∼ P for some distribution P on H × R .

Definition 1 (Data covariance and optimal parameter) . Let H := E [ xx ⊤ ] be the data covariance. Assume that tr( H ) and all entries of H are finite. Let ( λ i ) i ≥ 0 be the eigenvalues of H sorted in non-increasing order. Let w ∗ ∈ arg min w R ( w ) be the optimal model parameter 1 . Assume that ∥ w ∗ ∥ 2 H := ( w ∗ ) ⊤ Hw ∗ is finite.

We only assume access to M -dimensional sketched covariates and their responses, that is, ( Sx , y ) , where S ∈ R M × H is a fixed sketch matrix. We focus on the Gaussian sketch matrix 2 , that is, entries of S are independently sampled from N ( 0 , 1 /M ) . We then consider linear predictors with M trainable parameters given by

where v ∈ R M are the trainable parameters. Varying M should be viewed as a linear analog of varying the neural network model size. Our sketched linear regression setting is comparable to the teacher-student setting considered by Bahri et al. [4], Bordelon et al. [8].

We consider the training of f v via one-pass stochastic gradient descent (SGD), that is,

where ( x t , y t ) N t =1 are independent samples from P and ( γ t ) N t =1 are the stepsizes. We consider a popular geometric decaying stepsize scheduler [18, 44],

Here, the initial stepsize γ is a hyperparameter for the SGD algorithm. Without loss of generality, we assume the initial parameter is v 0 = 0 . The output of the SGD algorithm is the last iterate v N . Our proof techniques apply to other stepsize schedulers (e.g., polynomial decay) as well, but we focus on geometric decay as it is known to achieve near minimax-optimal excess risk for the last iterate of SGD [18].

Conditioning on a sketch matrix S ∈ R M × H , each parameter v ∈ R M induces a sketched predictor through x ↦→⟨ S ⊤ v , x ⟩ , and we denote its risk by

By increasing M and N , we have a sequence of datasets and trainable parameters of increasing sizes, respectively. This prepares us to study the scaling law (1) in the sketched linear regression problem, that is, to understand R M ( v N ) as a function of both M and N .

Risk decomposition. In a standard way, we decompose the risk achieved by v N , the last iterate of (SGD), to the sum of irreducible risk , approximization error , and excess risk as follows,

We emphasize that the irreducible risk is independent of M and N and thus can be viewed as a constant; the approximation error is determined by the sketch matrix S , thus depends on M but is independent of N ; the excess risk depends on both M and N as it is determined by the algorithm.

4 Scaling laws

We first demonstrate a scaling-law behavior when the data spectrum satisfies a power law.

Assumption 1 (Distributional conditions) . Assume the following about the data distribution.

Gaussian design. Assume that x ∼ N (0 , H ) .

Well-specified model. Assume that E [ y | x ] = x ⊤ w ∗ . Define σ 2 := E ( y -x ⊤ w ∗ ) 2 .

Parameter prior. Assume that w ∗ satisfies a prior such that E ( w ∗ ) ⊗ 2 = I .

Assumption 2 (Power-law spectrum) . There exists a > 1 such that the eigenvalues of H satisfy λ i ≂ i -a , i > 0 .

Theorem 4.1 (Scaling law) . Suppose that Assumptions 1 and 2 hold. Consider an M -dimensional sketched predictor trained by (SGD) with N samples. Let N eff := N/ log( N ) and recall the risk decomposition in (4) . Then there exists some a -dependent constant c > 0 such that when the initial stepsize γ ≤ c , with probability at least 1 -e -Ω( M ) over the randomness of the sketch matrix S , we have

Irreducible := R ( w ∗ ) = σ 2 .

E w ∗ Approx ≂ M 1 -a .

Suppose in addition σ 2 ≳ 1 . The expected excess risk ( Excess ) can be decomposed into a bias error ( Bias ) and a variance error ( Var ), namely,

where the expectation is over the randomness of w ∗ and ( x i , y i ) N i =1 . Moreover, Bias and Var satisfy

In all results, the hidden constants only depend on the power-law degree a . As a direct consequence, when σ 2 ≂ 1 , it holds with probability at least 1 -e -Ω( M ) over the randomness of the sketch matrix S that

where the expectation is over the randomness of w ∗ and ( x i , y i ) N i =1 .

Theorem 4.1 shows a sharp (up to constant factors) scaling law risk bound under an isotroptic prior assumption and the power-law spectrum assumption. We emphasize that the scaling law bound in Theorem 4.1 holds for every M,N ≥ 1 . We also remark that the sum of approximization and bias errors dominates E R M ( v N ) -σ 2 , whereas the variance error is of strict higher order in terms of both M and N , and is thus disappeared in the population risk bound.

Optimal stepsize. Based on the tight scaling law in Theorem 4.1, we can calculate the optimal stepsize that minimizes the risk. Specifically, the optimal stepsize is γ ≂ 1 when N eff ≲ M a and can be anything such that M a /N eff ≲ γ ≲ 1 when N eff ≳ M a . In both cases, choosing γ ≂ 1 is optimal. When the sample size is large such that N eff ≳ M a , the optimal stepsize is relatively robust and can be chosen from a range.

Allocation of data and model sizes. Following Hoffmann et al. [21], we measure the compute complexity by MN as (SGD) queries M -dimensional gradients for N times. Given a total compute budget of MN = C , from Theorem 6.1 and N eff := N/ log( N ) , we see that the best population risk is achieved by setting γ = Θ(1) , M = ˜ Θ( C 1 / ( a +1) ) , and N = ˜ Θ( C a/ ( a +1) ) . Our theory suggests setting a data size slightly larger than the model size when the compute budget is the bottleneck.

Comparison with [8]. The work by Bordelon et al. [8] considered the scaling law of batch gradient descent (or gradient flow) on a teacher-student model (see their equation (14)). Their teacher-student model can be viewed as our sketched linear regression model. However, we consider one-pass SGD, therefore in our setting the number of gradient steps is equivalent to the data size. When we equalize the number of gradient steps and the data size in their equation (14) and set the parameter prior as Assumption 1C, their prediction is consistent with ours. However, our analysis shows the computational advantage of SGD over batch GD since each iteration requires only 1 /N the compute. Bordelon et al. [8] obtained the limit of the population risk as two out of the data size, model size, and the number of gradient steps go to infinity based on statistical physics heuristics. In comparison, we obtain upper and lower risk bounds that hold for any finite M and N and match ignoring a constant factor depending only on the spectrum power-law degree a .

Average of the SGD iterates Results similar to Theorem 4.1 can also be established for the average of the iterates of online SGD with constant stepsize [34, 16, 25, 24, 48]. All results will be the same once replacing the effective sample size N eff in Theorem 4.1 to the sample size N . For more details see Theorem F.6 in Appendix F.

4.1 Scaling law under source condition

The isotropic parameter prior condition (Assumption 1C) in Theorem 4.1 can be generalized to the following anisotropic version [11].

Assumption 3 (Source condition) . Let ( λ i , v i ) i> 0 be the eigenvalues and eigenvectors of H with ( λ i ) i> 0 in non-increasing order. Assume w ∗ satisfies a prior such that

̸

A larger exponent b implies a faster decay of signal w ∗ and thus corresponds to a simpler task [11]. Note that Assumption 1C satisfies Assumption 3 with b = a .

Theorem 4.2 (Scaling law under source condition) . In Theorem 4.1, suppose Assumption 1C is replaced by Assumption 3 with 1 < b < a +1 . Then there exists some a -dependent constant c > 0 such that when γ ≤ c , with probability at least 1 -e -Ω( M ) over the randomness of the sketch matrix S , we have

where the expectation is over the randomness of w ∗ and ( x i , y i ) N i =1 , and Θ( · ) hides constants that may depend on ( a, b ) .

When 1 < b ≤ a , the tasks are relatively hard (compared to when b = a ), and the variance error is dominated by the sum of approximation and bias errors for all choices of M , N , and γ ≲ 1 . In this case, Theorem 4.2 gives the same prediction about optimal stepsize and optimal allocation of data and model sizes under compute budget as Theorem 4.1.

When a < b < a +1 , the tasks are relatively easy (compared to when b = a ), and variance remains dominated by the sum of approximation and bias error if the stepsize is optimally tuned. Recall that γ ≲ 1 , thus we can rewrite the risk bound in Theorem 4.2 as

Therefore the optimal stepsize and the risk under the optimal stepsize is

Under the optimally tuned stepsize, the population risk is in the form of

which is again in the scaling law form (1). This is expected since an optimally tuned stepsize controls the variance error by adjusting the strength of the implicit bias of SGD. Under a fixed compute budget C = MN , our theory suggests to assign M = ˜ Θ( C 1 / ( b +1) ) and N = ˜ Θ( C b/ ( b +1) ) , and set the stepsize to γ ≂ ˜ Θ( C ( a -b ) / ( b +1) ) .

When b ≥ a +1 , the tasks are even simpler. We provide upper and lower bounds in Appendix D.3. However, there exists a gap between the bounds, fixing which is left for future work.

Moreover, we note that in the comparable regimes where M is large, the results in Theorem 4.2 match existing bounds on the risk of SGD iterates and ridge estimators [33, 36].

4.2 Scaling law under logarithmic power law

We also derive the risk formula when the data covariance has a logarithmic power-law spectrum [5]. Assumption 4 (Logarithmic power-law spectrum) . There exists a > 1 such that the eigenvalues of H satisfy λ i ≂ i -1 log -a ( i +1) , i > 0 .

Theorem 4.3 (Scaling law under logarithmic power spectrum) . In Theorem 4.1, suppose Assumption 2 is replaced by Assumption 4. Then with probability at least 1 -e -Ω( M ) over the randomness of the sketch matrix S , we have

where the expectation is over the randomness of w ∗ and ( x i , y i ) N i =1 .

10 2

10 3

10 4

10 5

Number of effective samples (

N

eff

)

10

1

2×10 1

3×10 1

4×10 1

6×10 1

E

M

(

v

N

)

min

( )

a=1.5

Risk

(a)

a

= 1

.

5

10 1

10 2

Model size (M)

2×10 1

3×10 1

4×10 1

E

M

(

v

N

)

min

( )

a=1.5

k=-0.49

Risk

(b)

a

= 1

.

5

10 2

10 3

10 4

10 5

Number of effective samples (

N

eff

)

10

1

E

M

(

v

N

)

min

( )

a=2.0

k=-0.50

Risk

(c)

a

= 2

10 1

10 2

Model size (M)

10

1

E

M

(

v

N

)

min

( )

a=2.0

k=-0.98

Risk

(d)

a

= 2

Theorem 4.3 provides a scaling law under the logarithmic power-law spectrum. Similar to Theorem 4.1, the variance error is dominated by the approximation and bias errors for all choices of M , N , and γ , and thus disappeared from the risk bound. Different from Theorem 4.1, here the population risk is a polynomial of log( M ) and log( N eff γ ) .

5 Experiments

In this section, we examine the relation between the expected risk of the (SGD) output, the data size N , and the model size M when the covariates satisfy a power-law covariance spectrum. Although our results in Section 4 hold with high probability over S , for simplicity, we assume the expectation of the risk is taken over both w ∗ and S in our simulations. We adopt the model in Section 3 and train it using one-pass (SGD) with geometric decaying stepsize (3). We choose the dimension d sufficiently large to approximate the infinite-dimensional case, and the data are generated so that Assumption 1 is satisfied. Moreover, we choose the covariance H ∈ R d × d to be diagonal with H ii ∝ i a and tr( H ) = 1 for some a > 1 . From Figure 1, we observe that the risk indeed follows a power-law formula jointly in the number of samples and the number of parameters. In addition, the fitted exponents are aligned with our theoretical predictions ( a -1 , 1 -1 /a ) in Theorem 4.1. Figure 2 shows the scaling of the expected risk in data size (or model size) when the model size (or data size) is relatively large. We see that the expected risk also satisfies a power-law decay with exponents matching our predictions. It is noteworthy that our simulations demonstrate stronger observations than the theoretical results in Theorem 4.1, which only establishes matching upper and lower bounds up to a constant factor. Additional simulation results on the risk of the average of (SGD) iterates can be found in Appendix F.

6 Risk bounds under a general spectrum

In this section, we present some general results on the upper and lower bounds of the risk of the output of (SGD). Due to the rotational invariance of the sketched matrix S , without loss of generality, we assume the covariance H is diagonal with non-increasing diagonal entries. Our main results in Section 4 are directly built on the general bounds introduced here.

Assumption 5 (General distributional conditions) . Assume the following about the data distribution.

Hypercontractivity. There exists α ≥ 1 such that for every PSD matrix A it holds that

Misspecified model. There exists σ 2 > 0 such that E ( y -x ⊤ w ∗ ) 2 xx ⊤ ⪯ σ 2 H .

It is clear that Assumption 1 implies Assumption 5 with α = 3 .

Excess risk decomposition. Conditioning on the sketch matrix S , the training of the sketched linear predictor can be viewed as an M -dimensional linear regression problem. We can therefore

invoke existing SGD analysis [44, 45] to sharply control the excess risk by controlling the bias and variance errors. Specifically, let us define the ( w ∗ -dependent) bias error as

and the variance error as

where ( ˜ λ j ) M j =1 are eigenvalues of SHS ⊤ . We also let Bias := E Bias ( w ∗ ) , where the expectation is over the prior of w ∗ . Using the existing results on the output of (SGD) in Wu et al. [44, 45], we show that the excess risk in (4) can be exactly decomposed as the sum of bias and variance errors under weak conditions.

Theorem 6.1 (Excess risk decomposition) . Conditioning on the sketch matrix S , consider the excess risk in (4) induced by the output of (SGD) . Assume v 0 = 0 . Then for any w ∗ ∈ H ,

Under Assumptions 5A and 5B and suppose γ ≤ 1 / ( cα tr( SHS ⊤ ) ) for some constant c > 0 , we have

Under the stronger Assumptions 1A and 1B and suppose γ ≤ 1 / ( cα tr( SHS ⊤ )) for some constant c > 0 , we have

In both results, the expectations of Excess are taken over ( x t , y t ) N t =1 .

Assuming that the signal-to-noise ratio is upper bounded, that is, ∥ w ∗ ∥ 2 H /σ 2 ≲ 1 , then the biasvariance decomposition of the excess risk is sharp up to constant factors.

The variance error is in a nice form and can be computed using the following important lemma on the spectrum of SHS ⊤ . Similar results for logarithmic power-law are also established in Lemma G.6 in Appendix G.

Lemma 6.2 (Power law) . Under Assumption 2, it holds with probability at least 1 -e -Ω( M ) that

For any 0 ≤ k ∗ ≤ k † ≤ ∞ , let S k ∗ : k † ∈ R M × ( k † -k ∗ ) denote the matrix formed by the k ∗ +1 -k † -th columns of S . We also abuse the notation k † : ∞ for k † : d when d is finite. We let H k ∗ : k † ∈ R ( k † -k ∗ ) × ( k † -k ∗ ) be the submatrix of H formed by the k ∗ + 1 -k † -th eigenvalues. For the approximation and bias error, we use the following upper and lower bounds to compute their values.

Theorem 6.3 (A general upper bound) . Suppose Assumption 5 holds. Assume v 0 = 0 , r ( H ) ≥ 2 M and the initial stepsize satisfies γ < 1 / ( cα tr( SHS ⊤ )) for some constant c > 0 . Then for any k 1 , k 2 ≤ M/ 3 , with probability at least 1 -e -Ω( M )

Theorem 6.4 (A general lower bound) . Suppose Assumption 1 holds. Assume v 0 = 0 , r ( H ) ≥ M and the initial stepsize γ < 1 / ( c tr( SHS ⊤ ) ) for some constant c > 0 . Then

almost surely, where ( λ i ) d i =1 are eigenvalues of H in non-increasing order, ( ˜ λ i ) d i =1 are eigenvalues of SHS ⊤ in non-increasing order.

7 Conclusion

We analyze neural scaling laws in infinite-dimensional linear regression. We consider a linear predictor with M trainable parameters on the sketched covariates, which is trained by one-pass stochastic gradient descent with N data. Under a Gaussian prior assumption on the optimal model parameter and a power law (of degree a > 1 ) assumption on the spectrum of the data covariance, we derive matching upper and lower bounds on the population risk minus the irreducible error, that is, Θ( M -( a -1) + N -( a -1) /a ) . In particular, we show that the variance error, which increases with M , is of strictly higher order compared to the other errors, thus disappearing from the risk bound. We attribute the nice empirical formula of the neural scaling law to the non-domination of the variance error, which ultimately is an effect of the implicit regularization of SGD.

Many directions remain open for future study. First, our work is limited to the linear model; it would be interesting to see whether similar scaling laws can be derived in more complex models, such as random feature models or two-layer networks. Second, we focus on one-pass SGD training, and it is unclear if similar results hold for other optimization methods like accelerated SGD or Adam. Additionally, from a technical perspective, many results in our work depend on the Gaussian assumption and the source condition of the data. Investigating how these assumptions can be relaxed would also be valuable.

Bibliography

Ibrahim M Alabdulmohsin, Behnam Neyshabur, and Xiaohua Zhai. Revisiting neural scaling laws in language and vision. Advances in Neural Information Processing Systems , 35:2230022312, 2022.

Alexander B Atanasov, Jacob A Zavatone-Veth, and Cengiz Pehlevan. Scaling and renormalization in high-dimensional regression. arXiv preprint arXiv:2405.00592 , 2024.

Francis Bach and Eric Moulines. Non-strongly-convex smooth stochastic approximation with convergence rate o (1 /n ) . Advances in neural information processing systems , 26:773-781, 2013.

Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXiv:2102.06701 , 2021.

Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences , 2020.

Raphaël Berthier, Francis Bach, and Pierre Gaillard. Tight nonparametric convergence rates for stochastic gradient descent under the noiseless linear model. Advances in Neural Information Processing Systems , 33:2576-2586, 2020.

Tamay Besiroglu, Ege Erdil, Matthew Barnett, and Josh You. Chinchilla scaling: A replication attempt. arXiv preprint arXiv:2404.10102 , 2024.

Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. A dynamical model of neural scaling laws. arXiv preprint arXiv:2402.01092 , 2024.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Ethan Caballero, Kshitij Gupta, Irina Rish, and David Krueger. Broken neural scaling laws. arXiv preprint arXiv:2210.14891 , 2022.

Andrea Caponnetto and Ernesto De Vito. Optimal rates for the regularized least-squares algorithm. Foundations of Computational Mathematics , 7:331-368, 2007.

Luigi Carratino, Alessandro Rudi, and Lorenzo Rosasco. Learning with sgd and random features. Advances in neural information processing systems , 31, 2018.

Leonardo Defilippis, Bruno Loureiro, and Theodor Misiakiewicz. Dimension-free deterministic equivalents for random feature regression. arXiv preprint arXiv:2405.15699 , 2024.

Alexandre Défossez and Francis Bach. Averaged least-mean-squares: Bias-variance trade-offs and optimal sampling distributions. In Artificial Intelligence and Statistics , pages 205-213, 2015.

Aymeric Dieuleveut and Francis R. Bach. Non-parametric stochastic approximation with large step sizes. The Annals of Statistics , 2015.

Aymeric Dieuleveut, Nicolas Flammarion, and Francis Bach. Harder, better, faster, stronger convergence rates for least-squares regression. The Journal of Machine Learning Research , 18 (1):3520-3570, 2017.

Elvis Dohmatob, Yunzhen Feng, Pu Yang, Francois Charton, and Julia Kempe. A tale of tails: Model collapse as a change of scaling laws. arXiv preprint arXiv:2402.07043 , 2024.

Rong Ge, Sham M Kakade, Rahul Kidambi, and Praneeth Netrapalli. The step decay schedule: A near optimal, geometrically decaying learning rate procedure for least squares. Advances in neural information processing systems , 32, 2019.

Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, et al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010.14701 , 2020.

Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically. arXiv preprint arXiv:1712.00409 , 2017.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556 , 2022.

Marcus Hutter. Learning curve theory. arXiv preprint arXiv:2102.04074 , 2021.

Ayush Jain, Andrea Montanari, and Eren Sasoglu. Scaling laws for learning with real and surrogate data. arXiv preprint arXiv:2402.04376 , 2024.

Prateek Jain, Praneeth Netrapalli, Sham M Kakade, Rahul Kidambi, and Aaron Sidford. Parallelizing stochastic gradient descent for least squares regression: mini-batching, averaging, and model misspecification. The Journal of Machine Learning Research , 18(1):8258-8299, 2017.

Prateek Jain, Sham M. Kakade, Rahul Kidambi, Praneeth Netrapalli, Venkata Krishna Pillutla, and Aaron Sidford. A Markov Chain Theory Approach to Characterizing the Minimax Optimality of Stochastic Gradient Descent (for Least Squares). In 37th IARCS Annual Conference on Foundations of Software Technology and Theoretical Computer Science (FSTTCS 2017) , 2018.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020.

Aran Komatsuzaki. One epoch is all you need. arXiv preprint arXiv:1906.06669 , 2019.

Alexander Maloney, Daniel A Roberts, and James Sully. A solvable model of neural scaling laws. arXiv preprint arXiv:2210.16859 , 2022.

Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quantization model of neural scaling. Advances in Neural Information Processing Systems , 36, 2024.

Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning . MIT press, 2018.

Niklas Muennighoff, Alexander M Rush, Boaz Barak, Teven Le Scao, Aleksandra Piktus, Nouamane Tazi, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling data-constrained language models. arXiv preprint arXiv:2305.16264 , 2023.

Yoonsoo Nam, Nayara Fonseca, Seok Hyeong Lee, and Ard Louis. An exactly solvable model for emergence and scaling laws. arXiv preprint arXiv:2404.17563 , 2024.

Loucas Pillaud-Vivien, Alessandro Rudi, and Francis Bach. Statistical optimality of stochastic gradient descent on hard learning problems through multiple passes. Advances in Neural Information Processing Systems , 31, 2018.

Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM journal on control and optimization , 30(4):838-855, 1992.

Jonathan S Rosenfeld, Amir Rosenfeld, Yonatan Belinkov, and Nir Shavit. A constructive prediction of the generalization error across scales. arXiv preprint arXiv:1909.12673 , 2019.

Alessandro Rudi and Lorenzo Rosasco. Generalization properties of learning with random features. Advances in neural information processing systems , 30, 2017.

Utkarsh Sharma and Jared Kaplan. A neural scaling law from the dimension of the data manifold. arXiv preprint arXiv:2004.10802 , 2020.

William Swartworth and David P Woodruff. Optimal eigenvalue approximation via sketching. In Proceedings of the 55th Annual ACM Symposium on Theory of Computing , pages 145-155, 2023.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Aditya Vardhan Varre, Loucas Pillaud-Vivien, and Nicolas Flammarion. Last iterate convergence of SGD for least-squares in the interpolation regime. In Advances in Neural Information Processing Systems , 2021.

Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge University Press, 2019.

Alexander Wei, Wei Hu, and Jacob Steinhardt. More than a toy: Random matrix models predict how real-world neural representations generalize. In Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , 2022.

David P Woodruff et al. Sketching as a tool for numerical linear algebra. Foundations and Trends® in Theoretical Computer Science , 10(1-2):1-157, 2014.

Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham M. Kakade. Last iterate risk bounds of sgd with decaying stepsize for overparameterized linear regression. The 39th International Conference on Machine Learning , 2022.

Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham M. Kakade. The power and limitation of pretraining-finetuning for linear regression under covariate shift. The 36th Conference on Neural Information Processing Systems , 2022.

Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 12104-12113, 2022.

Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, Dean P Foster, and Sham Kakade. The benefits of implicit regularization from sgd in least squares problems. Advances in Neural Information Processing Systems , 34:5456-5468, 2021.

Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham M Kakade. Benign overfitting of constant-stepsize sgd for linear regression. Journal of Machine Learning Research , 24(326):1-58, 2023.

Appendix

Table of Contents

A Preliminary

In this section, we provide some preliminary discussions and a proof of Theorem 6.1. Concretely, in Section A.1 we discuss our data assumptions and introduce additional notations. In Section A.2, A.3 we derive intermediate results that contribute to the proof of Theorem 6.1. Finally, a complete proof of Theorem 6.1 is contained in Section A.4.

A.1 Additional notations and comments on data assumptions

Tensors. For matrices A , B , C , D , and X of appropriate shape, it holds that

and that

For simplicity, we denote

Comments on Assumption 2, 3 and 4 Due to the rotational invariance of the Gaussian sketched matrix S , throughout the appendix, we assume w.l.o.g. that the covariance of the input covariates H is diagonal with the ( i, i ) -th entry being the i -th eigenvalue. Specifically, Assumption 3 can be rewritten as

Assumption 6 (Source condition) . Assume H = ( h ij ) i,j ≥ 1 is a diagonal matrix with diagonal entries in non-increasing order, and w ∗ satisfies a prior such that

̸

Now that we assume H is diagonal. We make the following notations. Define

where 0 ≤ k ∗ ≤ k † are two integers, and we allow k † = ∞ . For example,

Similarly, for a vector w ∈ H , we have

A.2 Approximation error

Recall the risk decomposition in (4),

Lemma A.1 (Approximization error) . Conditional on the sketch matrix S , the minimizer of R M ( v ) is given by

and the approximation error in (4) is

Moreover, Approx ≤ ∥ w ∗ ∥ 2 H almost surely over the randomness of S .

Proof of Lemma A.1. Recall that the risk

is a quadratic function and that w ∗ is the minimizer of R ( · ) , so we have

Recall that the risk in a restricted subspace

is also a quadratic function, so its minimizer is given by

and

Therefore, the approximation error is

Finally, since

it follows that Approx ≤ ∥ w ∥ H

A.3 Bias-variance decomposition

The excess risk in (4) can be viewed as the SGD excess risk in an M -dimensional (misspecified) linear regression problem. We will utilize Corollary 3.4 in [45] to get a bias-variance decomposition of the excess risk. The following two lemmas check the related assumptions for Corollary 3.4 in [45] in our setup.

Lemma A.2 (Hypercontractivity and the misspecified noise under sketched feature) . Suppose that Assumptions 5A and 5B hold. Conditioning on the sketch matrix S , for every PSD matrix A ∈ R M × M , we have

Moreover, for the minimizer of R M ( v ) , that is, v ∗ defined in Lemma A.1, we have

The expectation in the above is over ( x , y ) .

Proof of Lemma A.2. The first part is a direct application of Assumption 5A:

For the second part, we first show that

where the last inequality is by Assumptions 5A and 5B. From the proof of Lemma A.1, we know that

So we have

Left and right multiplying both sides with S and S ⊤ , we obtain the second claim.

Lemma A.3 (Gaussianity and well-specified noise under sketched features) . Suppose that Assumptions 1A and 1B hold. Conditional on the sketch matrix S , we have

Moreover, for the minimizer of R M ( v ) , that is, v ∗ defined in Lemma A.1, we have

Proof of Lemma A.3. The first claim is a direct consequence of Assumption 1A. For the second claim, by Assumption 1A and Lemma A.1, we have

Notice that by Assumption 1A and that

therefore

Taking expectation over the second random vector in (8), we find

It remains to show

This follows from the proof of Lemma A.1. Specifically,

where the second equality is by the definition of Approx and the third equality is by Assumption 1B. We have completed the proof.

A.4 Proof of Theorem 6.1

We now use the results in [44, 45] for SGD to obtain the following bias-variance decomposition on the excess risk.

Theorem A.4 (Excess risk bounds) . Consider the excess risk in (4) induced by the output of (SGD) . Let

Then conditioning on the sketch matrix S , for any w ∗ ∈ H

Under Assumptions 5A and 5B, we have

when γ ≲ 1 cα tr( SHS ⊤ ) for some constant c > 0 .

Under Assumptions 1A and 1B, we have

when γ ≲ 1 c tr( SHS ⊤ ) for some constant c > 0 .

In both results, the expectation is over ( x t , y t ) N t =1 , and

where ( ˜ λ j ) j M =1 are eigenvalue of SHS ⊤ .

Theorem 6.1 follows immediately by Lemma A.1 and by setting v 0 = 0 and plugging the definition of Bias ( w ∗ ) and Var into Theorem A.4.

Proof of Theorem A.4. This follows from Corollary 3.4 in [45] for a linear regression problem with population data given by ( Sx , y ) . Note that the data covariance becomes SHS ⊤ and the optimal model parameter becomes v ∗ .

For the upper bound, Lemma A.2 verifies Assumptions 1A and 2 in [45], with the noise level being ˜ σ 2 = 2( σ 2 + α ∥ w ∗ ∥ 2 ) .

Then we can apply the upper bound in Corollary 3.4 in [45] (setting their index set K = ∅ ) to get

We verify that

which implies that

Substituting, we get the upper bound.

For the lower bound, Lemma A.3 shows Sx is Gaussian, therefore it satisfies Assumption 1B in Wu et al. [45] with β = 1 . Besides, Lemma A.3 shows that the linear regression problem is well-specified, with the noise level being

Although the lower bound in Corollary 3.4 in Wu et al. [45] is stated for Gaussian additive noise (see their Assumption 2'), it is easy to check that the lower bound holds for any well-specified noise as described by Lemma A.3. Using the lower bound in Corollary 3.4 in Wu et al. [45], we obtain

Plugging in ˜ σ 2 ≥ σ 2

gives the desired lower bound.

A.5 Proofs of Lemma 6.2, Theorem 6.3 and 6.4

Lemma 6.2 is proved in Lemma G.4. Theorem 6.3 follows from Lemma C.1 and D.1. Theorem 6.4 follows from Lemma C.2 and D.2.

B Proofs in Section 4

B.1 Proof of Theorem 4.1

Proof of part 1. By Assumption 1B and the definition of R ( · ) , we have

Note that the equality holds if and only if w = w ∗ . Therefore we have min R ( · ) = R ( w ∗ ) = σ 2 .

Proof of part 2. Part 2 of Theorem 4.1 follows immediately from Lemma C.2.

Proof of part 3. We choose Bias ( w ∗ ) , Var as defined in Eq. (5) and (6) and let Bias := E w ∗ Bias ( w ∗ ) . Part 3 of Theorem 4.1 follows directly from the decomposition of the excess risk in Theorem 6.1 (note that E ∥ w ∗ ∥ 2 H /σ 2 ≲ 1 ), and the matching bounds in Lemma D.3 and E.1.

It remains to verify the stepsize assumption required in Lemma D.3. Since we have from Lemma G.4 that

for some a -dependent constants c 1 , c 2 , c 3 > 0 with probability at least 1 -e -Ω( M ) , it follows that for any constant c > 1 , we can choose γ ≤ c 0 for some a -dependent c 0 such that γ ≤ 1 c tr( SHS ⊤ ) . Therefore, we have verified the stepsize assumption.

Finally, the last claim in Theorem 4.1 follows directly from combining the previous three parts and Theorem 6.1, noting σ 2 ≲ 1 . and

under the stepsize assumption γ ≲ 1 . Here the hidden constants may depend on a.

B.2 Proof of Theorem 4.2

Similar to the proof of Theorem 4.1, we have min R ( · ) = σ 2 under Assumption 1B. Moreover, by Lemma C.5, D.4 and E.1, we have with probability at least 1 -e -Ω( M ) that

when the stepsize γ ≤ c for some a -dependent constant c > 0 . Here the hidden constants in the bounds may depend only on ( a, b ) . Combining the bounds on Approx , Bias , Var and noting

yields Theorem 4.2. Here in the second inequality, we use the assumption b ≤ a .

B.3 Proof of Theorem 4.3

Similar to the proof of Theorem 4.1, we have min R ( · ) = σ 2 under Assumption 1B. Notice that we have γ ≲ 1 implies γ ≲ 1 / (tr( SHS ⊤ )) with probability at least 1 -e -Ω( M ) by Lemma G.6. It follows from Lemma C.6, D.5 and E.2 that

with probability at least 1 -e -Ω( M ) when the stepsize γ ≤ c for some a -dependent constant c > 0 . Since Var ≲ Bias and log 1 -a ( N eff γ ) ≲ log 1 -a M when ( N eff γ ) 1 /a ≳ M c , putting the bounds together gives Theorem 4.3.

C Approximation error

In this section, we derive upper and lower bounds for the approximation error in (4) (and 7). We will also show that the upper and lower bounds match up to constant factors in several examples.

C.1 An upper bound

LemmaC.1 (An upper bound on the approximation error) . Given any k ≤ d such that r ( H ) ≥ k + M , the approximation error in (4) (and 7) satisfies

almost surely, where A k := S k : ∞ H k : ∞ S ⊤ k : ∞ . If in addition k ≤ M/ 2 , then with probability 1 -e -Ω( M )

where ( λ i ) p i =1 are eigenvalues of H in non-increasing order.

Proof of Lemma C.1. Write the singular value decomposition H = UΛU ⊤ , where Λ := diag { λ 1 , λ 2 , . . . } with λ 1 ≥ λ 2 ≥ . . . ≥ 0 and UU ⊤ = I . Define ˜ S := SU , ˜ w ∗ := U ⊤ w ∗ . Then by Lemma A.1 the approximation error Approx = Approx ( S , H , w ∗ ) satisfies

Since ˜ S d = S by rotational invariance of standard gaussian variables, it suffices to analyze the case where H = Λ is a diagonal matrix, as the results may transfer to general H by replacing ˜ w ∗ with w ∗ .

Therefore, from now on we assume w.l.o.g. that H is a diagonal matrix with non-increasing diagonal entries. Define A := SHS ⊤ .

By definition of Approx , we have

Moreover, for any k ∈ [ p ]

Therefore

and hence

We claim the following results which we will prove at the end of the proof.

Note that in claim (11) the inverse A -1 k exists almost surely since r ( H k : ∞ ) ≥ r ( H ) -k ≥ M by our assumption and S k : ∞ ∈ R M × ( d -k ) is a random gaussian projection onto R M . First part of Lemma C.1 follows immediately from combining claim (10) and (11).

To prove the second part of Lemma C.1, first note that with probability 1 -e -Ω( M ) we have

forc some constant c > 0 by Lemma G.2. Moreover, by the concentration of the Gaussian variance matrix (see e.g., Theorem 6.1 in [41]), we have S ⊤ 0: k S 0: k ⪰ I k / 5 with probability 1 -e -Ω( M ) when M/k ≥ 2 . Combining the last two arguments, we obtain

and therefore

with probability 1 -e -Ω( M ) . Combining Eq. (12) with the first part of Lemma C.1 completes the proof.

Proof of claim (10) Note that

where the last inequality uses the fact that the norm of projection matrices is no greater than one. Therefore, we have ∥ W ∥ 2 ≤ 1 . Now, it remains to show

as claim (10) is a direct consequence of Eq. (13) and the fact that ∥ W ∥ ≤ 1 .

By definition of W in Eq. (9), we have

By definition of V in Eq. (9), we have

Since S 0: k H 0: k S ⊤ 0: k + A k = A , it follows that

Proof of claim (11) It suffices to show U 2 + V ⊤ V = [ H -1 0: k + S ⊤ 0: k A -1 k S 0: k ] -1 . Using the definition of U in Eq. (9), we obtain

where the second line uses Woodbury's matrix identity, namely

Continuing the calculation of U , we have

Therefore,

Since

by Woodbury's matrix indentity, it follows from the definition of V in Eq. (9) that

Combining Eq. (14) and (15) yields

and therefore

C.2 A lower bound

For the approximation error Approx , we have the following result.

Lemma C.2 (Lower bound on the approximation error) . When r ( H ) ≥ M , under Assumption 1C, the approximation error in (4) (and 7) satisfies

where ( λ i ) d i =1 are eigenvalues of H in non-increasing order.

Proof of Lemma C.2. For any k ≤ d , following the proof of Lemma C.1, we have

and

Therefore and

where the last line uses the fact that E w ∗ ( w ∗ ) ⊗ 2 = I d . Using Eq. (13) and (16) in the proof of Lemma C.1, we further obtain

where A k := S k : ∞ H k : ∞ S ⊤ k : ∞ . For T 3 , we further have

where the second line is due to A ⪰ A k (and hence -A -1 ⪰ -A -1 k ), the third line follows from Von-Neuman's inequality. Since M := I d -k -H 1 / 2 k : ∞ S ⊤ k : ∞ A -1 k S k : ∞ H 1 / 2 k : ∞ is a projection matrix such that M 2 = M and tr( I d -k -M ) = M , it follows that M has M eigenvalues 0 and d -k -M eigenvalues 1 . Therefore, we further have

for any k ≤ d . Letting k = 0 maximizes the lower bound and concludes the proof.

C.3 A lower bound under Assumption 3

Lemma C.3 (Lower bound on the approximation error under Assumption 3) . Under Assumption 3, the approximation error in (4) (and 7) satisfies

where ( λ i ) d i =1 are eigenvalues of H in non-increasing order and the inequality hides some ( a, b ) -dependent constant.

Proof of Lemma C.3. The proof is essentially the same as the proof of Lemma C.2 but we include it here for completeness. Let H w := E [ w ∗ w ∗⊤ ] be the covariance of the prior on w ∗ . Following the proof of Lemma C.2, we have

where the last line uses Assumption 3 and notice that H , H w are both diagonal. Next, similar to the proof of Lemma C.2, using Eq. (13) and (16), we derive

where A k := S k : ∞ H k : ∞ S ⊤ k : ∞ . For ˜ T 3 , following the same argument for T 3 in the proof of Lemma C.2, we have

for any k ≤ d where the last inequality uses Assumption 3. Setting k = 0 maximizes the lower bound and concludes the proof.

C.4 Examples on matching bounds for Approx

In this section, we derive matching upper and lower bounds for Approx (defined in Eq. 4 and 7) in three concrete examples: power-law spectrum (Lemma C.4), power-law spectrum with source condition (Lemma C.5) and logarithmic power-law spectrum (Lemma D.5).

Lemma C.4 (Bounds on Approx under the power-law spectrum) . Suppose Assumption 1C and 2 hold. Then with probability at least 1 -e -Ω( M ) over the randomness of S

Here, the hidden constants only depend on the power-law degree a.

Proof of Lemma C.4. For the upper bound, by Lemma C.1 and noting E w ∗ 2 i = 1 for all i , we have with probability at least 1 -e -Ω( M )

for any given k 1 ≤ M/ 2 . Here the hidden constants depend on a . Therefore, letting k 1 = M/ 2 in the upper bound yields

with probability at least 1 -e -Ω( M ) .

For the lower bound, we have from Lemma C.2 that

This completes the proof.

Lemma C.5 (Bounds on Approx under the source condition) . Suppose Assumption 3 hold. Then with probability at least 1 -e -Ω( M ) over the randomness of S

Here, the hidden constants only depend on the power-law degrees a, b.

Proof of Lemma C.5. For the upper bound, by Lemma C.1 and noting E w ∗ 2 i ≂ i a -b for all i , we have with probability at least 1 -e -Ω( M )

for any given k 1 ≤ M/ 2 . Here the hidden constants depend on a, b . Moreover, choosing k 1 = M/ 2 in the upper bound gives

with probability at least 1 -e -Ω( M ) .

For the lower bound, we have from Lemma C.3 that

This completes the proof.

Lemma C.6 (Bounds on Approx under the logarithmic power-law spectrum) . Suppose Assumption 4 hold. Then with probability at least 1 -e -Ω( M ) over the randomness of S

Here, the hidden constants only depend on the power-law degree a.

Proof of Lemma C.6. For the upper bound, by Lemma C.1 and noting E w ∗ 2 i = 1 for all i , we have with probability at least 1 -e -Ω( M )

for any given k 1 ≤ M/ 2 , where the third line uses ∑ i>k 1 λ 2 i ≲ 1 / ( k 1 log 2 a k 1 ) . Choosing k 1 = M/ 2 , we obtain

with probability at least 1 -e -Ω( M ) . Here the hidden constants depend on a, b . For the lower bound, we have from Lemma C.2 that

Therefore, we have established matching upper and lower bounds for Approx .

D Bias error

In this section, we derive upper and lower bounds for Bias ( w ∗ ) defined in Eq. (5). Moreover, we show that the upper and lower bounds match up to constant factors in concrete examples.

D.1 An upper bound

Lemma D.1 (Upper bound on the bias term) . Suppose the initial stepsize γ ≤ 1 c tr( SHS ⊤ ) for some constant c > 1 . Then for any w ∗ ∈ H and k ∈ [ d ] such that r ( H ) ≥ k + M , the bias term in (5) satisfies

Moreover, for any k ≤ M/ 3 such that r ( H ) ≥ k + M , the bias term satisfies

with probability 1 -e -Ω( M ) , where A k := S k : ∞ H k : ∞ S ⊤ k : ∞ , { µ i ( A k ) } i M =1 denote the eigenvalues of A k in non-increasing order for some constant c > 1 .

Proof of Lemma D.1. Similar to the proof of Lemma C.1, we can without loss of generality assume the covariance matrix H = diag { λ 1 , λ 2 , . . . , λ d } where λ i ≥ λ j for any i ≥ j . Let SH 1 / 2 = ˜ U ( ˜ Λ 1 / 2 0 ) ˜ V ⊤ be the singular value decomposition of SHS ⊤ , where ˜ Λ := diag { ˜ λ 1 , ˜ λ 2 , . . . , ˜ λ d } is a diagonal matrix diagonal entries in non-increasing order. Define A k := S k : ∞ H k : ∞ S ⊤ k : ∞ . Then it follows from similar arguments as in Lemma C.1 that A k is invertible. Since

for some constant c > 1 by the stepsize assumption, it follows that I M -γ t SHS ⊤ ≻ 0 M for all t ∈ [ N ] . Therefore, it can be verified that

and by definition of Bias ( w ∗ ) in Eq. (5), we have

Note that the eigenvalues of M are { ˜ λ i (1 -γ ˜ λ i ) 2 N eff } i M =1 . Since the function f ( x ) = x (1 -γx ) 2 N eff is maximized at x 0 = 1 / [(2 N eff +1) γ ] for x ∈ [0 , 1 /γ ] with f ( x 0 ) ≲ 1 / ( N eff γ ) , it follows that ∥ M ∥ 2 ≤ c/ ( N eff γ ) (19)

for some constant c > 0 . The first part of Lemma D.1 follows immediately.

Now we prove the second part of Lemma D.1. Recall that v ∗ = ( SHS ⊤ ) -1 SHw ∗ . Substituting SH = ( S 0: k H 0: k S k : ∞ H k : ∞ ) into v ∗ , we obtain

where

We claim the following results which we prove later. With probability 1 -e -Ω( M )

Combining Eq. (22a), (22b) gives the second part of Lemma D.1.

Proof of claim (22a) By definition of T 1 , we have

Moreover,

for some constant c > 0 , where the last line uses Eq. (19). It remains to show

for some constant c > 0 with probability 1 -e -Ω( M ) . Since SHS ⊤ = S 0: k H 0: k S ⊤ 0: k + A k , we have

where the second line uses Woodbury's identity. Since

for some constant c > 0 .

it follows that

Therefore, with probability at least 1 -e -Ω( M )

where the last inequality follows from the fact that ∥ S 0: k ∥ 2 = √ ∥ S ⊤ 0: k S 0: k ∥ 2 ≤ c for some constant c > 0 when k ≤ M/ 2 with probability at least 1 -e -Ω( M ) . Since S 0: k is independent of A k and the distribution of S 0: k is rotationally invariant, we may write S ⊤ 0: k A -1 k S 0: k = ∑ M i =1 1 ˆ λ M -i ˜ s i ˜ s ⊤ i ,

where ˜ s i iid ∼ N (0 , I k /M ) and ( ˆ λ i ) i M =1 are eigenvalues of A k in non-increasing order. Therefore, for k ≤ M/ 3

for some constant c > 0 with probability at least 1 -e -Ω( M ) , where in the last line we again use the concentration properties of gaussian covariance matrices (see e.g., Theorem 6.1 in [41]). As a direct consequence, we have

with probability at least 1 -e -Ω( M ) for some constant c > 0 . This concludes the proof.

Proof of claim (22b) By definition of T 2 in Eq. (21), we have

where the last line follows from

D.2 A lower bound

Lemma D.2 (Lower bound on the bias term) . Suppose w ∗ follows some prior distribution and the initial stepsize γ ≤ 1 c tr( SHS ⊤ ) for some constant c > 2 . Let H w := E w ∗ w ∗⊤ . Then the bias term in Eq. (5) satisfies

almost surely, where M N := SHS ⊤ ( I -2 γ SHS ⊤ ) 2 N eff .

Proof of Lemma D.2. Adopt the notations in the proof of Lemma D.1. By definition of the bias term, we have

where the third line uses I M -2 γ t SHS ⊤ ≻ 0 M for all t ∈ [ N ] established in the proof of Lemma D.1, ∑ N i =1 γ i ≤ 2 γN eff , and the fact that (1 -w )(1 -v ) ≥ 1 -w -v for w,v > 0 . Substituting the definition of v ∗ in Eq. (5) into the expression, we obtain

and

where the last line uses Von Neumann's trace inequality. Continuing the calculation, we have

where the first inequality uses µ M + i -1 ( A ) = µ -1 i ( A -1 ) for any positive definite matrix A ∈ R M × M , and the second line follows from the definition of M N and the fact that (1 -λγN eff ) -2 N eff ≲ 1 when λ < 1 / ( γN eff ) .

D.3 Examples on matching bounds for Bias ( w ∗ )

In this section, we derive matching upper and lower bounds for Bias ( w ∗ ) in (5) in three scenarios: power-law spectrum (Lemma D.3), power-law spectrum with source condition (Lemma D.4) and logarithmic power-law spectrum (Lemma D.5). Recall that we define Bias := E w ∗ Bias ( w ∗ ) .

Lemma D.3 (Bounds on Bias under the power-law spectrum) . Suppose Assumption 1C and 2 hold and the initial stepsize γ ≤ 1 c tr( SHS ⊤ ) for some constant c > 2 . Then with probability at least 1 -e -Ω( M ) over the randomness of S

when ( N eff γ ) 1 /a ≤ M/c for some constant c > 0 . Here, all the (hidden) constants depend only on the power-law degree a.

Proof of Lemma D.3. For the upper bound, using Lemma G.5, D.1 and the assumption that E w ∗ 2 i = 1 for all i > 0 , with probability at least 1 -e -Ω( M ) , we have

where in the last inequality, we choose k 2 = [ M/ 3] ∧ ( N eff γ ) 1 /a to minimize the upper bound.

When ( N eff γ ) 1 /a ≤ M/ 3 , combining Lemma D.2 and G.4 gives the lower bound

and

-Ω( M )

with probability at least 1 -e . Here, the hidden constants depend only on a.

Lemma D.4 (Bounds on Bias under the source condition) . Suppose Assumption 3 hold and the initial stepsize γ ≤ 1 c tr( SHS ⊤ ) for some constant c > 2 . Then with probability at least 1 -e -Ω( M ) over the randomness of S

when ( N eff γ ) 1 /a ≤ M/c for some constant c > 0 . In both results, the hidden constants depend only on a, b.

Proof of Lemma D.4. For the upper bound, using Lemma G.5, D.1 and the assumption that (w.l.o.g.) E w ∗ 2 i ≂ i a -b for all i > 0 , with probability at least 1 -e -Ω( M ) , we have

when b < a +1 , where in the last inequality, we choose k 2 = [ M/ 3] ∧ ( N eff γ ) 1 /a to minimize the upper bound.

When ( N eff γ ) 1 /a ≤ M/c for some large constant c > 0 , combining Lemma D.2 and G.4 yields the lower bound

with probability at least 1 -e -Ω( M ) . Here, the hidden constants depend only on a, b.

Upper bound when b ≥ a +1 . Following the previous derivations, when b = a +1 , we have with probability at least 1 -e -Ω( M )

where the last line follows by setting k 2 = [ M/ 3] ∧ ( N eff γ ) 1 /a . When b > a +1 , we have with probability at least 1 -e -Ω( M )

where the last follows by chooing k 2 = M/ 3 to minimize the upper bound.

Note that there exist non-constant gaps between the upper and lower bounds on the bias term in the simple regime where b ≥ a +1 . We leave a more precise analysis of the bias term for future work.

LemmaD.5 (Bounds on Bias under the logarithmic power-law spectrum) . Suppose Assumption 4 hold and the initial stepsize γ ≤ 1 c tr( SHS ⊤ ) for some constant c > 2 . Let ¯ k := inf { k : k log a k ≥ N eff γ } . Then with probability at least 1 -e -Ω( M ) over the randomness of S

and

when ( N eff γ ) ≤ M c for some sufficiently small constant c > 0 . Here, all constants depend only on the power-law degree a.

Proof of Lemma D.5. For the upper bound, using Lemma G.7, D.1 and the assumption that E w ∗ 2 i = 1 for all i > 0 , with probability at least 1 -e -Ω( M ) , we have

where in the last inequality, we choose k 2 = [ M/ 3] ∧ [ ( N eff γ ) / log a ( N eff γ ) ] to minimize the upper bound.

Recall k ∗ ≂ M/ log M (for example we may define k ∗ = inf { k : k log k ≥ M } ) in Lemma G.6. Combining Lemma D.2 and G.6 gives the lower bound

with probability at least 1 -e -Ω( M ) for some constant c 1 > 0 . Here, the (hidden) constants depend only on a. Therefore, when ( N eff γ ) 1 /a ≤ M c for some sufficiently small constant c > 0 , we have

with probability at least 1 -e -Ω( M )

.

E Variance error

In this section, we present matching upper and lower bounds on the variance term Var defined in (6) under the power-law or logarithmic power-law spectrum.

Lemma E.1 (Matching bounds on Var under power-law spectrum) . Under Assumption 2, Var defined in Eq. (6) satisfies

with probability at least 1 -e -Ω( M ) over the randomness of S . Here, the hidden constants only depend on a .

Proof of Lemma E.1. By the definition of Var in Eq. (6) and Lemma G.4, we have

with probability at least 1 -e -Ω( M ) over the randomness of S . Here the hidden constants may depend on a.

Lemma E.2 (Matching bounds on Var under logarithmic power-law spectrum) . Under Assumption 4, Var defined in Eq. (6) satisfies

with probability at least 1 -e -Ω( M ) over the randomness of S , where ¯ k := inf { k : k log a k ≥ ( N eff γ ) } and ≂ hides constants that only depend on a .

Proof of Lemma E.2. Define k ∗ = inf { k : k log k ≥ M } and let ˜ D := # { ˜ λ j ≥ 1 / ( N eff γ ) } + ( N eff γ ) 2 ∑ ˜ λ j < 1 / ( N eff γ ) ˜ λ 2 j . By the definition of Var in Eq. (6) and Lemma G.6, we have

with probability at least 1 -e -Ω( M ) over the randomness of S , where the second line follows from

F Expected risk of the average of (SGD) iterates

In this section, we study the expected risk of the average of (SGD) iterates. Namely, we consider a fixed stepsize (SGD) procedure where γ t = γ and define ¯ v N := ∑ N -1 i =0 v i /N . Our goal is to derive matching upper and lower bounds R (¯ v N ) in terms of the sample size N and model size M.

Compared with the last iterate of (SGD) with geometrically decaying stepsizes, we show that the average of (SGD) iterates with a fixed stepsize achieves a better risk, in the sense that the effective sample size N eff is replaced by N in the bounds (c.f. Theorem 4.1). This may give improvement up to logarithmic factors.

We start with invoking the following result in [48].

Theorem F.1 (A variant of Theorem 2.1 and 2.2 in [48]) . Suppose Assumption 1 hold. Consider an M -dimensional sketched predictor trained by fixed stepsize (SGD) with N samples. Let ¯ v N := ∑ N -1 i =0 v i /N be the average of the iterations of SGD. Assume v 0 = 0 and σ 2 ≳ 1 . Conditional on S and suppose the stepsize γ < 1 / ( c tr( SHS ⊤ )) for some constant c > 0 , then there exist Approx , Bias , Var such that

where the expectation of R M is over w ∗ and ( x i , y i ) N i =1 and

and

where ( ˜ λ j ) j M =1 are eigenvalue of SHS ⊤ .

See Section F.2.1 for the proof.

For T i ( i = 1 , 2 , 3 , 4) , we also have the following upper (and lower) bounds.

Lemma F.2 (Lower bound on T 1 ) . Under the assumptions and notations in Theorem F.1, we have

almost surely, where ( ˜ λ i ) N i =1 are eigenvalues of SHS ⊤ in non-increasing order.

See the proof in Section F.2.2.

Lemma F.3 (Upper bound on T 2 ) . Under the assumptions and notations in Theorem F.1, for any k ≤ M/ 3 such that r ( H ) ≥ k + M , we have with probability at least 1 -e -Ω( M ) that

where A k := S k : ∞ H k : ∞ S ⊤ k : ∞ .

See the proof in Section F.2.3.

Lemma F.4 (Lower bound on T 3 ) . Under the assumptions and notations in Theorem F.1, we have

almost surely, where ( ˜ λ i ) i M =1 are eigenvalues of SHS ⊤ in non-increasing order.

See the proof in Section F.2.4.

Lemma F.5 (Upper bound on T 4 ) . Under the assumptions and notations in Theorem F.1 and assume r ( H ) ≥ M , we have

almost surely, where A k := S k : ∞ H k : ∞ S ⊤ k : ∞ .

See the proof in Section F.2.5.

With these results at hand, we are ready to derive upper and lower bounds for the risk of the average of (SGD) iterates.

F.1 Matching bounds for the average of (SGD) iterates under power-law spectrum

In this section, we derive upper and lower bounds for the expected risk under the power-law spectrum. Our main result (Theorem F.6) follows directly from Theorem F.1 and the bounds on T i ( i = 1 , 2 , 3 , 4) in Lemmas F.2 to F.5.

Theorem F.6 (Scaling law for average iterates of SGD) . Suppose Assumption 1 and 2 hold and σ 2 ≲ 1 . Then there exists some a -dependent constant c > 0 such that when γ ≤ c , with probability at least 1 -e -Ω( M ) over the randomness of the sketch matrix S , we have

where the expectation is over the randomness of w ∗ and ( x i , y i ) N i =1 , and Θ( · ) hides constants that may depend on a .

See the proof in Section F.2.6.

Compared with Theorem 4.1, Theorem F.6 suggests that the average of (SGD) achieves a smaller risk in the sketched linear model-the ( N eff γ ) 1 /a is replaced by ( Nγ ) 1 /a in the bound for the bias term. This is intuitive since the sum of stepsizes ∑ t γ t ≂ N eff γ for the geometrically decaying stepsize scheduler while ∑ t γ t ≂ Nγ for the fixed stepsize scheduler.

We also verify the observations in Theorem F.6 via simulations. We adopt the same model and setup as in Section 5 but use the average of iterates of fixed stepsize (SGD) (denoted by ¯ v N ) as the predictor. From Figure 3 and 4 we see that the expected risk E R (¯ v N ) also scales following a power-law relation in both sample size N and model size M . Moreover, the fitted exponents match our theoretical predictions in Theorem F.6.

F.2 Proofs

F.2.1 Proof of Theorem F.1

Similar to the proof of Theorem A.4, we have the decomposition

Note that ( v t ) N t =1 can also be viewed as the SGD iterates on the model y = ⟨ Sx , v ∗ ⟩ + ξ + ϵ , where the noise satisfies

Therefore, the upper and lower bounds on Bias , Var follow directly from the proof of Theorem 2.1, 2.2 and related lemmas (Lemma B.6, B.11, C.3, C.5) in [48].

10 2

10 3

10 4

Number of samples (N)

10 1

10 2

Model size (M)

a=1.5

1.1

1.2

1.3

1.4

1.5

Risk

10 2

10 3

10 4

Number of samples (N)

10 1

10 2

10 3

Model size (M)

a=2.0

1.05

1.10

1.15

1.20

1.25

1.30

Risk

10 2

10 3

10 4

10 5

Number of samples (N)

10

1

E

M

(

v

N

)

min

( )

a=1.5

k=-0.32

Risk

(a)

a

= 1

.

5

10 1

10 2

Model size (M)

2×10 1

3×10 1

4×10 1

E

M

(

v

N

)

min

( )

a=1.5

k=-0.50

Risk

(b)

a

= 1

.

5

10 2

10 3

10 4

10 5

Number of samples (N)

10

2

10

1

E

M

(

v

N

)

min

( )

a=2.0

k=-0.49

Risk

(c)

a

= 2

10 1

10 2

Model size (M)

10

1

E

M

(

v

N

)

min

( )

a=2.0

k=-1.00

Risk

(d)

a

= 2

F.2.2 Proof of Lemma F.2

Let f 1 ( A ) := ( I -( I -γA ) N/ 4 ) 2 A -1 /γ 2 /N 2 for any positive definite matrix A ∈ R M × M . Since γ ≤ 1 / ( c tr( SHS ⊤ )) , we have f 1 ( SHS ⊤ ) ⪰ 0 . By definition of T 1 and recalling v ∗ = ( SHS ⊤ ) -1 SHw ∗ , we have with probability at least 1 -e -Ω( M ) that

Following the proof of Lemma D.2 (by Von Neumann's trace inequality), we have

where the third inequality is due to

when λ < 1 / ( Nγ ) .

F.2.3 Proof of Lemma F.3

By definition of T 2 , the fact that 1 -x N = (1 -x ) ∑ N -1 i =0 x i , and recalling v ∗ = ( SHS ⊤ ) -1 SHw ∗ , we have

where f 2 ( A ) := [ ∑ N -1 i =0 ( I -γA ) i ] 2 /A/N 2 for any symmetric matrix A ∈ R M × M . Moreover, we have

Using the assumption on the stepsize that γ ≤ 1 / ( c tr( SHS ⊤ )) , we have

Combining Eq. (28) with Eq. (23) in the proof of Lemma D.1 (note that we assume k ≤ M/ 3 ), we obtain

for some constant c > 0 with probability at least 1 -e -Ω( M ) . For T 22 , we have

Since ∥ f 2 ( SHS ⊤ ) SHS ⊤ ∥ = ∥ [ ∑ N -1 i =0 ( I -γ SHS ⊤ ) i ] 2 /N 2 ∥ ≤ 1 by the assumption γ ≤ 1 / ( c tr( SHS ⊤ )) , and

it follows that T 22 ≤ ∥ w ∗ k : ∞ ∥ 2 H k : ∞ . Combining the bounds on T 21 , T 22 completes the proof.

F.2.4 Proof of Lemma F.4

Let f 3 ( A ) := ( I -( I -γ A ) N/ 4 ) /γ/N 2 for any positive definite matrix A ∈ R M × M . Following the same arguments as in the proof of Lemma F.2, we have f 3 ( SHS ⊤ ) ⪰ 0 and

Moreover,

where the third inequality is due to

when λ < 1 / ( Nγ ) . Note that

We thus have

Combining Eq. (31) and (29) completes the proof.

F.2.5 Proof of Lemma F.5

Substituting v ∗ = ( SHS ⊤ ) -1 SHw ∗ in the expression of T 4 and noting v 0 = 0 , we have

Since

by our assumption on the stepsize, it follows that

Combining Eq. (32) and (33) we find

F.2.6 Proof of Theorem F.6

First, by Lemma G.4 we have 1 / tr( SHS ⊤ ) ≳ c 1 for some a -dependent c 1 > 0 with probability at least 1 -e -Ω( M ) . Therefore we may choose c sufficiently small so that γ ≤ c implies γ ≲ 1 / tr( SHS ⊤ ) with probability at least 1 -e -Ω( M ) . Now, suppose we have γ ≲ 1 / tr( SHS ⊤ ) . Following the notations in Theorem F.1, we claim the following bounds on Approx , Bias , Var :

with probability at least 1 -e -Ω( M ) . Putting the bounds together yields Theorem F.6.

Proof of claim (34a) Note that our definition of Approx in Thereom F.1 is the same as that in Eq. (4) (and 7). Therefore the claim follows immediately from Lemma C.4.

Proof of claim (34b) This follows from the proof of Lemma E.1 with N eff replaced by N.

Proof of claim (34c) By Theorem F.1, Lemma F.3 and F.5, we have

with probability at least 1 -e -Ω( M ) for any k 2 ≤ M/ 3 . Choosing k 2 = min { M/ 3 , ( Nγ ) 1 /a } and using Lemma G.4 and claim (34b), we obtain

Bias ≲ max { M 1 -a , ( Nγ ) 1 /a -1 } + min { M, ( Nγ ) 1 /a } N ≲ max { M 1 -a , ( Nγ ) 1 /a -1 } +( Nγ ) 1 /a -≲ max { M 1 -a , ( Nγ ) 1 /a -1 } with probability at least 1 -e -Ω( M ) .

Proof of claim (34d) By Theorem F.1 and Lemma F.2, we have

When ( Nγ ) 1 /a ≤ M/c for some large constant c > 0 , we have from Lemma G.4 that

with probability at least 1 -e -Ω( M ) .

G Concentration lemmas

G.1 General concentration results

Lemma G.1. Suppose that S ∈ R M × d is such that 3

Let ( λ i ) i ≥ 1 be the eigenvalues of H in non-increasing order. Let ( ˜ λ i ) i M =1 be the eigenvalues of SHS ⊤ in non-increasing order. Then there exists a constant c > 1 such that for every M ≥ 0 and every 0 ≤ k ≤ M , with probability ≥ 1 -e -Ω( M ) , we have

As a direct consequence, for k ≤ M/c 2 , we have

where c 1 = c +2 c 2

Proof of Lemma G.1. We have the following decomposition motivated by Swartworth and Woodruff [38] (see their Section 3.4, Proof of Theorem 1).

We remark that this decomposition idea has been implicitly used in Bartlett et al. [5] to control the eigenvalues of a Gram matrix. In fact, we will use techniques from Bartlett et al. [5] to obtain a sharper bound than that presented in Swartworth and Woodruff [38].

For the upper bound, we have

where the last inequality is by Lemma G.2. For j ≤ k , using Lemma G.3, we have

Putting these together, we have the following for every j = 1 , . . . , M :

Similarly, we can show the lower bound. By the decomposition, we have

For k < j ≤ M , we have

where the last inequality is by Lemma G.2. For j ≤ k , using Lemma G.3, we have

For k < j ≤ M , we have

where the last inequality is due to λ j ≤ λ k for j ≥ k . Putting these together, we have

So far, we have proved the first claim. To show the second claim, we simply apply

and in the first claim.

Lemma G.2 (Tail concentration, Lemma 26 in Bartlett et al. [5]) . For any k ≥ 1 , with probability at least 1 -e -Ω( M ) , we have

Moreover, the minimum eigenvalue of S k : ∞ H k : ∞ S ⊤ k : ∞ satisfies

with probability at least 1 -e -Ω( M ) .

Proof of Lemma G.2. The first part of Lemma G.2 is a version of Lemma 26 in [5] (see their proof). We provide proof here for completeness.

We write S ∈ R M × p as

Since Gaussian distribution is rotational invariance, without loss of generality, we may assume

Then we have

Fixing a unit vector v ∈ R M , then

where each s ⊤ i v is (1 /M ) -subGaussian. By Bernstein's inequality, we have, with probability ≥ 1 -δ ,

By a union bound and net argument on S M -1 , we have, with probability ≥ 1 -δ , for every unit vector v ∈ R M ,

So with probability at least 1 -e -Ω( M ) , we have

which completes the proof of the first part of Lemma G.2.

To prove the second part of Lemma G.2, it suffices to note that

for some constant c > 1 with probability at least 1 -e -Ω( M ) , where the last line follows from concentration properties of Gaussian covariance matrices (see e.g., Thereom 6.1 [41]).

Lemma G.3 (Head concentration) . With probability at least 1 -e -Ω( M ) , we have

Proof of Lemma G.3. Note that the spectrum of S 0: k H 0: k S ⊤ 0: k is indentical to the spectrum of H 1 / 2 0: k S ⊤ 0: k S 0: k H 1 / 2 0: k . We will bound the latter. We start with bounding the spectrum of S 0: k S ⊤ 0: k . To this end, we write S ⊤ 0: k ∈ R k × M as

Then repeating the argument in Lemma G.2, we have, with probability ≥ 1 -δ , for every unit vector v ∈ R k ,

So we have, with probability ≥ 1 -e -Ω( M ) ,

This implies that

and that

We have completed the proof.

G.2 Concentration results under power-law spectrum

Lemma G.4 (Eigenvalues of SHS ⊤ under power-law spectrum) . Suppose Assumption 2 hold. There exist a -dependent constants c 2 > c 1 > 0 such that

with probability at least 1 -e -Ω( M ) .

Proof of Lemma G.4. Let ( ˜ λ i ) i M =1 denote the eigenvalues of SHS ⊤ in an non-increasing order. Using Lemma G.1 with k = M/c for some sufficiently large constant c and noting that ∑ i>k i -a ≂ k 1 -a , we have

for every j ∈ [ M ] for some constants ˜ c i , i ∈ [2] with probability at least 1 -e -Ω( M ) . Therefore, for all j ≤ M/ ˜ c for some sufficiently large constant ˜ c > 1 , we have

with probability at least 1 -e -Ω( M ) for some constants ˜ c 3 , ˜ c 4 > 0 . For j ∈ [ M/ ˜ c, M ] , by monotonicity of the eigenvalues, we have

for some sufficiently large constant ˜ c 5 > ˜ c 4 with probability at least 1 -e -Ω( M ) . Moreover, using Lemma G.2 with k = 0 , we obtain

with probability at least 1 -e -Ω( M ) for some constants ˜ c 6 , ˜ c 7 , ˜ c 8 > 0 . Combining the bounds for j ≤ M/ ˜ c and j ∈ [ M/ ˜ c, M ] completes the proof.

Lemma G.5 (Ratio of eigenvalues of S k : ∞ H k : ∞ S ⊤ k : ∞ under power-law spectrum) . Suppose Assumption 2 hold. There exists some a -dependent constant c > 0 such that for any k ≥ 1 , the ratio between the M/ 2 -th and M -th eigenvalues

with probability at least 1 -e -Ω( M ) .

Proof of Lemma G.5. We prove the lemma under two scenarios where k is relatively small (or large) compared with M .

Let c > 0 be some sufficiently large constant. Applying Lemma G.1 with H k : ∞ replacing H , for k 0 = M/c , we have

with probability at least 1 -e -Ω( M ) for some constant c 1 > 0 .

Case 1 : k ≲ M From Lemma G.2, we have

with probability at least 1 -e -Ω( M ) . Therefore

with probability at least 1 -e -Ω( M ) when k/M ≲ 1 .

Case 2 : k ≳ M On the other hand, when k is relatively large, using Lemma G.1 with H k : ∞ replacing H again, we obtain

with probability at least 1 -e -Ω( M ) , where c 1 , c 2 , c 3 > 0 are some universal constants. Choosing k 0 = M/c 2 for some sufficiently large constant c > 0 , we further obtain

with probability at least 1 -e -Ω( M ) , where ( c i ) 7 i =4 are a -dependent constants. Since

when k is large, i.e., k/M > ˜ c for some sufficiently large a -dependent constant ˜ c > 0 that may depend on ( c i ) 7 i =1 , we have from Eq. (35) and (36) that

with probability at least 1 -e -Ω( M ) .

G.3 Concentration results under logarithmic power-law spectrum

Lemma G.6 (Proof of Theorem 6 in [5]) . Suppose Assumption 4 hold. Then there exist some a -dependent constants c, ˜ c > 0 such that, with probability at least 1 -e -Ω( M )

where k ∗ ≂ M/ log( M ) . Also, there exists some a -dependent constants c 1 , c 2 > 0 such that

with probability at least 1 -e -Ω( M ) .

Proof of Lemma G.6. The proof is adapted from the proof of Theorem 6 in [5]. We include it here for completeness.

First part of Lemma G.6. In Lemma G.1, for some constant c > 1 , choose

Then with probability ≥ 1 -e -Ω( M ) , we have:

where c 1 > 1 is a constant.

When λ j ≂ j -1 log -a ( j +1) , we have

and

Therefore, we have where k ∗ ≂ M/ log( M ) .

Second part of Lemma G.6. Let ¯ λ i denote the i -th eigenvalue of SH 2 S ⊤ for i ∈ [ M ] . Using Lemma G.1 with k = M/c for some sufficiently large constant c 0 and noting that ∑ i>k λ 2 i ≂ ∑ i>k i -2 log -2 a ( i +1) ≲ k -1 log -2 a k , we have

for every j ∈ [ M ] for some constants ˜ c i , i ∈ [2] with probability at least 1 -e -Ω( M ) . Therefore, for all j ≤ M/ ˜ c for some sufficiently large constant ˜ c > 1 , we have

with probability at least 1 -e -Ω( M ) for some constants ˜ c 3 , ˜ c 4 > 0 . For j ∈ [ M/ ˜ c, M ] , by monotonicity of the eigenvalues, we have

for some constants c 5 , c 6 > 0 with probability at least 1 -e -Ω( M ) . Moreover, using Lemma G.2 with k = 0 , we obtain

with probability at least 1 -e -Ω( M ) for some constants ˜ c 7 , ˜ c 8 > 0 when j ∈ [ M/ ˜ c, M ] . Combining the bounds for j ≤ M/ ˜ c and j ∈ [ M/ ˜ c, M ] completes the proof.

Lemma G.7 (Ratio of eigenvalues of S k : ∞ H k : ∞ S ⊤ k : ∞ under logarithmic power-law spectrum) . Suppose Assumption 4 hold. There exists some a -dependent constant c > 0 such that for any k ≥ 1 , the ratio between the M/ 2 -th and M -th eigenvalues

with probability at least 1 -e -Ω( M ) .

Proof of Lemma G.7. Similar to the proof of Lemma G.5, we prove the lemma under two scenarios where k is relatively small (or large) compared with M .

Let c > 0 be some sufficiently large constant. Applying Lemma G.1 with H k : ∞ replacing H , for k 0 = M/c , we have

with probability at least 1 -e -Ω( M ) for some constant c 1 > 0 .

with probability at least 1 -e -Ω( M ) . Therefore,

with probability at least 1 -e -Ω( M ) when k/M ≲ 1 .

Case 2 : k ≳ M . On the other hand, when k is relatively large, using Lemma G.1 with H k : ∞ replacing H and k 0 = M/c again, we obtain

with probability at least 1 -e -Ω( M ) , where c 1 , c 2 , c 3 > 0 are some a -dependent constants. Therefore,

with probability at least 1 -e -Ω( M ) when k/M ≳ 1 .

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 claims are accurate and reflect the contributions and scope of the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

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

Answer: [Yes]

Justification: We discuss the limitations of our results.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

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

Answer: [Yes]

Justification: We provide all the details for our theorems.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main 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: We provide the details for our experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

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: [No]

Justification: Our experiments are numerical simulations and easy to reproduce.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

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: We provide the details for our experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

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

Answer: [Yes]

Justification: We provide error bars.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

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: Our experiments are numerical simulations and can run with light compute resources.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

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 conform with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

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

Answer: [NA]

Justification: Our paper is theoretical and we do not expect direct societal impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

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: Our paper is theoretical and we do not expect risk of such.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

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

Answer: [NA]

Justification: Our paper is theoretical and does not use existing assets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

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

Answer: [NA]

Justification: Our paper is theoretical and does not create new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

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

Answer: [NA]

Justification: Our paper is theoretical and does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals 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: Our paper is theoretical and does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nLarge models often empirically satisfy the neural scaling law i.e. the error decreases when the number of data and the model size increase, however, this contradicts with the widely accepted belief in learning theory that the variance error (one of the decomposed errors) should increase with the model size. Authors investigate this in the infinite-dimensional linear regression setting and show that the variance error is dominated by other errors because of the implicit regularization of SGD. Authors also support their claim with numerical simulations.\n\nSTRENGTHS:\n- The authors reconcile empirically observed scaling law and the classical statistical learning theory, and offers a novel viewpoint to understand and explain their original \"conflictions\", which is novel and should be interesting for both learning theory and LLM community. \n- The authors offer a general framework in section 6 to incorporate the general spectrum case of data covariance matrix, and it gives a novel technical innovations to sharply control the effect of data sketch.\n- The paper is very well written and it is easy to understand\n\nWEAKNESSES:\n- The assumptions of the paper might be too strong with linear regression setting and Gaussian design assumption, is it possible to extend to the kernel setting and with relaxed assumption on features [1]?\n- While the current analysis is focusing on one-pass SGD, which is not often used in LLM training, can the result be extended to Adam?\n- The result is not very surprising \n\n[1] Barzilai et al.,  Generalization in Kernel Regression Under Realistic Assumptions\n\nQUESTIONS:\n- I am confused about the sketch matrix setting? Is the fixed sketch matrix analogous to the neural net with size M and why the teacher-student model in previous work can be viewed as sketched linear regression model\n- I am also confused about the Hypercontractivity assumption in Assumption 5, what's its intuition and why it is required in the proof?\n- What's the purpose for setting geometric decaying stepsize, what happens if it is fixed or it is polynomial decaying?\n- Is the regularization effect of SGD determined by the choice of stepsize, where the optimal choice controls the strength of the implicit bias of SGD?\n\nLIMITATIONS:\nSee weakness", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work examines neural scaling laws, in the simplified setting of linear regression trained by one-pass SGD. In particular, it attempts to explain the apparent mismatch between the statistical theory on one hand, which predicts that variance error increases with model size, and neural scaling laws on the other, which state that predictive performance should increase with model size. By considering sketched covariates of an infinite dimensional linear regression model with Gaussian prior under mild spectral assumptions, scaling laws are derived while keeping track of a variance term that increases with model complexity and corrects for a mismatch between the best-in-class predictor and the best possible predictor. However, this term is of higher order than the usual approximation and bias terms appearing in scaling laws, accounting for the empirical accuracy of scaling laws in practice. Numerical examples are provided for synthetic data.\n\nSTRENGTHS:\nThis is a good paper that addresses an important theoretical inconsistency, and provides progress towards more general results.\nThe paper is generally well written and communicated.\n\nWEAKNESSES:\nIt is not clear when the variance term is first given on page 2 that it is of higher order than the other terms. The reader has to find the more precise presentation in Theorem 4.1 in order to verify this for themselves. Clarification of this when the term is first introduced would be helpful and improve the paper.\nWhile the paper's contribution is primarily theoretical, the only examples given are synthetic. Validation of the theory on a large benchmark dataset, even if modified to make it appropriate to the assumptions, would strengthen the paper significantly.\n\nQUESTIONS:\nCan you please address the weaknesses outlined above?\nHow do you think this work can this be generalized to other settings of interest?\nUnder which settings should the variance term be accounted for?\n\nLIMITATIONS:\nSome discussion of the assumptions of their work and their limitations would be helpful. Broader societal impact is mostly not applicable.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

comments: SUMMARY:
This paper sheds new light on the bounds of linear regression by framing them in terms of scaling laws.

STRENGTHS:
- The set-up is very clear, the story is well-narrated and the model, including some sketch matrix to play the role of the model-size, is well-thought.

- The authors try to be exhaustive with the assumptions: by including in a second time, (i) the assumptions that make everything work (Assumptions A and B) and that allow to go beyond Gaussian data ; (ii) source condition, (iii) some other spectrum decay.

WEAKNESSES:
It is difficult to talk about a weakness: the paper is self-contained, the results are solid and the context very well explained. 

Perhaps the authors should put emphasis on the fact that there is **absolutely no technical novelty and that the contribution relies in re-framing known bounds under the perspective of scaling laws**.

QUESTIONS:
No question

LIMITATIONS:
No
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nMotivated by the recent neural scaling law literature, this work investigates the generalization error rate of a sketched linear regression model trained on Gaussian covariates with power-law spectrum under one-pass SGD.\n\nThe main result are lower and upper bounds with matching rates, providing a detailed characterization of the different contributions to the generalization error.\n\nSTRENGTHS:\nThe paper is well-written and easy to parse. The topic is of interest to the NeurIPS community, and the contribution is timely.\n\nWEAKNESSES:\nIn my reading, the two main weaknesses are:\n1. The lack of a wider context. While it makes precise connections to other recent works on the \"*\"theory of neural scaling laws\"*\" such as [Bahri et al., 2021, Maloney et al., 2022, Bordelon et al., 2024], it ignores that the study of generalization error rates in least-squares / ridge regression under source and capacity conditions is a classical topic with an extensive literature. For instance:\n\n     - [1,2,3] have studied both single and multiple pass SGD for least squares regression on Hilbert spaces under source and capacity conditions. This corresponds to the kernel limit $M\\gg N$ in this work. Borrowing the notation from [3] for concreteness, after a quick comparison with the identification of the capacity $\\alpha = a$ and source $r=\\frac{b-1}{2a}$, I believe the \"hard\" and \"easy\" rates from Theorems 4.1 and 4.2 can be retrieved from the one-pass rates in [3].\n     - [4] extends the discussion of [1-3] to the random features approximation of kernels.\n     - The rates from one-pass SGD on least squares can be identified with the rates of ridge regression with an appropriately chosen regularization, which has been exactly characterized in [7]. For example, the blue and orange regions in Figure 1 from [7] correspond to rates in Theorem 4.2 with $\\ell = 1$.\n     - A recent extension of the above to RFs case appeared in [6] (though note some of these rates were known from [7]).\n\n     From the discussion above, it is not really clear what rates are new in this submission. I am willing to change my evaluation if this can be clarified with an in-depth comparison with these works.  \n\n2. The organization of the manuscript in Theorems with \"settings of increasing generality\" can come through as deceptive. Aren't Theorems 4.3, 4.2 and 4.1 corollaries of Theorems 6.3 and 6.4?\n\n\n**References**\n\n- [1] Yao, Y., Rosasco, L. & Caponnetto, A. On Early Stopping in Gradient Descent Learning. Constr Approx 26, 289–315 (2007). https://doi.org/10.1007/s00365-006-0663-2\n- [2] Ying, Y., & Pontil, M. (2008). Online gradient descent learning algorithms. Foundations of Computational Mathematics, 8, 561-596.\n- [3] Pillaud-Vivien, L., Rudi, A., & Bach, F. (2018). Statistical optimality of stochastic gradient descent on hard learning problems\nthrough multiple passes. Advances in Neural Information Processing Systems, 31.\n- [4] Luigi Carratino, Alessandro Rudi, Lorenzo Rosasco. Learning with SGD and Random Features.  Part of Advances in Neural Information Processing Systems 31 (NeurIPS 2018).\n- [5] Hugo Cui, Bruno Loureiro, Florent Krzakala, Lenka Zdeborová. Generalization Error Rates in Kernel Regression: The Crossover from the Noiseless to Noisy Regime. Part of Advances in Neural Information Processing Systems 34 (NeurIPS 2021).\n- [6] Leonardo Defilippis, Bruno Loureiro, Theodor Misiakiewicz. Dimension-free deterministic equivalents for random feature regression. arXiv:2405.15699.\n- [7] Alessandro Rudi and Lorenzo Rosasco. Generalization properties of learning with random features. In Advances in Neural Information Processing Systems, pages 3215–3225, 2017.\n\nQUESTIONS:\n- L43-48:\n\n> \"*This difference must be reconciled, otherwise, the statistical learning theory and the empirical scaling law make conflic predictions: as the model size increases, the theoretical bound (2) predicts an increase of variance error that eventually causes an increase of the population risk, but the neural scaling law (1) predicts a decrease of the population risk. In other words, it remains unclear when to follow the prediction of the empirical scaling law (1) and when to follow that of the statistical learning bound (2).*\"\n\nGiven our current understanding of random features regression, this claim is inflated. Several of the works mentioned above [4,6,7], as well as many other works in the random features literature have made the observation that increasing the \"width\" $M$ does not necessarily lead to an increase of variance / overfitting, thanks to the implicit regularization.   \n\n\n- L273-275:\n\n> \"*It is noteworthy that our simulations demonstrate stronger observations than the theoretical results in Theorem 4.1, which only establishes matching upper and lower bounds up to a constant factor.*\"\n\nWhat do you exactly mean in this sentence? What in Figure 2 is stronger than the exponents reported in Theorem 4.1?\n\nLIMITATIONS:\nThe assumptions under which the Theorems hold are clearly stated, but I miss an explicit discussion of the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}

comments: SUMMARY:
This paper proves that the scaling laws that occur in deep learning also occur when solving linear regression with SGD.

Namely, suppose:
* that covariates x are Gaussian with  a power-law spectrum
* the true labels are given by <w^*, x> for unknown w^*
* we run one-pass SGD to do linear regression on an M-dimensional sketch of the covariates x
* we have N data points

Then the risk is upper-bounded by an expression that depends only on the "bias" and "approximation" terms, with inverse polynomial dependence as observed in neural scaling laws. The "variance" term is of lower order (provided the appropriate large step sizes are used).

Extensions to when the data follows a logarithmic power spectrum are also explored.

STRENGTHS:
The result is timely and I think will be of high interest to the community, since it sheds insight on why neural scaling laws might occur in practice.

The paper is well-presented, and the proofs are neatly written and easy to follow.

As opposed to previous works, this paper provides risk bounds for any network width M, and any number of data points N, instead of taking one of the parameters to infinity and deriving a scaling law in the other parameter.

WEAKNESSES:
Typo: "optimally tunned"

QUESTIONS:
1. How dependent are the results on a geometric step-size schedule? What if a constant schedule, or schedule decreasing as gamma_t = t^{-c} is used instead?

2. In line 232, it is stated that when $1 \leq b \leq a$ the tasks are "relatively hard", and in line 236 that when $a < b < a+1$ the tasks are "relatively easy". Is what is meant here that these tasks are "relatively harder/easier" than when a = b?

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-17