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: The paper proposes to use meta-learning to learn an optimizer that automatically seeks wide local minima without the need to explicitly compute the sharpness measure on the fly. When training the optimizer, the parameters of the optimizer are updated in such a way that the training loss and the sharpness measure at the end of the updated trajectory are minimized. Theoretical results were provided for the generalization guarantees of the learned optimizer under some technical conditions. Finally, the experimental results demonstrate that the sharpness-aware optimizer outperforms baselines on simple models.