W HEN D OES P RECONDITIONING H ELP OR H URT G ENERALIZATION ?
Anonymous authors
Paper under double-blind review
A BSTRACT
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 implicit bias of optimizers affects the comparison of generalization properties.
We provide an exact bias-variance decomposition of the generalization error of overparameterized ridgeless regression under a general class of preconditioner P , and consider the inverse population
Fisher information matrix (used in NGD) as a particular example.
We determine the optimal P for both the bias and variance, and ﬁnd 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 misspeciﬁed, or the signal is misaligned, NGD can achieve lower risk; conversely, GD generalizes better under clean labels, a well-speciﬁed model, or aligned signal.
Based on this analysis, we discuss approaches to manage the bias-variance tradeoff, and the beneﬁt of interpolating between ﬁrst- and second-order updates.
We then extend our analysis to regression in the reproducing kernel Hilbert space and demonstrate that preconditioned GD can decrease the population risk faster than GD.
Lastly, we empirically compare the generalization error of ﬁrst- and second-order optimizers in neural network, and observe robust trends matching our theoretical analysis.
I NTRODUCTION
We study the generalization property of an estimator ˆ θ obtained by minimizing the empirical risk (or the training error)
L ( f θ ) via a preconditioned gradient update with preconditioner P :
[Equation]
Setting P
= I recovers gradient descent (GD).
Choices of P which exploit second-order information include the inverse Fisher information matrix, which gives the natural gradient descent (NGD) (Amari, 1998); the inverse Hessian, which leads to Newton’s method; and diagonal matrices estimated from past gradients, corresponding to adaptive gradient methods (Duchi et al., 2011; Kingma & Ba, 2014).
These preconditioners often alleviate the effect of pathological curvature and speed up optimization , but their generalization properties has been under debate.
While several works reported that in neural network optimization, adaptive or second-order methods generalize worse compared to gradient descent (GD) (Wilson et al., 2017), other empirical studies suggested that second-order methods can 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., 2018a), i.e. preconditioning may lead to a different converged solution (with same training loss), as shown in Figure 1.
While many explanations have been proposed, our starting point is the well-known observation that GD implicitly regularizes the parameter (cid:96) 2 norm.
For instance in overparameterized least squares regression, GD and many ﬁrst-order methods ﬁnd the minimum (cid:96) 2 norm solution from zero initialization (without explicit regularization), but preconditioned updates often do not.
This being said, while the minimum norm solution may generalize well in the overparameterized regime (Bartlett et al., 2019), it is unclear whether
[Figure]
preconditioning leads to inferior solutions – even in the simple setting of overparameterized linear regression, quantitative understanding of how preconditioning affects generalization is large lacking.
Motivated by the observations above, in Section 3 we start with overparameterized least squares regression (unregularized) and analyze the stationary solution ( t → ∞ ) of update (1.1) under timeinvariant preconditioner.
Extending previous analysis in the proportional limit (Hastie et al., 2019), we consider a more general random effects model and derive the exact population risk in its biasvariance decomposition via random matrix theory.
We characterize the optimal 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 is the inverse population Fisher information matrix 1 .
We ﬁnd that the comparison of generalization performance is affected by the following factors:
Label Noise: Additive noise in the labels leads to the variance term in the risk.
We
show that NGD achieves the optimal variance among a general class of preconditioned updates.
Model Misspeciﬁcation: Under misspeciﬁcation, there does not exist f θ that perfectly learns the true function (target).
This contributes to the bias term, and we argue that its effect is similar to label noise, and thus NGD may also be beneﬁcial when the model is misspeciﬁed.
Data-Signal-Alignment: Alignment describes how the target signal distributes among input features 2 and affects the bias term.
GD achieves lower bias for isotropic signal, whereas NGD is preferred under “misalignment” — when the target function focuses on small feature directions.
Beyond the decomposition of stationary risk, our ﬁndings 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 in the RKHS and show that under early stopping, a preconditioned update interpolating between GD and NGD achieves minimax optimal convergence rate in much fewer steps, and thus reduces the population risk faster than GD.
• In Section 5 we empirically
test how well our ﬁndings in linear model carry over to neural networks: under a student-teacher setup, we compare the generalization of GD with preconditioned updates and illustrate the inﬂuence of all aforementioned factors.
The performance of neural networks under a variety of manipulations results in trends that align with our theoretical analysis.
B ACKGROUND AND R ELATED W ORKS
Natural Gradient Descent.
NGD is a second-order method originally proposed in Amari (1997).
Consider a data distribution p ( x ) on the space X , a function f θ : X → Z parameterized by θ , and a loss function L ( X , f θ ) = 1 n (cid:80) ni
=1 l ( y i , f θ ( x i )) , where l : Y × Z
→ R .
Also suppose a probability distribution p
( y | z ) =
p ( y | f θ ( x )) is deﬁned on the space of labels.
Then, the natural gradient is deﬁned as: ˜ ∇ θ L ( X , f θ ) =
F − 1 ∇ θ L ( X , f θ ) , where F = E [ ∇ θ log p ( x , y | θ ) ∇ θ log p ( x , y | θ ) (cid:62) ] is the Fisher information matrix , or simply the (population) Fisher.
Note that expectations in the
Fisher are under the joint distribution of the model p ( x , y | θ )
= p ( x ) p ( y | f θ ( x )) .
In the literature, the Fisher is sometimes deﬁned under the empirical data distribution { x i } n i =
(Amari et al., 2000).
We instead refer to this quantity as the sample Fisher , the properties of which inﬂuence optimization and have been studied in Karakida et al.
Kunstner et al.
; Thomas et al.
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 difﬁcult 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-efﬁcient (Martens & Grosse, 2015; Ollivier, 2015), and thus closing the gap between sample and population Fisher.
When the per-instance loss is the negative log-probability of an exponential family, the sample Fisher coincides with the generalized Gauss-Newton matrix (Martens, 2014).
In least squares regression, 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 ﬁndings to also translate to other second-order methods (not including adaptive gradient methods) in regression problems.
Analysis of Preconditioned Gradient Descent.
While Wilson et al.
(2017) outlined one example under ﬁxed 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 advantage of preconditioning (Levy & Duchi, 2019; Zhang et al., 2019a).
In addition, global convergence and generalization guarantees were derived for the sample Fisher-based update in neural networks in the kernel regime (Zhang et al., 2019b;
Cai et al., 2019).
Lastly, the generalization of different optimizers also connects to “sharpness” of the solution (Keskar et al., 2016;
Dinh et al., 2017), and it has been argued that second-order updates tend to ﬁnd 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 classiﬁcation problems.
We emphasize that instead of upper bounding the risk (e.g. Rademacher complexity), which may not decide the optimal P (for generalization), we compute the exact risk for least squares regression, which allows us to precisely compare different preconditioners.
A SYMPTOTIC R ISK OF R IDGELESS I NTERPOLANTS
In this section we consider the following setup: given n training samples { x i } ni =1 labeled by a teacher model (target function) f ∗ : R d → R with additive noise: y i =
f ∗ ( x i ) + ε i , we learn a linear student model f θ by minimizing the squared loss:
L ( X , f θ ) =
(cid:80) ni
=1 (cid:0)
y i − x (cid:62) i θ (cid:1) 2 .
We assume a random design: x i = Σ 1 / 2 X z i , where z i ∈ R
d is an i.i.d.
vector with zero-mean, unit-variance, and ﬁnite 12th moment, and ε is i.i.d.
noise independent to z with mean 0 and variance σ 2 .
Our goal is to compute the population risk R ( f ) =
E P X
[( f ∗ ( x ) − f ( x )) 2 ] in the proportional asymptotic limit:
• (A1)
Overparameterized Proportional Limit: n, d → ∞ , d/n → γ ∈ (1 , ∞ )
(A1) entails that the number of features (or parameters) is larger than the number of samples.
In this overparameterized setting, the population risk is equivalent to the generalization error, and there exist multiple empirical risk minimizers with potentially different generalization properties.
Denote X
=
[ x (cid:62) 1 , ..., x (cid:62) n ] (cid:62) ∈ R n × d the data matrix and y ∈ R n the corresponding label vector.
We optimize the parameters θ via a preconditioned gradient ﬂow with preconditioner P ( t ) ∈ R d × d ,
[Equation]
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 (cid:62) X /n up to variance scaling, whereas the population Fisher corresponds to the population covariance F = Σ X .
We thus limit our analysis to ﬁxed preconditioner of the form P ( t )
=: P .
Write parameters at time t under update (3.1) with ﬁxed P as θ P ( t ) .
For positive deﬁnite P , the stationary solution is given as:
[Equation]
y .
One may check that the discrete time gradient descent update (with appropriate step size) and other variants that do not alter the span of gradient (e.g. stochastic gradient or momentum) converge to the same solution as well.
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 (cid:96) 2 shrinkage (which GD implicitly provides) may not be the 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.
[Figure] [Figure]
Remark.
ˆ θ P is the minimum (cid:107) θ (cid:107) P − 1 norm interpolant: ˆ θ P = arg min θ (cid:107) θ (cid:107) P − 1 , s .
t .
Xθ =
y for positive deﬁnite P .
For GD this translates to the parameter (cid:96) 2 norm, whereas for NGD ( P = F − 1 = Σ − 1 X ), the implicit bias is the (cid:107) θ (cid:107) F norm.
Since E P X [ f ( x ) 2 ]= (cid:107) θ (cid:107) 2 Σ X , NGD ﬁnds an interpolating function with smallest norm under the data distribution.
We empirically observe this divide between small parameter norm and function norm in neural networks as well (Figure 1 and Appendix A.1).
We highlight the following choices of P and the corresponding stationary solution ˆ θ P as t → ∞ .
• Identity: P =
I d recovers GD that converges to the min (cid:96) 2 norm interpolant (also true for momentum GD and SGD), which we write as ˆ
θ I := X (cid:62) ( XX (cid:62) )
− 1 y and refer to as the GD solution .
• Population Fisher: P = F − 1 =
Σ − 1 X leads to the estimator ˆ θ F − 1 , which we refer to as the NGD solution .
• Sample Fisher: since the sample Fisher is rank-deﬁcient, we may add a damping term P =
( X (cid:62) X
+ λ I d ) − 1 or take the pseudo-inverse P =
( X (cid:62) X ) † .
In both cases, the gradient is still spanned by X , and thus the update ﬁnds the same min (cid:96) 2norm solution ˆ θ I (also true for full-matrix Adagrad (Agarwal et al., 2018)), although the trajectory differs (see Figure 3).
Remark.
The above choices reveal a gap between sample- and population-based P : while the sample Fisher accelerates optimization (Zhang et al., 2019b), the following sections demonstrate generalization properties only possessed by the population Fisher.
[Figure]
We compare the population risk of the GD solution ˆ θ I and NGD solution ˆ θ F − 1 in its bias-variance decomposition w.r.t.
label noise (Hastie et al., 2019) and discuss the two components separately,
[Equation]
Note that the bias term does not depend on the the label noise ε , whereas the variance term does not depend on the teacher f ∗ .
In addition, given that f ∗ can be independently decomposed into a linear component on the features x and a residual term: f ∗ ( x ) =
(cid:104)
x , θ ∗ (cid:105) + f ∗ c ( x ) , we can further decompose the bias into a well-speciﬁed component (cid:107) θ ∗ − E θ (cid:107) 2 Σ X , which captures the difﬁculty in learning the linear component of f ∗ , and a misspeciﬁed component (Hastie et al., 2019, Section 5), which corresponds to the error in ﬁtting f ∗ c (beyond the class of functions the student can represent).
3. T HE V ARIANCE T ERM : NGD IS O PTIMAL
We ﬁrst characterize the stationary variance which is independent to the teacher model f ∗ .
We restrict ourselves to preconditioners satisfying the following assumption on the spectral distribution:
• (A2) Converging Eigenvalues: P is positive deﬁnite and as n, d → ∞ , the spectral distribution of Σ XP :
= P 1 / 2 Σ X P 1 / 2 converges weakly to H XP supported on [ c, C ] for c, C > 0 .
The following theorem characterizes the asymptotic variance and the corresponding optimal P .
Theorem 1.
Given (A1-2), the asymptotic variance is given as
[Equation]
where m ( z ) > 0 is the Stieltjes transform of the limiting distribution of eigenvalues of 1 n XP X (cid:62) (for z beyond its support) deﬁned as the solution to m − 1 ( z )
= − z
+ γ (cid:82) τ (1+ τm ( z )) − 1 d H XP ( τ ) .
Furthermore, under (A1-2), V (ˆ θ P ) ≥ σ 2 ( γ − 1) − 1 , and the equality is achieved by P = F − 1 .
Formula (3.2) is a direct extension of Hastie et al.
(2019, Thorem 4), which can be obtained from Dobriban et al.
(2018, Thorem 2.1) or Ledoit & P´ech´e (2011, Thorem 1.2).
We note that the eigenvalue condition in (A2) may also be relaxed as in Xu & Hsu (2019).
Theorem 1 implies that preconditioning with the inverse population Fisher F results in the optimal stationary variance, which is supported by Figure 4(a).
In other words, when the labels are noisy so that the risk is dominated by the variance term, we expect NGD to generalize better upon convergence.
We emphasize that this advantage is only present when the population Fisher is used, but not its sample-based counterpart (which converges to ˆ θ I as commented above).
In Appendix A.3 we discuss the substitution error in replacing the population Fisher F with a sample-based estimate based on unlabeled data.
3.0 = d / n    v a r i a n c e
P = I d (GD) P = 1 X (NGD)
P = 1/2 X prediction (a) variance.
3.0 = d / n 0.
b i a s (b) well-speciﬁed bias (isotropic).
3.0 = d / n 0.
b i a s (c) well-speciﬁed bias (misaligned).
[Figure]
We now analyze the bias term when the teacher model is linear on the input features (hence wellspeciﬁed): f ∗ ( x ) =
x (cid:62) θ ∗ .
Extending the random effects hypothesis in Dobriban et al.
(2018), we consider a more general prior on θ ∗ : E [ θ ∗ θ ∗(cid:62) ]
= d − 1 Σ θ , and assume the following joint relations 3 :
• (A3) Joint Convergence: Σ X and P share the same eigenvector matrix U .
The empirical distributions of elements of ( e x , e θ , e xp ) jointly converge to random variables ( υ x , υ θ , υ xp ) supported on
[ c (cid:48) , C (cid:48) ] for c (cid:48) , C (cid:48) > 0 , where e x , e xp are eigenvalues of Σ X , Σ XP , and
e θ = diag ( U (cid:62) Σ θ U ) .
We remark that when P = I d , previous works (Hastie et al., 2019; Xu & Hsu, 2019) considered the special case of isotropic prior Σ θ = I d .
Our assumption thus allows for analysis of the bias term under much more general Σ θ , which gives rise to interesting phenomena that are not captured by simpliﬁed settings, such as non-monotonic bias and variance for γ > 1 (see Figure 15), and the epoch-wise double descent phenomenon (see Appendix A.2).
Under this general setup, we have the following characterization of the asymptotic bias and the corresponding optimal preconditioner:
Theorem 2.
Under (A1)(A3), the expected bias B (ˆ θ P ) :
=
E θ ∗ [ B (ˆ θ P )] is given as
[Equation]
where expectation is taken over υ and m ( z ) is the Stieltjes transform deﬁned in Theorem 1.
Furthermore, among all P satisfying (A3), the optimal bias is achieved by P = U diag [Equation]
[Figure] [Equation]
[Figure]
Note that the optimal P depends on the “orientation” of the teacher model Σ θ , which is usually not known in practice.
This result can thus be interpreted as a no-free-lunch characterization in choosing an optimal preconditioner for the bias term a priori .
As a consequence of the theorem, when the true parameters θ ∗ have roughly equal magnitude (isotropic), GD achieves lower bias (see Figure 4(b) where Σ θ = I d ).
On the other hand, NGD leads to lower bias when Σ X is “misaligned” with Σ θ , i.e. when θ ∗ focus on the least varying directions of input features (see Figure 4(c) where Σ θ = Σ − 1 X ), in which case learning is intuitively difﬁcult since the features are not useful 4 .
In the analogy of source condition (similar to E [ Σ − r/ 2 X θ ∗ ] < ∞ ), the coefﬁcient r can be interpreted as a measure of “misalignment”, which relates to the hardness of learning (see Section 4.3); the above discussion thus suggests that GD may be beneﬁcial in “easy” tasks, whereas NGD is preferable when the teacher is “difﬁcult”. In Appendix A.4 we elaborate this connection by analyzing the setting where Σ θ = Σ − r X , and provide a more precise comparison between the well-speciﬁed bias of GD and NGD in certain special cases.
3. M ISSPECIFICATION ≈ L ABEL N OISE
Finally, we consider the additional bias introduced by model misspeciﬁcation, under which there does not exist a linear student that perfectly recovers the teacher model f ∗ .
In this case, we may decompose the teacher into a linear component and its residual: f ∗ ( x ) =
x (cid:62) θ ∗ + f ∗ c ( x ) .
For simplicity, we ﬁrst consider f ∗ c to be a linear function on unobserved features (similar to Hastie et al.
): y i =
x (cid:62) i θ ∗ + x (cid:62) c,i θ c
+ ε i , where x c,i ∈
R d c is independent to x i with covariance Σ c X , and E [ θ c θ c (cid:62) ]
= d − 1 c Σ c θ .
In this setting, the misspeciﬁed bias behaves the same as the variance:
Proposition 3.
For the above unobserved features model, given (A1-3), the bias can be written as B (ˆ θ ) =
B θ (ˆ θ P ) + B c (ˆ θ P ) , where B θ is the well-speciﬁed bias in Thm. 2, and B c = d − 1 c tr( Σ c X Σ c θ )( V (ˆ θ P )+1) , where V (ˆ θ P ) is the variance in Thm. 1.
Misspeciﬁcation can thus be interpreted as additional label noise, for which NGD is advantageous by Theorem 1.
While Propositiondescribes one speciﬁc example of misspeciﬁcation, we expect such characterization to hold under broader settings.
In particular, Mei &
Montanari (2019, Remark 5) indicates that for many nonlinear f ∗ c , the misspeciﬁed bias is the same as variance due to label noise.
This result is only shown for isotropic data, but we empirically observe similar phenomenon under general covariances in Figure 6, in which f ∗ c is a quadratic function.
Observe that NGD leads to lower bias compared to GD as we further misspecify the teacher model.
[Figure]
c ( x ) =
α ( x (cid:62) x − tr( Σ X )) where α controls the extent of nonlinearity.
Predictions are generated by matching the noise level with 2nd moment of f ∗ c .
B IAS - VARIANCE T RADEOFF
Our characterization of the stationary risk suggests that preconditioners that achieve the optimal bias and variance are in general different.
This section discusses how the bias-variance tradeoff can be realized by interpolating between preconditioners or by early stopping.
In addition, we analyze regression in the RKHS and show that by balancing the bias and variance, a preconditioned update that interpolates between GD and NGD also decreases the population risk faster than GD.
4. I NTERPOLATING BETWEEN P RECONDITIONERS
[Figure]
r i a n c e
Depending on the orientation of the teacher model, we may expect a bias-variance tradeoff in choosing P .
Intuitively, given Pthat minimizes the bias and P 2 that minimizes the variance, it is possible that a preconditioner interpolating between P 1 and Pcan balance the bias and variance and thus generalize better.
The following proposition conﬁrms this intuition in a setup of general Σ X and isotropic Σ θ 5 , for which GD ( P = I d ) achieves optimal stationary bias and NGD ( P = F − 1 ) achieves optimal variance.
Proposition 4 (Informal) .
Let Σ X (cid:54) =
I d and Σ θ = I d .
Consider the following choices of interpolation scheme: (i) P α = α Σ − 1 X
+ (1 − α )
I d , (ii) P α
= ( α Σ X
+(1 − α )
I d ) − 1 , (iii) P α
= Σ − α X .
The stationary variance monotonically decreases with α ∈ [0 , 1] for all three choices.
For (i), the stationary bias monotonically increases with α ∈ [0 , 1] , whereas for (ii) and (iii), the bias monotonically increases with α in certain range that depends on Σ X .
In other words, as the signal-to-noise ratio (SNR) decreases (i.e. more label noise), one can increase α , which makes the update closer to NGD, to improve generalization, and vice versa 6 (small α entails GD-like update).
This intuition is supported by Figure 7 and 16(c): for certain SNR, interpolating between Σ − 1 X and Σ θ can lead to lower stationary risk than both GD and NGD.
Remark.
Two of the aforementioned interpolation schemes are analogous to common choices in practice: the additive interpolation (ii) corresponds to the damping term to stably invert the Fisher, whereas the geometric interpolation (iii) includes the “conservative” square-root scaling in adaptive gradient methods (Duchi et al., 2011; Kingma & Ba, 2014).
4. T HE R OLE OF E ARLY S TOPPING
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 brieﬂy comment on the effect of early stopping, starting from the monotonicity of the variance term.
Proposition 5.
For all P satisfying (A2), the variance V ( θ P ( t )) monotonically increases with time.
The proposition conﬁrms the intuition that early stopping reduces overﬁtting.
Variance reduction can beneﬁt GD in its comparison to NGD, which achieves lowest stationary variance: indeed, Figureand 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 always improve the well-speciﬁed bias.
While a complete analysis is difﬁcult partially due to the non-monotonicity of the bias term (see Appendix A.2), we speculate that previous observations on the stationary bias also translate to early stopping.
As a concrete example, we consider well-speciﬁed settings in which either GD or NGD achieves the optimal stationary bias, and demonstrate that such optimality is also preserved under early stopping:
Proposition 6.
Given (A1) and denote the optimal early stopping bias as B opt ( θ ) =
inf t ≥ 0 B ( θ ( t )) .
When Σ θ = Σ − 1 X , we have B opt ( θ P ) ≥ B opt ( θ F − 1 ) for all P satisfying (A3).
Whereas when Σ θ = I d , we have B opt ( θ F − 1 ) ≥ B opt ( θ I ) .
Figure 19 illustrates that the observed trend in the stationary bias (well-speciﬁed) is indeed preserved under optimal early stopping: GD or NGD achieves lower early stopping bias under isotropic or misaligned teacher model, respectively.
We leave a more precise characterization as future work.
4. F AST D ECAY OF P OPULATION R ISK
Our previous analysis suggests that certain preconditioners can achieve lower population risk, but does not address which method decreases the risk more efﬁciently.
Knowing that preconditioned updates may accelerate optimization, one natural question to ask is, is this speedup also present for generalization under ﬁxed dataset?
We answer this question in the afﬁrmative in a slightly different model: we study least squares regression in the RKHS, and show that a 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 detailed setup to Appendix D.8.
Let H be an RKHS included in L 2 ( P X ) equipped with a bounded kernel function k , and K x ∈ H be the Riesz representation of the kernel.
Deﬁne S as the canonical operator from H to L 2 ( P X ) , and write Σ
= S ∗ S and L = SS ∗ .
We aim to learn the teacher model f ∗ under the following standard regularity assumptions:
• (A4) Source Condition: ∃ r ∈ (0 , ∞ ) , M > 0 s.t.
f ∗ =
L r h ∗ for h ∗ ∈
L 2 ( P X ) and (cid:107) f ∗ (cid:107) ∞ ≤ M .
• (A5)
Capacity Condition: ∃ s > 1 s.t.
tr (cid:0) Σ 1 /s (cid:1) < ∞ and 2 r + s − 1 > 1 .
• (A6)
Regularity of RKHS: ∃ µ ∈ [ s − 1 , 1] , C µ > 0 s.t. sup x ∈ supp( P X ) (cid:13)(cid:13) Σ 1 / 2 − 1 /µ K x
(cid:13)(cid:13)
H ≤ C µ .
Note that in the source condition (A4), the coefﬁcient r controls the complexity of the teacher model and relates to the notions of model misalignment in Section 3.2: large r indicates a smoother teacher model which is “easier” to learn, and vice versa (Steinwart et al., 2009).
Given training points { ( x i , y i ) } ni =1 , we consider the following preconditioned update on the student model f t ∈ H :
[Equation]
where ˆΣ = 1 n (cid:80) ni
=1 K x i ⊗ K x i , ˆ S ∗ Y = 1 n (cid:80) ni
=
y i K x i .
In this setting, the population Fisher corresponds to covariance operator Σ , and thus (4.1) can be interpreted as additive interpolation between
GD and NGD: update with large α behaves like GD, and small α like NGD.
Related to our update is the FALKON algorithm (Rudi et al., 2017), which is a preconditioned gradient method for kernel ridge regression (under a different preconditioner).
The key distinction is that we consider optimizing the original objective (instead of a regularized version as in FALKON) under early stopping.
This is important since we aim to understand how preconditioning affects generalization , and therefore explicit regularization should not be taken into account (for more discussion see Appendix D.8.1).
The following theorem shows that with appropriately chosen α , the preconditioned update (4.1) leads to more efﬁcient decrease in the population risk, due to faster decay of the bias term.
Theorem 7 (Informal) .
Under (A4-6), the population risk of f t can be written as R ( f t ) =
(cid:107)
Sf t − f ∗ (cid:107) 2 L 2 ( P X ) ≤ B ( t )+ V ( t ) deﬁned in Appendix D.8.
Given r ≥ 1 / 2 or µ ≤ 2 r , preconditioned update (4.1) with α = n − 2 s 2 rs +1 achieves minimax optimal convergence rate R ( f t ) =
˜ O (cid:16) n − 2 rs 2 rs +1 (cid:17) in t =Θ(log n ) steps, whereas ordinary gradient descent requires t =Θ (cid:16) n 2 rs 2 rs +1 (cid:17) steps.
We comment that the optimal interpolation coefﬁcient α and stopping time t are chosen to balance the bias B ( t ) and variance V ( t ) .
Note that α depends on the teacher model in the following way: for n > 1 , α decreases as r becomes smaller, which corresponds to non-smooth and “difﬁcult” f ∗ , and vice versa.
This agrees with our previous observation that NGD is advantageous when the teacher model is difﬁcult to learn.
We defer empirical veriﬁcation of this result to Appendix C.
N EURAL N ETWORK E XPERIMENTS
Protocol.
We compare the generalization performance of GD and NGD in neural network settings and illustrate the inﬂuence of the following three factors: ( i ) label noise; ( ii ) misspeciﬁcation; ( iii ) signal misalignment.
We also verify the potential advantage of interpolating between GD and NGD.
We consider the MNIST and CIFAR-10 datasets.
To create a student-teacher setup, we split the training set into two halves, one of which ( pretrain split) along with the original labels is used to pretrain the teacher, and the other ( distill split) along with the teacher’s labels is used to distill (Hinton et al., 2015)
the student.
We take the teacher to 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 labels logits following Ba & Caruana (2014) before potentially adding label noise, and ﬁt the student by minimizing the L2 loss.
Student models are trained on a subset of the distill split with full-batch updates.
We implement NGD using Hessian-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 is 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.
5. E MPIRICAL F INDINGS
[Figure]
Label Noise.
We pretrain the teacher with the full pretrain split and use 1024 examples from the distill split to ﬁt the student.
For both the student and teacher, we use a two-layer ReLU net withhidden 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 4.
Misspeciﬁcation.
We use a ResNet-20 teacher and the same two-layer ReLU student from the label noise experiment.
We control the misspeciﬁcation level by varying amount of pretraining 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 (i.e. more misspeciﬁed).
Indeed, Figure 8(b) shows that NGD eventually achieves better generalization as the number of training steps for the teacher increases.
In Appendix A.5 we report a heuristic measure of model misspeciﬁcation that relates to the student’s NTK matrix (Jacot et al., 2018), and conﬁrm 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 models to be the same two-layer ReLU network.
We construct the teacher by perturbing the student’s initialization, the direction of which is given by F r , where F is the student’s Fisher (estimated from extra unlabeled data) and r ∈ [ − 1 , 0] .
Intuitively, as r approaches − 1 , the important parameters of the teacher (i.e. larger update directions) becomes misaligned with the student’s Hessian, and thus learning becomes more “difﬁcult”.
While this analogy is rather superﬁcial due to non-convexity, Figure 8(c) shows that as r becomes smaller (setup is more misaligned), NGD begins to generalize better than GD in terms of stationary risk.
Interpolating between Preconditioners.
We validate our observations in Section 3 and 4 on the difference between the sample Fisher and population Fisher, and the potential beneﬁt of interpolating between GD and NGD, in neural network experiments.
Figure 9(a) 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 under large noise.
This agrees with our remark on sample vs. population Fisher in Section 3 and Appendix A.1.
Figure 9(b)(c) supports the ﬁnding in Section 4.1 on the bias-variance tradeoff in neural network settings.
In particular, we interpret the left end to correspond to a bias-dominant regime (due to the same architecture of two-layer MLP for the student and teacher), and the right end to correspond to the variance-dominant regime (due to the added label noise).
Observe that at a certain SNR, a preconditioner that interpolates between GD and NGD achieves lower stationary risk.
[Figure] [Figure]
[Figure] [Figure]
[Figure]
D ISCUSSION AND C ONCLUSION
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 identiﬁed three factors that affect the comparison of generalization performance of different optimizers, the inﬂuence of which we also empirically observed in neural network 7 .
We then determined the optimal preconditioner for each factor.
While the optimal P is often not known in practice, we provided justiﬁcation for common algorithmic choices by discussing the bias-variance tradeoff.
Note that our current setup is limited to ﬁxed preconditioner or those that do not alter the span of the gradients, which does not cover many adaptive gradient methods; understanding 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) or algorithmic regularization (e.g. large step size and gradient noise).
T ABLE OF C ONTENTS A Discussion on Additional Results
A.
Implicit Bias of GD vs. NGD . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A.2 Non-monotonicity of the Bias Term w.r.t.
Time . .
A.3 Estimating the Population Fisher . . A.4 Bias Term Under Speciﬁc Source Condition . . . . . . . . . . . . . . . . . . . . . [Table]
B Additional Related Works
C Additional Figures
C.1 Overparameterized Linear Regression . . . . . . . . . . . . . . . . . . . . . . . .
C.2 RKHS Regression . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.3 Neural Networks . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D Proofs and Derivations
D.1 Missing Derivations in Section 3 . . . .
D.2 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.3 Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.4 Proof of Proposition 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.
Proof of Proposition 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.6 Proof of Proposition 5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.7 Proof of Proposition 6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.8 Proof of Theorem 7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.9 Proof of Proposition 8 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.10 Proof of Proposition 9 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.11 Proof of Corollary 10 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E Experiment Setup
E.1 Processing the Datasets . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.2 Setup and Implementation for Optimizers . . . . . . . . . . . . . . . . . . . . . .
E.3 Other Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A D ISCUSSION ON A DDITIONAL R ESULTS A. I MPLICIT B IAS OF GD VS . NGD
It is known that gradient descent is the steepest descent with respect to the (cid:96) 2 norm, i.e. the update direction is constructed to decrease the loss under small changes in the parameters measured by the (cid:96) 2 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 on the data distribution.
Intuitively, the above discussion suggests GD tend to ﬁnd solution that is close to the initialization in the Euclidean distance between parameters, whereas NGD prefers solution close to the initialization in terms of the function outputs on P X .
This observation turns out to be exact in the case of ridgeless interpolant under the squared loss, as remarked in Section 3.
Moreover, Figure 1 and 10 conﬁrms the same trend in neural network optimization.
In particular, we observe that
• GD results in small changes in the parameters (left), and NGD results in small changes in the learned function (right).
• preconditioning with the pseudo-inverse of the sample Fisher, i.e., P =
( J (cid:62) J ) † , leads to implicit bias similar to that of GD, but not NGD with the population Fisher.
• interpolating between GD and NGD ( P = F − 1 / 2 ) 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”, so that memorizing the noisy labels has small impact on the “global” shape of the learned function.
[Figure]
We note that the above observation also implies that wide neural networks trained with NGD (population Fisher) is less likely to 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 accurate.
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 with NGD 8 .
Implicit Bias of Interpolating Preconditioners.
We also expect that as we interpolate from GD to
NGD, the distance traveled by 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 linear model as well as neural network: we use the same two-layer MLP setup on MNIST as in Section 5.
Note that updates that are closer to GD result in smaller change in the parameters, whereas ones close to NGD lead to smaller change in the function outputs.
[Figure]
A. N ON - MONOTONICITY OF THE
[Figure]
IAS
[Figure]
[Figure]
Note that non-monotonicity
Many previous works on the high-dimensional characterization of linear regression assumed a random effects model with an isotropic prior on the true parameters (Dobriban et al., 2018;
Hastie et al., 2019; Xu & Hsu, 2019), which may not be realistic.
As an example of the limitation of this assumption, we note that when Σ θ = I
d , it can be shown that the expected bias B (ˆ θ ( t )) monotonically decreases through time (see proof of Proposition 6 for details).
In contrast, when the target parameters do not follow an isotropic prior, the bias of GD can exhibit non-monotonicity, which gives rise to the surprising “epoch-wise double descent” phenomenon also observed in deep learning (Nakkiran et al., 2019; Ishida et al., 2020).
We empirically demonstrate this non-monotonicity when the model is close to the interpolation threshold in Figure 12.
We set eigenvalues of Σ X to be two equally-weighted point masses with κ X = 32 , Σ θ = Σ − 1 X and γ = 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 conﬁrm in the proof of Proposition 6).
We remark that this mechanism of epoch-wise double descent may not be related to the empirical ﬁndings in deep neural networks (the robustness of which is also unknown), in which it is typically speculated that the variance term exhibits non-monotonicity.
A. E STIMATING THE P OPULATION F ISHER
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 3 and Section 4 are affected when the population covariance is approximated from N i.i.d.
(unlabeled) samples X u ∈ R N × d .
For the ridgeless interpolant we have the following result on the substitution error in replacing the true population covariance with the sample estimate.
Proposition 8.
Given (A1)(A3) and N/d → ψ > 1 as d → ∞ , let ˆ Σ X = X (cid:62) u X u /N , we have
(a) (cid:107) Σ X
− ˆ Σ X (cid:107)
= O ( ψ − 1 / 2 ) almost surely.
(b) Denote the stationary bias and variance of NGD (with the exact population Fisher) as B ∗ and V ∗ , respectively, and the bias and variance of the preconditioned update using the approximate Fisher ˆ F = ˆ Σ X as ˆ B and ˆ V , respectively.
Let 0 < (cid:15) < 1 be the desired accuracy.
Then ψ = Θ( (cid:15) − 2 ) sufﬁces to achieve | B ∗ − ˆ B | < (cid:15) and | V ∗ − ˆ V | < (cid:15)
Proposition 8 entails that when the preconditioner is a sample estimate ˆ F , we can approach the stationary bias and variance of the population Fisher at a rate of ψ − 1 / 2 as we increase the number of unlabeled data N linearly in the dimensionality d .
In other words, any accuracy (cid:15) such that 1 /(cid:15) is bounded can be achieved with ﬁnite ψ (to push (cid:15) → 0 , additional logarithmic factor is required, e.g. N =
O ( d log d ) , which is beyond the proportional limit).
On the other hand, for our result on RKHS regression in Section 4.3, Murata & Suzuki (2020,
Lemma A.5) directly implies that setting N = Θ( n s log n ) is sufﬁcient to approximate the population covariance operator (i.e., so that (cid:107) Σ 1 / 2 Σ − 1 / 2 N,λ (cid:107) =
O (1) ).
Finally, we note that our analysis above does not impose any structural assumptions on the estimated matrix.
When the Fisher exhibits certain structures (e.g. Kronecker factorization (Martens & Grosse, 2015)), then the estimation can be more sample-efﬁcient.
We leave characterization for this scenario as future work.
A. B IAS T ERM U NDER S PECIFIC S OURCE C ONDITION
Motivated by the connection between the notion of “alignment” and the source condition (Cucker &
Smale, 2002) in Section 3.2, we now consider a speciﬁc case of θ ∗ : Σ θ = Σ − r X , where r controls the extent of misalignment, and Theorem 2 implies that the optimal preconditioner for the bias term (well-speciﬁed) is P = Σ − r X .
Note that larger r corresponds to more misaligned and thus “difﬁcult” problem, and vice versa.
In this setup we have the following comparison between GD and NGD.
Proposition 9.
Under the setting of Theorem 2 and Σ θ = Σ − r X , for all r ≥ 1 we have B (ˆ θ F − 1 ) ≤
B (ˆ θ I ) , whereas for r ≤ 0 , we have B (ˆ θ F − 1 ) ≥ B (ˆ θ I ) ; the equality is achieved when the features are isotropic ( Σ X
= c I d ).
[Figure]
The proposition conﬁrms the intuition that when parameters of the teacher model θ ∗ are more “aligned” with x than the isotropic case ( r ≤ 0 ), then GD achieves lower bias than NGD; on the other hand, when Σ θ is more “misaligned” than Σ − 1 X , then NGD is advantageous for the bias term.
We therefore expect a transition from the GD-dominated to NGD-dominated regime for some r ∈ (0 , 1) .
The exact value of r for such transition depends on the spectral distribution of Σ X (and one would need to speciﬁcally evaluate (D.26)).
To give a concrete example, in the case where Σ X has a simple block structure, we can explicitly determine the the transition point r ∗ ∈ (0 , 1) , as shown by the following corollary.
Corollary 10.
Assume Σ
θ = Σ − r X , and eigenvalues of Σ X come from two equally-weighted point masses with κ (cid:44) λ max ( Σ X )
/λ min ( Σ X ) .
WLOG we take tr( Σ X ) /d
= 1 .
Then given r ∗ =
ln c κ,γ / ln κ (see Appendix D.11 for deﬁnition), we have B (ˆ θ I ) (cid:82) B (ˆ θ F − 1 ) if and only if r (cid:82) r ∗ .
Remark.
When γ = 2 , the transition happens at r ∗ = 1 / 2 which is independent of the condition number κ , as indicated by the dashed line in Figure 13.
However for other γ > 1 , r ∗ generally relates to both γ and κ .
Our characterization above is supported by Figure 13, in which we plot the bias term under varying extent of misalignment (controlled by r ) in the setting of Corollary 10.
Observe that as we construct the teacher model to be more “misaligned” (and thus difﬁcult to learn) by increasing r , NGD (blue) achieves lower bias compared to GD (red), and vice versa.
A. I NTERPRETATION OF (cid:112) y (cid:62) K − 1 y /n
[Figure]
[Figure]
As a heuristic measure of model misspeciﬁcation, in Figurewe report (cid:112) y (cid:62) K − 1 y /n studied in Arora et al.
(2019b), where y is the label vector and K is the student’s NTK matrix (Jacot et al., 2018).
This quantity relates to generalization of wide neural networks in the kernel regime, an can be interpreted as a proxy for measuring how much signal and noise are distributed along the eigendirections of the NTK (see Li et al.
Dong et al.
; Su & Yang (2019); Cao et al.
for detailed discussion).
Roughly speaking, large (cid:112) y (cid:62) K − 1 y /n implies that the problem is difﬁcult to learn by GD, vice versa.
Here we give a heuristic argument on how this quantity relates to label noise and misspeciﬁcation.
For the ridgeless regression model considered in Section 3, write y i =
f ∗ ( x i ) + f c ( x i ) + ε i , where f ∗ ( x ) =
x (cid:62) θ ∗ , f c is the misspeciﬁed component, and ε i is the label noise, then we have the following heuristic calculation:
[Equation]
where we heuristically replaced the misspeciﬁed component with i.i.d.
noise of the same variance σ 2 c (as argued in Section 3).
The ﬁrst term of (A.1) resembles an RKHS norm of the target θ ∗ , whereas the second term is small when the feature matrix is well-conditioned or when the level of label noise σ and misspeciﬁcation σ 2 c is small (these are conditions under which GD achieves good generalization by Theorem 1 and Proposition 3).
We expect similar trends for neural networks close to the kernel regime.
This provides a non-rigorous explanation of the trend we observed in Figureas we increase the level of label noise and model misspeciﬁcation.
B A DDITIONAL R ELATED W ORKS
Implicit Regularization in Optimization.
In overparameterized linear models, GD ﬁnds the minimum (cid:96) 2 norm solution under many loss functions.
For the more general mirror descent, the implicit bias is determined by the Bregman divergence of the update (Azizan & Hassibi, 2018;
Azizan et al., 2019; Gunasekar et al., 2018b; Suggala et al., 2018).
Under the exponential or logistic loss, recent works showed that GD ﬁnds 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 of the optimizer also relates to the model architecture; examples include matrix factorization (Gunasekar et al., 2017;
Saxe et al., 2013; Gidel et al., 2019; Arora et al., 2019a) and various types of neural network (Li et al., 2017; Gunasekar et al., 2018b; Williams et al., 2019;
Woodworth et al., 2020).
For wide neural networks in the kernel regime (Jacot et al., 2018), the implicit bias of GD relates to properties of the limiting neural tangent kernel (NTK)
(Xie et al., 2016; Arora et 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).
Asymptotics of Interpolating Estimators.
In Section 3 we analyzed overparameterized estimators that interpolate the training data.
Recent works have shown that interpolation may not lead to overﬁtting (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 overparameterized models has been characterized in various settings, such as linear regression (Karoui, 2013; Dobriban et al., 2018; Hastie et al., 2019), random features regression (Mei & Montanari, 2019; Gerace et al., 2020; Hu & Lu, 2020), max-margin classiﬁcation (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 results in random matrix theory developed in Rubio & Mestre (2011); Ledoit & P´ech´e (2011).
Similar tools can also be used to study the gradient descent dynamics of linear regression (Liao & Couillet, 2018; Ali et al., 2019).
C A DDITIONAL F IGURES C. O VERPARAMETERIZED L INEAR R EGRESSION
Non-monotonicity of the Risk.
Under our generalized (anisotropic) assumption on the covariance of the features and the target, both the bias and the variance term can exhibit non-monotonicity w.r
[Figure] t. the overparameterization level γ > 1 : in Figure 15 we observe two peaks in the bias term and three peaks in the variance term.
In contrast, it is known that when Σ X = I
d , both the bias and variance are monotonic for when γ > 1 .
[Figure]
Additional Figures for Section 3 and 4.
We include additional ﬁgures on (a) well-speciﬁed bias when Σ θ = I d (GD is optimal); (b) misspeciﬁed bias under unobserved features (predicted by
Proposition 3); (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.
[Figure]
[Figure]
[Figure] [Figure]
Early Stopping Risk.
Figure 18 compares the stationary risk with the optimal early stopping risk under varying misalignment level.
To increase the extent of misalignment, we set Σ θ = Σ − α X and vary α from 0 to 1: larger α entails more “misaligned” teacher, and vice versa.
Note that as the problem becomes more misaligned, NGD achieves lower stationary and early stopping risk.
Figure 19 reports the optimal early stopping risk under misspeciﬁcation (same trend can be obtained when the x-axis is label noise).
In contrast to the stationary risk (Figure 6), GD can be advantageous under early stopping even with large extent of misspeciﬁcation (for isotropic teacher).
This aligns with our ﬁnding in Section 4.2 that early stopping reduces the variance and the misspeciﬁed bias
[Figure] [Figure]
C. RKHS R EGRESSION
We simulate the optimization in the coordinates of RKHS via a ﬁnite-dimensional approximation (using extra unlabeled data).
In particular, we consider the teacher model in the form of f ∗ ( x ) =
(cid:80)
Ni =1 h i µ ri φ
i
( x ) for square summable { h i } Ni =1 , in which r controls the “difﬁculty” of the learning problem.
We ﬁnd { µ i } Ni =1 and { φ i } Ni =1 by solving the eigenfunction problem for some kernel k .
The student model takes the form of f ( x ) = (cid:80) Ni
=1 a i √ µ i φ i ( x )
and we optimize the coefﬁcients { a i } Ni =1 via the preconditioned update (4.1).
We set n = 1000 , d = 5 , N = 2500 and consider the inverse multiquadratic (IMQ) kernel: k ( x , y ) = 1 √ 1+ (cid:107) x
[Table]
Recall that Theorem 7 suggests that for small r , i.e. “difﬁcult” problem, the damping coefﬁcient λ would need to be small (which makes the update NGD-like), and vice versa.
This result is (qualitatively) supported by Figure 20, from which we can see that small λ is beneﬁcial when r is small, and vice versa.
We remark that this observed trend is rather fragile and sensitive to various hyperparameters, and we leave a comprehensive characterization of this observation as future work. steps p o p u l a t
i o n r i s k
=1 e 0 =1 e=1 e(a) r = 3 / 4 . steps p o p u l a t
i
o n r i s k
=1 e 0 =1 e=1 e(b) r = 1 / 4 .
C. N EURAL N ETWORKS
Label Noise.
In Figure 21, (a) we observe the same phenomenon on CIFAR-10 that NGD generalizes better as more label noise is added to the training data, 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
[Figure]
Misalignment.
We illustrate the ﬁnding in Proposition 6 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 in Section 5.
Figure 22 shows that as r approaches - (more “misaligned”), NGD achieves lower early stopping risk (b) [Figure] whereas GD dominates the early stopping risk in less misaligned setting (a).
We note that this phenomenon is difﬁcult to observe in practical neural network training on 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)
[Figure]
D P ROOFS AND D ERIVATIONS D. M ISSING D ERIVATIONS IN S ECTION
Gradient Flow of Preconditioned Updates.
Given positive deﬁnite P and γ > 1 , it is clear that the gradient ﬂow solution at time t can be written as
[Equation]
Taking t → ∞ yields the stationary solution ˆ θ P = P X (cid:62) ( XP X (cid:62) ) − 1 y .
We remark that the damped inverse of the sample Fisher P =
( XX (cid:62) + λ I d ) − 1 leads to the same minimum-norm solution as GD ˆ θ I
= X (cid:62) ( XX (cid:62) ) − 1 y since P X (cid:62) and X share the same eigenvectors.
On the other hand, when P is the pseudo-inverse of the sample Fisher ( XX (cid:62) ) † which is not full-rank, the trajectory can be obtained via the variation of constants formula:
[Equation]
for which taking the large t limit also yields the minimum-norm solution X (cid:62) ( XX (cid:62) )
− 1 y .
Minimum (cid:107) θ (cid:107) P − 1 Norm Interpolant.
For positive deﬁnite P and the corresponding stationary solution ˆ θ P = P X (cid:62) ( XP X (cid:62) ) − 1 y , note that given any other interpolant ˆ θ (cid:48) , we have (ˆ θ P − ˆ θ (cid:48) )
P − 1 ˆ θ P =
because both ˆ θ P and ˆ θ (cid:48) achieves zero empirical risk.
Therefore, (cid:107) ˆ θ (cid:48) (cid:107) 2 P − 1 − (cid:107) ˆ θ P (cid:107) 2 P − 1 = (cid:107) ˆ θ (cid:48) − ˆ θ P (cid:107) 2 P − 1 ≥ 0 .
This conﬁrms that ˆ θ P is the unique minimum (cid:107) θ (cid:107) P − 1 norm solution.
D. P ROOF OF T HEOREM
Proof.
By the deﬁnition of the variance term and the stationary ˆ θ ,
[Equation]
Write ¯ X = XP 1 / 2 .
Similarly, we deﬁne Σ XP = P 1 / 2 Σ X P 1 / 2 .
The equation above thus simpliﬁes to
[Equation]
The analytic expression of the variance term follows from a direct application of Hastie et al.
Thorem 4), in which conditions on the population covariance are satisﬁed by (A2).
Taking the derivative of m ( − λ ) yields
[Equation]
Plugging the quantity into the expression of the variance (omitting the scaling σ 2 and constant shift),
[Equation]
From the monotonicity of x 1+ x on x > 0 or the Jensen’s inequality we know that
[Equation]
Taking λ → 0 and omitting the scalar σ 2 , the RHS evaluates to 1 − 1 /γ .
We thus arrive at the lower bound V ≥ ( γ − 1) − 1 .
Note that the equality is only achieved when H XP is a point mass, i.e. P = Σ − 1 X .
In other words, the minimum variance is achieved by NGD.
As a veriﬁcation, the variance of the NGD solution ˆ θ F − 1 agrees with the calculation in Hastie et al.
(2019, A.3).
D. P ROOF OF T HEOREM
Proof.
By the deﬁnition of the bias term (note that Σ X , Σ θ , P are all positive semi-deﬁnite),
[Equation]
where we utilized (A3) and deﬁned ¯ X
= XP 1 / 2 , Σ XP
[Equation]
P − 1 / 2 Σ θ P − 1 / 2 in (i), applied the equality ( AA (cid:62) ) † A = lim λ →
[Equation]
( A (cid:62) A
+ λ I ) − 1 A in (ii), and deﬁned ˆ X = XP 1 / 2 Σ − 1 / 2 θ in (iii).
To proceed [Equation] we ﬁrst assume that Σ θ / P is invertible (i
[Equation]
e. λ min ( Σ θ / P ) is bounded away from 0) and observe the following relation via a leave-one-out argument similar to that in Xu & Hsu (2019),
[Equation]
) where (i) is due to the Woodbury identity and we deﬁned (cid:16) 1 n ˆ X (cid:62)
ˆ X + λ Σ − 1 θ / P (cid:17)
¬ i = 1 n ˆ X (cid:62) ˆ X − 1 n ˆ x i ˆ x (cid:62) i + λ Σ − 1 θ / P which is independent to ˆ x i (see Xu & Hsu (2019, Eq. 58) for details), and in (ii) we used (A2), the convergence to trace (Ledoit & P ´ ech ´ e, 2011, Lemma 2.1) and its stability under low-rank perturbation (e.g., see Ledoit & P ´ ech ´ e (2011, Eq. 18)) which we elaborate below.
In particular, denote ˆ Σ = 1 n ˆ X (cid:62)
ˆ X + λ Σ − 1 θ / P , for the denominator we have [Equation] [Table] [Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
where (i) is due to the deﬁnition of ˆ Σ ¬ i and (A1)(A3).
The result on the numerator can be obtained via a similar calculation, the details of which we omit.
Note that the denominator can be evaluated by previous results (e.g. Dobriban et al.
(2018, Theorem 2.1)) as follows,
[Equation]
On the other hand, following the same derivation as Dobriban et al.
; Hastie et al.
(2019), (D.1) can be decomposed as
[Equation]
We employ Rubio & Mestre (2011, Theorem 1) to characterize (D.4).
In particular, For any deterministic sequence of matrices Θ n ∈ R d × d with ﬁnite trace norm, as n, d → ∞ we have
[Equation]
in which c n ( z ) → − zm ( z ) for z ∈ C \ R + and m ( z ) is deﬁned in Theorem 1 due to the dominated convergence theorem.
By (A3) we are allowed to take Θ n = 1 d Σ θ / P .
Thus we have
[Equation]
in which (i) is due to (A3), the fact that the LHS is almost surely bounded for λ > − c l , where c l is the lowest non-zero eigenvalue of 1 n ¯ X (cid:62) ¯ X , and the application of the dominated convergence theorem.
Differentiating (D.5) (note that the derivative is also bounded on λ > − c l ) yields
[Equation]
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 Σ θ / P is not invertible, observe that if we increment all eigenvalues by some (cid:15) > 0 to ensure invertibility Σ (cid:15) θ / P = Σ θ / P + (cid:15) I
d , (D.6) is bounded and also decreasing w.r.t.
(cid:15) .
Thus by the dominated convergence theorem we take (cid:15) → 0 and obtain the desired result.
We remark that similar (but less general) characterization can also be derived based on Ledoit & P ´ ech ´ e (2011, Theorem 1.2) when the eigenvalues of Σ XP and Σ θ / P exhibit certain relations.
To show that P = U diag ( e θ ) U (cid:62) achieves the lowest bias, ﬁrst note that under the deﬁnition of random variables in (A3), our claimed optimal preconditioner is equivalent to υ xp a.s. =
υ x υ θ .
We therefore deﬁne an interpolation υ α = αυ x υ θ
+
(1 − α )
¯ υ for some ¯ υ and write the corresponding
Stieltjes transform as m α ( − λ ) and the bias term as B α .
We aim to show that argmin α ∈ [0 , 1] B α = 1 .
For notational convenience deﬁne g α (cid:44) m α (0) υ x υ
θ
and h α (cid:44) m α (0)
υ α .
One can check that
[Equation]
We now verify that the derivative of B α w.r.t.
α is non-positive for α ∈ [0 , 1] .
A standard simpliﬁcation of the derivative yields
[Equation]
E (cid:20) h
α (1 + h α )(cid:21)(cid:19)−(cid:18) E
(cid:20) g
α − h
α (1 + h α )(cid:21)(cid:19) 2 E (cid:20)
h 2 α (1 + h α )(cid:21)
+ 4 E (cid:20) h
[Equation]
( g
[Equation]
− h α ) (1 + h α )(cid:21) E (cid:20) g
[Table]
h α (1 + h α )(cid:21) E (cid:20) h
α (1 + h α )(cid:21) ( i ) ≤ −(cid:115)
E (cid:20) ( g α − h α )(1 + h α )(cid:21) E (cid:20)
h 2 α (1 + h α )(cid:21)(cid:18) E
(cid:20)
[Table]
h
[Table]
(1 + h
[Table]
)(cid:21)(cid:19)
[Table]
(cid:18)
[Table]
(cid:20) h
[Table]
(1 + h
[Table]
[Table]
(cid:21)(cid:19)+ 4 E (cid:20) h
[Table]
( g
[Table]
h
[Table]
) (1 + h α )(cid:21) E (cid:20) g
[Table]
h
[Table]
(1 + h α )
[Table]
(cid:21) E (cid:20) h
[Table] (1 + h α ) [Table] (cid:21) ( ii ) ≤ 0 , where (i) is due to AM-GM and (ii) due to Cauchy-Schwarz on the ﬁrst term.
Note that the two equalities hold when g α = h α , from which one can easily deduce that the optimum is achieved when υ xp a.s. =
υ x υ θ , and thus we know that the proposed P is the optimal preconditioner that is codiagonazable with Σ X .
D. P ROOF OF P ROPOSITION
Proof.
Via calculation similar to Hastie et al.
(2019, Section 5), the bias can be decomposed as
[Equation]
where we used the independence of x , ˆ x and θ ∗ , θ c in (i), and (A1-3) as well as the deﬁnition of the well-speciﬁed bias B θ (ˆ θ P ) and variance V (ˆ θ P ) in (ii).
D. P ROOF OF P ROPOSITION
Proof.
We ﬁrst outline a more general setup where P α = f
( Σ X ; α ) for continuous and differentiable function of α and f applied to eigenvalues of Σ x .
For any interval I ⊆
[0 , 1] , we claim
(a) Suppose all four functions 1 xf ( x ; α ) , f ( x ; α ) , ∂f ( x ; α ) ∂α /f
( x ; α ) and x ∂f ( x ; α ) ∂α are decreasing functions of x on the support of v x for all α ∈
I .
In addition, ∂f ( x ; α ) ∂α ≥ 0 on the support of v x for all α ∈ I .
Then the stationary bias is an increasing function of α on I .
(b) For all α ∈ I , suppose xf ( x ; α ) is a monotonic function of x on the support of v x and ∂f ( x ; α ) ∂α /f
( x ; α ) is a decreasing function of x on the support of v x .
Then the stationary variance is a decreasing function of α on I .
Let us verify the three choices of P α in Proposition 4 one by one.
• When P α = (1 − α ) I d + α ( Σ X ) − 1 , the corresponding f ( x ; α ) is (1 − α ) + αx .
It is clear that it satisﬁes all conditions in (a) and (b) for all α ∈ [0 , 1] .
Hence, the stationary variance is a decreasing function and the stationary bias is an increasing function of α ∈ [0 , 1] .
• When P α =
( Σ X ) − α , the corresponding f ( x ; α ) is x − α .
It is clear that it satisﬁes all conditions in (a) and (b) for all α ∈ [0 , 1] except for the condition that x ∂f ( x ; α ) ∂α =
− x 1 − α
ln
x is a decreasing function of x .
Note that x ∂f ( x ; α ) ∂α =
− x 1 − α
ln
x is a decreasing function of x on the support of v x only for α ≥ ln( κ ) − 1 ln( κ ) where κ = sup v x / inf v x .
Hence, the stationary variance is a decreasing function of α ∈ [0 , 1] and the stationary bias is an increasing function of α ∈ [max(0 , ln( κ ) − 1 ln( κ ) ) , 1] .
• When P α
=
( α Σ X
+(1 − α )
I d ) − 1 , the corresponding f ( x ; α ) is 1 / ( αx +(1 − α )) .
It is clear that it satisﬁes all conditions in (a) and (b) for all α ∈ [0 , 1] except for the condition that x ∂f ( x ; α ) ∂α
= x (1 − x ) ( αx +(1 − α )) 2 is a decreasing function of x .
Note that x ∂f ( x ; α ) ∂α [Equation] x ) [Equation]
1 −
[Equation]
is a decreasing function of x on the support of v x only for α ≥ κ − 2 κ − 1 .
Hence, the stationary variance is a decreasing function of α ∈ [0 , 1] and the stationary bias is an increasing function of α ∈ [max(0 , κ − 2 κ − 1 ) , 1]
To show (a) and (b), note that under the conditions on Σ x and Σ θ assumed in Proposition 4, the stationary bias B (ˆ θ P α ) and the stationary variance V (ˆ θ P α ) can be simpliﬁed to
[Equation]
where m α ( z ) and m (cid:48) α ( z ) satisfy
[Equation]
For notation convenience, let f α :=
v x f ( v x ; α ) .
From (D.8), we have the following equivalent expressions.
[Equation]
We ﬁrst show that (b) holds.
Note that from (D.10), we have
[Equation]
[Equation]
To calculate ∂m α ( z ) ∂α (cid:12)(cid:12)(cid:12) z =0 , we take derivatives with respect to α on both sides of (D.7),
[Equation]
(D [Equation]
[Equation]
Therefore, plugging (D.12) into (D.11) yields
[Equation]
Thus showing V (ˆ θ P α ) is a decreasing function of α is equivalent to showing that
[Equation]
Let µ x be the probability measure of v x .
We deﬁne a new measure ˜ µ x = [Equation] )) 2 , and let ˜ v x follow the new measure.
Since ∂f ( x ; α ) ∂α /f ( x ; α ) is a decreasing function of x and xf ( x ; α ) is a monotonic function of x ,
[Equation]
Changing ˜ v x back to v x , we arrive at (D.13) and thus (b).
For the bias term B (ˆ θ P α ) , note that from (D.7) and (D.9), we have
[Equation]
[Equation]
Similarly
[Equation]
we combine (D.12) and (D.14) and simplify the expression.
To verify B (ˆ θ P α ) is an increasing function of α
[Equation]
we need to show that
[Equation]
Let h α (cid:44) f
α m α (0) =
v x f ( v x ; α ) m α (0) and g α (cid:44)
∂f α ∂α =
v x ∂f ( v x ; α ) ∂α .
Then (D.15) can be further simpliﬁed to the following equation [Equation]
(1 + h α )(cid:124) (cid:123)(cid:122) (cid:125)
part+ E v x (1 + h α )
[Equation]
E g
[Equation]
h
[Equation]
(1 + h α )
[Equation]
E h
α
(1 + h α )
[Equation]
− E v x g α (1 + h α )
[Equation]
E h
α
(1 + h α ) 3 E h α (1 + h α )(cid:124) (cid:123)(cid:122) (cid:125)
part+ 2 E v x h α
(1 + h α )
[Equation]
E g
[Equation]
h
[Equation]
(1 + h
[Equation]
[Equation]
E h
[Equation] (1 + h [Equation] [Equation]
E v x g α (1 + h α ) 3 E
{+h 2 α (1 + h α ) 
+}
E h α
(1 + h α )(cid:124) (cid:123)(cid:122) (cid:125)
part+ E
[Equation]
[Equation]
[Equation]
E
[Equation]
[Equation]
E v
[Equation]
(1 + h α ) 3 E h 2 α
(1 + h α ) 3 E h 2 α
(1 + h α )(cid:124) (cid:123)(cid:122) (cid:125) part.
(D.16)
Note that under condition of (a), we know that both h α and v x /h
α are increasing functions of v x ; and both g α /h α
and g α are decreasing functions of v x .
Hence, with calculation similar to (D.13), we know part 1,2,3,4 in (D.16) are all non-negative, and therefore (D.16) holds.
Remark.
The above characterization provides sufﬁcient but not necessary conditions for the monotonicity 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. γ = 2 and the eigenvalues of Σ X are two equally weighted point masses, for which m α has a simple form and one may analytically check the monotonicity).
We conjecture that the bias is monotone for α ∈ [0 , 1] for a much wider class of Σ X , as shown in Figure 23.
[Figure]
(b) γ = 5 .
D. P ROOF OF P ROPOSITION
Proof.
Taking the derivative of V ( θ P ( t )) w.r
[Figure] t. time yields (omitting the scalar σ 2 ), d V ( θ P ( t ))
[Equation]
[Figure]
[Equation] [Figure] [Equation] where we deﬁned ¯ X = XP 1 / 2 and S P = XP X (cid:62) in (i), and (ii) is due to (A2-3) the inequality tr( AB ) ≥ λ min ( A )tr( B ) for positive semi-deﬁnite A and B .
D. P ROOF OF P ROPOSITION
Proof.
Recall the deﬁnition of the bias (well-speciﬁed) of ˆ θ P ( t ) ,
[Equation]
(D [Equation] where we deﬁned S P = XP X (cid:62) , W P ( t )
=
I n − exp (cid:0) − tn S P (cid:1) in (i), ¯ X = XP 1 / 2 in (ii), and (iii) is due to the inequality tr (cid:16) A (cid:62) A (cid:17) ≥ tr (cid:0)
A 2 (cid:1) .
When Σ X = Σ − 1 θ , i.e. NGD achieves lowest stationary bias, (D.17) simpliﬁes to
[Equation]
) where ¯ λ is the eigenvalue of S P .
On the other hand, since F = Σ X , for the NGD iterate ˆ θ F − 1 ( t ) we have [Equation] D
[Equation]
19) where ˆ X =
X Σ − 1 / 2 X
and ¯ λ is the eigenvalue of S F −
=
ˆ X ˆ X (cid:62) .
Comparing (D.18)(D.19), we see that given ˆ θ P ( t ) at a ﬁxed t, if we run NGD for time T > ¯ λ max ˆ λ min t (note that T/t
= O (1) by (A2-3)), then we have B ( θ P ( t )) ≥ B ( θ F − 1 ( T )) for any P satisfying (A3).
This thus implies that
B opt ( θ P ) ≥ B opt ( θ F − 1 ) .
On the other hand, when Σ θ = I d
[Equation]
we can show that the bias term of GD is monotonically decreasing through time by taking its derivative,
[Equation]
Similarly, one can verify that the expected bias of NGD is monotonically decreasing for all choices of Σ X and Σ θ satisfying (A2-4),
[Equation]
where (i) follows from calculation similar to (D.20).
Since the expected bias is decreasing through time for both GD and NGD when Σ θ = I d , and from Theorem 2 we know that B (ˆ θ I ) ≤ B (ˆ θ F − 1 ) , we conclude that B opt ( θ I ) ≤ B opt ( θ F − 1 ) .
D. P ROOF OF T HEOREMD.8. S ETUP AND M AIN R ESULT
We restate the setting and assumptions.
H is an RKHS included in L 2 ( P X ) equipped with a bounded kernel function k satisfying sup supp( P X ) k
( x , x ) ≤ 1 .
K x ∈ H is the Riesz representation of the kernel function k
( x, · ) , that is, k ( x , y ) = (cid:104) K x , K y (cid:105) H .
S is the canonical embedding operator from
H to L 2 ( P X ) .
We write Σ = S ∗ S :
H → H and L = SS ∗ .
Note that the boundedness of the kernel gives (cid:107) Sf (cid:107) L 2 ( P X ) ≤ sup x | f ( x ) | = sup x |(cid:104) K x , f (cid:105)| ≤ (cid:107) K x (cid:107) H (cid:107) f (cid:107) H ≤ (cid:107) f
(cid:107) H .
Hence we know (cid:107) Σ (cid:107) ≤ 1 and (cid:107) L (cid:107) ≤ 1 .
Our analysis will be made under the following assumptions.
• There exist r ∈ (0 , ∞ ) and M > 0 such that f ∗ =
L r h ∗ for some h ∗ ∈
L 2 ( P X ) and (cid:107) f ∗ (cid:107) ∞ ≤ M .
• There exists s > 1 s.t.
tr (cid:0) Σ 1 /s (cid:1) < ∞ and 2 r + s − 1 > 1 .
• There exist µ ∈ [ s − 1 , 1] and C µ > 0 such that sup x ∈ supp
( P X ) (cid:13) (cid:13) Σ 1 / 2 − 1 /µ K
x (cid:13) (cid:13) H ≤ C µ .
(A5)(A6) are standard regularity assumptions in the literature that provide capacity control of the RKHS (e
g., see Caponnetto & De Vito (2007) Pillaud-Vivien et al.
It is also worth noting that most previous works (including Rudi et al.
(2017)) consider r ≥ 1 / 2 which implies that f ∗ ∈H .
The training data is generated as y i =
f ∗ ( x i ) + ε i , where ε i is an i.i.d.
noise satisfying | ε i | ≤ σ almost surely.
Let y ∈ R n be the label vector.
We identify R n with L 2 ( P n ) and deﬁne
[Equation]
We consider the following preconditioned update on f t ∈ H :
[Equation]
We brieﬂy comment on how our analysis differs from Rudi et al.
(2017), which showed that a preconditioned 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” coefﬁcient ( α in our update, λ in FALKON).
In FALKON, λ turns the objective into kernel ridge regression; whereas in our (4.1), α controls the interpolation between GD and NGD.
As 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 ﬁrst point, since FALKON minimizes a regularized objective, it would not overﬁt even after large number of gradient steps, yet it is unclear how preconditioning impacts the generalization error (i.e., any preconditioner may generalize well with proper regularization).
In contrast, we consider the “unregularized” objective, and thus early stopping plays a crucial role in avoiding overﬁtting; this is different from most standard analysis on gradient descent.
• Algorithm-wise, the two updates employ different preconditinoers.
FALKON involves inverting the kernel matrix K deﬁned on the training points, whereas we 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, our analysis allows for r < 1 / 2 , whereas Rudi et al.
(2017) and many other previous works assumed r ∈ [1 / 2 , 1] , as commented above.
We aim to show the following theorem:
Theorem (Restatement of Theorem 7) .
Given (A4-6), if the sample size n is sufﬁciently large so that 1 / ( nλ ) (cid:28) 1 , then for η < (cid:107) Σ (cid:107) with ηt ≥ 1 and 0 < δ < 1 and 0 < λ < 1 , it holds that
[Equation]
with probability 1 − 3 δ , where C is a constant and
B ( t ) := exp( − ηt ) ∨ (cid:18)
[Equation]
V ( t ) :=
[Equation]
[Equation]
, in which
V
[Equation]
) :=
[Equation]
for β (cid:48)
= log (cid:18) 28 C
[Equation]
δ (cid:19) .
When r ≥ 1 / 2 , if we set λ = n − s 2 rs +
=: λ ∗ and t = Θ(log( n )) , then the overall convergence rate becomes
[Equation]
, which is the minimax optimal rate ( ˜ O p ( · ) hides a poly- log( n ) factor).
On the other hand, when r < 1 / 2 , the bound is also (cid:101)
O p (cid:16) n − 2 rs 2 rs +1 (cid:17) except the term V 1 ( t ) .
In this case, if 2 r ≥ µ holds additionally, we have V t ( t ) =
(cid:101) O p (cid:16) n −
[Equation] rs [Equation] 1 (cid:17) , which again recovers the optimal rate.
Note that if the GD (with iterates ˜ f t ) is employed, from previous work (Lin & Rosasco, 2017) we know that the bias term (cid:16) ληt (cid:17) 2 r is replaced by (cid:16) 1 ηt (cid:17) 2 r , and therefore the upper bound translates to (cid:107) S [Equation] P [Equation]
n
[Equation]
with high probability.
In other words, by the condition η
= O (1) , we need t =
Θ( n 2 rs 2 rs +1 ) steps to sufﬁciently diminish the bias term.
In contrast, the preconditioned update that interpolates between
GD and NGD (4.1) only require t =
O (log( n )) steps to make the bias term negligible.
This is because the NGD ampliﬁes the high frequency component and rapidly captures the detailed “shape” of the target function f ∗ .
D.8. P ROOF OF M AIN R ESULT
Proof.
We follow the proof strategy of Lin & Rosasco (2017).
First we deﬁne a reference optimization problem with iterates ¯ f t that directly minimize the population risk:
[Equation]
Note that E [ f t ] = ¯ f t .
In addition [Equation] we deﬁne the degrees of freedom and its related quantity as
[Equation]
We can see that the risk admits the following bias-variance decomposition
[Equation]
We upper bound the bias and variance separately.
Bounding the bias term B ( t ) : Note that by the update rule (4.1) [Equation] it holds that
[Equation]
Therefore, unrolling the recursion gives S ¯ f t − f ∗
=
( I − ηS (Σ + λI ) − 1 S ∗ )
t ( S ¯ f 0 − f ∗ ) =
( I − ηS (Σ + λI ) − 1 S ∗ ) t ( − f ∗ )
= − ( I − ηS (Σ + λI ) − 1 S ∗ ) t L r h ∗ .
Write the spectral decomposition of L as L = (cid:80) ∞
j =1 σ j φ j φ ∗ j for φ j ∈ L 2 ( P X ) for σ j ≥ 0 .
We have (cid:107) ( I − ηS (Σ + λI ) − 1 S ∗ )
t L r h ∗ (cid:107) L 2 ( P X ) =
(cid:80) ∞
j =1 (1 − η σ j σ j + λ ) 2 t σ 2 rj h 2 j , where h = (cid:80) ∞
j =1 h
j φ j .
We then apply
Lemma 11 to obtain
[Equation]
where C is a constant depending only on r .
Bounding the variance term V ( t ) : We now handle the variance term V ( t ) .
For notational convenience, we write A λ :=
A + λI for a linear operator A from a Hilbert space H to H .
By the deﬁnition of f t , we know
[Equation]
S ∗ Y = Σ −
[Equation]
λ η   [Equation] = ( I [Equation]
ˆΣΣ
[Equation]
−
[Equation]
S
[Equation] Y [Equation] : Σ −[-1 j =
-]
[Equation] G t Σ − [Equation] ˆ S
[Equation]
{+Y, where we deﬁned G t := η (cid:104)(cid:80) t − 1 j =
+}
( I − η Σ − 1 / 2 λ ˆΣΣ − 1 / 2 λ )
j (cid:105) .
Accordingly, we decompose V ( t ) as
[Equation]
(cid:125)
[Equation]
We bound ( a ) and ( b ) separately.
Step 1.
Bounding ( a ) .
Decompose ( a ) as
[Equation]
[Equation]
We bound the terms in the RHS individually.
[Equation]
Proposition 6 of Rudi & Rosasco (2017) and its proof implies that for λ ≤ (cid:107) Σ (cid:107) and 0 < δ < 1 , it holds that
[Equation]
with probability 1 − δ
[Equation]
where β = log
[Equation]
By Lemma 14,
[Equation]
and F ∞ ( λ ) ≤ λ − 1 .
Therefore, if λ = o ( n − 1 log( n )) and λ = Ω( n − 1 /s ) , the RHS can be smaller than 1 / 2 for sufﬁciently large n , i.e. Ξ n
= O ( (cid:112) log( n ) / ( nλ )) ≤ 1 / 2 .
In this case we have,
[Equation]
We denote this event as E 1 .
(iii) Note that
[Equation]
Thus, by Lemma 12 we have
[Equation]
(iv)
Note
that
[Equation]
Σ − 1 / 2 λ
[Equation]
First we bound the ﬁrst term of the RHS.
Let ξ i
= Σ − 1 / 2 λ
[Equation]
Then, { ξ i } ni =1 is an i [Equation] i
[Equation]
d. sequence of zero-centered random variables taking value in H and thus we have
[Equation]
The RHS can be bounded by using Bernstein’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
[Equation]
[Equation]
The sup-norm can be bounded as follows.
Observe that (cid:107) ¯ f t (cid:107) ∞ ≤ (cid:107) ¯ f t (cid:107) H , and thus by Lemma 13,
[Equation]
ξ
[Equation]
≤
[Equation]
Therefore, for 0 < δ < 1 , Bernstein’s inequality (see Proposition 2 of Caponnetto & De Vito (2007)) yields that
[Equation]
with probability 1 − δ where C is a universal constant.
We deﬁne this event as E 2 .
For the second term (cid:107) Σ − 1 / 2 λ [Equation]
we have
[Equation]
B
[Equation] Combining these evaluations, on the event E 2 where P ( E 2 ) ≥ 1 − δ for 0 < δ < 1 we have
[Equation]
where we used Lemma 14 in (i).
Step 2.
Bounding ( b ) .
On the event E 1 , the term ( b ) can be evaluated as
[Equation]
where we used Lemma 13 in the last inequality.
The term (cid:107) ( G t Σ − 1 / 2 λ ˆΣΣ − 1 / 2 λ −
[Equation]
Σ 1 / 2 λ f t (cid:107) H can be bounded as follows.
First, note that
[Equation]
Therefore
[Equation]
the RHS of (D.22) can be further bounded by (cid:107) ( I − η Σ
[Equation]
λ [Equation] = (cid:107) ( I − η Σ − 1 / 2 λ
[Equation] [Equation] [Table]
[Equation]
where (i) is due to exchangeability of Σ λ and Σ .
By Lemma 11, for the RHS we have
[Equation]
Next, as in the (D.21), by applying the Bernstein inequality for asymmetric operators (Corollary 3. of Minsker (2017) with the argument in its Section 3.2), it holds that
[Equation]
n
[Equation]
n
[Equation]
[Equation]
with probability 1 − δ , where
C (cid:48) is a universal constant and β (cid:48) ≤ log
[Equation]
(cid:19) . We also used the following bounds on the sup-norm and the second order moments: [Equation] λ
[Equation]
[Equation]
( 2nd order moment 1 ) [Equation] ( 2nd order moment 2 ) [Equation]
E x
[Equation]
Σ −
[Equation] r λ [Equation] =
[Equation]
We deﬁne this event as E 3 .
Therefore, the RHS of (D.23) can be further bounded by [
[Equation]
λ
[Equation] [Equation] rλ [Equation]
+ Ctη
[Equation]
Finally, note that when λ = λ ∗ and 2 r ≥ µ ,
[Equation]
Step 3.
Combining the calculations in Step 1 and 2 leads to the desired result.
D.8. A UXILIARY LEMMAS
Lemma 11.
For t ∈ N , 0 < η < 1 , 0 < σ ≤ 1 and 0 ≤ λ , it holds that
[Equation]
Proof.
When σ ≥ λ , we have
[Equation]
due to σ ≤ 1 .
On the other hand, note that
[Equation]
where we used sup x> 0 exp( − x ) x r = ( r/e ) r .
Lemma 12.
For t = N , 0 < η and 0 ≤ σ such that ησ < 1 , it holds that η (cid:80) t − 1 j =0 (1 − ησ )
j σ ≤ 1 .
Proof.
If σ = 0 , then the statement is obvious.
Assume that σ > 0 , then
[Equation]
This yields the desired claim.
Lemma 13.
Under (A5-7), for any 0 < λ < 1 and q ≤ r , it holds that
[Equation]
Proof.
Recall that
[Equation]
Therefore [Equation] we obtain the following [Equation]
=
[Equation]
λ
[Equation]
λ
[Equation]
[Table]
Lemma 14.
Under (A5-7) and for λ ∈ (0 , 1) , it holds that N ∞ ( λ ) ≤ tr (cid:0) Σ 1 /s (cid:1) λ − 1 /s , and F ∞ ( λ ) ≤ 1 /λ .
Proof.
For the ﬁrst inequality, we have
[Table]
[Equation]
As for the second inequality, note that
[Equation]
D. P ROOF OF P ROPOSITION
Proof.
Part (a) is a simple combination of Bai & Yin (2008, Theorem 2) and assumption (A3) [Equation] which implies (cid:107) Σ X (cid:107) 2 and (cid:107) Σ − 1 X (cid:107) 2 are both ﬁnite.
For part (b) [Equation] the substitution error for the variance term (ignoring the scalar σ 2 ) can be bounded as
[Equation]
where we deﬁned S = XF − 1 X (cid:62) and ˆ S =
X ˆ F − 1 X (cid:62) in (i) and applied tr( AB ) ≤ λ max ( A +
A (cid:62) )
tr( B ) for positive semi-deﬁnite B
[Equation] as well as tr( AB ) ≤ √ d (cid:107) A (cid:107) 2 (cid:107) B (cid:107) F [Equation] and (ii) is due to (A3), ψ > 1 , Wedin (1973, Theorem 4.1) and the following estimate,
[Equation]
where (i) again follows from Wedin (1973, Theorem 4.1), and (ii) is due to (A1)(A3) and ψ >(which implies (cid:107) n u S − 1 (cid:107) 2 and (cid:107) n u ˆ S − 1 (cid:107) 2 are bounded a.s.).
Finally, from part (a) we know that ψ =
Θ( (cid:15) − 2 ) sufﬁces to achieve (cid:15) -accurate approximation of F in spectral norm.
The substitution error for the bias term can be derived from similar calculation, the details of which we omit.
D. P ROOF OF P ROPOSITION
Proof.
Since Σ θ = Σ − r X , we can simplify the expressions by deﬁning υ x (cid:44) h and thus υ θ = h − r .
From Theorem 2 we have the following derivation of the GD bias under (A1)(A3),
[Equation]
{+where m 
+}
= lim λ
→ 0 + m ( − λ ) , and m satisﬁes
[Equation]
Similarly, for NGD ( P = Σ − 1 X ) we have
[Equation]
25) where standard calculation yields m 2 = ( γ − 1) − 1 , and thus B (ˆ θ F − 1 ) → (1 − γ − 1 )
E h 1 − r .
To compare the magnitude of (D.24) and (D.25), observe the following equivalence.
[Equation]
where (i) follows from the deﬁnition of m 1 and we deﬁned ζ (cid:44) h · m 1 .
Note that when r ≥ 1 and h is not a point mass, we have
[Equation]
On the other hand, when r ≤ 0 , following the exact same procedure we get
[Equation]
Combining the two cases completes the proof.
D. P ROOF OF C OROLLARY
Proof.
Note that in this setting υ x takes value of 21+ κ and 2 κ 1+ κ with probability 1/2 each.
From (D.25) one can easily verify that for NGD,
[Equation]
For GD
[Equation]
the bias formula (D.24) can be simpliﬁed as
[Equation]
On the other hand
[Equation]
from standard numerical calculation one can show that when γ > 1 , κ ≥ 1 ,
[Equation]
Setting B (ˆ θ I ) =
B (ˆ θ F − 1 ) and solve for r
[Equation]
we have
[Equation]
where
[Equation]
Hence, Proposition 9 (from which we know r ∗ ∈ (0 , 1) ) and the uniqueness of (D.27) implies that when r ≥ r ∗ , B (ˆ θ I ) ≥ B (ˆ θ F − 1 ) , and vice versa.
Finally, observe that in the special case of γ = 2 , m 1 = κ +12 √ κ .
Therefore, one can check that constants in (D.27) simpliﬁes to
[Equation] which implies that r ∗ = 1 / 2 . E E XPERIMENT S ETUP E. P ROCESSING THE D ATASETS To obtain extra unlabeled data to estimate the Fisher, we zero pad pixels on the boarders of each image before randomly cropping; a random horizontal ﬂip is also applied for CIFAR10 images (Krizhevsky et al., 2009).
We preprocess all images by dividing pixel values by 255 before centering them to be located within [ − 0 .
5] with the subtraction by 1 / 2 .
For experiments on
CIFAR10, we downsample the original images using a max pooling layer with kernel size 2 and stride 2.
E. S ETUP AND I MPLEMENTATION FOR O PTIMIZERS
In all settings, GD uses a learning rate of 0 .
01 that is exponentially decayed every 1k updates with the parameter value 0 .
For NGD, we use a ﬁxed learning rate of 0 . 03 .
Since inverting a parameter-by-parameter-sized Fisher estimate per iteration would be costly, we adopt the Hessian free approach (Martens, 2010) which computes approximate matrix-inverse-vector products using the conjugate gradient (CG) method (Nocedal & Wright, 2006; Boyd et al., 2004).
For each approximate inversion, we run CG for 200 iterations starting from the solution returned by the previous CG run.
The precise number of CG iterations and the initialization heuristic roughly follow Martens &
Sutskever (2012).
For the ﬁrst run of CG, we initialize the vector from a standard Gaussian, and run
CG for 5k iterations.
To ensure invertibility, we apply a very small amount of damping ( 0 . 00001 ) in most scenarios.
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. O THER D ETAILS
For experiments in the label noise and misspeciﬁcation sections, we pretrain the teacher using the
Adam optimizer (Kingma & Ba, 2014) with its default hyperparameters and a learning rate of 0 . 001 .
For experiments in the misalignment section, we downsample all images twice using max pooling with kernel size 2 and stride 2.
Moreover, only for experiments in this section, we 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.