Generalizable Learning to Optimize into Wide Valleys

Learning to optimize ( L2O ) has gained increasing popularity in various optimization tasks , since classical optimizers usually require laborious , problem-specific design and hyperparameter tuning . However , current L2O approaches are designed for fast minimization of the objective function value ( i.e. , training error ) , hence often suffering from poor generalization ability such as in training deep neural networks ( DNNs ) , including ( i ) disappointing performance across unseen optimizees ( optimizer generalization ) ; ( ii ) unsatisfactory test-set accuracy of trained DNNs ( optmizee generalization ) . To overcome the limitations , this paper introduces flatness-aware regularizers into L2O for shaping the local geometry of optimizee ’ s loss landscape . Specifically , it guides optimizee to locate wellgeneralizable minimas in large flat regions of loss surface , while tending to avoid sharp valleys . Such optimizee generalization abilities of flatness-aware regularizers have been proved theoretically . Extensive experiments consistently validate the effectiveness of our proposals with substantially improved generalization on multiple sophisticated L2O models and diverse optimizees . Our theoretical and empirical results solidify the foundation for L2O ’ s practically usage . All codes and pre-trained models will be shared upon acceptance . 1 INTRODUCTION One cornerstone of deep learning ’ s success is arguably the stochastic gradient-based optimization methods , such as SGD ( Robbins & Monro , 1951 ) , Adam ( Kingma & Ba , 2014 ) , AdaGrad ( Duchi et al. , 2011 ) , RProp ( Riedmiller & Braun , 1993 ) , and RMSProp ( Tieleman & Hinton , 2012 ) . The performance of deep neural networks ( DNNs ) hinges on the choice of optimization methods and the corresponding parameter settings . Thus , intensive human labor is often required to empirically select the best optimization method and its parameters for each specific problem . A seemingly promising data-driven approach , learning to optimize ( L2O ) , arose from the meta learning community to alleviate this issue . It aims to replace traditional optimizers tuned by human hand with neu- ral network based optimizers that can learn update rules from data . Existing works have demonstrated that a learned optimizer is not only able to decrease the objective function faster , but also able to tremendously reduce the required human labor . Andrychowicz et al . ( 2016a ) first proposed to parameterize the update rules using a long short-term memory ( LSTM ) network . The LSTM optimizer tries to simulate the behavior of iterative methods by unrolling . By aggregating the loss of the function to be optimized ( optimizee ) at each time step , it aims to minimize the overall loss along the optimization path . Wichrowska et al . ( 2017 ) enlarged the optimizer neural network to a hierarchical recurrent neural network ( RNN ) architecture , which improves its capability on larger or unseen optimization problems . Li & Malik ( 2016 ) also proposed a reinforcement learning based approach for this sub-field . Although L2O methods are able to achieve better performance than analytical optimizers in many traditional optimization tasks , they are yet mature to serve as practical optimizers for deep neural networks . The main reason is that all existing L2O methods are designed for solving traditional optimization problems where the generalization ability is not a concern . Generalization ability , one of the core problems in machine learning , is neither guaranteed for deep neural networks optimized by L2O , nor studied by literature in the context of L2O before . It has been shown that for today ’ s heavily overparameterized networks , it is easy to memorize the entire training data which lead to zero training loss ( Zhang et al. , 2021 ) . Therefore , for a L2O method , fast minimization of the training error can not really lead to generalization ability of the optimizee . The generalization ability of machine learning methods has been extensively studied both theoretically and empirically in the past decades ( Keskar et al. , 2017 ; Hochreiter & Schmidhuber , 1997 ; 1994 ; Jiang et al. , 2020 ; Chaudhari et al. , 2017 ; Damian et al. , 2021 ) . Among them , a prevailing theory that links the generalization ability of models to the geometry of the loss landscape has recently shown its empirical effectiveness . Most methods taking this idea measure the geometry of the loss landscape using Hessian spectrum , while some others use the local entropy ( Damian et al. , 2021 ; Chaudhari et al. , 2017 ) . Further , Foret et al . ( 2021 ) proposes the SAM method , which minimizes the loss value and the loss sharpness simultaneously . Specifically , an optimization method that prefers to converge to wide valleys ( i.e. , flat basin ) in loss landscape shows better generalization ability . Inspired by this , we propose flatness-aware regularizers i.e. , Hessian regularizer and Entropy regularizer for learning to optimize . The Hessian regularization is the pseudonorm of the optimizee ’ s Hessian matrix and the Entropy regularization measures the local Entropy of the optimizee . Both of them characterize the geometry of the loss landscape which intends to teach optimizer favor update rules that lead to wide valleys in loss landscape . To summarize , the contributions of this paper can be outlined below : • We propose to use the flatness-aware regularizers in the training of L2O optimizers . The update rules learned with such regularizers would favor converging to wide valleys in loss landscape when minimizing loss functions . Thus , L2O optimizers trained with flatnessaware regularizers can be deemed as a plug-and-play optimizer that favors generalization and requires no more time calculating Hessian or Entropy information while in use . • We demonstrate theoretically that adopting flatness-aware regularizers in L2O can enhance the generalization ability of optimizees trained by regularized optimizers . Note that the theoretical result of Entropy regularizer implies Entropy-SGD favors wide valleys in original loss landscape rather than only in Entropy energy landscape . • Comprehensive experiments over various tasks demonstrate the effectiveness of our methods . Empirical results show that our methods significantly improve the generalization ability of existing L2O methods , enabling them to outperform current state-of-the-art by a large margin . 2 RELATED WORK . Learning to Optimize ( L2O ) As a special case of learning to learn , L2O has been widely investigated in various machine learning problems , including black-box optimization ( Chen et al. , 2017 ) , Bayesian swarm optimization ( Cao et al. , 2019 ) , minmax optimization ( Shen et al. , 2021 ) , domain adaptation ( Li et al. , 2020 ; Chen et al. , 2020b ) , adversarial training ( Jiang et al. , 2018 ; Xiong & Hsieh , 2020 ) , graph learning ( You et al. , 2020 ) , and noisy label training ( Chen et al. , 2020c ) . The first L2O framework dates back to Andrychowicz et al . ( 2016b ) , in which the gradients and update rules of optimizee are formulated as the input features and outputs for a RNN optimizer , respectively . It proposes a coordinate-wise design which enables trained L2O to be applicable for different neural networks with diverse amount of parameters . Li & Malik ( 2016 ) proposes an alternative reinforcement learning frameworks for L2O , leveraging gradient history and objective values as observations and step vectors as actions . Later on , more advanced variants arise to power up the generalization ability of L2O . For example , ( i ) regularizers like random scaling , objective convexifying ( Lv et al. , 2017 ) , and Jacobian regularization ( Li et al. , 2020 ) , ( ii ) enhanced L2O model like hierarchical RNN architecture ( Wichrowska et al. , 2017 ) , and ( iii ) improved training techniques like curriculum learning and imitation learning ( Chen et al. , 2020a ) . A more comprehensive literature is referred to a recent survey paper of learning to optimize . Flatness on Generalization of Neural Network Generalization analysis of neural networks has been widely studied by various methods , including VC-dimension ( Bartlett et al. , 2019 ) , covering number ( Bartlett et al. , 2017 ) , stability ( Hardt et al. , 2016 ; Zhou et al. , 2018a ) , Rademacher complexity ( Golowich et al. , 2018 ; Ji & Liang , 2018 ; Ji et al. , 2021 ; Arora et al. , 2018 ; 2019 ) , etc . In particular , the landscape flatness has been known to be associated with better generalization . On the empirical side , Keskar et al . ( 2017 ) and He et al . ( 2019 ) showed that minima in wide valleys often generalize better than those in sharp basins . Further , Wilson et al . ( 2017 ) and Keskar & Socher ( 2017 ) showed empirically that SGD favors better generalization solutions than Adam . On the theory side , Zhou et al . ( 2020 ) showed that SGD is more unstable at sharp minima than Adam and explained why SGD generalize better than Adam theoretically and Zou et al . ( 2021 ) explained that the inferior generalization performance of Adam is connected to nonconvex loss landscape . To improve the generalization performance , Entropy-SGD was introduced in Chaudhari et al . ( 2017 ) which was shown to outperform SGD in terms of the generalization error and the training time . Meanwhile , the spectral norm regularization has been proposed in Yoshida & Miyato ( 2017 ) to improve the generalization ability of neural networks empirically . 3 METHODOLOGY . In this section , we provide basic notations and the detailed formulations about our flatness-aware regularizers , i.e. , Hessian and Entropy regularizers . 3.1 PRELIMINARY . We define F ( 1 ) ( θ ; ξ ) and F ( 2 ) ( θ ; ζ ) respectively as the non-negative meta-traning and meta-testing functions , where θ ∈ Rp is the optimizee parameter , and ξ and ζ respectively denote training and testing data samples . Suppose there are N training data samples ξ ∈ { ξi , i = ( 1 , . . . , N ) } and M testing data samples ζ ∈ { ζj , j = ( 1 , . . . , M ) } . Then we define the empirical meta-training and meta-testing functions and their corresponding population risk functions as follows : ( Meta-Training ) L ( 1 ) N ( θ ) = 1 N N∑ i=1 F ( 1 ) ( θ ; ξi ) , L ( 1 ) ( θ ) = EξF ( 1 ) ( θ ; ξ ) , ( 1 ) ( Meta-Testing ) L ( 2 ) M ( θ ) = 1 M M∑ j=1 F ( 2 ) ( θ ; ζj ) , L ( 2 ) ( θ ) = EζF ( 2 ) ( θ ; ζ ) . ( 2 ) An L2O algorithm aims to learn an update rule for optimizee θ based on the meta-training function . An update rule can be expressed as θ ( 1 ) t+1 ( φ ) = θ ( 1 ) t ( φ ) + m ( z ( 1 ) t ; φ ) , where t = 0 , 1 , . . . , T − 1 denotes the iteration index over one epoch , the variable z captures the information ( e.g. , loss values , gradients ) that we collect on the optimization path , and the optimizer function m ( z ; φ ) is parameterized by φ and captures how the update of the optimizee parameter θ depends on the loss landscape information included in z . In order to find a desirable optimizer parameter φ , L2O solves the following meta-training problem : min φ { L ( 1 ) N ( θ ( 1 ) T ( φ ) ) } where θ ( 1 ) t+1 ( φ ) = θ ( 1 ) t ( φ ) +m ( z ( 1 ) t ; φ ) . ( 3 ) A popular L2O meta-training algorithm applies the gradient descent method , which updates φ based on the gradient of the objective function L ( 1 ) N ( θ ( 1 ) T ( φ ) ) with respect to φ . As suggested by eq . ( 3 ) , each update of φ requires T iterations of the optimizee parameter θ ( 1 ) 0 ( φ ) to obtain θ ( 1 ) T ( φ ) . In meta-testing , we apply the output φ of meta-training and its corresponding optimizer to update the optimizee as θ ( 2 ) t+1 ( φ ) = θ ( 2 ) t ( φ ) + m ( z ( 2 ) t ; φ ) ( t = 0 , 1 , . . . , T − 1 ) . Note that we differentiate the optimizee updates in training and testing by superscripts ( 1 ) and ( 2 ) , respectively .

Summary: This paper aims to solve the poor generalization performance of Learning to Optimize (L2O) methods by introducing flatness-aware regularizers. The idea is to add a regularization term to the meta-training objective that encourages the final iterate obtained by applying the learned optimizer, to lie in a flat region (measured by the spectral norm of the Hessian, or the local entropy function). The authors prove a bound for the generalization error, and show empirically that adding flatness-aware regularizers to existing L2O methods, can improve generalization or both the learned optimizer (across unseen tasks) and the solution (across dataset splits).