G ENERALIZABLE L EARNING TO O PTIMIZE INTO W IDE V ALLEYS
Anonymous authors
Paper under double-blind review
A BSTRACT
Learning to optimize (L2O) has gained increasing popularity in various optimization tasks, since classical optimizers usually require laborious, problem-speciﬁc 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 ﬂatness-aware regularizers into L2O for shaping the local geometry of optimizee’s loss landscape.
Speciﬁcally, it guides optimizee to locate wellgeneralizable minimas in large ﬂat regions of loss surface, while tending to avoid sharp valleys.
Such optimizee generalization abilities of ﬂatness-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.
I NTRODUCTION
[Figure]
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 speciﬁc 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 neural 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)
ﬁrst 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-ﬁeld.
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 cannot 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.
Speciﬁcally, an optimization method that prefers to converge to wide valleys (i.e., ﬂat basin) in loss landscape shows better generalization ability.
Inspired by this, we propose ﬂatness-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 ﬂatness-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 ﬂatnessaware 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 ﬂatness-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 signiﬁcantly improve the generalization ability of existing L2O methods, enabling them to outperform current state-of-the-art by a large margin.
R ELATED W ORK
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 ﬁrst 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 ﬂatness 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.
M ETHODOLOGY
In this section, we provide basic notations and the detailed formulations about our ﬂatness-aware regularizers, i.e., Hessian and Entropy regularizers.
3. P RELIMINARY
We deﬁne F (1) ( θ ; ξ ) and F (2) ( θ ; ζ ) respectively as the non-negative meta-traning and meta-testing functions, where θ ∈ R p 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 deﬁne the empirical meta-training and meta-testing functions and their corresponding population risk functions as follows:
[Equation]
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 ﬁnd a desirable optimizer parameter φ , L2O solves the following meta-training problem:
[Equation]
where
[Equation]
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.
3. H ESSIAN R EGULARIZER
Motivated by the idea that the optimizee in a ﬂat area of the training objective has superior generalization capability, we propose to incorporate ﬂatness-aware regularizers into L2O meta-training, in order to learn optimizers that favors to land the optimizee into a ﬂat region.
We introduce two such regularizers in this and next subsection.
The ﬁrst regularizer we introduce is based on the spectral norm of the Hessian , smaller values of which corresponds to a ﬂatter landscape.
Thus, the new L2O meta-training objective is given by:
[Equation]
where
[Equation]
where λ is the regularizer hyperparameter.
Note that the Hessian regularizer is adopted for training the optimizer parameter φ , and its impact on the update rule m ( z (1) t ; φ ) is only through φ , i.e., the information in z t does not include such regularization.
We let φ ∗ be the optimal optimizer parameter, which can be written as
[Equation]
Due to the computational intractability of directly penalizing ∇ 2 θ L (1) N ( θ (1) T ( φ )) , we investigate three approximation variants in the implementation.
(cid:172)
Hessian EV : the eigenvalue of largest module of
Hessian matrix, computed by power iteration (Yao et al., 2020); (cid:173) Hessian Trace : the trace of Hessian matrix, calculated via Hutchinson method (Yao et al., 2020); (cid:174) Jacobian Trace : the trace of
Hessian’s Jacobian approximation ∇ θ L (1) N ( θ (1) T ( φ )) (cid:62) ∇ θ L (1) N ( θ (1) T ( φ )) .
Note that such Hessian approximation methods do not involve computing Hessian explicitly which helps to reduce the memory and computational cost.
In our case, we perform 10 iterations for Hessian norms’ approximation.
3. E NTROPY R EGULARIZER
The second ﬂatness-aware regularizer we incorporate to L2O is based on the local entropy function G [Equation] [Equation]
N [Equation] = log [Equation]
[Equation]
) − [Equation] d θ (cid:48) proposed in Chaudhari et al. Due to the exponential decay with respect to (cid:107) θ − θ (cid:48) (cid:107) 2 , the integral mainly captures the value of the loss function L (1) N ( θ (cid:48) ) over the neighborhood of θ .
Thus, the value of G (1) N ( θ ; γ ) measures the ﬂatness of the local area around θ .
Thus, the L2O meta training objective with Entropy regularizer is given by:
[Equation]
where
[Equation]
[Equation]
We let φ ∗ be the optimal optimizer parameter, which can be written as
[Equation]
In order to implement the gradient descent algorithm for meta-training, the gradient −∇ φ G N ( θ (1) T ( φ ); γ ) can be calculated by the entropy gradient −∇ θ G (1) N ( θ ; γ ) and the chain rule.
In particular, as given in Chaudhari et al.
(2017), the entropy gradient takes the following form
[Equation]
(8) where ξ ∈ { ξ i , i = (1 , . . .
, N ) } are training samples and the distribution of θ (cid:48) is given by P ( θ (cid:48) ; θ, γ )
[Equation]
T HEORETICAL A NALYSIS
In this section, we ﬁrst introduce several assumptions, and then present the generalization analysis of L2O with Hessian and Entropy regularizers.
4. A SSUMPTIONS
We ﬁrst deﬁne the local basin of θ with the radius d as D d ( θ ) =
{ θ (cid:48) : (cid:107) θ − θ (cid:48) (cid:107) 2 ≤ d } .
As have been observed widely in training a variety of machine learning objectives, the convergent point enters into a local neighborhood where the strong convexity (or similar properties such as gradient dominance condition, reguarity condition, etc) holds (Du et al., 2019;
Li & Yuan, 2017; Zhou et al., 2018b;
Safran & Shamir, 2016; Milne, 2019).
We thus make the following assumption on the geometry of the meta-training function.
Assumption 1.
We assume that there exist a a local basin D d ( θ (1) T ( φ ))
( d > 0) of the convergence point θ (1) T ( φ ) that in such local basin, L (1) ( θ ) and L (1) N ( θ ) are µ -strongly convex w.r.t.
θ .
Futhermore, there exist a unique optimal point θ ∗ (1) of function L (1) ( θ ) and a optimal point θ ∗ (1) N of function L (1) N ( θ ) in local basin D d ( θ (1) T ( φ ∗ )) .
Assumption 2.
We assume that L (1) ( θ ) function is M -Lipschitz; ∇ θ L (1) ( θ ) and ∇ θ L (2) ( θ ) are
L -Lipschitz; ∇ 2 θ L (1) ( θ ) , ∇ 2 θ L (2) ( θ ) , ∇ 2 θ L (1) N ( θ ) and ∇ 2 θ L (2) M ( θ ) are ρ -Lipschitz.
We further adopt the following assumptions introduced in Mei et al.
(2018), in order to guarantee the similarity between the landscape of the empirical and population objective functions.
Assumption 3.
The gradient of the training loss ∇ F (1) ( θ ; ξ ) is τ 2 -sub-Gaussian and the Hessian of the loss function is τ 2 -sub-exponential.
See Appendix B.1 for more details.
Assumption 4.
There exist a basin D r (0) that the meta-training functions L (1) ( θ ) is ( (cid:15), η ) -strongly
Morse in D r (0) and the local basins D d ( θ ( i ) T ( φ ∗ )) for i = 1 , 2 of convergence points θ ( i ) T ( φ ∗ )( i = 1 , 2) are in D r (0) .
See Appendix B.1 for more details.
4. G ENERALIZATION A NALYSIS OF H ESSIAN - REGULARIZED L2O
In this section, we adopt the optimizer learned by regularized L2O to train a new optimizee, and analyze the advantages of the Hessian regularizer on the optimizee generalization ability.
We ﬁrst note that by optimization theory, the regularized optimization problem eq.
(5) is equivalent to the following constrained optimization
[Equation]
where
[Equation]
where B Hessian ( λ ) is the constraint bound on the Hessian determined by λ .
Thus, the optimizer parameter φ ∗ learned by the Hessian-regularized L2O meta-training in eq.
(5) is also a solution to eq.
(9), i.e., its Hessian satisfy the constraint.
Then let θ (2) T ( φ ∗ ) denote the optimizee parameters trained by optimizer φ ∗ in meta-testing, and θ ∗ (2) denote the optimal point of the population meta-testing function L (2) ( · ) .
We then characterize the generalization error as
L (2) ( θ (2) T ( φ ∗ ))
− L (2) ( θ ∗ (2) ) , which capture how well the optimizer φ ∗ performs on a testing task with respect to the best possible testing loss value.
The following theorem characterizes the generalization performance of the optimizee trained with
Hessian regularized optimizer as deﬁned above.
Theorem 1 (Generalization Error of Hessian-Regularized L2O) .
Suppose Assumptions 1, 2, 3,hold.
We let N ≥ max { 4 Cp log N/η 2 ∗ ,
Cp log p } where C = C 0 max { c h , 1 , log( rτδ ) } , η 2 ∗ =
min { (cid:15) 2 τ 2 , η 2 τ 4 , η 4 ρ 2 τ 2 } and C 0 is an universal constant.
Then, with probability at least 1 − 2 δ we have
[Equation]
(10) where A 1 = B Hessian ( λ ) + ρ ∆ ∗ T + ρ ∆ ∗ θ + ∆ ∗ H + O ( w T − T (cid:48) )
+ O ( (cid:113) C log N N ) , A 2 = ∆ ∗ T + ∆ ∗ θ + O ( w T − T (cid:48) ) + O ( (cid:113) C log N N ) , with w = L − µ L + µ , ∆ ∗ H =
(cid:107)∇ 2 θ L (2) ( θ ∗ (2) ) − ∇ 2 θ L (1) ( θ ∗ (1) ) (cid:107) , ∆ ∗ θ = (cid:107) θ ∗ (1) − θ ∗ (2) (cid:107) , ∆ ∗ T = (cid:107) θ (2) T ( φ ∗ ) − θ (1) T ( φ ∗ ) (cid:107) , and T (cid:48) is the minimum gradient descent iterations for θ (1) T (cid:48)
( GD ) to enter into the local basin of θ (1) T ( φ ∗ ) .
In Theorem 1 the generalization error is bounded by the terms A 1 and A 2 , where the Hessian regularizer affects the generalization error through the constraint B Hessian ( λ ) in A 1 .
Clearly, by choosing the regularization hyperparameter λ , we control the value of B Hessian ( λ ) , which then makes an impact on the generalization.
Speciﬁcally, larger λ corresponds to smaller B Hessian ( λ ) and hence smaller generalization error.
This also explains that ﬂatter landscape (i.e., smaller B Hessian ( λ ) on Hessian) yields better generalization performance (i.e., smaller generalization error).
The generalization error in Theorem 1 also contains other terms which we explain as follows: (a) ρ ∆ ∗ T + ρ ∆ ∗ θ + ∆ ∗ H capture the similarity between the training and testing tasks; more similar tasks yields better generalization; (b) O ( w T − T (cid:48) ) captures the exponential decay rate of the optimizee’s iteration due to the strong convexity, and vanishes for large T ; and (c) O ( (cid:113) C log N N ) arises due to the difference between the empirical and population loss functions, and vanishes as the sample size
N gets large.
4. G ENERALIZATION A NALYSIS OF E NTROPY - REGULARIZED L2O
In this section, we analyze the generalization error of the Entropy regularizer on the optimizee generalization ability.
Similarly to the Hessian regularizer, the regularized optimization problem eq.
(6) is equivalent to the following constrained optimization:
[Equation]
where B Entropy ( λ ) is the constraint bound on the Entropy determined by λ .
Thus, the optimizer φ ∗ learned by the Entropy-regularized L2O meta-training in eq.
(6) is also a solution to eq.
(11), i.e., the local entropy satisﬁes the constraint.
In the following, we ﬁrst establish an important connection between the Entropy constraint bound
B Entropy ( λ ) and the Hessian (cid:107)∇ 2 L (1) N ( θ (1) T ( φ )) (cid:107) of the optimizee.
Theorem 2 (Connection between Hessian and Entropy Regularizer) .
Suppose Assumptions 1 andhold.
We deﬁne
[Equation]
exp( − γ 2 (cid:107) θ − θ (cid:48) (cid:107) 2 ) d θ (cid:48) where m is a constant and θ ∈ R p .
Then, we have
[Equation]
where D ( x ) =
L (1) N ( θ (1) T ( φ ))
+ ( p − 1) log( γ + µ ) − mM − p 2 log(2 π ) − 12 ρm 3 − C ( γ, p, m ) + log( x + γ ) and D − 1 ( x ) denotes the inverse function of D ( x ) .
It can be observed that the function D ( x ) in monotonically increasing w.r.t.
x and so is its inverse
D − 1 ( x ) , as determined by the only x -dependent term log( x + γ ) .
Hence, Theorem 2 shows that the entropy bound constraint B Entropy ( λ ) implies a corresponding Hessian constraint D − 1 ( B Entropy ( λ )) .
Thus, Theorem 1 can be applied to provide a bound on the generalization error here.
Corollary 1 (Generalization Error of Entropy-Regularized L2O) .
Suppose the same conditions of
Theorem 1 hold.
Then the generalization error of L2O with Entropy regularizer takes the bound in eq.
(10) with B Hessian ( λ ) in A 1 being replaced by D − 1 ( B Entropy ( λ )) .
Theorem 2 and Corollary 1 establish that the bound D − 1 ( B Entropy ( λ )) serves the same role as the
Hessian bound in the generalization performance.
Thus, by controlling the hyperparameter λ to be large enough in the L2O training, B Entropy ( λ ) as well as D − 1 ( B Entropy ( λ )) and Hessian can be controlled to be sufﬁciently small.
In this way, the optimizee will be landed into a ﬂat basin (due to small Hessian) to enjoy better generalization.
To compare with the result in Chaudhari et al.
(2017), we note that
Chaudhari et al.
(2017) proposed the Entropy-SGD method and showed that Entropy-SGD favors better generalization solutions in terms of the Entropy energy landscape.
As a comparison, Theorem 2 and Corollary 1 establish the generalization error for Entropy-SGD in terms of the un-regularized loss
L (2) ( θ (2) T ( φ ∗ ))
− L (2) ( θ ∗ (2) ) in original loss landscape, which is the ultimate goal of generalization.
Such a favorable result is established by exploiting the equivalence between the regularized optimization and the un-regularized constrained optimization problems.
We further note that Dinh et al.
(2017) theoretically shows that sharp minimas can also generalize well for deep neural networks.
However, such a result does not contradict the fact that ﬂat minima generalizes well, which has strong evidence (He et al., 2019; Keskar et al., 2017), and is the property that we exploit in the paper.
[Figure]
[Figure]
E XPERIMENT
In our experiments [Figure] we consider two advanced L2O algorithms, i.e., L2O-DM-CL 1 (Chen et al., 2020a) and L2O-Scale (Wichrowska et al., 2017) on four diverse meta testing optimizees.
Meta Training Optimizees.
For training L2O-Scale, we use a three-layer convolutional neural network (CNN) which has one fully-connected layer, and two convolutional layers with eight 3 ×and 5 × 5 kernels respectively.
For training L2O-DM, we adopt the same meta training optimizee from Andrychowicz et al.
(2016b), which is a simple Multi-Layer Perceptron (MLP) with one hidden layer of 20 dimensions and the sigmoid activation function.
MNIST (LeCun et al., 1998) dataset is used for all the meta-training.
Meta Testing Optimizees.
We select four distinct and representative meta testing optimizees from
Andrychowicz et al.
(2016b) and Chen et al.
(2020a) to evaluate the generalization ability of the
Optimizees (cid:172) , (cid:173) , and (cid:174) are for evaluating the generalization of L2O across network architectures.
Then, (cid:175) evaluates the generalization of L2O across both network architectures and datasets.
Training and Evaluation details.
During the meta training stage of L2O, L2O-Scale is trained with 5 epochs, where the number of each epoch’s iteration is drawing from a heavy tailed distribution (Wichrowska et al., 2017).
L2O-DM-CL is trained with a curriculum schedule of training epochs and iterations, following the default setup in Chen et al.
(2020a).
RMSprop with the learning rate 1 × 10 − 6 is used to update L2Os.
For the { Hessian, Entropy } regularization coefﬁcients { λ Hessian , λ Entropy , γ } , we perform a grid search and choose { 5 × 10 − 5 ,-,- } / { 1 × 10 − 8 , 1 × 10 − 6 , 1 × 10 − 4 } for L2O-Scale/L2O-DM-CL.
In the meta testing stage of L2O, we compare our methods with classical optimizers like SGD and
Adam, and state-of-the-art (SOTA) L2Os such as L2O-Scale and L2O-DM-CL.
Hyperparameters of both classical optimizers and L2O baselines are carefully tuned through the grid search and all other irrelevant variables are strictly controlled for a fair comparison.
We run 10 , 000 iterations for the meta testing, and the corresponding training loss and test accuracy on all unseen optimizees are collected to evaluate the optimizer and optimizee generalization .
Note that the training loss corresponds to meta testing L (2) M and test accuracy corresponds to L (2) in Section 3.1.
We conduct ten independent replicates with different random seeds and report the average performance.
All of our experiments are conducted on a computing facility of NVIDIA GeForce GTX 1080Ti GPUs.
5. I MPROVE G ENERALIZATION WITH H ESSIAN R EGULARIZATION
In this section, we conduct extensive evaluations of our proposed Hessian regularization on previous state-of-the-art L2O methods, i.e., L2O-Scale
(Wichrowska et al., 2017) and L2O-DM-CL (Chen et al., 2020a).
Achieved training loss and testing accuracy are collected in Figure 2 and 3 which also include comparisons with representative analytical optimizers like SGD (Ruder, 2016) and
Adam (Kingma & Ba, 2014).
Several consistent observations can be drawn from our results: (cid:182) Hessian Trace regularizer consistently enhances the generalization abilities of learned L2Os and trained optimizees.
Speciﬁcally, L2Os with Hessian Trace enable fast training loss decay and much lower ﬁnal loss on all four unseen meta-testing optimizees, demonstrating the improved optimizer generalization ability.
Furthermore, all unseen optimizees trained by Hessian regularized L2Os enjoy substantial testing accuracy which boosts up to 31% , showing the enhanced optimizee generalization ability.
Such impressive performance gains effective evidence of our proposal, which again suggests that Hessian regularization potentially leads to well-generalizable minimas in wide valleys in loss landscape.
(cid:183)
Adopting vanilla L2O-DM-CL to train meta-testing optimizees (e.g., Conv-MNIST and Conv-CIFAR ) suffers from instability as shown in Figure 3, and it can be signiﬁcantly mitigated by introducing our ﬂatness-aware regularization.
Conv-Large-MNIST is an exception, where the L2O-DM-CL fails to train this optimizee and ends up with random guessed accuracy, i.e., 10% .
Although plugging Hessian Trace into L2O-DM-CL greatly improves its test accuracy from 10% to 95%+ , it still undergoes a unsatisfactory training loss.
Potential reasons may lie in the rough model architecture and limited input features of L2O-DM-CL, coincided with the ﬁndings in Chen et al.
(2020a).
We will further investigate this interesting phenomenon in the future.
(cid:184)
For advanced
L2O-Scale, Hessian Trace regularization facilitates it to converge a signiﬁcantly lower minima and obtain considerable accuracy improvements.
It enlarges the advantages of L2O methods compared to analytical optimizers, SGD and Adam, unleashing the power of parameterized optimizers.
5. I MPROVE G ENERALIZATION WITH E NTROPY R EGULARIZATION
We investigate the generalization ability improvements from the Entropy regularization.
Generally, it boosts both optimizee and optimizer generalizations of L2O in most cases, as shown in Figure 3.
Hessian v.s. Entropy Regularization.
We compare our two kinds of ﬂatness-aware regularizers from both computational cost and performance beneﬁts perspectives.
(cid:182)
In order to calculate the local entropy’s gradient in eq.
(8), it involves gradients from multiple unroll steps for the estima- [Figure] Trace . tion (Chaudhari et al., 2017), leading to extra memory and computing outlays.
Compared to Hessian augmented L2O, it costs ∼ 2 .
6 x memory and ∼ 3 x running time for L2O-DM-CL experiments 2 .
(cid:183)
As for generalization gains, Entropy regularizer performs slightly better on Conv-MNIST and
Conv-CIFAR , while behaves marginally worse on MLP-ReLU and Conv-Large-MNIST compared to Hessian regularizer.
We would like to draw reader’s attention to Conv-Large-MNIST , in which Entropy regularized L2O-DM-CL is capable of decaying the training loss and ﬁnding a much lower minima than Adam.
Note that on this optimizee, both L2O-DM-CL and its Hessian variant can not decrease the training loss.
The possible reason is that multi-layer convolutional neural network without BN cannot be stably trained on MNIST. However, our L2O-DM-CL+Entropy is more stable in training and improves testing accuracy compared with L2O-DM-CL. This indicates that L2O-DM-CL + Entropy may also produce a more trainable loss surface for optimizees.
Based on the above experiments as well as the experiment on ResNet20 in Appendix A.1, we observe that L2O+Entropy is preferred when we adopt L2O to train large neural networks, where L2O+Entropy yields better optimizers and optimizee generalization abilities. On the other hand,
L2O+Hessain optimizer requires less time per iteration to train and achieves lower training loss as well as higher test accuracy than L2O+Entropy in small networks, e.g. MLP.
The possible reason is that Entropy takes account of the landscape over a large range of loss to measure the ﬂatness, and can hence capture complex landscape information in large neural networks. On the other hand, Hessian regularizer captures the ﬂatness information only for the individual point, but in a more accurate manner, and thus is more suitable to smaller neural networks with a relative simple landscape.
5. A BLATION AND V ISUALIZATION
In this section, we carefully examine the effect of Hessian regularization’s different approximation variants, including Hessian EV , Hessian Trace , and Jacobian Trace .
Results are presented in Figure 4.
We ﬁnd that Hessian Trace regularizer achieves the most stable and substantial performance boosts across all optimizees.
Jacobian Trace performs the worst which is within expectation since it provides the roughest estimation of Hessian.
C ONCLUSION
In this paper, we propose several ﬂatness-aware regularizers to improve both optimizer and optimizee generalization abilities of current state-of-the-art
L2O approaches.
Such regularizers are capable of shaping the local geometry of optimizee’s loss surface, and leading to well-generalizable minimas in wide valleys which have been proved theoretically.
Our empirical results validate the effectiveness of our proposal, taking a further step for L2O’s practically usage in real-world scenarios.
R EPRODUCIBILITY C HECKLIST
To ensure reproducibility, we use the Machine Learning Reproducibility Checklist v2.0, Apr.2020 (Pineau et al., 2021).
An earlier version of this checklist (v1.2) was used for NeurIPS 2019 (Pineau et al., 2021).
• For all models and algorithms presented, – A clear description of the mathematical settings, algorithm, and/or model.
We clearly describe all of the settings, formulations, and algorithms in Section 3.
– A clear explanation of any assumptions.
All assumptions are stated in Section 4.
and details are clearly explained in Appendix B.1.
– An analysis of the complexity (time, space, sample size) of any algorithm.
We do not make the analysis.
• For any theoretical claim , – A clear statement of the claim.
A clear statement of theoretical claims are made in Section 4.2 and Section 4.3.
– A complete proof of the claim.
The complete proofs of all claims are available in Appendix B and Appendix C. • For all datasets used, check if you include: – The relevant statistics, such as number of examples.
We use widely adopted datasets MNIST and CIFAR-10 in Section 5.
The related statistics can be seen at http://yann.lecun.com/exdb/mnist/ and https://www.cs.
toronto.edu/˜kriz/cifar.html .
– The details of train/validation/test splits.
We give this information in our repository in the supplementary material.
– An explanation of any data that were excluded, and all pre-processing step.
We did not exclude any data or perform any pre-processing.
– A link to a downloadable version of the dataset or simulation environment.
Our repository contains all instructions to download and run experiments on the datasets.
– For new data collected,a complete description of the data collection process, such as instructions to annotators and methods for quality control.
We do not collect or release new datasets.
• For all shared code related to this work, check if you include: – Speciﬁcation of dependencies.
We give installation instructions in the README of our repository.
– Training code.
The training code is available in our repository.
– Evaluation code.
The evaluation code is available in our repository.
– (Pre-)trained model(s).
We do not release any pre-trained models.
– README ﬁle includes table of results accompanied by precise command to run to produce those results.
We include a README with detailed instructions to reproduce all of our experimental results.
• For all reported experimental results , check if you include: – The range of hyper-parameters considered, method to select the best hyperparameter conﬁguration, and speciﬁcation of all hyper-parameters used to generate results.
We provide all details of the hyper-parameter tuning in Section 5.
– The exact number of training and evaluation runs.
The details about training and evaluation can be seen in Section 5.
– A clear deﬁnition of the speciﬁc measure or statistics used to report results.
We use the classiﬁcation accuracy on test-set and the loss on the train-set.
– A description of results with central tendency (e.g. mean) & variation (e.g. error bars).
We do not report the mean and standard deviation for experiments.
– The average runtime for each result, or estimated energy cost.
We do not report the running time or energy cost.
– A description of the computing infrastructure used.
All detailed descriptions are presented in Section 5.
Supplementary Materials A A DDITIONAL E XPERIMENTAL R ESULTS A. R ES N ET 20 E XPERIMENTS
In this section, we evaluate the performance of our trained optimizers on larger neural networks
ResNet-20 on CIFAR-10 dataset.
The training loss and testing accuracy are plotted in Figure 5.
We can see that the Entropy regularizer is able to outperform other methods in both training loss and testing accuracy, demonstrating its generalization ability on large unseen models.
Further note that although the Hessian regularizer may not be preferred in large neural networks, it does perform better than the Entropy regularizer in small networks as we have shown in Figure
[Figure]
Training Iterations
[Figure]
i n i n g L o ss
ResNet-20-CIFAR
[Figure]
Training Iterations
[Figure]
A. W ALL C LOCK C OMPARISON BETWEEN DIFFERENT ALGORITHMS
We further conduct an optimizee training time comparison between our methods and analytical optimizers, L2O-DM-CL, and Entropy-SGD (Chaudhari et al., 2017) in Table 1.
Note that L2ODM-CL+Hessian and L2O-DM-CL+Entropy share the same time to train optimizee as L2O-DMCL.
From Table 1, we can see that trained L2O-DM-CL requires only ∼ 1 .
5 x time than analytical optimizers in terms of inference time, which is thus time efﬁcient for practical usage.
However,
Entropy-SGD requires ∼ 21 x time than analytical optimizers to train optimizees.
Such cost is because Entropy-SGD requires multiple Langevin dynamic steps per iteration to estimate the local entropy.
[Table]
A. A CCURACY C OMPARISON BETWEEN DIFFERENT ALGORITHMS
We also compare the testing accuracy (%) of our proposed methods with Entropy-SGD (Chaudhari et al., 2017) and SGD with Hessian regularization.
The Conv-MNIST results shown in Table 2 are evaluated on L2O-DM-CL and the Conv-CIFAR results shown in Table 3 are evaluated on L2OScale.
We adopt the same experimental setting as in Section 5 for training except that the running epochs are limited to 100 to investigate whether the performance of trained optimizers would persist in long term.
[Table]
[Table]
From these comparisons, we can see that our proposed optimizers (L2O+Hessian, L2O+Entropy) achieve the best performance compared with regularized analytical optimizers.
Speciﬁcally, in
Conv-CIFAR setting as shown in Table 3, our algorithm L2O+Hessian outperforms SGD+Hessian and Entropy-SGD.
In Conv-MNIST setting as shown in Table 2, the performances of top three algorithms, i.e. L2O+Entropy, Entropy-SGD and L2O+Hessian, are similar and much better than the performances of L2O and SGD+Hessian.
Among the top three algorithms, the iteration running time for Entropy-SGD is 0.958 secs while L2O+Hessian and L2O+Entropy only take 0.067 secs as shown in Table 1.
Such wall clock comparison shows that L2O+Hessian and L2O+Entropy are more time efﬁcient than Entropy-SGD while achieving the high accuracy, which are preferred for practical usage.
B P ROOF OF T HEOREMB. R ESTATEMENT OF A SSUMPTIONS
Assumption 5 (Restatement of Assumption 2) .
Lipschitz properties are assumed on functions
L (1) ( θ ) and L (2) ( θ ) .
a) L (1) ( θ ) function is M -Lipschitz, i.e., for any θ 1 and θ 2 , (cid:107) L (1) ( θ 1 ) − L (1) ( θ 2 ) (cid:107) ≤ M (cid:107) θ 1 − θ 2 (cid:107) .
b) ∇ θ L (1) ( θ ) and ∇ θ L (2) ( θ ) are L -Lipschitz,
i
e., for any θ 1 and θ 2 , (cid:107)∇ θ L ( i ) ( θ 1 ) − ∇ θ L ( i ) ( θ 2 ) (cid:107) ≤ L (cid:107) θ 1 − θ 2 (cid:107) ( i = 1 , 2) .
c) ∇ 2 θ L (1) ( θ ) and ∇ 2 θ L (2) ( θ ) are ρ -Lipschitz, i.e., for any θ 1 and θ 2 , (cid:107)∇ 2 θ L ( i ) ( θ 1 ) − ∇ 2 θ L ( i ) ( θ 2 ) (cid:107) ≤ ρ (cid:107) θ 1 − θ 2 (cid:107) ( i = 1 , 2) .
This assumption also holds for stochastic ∇ 2 θ L (1) N ( θ ) and ∇ 2 θ L (2) M ( θ ) .
Assumption 6 (Restatement of Assumption 3) .
Similarly as in Mei et al.
(2018), we assume the loss gradient ∇ F (1) ( θ ; ξ ) is τ 2 -sub-Gaussian, i.e., for any (cid:37) ∈ R p , and θ ∈ D r (0) where D r (0) ≡ { θ ∈ R p , (cid:107) θ (cid:107) 2 ≤ r } ,
[Equation]
Meanwhile, we assume the loss Hessian is τ 2 -sub-exponential, i.e., for any (cid:37) ∈ D 1 (0) , and θ ∈
D r (0) , ξ
[Equation]
∇
[Equation]
and there exists a constant c h such that L ≤ τ 2 p c h , ρ ≤ τ 3 p c h .
Assumption 7 (Restatement of Assumption 4) .
We assume functions L (1) ( θ ) is ( (cid:15), η ) -strongly
Morse in D r (0) , i.e., if (cid:107)∇ L (1) ( θ ) (cid:107) 2 > (cid:15) for (cid:107) θ (cid:107) 2 = r and, for any θ ∈ R p , (cid:107) θ (cid:107) 2 < r , the following holds:
[Equation]
where λ i ( ∇ 2 L (1) ( θ )) denotes the i -th eigenvalue of ∇ 2 L (1) ( θ ) .
We further make the assumption that the local basins D d ( θ ( i ) T ( φ ∗ ))
( i = 1 , 2) of convergence points θ ( i ) T ( φ ∗ )( i = 1 , 2) are in D r (0) .
B. P ROOF OF S UPPORTING L EMMAS
Lemma 1 (Restatement of Theorem 1(b) in Mei et al.
We assume θ ∗ corresponding to θ ∗ N in local basin.
Based on Assumptions 2 and 3, there exists a universal constant C 0 , and we let
C = C 0 max { c h , log( rτ/δ ) , 1 } .
If N ≥ Cp log p , then we have
[Equation]
with probability at least 1 − δ .
Lemma 2 (Restatement of Theorem 2 in Mei et al.
Based on Assumptions 2, 3 and 4, we set C as in Lemma 1, assume that θ ∗ is corresponding to θ ∗ N , and let N ≥
Cp log N/η 2 ∗ where η 2 ∗ =
min { ( (cid:15) 2 /τ 2 ) , ( η 2 /τ 4 ) , ( η 4 / ( L 2 τ 2 )) } .
Then, for each corresponding θ ∗ N and θ ∗ , we have [Equation] N
[Equation]
with probability at least 1 − δ .
Lemma 3.
Suppose Assumptions 1 and 2 hold.
Then, we have
[Equation]
where T (cid:48) is the minimum that after T (cid:48) gradient descent updates, the updated optimizee parameter θ (1) T (cid:48) ( GD ) locates into the local basin of θ (1) T ( φ ∗ ) .
Proof.
Since the local basin is µ -strongly convex and θ ∗ (1) N is the optimal point of smooth function
L (1) N ( θ ) in local basin.
Then, we have
[Equation]
Furthermore
[Equation]
we rearrange the terms and obtain
[Equation]
where ( i ) follows because φ ∗ =
arg min
φ L (1) N ( θ (1) T ( φ )) and θ (1) T ( GD ) locates in the local basin of θ ∗ (1) N , ( ii ) follows from Assumptioin 2 and the fact that θ ∗ (1) N
= arg min θ L (1) N ( θ ) , and ( iii ) follows if we set step size of GD as 2 µ + L .
Lemma 4.
Based on Assumptions 1, 2, 3 and 4, we let N ≥ max
{ Cp log p, 4 Cp log N/η 2 ∗ } where
C = C 0 max { c h , 1 , log ( rτδ ) } , η 2 ∗ =
min { (cid:15) 2 τ 2 , η 2 τ 4 , η 4 ρ 2 τ 2 } , C 0 is an universal constant.
Then, with probability at least 1 − 2 δ we have
[Equation]
where ∆ ∗ H =
(cid:107)∇ 2 θ L (2) ( θ [Equation]
− ∇ 2 θ L (1) ( θ ∗ (1) ) (cid:107) and T (cid:48) is deﬁned in Lemma 3.
Proof.
Firstly, we bound (cid:107)∇ 2 θ L (2) ( θ ∗ (2) )
(cid:107) as following:
[Equation]
where θ ∗ (1) is corresponding to θ ∗ (1) N in the same local basin of [Equation]
Based on the constrained problem formulation in eq.
(9), the optimal optimizer parameter φ ∗ is equivalent to the following:
[Equation]
Thus, we obtain (cid:107)∇ 2 θ L (1) N ( θ (1) T ( φ ∗ )) (cid:107) ≤ B Hessian ( λ ) .
Furthermore, if we let N ≥ Cp log p where
C = C 0 max { c h , 1 , log ( rτδ ) } and C 0 is an universal constant, based on Lemmas 1, 2 and 3, and
Assumptions 2, we have
[Equation]
with probability at least 1 − δ .
Furthermore, if we assume N ≥ max { 4 Cp log N/η 2 ∗ , Cp log p } where η 2 ∗ =
min { (cid:15) 2 τ 2 , η 2 τ 4 , η 4 ρ 2 τ 2 } , based on Lemma 2, we have
[Equation]
with probability at least 1 − 2 δ .
Lemma 5.
Based on Assumptions 1, 2, 3, and 4, we let N ≥ 4 Cp log N/η 2 ∗ where C and η 2 ∗ are deﬁned in Lemma 2.
Then, with probability at least 1 − δ , we have
[Equation]
where ∆ ∗ θ = (cid:107) θ ∗ (1) − θ ∗ (2) (cid:107) , ∆ ∗ T = (cid:107) θ (2) T ( φ ∗ ) − θ (1) T ( φ ∗ ) (cid:107) and T (cid:48) is deﬁned in Lemma 3.
Proof.
Based on triangle inequality, we obtain
[Equation]
[Equation]
[Table]
[Equation] [Table] [Equation] where ( i ) follows from deﬁnition of ∆ ∗ T , ( ii ) follows from Lemma 3 and deﬁnition of ∆ ∗ θ .
Based on Lemma 2, if we let N ≥ 4 Cp log N/η 2 ∗ .
Then, with probability at least 1 − δ , we have
[Equation]
where w = L − µ L + µ .
B. P ROOF OF T HEOREM
Generalization loss is deﬁned as below:
[Equation]
where ( i ) follows from Taylor expansion and θ (cid:48) follows from the conditions that (cid:107) θ (cid:48) − θ (2) T ( φ ∗ ) (cid:107) ≤ (cid:107) θ ∗ (2) − θ (2) T ( φ ∗ ) (cid:107) and (cid:107) θ (cid:48) − θ ∗ (2) (cid:107) ≤ (cid:107) θ ∗ (2) − θ (2) T ( φ ∗ ) (cid:107) , ( ii ) follows because ∇ θ L (2) ( θ ∗ (2) )
= 0 , and ( iii ) follows from Assumption 2 and the fact that (cid:107) θ (cid:48) − θ ∗ (2) (cid:107) ≤ (cid:107) θ ∗ (2) − θ (2) T ( φ ∗ ) (cid:107) .
Based on Lemmas 4 and 5, if we let N ≥ max { 4 Cp log N/η 2 ∗ , Cp log p } where C and η 2 ∗ are deﬁned in Lemma 4.
Then, with probability at least 1 − 2 δ , we have
[Equation]
N
[Equation]
where w = L − µ L + µ , ∆ ∗ H
=
(cid:107)∇ 2 θ L (2) ( θ ∗ (2) )
− ∇ 2 θ L (1) ( θ ∗ (1) ) (cid:107) , ∆ ∗ θ = (cid:107) θ ∗ (1) − θ ∗ (2) (cid:107) , ∆ ∗ T
= (cid:107) θ (2) T ( φ ∗ ) − θ (1) T ( φ ∗ ) (cid:107) .
Furthermore, We let A 1 = B Hessian ( λ ) + ρ ∆ ∗ T + ρ ∆ ∗ θ + ∆ ∗ H + O ( w T − T (cid:48) )
+ O ( (cid:113) C log N N ) , A 2 = ∆ ∗ T + ∆ ∗ θ + O ( w T − T (cid:48) )
+ O ( (cid:113) C log N N ) .
Then, we have
L (2) ( θ (2) T ( φ ∗ ))
− L (2) ( θ ∗ (2) ) ≤ 1 2 (cid:107) θ (2) T ( φ ∗ ) − θ ∗ (2) (cid:107) 2 ( ρ (cid:107) θ (2) T ( φ ∗ ) − θ ∗ (2) (cid:107)
+ (cid:107)∇ 2 θ L (2) ( θ ∗ (2) ) (cid:107) ) ≤2 A 22 A 1 , with probability at least 1 − 2 δ .
Then, the proof is complete.
C P ROOF OF T HEOREMC. P ROOF OF S UPPORTING L EMMA
Lemma 6.
Based on Assumptions 1 and 2, in terms of entropy regularizer constraint B Entropy ( λ ) , we have
[Equation]
where m is a constant, C ( γ, p, m )
= log (cid:82) θ (cid:48) :
(cid:107) θ (cid:48) − θ (cid:107) >m exp (cid:0) − γ 2 (cid:107) θ − θ (cid:48) (cid:107) 2 (cid:1)
d θ (cid:48) and θ ∈ R p .
Proof.
We ﬁrstly split the integral area θ (cid:48)
∈
R p into two parts: { θ (cid:48) :
(cid:107) θ (cid:48) − θ (cid:107) ≤ m } and { θ (cid:48) :
(cid:107) θ (cid:48) − θ (cid:107) > m } .
Based on the deﬁnition of G (1) N ( θ ; γ ) , we have
[Equation]
where ( i ) follows from the fact that L (1) N ( θ (cid:48) ) is non-negative, ( ii ) follows from the fact that θ ∈ R p and the deﬁniton of C ( γ, p, m ) , ( iii ) follows from Taylor expansion.
Note that θ (cid:48)(cid:48) satisﬁes [Equation]
and (cid:107) θ (cid:48)(cid:48) − θ (cid:48)
(cid:107) ≤ (cid:107) θ − θ (cid:48) (cid:107) .
Based on Assumption 2, − ( θ (cid:48) − θ ) T ∇ L (1) N ( θ ) ≤ m (cid:107)∇ L (1) N ( θ ) (cid:107) ≤ mM .
Then, we obtain
G (1) N ( θ ; γ )
[Table]
[Equation]
where ( i ) follows from Assumption 2 and the fact that (cid:107) θ (cid:48)(cid:48) − θ
[Equation]
≤
[Equation]
θ − θ
[Equation]
(cid:107) and ( ii ) follows because exp [Equation]
≥ 0 .
Based on Assumption [Equation] we have the fact that ( ∇ 2 L (1) N ( θ (1) T ( φ )) + γI ) is a symmetric and positivedeﬁnite matrix.
Hence, we obtain
[Equation]
We rearrange the terms and get
[Equation]
[Equation] Since − G (1) N ( θ (1) T ( φ ); γ ) ≤ B Entropy ( λ ) .
Then, we obtain
[Equation]
C. P ROOF OF T HEOREM
Based on Lemma 6, we have
[Equation]
Since ∇ 2 L (1) N ( θ (1) T ( φ )) + γI is positive deﬁnite and λ i ( ∇ 2 L (1) N ( θ (1) T ( φ )) + γI ) ≥ γ + µ for any i = 1 , . .
, p .
Then, based on the deﬁnition of Matrix norm (cid:107)∇ 2 L (1) N ( θ (1) T ( φ )) + γI (cid:107) =
λ max ( ∇ 2 L (1) N ( θ (1) T ( φ )) + γI ) , we have
[Equation]
φ
[Equation]
+ γI (cid:107) .
Note that we use λ
i ( H ) to denote the i -th eigenvalue of matrix H .
Then,
[Equation]
Then, we can obtain
[Equation]
Hence, we can get a new function D ( x ) that
[Equation]
, where D ( x ) =
L (1) N ( θ (1) T ( φ ))
+ ( p − 1) log( γ + µ ) − mM − p 2 log(2 π ) − 12 ρm 3 − C ( γ, p, m ) + log( x + γ ) .
Then, the proof is complete.