Year: 2021

Abstract:
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.

text: 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 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. 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHODOLOGY
text: In this section, we provide basic notations and the detailed formulations about our flatness-aware regularizers, i.e., Hessian and Entropy regularizers.

heading: 3.1 PRELIMINARY
text: 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.

heading: 3.2 HESSIAN REGULARIZER
text: Motivated by the idea that the optimizee in a flat area of the training objective has superior generalization capability, we propose to incorporate flatness-aware regularizers into L2O meta-training, in order to learn optimizers that favors to land the optimizee into a flat region. We introduce two such regularizers in this and next subsection.
The first regularizer we introduce is based on the spectral norm of the Hessian, smaller values of which corresponds to a flatter landscape. Thus, the new L2O meta-training objective is given by:
min φ {L(1)N (θ (1) T (φ)) + λ‖∇ 2 θL (1) N (θ (1) T (φ))‖} where θ (1) t+1(φ) = θ (1) t (φ) +m(z (1) t ;φ), (4)
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 zt does not include such regularization. We let φ∗ be the optimal optimizer parameter, which can be written as
φ∗ = arg min φ {L(1)N (θ (1) T (φ)) + λ‖∇ 2 θL (1) N (θ (1) T (φ))‖}. (5)
Due to the computational intractability of directly penalizing∇2θL (1) N (θ (1) T (φ)), we investigate three approximation variants in the implementation. ¬ Hessian EV: the eigenvalue of largest module of Hessian matrix, computed by power iteration (Yao et al., 2020);  Hessian Trace: the trace of Hessian matrix, calculated via Hutchinson method (Yao et al., 2020); ® Jacobian Trace: the trace of Hessian’s Jacobian approximation ∇θL(1)N (θ (1) T (φ)) >∇θ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.

heading: 3.3 ENTROPY REGULARIZER
text: The second flatness-aware regularizer we incorporate to L2O is based on the local entropy function G
(1) N (θ; γ) = log ∫ θ′ exp ( −L(1)N (θ′)− γ 2 ‖θ − θ ′‖2 ) dθ′ proposed in Chaudhari et al. (2017). Due to the exponential decay with respect to ‖θ− θ′‖2, the integral mainly captures the value of the loss function L(1)N (θ
′) over the neighborhood of θ. Thus, the value of G(1)N (θ; γ) measures the flatness of the local area around θ. Thus, the L2O meta training objective with Entropy regularizer is given by:
min φ {L(1)N (θ (1) T (φ))− λG (1) N (θ (1) T (φ); γ)} where θ (1) t+1(φ) = θ (1) t (φ) +m(z (1) t ;φ). (6)
We let φ∗ be the optimal optimizer parameter, which can be written as
φ∗ = arg min φ L (1) N (θ (1) T (φ))− λG (1) N (θ (1) T (φ); γ). (7)
In order to implement the gradient descent algorithm for meta-training, the gradient −∇φGN (θ(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
−∇θG(1)N (θ; γ) = γ(θ − E[θ ′; ξ]), (8)
where ξ ∈ {ξi, i = (1, . . . , N)} are training samples and the distribution of θ′ is given by P (θ′; θ, γ) ∝ exp [ −L(1)N (θ′)− γ 2 ‖θ − θ ′‖2 ] .

heading: 4 THEORETICAL ANALYSIS
text: In this section, we first introduce several assumptions, and then present the generalization analysis of L2O with Hessian and Entropy regularizers.

heading: 4.1 ASSUMPTIONS
text: We first define the local basin of θ with the radius d as Dd(θ) = {θ′ : ‖θ− θ′‖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 Dd(θ(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 Dr(0) that the meta-training functions L(1)(θ) is ( , η)-strongly Morse in Dr(0) and the local basins Dd(θ(i)T (φ ∗)) for i = 1, 2 of convergence points θ(i)T (φ ∗)(i = 1, 2) are in Dr(0). See Appendix B.1 for more details.

heading: 4.2 GENERALIZATION ANALYSIS OF HESSIAN-REGULARIZED L2O
text: 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 first note that by optimization theory, the regularized optimization problem eq. (5) is equivalent to the following constrained optimization
min φ {L(1)N (θ (1) T (φ))} where θ (1) t+1 = θ (1) t +m(z (1) t ;φ)
subject to ‖∇2θL (1) N (θ (1) T (φ))‖ ≤ BHessian(λ), (9)
where BHessian(λ) 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 defined above. Theorem 1 (Generalization Error of Hessian-Regularized L2O). Suppose Assumptions 1, 2, 3, 4 hold. We let N ≥ max{4Cp logN/η2∗, Cp log p} where C = C0 max{ch, 1, log( rτδ )}, η 2 ∗ = min{ 2
τ2 , η2 τ4 , η4
ρ2τ2 } and C0 is an universal constant. Then, with probability at least 1 − 2δ we have
L(2)(θ (2) T (φ ∗))− L(2)(θ∗(2)) ≤ 1 2 A22A1, (10)
where A1 = BHessian(λ) + ρ∆∗T + ρ∆ ∗ θ + ∆ ∗ H + O(wT−T ′ ) + O( √ C logN N ), A2 = ∆ ∗ T + ∆ ∗ θ +
O(wT−T ′) + O( √
C logN N ), with w = L−µ L+µ , ∆ ∗ H = ‖∇2θL(2)(θ∗(2)) − ∇2θL(1)(θ∗(1))‖, ∆∗θ =
‖θ∗(1) − θ∗(2)‖, ∆∗T = ‖θ (2) T (φ ∗) − θ(1)T (φ∗)‖, and T ′ is the minimum gradient descent iterations for θ(1)T ′ (GD) to enter into the local basin of θ (1) T (φ ∗).
In Theorem 1 the generalization error is bounded by the termsA1 andA2, where the Hessian regularizer affects the generalization error through the constraintBHessian(λ) inA1. Clearly, by choosing the regularization hyperparameter λ, we control the value of BHessian(λ), which then makes an impact on the generalization. Specifically, larger λ corresponds to smaller BHessian(λ) and hence smaller generalization error. This also explains that flatter landscape (i.e., smaller BHessian(λ) 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(wT−T ′) captures the exponential decay rate of the optimizee’s iteration due to the strong convexity, and vanishes for large T ; and (c) O( √ C logN N ) arises due to the difference between the empirical and population loss functions, and vanishes as the sample size N gets large.

heading: 4.3 GENERALIZATION ANALYSIS OF ENTROPY-REGULARIZED L2O
text: 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:
min φ L (1) N (θ (1) T (φ)) where θ (1) t+1 = θ (1) t +m(z (1) t ;φ)
subject to −G(1)N (θ (1) T (φ); γ) ≤ BEntropy(λ), (11)
where BEntropy(λ) 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 satisfies the constraint.
In the following, we first establish an important connection between the Entropy constraint bound BEntropy(λ) and the Hessian ‖∇2L(1)N (θ (1) T (φ))‖ of the optimizee. Theorem 2 (Connection between Hessian and Entropy Regularizer). Suppose Assumptions 1 and 2 hold. We define C(γ, p,m) = log ∫ θ′:‖θ′−θ‖>m exp(− γ 2 ‖θ − θ
′‖2)dθ′ where m is a constant and θ ∈ Rp. Then, we have
‖∇2L(1)N (θ (1) T (φ))‖ ≤ D −1(BEntropy(λ))
where D(x) = L(1)N (θ (1) T (φ)) + (p − 1) log(γ + µ) −mM − p 2 log(2π) − 1 2ρ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 BEntropy(λ) implies a corresponding Hessian constraint D−1(BEntropy(λ)). 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 BHessian(λ) in A1 being replaced by D−1(BEntropy(λ)).
Theorem 2 and Corollary 1 establish that the bound D−1(BEntropy(λ)) 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, BEntropy(λ) as well as D−1(BEntropy(λ)) and Hessian can be controlled to be sufficiently small. In this way, the optimizee will be landed into a flat 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 flat 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.

heading: 5 EXPERIMENT
text: In our experiments, we consider two advanced L2O algorithms, i.e., L2O-DM-CL1 (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× 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
1It is an enhanced version of the earliest L2O-DM introduced by DeepMind Andrychowicz et al. (2016b). We choose it as a much stronger baseline instead of the vanilla L2O-DM with “poor-generalization baseline” (Wichrowska et al., 2017; Lv et al., 2017).
learned optimizer. Specifically, ¬ MLP-ReLU: a single layer MLP with 20 neurons and the ReLU activation function on MNIST.  Conv-MNIST: a CNN has one fully-connected layer, and two convolutional layers with 16 3 × 3 and 32 5 × 5 kernels on MNIST. ® Conv-Large-MNIST: a large CNN has one fully-connected layer, and four convolutional layers with two 32 3 × 3 and two 32 5 × 5 kernels on MNIST. ¯ Conv-CIFAR: a CNN has one fully-connected layer, and two convolutional layers with 16 3×3 and 32 5×5 kernels on CIFAR-10 (Krizhevsky & Hinton, 2009). Optimizees ¬, , and ® are for evaluating the generalization of L2O across network architectures. Then, ¯ 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 coefficients {λ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.

heading: 5.1 IMPROVE GENERALIZATION WITH HESSIAN REGULARIZATION
text: 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: ¶ Hessian Trace regularizer consistently enhances the generalization abilities of learned L2Os and trained optimizees. Specifically, L2Os with Hessian Trace enable fast training loss decay and much lower final 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. · 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 significantly mitigated by introducing our flatness-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 findings in Chen et al. (2020a). We will further investigate this interesting phenomenon in the future. ¸ For advanced L2O-Scale, Hessian Trace regularization facilitates it to converge a significantly 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.

heading: 5.2 IMPROVE GENERALIZATION WITH ENTROPY REGULARIZATION
text: 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 flatness-aware regularizers from both computational cost and performance benefits perspectives. ¶ In order to calculate the local entropy’s gradient in eq. (8), it involves gradients from multiple unroll steps for the estima-
tion (Chaudhari et al., 2017), leading to extra memory and computing outlays. Compared to Hessian augmented L2O, it costs ∼ 2.6x memory and ∼ 3x running time for L2O-DM-CL experiments2. · 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 finding 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 flatness, and can hence capture complex landscape information in large neural networks. On the other hand, Hessian regularizer captures the flatness 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.

heading: 5.3 ABLATION AND VISUALIZATION
text: 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 find 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.

heading: 6 CONCLUSION
text: In this paper, we propose several flatness-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.
2We conduct entropy-related experiments on light-weight L2O-DM-CL rather heavy L2O-Scale models, since RTX TITAN with 24G memory is the largest GPU we can access and afford.

heading: 7 REPRODUCIBILITY CHECKLIST
text: To ensure reproducibility, we use the Machine Learning Reproducibility Checklist v2.0, Apr. 7 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.1
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: – Specification 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 file 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 hyper-
parameter configuration, and specification 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 definition of the specific measure or statistics used to report results. We use the classification 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.

heading: Supplementary Materials

heading: A ADDITIONAL EXPERIMENTAL RESULTS

heading: A.1 RESNET20 EXPERIMENTS
text: 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 3.

heading: A.2 WALL CLOCK COMPARISON BETWEEN DIFFERENT ALGORITHMS
text: 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.5x time than analytical optimizers in terms of inference time, which is thus time efficient for practical usage. However, Entropy-SGD requires ∼ 21x 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.

heading: A.3 ACCURACY COMPARISON BETWEEN DIFFERENT ALGORITHMS
text: 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.
From these comparisons, we can see that our proposed optimizers (L2O+Hessian, L2O+Entropy) achieve the best performance compared with regularized analytical optimizers. Specifically, 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 efficient than Entropy-SGD while achieving the high accuracy, which are preferred for practical usage.

heading: B PROOF OF THEOREM 1

heading: B.1 RESTATEMENT OF ASSUMPTIONS
text: 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, ‖L(1)(θ1)−L(1)(θ2)‖ ≤M‖θ1− θ2‖.
b) ∇θL(1)(θ) and ∇θL(2)(θ) are L-Lipschitz, i.e., for any θ1 and θ2, ‖∇θL(i)(θ1) − ∇θL(i)(θ2)‖ ≤ L‖θ1 − θ2‖(i = 1, 2).
c) ∇2θL(1)(θ) and ∇2θL(2)(θ) are ρ-Lipschitz, i.e., for any θ1 and θ2, ‖∇2θL(i)(θ1) − ∇2θL(i)(θ2)‖ ≤ ρ‖θ1 − θ2‖(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 % ∈ Rp, and θ ∈ Dr(0) where Dr(0) ≡ {θ ∈ Rp, ‖θ‖2 ≤ r},
E{exp(〈%,∇F (1)(θ; ξ)− E[∇F (1)(θ; ξ)]〉)} ≤ exp ( τ2‖%‖2
2
) .
Meanwhile, we assume the loss Hessian is τ2-sub-exponential, i.e., for any % ∈ D1(0), and θ ∈ Dr(0),
ξ%,θ ≡ 〈%,∇2F (1)(θ; ξ)%〉, E { exp ( 1
τ2 |ξ%,θ − Eξ%,θ|
)} ≤ 2,
and there exists a constant ch such that L ≤ τ2pch , ρ ≤ τ3pch . Assumption 7 (Restatement of Assumption 4). We assume functions L(1)(θ) is ( , η)-strongly Morse in Dr(0), i.e., if ‖∇L(1)(θ)‖2 > for ‖θ‖2 = r and, for any θ ∈ Rp, ‖θ‖2 < r, the following holds:
‖∇L(θ)‖2 ≤ ⇒ min i∈[p] |λi(∇2L(1)(θ))| ≥ η,
where λi(∇2L(1)(θ)) denotes the i-th eigenvalue of ∇2L(1)(θ). We further make the assumption that the local basins Dd(θ(i)T (φ ∗))(i = 1, 2) of convergence points θ(i)T (φ ∗)(i = 1, 2) are in Dr(0).

heading: B.2 PROOF OF SUPPORTING LEMMAS
text: Lemma 1 (Restatement of Theorem 1(b) in Mei et al. (2018)). We assume θ∗ corresponding to θ∗N in local basin. Based on Assumptions 2 and 3, there exists a universal constant C0, and we let C = C0 max{ch, log(rτ/δ), 1}. If N ≥ Cp log p, then we have
sup θ∈Dp(r)
‖∇2LN (θ)−∇2L(θ)‖ ≤ τ2 √ Cp logN
N ,
with probability at least 1− δ. Lemma 2 (Restatement of Theorem 2 in Mei et al. (2018)). Based on Assumptions 2, 3 and 4, we set C as in Lemma 1, assume that θ∗ is corresponding to θ∗N , and let N ≥ 4Cp logN/η2∗ where η2∗ = min{( 2/τ2), (η2/τ4), (η4/(L2τ2))}. Then, for each corresponding θ∗N and θ∗, we have
‖θ∗N − θ∗‖2 ≤ 2τ
η
√ Cp logN
N ,
with probability at least 1− δ. Lemma 3. Suppose Assumptions 1 and 2 hold. Then, we have
‖θ(1)T (φ ∗)− θ∗(1)N ‖ ≤
√ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖, (12)
where T ′ is the minimum that after T ′ gradient descent updates, the updated optimizee parameter θ (1) T ′ (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
L (1) N (θ (1) T (φ ∗))− L(1)N (θ ∗(1) N ) ≥
µ 2 ‖θ(1)T (φ ∗)− θ∗(1)N ‖ 2.
Furthermore, we rearrange the terms and obtain
‖θ(1)T (φ ∗)− θ∗(1)N ‖ ≤
√ 2
µ (L
(1) N (θ (1) T (φ ∗))− L(1)N (θ ∗(1) N ))
(i) ≤ √ 2
µ (L
(1) N (θ (1) T (GD))− L (1) N (θ ∗(1) N ))
(ii)
≤
√ 2
µ
L 2 ‖θ(1)T (GD)− θ ∗(1) N ‖2
≤
√ L
µ ‖θ(1)T (GD)− θ ∗(1) N ‖
(iii)
≤
√ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖,
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, 4Cp logN/η2∗} where C = C0 max{ch, 1, log( rτδ )}, η 2 ∗ = min{ 2 τ2 , η2 τ4 , η4
ρ2τ2 }, C0 is an universal constant. Then, with probability at least 1− 2δ we have
‖∇2θL(2)(θ∗(2))‖ ≤ρ
( 2τ
η
√ Cp logN
N +
√ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖ )
+ ∆∗H + τ 2
√ Cp logN
N +BHessian(λ),
where ∆∗H = ‖∇2θL(2)(θ∗(2))−∇2θL(1)(θ∗(1))‖ and T ′ is defined in Lemma 3.
Proof. Firstly, we bound ‖∇2θL(2)(θ∗(2))‖ as following:
‖∇2θL(2)(θ∗(2))‖ ≤‖∇2θL(2)(θ∗(2))−∇2θL(1)(θ∗(1))‖+ ‖∇2θL(1)(θ∗(1))−∇2θL (1) N (θ ∗(1))‖
+ ‖∇2θL (1) N (θ ∗(1))−∇2θL (1) N (θ ∗(1) N )‖+ ‖∇ 2 θL (1) N (θ ∗(1) N )−∇ 2 θL (1) N (θ (1) T (φ ∗))‖
+ ‖∇2θL (1) N (θ (1) T (φ ∗))‖,
where θ∗(1) is corresponding to θ∗(1)N in the same local basin of θ (1) T (φ ∗).
Based on the constrained problem formulation in eq. (9), the optimal optimizer parameter φ∗ is equivalent to the following:
φ∗ = arg min φ L (1) N (θ (1) T (φ)) subject to∇ 2 θL (1) N (θ (1) T (φ)) ≤ BHessian(λ).
Thus, we obtain ‖∇2θL (1) N (θ (1) T (φ ∗))‖ ≤ BHessian(λ). Furthermore, if we let N ≥ Cp log p where C = C0 max{ch, 1, log( rτδ )} and C0 is an universal constant, based on Lemmas 1, 2 and 3, and Assumptions 2, we have
‖∇2θL(2)(θ∗(2))‖ ≤ρ ( ‖θ∗(1) − θ∗(1)N ‖+ √ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖ )
+ ∆∗H + τ 2
√ Cp logN
N +BHessian(λ),
with probability at least 1− δ.
Furthermore, if we assume N ≥ max{4Cp logN/η2∗, Cp log p} where η2∗ = min{ 2 τ2 , η2 τ4 , η4
ρ2τ2 }, based on Lemma 2, we have
‖∇2θL(2)(θ∗(2))‖ ≤ρ
( 2τ
η
√ Cp logN
N +
√ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖ )
+BHessian(λ) + ∆ ∗ H + τ 2
√ Cp logN
N ,
with probability at least 1− 2δ.
Lemma 5. Based on Assumptions 1, 2, 3, and 4, we let N ≥ 4Cp logN/η2∗ where C and η2∗ are defined in Lemma 2. Then, with probability at least 1− δ, we have
‖θ(2)T (φ ∗)− θ∗(2)‖ ≤ ∆∗T +
√ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖+ 2τ η √ Cp logN N + ∆∗θ,
where ∆∗θ = ‖θ∗(1) − θ∗(2)‖, ∆∗T = ‖θ (2) T (φ ∗)− θ(1)T (φ∗)‖ and T ′ is defined in Lemma 3.
Proof. Based on triangle inequality, we obtain
‖θ(2)T (φ ∗)− θ∗(2)‖
≤ ‖θ(2)T (φ ∗)− θ(1)T (φ ∗)‖+ ‖θ(1)T (φ ∗)− θ∗(1)N ‖+ ‖θ ∗(1) N − θ ∗(1)‖+ ‖θ∗(1) − θ∗(2)‖ (i) ≤ ∆∗T + ‖θ (1) T (φ ∗)− θ∗(1)N ‖+ ‖θ ∗(1) N − θ
∗(1)‖+ ‖θ∗(1) − θ∗(2)‖ (ii)
≤ ∆∗T +
√ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖+ ‖θ ∗(1) N − θ ∗(1)‖+ ∆∗θ,
where (i) follows from definition of ∆∗T , (ii) follows from Lemma 3 and definition of ∆ ∗ θ . Based on Lemma 2, if we let N ≥ 4Cp logN/η2∗. Then, with probability at least 1− δ, we have
‖θ(2)T (φ ∗)− θ∗(2)‖ ≤∆∗T +
√ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖+ 2τ η √ Cp logN N + ∆∗θ
=∆∗T + ∆ ∗ θ +O(wT−T ′ ) +O(
√ C logN
N ),
where w = L−µL+µ .

heading: B.3 PROOF OF THEOREM 1
text: Generalization loss is defined as below:
L(2)(θ (2) T (φ ∗))− L(2)(θ∗(2)) (i) = (θ
(2) T (φ ∗)− θ∗(2))T∇θL(2)(θ∗(2)) + 1
2 (θ
(2) T (φ ∗)− θ∗(2))T∇2θL(2)(θ′)(θ (2) T (φ ∗)− θ∗(2))
(ii) =
1 2 (θ (2) T (φ ∗)− θ∗(2))T∇2θL(2)(θ′)(θ (2) T (φ ∗)− θ∗(2))
= 1
2 (θ
(2) T (φ ∗)− θ∗(2))T (∇2θL(2)(θ′)−∇2θL(2)(θ∗(2)) +∇2θL(2)(θ∗(2)))(θ (2) T (φ ∗)− θ∗(2))
≤ 1 2 ‖θ(2)T (φ ∗)− θ∗(2)‖2(‖∇2θL(2)(θ′)−∇2θL(2)(θ∗(2))‖+ ‖∇2θL(2)(θ∗(2))‖)
(iii)
≤ 1 2 ρ‖θ(2)T (φ ∗)− θ∗(2)‖3 + 1 2 ‖∇2θL(2)(θ∗(2))‖‖θ (2) T (φ ∗)− θ∗(2)‖2
≤ 1 2 ‖θ(2)T (φ ∗)− θ∗(2)‖2(ρ‖θ(2)T (φ ∗)− θ∗(2)‖+ ‖∇2θL(2)(θ∗(2))‖),
where (i) follows from Taylor expansion and θ′ follows from the conditions that ‖θ′ − θ(2)T (φ∗)‖ ≤ ‖θ∗(2)− θ(2)T (φ∗)‖ and ‖θ′− θ∗(2)‖ ≤ ‖θ∗(2)− θ (2) T (φ
∗)‖, (ii) follows because∇θL(2)(θ∗(2)) = 0, and (iii) follows from Assumption 2 and the fact that ‖θ′ − θ∗(2)‖ ≤ ‖θ∗(2) − θ(2)T (φ∗)‖.
Based on Lemmas 4 and 5, if we let N ≥ max{4Cp logN/η2∗, Cp log p} where C and η2∗ are defined in Lemma 4. Then, with probability at least 1− 2δ, we have
ρ‖θ(2)T (φ ∗)− θ∗(2)‖+ ‖∇2θL(2)(θ∗(2))‖
≤ρ (
∆∗T +
√ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖+ 2τ η √ Cp logN N + ∆∗θ ) + τ2 √ Cp logN
N + ρ
( 2τ
η
√ Cp logN
N +
√ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖ ) + ∆∗H +B(λ)
=ρ∆∗T + 2ρ
√ L
µ ( L− µ L+ µ )T−T ′ ‖θ(1)T ′ (GD)− θ ∗(1) N ‖+ ( 4ρτ η + τ2 )√ Cp logN N
+ ρ∆∗θ + ∆ ∗ H +BHessian(λ)
=BHessian(λ) + ρ∆ ∗ T + ρ∆ ∗ θ + ∆ ∗ H +O(wT−T ′ ) +O(
√ C logN
N ),
where w = L−µL+µ , ∆ ∗ H = ‖∇2θL(2)(θ∗(2)) − ∇2θL(1)(θ∗(1))‖, ∆∗θ = ‖θ∗(1) − θ∗(2)‖, ∆∗T = ‖θ(2)T (φ∗)− θ (1) T (φ ∗)‖.
Furthermore, We let A1 = BHessian(λ) + ρ∆∗T + ρ∆ ∗ θ + ∆ ∗ H +O(wT−T ′ ) +O( √ C logN N ), A2 =
∆∗T + ∆ ∗ θ +O(wT−T ′ ) +O( √ C logN N ). Then, we have
L(2)(θ (2) T (φ ∗))− L(2)(θ∗(2)) ≤ 1 2 ‖θ(2)T (φ ∗)− θ∗(2)‖2(ρ‖θ(2)T (φ ∗)− θ∗(2)‖+ ‖∇2θL(2)(θ∗(2))‖)
≤ 1 2 A22A1,
with probability at least 1− 2δ. Then, the proof is complete.

heading: C PROOF OF THEOREM 2

heading: C.1 PROOF OF SUPPORTING LEMMA
text: Lemma 6. Based on Assumptions 1 and 2, in terms of entropy regularizer constraint BEntropy(λ), we have
BEntropy(λ) +mM + p
2 log(2π) +
1 2 ρm3 + C(γ, p,m) ≥ log(det(∇2L(1)N (θ (1) T (φ)) + γI))
+ L (1) N (θ (1) T (φ)), where m is a constant, C(γ, p,m) = log ∫ θ′:‖θ′−θ‖>m exp ( −γ2 ‖θ − θ ′‖2 ) dθ′ and θ ∈ Rp.
Proof. We firstly split the integral area θ′ ∈ Rp into two parts: {θ′ : ‖θ′ − θ‖ ≤ m} and {θ′ : ‖θ′ − θ‖ > m}. Based on the definition of G(1)N (θ; γ), we have
G (1) N (θ; γ)
= log ∫ θ′ exp ( −L(1)N (θ ′)− γ 2 ‖θ − θ′‖2 ) dθ′
= log ∫ θ′:‖θ′−θ‖≤m exp ( −L(1)N (θ ′)− γ 2 ‖θ − θ′‖2 ) dθ′
+ log ∫ θ′:‖θ′−θ‖>m exp ( −L(1)N (θ ′)− γ 2 ‖θ − θ′‖2 ) dθ′
(i) ≤ log ∫ θ′:‖θ′−θ‖≤m exp ( −L(1)N (θ ′)− γ 2 ‖θ − θ′‖2 ) dθ′
+ log ∫ θ′:‖θ′−θ‖>m exp ( −γ 2 ‖θ − θ′‖2 ) dθ′
(ii) = log ∫ θ′:‖θ′−θ‖≤m exp ( −L(1)N (θ ′)− γ 2 ‖θ − θ′‖2 ) dθ′ + C(γ, p,m)
(iii) = log ∫ θ′:‖θ′−θ‖≤m exp ( − L(1)N (θ)− (θ ′ − θ)T∇L(1)N (θ)− 1 2 (θ′ − θ)T∇2L(1)N (θ ′′)(θ′ − θ)
− γ 2 ‖θ − θ′‖2
) dθ′ + C(γ, p,m),
where (i) follows from the fact that L(1)N (θ ′) is non-negative, (ii) follows from the fact that θ ∈ Rp and the definiton of C(γ, p,m), (iii) follows from Taylor expansion. Note that θ′′ satisfies ‖θ′′ − θ‖ ≤ ‖θ − θ′‖ and ‖θ′′ − θ′‖ ≤ ‖θ − θ′‖.
Based on Assumption 2, −(θ′ − θ)T∇L(1)N (θ) ≤ m‖∇L (1) N (θ)‖ ≤ mM . Then, we obtain
G (1) N (θ; γ)
≤− L(1)N (θ) +mM + log ∫ θ′:‖θ′−θ‖≤m exp ( − 1 2 (θ′ − θ)T∇2L(1)N (θ ′′)(θ′ − θ)
− γ 2 ‖θ − θ′‖2
) dθ′ + C(γ, p,m)
=− L(1)N (θ) +mM + log ∫ θ′:‖θ′−θ‖≤m exp ( − 1 2 (θ′ − θ)T (∇2L(1)N (θ ′′) + γI)(θ′ − θ) ) dθ′
+ C(γ, p,m)
=− L(1)N (θ) +mM + C(γ, p,m)
+ log ∫ θ′:‖θ′−θ‖≤m exp ( − 1 2 (θ′ − θ)T ( ∇2L(1)N (θ ′′)−∇2L(1)N (θ) +∇ 2L (1) N (θ)
+ γI ) (θ′ − θ) ) dθ′
(i) ≤ − L(1)N (θ) +mM + C(γ, p,m) + 1
2 ρm3
+ log ∫ θ′:‖θ′−θ‖≤m exp ( − 1 2 (θ′ − θ)T ( ∇2L(1)N (θ) + γI ) (θ′ − θ) ) dθ′
(ii)
≤ − L(1)N (θ) +mM + 1
2 ρm3 + log ∫ θ′ exp ( − 1 2 (θ′ − θ)T ( ∇2L(1)N (θ) + γI ) (θ′ − θ) ) dθ′
+ C(γ, p,m),
where (i) follows from Assumption 2 and the fact that ‖θ′′−θ‖ ≤ ‖θ−θ′‖ and (ii) follows because exp(− 12 (θ ′ − θ)T (∇2L(1)N (θ) + γI)(θ′ − θ)) ≥ 0.
Based on Assumption 1, we have the fact that (∇2L(1)N (θ (1) T (φ)) + γI) is a symmetric and positivedefinite matrix. Hence, we obtain
G (1) N (θ (1) T (φ); γ) ≤ −L (1) N (θ (1) T (φ))− log(det(∇ 2L (1) N (θ (1) T (φ)) + γI)) + C(γ, p,m)
+mM + 1
2 ρm3 +
p 2 log(2π).
We rearrange the terms and get
−G(1)N (θ (1) T (φ); γ) ≥ L (1) N (θ (1) T (φ)) + log(det(∇ 2L (1) N (θ (1) T (φ)) + γI))− C(γ, p,m)
−mM − p 2 log(2π)− 1 2 ρm3.
Since −G(1)N (θ (1) T (φ); γ) ≤ BEntropy(λ). Then, we obtain
BEntropy(λ) +mM + p
2 log(2π) +
1 2 ρm3 + C(γ, p,m) ≥ log(det(∇2L(1)N (θ (1) T (φ)) + γI))
+ L (1) N (θ (1) T (φ)).

heading: C.2 PROOF OF THEOREM 2
text: Based on Lemma 6, we have
BEntropy(λ) +mM + p
2 log(2π) +
1 2 ρm3 + C(γ, p,m) ≥ log(det(∇2L(1)N (θ (1) T (φ)) + γI))
+ L (1) N (θ (1) T (φ)).
Since ∇2L(1)N (θ (1) T (φ)) + γI is positive definite and λi(∇2L (1) N (θ (1) T (φ)) + γI) ≥ γ + µ for any i = 1, . . . , p. Then, based on the definition of Matrix norm ‖∇2L(1)N (θ (1) T (φ)) + γI‖ =
λmax(∇2L(1)N (θ (1) T (φ)) + γI), we have
‖∇2L(1)N (θ (1) T (φ)) + γI‖ p ≥ det(∇2L(1)N (θ (1) T (φ)) + γI)) ≥ (γ + µ) p−1‖∇2L(1)N (θ (1) T (φ)) + γI‖.
Note that we use λi(H) to denote the i-th eigenvalue of matrix H . Then,
log(det(∇2L(1)N (θ (1) T (φ)) + γI))) ≥(p− 1) log(γ + µ) + log ‖∇ 2L (1) N (θ (1) T (φ)) + γI‖
=(p− 1) log(γ + µ) + log(‖∇2L(1)N (θ (1) T (φ))‖+ γ).
Then, we can obtain
BEntropy(λ)+mM + p
2 log(2π) +
1 2 ρm3 + C(γ, p,m) ≥ L(1)N (θ (1) T (φ)) + (p− 1) log(γ + µ)
+ log(‖∇2L(1)N (θ (1) T (φ))‖+ γ).
Hence, we can get a new function D(x) that
‖∇2L(1)N (θ (1) T (φ))‖ ≤ D −1(BEntropy(λ)),
where D(x) = L(1)N (θ (1) T (φ)) + (p − 1) log(γ + µ) −mM − p 2 log(2π) − 1 2ρm 3 − C(γ, p,m) + log(x+ γ). Then, the proof is complete.

Review summaries:
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.

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).

This work proposes to apply flatness-aware regularizations to learning to optimize methods. Specifically, the approximate hessian spectrum and entropy based regularizers are added to the L2O objective. They performed a theoretical analysis of the regularized L2O method and evaluated different variants of the regularizer on simple CNN and MLP models trained on CIFAR and MNIST. Similar to the effect of flatness-aware regularization on the ordinary training, the regularizer improved the learned optimizers to find minima that generalizes better.  

The paper aims at improving the generalization of the learned optimizers. Some regularizers to induce the flatness of local minimas including Hessian spectrum and the entropy function are proposed. Both theoretical analysis and empirical results are presented to demonstrate that flatness-aware regularizers can enhance the generalization ability of optimizees.