When does preconditioning help or hurt generalization?

1 INTRODUCTION . 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 : θt+1 = θt − ηP ( t ) ∇θtL ( fθt ) , t = 0 , 1 , . . . ( 1.1 ) 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 ( LeCun et al. , 2012 ) ; and diagonal matrices estimated from past gradients , which include various 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 impact on generalization has been under debate : Wilson et al . ( 2017 ) reported that in neural network optimization , adaptive or secondorder 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. , 2018a ) , i.e . preconditioning may lead to a 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 ` 2 norm . For instance in overparameterized least squares regression , GD and many first-order methods find the minimum ` 2 norm solution from zero initialization ( without explicit regularization ) , but preconditioned updates may not . This being said , while the minimum ` 2 norm solution can generalize well ∗Alphabetical ordering . Correspondence to : Denny Wu ( dennywu @ cs.toronto.edu ) . in the overparameterized regime ( Bartlett et 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 3 we start with overparameterized least squares regression ( unregularized ) and analyze the stationary solution ( t→∞ ) of update ( 1.1 ) under time-invariant preconditioner . Extending previous analysis in the proportional limit ( Hastie et al. , 2019 ) , we consider a more general random design setting and derive the exact population risk in its bias-variance decomposition . 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 matrix1 . We find that the comparison of generalization is affected by the following factors : 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 a perfect fθ that recovers the true function ( target ) . We argue that this factor is similar to additional label noise , and thus NGD may also be beneficial when the model is misspecified . 3 . Data-Signal-Alignment : Alignment describes how the target signal distributes among the input 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 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 findings 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 influence of all aforementioned factors . The performance of neural networks under a variety of manipulations results in trends that align with our theoretical analysis . 2 BACKGROUND AND RELATED WORKS . 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θ ) = 1n ∑n i=1 l ( yi , fθ ( xi ) ) , where l : Y ×Z → R. Also suppose a probability distribution p ( y|z ) = p ( y|fθ ( x ) ) is defined on the space of labels . Then , the natural gradient is defined as : ∇̃θL ( X , fθ ) = F−1∇θL ( X , fθ ) , where F = E [ ∇θ log p ( x , y|θ ) ∇θ log p ( x , y|θ ) > ] 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 defined under the empirical data distribution { xi } ni=1 ( Amari et al. , 2000 ) . We instead refer to this quantity as the sample Fisher , the properties of which influence optimization and have been studied in 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 . 1From now on we use NGD to denote the population Fisher-based update , and we write “ sample NGD ” when P is the inverse or pseudo-inverse of the sample Fisher ; see Section 2 for discussion . 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 findings 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 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 advantage 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 updates 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 . 3 ASYMPTOTIC RISK OF RIDGELESS INTERPOLANTS . In this section we consider the following setup : given n training samples { xi } ni=1 labeled by a teacher model ( target function ) f∗ : Rd→R with additive noise : yi = f∗ ( xi ) + εi , we learn a linear student model fθ by minimizing the squared loss : L ( X , fθ ) = ∑n i=1 ( yi − x > i θ ) 2 . We assume a random design : xi = Σ 1/2 X zi , where zi ∈ Rd is an i.i.d . vector with zero-mean , unit-variance , and finite 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 ) = Ex [ ( 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 , and there exist multiple empirical risk minimizers with potentially different generalization properties . Denote X = [ x > 1 , ... , x > n ] > ∈ Rn×d the data matrix and y ∈ Rn the corresponding label vector . We optimize the parameters θ via a preconditioned gradient flow with preconditioner P ( t ) ∈ Rd×d , ∂θ ( t ) ∂t = −P ( t ) ∂L ( θ ( t ) ) ∂θ ( t ) = 1 n P ( t ) X > ( y −Xθ ( t ) ) , θ ( 0 ) = 0 . ( 3.1 ) 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 F = ΣX . We thus limit our analysis to fixed preconditioner of the form P ( t ) = : P . Write parameters at time t under update ( 3.1 ) with fixed P as θP ( t ) . For positive definite P , the stationary solution is given as : θ̂P : =limt→∞ θP ( t ) =PX > ( XPX > ) −1y . One may check that 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 ` 2 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 . Remark . θ̂P is the minimum ‖θ‖P−1 norm interpolant : θ̂P = arg minθ‖θ‖P−1 , s.t.Xθ=y for positive definite P . For GD this translates to the parameter ` 2 norm , whereas for NGD ( P =F−1 ) , the implicit bias is the ‖θ‖ΣX norm . Since E [ f ( x ) 2 ] =‖θ‖2ΣX , NGD finds 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 ( see Figure 1 and Appendix A.1 ) . We highlight the following choices of P and the corresponding stationary solution θ̂P as t→∞ . • Identity : P =Id recovers GD that converges to the min ` 2 norm interpolant ( also true for momentum GD and SGD ) , which we write as θ̂I : = X > ( XX > ) −1y and refer to as the GD solution . • Population Fisher : P = F−1 = Σ−1X leads to the estimator θ̂F−1 , which we refer to as the NGD solution . • Sample Fisher : since the sample Fisher is rank-deficient , we may add a damping term P = ( X > X+λId ) −1 or take the pseudo-inverse P = ( X > X ) † . In both cases , the gradient is still spanned by X , and thus the update finds the same min ` 2- norm 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 . 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 , R ( θ ) = Ex [ ( f∗ ( x ) − 〈x , Eε [ θ ] 〉 ) 2 ] ︸ ︷︷ ︸ B ( θ ) , bias + tr ( Cov ( θ ) ΣX ) ︸ ︷︷ ︸ V ( θ ) , variance . Note that the bias does not depend on label noise ε , 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 x and a residual : f∗ ( x ) = 〈x , θ∗〉 + f∗c ( x ) , we can separate the bias term into a wellspecified component ‖θ∗ − Eθ‖2ΣX , which captures the difficulty in learning θ ∗ , and a misspecified component , which corresponds to the error due to fitting f∗c ( beyond what the student can represent ) .

Summary: The 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.