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 authors theoretically study the prediction performance of pre-conditioned gradient descent/flow with linear models and squared loss aligning in the setting of least squares regression and non-parametric regression. For parametric least squares, the predication performance of the limiting solution for preconditioned gradient flow i.e. time goes to infinity, is studied in an asymptotic regime where both the number of samples and dimension go to infinity in proportion to one another. Meanwhile for non-parametric regression, source and capacity assumptions are leveraged to achieve finite sample guarantees. Experiments are also conducted on neural networks in a student and teacher setup.