Optimizer Amalgamation

1 INTRODUCTION . 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 specific 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-specific 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 ? Specifically , we wish for our combined optimizer to be adaptive both per-parameter and per-iteration , and exploit problem-specific 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 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 define 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 define as finding 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 . Specifically , we find 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 significantly better than previous methods on all problems , with few exceptions . 2 RELATED WORKS . Knowledge Distillation and Amalgamation The prototype of knowledge distillation was first 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 significant performance gains as a result . The success of knowledge distillation spurred significant 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 benefits of knowledge distillation to the improved smoothness of loss surfaces , which has been demonstrated to help adversarial training Cohen et al . ( 2019 ) ; 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 final 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 first 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 ) , and adversarial training ( Jiang et al. , 2018 ; Xiong & Hsieh , 2020 ) . Moving away from gradient-based optimization , black-box optimization has also been explored ( Chen et al. , 2017 ; Cao et al. , 2019 ; Shen et al. , 2021 ) . 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 . ( 2020 ) ; Simsekli et al . ( 2019 ) which is often highly non-Gaussian and heavy-tail in practice ; the random initialization and ( often non-optimal ) hyperparameter configuration ; 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 . ( 2017 ) . 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 . ( 2016 ) . Recent work ( He et al. , 2019 ; Wu et al. , 2020 ) extends input perturbations to weight perturbations that explicitly regularize the flatness 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 ) . 3 OPTIMIZER AMALGAMATION . 3.1 MOTIVATION . Optimizer selection and hyperparameter optimization is a difficult task even for experts . With a vast number of optimizers to choose from with varying performance dependent on the specific 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 benefit 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 sufficiently 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 . ( 2016 ) ; Lv et al . ( 2017 ) , recurrent neural networks with hierarchical architectures Wichrowska et al . ( 2017 ) ; Metz et al . ( 2019 ) , and symbolically in terms of predefined blocks Bello et al . ( 2017 ) . 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.2 .

Summary: This paper proposed a new problem called Optimizer Amalgamation and made an attempt to obtain a more powerful learned optimizer from several analytical optimizers. More concretely, three amalgamation losses are designed to train the amalgamated optimizer. In the meanwhile, two types of noise, random gaussian perturbation and projected gradient descent are incorporated into the training objective to increase the stability of the optimizer. The evaluation part compared the amalgamated optimizer with its original teachers, i.e., those analytical optimizers, and some L2O baselines on image classification tasks.