O PTIMIZER A MALGAMATION
A BSTRACT
Selecting an appropriate optimizer for a given problem is of major interest for researchers and practitioners.
Many analytical optimizers have been proposed using a variety of theoretical and empirical approaches; however, none can offer a universal advantage over other competitive optimizers.
We are thus motivated to study a new problem named Optimizer Amalgamation : how can we best combine a pool of “teacher” optimizers into a single “student” optimizer that can have stronger problem-speciﬁc performance?
In this paper, we draw inspiration from the ﬁeld of ”learning to optimize” to use a learnable amalgamation target.
First, we deﬁne three differentiable amalgamation mechanisms to amalgamate a pool of analytical optimizers by gradient descent.
Then, in order to reduce variance of the amalgamation process, we also explore methods to stabilize the amalgamation process by perturbing the amalgamation target.
Finally, we present experiments showing the superiority of our amalgamated optimizer compared to its amalgamated components and learning to optimize baselines, and the efﬁcacy of our variance reducing perturbations.
Our code and pre-trained models are publicly available at http://github.com/VITA-Group/OptimizerAmalgamation .
I NTRODUCTION
Gradient-based optimization is ubiquitous in machine learning; accordingly, a cottage industry of gradient-based optimizer design has emerged (Schmidt et al., 2020).
These optimizers generally propose algorithms that aim to make the “best” parameter update for a computed gradient (Kingma &
Ba, 2017; Liu et al., 2020), with some also modifying the location where the parameters are computed (Zhang et al., 2019b).
However, each gradient-based optimizer claim speciﬁc problems where they hold performance advantages, none can claim to be universally superior.
Due to the “No Free Lunch” theorem for optimization (Wolpert & Macready, 1997), no optimizer can provide better performance on a class of problems without somehow integrating problem-speciﬁc knowledge from that class.
Furthermore, problems such as training neural networks are not homogeneous.
In the spatial dimension, different layers or even parameters can have different behavior (Chen et al., 2020b).
Also, as evidenced by the popularity of learning rate schedules, neural network optimization also behaves very differently in the temporal dimension as well (Golatkar et al., 2019).
This implies that no optimizer can provide the best performance for all parameters on a single problem or best performance over the entire optimization process.
In order to build a stronger optimizer, we propose the new problem of optimizer amalgamation : how can we best combine a pool of multiple “teacher” optimizers, each of which might be good in certain cases, into a single stronger “student” optimizer that integrates their strengths and offsets their weaknesses?
Speciﬁcally, we wish for our combined optimizer to be adaptive both per-parameter and per-iteration, and exploit problem-speciﬁc knowledge to improve performance on a class of problems.
To “amalgamate” an optimizer from a pool of optimizers, we draw inspiration from recent work in
Learning to Optimize (Chen et al., 2021a) which provides a natural way to parameterize and train optimization update rules.
In Learning to Optimize, optimizers are treated as policies to be learned from data.
These “learned” optimizers are typically parameterized by a recurrent neural network (Andrychowicz et al., 2016;
Lv et al., 2017); then, the optimizer is meta-trained to minimize the loss of training problems, or “optimizees”, by gradient descent using truncated back-propagation through time.
Yet to our best knowledge, no existing work has leveraged those learnable parameterizations to amalgamate and combine analytical optimizers.
For our proposed formulation of optimizer amalgamation, we treat the learned optimizer as the amalgamation target.
Then, we deﬁne amalgamation losses which can be used to combine feedback from multiple analytical optimizers into a single amalgamated optimizer, and present several amalgamation schemes.
Finally, we explore smoothing methods that can be used during the amalgamation process to reduce the variance of the amalgamated optimizers.
Our contributions are outlined below:
• We formulate the new problem of optimizer amalgamation, which we deﬁne as ﬁnding a way to best amalgamate a pool of multiple analytical optimizers to produce a single stronger optimizer.
We present three schemes of optimizer amalgamation: additive amalgamation, min-max amalgamation, and imitation of a trained choice.
• We observe instability during the amalgamation process which leads to amalgamated optimizers having varied performance across multiple replicates.
To mitigate this problem, we explore ways to reduce amalgamation variance by improving smoothness of the parameter space.
We propose smoothing both by random noise or adversarial noise.
• We present experiments showing extensive and consistent results that validate the effectiveness of our proposal.
Speciﬁcally, we ﬁnd that more advanced amalgamation techniques and weight space training noise lead better average case performance and reduced variance.
We also show that our amalgamation method performs signiﬁcantly better than previous methods on all problems, with few exceptions.
R ELATED W ORKS
Knowledge Distillation and Amalgamation The prototype of knowledge distillation was ﬁrst introduced by (Bucilua et al., 2006), which used it for model compression in order train neural networks (“students”) to imitate the output of more complex models (“teachers”).
Knowledge distillation was later formalized by (Hinton et al., 2015), who added a temperature parameter to soften the teacher predictions and found signiﬁcant performance gains as a result.
The success of knowledge distillation spurred signiﬁcant efforts to explain its effectiveness.
Notably,
Chen et al.
(2020c)
; Yuan et al.
(2020) discovered that trained distillation teachers could be replaced by hand-crafted distributions.
(Yuan et al., 2020) provided further theoretical and empirical explanation for this behavior by explicitly connecting Knowledge distillation to label smoothing, and (Ma et al.;
Chen et al., 2021b) further credited the beneﬁts of knowledge distillation to the improved smoothness of loss surfaces, which has been demonstrated to help adversarial training Cohen et al.
Lecuyer et al.
(2019) and the training of sparse neural networks Ma et al..
The potential of knowledge distillation to improve the training of neural networks also spurred diverse works extending knowledge distillation.
For example, (Romero et al., 2015; Wang et al., 2018;
Shen et al., 2018; 2019b; Ye et al., 2020b) propose using intermediate feature representations as distillation targets instead of just network outputs, and (Tarvainen & Valpola, 2017; Yang et al., 2018; Zhang et al., 2019a) unify student and teacher network training to reduce computational costs.
Knowledge distillation has also been extended to distilling multiple teachers, which is termed
Knowledge Amalgamation (Shen et al., 2019a;
Luo et al., 2019; Ye et al., 2019; 2020a).
Although using output logits from pre-trained networks has been extensively explored in knowledge distillation, we study a new direction of research distilling optimization knowledge from sophisticated analytical optimizers to produce stronger “learned” optimizers, hence the name “optimizer amalgamation”.
Not only this is a new topic never studied by existing knowledge distillation literature, but also it needs to distill longitudinal output dynamics — not one ﬁnal output — from multiple teachers.
Learning to optimize Learning to Optimize is a branch of meta learning which proposes to replace hand-crafted analytical optimizers with learned optimizers trained by solving optimization problems, or optimizees .
The concept was ﬁrst introduced by (Andrychowicz et al., 2016), who used a Long
Short-Term Memory (LSTM) based model in order to parameterize gradient-based optimizers.
This model took the loss gradient as its input and output a learned update rule which was then trained by gradient descent using truncated backpropagation through time.
(Andrychowicz et al., 2016) also established a coordinate-wise design pattern, where the same LSTM weights are applied to each parameter of the optimizee in order to facilitate generalization to models with different architectures.
Building on this architecture, Wichrowska et al.
(2017) and Lv et al.
(2017) proposed improvements such as hierarchical architectures connecting parameter RNNs together and augmenting the gradient with additional inputs.
Many methods have also been proposed to improve the training of learned optimizers such as random scaling and convex augmentation (Lv et al., 2017), curriculum learning and imitation learning (Chen et al., 2020a), and Jacobian regularization (Li et al., 2020).
Notably,
Chen et al.
(2020a) also proposed a method of imitation learning, which can be viewed as a way of distilling a single analytical optimizer into a learned parameterization.
Learning to Optimize has been extended to a variety of other problems such as graph convolutional networks (You et al., 2020), domain generalization (Chen et al., 2020b), noisy label training (Chen et al., 2020c), adversarial training (Jiang et al., 2018; Xiong & Hsieh, 2020), and mixmax optimization (Shen et al., 2021).
Moving away from gradient-based optimization, black-box optimization has also been explored (Chen et al., 2017; Cao et al., 2019).
For a comprehensive survey with benchmarks, readers may refer to Chen et al.
(2021a).
Perturbations and Robustness
The optimization process is naturally subject to many possible sources of noise, such as the stochastic gradient noise Devolder et al.
(2011); Gorbunov et al.
Simsekli et al.
(2019) which is often highly non-Gaussian and heavy-tail in practice; the random initialization and (often non-optimal) hyperparameter conﬁguration; the different local minimum reached each time in non-convex optimization Jain & Kar (2017); and the limited numerical precision in implementations De Sa et al.
The seen and unseen optimizees also constitute domain shifts in our case.
In order for a consistent and reliable amalgamation process, the training needs to incorporate resistance to certain perturbations of the optimization process.
We draw inspiration from deep learning defense against various random or malicious perturbations.
For example, stability training Zheng et al.
(2016) stabilizes deep networks against small input distortions by regularizing the feature divergence caused by adding random Gaussian noises to the inputs.
Adversarial robustness measures the ability of a neural network to defend against malicious perturbations of its inputs (Szegedy et al., 2013; Goodfellow et al., 2014).
For that purpose, random smoothening (Lecuyer et al., 2019; Cohen et al., 2019) and adversarial training (Madry et al., 2017) have been found to increase model robustness with regard to random corruptions or worst-case perturbations; as well as against testing-time domain shifts Ganin et al.
Recent work (He et al., 2019; Wu et al., 2020) extends input perturbations to weight perturbations that explicitly regularize the ﬂatness of the weight loss landscape, forming a double-perturbation mechanism for both inputs and weights.
Other Approaches
The problem of how to better train machine learning models has many diverse approaches outside the Learning to Optimize paradigm that we draw from, and forms the broader
AutoML problem Hutter et al.
(2018) together with model selection algorithms.
Our approach falls under meta-learning, which also includes learned initialization approaches such as MAML (Finn et al., 2017) and Reptile (Nichol et al., 2018).
Other optimizer selection and tuning methods include hypergradient descent (Baydin et al., 2017) and bayesian hyperparameter optimization (Snoek et al., 2012).
Similar to our knowledge amalgamation approach, algorithm portfolio methods (where many algorithms are available, and some subset is selected) have also been applied to several problem domains such as Linear Programming (Leyton-Brown et al., 2003) and SAT solvers
(Xu et al., 2008).
O PTIMIZER A MALGAMATION 3. M OTIVATION
Optimizer selection and hyperparameter optimization is a difﬁcult task even for experts.
With a vast number of optimizers to choose from with varying performance dependent on the speciﬁc problem and data (Schmidt et al., 2020), most practitioners choose a reasonable default optimizer such as
SGD or Adam and tune the learning rate to be “good enough” following some rule of thumb.
As a consequence of the No Free Lunch theorem (Wolpert & Macready, 1997), the best optimizer to use for each problem, weight tensor within each problem, or each parameter may be different.
In practice, different layers within a given neural network can beneﬁt from differently tuned hyperparameters, for example by meta-tuning learning rates by layer (Chen et al., 2020b).
Accordingly, we wish to train an optimizer which is sufﬁciently versatile and adaptive at different stages of training and even to each parameter individually.
Many methods have been proposed to parameterize optimizers in learnable forms including coordinate-wise LSTMs Andrychowicz et al.
Lv et al.
(2017), recurrent neural networks with hierarchical architectures Wichrowska et al.
Metz et al.
(2019), and symbolically in terms of predeﬁned blocks Bello et al.
Due to its high expressiveness and relative ease of training, we will use the workhorse of LSTM-based
RNNProp architecture described by Lv et al.
(2017) as our amalgamation target; more details about this architecture can be found in Appendix C.1.
3. T HE B ASIC D ISTILLATION M ECHANISM
Knowledge distillation can be viewed as regularizing the training loss with a distillation loss that measures the distance between teacher and student predictions (Hinton et al., 2015).
In order to distill a pool of teacher optimizers T =
T 1 , T 2 , . . .
T k into our target policy P by truncated backpropagation (Appendix A: Algorithm 1), we start by deﬁning a training loss and amalgamation loss.
Meta Loss
In the context of training optimizers, the training loss is described by the meta loss , which is a function of the optimizee problem loss at each step (Andrychowicz et al., 2016).
Suppose we are training a policy P with parameters φ on a problem M : X → R whose output is a loss for each point in data domain X .
During each iteration during truncated backpropagation through time, P is used to compute parameter updates for M to obtain a trajectory of optimizee parameters θ 1 , θ 2 , . . .
θ N where for the i th data batch x
i and parameters θ i at step i , i.e. θ
i
+
=
θ i − P ( ∇ θ i M
( x i , θ i )) .
For some weighting function f 1 , f 2 , . . .
f N , the meta loss is L meta ( x , θ i ; φ ) =
(cid:80) Ni =
1 f i ( M ( x i , θ i )) ; speciﬁcally, we will use the scaled log meta loss
f i ( m ) = log( m ) − log( M ( x i , θ 0 )) , which can be interpreted as the “mean log improvement.”
Distillation Loss
The distillation loss in knowledge distillation measures the distance between teacher predictions and student predictions.
In training optimizers, this corresponds to the distance between the optimization trajectories generated by the teacher and student.
Suppose we have optimizee parameter trajectories θ i
= ( θ ( P [Equation] i , θ ( T ) i ) generated by the teacher and student, respectively.
Then, our distillation loss L T for teacher T is given by the l 2 log-loss:
[Equation]
While knowledge distillation generally refers to imitating a model and imitation learning imitating a policy, the optimizer in our case can be regarded as both a model and a policy.
As such, our loss function is similar to the imitation loss mechanism used by Chen et al.
(2020a), which can be thought of as a special case of optimizer amalgamation where only a single teacher is used.
3. A MALGAMATION OF M ULTIPLE T EACHER O PTIMIZERS : T HREE S CHEMES
Now, what if there are multiple teachers that we wish to amalgamate into a single policy?
How to best combine different knowledge sources is a non-trivial question.
We propose three mechanisms:
(1) Mean Amalgamation : adding distillation loss terms for each of the optimizers with constant equal weights.
(2) Min-max Amalgamation : using a min-max approach to combine loss terms for each of the optimizers, i e., “the winner (worst) takes all”
(3) Optimal Choice Amalgamation : First training an intermediate policy to choose the best optimizer to apply at each step, then distilling from that “choice optimizer”.
Mean Amalgamation In order to amalgamate our pool of teachers T = T 1 , . . .
T | T | , we generate | T | + 1 trajectories θ i
= ( θ ( P ) i , θ ( T 1 ) i . . .
θ ( T | T | ) i ) and add distillation losses for each teacher:
[Equation]
If we view knowledge distillation as a regularizer which provides soft targets during training, mean amalgamation is the logical extension of this by simply adding multiple regularizers to training.
An interesting observation is: when multiple teachers diverge, mean amalgamation loss tends to encourage the optimizer to choose one of the teachers to follow, potentially discarding the inﬂuence of all other teachers.
This may occur if one teacher is moving faster than another in the optimizee space, or if the teachers diverge in the direction of two different minima.
As this choice is a local minimum with respect to the mean log amalgamation loss, the optimizer may “stick” to that teacher, even if it is not the best choice.
Min-Max Amalgamation In order to address this stickiness, we propose a second method: minmax amalgamation, where, distillation losses are instead combined by taking the maximum distillation loss among all terms at each time step:
[Equation]
This results in a v-shaped loss landscape which encourages the amalgamation target to be between the trajectories generated by the teacher pool and prevents the optimizer from “sticking” to one of the teachers.
One weakness shared by both mean and min-max amalgamation is memory usage.
Both require complete training trajectories for each teacher in the pool to be stored in memory, resulting in memory usage proportional to the number of teachers, which limits the number of teachers that we could amalgamate from in one pool.
Min-max amalgamation also does not fully solve the problem of diverging teachers.
While minmax amalgamation does ensure that no teacher is ignored, it pushes the amalgamation target to the midpoint between the optimizee weights of the two teachers, which does not necessarily correspond to a good optimizee loss.
In fact, when teachers diverge into multiple local minima, any solution which considers all teachers must necessarily push the learned optimizer against the gradient, while any solution which allows the learned optimizer to pick one side must discard a number of teachers.
Optimal Choice Amalgamation To fully unlock the power of knowledge amalgamation, we propose to solve the teacher divergence problem by ﬁrst training an intermediate amalgamation target.
By using only one teacher for a ﬁnal distillation step, we remove the possibility of multiple teachers diverging while also allowing us to use more teachers without a memory penalty.
For optimizer pool T , we deﬁne an choice optimizer C which produces choices c 1 , c 2 , . . .
c N of which optimizer in the pool to apply at each time step, producing updates θ ( C )
i +
= θ ( C )
i − T c
i ( g i ) .
The objective of the choice optimizer is to minimize the meta loss L meta ( C ; x ) with respect to these choices c 1: N .
We parameterize the choice function C as a small two-layer LSTM, and train it by gradient descent.
The LSTM takes the outputs of each optimizer in the pool, the layer type, and time step as inputs; more details are provided in Appendix C.1.
To make it easier to train C by truncated back-propagation through time, we relax the choices c 1: N to instead be soft choices c i ∈ R | T | : c i ≥ 0 ,
|| c
i || 1 = 1 , resulting in the policy θ ( C ) i +
= θ ( C )
i −
[Equation]
| T | j =1 c ( j ) i
T j ( g i ) .
Now, we use
C as a teacher to produce our ﬁnal loss:
[Equation]
S TABILITY -A WARE O PTIMIZER A MALGAMATION 4. M OTIVATION
Modern optimization, even analytical, is subject to various forms of noise.
For example, stochastic ﬁrst-order method are accompanied with gradient noise (Devolder et al., 2011; Gorbunov et al., 2020; Simsekli et al., 2019) which is often highly non-Gaussian and heavy-tail in practice.
Any non-convex optimization could reach different local minimum when solving multiple times (Jain &
Kar, 2017).
When training deep neural networks, thousands or even millions of optimization steps are typically run, and the ﬁnal outcome can be impacted by the random initialization, (often non-optimal) hyperparameter conﬁguration, and even hardware precision (De Sa et al., 2017).
Hence, it is highly desirable for optimizers to be stable: across different problem instances, between multiple training runs for the same problem, and throughout each training run (Lv et al., 2017).
Meta-training optimizers tends to be unstable.
During the amalgamation process, we encounter signiﬁcant variance where identically trained replicates achieve varying performance on our evaluation problems; this mirrors problems with meta-stability encountered by Metz et al.
While amalgamation variance can be mitigated in small-scale experiments by amalgamating many times and using the best one, that variance represents a signiﬁcant obstacle to large-scale training (i.e. on many and larger problems) and deployment of amalgamated optimizers.
Thus, besides the aforementioned optimization stability issues, we also need to consider meta-stability , denoting the relative performance of optimizers across meta-training replicates.
In order to provide additional stability to the amalgamation process, we turn to adding noise during training, which is known to improve smoothness (Chen & Hsieh, 2020; Lecuyer et al., 2019; Cohen et al., 2019) and in turn improve stability (Miyato et al., 2018).
Note that one can inject either random noise or adversarial perturbations onto either the input or the weight of the learnable optimizer.
While perturbing inputs is more common, recent work (Wu et al., 2020) identiﬁed that a ﬂatter weight loss landscape (loss change with respect to weight) leads to smaller robust generalization gap in adversarial training, thanks to its more “global” worst-case view.
We also discover in our experiments that perturbing inputs would make the meta-training hard to converge, presumably because the inputs to optimizers (gradients, etc.) already contain large amounts of batch noise and do not tolerate further corruption.
We hence focus on perturbing optimizer weights for smoothness and stability.
4. W EIGHT S PACE P ERTURBATION FOR S MOOTHNESS
Weight space smoothing produces a noised estimate of the loss ˜ L by adding noise to the optimizer parameters φ .
By replacing the loss L ( φ, x ) with a noisy loss ˜ L
= L (˜ φ, x ) , we encourage the optimizer to be robust to perturbations of its weights, increasing the meta-stability.
We explore two mechanisms to increase weight space smoothness during training, by adding (1) a random perturbation to the weights as a gradient estimator, and (2) an adversarial perturbation in the form of a projected gradient descent attack (PGD).
Though new to our application, these two mechanisms have been adopted for other problems where smoothness is important such as neural architecture search (Chen & Hsieh, 2020) and adversarial robustness (Lecuyer et al., 2019; Cohen et al., 2019).
Random Gaussian Perturbation
In the ﬁrst type of noise, we add gaussian noise with variance σ 2 to each parameter of the optimizer at each iteration,
˜ φ
=
φ + N (0 , σ 2 I ) .
Since optimizer weights tend to vary largely in magnitude especially between different weight tensors, we modify this gaussian noise to be adaptive to the magnitude of the l 2 norm of each weight tensor
˜ φ ( w ) .
For tensor size | φ ( w ) | , the added noise is given by
[Equation]
Projected Gradient Descent For the second type of noise, we use adversarial noise obtained by projected gradient descent (Appendix A, Algorithm 2).
For A adversarial steps, the noised parameters are given by ˜
φ = φ + ψ A , where ψ 0 = 0 , and ψ i +1 =
ψ i + η clip ε ( ∇ ˜ ψ i L ) for optimizer loss L .
As with random Gaussian perturbations, we also modify the adversarial perturbation to be adaptive with magnitude proportional to the l 2 norm of each weight tensor φ .
Here, the adversarial attack step for weight tensor w is instead given by
[Equation]
E XPERIMENTS
Optimizee Details All optimizers were amalgamated using a 2-layer convolutional neural network (CNN) on the MNIST (LeCun & Cortes, 2010) dataset (shortened as “Train”) using a batch size of 128.
During evaluation, we test the generalization of the amalgamated optimizer to other problems: (cid:182)
Different Datasets : FMNIST (Xiao et al., 2017) and SVHN (Netzer et al., 2011).
We also run experiments on CIFAR (Krizhevsky et al., 2009); since the Train network is too small to obtain reasonable performance on CIFAR, we substitute it for the Wider architecture and a 28-layer ResNet (He et al., 2015) labelled “CIFAR” and “ResNet” respectively.
(cid:183) Different Architectures : a 2-layer MLP (MLP), a CNN with twice the number of units in each layer (Wider), and a deeper CNN (Deeper) with 5 convolutional layers.
(cid:184) Training settings : training with a smaller batch size of(Small Batch).
We also try a new setting of training with differential privacy (Abadi et al., 2016)
(MNIST-DP).
Appendix B provides full architecture and training speciﬁcations.
Optimizer Pool We use two different optimizer pools in our experiment: “small,” which consists of
Adam and RMSProp, and “large,” which also contains SGD, Momentum, AddSign, and PowerSign.
Each optimizer has a learning rate tuned by grid search over a grid of { 5 × 10 − 4 , 1 × 10 − 3 , 2 × 10 − 3 , . . .
The selection criteria is the best validation loss after 5 epochs for the Train network on MNIST, which matches the meta-training settings of the amalgamated optimizer.
Appendix C. describes the optimizers used and other hyperparameters.
Baselines First, we compare our amalgamated optimizer against our analytical optimizer teachers which are combined into a “oracle optimizer,” which is the optimizer in our pool of teachers with the best validation loss.
We also compare against the optimal choice optimizer used in amalgamation, which functions like a per-iteration trained approximation of the oracle optimizer.
Then, we evaluate previous learned optimizer methods: the original “Learning to Learn by Gradient Descent by Gradient
Descent” optimizer Andrychowicz et al.
(2016) which we refer to as “Original”, RNNProp
(Lv et al., 2017), a hierarchical architecture presented by “Learned Optimizers that Scale and Generalize” (Wichrowska et al., 2017) which we refer to as “Scale”, and the best setup from Chen et al.
(2020a) which shorten as “Stronger Baselines.”
Training and Evaluation Details
The RNNProp amalgamation target was trained using truncated backpropagation though time with a constant truncation length of 100 steps and total unroll of up to 1000 steps and meta-optimized by Adam with a learning rate of 1 × 10 − 3 .
For our training process, we also apply random scaling (Lv et al., 2017) and curriculum learning (Chen et al., 2020a); more details about amalgamation training are provided in Appendix C.3.
Amalgamation takes up to 6.35 hours for optimal choice amalgamation using the large pool and up to 10.53 hours when using adversarial perturbations; a full report of training times is provided in Appendix C.4.
For each optimizer amalgamation conﬁguration tested, we independently trained 8 replicate optimizers.
Then, each replicate was evaluated 10 times on each evaluation problem, and trained to a depth of 25 epochs each time.
Finally, we measure the stability of amalgamated optimizers by deﬁning three notions of stability for meta-trained optimizers: (cid:182) Optimization stability : the stability of the optimizee during the optimization process.
Viewing stability of the validation loss as a proxy for model stability with respect to the true data distribution, we measure the epoch-to-epoch variance of the validation loss after subtracting a smoothed validation loss curve (using a Gaussian ﬁlter).
(cid:183) Evaluation stability : the variance of optimizer performance across multiple evaluations.
We ﬁnd that the evaluation stability is roughly the same for all optimizers (Appendix E.1).
(cid:183) Meta-stability : the stability of the amalgamation process, i.e. the variance of amalgamation replicates after correcting for evaluation variance.
Meta-stability and evaluation stability are jointly estimated using a linear mixed effects model.
The stability is reported as a standard deviation. More details are in Appendix D. 5. O PTIMIZER A MALGAMATION
Amalgamation Methods Figure 1 compares the mean performance of the three amalgamation methods with the small pool and Choice amalgamation with the large pool.
Mean and min-max amalgamation were not performed on the large pool due to memory constraints.
The amalgamated optimizers using optimal choice amalgamation perform better than Mean and Min-Max amalgamation.
The size of the optimizer pool does not appear to have a signiﬁcant effect in Optimal Choice amalgamation, with small pool and large pool amalgamated optimizers obtaining similar results.
[Figure]
R e l a t i v e
B e s t
L o g
V a li d a t i o n L o ss
[Figure]
[Figure]
l
[Figure]
t i v e
B
[Figure]
s t
L o g
V a l L
[Figure]
[Figure]
[Figure]
optimizers in the shown Previous Learned Optimizers Figure 2 compares the amalgamated optimizer against the baselines from Learning to Optimize.
Optimizer amalgamation performs significantly better than all previous methods on all problems, with few exceptions (where it performs better but not signiﬁcantly better).
Analytical Optimizers In Figure 3, we compare the best replicate amalgamated from the large pool using
Choice amalgamation with the “oracle optimizer” described above.
The amalgamated optimizer achieves similar or better validation losses than the best analytical optimizers, indicating that our amalgamated optimizer indeed captures the “best” loss-minimization characteristics of each optimizer.
The amalgamated optimizer also beneﬁts from excellent optimization stability, meeting or exceeding the optimization stability of the best analytical optimizers in the large [Figure] t
i [Figure] i z [Figure]
t i [Figure] n
[Figure]
t
[Figure]
pool (Figure 5).
Comparing analytical optimizers, we observe a general inverse relationship between optimization performance and optimization stability: in order to achieve better optimization, an optimizer typically sacriﬁces some optimization stability in order to move faster through the optimizee weight space.
By integrating problem-speciﬁc knowledge, the amalgamated optimizer is able to combine the best optimization performance and optimization stability characteristics (Figure 4).
5. S TABILITY -A WARE O PTIMIZER A MALGAMATION
Input Perturbation While we also tested perturbing the inputs of the optimizer during amalgamation, we were unable to improve stability.
These experiments are included in Appendix E.4.
Random Perturbation Min-max amalgamation was trained on the small optimizer pool with random perturbation relative magnitudes of ε = { 5 × 10 − 4 , 10 − 3 , 2 × 10 − 3 , 5 × 10 − 3 , 10 − 2 } .
ε = 10 − 1 was also tested, but all replicates tested diverged and are not reported here.
[Figure]
Comparing perturbed amalgamation against the nonperturbed baseline ( ε = 0 ), we observe that perturbations increase meta-stability up to about ε = 10 − 3 (Figure 6).
For larger perturbation magnitudes, meta-stability begins to decrease as the perturbation magnitude overwhelms the weight “signal,” eventually causing the training process to completely collapse for larger perturbation values.
While the stability with random perturbation ε = 10 − 2 is better than 10 − 3 , this is likely due to random chance, since we use a small sample size of 8 replicates.
Adversarial Perturbation Since adversarial perturbation is more computationally expensive than random perturbations, min-max amalgamation was tested on a coarser grid of relative magnitudes ε = { 10 − 4 , 10 − 3 , 10 − 2 } , and to an adversarial attack depth of 1 step.
These results are also reported in Figure 6, with ε = 10 − 2 omitted since all replicates diverged during training.
From our results, we observe that adversarial perturbations are about as effective as random perturbations.
We also observe that the maximum perturbation magnitude that the amalgamation process can tolerate is much smaller for adversarial perturbations compared to random perturbations, likely because adversarial perturbations are much “stronger.”
Due to the signiﬁcantly larger training cost of adversarial perturbations, we recommend random perturbations for future work.
Application to Other Methods Random and Adversarial perturbations can be applied to any gradient-based optimizer meta-training method, including all of our baselines.
An experiment applying Gaussian perturbations to the RNNProp baseline can be found in Appendix E.5.
C ONCLUSION
We deﬁne the problem of optimizer amalgamation, which we hope can inspire better and faster optimizers for researchers and practitioners.
In this paper, we provide a procedure for optimizer amalgamation, including differentiable optimizer amalgamation mechanisms and amalgamation stability techniques.
Then, we evaluate our problem on different datasets, architectures, and training settings to benchmark the strengths and weaknesses of our amalgamated optimizer.
In the future, we hope to bring improve the generalizability of amalgamated optimizers to even more distant problems.
A A LGORITHMS
In this section, we provide a detailed description of the key algorithms used in our paper.
Truncated Back-propagation : Algorithm 1 shows truncated back-propagation applied to optimizer amalgamation. For an unrolling length N , N data points (batches, in the case of mini-batch SGD) are sampled, which are split into N/t truncations of length t . Note that this requires N to be divisible by t ; in our implementation, we require t and N/t to be speciﬁed as integers. For each truncation, the optimizee and teachers are trained for t iterations, and meta-gradients are computed over that truncation and applied.
Adversarial Weight Perturbation : Algorithm 2 shows adversarial perturbations applied to optimizer amalgamation.
For each adversarial attack step, meta-gradients are taken with respect to the parameters, and are normalized for each tensor with respect to its tensor norm before being applied as an adversarial perturbation.
Algorithm 1:
Distillation by Truncated Back-propagation
Inputs : Amalgamation loss L a
Policy P with parameters φ
Teacher policies T = T 1 , . . .
T | T |
Optimizee M , X , θ
Unrolling, truncation lengths N , t
Outputs : Updated policy parameters φ
Sample N data points x 1 , . . .
x N from X .
θ ( P )
=
θ ( T 1 )
= . . .
= θ ( T | T | )
= θfor i = 1 , 2 , . . .
N/t do for j = 1 , 2 , . . .
t do n =
( i − 1) t + j
Update optimizee for P : θ ( P ) n +1 ← θ ( P ) n − P (cid:104) ∇M ( x n , θ ( P ) n ) (cid:105)
Update optimizees for each teacher: for k = 1 , . . .
| T | do θ ( T k ) n +1 ← θ ( T k ) n − T k (cid:104) ∇M ( x n , θ ( T k ) n ) (cid:105) end end
Compute distillation loss:
L i ← L a ( x [( i − 1) t : it ] , θ [( i − 1) t : it ] ; φ ) Update φ using ∇L i end
Inputs : Truncated back-propagation parameters L a , P , φ , T , M , X , θ 0 , N , t
Adversarial attack steps A
Outputs : Updated policy parameters φ
Sample N data points and initialize optimizee parameters for i
= 1, 2, ...N/t do ψ 0 ←for a = 1, 2, ...
A do
Compute trajectories θ [( i − 1) t : it ] for P and T
L ( a ) i ← L a ( x [( i − 1) t : it ] , θ [( i − 1) t : it ] ; φ + ψ a ) for each weight tensor w do γ ← ∇ ψ ( w ) i L ( a ) i ||∇ ψ ( w ) i L ( a ) i ||ψ ( w ) a ← ψ a − 1 +
ε
||
φ || 2 γ.
end end
L i ← L a ( x [( i − 1) t : it ] , θ [( i − 1) t : it ] ; φ + ψ A ) Update φ using ∇L i end
B O PTIMIZEE D ETAILS
Table 1 shows a summary of the training problems used.
While all training is performed on a 2-layer CNN on MNIST, we evaluated our optimizer on 4 different datasets described in (B.1) and 5 different architectures (described in B.2). We also experiment with different training settings, which are described in B.3.
B. D ATASETS
All datasets used are classiﬁcation datasets, with cross entropy used as the training loss.
The MNIST dataset (LeCun & Cortes, 2010) is used during training; the other datasets are, from most to least similar, are:
[Table]
• FMNIST: Fashion MNIST (Xiao et al., 2017).
FMNIST is also a drop-in replacement for MNIST with 10 classes and 28x28 grayscale images.
Unlike MNIST or KMNIST, it features images of clothing instead of handwritten characters.
• SVHN: Street View House Numbers, cropped (Netzer et al., 2011).
While SVHN also has 10 classes of numerical digits, the images are 32x32 RGB, and have signiﬁcantly more noise than MNIST including “distraction digits” on each side.
• CIFAR-10 (Krizhevsky et al., 2009): the least similar dataset.
While CIFAR-10 still hasclasses and 32x32 RGB images, it has much higher noise and within-class diversity.
Sample images from these datasets are shown in Figure 7.
All datasets were accessed using TensorFlow Datasets and have a CC-BY 4.0 license.
[Figure]
[Table]
A RCHITECTURES
The Train convolutional network (Table 2a) has one convolution layer with 16 3x3 ﬁlters and one convolution layer with 32 5x5 ﬁlters.
Each convolution layer uses ReLU activation, has stride 1x1, and is followed by a max pooling layer with size 2x2.
Finally, a fully connected softmax layer is used at the output.
The four architectures evaluated are:
MLP: a 2-layer MLP with 20 hidden units and sigmoid activation 2.
Wider: a modiﬁed version of Train with double width on each layer (Table 2b) 3.
Deeper: a deeper network with 5 convolutional layers instead of 2 (Table 2c), again using ReLU activation and 1x1 stride 4.
ResNet: a 28-layer ResNet (He et al., 2015) (Table 2d) [Table]
B. O PTIMIZEE T RAINING
During training, a batch size of 128 is used except for the Small Batch evaluation, which has a batch size of 32.
During training and evaluation, datasets are reshufﬂed each iteration.
To match the warmup process used in meta-training, warmup is also applied during evaluation.
The
SGD learning rate is ﬁxed at 0.01, which is a very conservative learning rate which does not optimize quickly, but is largely guaranteed to avoid divergent behavior.
For differentially private training, we implement differentially private SGD (Abadi et al., 2016).
In differentially private SGD, gradients are ﬁrst clipped to a ﬁxed l 2 norm ε on a per-sample basis; then, gaussian noise with standard deviation σε where σ > 1 is added to the aggregated batch gradients.
In our experiments, we use clipping norm ε = 1 . 0 and noise ratio σ = 1 . 1 .
Both MNIST and KMNIST are used as training sets in order to simulate transfer from a non-private dataset (MNIST) used for meta-training to a private dataset (KMNIST).
C A MALGAMATION D ETAILS C. O PTIMIZER A RCHITECTURES
In this section, we provide the exact architecture speciﬁcations and hyperparameters of our amalgamated optimizer along with other training details and training time.
Our implementation is open source, and can be found here: http://github.com/VITA-Group/OptimizerAmalgamation .
C.1. RNNP ROP A RCHITECTURE
For our amalgamation target, we use RNNProp architecture described by Lv et al.
For each parameter on each time step, this architecture takes as inputs RMSProp update g/ √ ˆ v and Adam update ˆ m/ √ ˆ v using momentum ( ˆ m ) decay parameter β 1 = 0 .
9 and variance ( ˆ v ) decay parameter β 2 = 0 .
999 , matching the values used for our analytical optimizers.
These values pass through a 2-layer LSTM with tanh activation, sigmoid recurrent activation, and 20 units per layer.
The output of this 2-layer LSTM is passed through a ﬁnal fully connected layer with tanh activation to produce a scalar ﬁnal update for each parameter.
C.1. C HOICE N ETWORK A RCHITECTURE
Our Choice network for Optimal Choice Amalgamation is a modiﬁed RNNProp architecture.
The update steps for each analytical optimizer are given as inputs to the same 2-layer LSTM used in
RNNProp.
Additionally, the current time step and tensor number of dimensions are provided, with the number of dimensions being encoded as a one-hot vector.
Then, instead of directly using the output of a fully connected layer as the update, LSTM output passes through a fully connected layer with one output per optimizer in the pool.
This fully connected layer has a softmax activation, and is used as weights to combine the analytical optimizer updates.
C. O PTIMIZER P OOL
We consider six optimizers as teachers in this paper: Adam, RMSProp, SGD, Momentum, AddSign, and PowerSign.
These optimizers are summarized in table 3.
[Table]
g
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
Joining the popular hand-crafted optimizers Adam, RMSProp, SGD, and Momentum, AddSign and PowerSign are two optimizers discovered by neural optimizer search (Bello et al., 2017).
These two optimizers share the design principle that update steps should be larger when the momentum and gradient are in agreement:
[Equation]
Here, g represents the gradient, ˆ m an exponential moving average of the gradient.
In order to use AddSign and Powersign as teachers for gradient-based distillation, we modify them to be differentiable by replacing the sign function with a scaled tanh with magnitudes normalized by √ ˆ
[Equation]
By dividing by √ ˆ v , we provide a consistent magnitude to the tanh function so that sign agreement mechanism is not affected by overall gradient magnitudes.
For all optimizers, the momentum decay parameter is set to β 1 = 0 .
9 , the variance decay parameter is set to β 2 = 0 . 999 , and the learning rate multiplier is found by grid search on the Train optimizee over a grid of { 5 × 10 − 4 , 1 × 10 − 3 , 2 × 10 − 3 , . . .
C. A DDITIONAL T RAINING D ETAILS
During amalgamation, we apply a number of techniques from previous Learning to Optimize literature in order to boost training:
• Curriculum Learning : We apply curriculum learning (Chen et al., 2020a) to progressively increase the unrolling steps across a maximum of 4 stages with length 100, 200, 500, and 1000.
During curriculum learning, checkpoints are saved and validated every 40 “meta-epochs,” which refers to a single optimizee trajectory trained with truncated back-propagation.
• Random Scaling : We apply random scaling (Lv et al., 2017) to reduce overﬁtting to the gradient magnitudes of the training problem.
This random scaling is only applied to the amalgamation target; amalgamation teachers receive “clean” (unscaled) gradients.
• Warmup : Instead of initializing each training optimizee with random weights, we ﬁrst applysteps of SGD optimization as a “warmup” to avoid the turbulent initial phase of optimizing neural networks.
A SGD learning rate of 0.01 is used during this period, and was chosen to be very conservative on all optimizees tested.
These techniques are also applied to all of our baselines, except for Random Scaling is only applied to baselines using the RNNProp architecture since we ﬁnd that it harms the performance of other optimizer architectures.
C. T RAINING C OST
Table 4 provides a summary of the training costs for each amalgamation method and baseline are provided.
For optimal choice amalgamation, this includes both training the optimal choice optimizer and amalgamation training.
All values are reported as the mean across 8 replicates.
[Table]
All experiments were run on single nodes with 4x Nvidia 1080ti GPUs, providing us with a metabatch size of 4 simultaneous optimizations.
In order to replicate our results, GPUs with at least 11GB of memory are required, though less memory can be used if the truncation length for truncated back-propagation is reduced.
D S TABILITY D EFINITIONS
In this section, we provide the mathematical deﬁnition and measurement details of meta-stability, evaluation stability, and optimization stability.
D. M ETA - STABILITY AND E VALUATION S TABILITY
In order to quantify meta-stability and evaluation stability, we ﬁrst summarize the performance of each evaluation using the best validation loss obtained and the training loss of the last epoch.
Then, we model the best validation loss Y ( val ) ij and ﬁnal training loss Y ( train ) ij for replicate i and evaluation j with the linear mixed effect model
[Equation]
(9) where µ is the true mean, α i are IID random variables representing the meta-stability of the amalgamated optimizer, and ε ij are IID random variables representing the evaluation stability of each replicate.
The meta-stability and evaluation stability are then quantiﬁed by standard deviations σ α and σ ε .
D. O PTIMIZATION S TABILITY
To measure optimization stability, we model the validation loss L ij ( t ) at epoch t for replicate i and evaluation j as
[Equation]
for smooth function β ij ( t ) which represents the behavior of the evaluation and random variable η ( t ) ij which captures the optimization stability; we assume that η ( t ) ij is IID with respect to t .
In order to estimate σ η , we ﬁrst estimate β ij ( t ) by applying a Gaussian ﬁlter with standard deviation σ = 2 (epochs) and ﬁlter edge mode “nearest,” and ˆ σ ( ij ) η is calculated accordingly.
Finally, σ ( ij ) η is treated as a summary statistic for each evaluation, and the mixed effect model described previously (Equation 9) is ﬁt to obtain a ﬁnal conﬁdence interval for the mean optimization stability.
E A DDITIONAL R ESULTS E. E VALUATION S TABILITY
Table 5 summarizes the evaluation stability of analytical and amalgamated optimizers.
All optimizers obtain similar evaluation stability, except for cases where an optimizer cannot reliably train the optimizee at all such as Momentum, AddSign
[Table]
and PowerSign on the Deeper CNN.
In these cases, the optimizer consistently learns a constant or random classiﬁer, which in very low variance and high “stability.”
[Table]
E. L ARGER E VALUATIONS
In order to explore the limits of our optimizer, we evaluated the amalgamated optimizer with a 52-layer ResNet (763882 parameters — 40x the train network size), and the same 52-layer ResNet on CIFAR-100 instead of CIFAR-10.
These results are compared to CIFAR-10 on a 2-layer network and CIFAR-10 on a 28-layer ResNet using Adam as a single baseline (Figure 8).
While our amalgamated optimizer has signiﬁcant performance advantages on the shallow CIFARnetwork in our original evaluations and achieves performance parity in the 28-layer ResNet, the amalgamated optimizer can no longer perform as well as the oracle once we reach 52 layers and change to CIFAR-100.
[Figure]
E. A DDITIONAL P LOTS
In this section, we include plots providing alternate versions of Figures 2, 3, and 5 in the main text which had some outliers cropped out in order to improve readability.
[Figure]
l a t i v e
[Figure]
s t V a li
[Figure]
a t i o n
[Figure]
[Figure]
[Figure]
E. I NPUT P ERTURBATION
When applying input perturbations, we perturb the inputs to the optimizer, or the optimizee gradients, instead of the optimizer weights:
[Equation]
[Table]
We tested magnitudes σ = 10 − 1 and σ = 10 − 2 on a smaller experiment size of 6 replicates using Choice amalgamation on the small pool as a baseline; these results are given in Table 6.
Many experiment variants remain relating to input noise such as adding noise proportional to parameter norm or gradient norm or trying smaller magnitudes of noise, and this may be a potential area of future study.
However, we believe that input noise is generally not helpful to optimizer amalgamation, and did not study it further.
E. B ASELINES WITH R ANDOM P ERTURBATION
Our perturbation methods can be applied to any gradient-based optimizer meta-training method, including all of our baselines.
To demonstrate this application, we trained 8 RNNProp replicates with
Gaussian perturbations with magnitude 1 × 10 − 4 ; all other settings were identical to the RNNProp baseline.
With perturbations, the RNNProp baseline is signiﬁcantly improved, though not enough to match the performance of our amalgamation method.
[Figure]