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 discusses the problem of selecting a neural network optimizer from a pool of possible optimizers. They propose three variants of a meta-algorithm which combines optimizers from the pool, whereby a differentiable meta-loss is defined on the training loss achieved by selection protocol. They show that their algorithm, equipped with weight space training noise leads to better performance on a variety of problems.