When does preconditioning help or hurt generalization?
Authors: Shun-ichi Amari
Jimmy Ba
Roger Baker Grosse
Xuechen Li
Atsushi Nitanda
Taiji Suzuki
Denny Wu
Ji Xu
Venue: ICLR 2021 Poster (2021)
Keywords: generalization
second-order optimization
natural gradient descent
high-dimensional asymptotics

Abstract: While second order optimizers such as natural gradient descent (NGD) often speed up optimization, their effect on generalization has been called into question. This work presents a more nuanced view on how the \textit{implicit bias} of optimizers affects the comparison of generalization properties. 
We provide an exact asymptotic bias-variance decomposition of the generalization error of preconditioned ridgeless regression in the overparameterized regime, and consider the inverse population Fisher information matrix (used in NGD) as a particular example. We determine the optimal preconditioner $\boldsymbol{P}$ for both the bias and variance, and find that the relative generalization performance of different optimizers depends on label noise and ``shape'' of the signal (true parameters): when the labels are noisy, the model is misspecified, or the signal is misaligned with the features, NGD can achieve lower risk; conversely, GD generalizes better under clean labels, a well-specified model, or aligned signal. 
Based on this analysis, we discuss several approaches to manage the bias-variance tradeoff, and the potential benefit of interpolating between first- and second-order updates. We then extend our analysis to regression in the reproducing kernel Hilbert space and demonstrate that preconditioning can lead to more efficient decrease in the population risk. Lastly, we empirically compare the generalization error of first- and second-order optimizers in neural network experiments, and observe robust trends matching our theoretical analysis. 

Published as a conference paper at ICLR 2021
WHEN DOES PRECONDITIONING HELP OR HURT GEN- ERALIZATION?
"Shun-ichi Amari', Jimmy 2, Roger Grosse233 Xuechen Li', Atsushi Nitanda Taiji Suzuki Denny 2 Xu'
RIKEN CBS, University of Toronto, Vector Institute, *Google Research, Brain Team, 6 University of Tokyo, RIKEN AIP, 'Columbia University
amari@brainn riken. jp. {iba, rgrosse, lxuechen, dennywu 8cs toronto.ooo..................... eddu, (nitanda, taiji]mmistt u-tokyo, ac. ip, jixuêcs. columbia.
ABSTRACT
While second order optimizers such natural gradient descent (NGD) often speed up optimization, their effect generalization has been called into question. This work presents more nuanced view on how the implicit bias of optimizers affects the comparison of generalization properties. We provide an exact asymptotic bias- variance decomposition of the generalization error of preconditioned ridgeless re- gression in the overparameterized regime, and consider the inverse population Fisher information matrix (used in NGD) as a particular example. We determine the optimal preconditioner P for both the bias and variance, and find that the rela- tive generalization performance of different optimizers depends on label noise and "shape" of the signal (true parameters): when the labels are noisy, the model is misspecified, or the signal is misaligned with the features, NGD can achieve lower risk; conversely, GD generalizes better under clean labels, well speci-ee model, or aligned signal. Based on this analysis, we discuss several approaches to man- age the bias-variance tradeoff, and the potential benefit interpolating between first- and second-order updates. We then extend our analysis to regression in the reproducing kernel Hilbert space and demonstrate that preconditioning can lead to more efficient decrease in the population risk. Lastly, empirically compare the generalization error of first- and second-order optimizers in neural network experiments, and observe robust trends matching our theoretical analysis.
INTRODUCTION
We study the generalization property of an estimator â obtained by minimizing the empirical risk (or the training error) L(fe) ) via a preconditioned gradient update with preconditioner P:
0++1
= 0,
(1.1)
Setting P I recovers gradient descent (GD). Choices of P which exploit second-order informa- tion include the inverse Fisher information matrix, which gives the natural gradient descent (NGD) (Amari, 1998); the inverse Hessian, which leads to Newton's method (LeCun et al., 2012); and diagonal matrices estimated from past gradients, which include various adaptive gradient meth- ods (Duchi et al., 2011; Kingma & Ba, 2014). These preconditioners often alleviate the effect of pathological curvature and speed up optimization, but their impact on generalization has been under debate: Wilson et al. (2017) reported that in neural network optimization, adaptive second- order methods generalize worse compared to gradient descent (GD), whereas other empirical studies showed that second-order methods achieve comparable, if not better generalization (Xu et al., 2020).
The generalization property of optimizers relates to the discussion of implicit bias (Gunasekar et al., 201 8a), i.e. preconditioning may lead to different converged solution (with potentially the same training loss), as illustrated in Figure 1. While many explanations have been proposed, our starting point is the well-known observation that GD often implicitly regularizes the parameter & norm. For instance in overparameterized least squares regression, GD and many first-order methods find the minimum E norm solution from zero initialization (without explicit regularization). but precondi- tioned updates may not. This being said, while the minimum & norm solution can generalize well
Alphhbbttical ordering. Correspondence to: Denny Wu (dennywu8cs toronto.

---

Published as a conference paper at ICLR 2021
in the overparameterized regime (Bartlett al., 2019). it is unclear whether preconditioning leads to inferior solutions- even in the simple setting of overparameterized linear regression, quantitative
understanding of how preconditioning affects generalization is largely lacking.
Motivated by the observations above, in Section we start with overparameterized least squares regression (unregularized) and analyze the stationary solution (t -> oo) of update (1.1) under
label (model outpur)
time-invariant preconditioner. Extending previous analysis in the proportional limit (Hastie et al., 2019), we consider more general random design setting and derive the exact population risk in its bias-variance decomposition. We characterize the op- timal P within a general class of preconditioners for both the bias and variance, and focus on the comparison between GD, for which P is identity, and NGD, for which P the inverse population Fisher information matrix'. We find that the com- parison of generalization is affected by the following factors:
0.0
-0.2
NGD
(s2t l2aton
0atiing 03ta
Input
Figure 1: ID illustration of different implicit biases: two-layer sigmoid net- work trained with preconditioned GD,
1. Label Noise: Additive noise in the labels leads to the variance term in the risk. We prove that NGD achieves the optimal variance among a general class of preconditioned updates.
2. Model Misspecification: Under misspecification, there does not exist perfect fe that recovers the true function (target). We argue that this factor similar to additional label noise, and thus NGD may also be beneficial when the model is misspecified.
Data--iggal-Aligammentt Alignment describes how the target signal distributes among the in- put features. We show that GD achieves lower bias when signal is isotropic, whereas NGD is preferred under "misalignment" - when the target function focuses on small feature directions.
Beyond the decomposition of stationary risk, our findings in Section 4 and 5 are summarized as:
In Section 4.1 and 4.2 we discuss how the bias-variance tradeoff can be realized by different choices of preconditioner P (e.g. interpolating between GD and NGD) or early stopping.
In Section 4.3 we extend our analysis to regression the RKHS and show that under early stop- ping, preconditioned update interpolating between GD and NGD achieves minimax optimal convergence rate much fewer steps, and thus reduces the population risk faster than GD.
In Section 5 empirically test how well our findings in linear model carry over to neural net- works: under student-teaches setup, we compare the generalization of GD with preconditioned updates and illustrate the influence of all aforementioned factors. The performance of neural net- works under a variety of manipulations results in trends that align with our theoretical analysis.
2 BACKGROUND AND RELATED WORKS
Natural Gradient Descent. NGD is second-order method originally proposed in Amari (1997).
Consider a data distribution p(æ) on the space X, a function fo 2' - parameterized by 0, and a loss function L(X, fo) 1
n Ei [[,, fe(z:)), where 1 y Z + R. Also suppose a probability distribution p(y|z) = P(u(fe(z) is defined on the space of labels. Then, the natural gradient is defined as: VoL(X, fa) = F F- VLL fo), where F = E o log v/0) p(z,( y ()) T] the Fisher information matrix, or simply the (population) Fisher. Note that expectations in the Fisher are under the joint distribution of the model p(ær, y 0) = p(æ))((y) ʃo (xx)) In the literature, the Fisher is sometimes defined under the empirical data distribution { [x, ) (Amari et al., 2000). We instead refer to this quantity as the sample Fisher, the properties of which influence optimization and have been studied Karakida et al. (2018); Kunstner et al. (2019); Thomas et al. (2020). We remark that in linear and kernel regression under squared loss, sample Fisher-based updates give the same stationary solution as GD (see Section 3), whereas population Fisher-based update may not.
While the population Fisher is typically difficult to obtain, extra unlabeled data can be used in its estimation, which empirically improves generalization (Pascanu & Bengio, 2013). Moreover, under structural assumptions, parametric approaches to estimate F can be more sample-efficient (Martens & Grosse, 2015; Ollivier, 2015), and thus closing the gap between sample and population Fisher.
From now use NGD to denote the population Fisher- update, and we write "sample NGD" when P is the inverse or pseudo-inverse of the sample Fisher; see Section for discussion.
2

---

Published as a conference paper at ICLR 2021
When the per-instance loss is the negative log- probability of exponential family, the sample Fisher coincides with the generalized which is the focus of this work, the quantity also coincides with the Hessian. We thus take NGD as a representative example of preconditioned update, and we expect our findings to also translate to other second-order methods (not including adaptive gradient methods) in regression problems.
matrix (Martens, 2014). In least squares regression,
Analysis of Preconditioned Gradient Descent. While Wilson et al. (2017) outlined one example under fixed training data where GD generalizes better than adaptive methods, in the online learning setting, for which optimization speed relates to generalization, several works have shown the advan- tage of preconditioning (Levy & Duchi, 2019; Zhang et al., 2019a). In addition. Zhang et al. (2019b); Cai et al. (2019) established convergence and generalization guarantees of sample Fisher-based up- dates for neural networks in the kernel regime. Lastly, the generalization of different optimizers relates to the notion of "sharpness" (Keskar et al., 2016; Dinh et al., 2017), and it has been argued that second-order updates tend to find sharper minima (Wu et al., 2018).
We note that two concurrent works also discussed the generalization performance of preconditioned updates. Wadia et al. (2020) connected second-order methods with data whitening in linear models, and qualitatively showed that whitening (thus second-order update) harms generalization in certain cases. Vaswani et al. (2020) analyzed the complexity of the maximum P-margin solution in linear classification problems. We emphasize that instead of upper bounding the risk (e.g. Rademacher complexity), which may not decide the optimal P for generalization error, we compute the exact risk for least squares regression, which allows us to precisely compare different preconditioners.
ASYMPTOTIC RISK OF RIDGELESS INTERPOLANTS
In this section consider the following setup: given n training samples {.;} labeled by teacher model (target function) f" - with additive noise: y/ (a;) + 6i. we learn linear (v: student model fe by minimizing the squared loss: L(X, = Σ. 1 x T))² We assume a random design: æ, Σx²² z;. where 2; R is i.i.d. vector with zero-mean, unit-variance, and finite 12th moment, and ε is i.i.d. noise independent
with mean 0 and variance o². Our goal is to compute the population risk R(ʃ) = [0 (w) - [x)) in the proportional asymptotic limit:
(A1) Overparameteriized Proportional Limit: n, d oc, d/n + 7 (1,
(A1) entails that the number of features (or parameters) is larger than the number of samples, and there exist multiple empirical risk minimizers with potentially different generalization properties.
Denote X [ TT] E Ra the data matrix and E Rn the corresponding label vector. We optimize the parameters 0 via a preconditioned gradient flow with preconditioner P(I) Rdxd
a0(t) ØL
P(I) aL(O()) - P( P(X x y XO(I), 0(0) = 0.
(3.1)
(0(1) Y1
In this linear setup, many common choices of preconditioner do not change through time: under Gaussian likelihood, the sample Fisher (and also Hessian) corresponds to the sample covariance X X/n up to variance scaling, whereas the population Fisher corresponds to the population covariance = Σx- We thus limit our analysis to fixed preconditioner of the form P(t) =: P.
#box 11
y=X0
initialization
Precond.
Write parameters at time under update (3.1) with fixed P as (t()). For positive definite P. the stationary solution is given as: 0p:=limg (1) PX (XPX - One may check that discrete time gradient descent update (with appropriate step size) and other variants that do not alter the span of gradient (c.g. stochas- tic gradient momentum) converge to the same solution well.
Geometric illustration (2D) how the interpolating O depends on the preconditioner.
Figure 2:
Intuitively speaking, if the data distribution (blue contour in Figure 2) is not isotropic, then certain directions will be more "important" than others. In this case uniform la shrinkage (which GD implicitly provides) may not be most desirable, and certain P that takes data geometry into account may lead to better generalization instead. The above intuition will be made rigorous in this section.
3

---

Published as a conference paper at ICLR 2021
Remark. 0p is the minimum ||0|| norm interpolant: 0p arg min 0 s.t. y for positive definite P. For GD this translates to the parameter l2 norm, whereas for NGD (P=F-1), the implicit bias is the x10||x norm. Since (x) = Ex NGD finds an interpolating function with smallest norm under the data distribution. We empirically observe this divide between small purameter norm and function norm in neural networks as well (see Figure I and Appendix A.
We highlight the following choices of P and the corresponding stationary solution Op as t DC.
Identity:
P= d recovers GD that converges to the min lz norm interpolant (also true for momentum GD and SGD), which we write as 0 == X
0.75
6D
0.70
NGD Sample NGD
[X x ) y and refer to as the GD solution.
0.65
Population Fisher: P === F = Σx leads to the estimator
0.60
0--1, which we refer to as the NGD solution.
0.55
0.50
Sample Fisher: since the sample Fisher is rank-deficient, we may add a damping term P (X' X X+ - or take the pseudo-inverse P = (x X)i. In both cases, the gradient is still spanned by X, and thus the update finds the same min lg- norm solution 0, (also true for full-matrix Adagrad (Agarwal et al., 2018)), although the trajectory differs (see Figure 3).
0.45
Umme
Figure 3: Population risk of preconditioned linear regression vs. time with the following P: I (red), Ex (blue) and (X X)1 (cyan). Time rescaled differ- ently for each curve (convergence speed is not comparable). Ob- serve that GD and sample NGD give the same stationary risk.
Remark. The above choices reveal a gap between sample- and population-based P: while the sample Fisher accelerates opti- mization (Zhang et al., 2019b), the following sections demonstrate generalization properties only possessed by the population Fisher.
We compare the population risk of the GD solution 0 ı and NGD solution 0 1 in its bias-variance decomposition w.r.t. label noise (Hastie et al., 2019) and discuss the two components separately,
R( = Exe [0 - (z,, E. (0))) + (r(Co)(0)Ex)
BB)) bias
V(0), variance
Note that the bias does not depend on label noise E, and the variance does not depend on the teacher model f Additionally, given that f" can be independently decomposed into a linear component on features a and a residual: I (ar) = (ær, 0²) + f((x)) we can separate the bias term into well- specified component ||0" |0l| 2 x which captures the difficulty in learning 0", and misspecified component, which corresponds to the error due to fitting f (beyond what the student can represent).
3.1 THE VARIANCE TERM: NGD IS OPTIMAL
We first characterize the stationary variance which is independent to the teacher model S", We restrict ourselves to preconditioners satisfying the following assumption on the spectral distribution:
(A2) Converging Eigenvalues: P is positive definite and as n, d + co, the spectral distribution of Σxp= := X' P1/2 converges weakly to H x supported on [c, CJ for C'> 0.
The following theorem characterizes the asymptotic variance and the corresponding optimal P. Theorem 1. Given (AI-2), the asymptotic variance is given as
v(Âp) & o² lim m (- x-50
(3.2)
where m(z) 0 the Stieltjes transform of the limiting distribution of eigenvalues of 1 1XPX satisfying the self-consistent equation m (2) =z 7(1 +
n x
Furthermore, under (A1-2), V (ÔOp) Z 0²(7 - 1)- and the equality is achieved by P=F-!
Formula (3.2) is a direct extension of (Hastie et al., 2019, Thorem 4). which can be obtained from the general results of Ledoit & Péché (2011); 1);; Dobriban et al. (2018). Theorem implies that precon- ditioning with the inverse population Fisher F results in the optimal stationary variance, which is supported by Figure 5(a). In other words, when the labels noisy that the risk is dominated by
4

---

Published as a conference paper at ICLR 2021
Pa (GD)
0.5
1.2
(NGD)
1.0
Pw 12 prediction
0.4
veriance
0.8
0.6
02
0.4
0.2
1.3
2.0
2.5 y- din
3.0
0.0
1.5
2.0
2.5
3.0
0.0
1.5
2.0
2.5
3,0
ym dyn
y y-ddin
(a) variance.
(b) well-specified bias (isotropic). (c) well-specified bias (misaligned).
Figure 5: We eigenvalues of Σx two point masses with Kx 20 and |Ex| 1|2 d: empirical values (dots) computed with 300. (a) NGD (blue) achieves minimum variance. (b) GD (red) achieves lower
bias under isotropic signal: Ne (c) NGD achieves lower bias under "misalignment". Σx 2
the variance, we expect NGD to generalize better upon convergence. emphasize that this advan- tage is only present when the population Fisher is used, but not its sample- cample-hased counterpart (which gives 02). In Appendix A.3 discuss the sample complexity of estimating F from unlabeled data.
Misspecification Label Noise. Under model misspecification, there does not exist a linear stu- dent that perfectly recovers the teacher f* which may decompose as: f*)x x 0 f2 (z)). In the simple case where L is an independent linear function on unobserved features (Hastie et al., 2019, Section 5): Yr x; 0 0"+ 1 + where æe,i Rd is independent 2; the additional error in the bias term due to misspecification is analogous to the variance term above:
show that
Corollary Under (A1)AA2), for the above unobserved features model with [[x"x=T] Exx and E[ d= -g the addi- tional error due to misspecification can be written as (0(0p) tr=x Xx))(((0))) 1). where p(0p) is the variance in (3.2).
P- Eg prediction
1.
1.0
In this case, misspecification can be interpreted as additional label noise, for which NGD is optimal by Theorem 1. While Corollary 2 may expect
0.5
nonimerrtty
describes one specific example of misspecification, similar outcome under broader settings. In particular, (Mei & Mon- tanari, 2019, Remark 5) indicates that for many nonlinear fä the misspecified bias is same as variance due to label noise. We empir- ically observe similar findings under general covariance in Figure 4, in which f2 is quadratic function. Observe that NGD leads to lower bias compared to GD as we further misspecify the teacher model.
Figure 4: Misspecified bias with Eg Ia (favors GD) and S: (z) a(æ æ --t==>>xxxx )), where a controls the extent of nonlinearity. Predictions are generated by matching of with second moment fa.
3.2 THE BIAS TERM: ALIGNMENT AND "DIFFICULTY" OF LEARNING
We now analyze the bias term when the teacher model is linear on the input features (hence well- specified); (x: x 0" Extending the random effects hypothesis in Dobriban et al. (2018), we d d'' Eg. and assume the following joint relations?:
consider more general prior on 0" E0000"
(A3) Joint Convergence: Σx and P share the same eigenvector matrix U. The empirical distri- butions of elements of (ez, Co, Ezp) jointly converge random variables (Uz, Vp, Uxp) supported on [d', C*] for é²,'' > O, where Ez, Eup are eigenvalues Ex. Exp, and eœ = diag (U EgU).
We remark that when P Is, previous works (Hastie al., 2019; Xu & Hsu, 2019) considered the special case of isotropic prior Eg = Ig. Our assumption thus allows for analysis of the bias term under much more general 3 which gives rise interesting phenomena that are not captured by simplified settings, such non-monen non-monotonic bias and variance for epoch-wise double descent phenomenon (see Appendix A.5). Under this general setup, we have the following characterization of the asymptotic bias and the corresponding optimal preconditioner:
(see Figure 15), and the
Theorem 3. Under (AI))A3), the expected bias B(Op) Eg ( } B(Op) 2 lim m' m'(-A)m -2 (--A)E U2 (o Ue "], A
given
(3.3)
Nooee that (A2)(A3) covers many common choices of preconditioner. such the population Fisher and variants the sample Fisher (which is degenerate hut leads the minimum solution GD). 'Two concurrent works (Wu & Xu, 2020; Richards et al., 2020) also considered similar relaxation of Ee in the context of ridge regression.
5

---

Published as a conference paper at ICLR 2021
where expectation is taken over and m(2) is the Stieltjes transform defined in Theorem 1. Furthermore, among all P satisfying (A3), the optimal bias is achieved by P = U diag (eo)UT.
Note that the optimal P depends on the "orientation" of the teacher model Σg. which is usually not known in practice. This result can thus be interpreted as a no-free-lunch characterization in choosing an opti- mal preconditioner for the bias term priori. As a consequence of the theorem, when the true parameters 0* have roughly equal magnitude (isotropic), GD achieves lower bias (see Figure 5(b) where Eg aa). On the other hand, NGD leads to lower bias when Σx is "misaligned" with Eo. i.e. when 0* focus on the least varying directions of input features (see Figure 5(c) where Σg E Σx', ), in which case learning is intuitively difficult since the features are not useful.
Exr!! isooopic misaligned o*
Illustration of isotropic and misaligned 0*
Figure 6:
Connection to Source Condition. The "difficulty" of learning above relates to the source condi- tion in RKHS literature (Cucker & Smale, 2002) (i.e., B|Ex 0* oo, see (A4) in Section 4.3), in which the coefficient can be interpreted as a measure of "misalignment". To elaborate this con- nection, we consider the setting of Σn X': note that as r decreases, the teacher 0" focuses more on input features with small magnitude, thus the learning problem becomes harder, and vice versa. In this case Proposition 4 (Informal). When Eg Ex. there exists a transition point r" E (-1,0) such that GD achieves lower (higher) stationary bias than NGD when > (<)r".
show clear transition in for the comparison between GD and NGD.
The above proposition confirms that for the stationary bias (well-specified), NGD outperforms GD in the misaligned setting (i.e., small r), whereas GD has advantage when the signal is aligned (large r). For formal statement and more discussion on the transition point r" see Appendix A.2.
4 BIAS-VARIANCE TRADEOFF
Our characterization of stationary risk suggests that preconditioners that achieve the optimal bias and variance are generally different. This section discusses how bias-variance tradeoff can be realized by interpolating between preconditioners or by early stopping. Additionally, we analyze the non- parametric least squares setting and show that by balancing the bias and variance, preconditioned update that interpolates between GD and NGD decreases the population risk faster than GD.
4.1 INTERPOLATING BETWEEN PRECONDITIONERS
Depending on the orientation of the teacher model, we may ex- pect bias-variance tradeoff in choosing P. Intuitively, given P1
0.40
Biss
vartance
that minimizes the bias and P2 that minimizes the variance, it is
0.35
geometric
additive
possible that a preconditioner interpolating between P1 and P can balance the bias and variance and thus generalize better. The following proposition confirms this intuition in a setup of general Ex and isotropic Σn¹, for which GD (P= Ia) achieves optimal stationary bias and NGD (P=F-) achieves optimal variance.
0.30
0.25
0.20
0.15
GD
interpolation confficient NGD
Proposition 5 (Informal). Let Σx & Ia and Eg = Ia- Consider the following choices of interpolation scheme: (i) P.. = a' (I-a)I (ii) Pa = (aEx +1 (a) (iii) P. Tx=bia x The stationary variance monotonically decreases with a [0, for all three choices. For (i), the stationary bias monotonically increases with a [0, 1), whereas for (ii) and (iii), the bias monotonically increases with a in certain range that depends on Σx.
Figure 7: Bias-variance tradeoff
with кX = 25. Σe = Ia and SNR=32/5. As additively (i)) or geometrically (ii) interpolate from GD to NGD (left to right), the stationary bias (blue) increases and the stationary variance (or- ange) decreases.
In other words, as the signal-to-now ratio (SNR) decreases, one can increase a, which makes the update closer NGD, to improve generalization, and vice versa. Indeed as shown in Figure and 16(c), a certain SNR, interpolating between Σx' and Eg can improve the stationary risk.
Note that this reduces the random effects model studied in Dobrihan al. (2018).

---

Published as a conference paper at ICLR 2021
Remark. Two of the aforementioned interpolation schemes reflect common practical choices: addi- rive interpolation (ii) corresponds to the damping term to stably invert the Fisher, whereas geometric interpolation (iii) resembles the "conservative" square-root scaling in adaptive gradient methods.
4.2 THE ROLE OF EARLY STOPPING
Thus far we considered the stationary solution of the unregularized objective. It is known that the bias-variance tradeoff can also be controlled by either explicit or algorithmic regularization. We briefly comment on the effect of early stopping, starting from the monotonicity of the variance term. Proposition 6. For all P satisfying (A2), the variance V (t)( monotonically increases with time.
The proposition confirms the intuition that early stopping reduces overfitting. Variance reduction can benefit GD in its comparison to NGD, which achieves lowest stationary variance: indeed, Figure and 19 show that under early stopping, GD may be favored even if NGD has lower stationary risk.
On the other hand, early stopping may not improve the bias term. While a complete analysis is difficult partially due to the potential non-monotonicity of the bias term (see Appendix A.5), we speculate that previous findings for the stationary bias also translate to early stopping. As a con- crete example, we consider well-specified settings in which either GD or NGD achieves the optimal stationary bias, and demonstrate that such optimality is also preserved under early stopping:
Proposition 7. Given (AI) and denote the optimal early stopping bias as Bopt (0) = infr)o
When Ee = Σx we have Bopt (p)) 2 (0 for all P satisfying (A3). Whereas when Eg = de we have ( p-1)
Bopt
Figure 19 illustrates that the observed trend in the stationary bias (well-specified) is indeed preserved under optimal early stopping: GD or NGD achieves lower early stopping bias under isotropic or misaligned teacher model. respectively. We leave more precise characterization as future work.
4.3 FAST DECAY OF POPULATION RISK
Our previous analysis suggests that certain preconditioners achieve lower population risk. but does not address which method decreases the risk more efficiently. Knowing that preconditioned updates may accelerate optimization, one natural question ask is, is this speedup also present for generalization under fixed dataset? We answer this question in the affirmative in slightly different model: we study least squares regression the RKHS, and show that preconditioned update that interpolates between GD and NGD achieves the minimax optimal rate in much fewer steps than GD.
We provide a brief outline and defer the details to Appendix 1.9.1. Let be an RKHS included in (2x)) equipped with bounded kemel function k. and K₂ E H be the Riesz representation of the kernel function. Define 5 as the canonical operator from H to L2 (Px), and write Σ = 5 S and We aim to learn the teacher model f", under the following standard regularity conditions:
I, = SS*
(A4) Source Condition:
3r (0, 00), M > s.t. = for h E L2 and II S M.
(A5) Capacity Condition: As > 1 s.t. tr //--) < 0o and 2r + s-1 1 > 1.
(A6) Regularity of RKHS: 3y E [s-1,1], Cp > 0 s.t. suprssupp( Px) Σ 11/2-1/WK
11x S CCp-
Note that in the source condition (A4), the coefficient r controls the complexity of the teacher model and relates to the notions of model misalignment in Section 3.2: large indicates a smoother teacher model which is "casier" to learn, and vice versa" (Steinwart et al., 2009). Given training points ((x,, y:) H we consider the following preconditioned update on the student model ft H:
f= m/E (E) -S =
(4.1)
where -= E K2, & Kz,, SYY responds to the covariance operator Σ, and thus (4.1) can be interpreted as additive interpolation between GD and NGD: update with large behaves like GD, and small a like NGD. Related to our update is the FALKON algorithm (Rudi et al., 2017) a preconditioned gradient method for kernel ridge regression. The key difference is that we consider optimizing the original objective (instead of
; ΣT y4 Kz.- In this setup, the population Fisher cor-
We remark that most previous works considered the case where 1/2 which implics f EHH

---

Published as a conference paper at ICLR 2021
regularized version in FALKON) under early stopping. Importantly, since we aim understand how preconditioning affects generalization, explicit regularization should not be taken into account.
The following theorem shows that with appropriately chosen a, the preconditioned update (4.1) leads to more efficient decrease in the population risk, due to faster decay of the bias term.
Theorem 8 (Informal). Under (A4-6), the population risk of Se
be writlen as R(Se)
f[S - (' S ((()+ defined in Appendix D.9. Given r r 21/2 or u 2r, preconditioned 209TI achieves minimux optimal convergence rate R(J.) Ö I
update (4.1) with
in (logg) steps, whereas ordinary gradient descent requires 1 0 276+1 steps.
We comment that the optimal interpolation coefficient and stopping time t chosen to balance the bias B()) and variance (1). Note that a depends on the teacher model in the following way: for 1, a decreases r becomes smaller, which corresponds non-smooth and "difficult" f" and vice versa. This agrees with our previous observation that NGD is ad advantageous when the teacher model is difficult learn. We defer empirical verification of this result Appendix C.
NEURAL NETWORK EXPERIMENTS
Protocol. We compare the generalization performance of GD and NGD in neural network settings and illustrate the influence the following three factors: (i) label noise; (ti) misspecification; (iti) signal misalignment. We also verify the potential advantage of interpolating between GD and NGD.
We consider the MNIST and CIFAR-10 datasets. To create student-teacher setup, we split the training set into two halves, pretrain the teacher, and the other (distill split) along with the teacher's labels used to distill (Hin- et al., 2015) the student. We take the teacher be either a two-layer fully-connected ReLU network or a ResNet (He et al., 2016), and the student is a two-layer ReLU network. We normalize the teacher's s labels logits following Ba & Caruana (2014) before potentially adding label noise, and fit the student by minimizing the L.2 loss. Student models are trained on subset of the distill split with full-batch updates. We implement NGD using Hessian-free free optimization (Martens, 2010). We use 100k unlabeled data (possibly applying data augmentation) to estimate the population Fisher. We report the test error when the training error below 0.2% of its initial value as a proxy for the stationary risk. We defer detailed setup to Appendix E and additional results to Appendix C.
of which (pretrain split) along with the original labels is used to
5.1
EMPIRICAL FINDINGS
NGD 1100000) +
NGD (100000)
10-2
NGD (60000)
Stationary Risk
12.
Stationary Risk
Stationary Risk
20-
10
0.000250500.75 1,00 1.501.552.00 Noise Standard Deviation
10
10'
-1.0 1.0-0.9-..----00 -10-00-0---------------22
Teacher Pretrain Iters.
(a) label noise (MNIST).
(b) misspecification (CIFAR-10), 10).
(c) misalignment (MNIST).
Figure 8: Comparison between NGD and GD. Error bar is standard deviation away from over five
independent runs. Numbers parentheses denote amount unlabeled examples for estimating the Fisher.
Label Noise. We pretrain the teacher with the full pretrain split and use 1024 examples from the distill split to fit the student. For both the student and teacher, we use two-layer ReL.U net with 80 hidden units. We corrupt the labels with isotropic Gaussian noise of varying magnitude. Figure 8(a) shows that as the noise level increases (the variance term begins to dominate), the stationary risk of both NGD and GD worsen, with GD worsening faster, which aligns with our observation in Figure 5.
Misspecification. We use a ResNet-20 teacher and the same two-layer ReLU student from the label noise experiment. We control the misspecification level by varying amount of pretraining
8

---

Published as a conference paper at ICLR 2021
of the teacher. Intuitively, large teacher models that are trained longer should be more complex and thus likely to be outside of functions that a small two-layer student can represent (hence the problem becomes misspecified). Indeed, Figure 8(b) shows that NGD eventually achieves better generalization as the number of training steps for the teacher increases. In Appendix A.4 we report a heuristic measure of model misspecification that relates to the NTK (Jacot al., 2018), and confirm that the quantity increases as more label noise is added or as the teacher model is trained longer.
Misalignment. We set the student and teacher to be the same two-layer ReLU network. We con- struct the teacher model by perturbing the student's ' initialization, the direction of which given by F", where F is the Fisher of the student model and r E [-11, [-1,0) O). Intuitively, as r decreases, the impor- tant parameters of the teacher (i.e. larger update directions) becomes misaligned with the student's s gradient, and thus learning becomes more "difficullt. While this analogy is rather superficial due to the non-convexity of neural network optimization, Figure 8(c) shows that as r becomes smaller (setup is more misaligned), NGD begins generalize better than GD (in terms of stationary risk).
Interpolating between Preconditioners. We validate our observations in Section and 4 on the difference between the sample Fisher and population Fisher, and the potential benefit of interpolating between GD and NGD, in neural network experiments. Figure 9() shows that as we decrease the number of unlabeled data in estimating the Fisher, which renders the preconditioner closer to the sample Fisher, the stationary risk becomes more akin to that of GD, especially in the large noise setting. This agrees with our remark on sample vs. population Fisher in Section 3 and Appendix A.
Figure {b)(c) supports the haasinarancce tradeoff discussed in Section 4. in neural network settings. In particular, we interpret the left end of the figure to correspond to the bias-dominant regime (due to the same architecture for the student and teacher), and the right end be the variance-dominant regime (due to large label noise). Observe that at certain SNR, a preconditioner that interpolates (additively or geometrically) between GD and NGD can achieve lower stationary risk.
Stationary Risk
MSD (10000)
GB
MGD (24000)
20.0
17.5-
damp--00.0
damp-0.1
3.0
Stattensry RIRK
50
6:
NGD (54001)
15.0
damppe0.01
NGD (200000)
Statonarry
12.5
dampp0
dampp 2.001
20
2.5
5.00 0.25 0.50 075 1:00 1-b 152 1.5 200
0.0 0.5 0.39 5/5 1000 125 150 177 205
5.00 025 3.30 0.5 100 125 130 1./5 200
Noise Standard Deviation
Noise Standard Devialion
Noise Standard Deviation
(a) interpolation between sample and population Fisher (CIFAR-1O).
(b) additive interpolation (ii)
(c) geometrie interpolation (iii)
between GD and NGD (MNIST).
between GD and NGD (MNIST).
Figure 9: (a) numbers in parentheses indicate the amount of unlabeled data used in estimating the Fisher F; we expect the estimated Fisher he closer the sample Fisher when the number of unlabeled data is small. (a)
additive interpolation P (F+ + aTa) aa=' )); larger damping parameter yields update closer to GD. (b) geometric interpolation P= F F= larger a parameter yields update closer to that of NGD (blue).
DISCUSSION AND CONCLUSION
We analyzed the generalization properties of a general class of preconditioned gradient descent in overparameterized least squares regression, with particular emphasis on natural gradient descent. We identified three factors that affect the comparison of generalization performance of different optimizers, the influence of which we also empirically observed in neural network We then de- termined the optimal preconditioner for each factor. While the optimal P is usually not known in practice, we provided justification for common algorithmic choices by discussing the bias-variance tradeoff. Note that our current theoretical setup is limited to fixed preconditioners or those that do not alter the span of gradient, and thus does not cover many adaptive gradient methods; understand- ing these optimizers in similar setting would be an interesting future direction. Another important problem is to further characterize the interplay between preconditioning and explicit (e.g. weight decay") algorithmic regularization (e.g. large step size and gradient noise).
We however note that observations in linear model may not translate to neural network: many works have illustrated such a gap (c.g., Ghorbani et al. (2019); Allen-Zhn & Li (2019); Suzuki (2020); Yang & Hu (2020). In companion work (Wu& & Xu, 2020) we characterized the impact of E2 regularization in similar settings.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGEMENT
The authors would like to thank Murat A. Erdogdu, Fartash Faghri, Ryo Karakida, Yiping Lu, Jiaxin Shi, Shengyang Sun, Yusuke Tsuzuku, Guodong Zhang, Michael Zhang, Tianzong Zhang, and anonymous ICLR reviewers for helpful feedback. The authors are also grateful to Tomoya Murata for his contribution preliminary studies on the nonparametric least squares setting.
JB and RG were supported by the CIFAR AI Chairs program. JB and DW were partially supported by I.G Electronies and NSERC. AN was partially supported by JSPS Kakenhi (19K20337) and IST-PRESTO. TS was partially supported by JSPS Kakenhi (26280009, 15H05707 and 18H03201), Japan Digital Design and JST-CREST. JX was supported by a Cheung- Kong Graduate School of Business Fellowship. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government Canada through CIFAR, and companies sponsoring the Vector Institute.
REFERENCES
Ben Adlam and Jeffrey Pennington. Understanding double descent requires fine-grained bias- variance decomposition. arXiv preprint arXiv:201..0332/ 2020.
Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. The case for full-matrix adaptive regularization. arXiv preprint arXiv: 1806.02958, 2018.
Alnur Ali, J Zico Kolter, and Ryan J Tibshirani. A continuous-time view of early stopping for least squares. International Conference on Artificial Intelligence and Statistics, volume 22, 2019.
Zeyuan Allen and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? arXiv preprint arXiv:1905.1033 2019.
Shun-ichi Amari. Neural learning in structured parameter natural riemannian gradient. In Advances in neural information processing systems, pp. 127-133, 1997.
Shun-Ichi Amari Natural gradient works efficiently in learning. Neural computation, 10(2):251- 276, 1998.
Shun-Ichi Amari, Hyeyoung Park, and Kenji Fukumizu. Adaptive method of realizing natural gra- dient learning for multilayer perceptrons. Neural computation, 12(6): 1999-149,, 2000.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Advances in Neural Information Processing Systems, pp. 7411-7422 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. arXiv preprint arXiv:1901.08584, 2019b.
Navid Azizan, Sahin Lale, and Babak Hassibi. Stochastic mirror descent on overparameterized nonlinear models: Convergence, implicit regularization, and generalization. arXiv preprint arXiv:1906.038300 03830, 2019.
Jimmy Ba and Rich Caruana. Do deep nets really need to be deep? In Advances in neural informa- fion processing systems, pp. 2654-2662, 2014.
Jimmy Ba, Murat A. Erdogdu, Taiji Suzuki, Denny Wu, and Tianzong Zhang. Generalization of two-layer neural networks: An asymptotic viewpoint. International Conference Learning Representations, 2020.
Zhi-Dong Bai and Yong- Qua Yin. Limit of the smallest eigenvalue of large dimensional sample covariance matrix. In Advances In Statistics, pp. 107-127. World Scientific, 2008
Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. arXiv preprint arXiv: 1906. 1300, 2019.
10

---

Published as a conference paper at ICLR 2021
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine learning and the bias-variance trade-off. arXiv preprint arXiv:1812.1111,, 2018a.
Mikhail Belkin, Daniel J Hsu, and Partha Mitra. Overfitting or perfect fitting? risk bounds for classification and regression rules that interpolate. Advances in neural information processing systems. pp. 2300-2311 2018b.
Mikhail Belkin, Alexander Rakhlin, and Alexandre Tsybakov. Does data interpolation contradict statistical optimality? arXiv preprinl arXiv: urXiv:/806.0000 17880000977, 2018c.
Alberto Bietti and Julien Mairal. On the inductive bias of neural tangent kernels. In Advances in Neural Information Processing Systems, pp. 12873-12888 12884, 2019.
Tianle Cai, Ruiqi Gao, Jikai Hou, Siyu Chen, Dong Wang. Di He, Zhihua Zhang, and Liwei Wang. A gram-gauss gauss-newton method learning overparameterized deep neural networks for regression problems. arXiv preprint arXiv: 19t.:::995.1665,,, 2019.
Andrea Caponnetto and Emesto De Vito. Optimal rates for the regularized least-squares algorithm. Foundations of Computational Mathematics, 7(3):331-368, 2007.
Lenaic Chizat and Francis Bach. Implicit bias of descennt descent for wide two-layer neural networks trained with the logistic loss. arXiv preprini arXiv:2002.04486, 2020.
Nello Cristianini, John Shawe-taylor, André Elisseeff, and Jaz Kandola. On kernel-target alignment. In Advances in Neural Information Processing Systems 14. Citeseer, 2001.
Felipe Cucker and Steve Smale. On the mathematical foundations of learning. Bulletin of the American mathematical society, 39(1):1-49, 2002.
Assaf Dauber, Meir Feder, Tomer Koren, and Roi Livni. Can implicit bias explain generalization?
stochastic convex optimization case study. arXiv preprint arXiv:2003.06/52: 2020.
Zeyu Deng, Abla Kammoun, and Christos Thrampoulidis. A model of double descent for high- dimensional binary linear classification. arXiv preprint arIiv: 1at019911.052,,, 2019.
Guillaume Desjardins, Razvan Pascanu, Aaron Courville, and Yoshua Bengio. Metric-free natural 2013.
gradient for joint-training of boltzmann machines. arXiv preprint
Oussama Dhifallah and Yue M Lu. A precise performance analysis of learning with random features. arXiv preprint arXiv:2008. /904, 2020.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. arXiv preprint arXiv:1703.04933: 2017.
Edgar Dobriban, Stefan Wager, et al. High-dimensional asymptotics of prediction: Ridge regression and classification. The Annals of Statistics, 46(1):247-279, 2018.
Bin Dong, Jikai Hou, Yiping Lu, and Zhihua Zhang. Distillation x early stopping? harvesting dark knowledge utilizing anisotropic information retrieval for overparameterized neural network. arXiv preprint arXiv: 1910.01255, 2019.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(Jul):121-2155) 2011.
Federica Gerace, Bruno Loureiro, Florent Krzakala, Mare Mézard, and Lenka Zdeborová. Gener- alisation error in learning with random features and the hidden manifold model. arXiv preprint 0033939002..9339 2020.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural networks. arXiv preprint wxXiv:196.08889, 2019.
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julins Julien. implicit of discrete gradient dynamics in deep linear neural networks, arXiv preprint arXiv: 19X4132994.1222.. 2019.
11

---

Published as a conference paper at ICLR 2021
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro.
Implicit regularization in matrix factorization. In Advances in Neural Information Processing Systems, pp. 6151-6159, 2017.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. arXiv preprint urXiv:/802.08246. 2018a.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. In Advances in Neural Information Processing Systems, pp. 9461-9471, 2018b.
Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan J Tibshirani. Surprises in high- dimensional ridgeless least squares interpolation. arXiv preprint arXiv:1903.08566 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Magnus Rudolph Hestenes, Eduard Stiefel, et Methods of conjugate gradients for solving linear systems, volume 49, 1952.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.022 02531, 2015.
Takashi Ishida, Ikko Yamane, Tomoya Sakai, Gang Niu, and Masashi Sugiyama. Do we need zero training loss after achieving zero training error? arXiv preprint arXiv:2002.08709, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, pp. 8571- 8580, 2018.
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. arXiv preprint arXiv:1810.0203: 2018.
Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Conference Learning Theory, pp. 1798, 2019.
Ryo Karakida and Kazuki Osawa. Understanding approximate fisher information for fast conver- gence of natural gradient descent in wide neural networks. Advances in Neural Information Pro- cessing Systems, 33, 2020.
Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. Universal statisties of fisher information in deep neural networks: Mean field approach. arXiv preprint arXiv:1806.01316, 2018.
Noureddine El Karoui. Asymptotic behavior of unregularized and ridge-regularized high- dimensional robust regression estimators: rigorous results. arXiv preprint arXiv:1311.2445 2013.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint reXivv16699 04666, 2016.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Frederik Kunstner, Lukas Balles, and Philipp Hennig. Limitations of the empirical fisher approxi- mation. arXiv preprint arXiv:1905.12558 2019.
Yann A LeCun, Léon Bottou, Genevieve B Orr, and Robert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pp. 9-48. Springer, 2012.
Olivier Ledoit and Sandrine Péché. Eigenvectors some large sample covariance matrix ensembles. Probability Theory and Related Fields, (51(1)2)233320)) 264, 2011.
12

---

Published as a conference paper at ICLR 2021
Daniel Levy and John C Duchi. Necessary and sufficient geometries for gradient methods. In Advances in Neural Information Processing Systems, pp. 11491-11501, 2019.
Mingchen Li, Mahdi Soltanolkotabi, and Samet Oymak. Gradient descent with early stop- ping is provably robust to label noise for overparameterizee neural networks. arXiv preprint arXhv:1903.11680 2019.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-parameteriized matrix sensing and neural networks with quadratic activations. arXiv preprint arXiv: 2017.
Tengyuan Liang and Alexander Rakhlin. Just interpolate: Kernel" ridgeless" regression can gener- alize. arXiv preprint
2018.
Zhenyu Liao and Romain Couillet. The dynamics of learning: random matrix approach. arXiv preprint arXiv: /9/7, 2018.
Junhong Lin and Lorenzo Rosasco. Optimal rates for multi-pass stochastic gradient methods. The Journal of Machine Learning Research, 18(1):3375-3411, 2017.
Cosme Louart, Zhenyu Liao, Romain Couillet, et al. random matrix approach to neural networks. The Annals of Applied Probability, 28(2): 190-1248, 2018.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint (003vv:900008999 2019.
James Martens. Deep learning via hessian-free optimization. In ICML, volume 27, pp. 735-742, 2010.
James Martens. New insights and perspectives on the natural gradient method. arXiv preprint arXiv:1412.1193, 2014.
James Martens and Roger Grosse. Optimizing neural networks with cronecker-factoree approximate curvature. In International conference machine learning, pp. 2408-2417, 2015.
James Martens and Ilya Sutskever. Training deep and recurrent networks with hessian-free opti- mization. In Neural networks: Tricks of the trade, pp. 479-535. Springer, 2012.
Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and double descent curve. arXiv preprint arXiv: 7900.9088000 05355, 2019.
Stanislav Minsker. On some extensions of Bernstein's 's inequality for self-adjoint operators. Statistics & Probability Letters, 127:1 11-19,, 2017.
Andrea Montanari, Feng Ruan, Youngtak Sohn, and Jun Yan. The generalization error of max- margin linear classifiers: High-dimension asymptotics in the overparametrized regime. arXiv preprint arXiv:1911.01544 2019.
Tomoya Murata and Taiji Suzuki. Gradient descent in rkhs with importance labeling. arXiv preprint arXiv:2006. 10925, 2020.
Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: Where bigger models and more data hurt. arXiv preprint arXiv: 2019.
Yann Ollivier. Riemannian metrics for neural networks i: feedforward networks. Information and Inference: Journal of the IMA, 4(2):108-153, 2015.
Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. arXiv preprint arXiv:/301.. /301.3544 2013.
Loucas Pillaud- Vivien, Alessandro Rudi, and Francis Bach. Statistical optimality of stochastic gra- dient descent on hard learning problems through multiple passes. In Advances in Neural Infor- mation Processing Systems, pp. 81 14-8124, 2018.
13

---

Published as a conference paper at ICLR 2021
Qian Qian and Xiaoyuan Qian. The implicit bias of adagrad separable data. In Advances in Neural Information Processing Systems, pp. 7759-7767, 2019.
Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms, arXiv preprint arXiv:2005.06338 2020.
Dominic Richards, Jaouad Mourtada, and Lorenzo Rosasco. Asymptotics of ridge (less) regression under general source condition. arXiv preprint arXiv:2006.06.86 2020.
Francisco Rubio and Xavier Mestre. Spectral convergence for a general class of random matrices. Statistics & probability letters, 81(5):592-60 602, 2011.
Alessandro Rudi and Lorenzo Rosasco. Generalization properties of learning with random features. In Advances in Neural Information Processing Systems, pp. 3215-3225, 2017.
Alessandro Rudi, Luigi Carratino, and Lorenzo Rosasco. Falkon: An optimal large scale kernel method. In Advances in neural information processing systems, pp. 3888-3898, 2017.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynam- ics learning in deep linear neural networks. arXiv preprint 1rXiv:1312.6122 2013.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19 (1):2822-2878, 2018.
Ingo Steinwart, Don R Hush, and Clint Scovel. Optimal rates for regularized least squares regres- sion. In COLT, pp. 79-93, 2009.
Lili Su and Pengkun Yang. On learning werrpaaametterized neural networks: A functional approx- imation perspective. In Advances in Neural Information Processing Systems, pp. 2637-2646, 2019.
Arun Suggala, Adarsh Prasad, and Pradeep K Ravikumar. Connecting optimization and regulariza- tion paths. Advances in Neurul Information Processing Systems, pp. 10608-10619, 2018
Taiji Suzuki. Generalization bound of globally optimal non-convex neural network training: Transportation map estimation by infinite dimensional langevin dynamics. arXiv preprint arXiv:2007.05824. 2020.
Valentin Thomas, Fabian Pedregosa, Bart Merriënboer, Pierre-Antoine Manzagol, Yoshua Bengio, and Nicolas Le Roux. On the interplay between noise and curvature and its effect optimization and generalization. In International Conference on Artificial Intelligence and Statistics, pp. 3503- 3513, 2020.
Sharan Vaswani, Reza Babanezhad, Jose Gallego. Aaron Mishkin, Simon Lacoste-Julien, and Nicolas Le Roux. To each optimizer a norm, to each norm its generalization. arXiv preprint arXiv:2006.06822 2020.
Neha S Wadia, Daniel Duckworth, Samuel S Schoenholz, Ethan Dyer, and Jascha Sohl-Dickstein. Whitening and second order optimization both destroy information about the dataset, and can make generalization impossible. arXiv preprint arXiv:2008.07545 2020.
Per-Ake Wedin. Perturbation theory for pseudo-inversed BIT Numerical Mathematics, 13(2):217- 232, 1973.
Francis Williams, Matthew Trager, Claudio Silva, Daniele Panozzo, Denis Zorin, and Joan Bruna. Gradient dynamics of shallow univariate relu networks. arXiv preprint /006.07990.78888 2019.
Ashia C Wilson, Rebecca Roclofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems, pp. 4148 4184115,11 2017.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. arXiv preprint ar arXiv:2002.0922 2020.
14

---

Published as a conference paper at ICLR 2021
Denny Wu and Ji Xu. On the optimal weighted lg regularization in overparameterized linear regres- sion. Advances in Neural Information Processing Systems, 33, 2020.
Lei Wu, Chao Ma, and E Weinan. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. In Advances in Neural Information Processing Systems, pp. 8279-8288, 2018.
Bo Xie, Yingyu Liang, and Le Song. Diverse neural network learns true target functions. arXiv preprint 1XXv::611.033333 2016.
Xu and Daniel Hsu. How many variables should be entered in a principal component regression equation? arXiv preprint arXiv: 139, 2019.
Peng Xu, Fred Roosta, and Michael W Mahoney. Second-order optimization for non-convex chine learning: An empirical study. In Proceedings of the 2020 SIAM International Conference on Data Mining. pp. 199-207. SIAM, 2020.
Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011. 1. 14522,, 2020.
Guodong Zhang. Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger Grosse. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. In Advances in Neural Information Processing Systems, pp. 8194- 8205, 2019a.
Guodong Zhang, James Martens, and Roger B Grosse. Fast convergence of natural gradient descent for wee-parramterrized neural networks. In Advances Neural Information Processing Systems, pp. 1080-8091, 2019b.
15

---

Published as a conference paper at ICLR 2021
TABLE OF CONTENTS
Discussion on Additional Results
17
A.1 Implicit Bias of GD vs. NGD
17
A.2 Bias Term Under Specific Source Condition
18
A.3 Estimating the Population Fisher
19
A.4 Interpretation of Vy K y/n
19
A.5 Non-monotonicity of Bias Term w.r.t. Time
20
B Additional Related Works
20
Additional Figures
21
C.I Overparameterizee Linear Regression
21
C.2 RKHS Setting
22
C.3 Neural Network
23
D Proofs and Derivations
24
D.1 Missing Derivations in Section 3
24
D.2 Proof of Theorem
24
D.3 Proof of Corollary 2
25
D.4 Proof of Theorem 3
25
D.5 Proof of Proposition 4.
27
D.6 Proof of Proposition 5
28
D.7 Proof of Proposition 6
30
D.8 Proof of Proposition
31
D.9 Proof of Theorem
32
D.10 Proof of Corollary 9
39
D.11 Proof of Proposition 10
39
E Experiment Setup
40
E.1 Processing the Datasets
40
E.2 Setup and Implementation for Optimizers
40
E.3 Other Details
40
16

---

Published as a conference paper at ICLR 2021
A DISCUSSION ON ADDITIONAL RESULTS
A.1 IMPLICIT BIAS OF GD vs. NGD
It is known that gradient descent is the steepest descent with respect to the l2 norm, i.e. the update direction is constructed to decrease the loss under small changes the parameters measured by the l2 norm (Gunasekar et al., 2018a), Following this analogy, NGD is the steepest descent with respect to the KL divergence on the predictive distributions (Martens, 2014); this can be interpreted as a proximal update which penalizes how much the predictions change the data distribution.
Intuitively, the above discussion suggests GD tend to find solution that is close to the initialization in the Euclidean distance between parameters, whereas NGD prefers solution close the initialization in terms of function outputs on the data distribution. This observation turns out to be exact in the case of ridgeless interpolant under the squared loss, as remarked in Section 3. Moreover, Figure and 10 confirms the same trend in the optimization of overparameterized neural network. In particular,
GD results in small changes in parameters, whereas NGD results in small changes in the function.
Preconditioning with the pseudo-inverse of the sample Fisher, i.e., P p= (JT J)1, leads implicit bias similar to that of GD (also noted in (Zhang et al., 2019b)), but different than NGD with the population Fisher.
"Interpolating" between GD and NGD (green) results in properties in between GD and NGD.
Remark. Qualitatively speaking, the small change in the function output is the essential reason that NGD performs well under noisy labels in the interpolation setting: NGD seeks to interpolate the training data by changing the function only "locally", that memorizing the noisy labels has small impact on the "global" shape of the learned function (see Figure l).
10a
10-1
[[W00) -W00112
R011
10-
20-2
10
10-3
[ite
NGD
10-4
sqrt NGD
10-3
sample NGD
steps
steps
(a) Difference Parameters.
(b) Difference Function Values.
Figure 10:
Illustration of implicit bias of GD and NGD. We set n = 100, d = 50, and regress a two-layer ReLU network with 50 hidden units towards teacher model of the same architecture on Gaussian input. The x-axis is rescaled for each optimizer such that the final training error is below 10 GD finds solution with small changes the parameters, whereas NGD finds solution with small changes in the function. Note that the sample Fisher (cyan) has implicit bias similar to GD and does not resemble NGD (population Fisher).
We note that the above observation also implies that wide neural networks trained with NGD (popu- lation Fisher) is less likely stay in the kernel regime: the distance traveled from initialization can be large (see Figure 10(a)) and thus the Taylor expansion around the initialization is no longer ac- curate. In other words, the analogy between wide neural net and its linearized kernel model (which we partially employed in Section 5) may not be valid in models trained by NGDS.
Implicit Bias of Interpolating Preconditioners. We also expect that as we interpolate from GD to NGD, the distance traveled in the parameter space would gradually increase, and distance traveled in the function space would decrease. Figure 11 demonstrate that this is indeed the case for neural networks: we use the same two-layer MLP setup on MNIST as in Section 5. Observe that updates closer GD result in smaller change in the parameters, and ones close NGD lead smaller change in the function outputs.
Note that this gap is only present when the population Fisher is used; previous works have shown NTK-type global convergence for sample Fisher-related update (Zhang et al., 2019h; Cai al., 2019).
17

---

point Published as a conference paper at ICLR 2021
Damping
Damsing
(a) additive interp:; (b) additive interp.; (c) geometric interp.;
(d) geometrie interp.; difference in functions.
difference in parameters. difference in functions.
difference in parameters.
Illustraioon of the implicit bias of preconditioned gradient descent that interpolates between GD and NGD on MNIST. Note that the update becomes more similar NGD (smaller damping larger a), the distance traveled the parameter space increases, where as the distance traveled on the output space decreases.
Figure 11:
Two Kinds of Second-order Optimizer. Note that our optimal preconditioner derived in Section requires knowledge of population second-order statistics, which we empirically approximate using extra unlabeled data. Consequently, our results suggest that different "types" of second-order infor- mation (sample vs. population) may affect generalization differently. Broadly speaking, there two types of practical approximate second-order optimizers for neural networks. Some algorithms, such as Hessian-free optimization (Martens, 2010; Martens & Sutskever, 2012: Desjardins et al., 2013), approximate second-order matrices (typically the Hessian or Fisher) using the exact matrix on finite training examples. In high-dimensional problems, this sample- sample-based approximation can be very different from the population quantity (e.g. it is degenerate in the overparameterized regime). Other algorithms fit a parametric approximation to the Fisher, such as diagonal (Duchi et al., 2011: I; Kingma & Ba, 2014), quasi-diagonal (Ollivier, 2015), or Kronecker-factored (Martens & Grosse, 2015). If the parametric assumption is accurate, these approximations are more statistically efficient and thus may lead to better approximation to the population Fisher. Our analysis reveals a difference (in generalization properties) hetween sample- and population-based preconditioned updates, which may also suggest separation between the two kinds of approximate second-order optimizers. As future work, we intend to investigate this discrepancy in real-world problems.
A.2 BIAS TERM UNDER SPECIFIC SOURCE CONDITION
Motivated by the connection between the notion of "alignment" and the source condition in Sec- tion 3.2, we consider a specific case of 0": Eg Ex, where controls the extent of misalignment, and Theorem implies that the optimal preconditioner for the bias term (well-specified) P Exx Note that smaller r corresponds to more misaligned and thus "difficult" problem, and vice versa. In this setup we have the following comparison between GD and NGD.
Proposition (Formal Statement of Proposition 4). Consider the setting of Theorem 3 and Eg = Exx then for all r S -1 we have B(0p-:) S B(θr), whereas for all r 2 0, we have B(0p-1) 2 B(Or); the equality is achieved when features are isotropic (ie.. Σx = acaa).
The proposition confirms the intuition that when parameters of the teacher model 0* are more "aligned" with features æ than the isotropic setting (r 2 0), then GD achieves lower bias than NGD; on the other hand, when Eg is more "misaligned" than the Σx' case (r < -1), then NGD guaranteed to be advantagcous for the bias term. We therefore expect tran- sition from the NGD-dominated to GD-dominated regime for some (-1,0), The exact value of for such transition depends on the spectral distribution of Σx and varies case- by-case (one would need specifically evaluate the equality (D.9)). To give a concrete example, when Σx has a simple block structure,
(GD)
1.2
P- E2² (NGD)
P = 5,10
1.0
prediction
i 0.8
0.6
0.4
0.0
0.2
0.4
0.6
0.8 1.0
extent of misalignmert (-r)
Figure 12: We set Eg Σx Y 20, and plot the stationary bias (well-specified) under varying r.
explicitly determine the the transition point r" E (-1, 0), as shown by the following corollary.
Corollary 9. Assume Eo Exx and eigenvalues of Ex come from two masses with Amaxe MaaxxExx (Ex)/ Amin (Ex). WLOG we take tr (xxxx// )/d 1. Then given r" In C In (see Appendix D). 10 for definition), we have B(O,) 2 B(@ F-1) if and only if - rr"
18

---

Published as a conference paper at ICLR 2021
Remark. When 2, the transition happens r' -1/2 which is independent of the condition number R, as indicated by the dashed line in Figure 12. However for other y > 1, r* generally relates to both and
Our characterization above is supported by Figure 12, in which we plot the bias term under varying extent of misalignment (controlled by r*) in the setting of Corollary 9. Observe that as we construct the teacher model to be more "misaligned" (and thus difficult to learn) by decreasing r, NGD (blue) achieves lower bias compared to GD (red), and vice versa.
A.3 ESTIMATING THE POPULATION FISHER
Our analysis on linear model considers the idealized setup with access to the exact population Fisher, which can be estimated using additional unlabeled data. In this section we discuss how our result in Section and Section 4 are affected when the population covariance is approximated from N ii.d. (unlabeled) samples X IR Nxd, For the ridgeless interpolant have the following result on the substitution error in replacing the true population covariance with the sample estimate.
Proposition 10. Given (A (A3) and N/d + dp as + oc, let Σx X Xx X./N, /N, we have
(a) |Σx =
almost surely.
(b) Denote the stationary bias and variance of NGD (with the exact population Fisher) B" and V*, respectively, and the bias and variance of the preconditioned update using the approximate Fisher F = Σx x as B and V, respectively. Let O < c < 1 be the desired accuracy. Then
- = f 2) suffices to achieve 1B* B) e and V - V<< e.
Proposition 10 entails that when the preconditioner is a sample estimate of the Fisher F (based on unlabeled data), we approach the stationary bias and variance of the population Fisher at rate of 1-1/ we increase the number of unlabeled data N linearly in the dimensionality d. In other words, any non-vanishing accuracy e can be achieved with finite & (to push e 0, additional logarithmic factor is required, c.g. N logdd, which is beyond the proportional limit).
On the other hand, for our result in Section 4.3, (Murata & Suzuki, 2020, Lemma A.5) directly im- plies that setting N Onn" log n. is sufficient to approximate the population covariance operator (i.c.. so that || || 1/2 = (1)). Finally, we remark that our analysis above does not impose any structural assumptions on the estimated matrix. When the Fisher exhibits certain structures (e.g. Kro- necker factorization (Martens & Grosse, 2015)), then estimation be more sample-efficient For analysis on such structured approximations of the Fisher see Karakida & Osawa (2020).
A.4 INTERPRETATION OF VUT K yyn
Teacher Pretrain Iers.
10a
10f
As heuristic measure of model misspecification, in Figure 13 we report v Ki--k-- "y/n studied in Arora et al. (2019b), where is the label vector and K is the NTK matrix (Jacot et al., 2018) of the student model. This quantity relates to kernel- based alignment measures (Cristianini et al., 2001), and in the context of neural network optimization, it can be interpreted as proxy for measuring how much signal and noise are dis- tributed along the eigendirections of the NTK (e.g., see Li et al. (2019); Dong et al. (2019); Su & Yang (2019)). Roughly speaking, large Vy K K- y/n implies that the problem is "dif- ficult" learn by the student model via GD, and vice versa.
5.0
3.0-
2.5
2.0
Misspedification
3.0
0.2
2.4
0.8
0.8
1.0
Noise Standard Deviation on
Figure 13:
V K k'yy/n
layer neural network (CIFAR-IO).
Here we give heuristic argument on how this quantity relates to label noise and misspecification. For the ridgeless regression model considered in Section 3, write y; = (xx) I (xx) x 0", fe is the misspecified component, and E; is the label noise, then we have
[[(x:)+7] +
fxx.)) + where
E y K x--yl
yy
|| xxx
(x)) + ++)
Str @
0 x (XXX (0² + toy (XXT)-1).
(A.1)
19

---

Published as a conference paper at ICLR 2021
where we heuristically replaced the misspecified component with i..d. noise of the same variance o The first term of (A.1) resembles an RKHS norm of the target 0", whereas the second term is small when the feature matrix is well-conditioned or when the level of label noise a and misspecification of is small (note that these are conditions under which GD achieves good generalization). We may expect similar behavior for neural networks close the kernel regime. This provides non-rigorous explanation of the trend observed in Figure 13: as we increase the level of label noise or model misspecification (by pretraining the teacher for more steps), the quantity of interest becomes larger.
A.5 NON-MONOTONICITY OF BIAS TERM W.R.T. TIME
Many previous works on the high- high-dimensional characteri- zation of linear regression assumed a random effects model with an isotropic prior on the true parameters (Dobriban et al., 2018; Hastie et al., 2019). which may not hold in practice. As an example of the limitation of this assump- tion, note that when Σg- = Ia, it can be shown that the ex- pected bias (()( monotonically decreases through time (see proof of Proposition 7). In contrast, when the target pa- rameters do not follow an isotropic prior, the bias of GD can exhibit non-monotonieityy which gives rise to the "epoch- wise double descent" phenomenon also observed in deep learning (Nakkiran et al., 2019; Ishida et al., 2020).
1.00
1.0
0.95
NGD
0.8
0.90
0.05
nisk
0.80
0.4
0.75
0.2
0.70
time
Figure 14: Epoch-wise double descent. Note that non non-monotonicity of the bias w.r.t. time is present in GD but not NGD.
We empirically demonstrate this non-monotonicity when the model is close to the interpolation threshold in Figure 14. We set eigenvalues of Σx to be two point masses with xx = 32, Eg Σx and y = 16/15. Note that the GD trajectory (red) exhibits non-monotonicity in the bias term, whereas for NGD the bias is monotonically decreasing through time (which we confirm in the proof of Proposition 7). We remark that this mechanism of epoch-wise double descent may not relate the empirical findings in deep neural networks (the robustness of which is also largely unknown), in which it is typically speculated that the variance term exhibits non-monotonicity.
B ADDITIONAL RELATED WORKS
Implicit Regularization in Optimization. In overparameterized linear models, GD finds the min- imum & norm solution under many loss functions. For the more general mirror descent, the implicit bias is determined by the Bregman divergence of the update (Gunasekar et al., 2018b; Suggala al., 2018; Azizan et al., 2019). Under the exponential or logistic loss, recent works showed that GD finds the max-margin direction in various models (Ji & Telgarsky, 2018; 2019; Soudry et al., 2018: Lyu & Li, 2019; Chizat & Bach, 2020). The implicit bias of Adagrad has been analyzed under similar setting (Qian & Qian, 2019). Implicit regularization also relates to the model architecture; examples include matrix factorization (Gunasekar et al., 2017: Saxe et al., 2013; Gidel al., 2019; Arora et al., 2019a) and certain stylized neural networks (Li et al., 2017: Gunasekar et al., 2018h; 88;; Williams et al., 2019; Woodworth et al., 2020). For wide networks in the kernel regime (Jacot al., 2018 8), the implicit bias of GD relates to properties of the neural tangent kernel (NTK) (Xie et al., 2016; Arora al., 2019b; Bietti & Mairal, 2019). Finally, we note that the implicit bias of GD is not always explained by the minimum norm property (Razin & Cohen, 2020; Dauber et al., 2020).
Asymptotics of Interpolating Estimators. In Section 3 we analyzed overparameterized estima- tors that interpolate the training data. Recent works have shown that interpolation may not lead to overfitting (Liang & Rakhlin, 2018: Belkin et al., 2018c:b; Bartlett et al., 2019), and the optimal risk may be achieved under no regularization and extreme overparameterization (Belkin et al., 2018a; Xu & Hsu, 2019). The asymptotic risk of overparameterizee models has been characterized in various settings, such as linear regression (Karoui, 2013; Dobriban et al., 2018: Hastic et al., 2019), random features regression (Mei & Montanari, 2019; Gerace et al., 2020; Dhifallah & Lu, 2020; Adlam & Pennington, 2020), max- mar-in classification (Montanari et al., 2019; Deng et al., 2019), and certain neural networks (Louart et al., 2018; Ba et al., 2020). Our analysis is based on random matrix theory results developed in Rubio & Mestre (201 1); Ledoit & Péché (2011), 1). Similar tools can also be used to study the gradient descent dynamics of linear regression (Liao & Couillet, 2018; Ali al., 2019).
20

---

Published as a conference paper at ICLR 2021
C ADDITIONAL FIGURES
C.I OVERPARAMETERIZEE LINEAR REGRESSION
Non-monotonicity of the Risk. Under generalized (anisotropic) assumption the covariance of the features and the target, both the bias and the variance term can exhibit non-monoue monooonotynicity the overparameterization level y 1: in Figure 15 observe peaks in the bias term and three peaks the variance term. In contrast, known that when Σx = La. both the bias and variance are monotone in the overparameterizze regime (e.g., Hastie et al. (2019)).
10
P P-ly (GD)
P Ey² (NGD)
2.5
P Ex prediction
varlance
2.0
1.0
2
0.5
1.5
2.0
2.5
3.0
3.5
1.5
2.0
2.5
3.0
3.5
Y din
d/n
(a) variance.
(b) bias (well-specified).
Figure 15:
Illustration of the "multiple-descent" curve of the risk for ? > 1. We take n = 300, eigenvalues of Ex as three equally-spaee point masses with xx 5000 and ||Ex 1 = d, Ee Σx (misaligned). Note that for GD, both the bias and the variance are highly non-monotonic for ? >
Additional Figures for Section and We include additional figures on (a) well-specified bias when Eg La (GD is optimal); (b) misspecified bias under unobserved features (predicted by
Corollary 2); (c) bias-variance tradeoff by interpolating between preconditioners (SNR=5). As shown in Figure 16 and 17, in all cases the experimental values match the theoretical predictions.
0.6
1.75
Pa
0.5
1.50
P 22² prediction
0.75
Pi 1
0.4
1.25-
1sk
0.70
0.3
1.00
0.2
0.75
01
Ps 2,
0.50
0.60
0.0
din
magnitude undaserved features
(a) well-specified bias (aligned).
(b) misspecified bias (unobserved features).
(c) bias-variance tradeoff.
Figure 16: We set cigenvalues of Σx uniform distribution with xx 20 and |Ex 11 = d.
0.150
Psi
Pw E prediction
0.6
Pe Ij (NGD)
P- (GP)
0.125
0.5
Pe to 0.50* prediction
0.100
bloe
1.0
0.075
0.050
0.5
0.3
0.2
0.000
1.3
2.0
2.5
3.0
0.0
1.5 magnitude of unotserved features
0.5
1.0
1.3
2.0
2.5
3.0
y= din
y= din
(a) well-specified bias (aligned).
misspecified bias (unobserved features).
(c) bias-variance tradcoff.
Figure 17: We construct eigenvalues of Ex with polynomial decay: {xEEx) = -1 and rescale cigenvalues such that wx 500 and Σx = d.
21

---

Published as a conference paper at ICLR 2021
Early Stopping Risk. Figure 18 compares the stationary risk with the optimal early stopping risk under varying misalignment level. We set Eg = Ex and vary 7 from 0 to -1. As discussed in Section 3.2 smaller entails more "misaligned" teacher, and vice versa. Note that the problem becomes more misaligned, NGD achieves lower stationary and early stopping risk.
Figure 19 reports the optimal early stopping risk under misspecification (same trend be obtained when the x-axis label noise). In contrast to the stationary risk (Figure 4), GD can be advantageous under early stopping even with large extent of misspecification (for isotropic teacher). This aligns with our finding in Section 4.2 that early stopping reduces the variance and the misspecified bias.
P xy²
Zx
P= Iyll
0.6
P
Igl2
statlonary
prediction
carty dopping
stationary
0.5
0.4
0.4
0.3
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4 0.6
0.8
1.0
extent of misalignment
extent of misalignment (-r)
(a) stationary risk.
(b) optimal early stopping risk.
Figure 18: Well-specified bias against different extent of "alignment". We set n = 300, eigenvalues of Σx as two point masses with кх 20, and take Ee = Σx and vary from -1 to 0. (a) GD achieves lower bias when Eg isotropic, NGD dominates when Ex = EØ whereas P = Σx 1/1 (interpolates between GD and NGD) is advantageous in between. (b) optimal early stopping bias follows similar trend as stationary bias.
1.0
--.
P-1
p-1
early stopping Hok
P= 2z
arl sropping rick
1.2
P
23
0.8
Pa 12
P 20
11.0
0.6
0.8
0.4
0.6
0.2
maspecification
maspeciffcation
(a) optimal early stopping risk (aligned misspecified).
(b) optimal early stopping risk (misaligned misspecified).
Figure 19: Optimal early stopping risk vs. increasing model misspecification. Figure 5(c). (a) Eg Ia (favors GD); unlike Figure 5(c). GD has lower early stopping risk even under large extent of misspecification. (b) Eg = Σx (favors NGD); NGD is also advantagcous under carly stopping.
follow same setup as
C.2 RKHS SETTING
We simulate the optimization in the coordinates of RKHS via dinteeddimesionall approximation (using extra unlabeled data). In particular, we consider the teacher model in the form f*(xx)
ΣN hap (ar) for square summable [ in which r controls the "difficulty" of the learning problem. We find {u and [ by solving the eigenfunction problem for some kernel k. The student model takes the form of f(x) Σw a (x)) and we optimize the coefficients (a via the preconditioned update (4. We set 1 1000. d = 5. N = 2500 and consider the inverse multiquadratic (IMQ) kernel: k(x, y)
vor
/1+
Recall that Theorem 8 suggests that for small i.e. "difficult" problem, the damping coefficient À would need to small (which makes the update NGD-like), and vice versa. This result is (quali- that small À is beneficial when is small,
tatively) supported by Figure 20, from which and vice versa. We remark that this observed trend is rather fragile and sensitive to various hyperpa- rameters, and leave comprehensive characterization of this observation future work.
22

---

Published as a conference paper at ICLR 2021
A= le0
A= le0
= le -
A= le -4
population risk
A 1=1e-- 8
population risk
4 le 8
steps
steps
(a) r = 3/4.
(b) = 1/4.
Population risk of the preconditioned update in RKHS that interpolates between GD and NGD. We use the IMQ kernel and n 1000, each curve and thus convergence speed is not directly comparable. Note that (a) large X (i.e., GD-like update) is beneficial when r is large, and (b) small X (i.e., NGD-like update) is beneficial when r is small.
Figure 20:
N 2500, o² 5 10 The x-axis has been rescaled for
C.3 NEURAL NETWORK
Label Noise. In Figure 21, (a) we observe the same phenomenon on CIFAR-10 that NGD gener- alizes better as more label noise is added to the labels, and vice versa. Figure 21 (b) shows that in all cases with varying amounts of label noise, the early stopping risk is however worse than that of GD. This agrees with the observation in Section 4 and Figure 19(a) that early stopping can potentially favor GD due to the reduced variance.
Misalignment. In Figure 21(c)(d) we confirm the finding in Proposition 7 and Figure 18(b) in neural networks under synthetic data: we consider 50-dimensional Gaussian input, and both the teacher and the student model are two-layer ReLU networks with 50 hidden units. We construct the teacher by perturbing the initialization of the student as described Section 5. As r approaches -1 (problem more "misaligned"). NGD achieves lower carly stopping risk (Figure 21(d)), whereas GD dominates the early stopping risk in less misaligned setting (2)())). We note that this phenomenon is difficult observe in practical neural network training real-world data, which may be partially due to the fragility of the analogy between neural nets and linear models, especially under NGD (discussed in Appendix A.
NCD
GD
MED 20:000 METS narnes
prpalation rik
prpaleeion
60
NGD
38 15 si 3% SEa 335 LSE Li 200
10
23
39 Deviation
kosse
Steps
(a) stationary risk (CIFAR).
(b) optimal early stopping
r-/2 (synthetic).
(d) r -3/4
risk (CIFAR).
(synthetic).
Figure 21: (b)(b) Additional label noise experiment on CIFAR-10. (cXd) Population risk of two-layer neural networks in the misalignment setup with synthetic Gaussian data. We set n 200, d 50, the damping coefficient 10-8, and both the student and the teacher are two-layer ReLU networks with 50 hidden units. The x-axis and the learning rate have been rescaled for each curve (i.e., optimization speed not comparable). When r is sufficiently small, NGD achieves lower carly stopping risk, whereas GD is beneficial for larger
23

---

Published as a conference paper at ICLR 2021
D PROOFS AND DERIVATIONS
D.1 MISSING DERIVATIONS IN SECTION
Gradient Flow of Preconditioned Updates. Given positive definite Pand y 1, may check that the gradient flow solution at time i can be written as
0 (l(() = PX I exp ( (- x
! ) 1
Taking t - oo yields the stationary solution 0p = PX (XPX")- We remark that the damped inverse of the sample Fisher P = (xx + ALa) leads to the same minimum-norm solution as GD 0, x (xxXX"" since PX and X share the same eigenvectors. On the other hand, when P is the pseudo-inverse of the sample Fisher (XXT) which is not full-rank, the trajectory can be obtained via the variation constants formula:
- (XxXxxx
((t) =
X T (XX
II!
k=
for which taking the large t limit also yields the minimum- cormmummm--ommmmmo norm solution X (XXT)-1 y.
Minimum ||0||p Norm Interpolant. For positive definite P and the corresponding stationary
solution
0p = PX P (XPXT)- note that given any other interpolant 0 we have (Op
0') PP 0p because both 0 and 0' achieves zero empirical risk. Hence ||0 1| -
100 0pl2 2 P 2 This confirms that 0p is the unique minimum ||0]|-- norm solution.
D.2 PROOF OF THBOREM 1
Proof. By the definition of the variance term and the stationary 0,
v(0) tr (Covv())Ex) o²tr Px x
Write X = Xp1 Similarly, we define Exp = pl/2Ex p1/2 It is clear that the equation
above simplifies
v(pp)= = o (XXx 2 XExP).
The analytic expression of the variance term follows from direct application of (Hastie al., 2019, Thorem 4), in which conditions on the population covariance are satisfied by (A2).
Taking the derivative of m( (e(-) A) yields
(1+ +m(mm))) dH m' (-\) - ( d
,2
(-A)
Plugging the quantity into the expression of the variance (omitting the scaling o² and constant shift),
m' (-x) = (-mm)((->) 1-- ʃ
(1+ (m( - -))²² (xx((T)
m
From the monotonicity of x I+z on r > 0 or the Jensen's inequality we know that
Tm(-A) 1 (1 rm((>)) dHxp(T) S1 -1
(/ 3ml 1+m(-- dH Hxp(())
rm(->) (/)mm(->)
1+ Taking 0 and omitting the scalar σ², the RHS evaluates lower bound V 2 (y - 1)- Note that the equality is only achieved when H HxP is a point mass, i.c. P = Σx In other words, the minimum variance is achieved by NGD. As a verification, the variance of the NGD solution θ p--1 agrees with the calculation for the case were the features isotropic (Hastic et al., 2019, A.3).
1/7. We thus arrive at the
24

---

Published as a conference paper at ICLR 2021
D.3 PROOF OF COROLLARY 2
Proof.
Via calculation similar to (Hastie et al., 2019, Section 5), the bias can be decomposed as
x [aJXX*] [XPX]] (X0" + + - (x 00 +x x
27
E [(0p))
- -m ,00
x: (xTXXT (XPXT) (xXXXT(XP"")
X0- -x + Burnor (ssTo999
( (xPXXT (XPX(XXXT) (XPXT) x x- (XPXT) -1 XPx)
+ Eu,0
(4) Bo (o() d'
(1++ V(0p)).
where we used the independence of à and 0-, 0 in (i), and (A1-3) as well as the definition of the well- bias p0p) and variance V(Op) in (ii).
D.4 PROOF OF THEOREM 3
Proof. By the definition of the bias term (note that Ex, Eg. P are all positive semi-definite)
B(0p) = Rg.
PX (XX((XXT)) xo. 0:
Ex
T Ex [aa PI--X
Eo La- PX 1 (XPX )xx)
d tr
(
d
Ir
E
Ia X ((xX] )11x)
Exp Exp(I1 11- x (xxx)-x
(61)
X2
1 _x 8 + +AAa
_x x + +la)
X-4 d Ir
lim
Σo/p
Exp
71
1
(152)
lim
2²
(1'x 72
x x + x
-2
E-1 Exp
ExpE 1/2
2-0
d
0/P
where we utilized (A3) and defined X
XP1/2 Exp = p1/2XX p1/2 Egjp
p 0-1/2 p 1/2 in (i), applied the equality (AA ) lima-o A + X1) in (ii),
and defined X- XP'/22/1/" in (iii. To proceed, we first assume that Σo/p is invertible (i.e. Amnna (E) p) is bounded away from 0) and observe the following relation via a leave-one-out argument similar to that in Xu & Hsu (2019),
tr
x
x
X
X
+
AE
8/P
(D.1)
12
71
n #
1 T
T
X + 1301P
2;
2
+ #
Lat
(+x)) x + +X00/P1
23
1
[ x + xا
DOTP Ex
1/2
B//P
(6)
(D.2)
1 + ttr (#x
Exp
where (i) due to the Woodbury identity and we defined (+x* T x + +1-11 )
=
0/p
I [ T >EEIP which is independent i; (see (Xu & Hsu, 2019, Eq. 58) for details), and in
25

---

Published as a conference paper at ICLR 2021
(ii) used (A3), the convergence to trace (Ledoit & Péché, 2011, Lemma 2.1) and its stability under low-rank perturbation (c.g., see (Ledoit & Péché, 2011, Eq. 18)) which we elaborate below.
for the denominator we have In partieular, denote Σ = + , x + A.
sup
Σ
1/2 E Exp 1/2)
-11(5120010 201P 1/2 E Xx x 1/2 0/
1/2
1/2
-1
( E )
E x 0/Pp
0/1
2
sup
tr
20 1/2 0 Ex
x
1/2 1
|111
sup
tr
و شيح ) Op 1
where (i) is due the definition of E-, and (AI)A33). The result on the numerator be obtained via a similar calculation, the details of which we omit.
We first note that the denominator can be evaluated by previous results (e.g. (Dobriban et al., 2018, Theorem 2.))) as follows,
1 n
(+xTx+Ma) Exp
e.g. Am (-))
1.
(D.3)
On the other hand, following the same derivation as Dobriban et al. (2018); Hastie et al. (2019), (D.I) can be decomposed
-2
d tr
x
X
X + A0OLP
n2
-1
-2 x X x+NIs
(=xT X + -Ja
tr
Egip
drr tr
Ed/p
dth tr ( x Ala
Eaip
+ A ddl
tr
x
Eoip
(D.4)
We employ (Rubio & Mestre, 2011, Theorem 1) to characterize (D.4). In particular, For any deter- ministic sequence of matrices O E Rddd with finite trace norm, as n, d + oo we have
(
tr
x x n
Itd)
(n)( Exp 2I)
a 0,
in which en (2) -zm(z) for E CVE and m(2) is defined in Theorem due to the dOminated convergence theorem. By (A3) we are allowed to take θ, 1 Lo/pp. p. Thus we have
[ x +da )")
d tr
Co:p
xxxx+N4
tr Eo - ->)) xxa)) +Xaa
(Amp
Uz U
AA -
-Cs
(D.5)
+ ml A)U
in which (i) is due to (A3), the fact that the LHS is almost surely bounded for À > -c, where c is the lowest non-zero eigenvalue of 1 X X, and the application of the dominated convergence n theorem. Differentiating (D.5) (note that the derivative is also bounded on À > -c,) yields
tr A ur (1*xxxx X (xxX++1X) (
Uy U 2p
m'(-X)uzua
Eo/p|
+ E
(D.6)
d d
X(1 ml uA))pp))
(1+ (1+m(-1) m(=))uupp?
Note that the numerator of (D.2) is the quantity of interest. Combining (D.1) (D.2) (D.3) (D.4) (D.5) (D.6) and taking À + 0 yields the formula of the bias term. Finally, when Ep/p not in- vertible, observe that if we increment all eigenvalues by some small > ensure invertibility
26

---

Published as a conference paper at ICLR 2021
Epyp= Ea/p cld. (3.3) is bounded and also decreasing w.r.t e. Thus by the dominated con- vergence theorem we take - 0 and obtain the desired result. We remark that similar (but less general) characterization can also be derived based on (Ledoit & Péché, 2011, Theorem 1.2) when the cigenvalues of Exp and Zg/p exhibit certain relations.
To show that P U diag (eo)UT achieves the lowest bias, first note that under the definition of random variables in (A3), our claimed optimal preconditioner is equivalent to Upp Ug228. We therefore define interpolation U, = aVgU0 + (1 a) for some and write the corresponding Stieltjes transform as m. ( X) and the bias term as B., We aim to show that argmin) B, 1.
For notational convenience define Ja - ma (0) and hn - ma (0) One check that
ho 4o (a)
Un
ho
dm (-
(I+ha) 1
B, = E (1 ha)
ha )2
da 1-+0
(1 a)
h (1+h,
We now verify that the derivative of Ba w.r.t. is non-positive for E [0, 1]. A standard simplifi- cation the derivative yields
dB da
(ga hay )
ha
Ja
ha h )2
2
h2
2E
E
2 E
E
(1 ha)
(1 + h )2
(1
heg
Na (r ho)
9a ha
(1 ha
+ 4E
E
(1 ha)a
+ Na
)2
(ga ha)²
hg
Ja
ha
ha
AV E
E
E
(1 ha)²
(1 + ha a
(1 + ha) )
(1
hay
Nn (gte ha)) )
Do
ha
ho
(18)
+ 4E
E
S 0,
(1 + ha)
(1 + hx)
(1 + ha)
where (i) due AM-GM and (ii) due to Cauchy-Schwarz on the first term. Note that the two equalities hold when Jo. = has from which one can easily deduce that the optimum is achieved e. U,U0, and thus know that the proposed P is the optimal preconditioner that is codiagonazable with Σx.
when Usp
D.5 PROOF OF PROPOSITION 4
Proof. Since Eg Exx we can simplify the expressions by defining V. & h and thus uo = h². From Theorem we have the following derivation the GD bias under (A1)(A3),
E h
B(θr)
m mi (1 h 22
E
hh
h.m.)
(D.7)
7
(homy) hommy
where my limx-o, m(-A), and m satisfies
1
h
+ 7E
ml X)
1 1 h ml (2)
Similarly, for NGD (P xx²) have
E Altr
B(0
ml
E
h- h²
[1+ma)]
Eh Eh 1+r
(D.8)
m (1 + m2)
1 EE (1+
m2
(1 + m2) ym3
where standard calculation yields m2 (v 1)-1 and thus p1) + (1 (1 E
To compare the magnitude of (D.7) and (D.8), observe the following equivalence.
B(0r) $ B(0
27

---

Published as a conference paper at ICLR 2021
hltr
(h.m,) m1)
+ E
y-1 $
1-1E 1- (1+ h (..mm)
Eh Eh
(1 + h m
S
1
E (1+
$
E
BD +
(D.9)
1+5
()+
1+
where (i) follows from the definition m, and we defined h is not point mass, we have
Ç - h m. Note that when r S and
clt E S E C) I+C
E
E B
E I I+c
(1+
I+C
(1+
On the other hand, when 2 0, following the exact same procedure we get
5 E
E (1 (1+
1+ç
<
()
B
1 1+
Combining the two cases completes the proof.
D.6 PROOF OF PROPOSITION 5
Proof. We first outline a more general setup where Pa = f(Ex;a) for continuous and differen- tiable function of and f applied to eigenvalues of Eg. For any interval I C [0, 1], we claim
(a) Suppose all four functions =f((a)))' f(z; 0f(z:))
da /f(x; and x af(aa)))))
da are decreasing
functions of the support of eg for all a E I. In addition, @f(xx)) > 0 the support do
of Uz for all a E Then the stationary bias an increasing function of a on
(b) For all I, suppose x [x;; a) is monotonic function of on the support of u2 and ( do ff;;; a:) is a decreasing function of x on the support of tg. Then the stationary variance is decreasing function of on I.
Let us verify the three choices of Pa in Proposition 5 one by one.
the corresponding f(z; a is (1-a) a+) + ar. This satisfies all conditions in (a) and (b) for all a e [0, 1]. Hence, the stationary variance is decreasing function and the stationary bias is an increasing function of a E [0, 1].
When Pa = (I-a)aa+aa (xxx) >'
When P. = (Ex)
the corresponding f(z; a) is x a is clear that it satisfies all conditions
in (a) and (b) for all E [0, 1] except for the condition that 0f(za)) ----a In x is decreasing function of x. Note that x Ba
do
-x 1-a -x'-o In x is decreasing function of x on the support of Ur only for a 2 In(x)=1 where k = sup / / inf ug. Hence, the stationary variance a decreasing
(f(za))
Im((s)
function of a E [0, and the stationary bias is an increasing function of a E [max(() |nin)))) ), 1].
When P. (1Ex+(1-)) the corresponding f(z; a) is 1/(az + (1 ()) which satisfies all conditions in (a) and (b) for all a [0. 1] except for the condition that x (f:o)) Da
x(1-x) (1=)( is decreasing function of z. Note that x: [f(x;a)) (a=+(1aa)² is decreasing
=(1=x)
function z on the support of eg only for a 2 x-2
x-1 Hence, the stationary variance is a decreasing
function of [0, 1] and the stationary bias is an increasing function of [mnax(0, x=)), 1].
To show (a) and (b), note that under the conditions on Eg and Σg assumed in Proposition 5, the stationary bias B(Âp.) and the stationary variance ((Qp.)) can be simplified to
B(Op.) = m (0) E
2x
and V(θp.) = =²
m', (0) m3 (0) - -1).
(0) (1 + Uz am)
where (2) and m/, satisfy
Uz uef( (
= === + ,,
(D.10)
U2 (Uaj (2)
28

---

Published as a conference paper at ICLR 2021
m (z)
(D.11)
m (z)
>E
f(uz:a mm. (x)
2"
I+J (+
For notation convenience, let fa Un (Vo) a). From (D.11), we have the following equivalences.
B(Qp.)= 1 y1F
(I+Jama (0))²
(D.12)
fa (0) [ala
v(Op.) 0²
1
(D.13)
1
7/3
Lama (0)
[+Jama(0)
We first show that (b) holds. Note that from (D.13), we have
2
OV (V =7o²
2fa ma
(
dma(z) da ==0 da +
0fa (m
maa(()
da
1 1->E
Sama (0)
(1+ (m)
Lo (0)
(D.14)
To calculate ama (z)
take derivatives with respect to both sides of (D.10),
==
(Ia 0ma(z)
0fa
ma (0) [0)]
(D.15)
(1 fama (0))²
da
=== Da
Therefore, plugging (D.15) into (D.14) yields
2
(VP.).)
ma(0)
fa
-1
=270²
da
(u)
2
(1 + ma (0))
yB
I+Jama
fa
8fa öa
fa
8fa
E
E (1 Sam. (0))
(1 + Sam,, (0))
(1+ fama(0)) (0))
(1+ sa m. (0)2
Thus showing V(Âp.) ) is decreasing function of a is equivalent to showing that
fa (1 Let Pl be the probability measure of Ug-
Bfa
af fa
fa
E
E
da
2E E (1 + fam. (0)) (1 m. (0))
(D.16)
fam..
(1 f. (0))
We define a new measure D [I+ ( and let
fa
C follow the new measure. Since 8f(m)) monotonic function of x,
/f(r; is decreasing function of x and ff a) is
da
(
86, ((0:ia)
88, ((02:0))
1 + Uz (fa (0)
02 (ūx; a)
1 + 02 (zi mm
Changing üz back to ex,, we arrive at (D.16) and thus (b).
For the bias term B(@P.). note that from (D.10) and (D.12), we have
0B(0P.)
1
fa
E
da
7
?
sama(()
(I.
@m.(2)
Ifa
mm.((())) (0) (01)
fama
EE
2
1=- +
mn
E (1
(1 fa m. (0))
da
da
(0))2
+ I
Ux
E 2
(a)
Sa
dma(2)
+
0fa
ma (0)
(1 + m (0))
[
(1 fama
da
1=
aa
(D.17)
29

---

Published as a conference paper at ICLR 2021
Similarly, combine (D.15) and (D.17) and simplify the expression. To verify B(Op.) is an increasing function of a, we need to show that
Uz ma (0)
a/a
Up
fa
fama (0)
E
E
E
da
E
E
(1 fa
(L + fa (0))
(ama(0))
(1 fama (0))²
(1 + fama (0))²
E
Ux
E
(fama (0))² E
af
de
E fam (0)
Bka 0n
E
fa (0)
(1 + (O))(
(1+ Sam (0))3
(1
sa (0))²
(1 + La (0))3
Samn (0))²
(D.18)
Let Na A Sam, (0) (z a)ma (0) and ga simplified to the following equation
da
Ux
0f(r2))) da
Then (D.18) can be further
hh E (1 + ho)
Ja
ha
Up
Ja
0 S
E
E
E
E
)
(1
+ ha)
(1 + ha)s
(1
+ ha)
(1 + ha)
(1 + ha )3
part
Ux
E Jaha
ha
Vxaa
ha
ha
+ E
E
E
E
E
(1 Nay )3
(1 ha) )3 (1 + ha)
(1 hay )3
(1 ha)
(1 Ma
part
Jaha E
ha
Ug
h2
ha
+ 2E
E
2E
E
E
(1 ha)
(1 ha)3 (1 + ha)3
ha) 33
(1 + ha)
)3
(1 + ha) 3
part
+ E Vsha
E Jaha
h
UrBa
h2
h2
E
E
E
E
(D.19)
(1 + ha) )3 (1 ha )3 (1 ha)
(1
hay
+ ha
+ ha
part
Note that under condition of (a), we know that both ha and U ha are increasing functions of Uzi and both 9a/ h, and Ja are decreasing functions of ug. Hence, with calculation similar to (D.16), we know part 1,2,3,4 in (D.19) are all non non-negative, and therefore (D.19) holds.
Remark. The above characterization provides sufficient but not necessary conditions for the mono- tonicity of the bias term. In general, the expression of the bias is rather opaque, and determining the sign of its derivative can be tedious, except for certain special cases (e.g., Y = 2 and the eigen- values of Ex are equally weighted point masses). We conjecture that the bias is monotone for a E [0, 1] for a much wider class of Ex. as shown in Figure 22.
0.65
0.3
0.55
20
-0
0.50
0.1
unform
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.6
1.0 NGD
GD
interpolation coefficient
NGD
interpolation coefficient
(a) 7 2.
(b) 7 5.
Figure 22: Illustration the monotonicity of the term under Eg = Id. We consider two distributions of eigenvalues for Xx: : two equally weighted point masses (circle) and a uniform distribution (star), and vary the condition number xx and overparameterization level y. In all cases the bias in monotone in E [0, 1].
D.7 PROOF OF PROPOSITION 6
Proof. Taking the derivative of V (00((t)) w.n.t time yields (omitting the scalar g²),
dt (In-eep(- _ =xxxx)) ( (xpx"))
dV (v)))(t)) d 11/ px 7 (en-esp/
dt
F
30

---

Published as a conference paper at ICLR 2021
Exp X x0 Sp cxp ( Sp SP 2
t (-1Spp
(68)
I, exp
Sp
X
0,
n
n
n2
p.s.d.
where we defined X = Xppl/2 and Sp = xPx T in (i), and (ii) is due to (A2-3) the inequality
tr(AB) 2 Amin (A)tr(B) for positive semi-definite A and B.
D.8 PROOF OF PROPOSITION 7
Proof. Recall the definition of the bias (well-specified) of 0p(t).
(t))
(
tr
(o(()sp'x) P Eg Ia Px
Ex (I=
(a-x' (P((s)p')) (I==X\W((Ssx')) Exp Exp(Ia- (a-x x x WP(()5pxX)
d tr
Eo/p
(aa) 1
(>>/))
Ia - X (P))p X -//p7) )
2 d tr
(D.20)
where we defined Sp XPX W (t) = I, exp( -p)pp))pppppppppppppp)p)pp) in (i), X Xp1/2 in (ii) and
(iii) is due to the inequality tr A A 2 tr (A*).
When Ex= Eg i.e. NGD achieves lowest stationary bias, (D.20) simplifies to
((0p(1)) 2 ( (I_=X\(005xx))
Σ
--xx
2
exp
(D.21)
t=1
where i
is the eigenvalue Sp. On the other hand, since F = Ex, for the NGD iterate 2
1 jor B(O (())=
x WF WF 1
- (D.22)
1
exp
i=1
where X X 1/2 and Ā is the eigenvalue of Sp-1 X x
Comparing (D.21)(D.2)), Amaa t (note that T/t = O(1) by Aman
see that given (t()) at fixed t, if we run NGD for time T >
(A2-3)), then we have (()) 2 (T)) for any P satisfying (A3). This thus implies that Bopt (0p) 2 i..
On the other hand, when Eg Ls. we can show that the bias term of GD monotonically decreas- ing through time by taking its derivative,
d dt
tr (Ia- x W x) Ex IIa x W ((())'')))
((Or(t)) = d dt
1 =- tr nd
(Ia-)xxx((xxx' x (Ia-x>W(())5)> x) s x (1a
Ex x
Sexp
<0.
(D.23)
n
p.s.d.
Similarly, one can verify that the expected bias of NGD is monotonically decreasing for all choices of Ex and Eg satisfying (A2-4),
d tr Eo Ia II-=F-1 x W -t) xx) Ex (Ia- - 1i---lx x W (t)55))p"
dt
d dt
tr Exo Ia (Ia-X x W (1)5 -1))5(11)))
x) [Ia (Ia- x Wr (1155 .x) 2 00
31

---

Published as a conference paper at ICLR 2021
where (i) follows from calculation similar to (D.23). Since the expected bias is decreasing through time for both GID and NGD when Eg Id, and from Theorem 3 know that B(0r) S B P- we conclude that BoDs (01) I Bopt(O p-1).
D.9 PROOF OF THEOREM 8
D.9.1 SETUP AND MAIN RESULT
We restate the setting and assumptions. H an RKHS included in L2(Px) equipped with a bounded kernel function & satisfying suPsupp(Px) k( )) S 1. K E H the Riesz representation of the kernel function k(z, -), that is, k(x, y) (K2, Ky )
H- S is the canonical embedding operator from
H to L2(Px), We write Σ S: H H and L S8* Note that the boundedness of the kernel gives ( S sup= [=)) sup= (K2, f' s || ///l///// < || Hence we know
||>]|| S 1 and |A]| S 1. Our analysis will be made under the following regularity assumptions.
(A4). There exist E (0, oo) and Ost. f = L for some h" L2(x) and Il S.M
(A5). There exists s s.t. tr [ /+) < oo and 2r + s- > 1.
(A6). There exist μ [s 1,1,1 1] and C₁ 0 s.t. supsesupp(Px)) | ||
S C.p.
(A5)A66 are standard regularity assumptions in the literature that provide capacity control of the RKHS (e.g., see Caponnetto & De Vito (2007); Fillaud-Vivies et al. (2018)). For (A4), it is worth noting that previous works mostly consider 2 1/2 which implies that EH.
The training data is generated as y/ = (ax) + E where & is an i.i.d. noise satisfying |e,| S almost surely. Let y E 12 be the label vector. We identify RR' with (P,x) and define
Ka; & Kz, 7 - + H, SYY
Σ 1=1
YKK,, (Y (2(P))).
We consider the following preconditioned update on ft E Ht:
f- - X1) 1==-----5------00000--------
fo
We briefly comment on how our analysis differs from Rudi et al. (2017), which showed that a precon- ditioned update (the FALKON algorithm) for kernel ridge regression can also achieve accelerated convergence in the population risk. We emphasize the following differences.
The two algorithms optimize different objectives, as highlighted by the different role of the "ridge" coefficient X. In FALKON, À turns the objective into kernel ridge regression; whereas in our (4.1), X controls the interpolation between GD and NGD. we aim to study how the preconditioner affects generalization, it is important that we look at the objective in its original (instead of regularized) form.
To elaborate on the first point, since FALKON minimizes regularized objective, it would not overfit even after large number of gradient steps, but it is unclear how preconditioning impacts generalization (i.e., any preconditioner may generalize well with proper regularization). In con- trast, we consider the unregularized objective, and thus early stopping plays a crucial role in avoiding overfitting this differs from most standard analysis of GD.
Algorithm-wise, the two updates employ different preconditioners. FALKON involves inverting the kernel matrix K defined the training points, whereas consider the population covariance operator Σ, which is consistent with our earlier discussion on the population Fisher in Section 3.
In terms of the theoretical setup, analysis allows for < 1/2, whereas Rudi al. (2017) and many other previous works assumed r E [1/2, 1), as commented in Section 4.3.
We aim show the following theorem: Theorem (Formal Statement of Theorem 8). Given (A4-6), if the sample size n is sufficiently large so that 1/(mX) < then 17 < |Σ| with nt 2 and 8 1 and 0 < X holds that
|| 100---1l 1' C(B(() V()),
32

---

Published as a conference paper at ICLR 2021
with probability - 38, where C is constant and
B(t) exp(-ntt) (7)
2r
7/6
B11 o Ez )=
M1
V(O) = +1
+
log( (
n²
in which
A 7//
A 2r
B1 12r-2))r
Σ x-
B
x (1 x
2 ") )
V1(()= =
exp(-nn)
+ (tn))
+
(1 tn)
n1
n²
(/)))))))
for 3' = log
When 2 1/2, if we ser X -.ाा =: X and
t ((log(n)), then the overall convergence rate becomes
Sg - 1 Öp (n Iree1
which is the minimax optimal rate (Öp() hides poly-log(n) factor). On the other hand, when except the term (t(t). In this case, if 2r p holds which again recovers the optimal rale.
< 1/2, the bound is also Öp n I
additionally, we have V.(1) Op
2
Note that if the GD (with iterates fi) is employed, from previous work Lin & Rosasco (2017) 2r is replaced by and therefore the upper bound translates to
2r
know that the bias term
T/
J[Sf. - 1' S &ch C + ((x cr (x(())/6) (1t)2//
(nt))
+
ntt n2
0²
+
2r
M (nt) (2r-
rpt
with high probability. In other words, by the condition 7 = 0(),, need (n steps to sufficiently diminish the bias term. In contrast, the preconditioned update that interpolates between GD and NGD (4.1) only require & (l)g(nn) steps make the bias term negligible. This is because the NGD amplifies the high frequency component and rapidly captures the detailed "shape" of the target function f*.
D.9.2 PROOF OF MAIN RESULT
Proof. We follow the proof strategy of Lin & Rosasco (2017). First we define reference optimiza- tion problem with iterates Se that directly minimize the population risk:
J, J_ </E AT) -
Jo = 0.
Note that B[s{] = Jr. In addition, define the degrees of freedom and its related quantity
Noc (A) = Eg (KK, VI Ka)Hd =r(2)))) (A) ==
|1= K 1. (P))
sup
We can see that the risk admits the following bias-variance decomposition
S 1' $ 2([sSt- - + lIf f 1 (2
((t), variance
B(t). bias
We upper bound the bias and variance separately.
33

---

Published as a conference paper at ICLR 2021
Bounding the bias term Btt): Note that by the update rule (4.1), it holds that
S/. - /-
- - (*)
Sf - fr = -1 (4-15(2++x44-))(ff===
Hence, unrolling the recursion gives -
(T-
XT) (-)="))((-")) = - n + \ʃ) - Write the spectral decomposition of as L L= - 05
3 for Oj L2 foro 2 0. We have (14=n)(Z++V) y 11th"lleaa
L (l))))(Px) =
E)=(( To,4x )210 2010gh1g where = = Ky0z. We then apply Lemma 11 to obtain
B(I S
Σ
R2 +
2r Tp
2r
n j.yyA
Σ
exp(-ntt) v (
|| 11112(20))
77t
joy2A
where C' a constant depending only on r.
Bounding the variance term V(t): We now handle the variance term V(t). For notational con- venience, we write A₂ == A + AI for a linear operator A from a Hilbert space H to H. By the definition of ft, know
ft = (1 - -- + 2) - f- (
1
=> (I n/ -1
3=n
1/,
Σ (1 11ー
Σn 112 =.
1 1/ أエ 1/2y]. Accordingly, we decompose V(t) as where we defined G2 := 1) Ej=b( -
[| - Sfi (2x)())) <2([S(ffi Σx 1/-1/22 G,,,, 12E
(a)
||(( 1/2 ",
G, 1/2 - F00)
We bound (a) and (b) separately.
Step 1. Bounding (a). Decompose (a) as
||( - -x 1/1 "3 - EFA)lL S1 1/2 " We bound the terms in the RHS individually.
- E
(i) || 1/2 12 = |> 1/2 ン 11
1.
X
(ii) Note that 1 - 1/2 = II
(1 "" (1-
and 0 5 1 Proposition 6 of Rudi & Rosasco (2017) and its proof implies that for À S |Σ]
(A) 2 + Fx)) E
-
1/2 1
(D.24)
n
3n
with probability 1 - 5, where 3 log
[tt((EE5")
log LN (N)
( By Lemma 14, 3 s
8
4er Atr(x:/>))
1/a and Foc (A) S A Therefore, if X o(n -1 log(n)) and A Q(n 1/*),
log
the RHS can be smaller than 1/2 for sufficiently large n, i.c. E, = 0 ((Voog(n\\(m))) < 1/2; thus,
Σx 1 1/2 2 2 I.
34

---

Published as a conference paper at ICLR 2021
We denote this event E1- (iit) Note that
G.E Ex
Σ 11 1- 1 イイ 1/255
Σx 1/2 1/ Ex 1/2
"
"
Σ
[I - 1
1 1/200
(Ex 1 1/2 + -1ズz))
Thus, by Lemma 12 we have
|G
1/2
xx
1/2
A
$
17
ΣΣ 1/20 Dx
--1---111 1/2
+ 2
(I
-
1/2 -
1/29
NET
1=0
1 (due Lemma 12)
[-1
$1 + 1
|1(I -
12
s 1 nt.
j=0
(iv) Note that
|Σ Y
- STO s 22 -5.) - +
1/2 18
First we bound the first term of the RHS. Let & = 1/ Then, (5 is an i.i.d. sequence of zero-centered random variables taking value in H and thus
1 - -
1 &
Σ
n
The RHS can be bounded by using Bernstein's 's inequality in Hilbert space Caponnetto & De Vito (2007). To apply the inequality, we need to bound the variance and sup-norm of the random variable. The variance can be bounded as
|| S ((
1/2 (Kn( - [((=)) + Ket)1
< 2 E (m.y)
|1
1/2 (K2 (x) - f(())) + || 1/2 (K )
1/2
22
sup
| 1/2 K 11 - St S + +o
+o²tr(
#Esupp(Px)
< 2 (A) + + (
s B(t) +
The sup-norm can be bounded as follows. Observe that fe S fell and thus by Lemma 13,
| 2
sup zEsupp Px)
| 1/2 Ke (0 + |f + llfelloo)
( + M + (1 + tn) t7)
1/2 M + bn) -112--))))
Hence, for 0 < < 1, Bernstein's inequality (Proposition 2 of Caponnetto & De Vito (2007)) gives
2
X-1 + (
X + (1 tn)
0
log(1 (
n
35

---

Published as a conference paper at ICLR 2021
with probability 1 - 5 where C' is universal constant. We define this event as Eg.
For the second term 1 Efill3 we have
(5*
1 - - Sfell
S B(t).
Combining these evaluations, on the event Eg where P(E2) - 8 for < 8 < we have
|Σ 1 - -
2 log(1/0) B(I)
( 30
A-1 + ( A 1- M (1 tm) (
where we used Lemma 14 in (i).
Step 2. Bounding (b). On the event & the term (b) can be evaluated as
1/ f()
< 1
1:)
</E01/23 (GEE // |) fell
< </ 1/2 |1[ 13555 1/2 T) 1/2 felli
< <| 1/2 - ""
-1)= fell
(D.25)
where we used Lemma 13 in the last inequality. The term be bounded follows. First, note that
[| "/> - 1)>: fell can
[-1
(G.E,12350 1/2 1/
Σ
T - 1 1/2 1/355 1/3gs
Σx 1/2555555 1/2 I 1/2
1
(I-nz) 1/2 "
1/y
Therefore, the RHS of (D.25) can be further bounded by
- - 1/2 12',, // /2 fill
- - 1/2 Σソン
+ + 1/2(2 - 1 frllxx
fellx
[-1
=| -1 (T - E
1/2 Σイ 1/2 1/2 (n) 1/2 -Σ1Σ 1/2)(1
- 1/))1//((1 1) -(I (I- -n " ' fell
k=0
[] 1 " 1/2 2
[-1
+1 Σ ||(1 n2;
125,
(1 1/2
1/2+r
- fell
k=0
- 3| - 1/2+ /1|/
=||(1 - 1|
- ΣدΣ 1/2+ ||1|
fellx
SII(T - ~ ーリ + 10 1/2+r [1 + + (11+t)\lA*"|>a(( (<x(()))
(D.26)
where (i) is due exchangcability of Σx and Σ. By Lemma 11, for the RHS we have
1 (1A)²
|| - (リーーーーーン)ンN S <xpp) V e mt
Next, in the (D.24), by applying the Bernstein inequality for asymmetric operators (Corollary 3.1 of Minsker (2017) with the argument in its Section 3.2), it holds that
| 1/2 , 1/2+r
36

---

Published as a conference paper at ICLR 2021
V + (((1+A)) + 222 A V X- 1/23
SO
=
n
n
with
probability 1 (2
8, where
C' is a universal
constant
and
8
"VA e ( ( 1/1
log
We also used the following bounds on the sup-norm
and the second order moments:
(sup-norm)
1 - - 1/2+
< $1> Ko 1/22 || + |
///"//20//0 1/2 K K3 1/2111222//222 E + |Σ
$ O2, (1 V X /121 + (1 2) (a.s.).
(2nd order moment 1)
| 1=5 1/2 ( K2 Σ>ΣΙ 1+2 KK ΣΙΙΣ 1/2
| 1/2 ΣΣE 1/21
K 1/2 //2 -n Σ 1/2+1/2 K
sup *Esupp(Px)
(2 V )²
(2nd order moment 2)
Ex ND 1/2tr (K K2 - Σ 1/2 //- (Ka ΣソΣズ 1/2t"]1
< || - 1/22t K K k3 K - 1/2tT]]
S C2 V 2 EB K K2 Kz]
C2 2- )tr ΣΣ,
C2 2 1) Woo
We define this event Eg. Therefore, the RHS of (D.26) can be further bounded by
[I| n2 2) Citll Σ 1/2 (Σ وز 1/2+7 II[(1 + bn) [1es(())
(1 + t7)l1h"|1 exp(-nt/2) v (1))) + tn
nt
Finally, note that when x == X* and 2r 2 14
-/ ò 1-2r-H-11 ( )=
S Ö(n 2
o(en 22)+2 2
S Ö(n - 20071) ).
n²
Step 3. Combining the calculations in Step and leads to the desired result.
D.9.3 AUXILIARY LEMMAS
Lemma 11, For t EN and 0 & X,
holds that
1 1 (1-1 o (1-1otN²
exp(-n!!
(a 1 A) <<< A)
o" <
2r 2 T/
Proof. When o 2 A, we have
1 1- (1- o" < (1-1 20
o" 00000"
(1 = cxx((tnn22j0 S (xx((-tn/2)
due to σ S 1. On the other hand, note that
(1 n σ )' o" s exp
o
ont
( (I+ta)
-1pl
x
x
o
nt
37

---

Published as a conference paper at ICLR 2021
0 (
( )
< sup expl (u)))
17 x
2>U
n/
where we used sup=>0
(epp=>xepppppppp) exp(-x)x - (r/e)*.
Lemma 12. Fort t = N. 0 < 1j and 0 S a such that D10 < 1, it holds that
not s 1.
Proof. If o = 0, then the statement is obvious. Assume that o > 0, then
11 - (1-177) 0= == [1 (1 >
Σ (1-noo 1- n0) 17
3=0
This yields the desired claim.
Lemma 13. Under (A4-6), for any < and y S r, holds that
" (1 + A (1/22((00-1))) + -tppattnn/----- -3//2+(9--1)+) (1)(()))
Proof. Recall that
ft - 1
Σ (
Therefore, we obtain the following
|Σ "Ex"fellee follx = n/l 2-1
1=0
(I- - m 1,,
+ Xml
プリー
Σ
(I E
"
Anl
Σ
(I
S'L''l
1=0
1=0
<|Σ 1 h"' +
L115" q- 1 1 + Atnl -- 1|
S (h', Ln," 1+ S8
2tr 120) L2(xx) + \tn/ (h,, L,,""""" 2 LS 9- 2tt
2tr
- (h', -V 1, 1 LL
1-1 + Atn (1i, 11', 21, 9 Lz(Px)
<(1 112--1001-1 Xtya
3/22 (a- (1 S (1+ - 1/2- (1- -7) || [x()))
Lemma 14. Given (A4-6) and X (0, 1), we have (o) S tr(2:1/))> x- 1/5 and Foo S 1/2
Proof. For the first inequality, we have
(A) =t((Ex=)) = tr 11// 1/o) str <tr(2>1//0) //:1////1111//
(1-1111) (11//0) > x 1/1 < tr (2:1// -1/0
As for the second inequality, note that
Fox (A) = sup(Kz Ej Ka)# S sup - (Kn: Ka) S sup x) & XA-
38

---

Published as a conference paper at ICLR 2021
D.10 PROOF OF COROLLARY 9
Proof. Note that in this setting Uz takes value of T=x
2_ and I+E with probability 1/2 each. From
2x
(D.8) in the proof of Proposition 4 one can easily verify that for NGD,
B(O 1(0p-1) 21 + n²tr))
(1 (1-
+ +1+
For GD, the bias formula (D.7) can be simplified as
2K
B(θ,)
1 -1 (1+ 2m (1 + K 2xm²)
I+h
+
I+x
(1 K + 2m1)²
me
+
(1 K + 2km1)²
P377
where m, is the Stieltjes transform defined after Equation (D.7). From standard numerical calcula- tion one can show that when y > 1, K 2 1.
(n + 10+1) + 1)2 + 4(1 7) - 1)a + (2 - -16 (+)+))))) 1)²
8(7 1)k
Setting B(Or)= B(0 and solve for r, we have
C1 C2
r" In Cx.7/ In K. CN,y 01 63
(D.27)
where we defined
6 1
02
K ++1
(1+ + 2xm;)²
C3 mi (1 + + 2km,) + km, (1 + k. + 2m
s(1 + K + 2m1)²
C1 mi m1 2km,)² )2 + km + K + 2m
Hence, Proposition 4 (from which we know r' E (-1, 0)) and the uniqueness of (D.27) implies that when r 2 r-, B(O) S B0O, 1), and vice versa. Finally, observe that in the special case of 2, Therefore, one can check that constants in (D.27) simplify to
= x#1 2/R"
1- VR.
- = VE(VE -1)
C1 C3 =
+1'
2(x + 1)
which implies that "" 1/2.
D.11 PROOF PROPOSITION 10
Proof. Part (a) is simple combination of (Bai & Yin, 2008, Theorem 2) and assumption (A3),
which implies || |Xx llz and |Σ ||2 are both finite. For part (b), the substitution error for the variance
term (ignoring the scalar o² can be bounded
V
VI
-xx
- ir F
[XF
xT) -2 Σxx
(6) F F 1||
_0(
tr x S XFEX)
x + Vdl x S- xx]
ΣxF (1,))
+
XF
ExF 5-2 S 3/1 O(e)
39

---

Published as a conference paper at ICLR 2021
where we defined = XF ' S = X in (i) and applied tr(AB) 3 Amax +
A (tr(B) for positive semi-definite B, as well as tr(AB) S vdI B and (ii) is due to (A3), vb 1, (Wedin, 1973, Theorem 4.1) and the following estimate,
n -
< In.ss S
nn
1-
nbu S
112)
0() na S 1
S
S/no S S/nu
1
(1)) O(c).
where (i) again follows from (Wedin, 1973, Theorem 4.1), and (ii) is due to (A1)(A3) and & > (which implies ||n,,--12 la and |n... 2l are bounded a.s.). Finally, from part (a) we know that ( suffices achieve c-accurate approximation of F in spectral norm. The substitu- tion error for the bias term can be derived from similar calculation, the details of which we omit. O
E EXPERIMENT SETUP
E. PROCESSING THE DATASETS
To obtain extra unlabeled data to estimate the Fisher, we zero pad pixels on the boarders of each image before randomly cropping; random horizontal flip is also applied for CIFARIO im- ages (Krizhevsky et al., 2009). We preprocess all images by dividing pixel values by 255 before centering them be located within [-0.5, 0.5) with the subtraction by 1/2. For CIFAR10, downsample the original images using max pooling layer with kernel size 2 and stride 2.
E.2 SETUP AND IMPLEMENTATION FOR OPTIMIZERS
In all settings, GD uses learning rate of 0.01 that is exponentially decayed every 1 updates with the parameter value 0.999. For NGD, we use fixed learning rate of 0.03. Since inverting a
parapettttttte-szeeeee Fisher estimate per iteration would be costly, we adopt the Hessian free approach (Martens, 2010) which computes approximate matrix-inverse-
products using the conjugate gradient (CG) method (Hestenes et al., 1952). For each approximate inversion, we CG for 200 iterations starting from the solution returned by the previous CG run. The pre- cise number of CG iterations and the initialization heuristic roughly follow Martens & Sutskever (2012). For the first of CG, we initialize the vector from standard Gaussian, and CG for 5k iterations. To ensure invertibility, we apply very small amount of damping (0. (0.00001) in most sce- narios. For geometric interpolation experiments between GD and NGD, we use the singular value decomposition to compute the minus power of the Fisher, as CG is not applicable in this scenario.
E.3
OTHER DETAILS
For experiments in the label noise and misspecification sections, we pretrain the teacher using the Adam optimizer (Kingma & Ba, 2014) with its default hyperparameters and learning rate of 0.001.
For experiments in the misalignment section, we downsample all images twice using max pooling with kemnel size 2 and stride 2. Moreover, only for experiments in this section, implement natural gradient descent by exactly computing the Fisher on a large batch of unlabeled data and inverting the matrix by calling PyTorch's torch. inverse before right multiplying the gradient.
40


Reviews:
[{"review_id": "VW-v5iTMD12", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \n\nThe authors theoretically study the prediction performance of pre-conditioned gradient descent/flow with linear models and squared loss aligning in the setting of least squares regression and non-parametric regression. For parametric least squares, the predication performance of the limiting solution for preconditioned gradient flow i.e. time goes to infinity, is studied in an asymptotic regime where both the number of samples and dimension go to infinity in proportion to one another. Meanwhile for non-parametric regression, source and capacity assumptions are leveraged to achieve finite sample guarantees. Experiments are also conducted on neural networks in a student and teacher setup. \n\nSummary of main Contributions: \n\nA1) In the case of parametric least squares, an asymptotic characterisation of the test risk is utilised to study the limiting solution of preconditioned gradient flow. Preconditioning with the inverse Fisher information matrix (covariates population covariance) is shown to achieve the optimal variance among preconditioned updates (Theorem 1).  Meanwhile for the asymptotic bias, the optimal pre-conditioner depends upon the covariance of the ground truth parameter. In a mis-aligned case, where the ground truth covariance is equal to the inverse of the population covariates covariance, the optimal pre-conditioner for the bias aligns with the inverse Fisher information matrix (Theorem 2).   \n\nA2) In the case of an Isotropic covariance for the ground truth parameter, it is found that the Bias and Variance can be traded-off by interpolating between the two aforementioned pre-conditioners (Proposition 4).\n\nA3) For non-parametric regression, gradient descent pre-conditioned with the inverse regularised population covariates covariance is considered. Mini-max optimal statistical rates are achieved with a number of iterations that grows logarithmically in the data set size i.e. linear convergence (Theorem 7).\n\nA4) Experiments for neural networks are conducted in support of A1). Specifically, gradient descent pre-conditioned with the Fisher information matrix achieves better generalisation performance when the noise is large or the model is misaligned (Section 5). \n\nA5) For parametric least squares with a mis-aligned ground truth parameter, it is shown that early stopping with NGD achieves lower Bias than any other pre-conditioned gradient descent (Proposition 6).  \n\nPros: \nB1) I feel contribution A1) in conjunction with A5) is novel and offers a precise interpretation of when pre-conditioning with the inverse Fisher information matrix can yield an improvement in performance. \n\nB2) Contribution A2) is also interesting and can point towards understanding and controlling the implicit bias of gradient descent through the pre-conditioner i.e. taking a linear combination of two pre-conditioners. \n\nB3) Contribution A4) supports the findings in A1) in a setting beyond least squares.\n\n\n\nConcerns:\n\nC1) The authors do not compare their theoretical results for non-parametric regression (contribution A3) ) to prior work within the literature. Specifically, reference [1] where the generalisation performance of a pre-conditioned gradient method is considered. To remedy this, I feel the authors should discuss how their theoretical results and proof method differ from [1] as well as the novelty of their approach. \n\nC2) The theoretical results and discussion focus on a particular type of pre-conditioner: the inverse population covariates covariance and transforms thereof. This limits the applicability of the insights as this quantity is often not known in practice. Similarly, the experiments are in a setting where Fisher information is estimated accurately using 100,000 samples while training uses 1024 samples. In contrast, prior work for non-parametric regression considers pre-conditioners involving estimates of the population covariance [1]. To remedy this, I feel the authors should include a discussion on how their insights i.e. A1), A2) are impacted when the population covariance is swapped for an estimate (using unlabelled data).    \n\nC3) The manuscript can be difficult to read. For instance, the authors start with a time varying pre-conditioner while all pre-conditioners considered are constant in time. Tools from random matrix theory and regularity assumptions for non-parametric regression are introduced with little discussion. Section 3.3 \"Misspecficiation \\approx Label Noise\" considers misspecification that is independent and gets interpreted as additional noise. It is not clear what this brings to the manuscript in terms of insights and introduces another layer of complexity.   \n\n C4) For parametric least squares regression the results focus on three cases for the ground truth covariance: well-aligned (where it equals the covariates population covariance), mis-aligned (where it equals the inverse covariance population covariance) and Isotropic. Whereas the theoretical results allow for a more general ground truth covariance to be considered. It would be natural to follow the source conditions from non-parametric regression and investigate natural gradient descent when the ground truth covariance is not fully well- or mis-aligned.   \n\n\nGeneral Comments: \n\n-Remark on page 4 states \"we demonstrate generalisation properties only possed by the population Fisher\", clarify which properties are /only/ held the population Fisher versus Sample Fisher. \n\n-In Proposition 4 possibly change the description \"interpolating preconditioners\" as all the preconditioned methods are interpolating the data, and thus, can be confusing.\n\n-Proposition 4 states for pre-conditioners (ii) and (iii) the bias is monotone for $\\alpha$ in some range depending upon the covariates population covariance. What is the range of $\\alpha$ and is the risk increasing or decreasing? What conclusions are we to draw from this part of the result? \n\n-In Figure 6 and Figure 23 how is \"geometric\" and \"additive\" interpolation defined ?    \n \n-More discussion around Proposition 6 would be helpful. For instance, in the statement of the result what is choice of P ? The analysis is described as difficult, although no details are provided into how this result was obtained. Within the proof why is the ratio of Eigenvalues \\overline{\\lambda}_{min}/\\widehat{\\lambda}_{min} is bounded, and how many iterations are required until NGD is below, say, standard gradient descent?  \n\n-In Section 5, the misalignment experiment in Figure 7 is conducted for MNIST but not CIFAR-10, with no discussion in the main body of the manuscript for why this is. Although, paragraph \"Misalignment\" in Appendix C.3 states the phenomena of NGD outperforming GD in the misaligned case is \"... difficult to observe in practical neural network training on real-world data\". The authors then go on to state that, in short, this is due to (see Appendix A) NGD moving parameters further from initialisation, and thus, no longer well described by a linear model i.e. NTK.  Is there a link between this discussion within the Appendix and the experiments?   \n\n-A Summary at the start of Appendix A to describe contents of A1-A4 would improve readability.  Similarly, for Appendix C and D.\n\n-In the proof of Theorem 2 (Appendix D.2) some details on how to get from (ii) to (iii). \n\n\n[1] - Rudi, A., Carratino, L., and Rosasco, L. \"Falkon: An optimal large scale kernel method\",  Advances in Neural Information Processing System 2017.\n\n\nPOST REBUTTAL EDIT: \n\nI thank the authors for providing detailed answers regarding my concerns. I have updated my score in light these comments. Below are some additional comments in response. \n\nResponse to comments regarding C1) and C2): \nWhile early stopping with pre-conditioned updates differentiates this work from (A. Rudi et. al 2019), the analysis still requires the knowledge of the population covariance. Indeed, while the authors have included a section (Appendix A.3) showing that the operator norm of the population and the inverse regularised empirical covariance can be controlled, it would be insightful to discuss to what extent this allows the analysis for the pre-conditioned gradient descent to be extended to an approximated population covariance.  \n\nResponse to comment regarding C3): \nI am inclined to agree with reviewer 3, in that the manuscript is difficult to read due to the larger number of fragmented results. In this regard, I feel the authors should focus on a single phenomenon that is supported by both the parametric and non-parametric aspects of the paper, for instance, how pre-conditioning helps against misalignment.      \n\nResponse to comment regarding different prior on ground truth (point 4. third bullet point): \nNote that some concurrent works have studied the case of different priors on the ground truth [2,3], which are likely relevant in this case. \n\nMinor Comment: The pre-conditioned updates for non-parametric regression (4.1) use notation $\\alpha$ where as Appendix D.8.1 uses notation $\\lambda$, with the discussion then switching back to using $\\alpha$ and $\\lambda$ being used in reference to the regularisation used within FALKON. The switching of notation is possibly confusing here. \n\n[2] - D. Richards, J. Mourtada, L. Rosasco \"Asymptotics of Ridge (less) Regression under General Source Condition\", arXiv:2006.06386 (2020) \n\n[3] - Wu, D. and Xu, J. \"On the Optimal Weighted $\\ell_2 $ Regularization in Overparameterized Linear Regression\" NeurIPS 2020\n\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "GkMGr9IRM-", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper studies generalization properties of preconditioned gradient descent on linear/kernel regression problems. The main preconditioner that is studied in addition to vanilla GD is the (population) Fisher matrix (natural gradient descent or NGD), its empirical counterpart, and its interpolation with GD.\nThe authors first consider the \"ridgeless\" regression setup in high-dimension, where the estimator corresponds to the limiting gradient flow iterate, and show that NGD leads to a smaller (and optimal) variance term, and can improve the bias term compared to GD particularly in the presence of strong misspecification. Among others, the authors also consider early-stopping in a non-parametric RKHS setup, showing that an appropriate interpolation between NGD and GD achieves optimal rates with a much smaller number of steps compared to GD, a difference which becomes larger for \"difficult\" problems (which require more weight on the Fisher preconditioner). The findings are further illustrated with simple experiments on neural networks.\n\nOverall, the paper provides a comprehensive study of the impact of preconditioning/second-order methods/natural gradient on generalization by giving a precise analysis in tractable regression settings, which illustrate conditions under which preconditioning is or is not useful for better generalization. This makes the paper a strong contribution, and I am in favor of acceptance.\n\ncomments/typos:\n- section 3: 'population risk' -> should this be excess risk given the presence of noise? add a reference or some more details on the bias-variance decomposition?\n- the last sentence in section 3.2 \"in the analogy...\" could be clarified\n- end of p.5 \"lower bias compare to\" -> \"compared to\"\n- Prop. 6: first part with theta_P holds for any P? please specify\n- theorem 7: specify conditions on eta?\n- some comments on computational difficulties of the full preconditioner would be welcome. Would a diagonal preconditioner, as often used in deep learning, provide any (partial) benefits as in the full-matrix case presented here?\n\n\n### Update after rebuttal\nThank you for the clarifications.\nA couple minor comments:\n- regarding theorem 7, my comment was that it would be useful to include the conditions on eta in the theorem statement in the main text (though I do not feel strongly about it)\n- regarding \"misalignment\" and the relationship between the random effects model and the source condition, I appreciate the improved explanation of this analogy, but I still find that the last paragraph in section 3.2 could do a better job at providing the right intuition (skimming through the Richards et al. reference pointed out by R4 gave me a better intuition).", "extra_scores": {}, "extra_text": {}}, {"review_id": "5dGfAHDbgV", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThe paper studies the effects of preconditioning on generalization properties in deep learning. By using a bias-variance decomposition of the expected risk, the paper determines optimal precondition matrix $P$ for bias and variance. Then the paper analyzes the generalization performance via the aspects: clean labels, well-specified model and aligned signal. Finally, it extends the analysis to the reproducing kernel Hilbert.\n\nPros:\n\nThe theoretical results provide guidelines of choosing precondition matrix for practical problems. In particular, by decomposing the risk into a sum of a bias and a variance, the paper addresses the following points:\n\n1. The asymptotic result on the variance (Theorem 1) implies that NGD achieves the minimal variance at stationary points, that suggests using NGD in the case where the variance term dominates.\n\n2. Theorem 2, on the other hand, provides the asymptotic result on the bias and the optimal precondition matrix for the bias to reach minimal value at stationary.\n\n3. Based on the results on the variance and the bias, Proposition 4 suggests an interpolating scheme between NGD and GD that aim at achieving better stationary risk than NGD or GD. The efficiency of this scheme is demonstrated in a least squares regression with the regular RKHS, where the interpolating scheme achieves the optimal convergence rate in fewer step than GD.\n\nCons:\n\n1. The paper contains a number of unclear / undefined terms such as well-specified and aligned signal, that make it difficult to read.\n\n2. The paper uses a lot of vague and unverified claims / statements which are usually the explanations after each theorem / proposition. For example, after theorem 1, it says that \"Theorem 1 implies that preconditioning with the inverse population Fisher results in the optimal stationary variance... In other words, when the labels are noisy so that the risk is dominated by the variance term... We emphasize that this advantage is only present when the population Fisher is used, but not its sample-based counterpart\". For me, it would be more clear if these statements could be explained in detail.\n\n3. The paper is not well-organized. For me, it is a collection of results that are unconnected. For example, after reading the analyses of bias and variance, I have no idea how they support the study of generalization or why section \"3.3 misspecification\" is placed along with bias and variance analyses, etc. I am not saying these results are irrelevant, however, there should be a better way of arranging / writing them so that they can support well the ideas of the paper.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper provides a study of the impact of preconditioning/second-order methods on generalization by giving a precise analysis in tractable regression settings.
It illustrates conditions under which preconditioning might be useful for better generalization. 
The readability issues raised by the reviewers have been taken into account, as well as some missing references, except

Wu, D. and Xu, J. "On the Optimal Weighted Regularization in Overparameterized Linear Regression" NeurIPS 2020, raised by reviewer (though it is a really recent reference).
Overall the contributions are significant enough to accept the paper for publication.
Decision: Accept (Poster)