Learned optimizers that outperform on wall-clock and validation loss

1 INTRODUCTION . Gradient based optimization is a cornerstone of modern machine learning . Improvements in optimization have been critical to recent successes on a wide variety of problems . In practice , this typically involves analysis and development of hand-designed optimization algorithms ( Nesterov , 1983 ; Duchi et al. , 2011 ; Tieleman & Hinton , 2012 ; Kingma & Ba , 2014 ) . These algorithms generally work well on a wide variety of tasks , and are tuned to specific problems via hyperparameter search . On the other hand , a complementary approach is to learn the optimization algorithm ( Bengio et al. , 1990 ; Schmidhuber , 1995 ; Hochreiter et al. , 2001 ; Andrychowicz et al. , 2016 ; Wichrowska et al. , 2017 ; Li & Malik , 2017b ; Bello et al. , 2017 ) . That is , to learn a function to perform optimization , targeted to particular problems of interest . In this way , the algorithm may learn task specific structure , enabling dramatic performance improvements over more general optimizers . However , training learned optimizers is notoriously difficult . Existing work in this vein can be classified into two broad categories . On one hand are black-box methods such as evolution ( Goldberg & Holland , 1988 ; Bengio et al. , 1992 ) , random search ( Bergstra & Bengio , 2012 ) , reinforcement learning ( Bello et al. , 2017 ; Li & Malik , 2017a ; b ) , or Bayesian optimization ( Snoek et al. , 2012 ) . However , these methods scale poorly with the number of optimizer parameters . The other approach is to use first-order methods , by computing the gradient of some measure of optimizer effectiveness with respect to the optimizer parameters . Computing these gradients is costly as we need to iteratively apply the learned update rule , and then backpropagate through these applications , a technique commonly referred to as “ unrolled optimization ” ( Bengio , 2000 ; Maclaurin et al. , 2015 ) . To address the problem of backpropagation through many optimization steps ( analogous to many timesteps in recurrent neural networks ) , many works make use of truncated backpropagation though time ( TBPTT ) to partition the long unrolled computational graph into separate pieces ( Werbos , 1990 ; Tallec & Ollivier , 2017 ) . This yields computational savings , but at the cost of increased bias ( Tallec & Ollivier , 2017 ) and/or exploding gradients due to many iterated update steps ( Pascanu et al. , 2013 ; Parmas et al. , 2018 ) . Existing methods address the bias at the cost of increased variance or computa- tional complexity ( Williams & Zipser , 1989 ; Ollivier et al. , 2015 ; Tallec & Ollivier , 2017 ) . Previous techniques for training RNNs via TBPTT have thus far not been effective for training optimizers . In this paper , we analytically and experimentally explore the debilitating role of bias and exploding gradients on training optimizers ( §2.3 ) . We then show how these pathologies can be remedied by optimizing the parameters of a distribution over the optimizer parameters , known as variational optimization ( Staines & Barber , 2012 ) ( §3 ) . We define two unbiased gradient estimators for this objective : a reparameterization based gradient , and evolutionary strategies ( Rechenberg , 1973 ; Nesterov & Spokoiny , 2011 ) . By dynamically reweighting the contribution of these two gradient estimators , we are able to avoid strongly biased or exploding gradients , and thus stably and efficiently train learned optimizers . We demonstrate the utility of this approach by training a learned optimizer to target optimization of convolutional networks on image classification ( §4 ) . On the targeted task distribution , this learned optimizer achieves better validation loss , and is five times faster in wall-clock time , compared to well tuned hand-designed optimizers such as SGD+Momentum , RMSProp , and ADAM ( Figure 1 ) . While not the explicit focus of this work , we also find that the learned optimizer demonstrates promising generalization ability on out of distribution tasks ( Figure 6 ) . 2 UNROLLED OPTIMIZATION FOR LEARNING OPTIMIZERS . 2.1 PROBLEM FRAMEWORK . Our goal is to learn an optimizer which is well suited to some set of target optimization tasks . Throughout the paper , we will use the notation defined in Figure 2 . Learning an optimizer can be thought of as a bi-level optimization problem ( Franceschi et al. , 2018 ) , with inner and outer levels . The inner minimization consists of optimizing of the weights ( w ) of a target problem ` ( w ) by the repeated application of an update rule . The update rule is a parameterized function that defines how to map the weights at iteration t to iteration t + 1 : w ( t+1 ) = u ( w ( t ) ; θ ) . Here , θ represents the parameters of the learned optimizer . In the outer loop , these optimizer parameters ( θ ) are updated so as to minimize some measure of optimizer performance , the outer-objective L ( θ ) . Our choice for L will be the average value of the target loss ` measured over either training or validation data . Throughout the paper , we use inner- and outer- prefixes to make it clear when we are referring to applying a learned optimizer on a target problem ( inner ) versus training a learned optimizer ( outer ) . 2.2 UNROLLED OPTIMIZATION . In order to train an optimizer , we wish to compute derivatives of the outer-objective L with respect to the optimizer parameters , θ . Doing this requires unrolling the optimization process . That is , we can form an unrolled computational graph that consists of iteratively applying an optimizer ( u ) to optimize the weights ( w ) of a target problem ( Figure 2 ) . Computing gradients for the optimizer parameters involves backpropagating the outer loss through this unrolled computational graph . This is a costly operation , as the entire inner-optimization problem must be unrolled in order to get a single outer-gradient . Partitioning the unrolled computation into separate segments , known as truncated backpropagation , allows one to compute multiple outer-gradients over shorter segments . That is , rather than compute the full gradient from iteration t = 0 to t = T , we compute gradients in windows from t = a to t = a+ τ . Below , we analyze two key problems with gradients computed via truncated unrolled optimization : bias and exploding norm . 2.3 EXPONENTIAL EXPLOSION OF GRADIENTS WITH INCREASED SEQUENCE LENGTH . We can illustrate the problem of exploding gradients analytically with a simple example : learning a learning rate . Following the notation in Figure 2 , we define the optimizer as w ( t+1 ) = u ( w ( t ) ; θ ) = w ( t ) − θ∇ ` ( w ( t ) ) , where θ is a scalar learning rate that we wish to learn for minimizing some target problem ` ( w ( t ) ) . For simplicity , we assume no minibatch x . The quantity we are interested in is the outer-gradient , the derivative of the loss after T steps of gradient descent with respect to θ . We can compute this outer-gradient ( see Appendix A ) as : d ` ( w ( T ) ) dθ = 〈 −g ( T ) , T−1∑ i=0 T−1∏ j=1 ( I − θH ( j ) ) g ( i ) 〉 , where g ( i ) and H ( j ) are the gradient and Hessian of the target problem ` ( w ) at iteration i and j , respectively . We see that this equation involves a sum of products of Hessians . In particular , the first term in the sum involves a product over the entire sequence of Hessians observed during training . That is , the outer-gradient becomes a matrix polynomial of degree T , where T is the number of gradient descent steps . Thus , the outer-gradient grows exponentially with T . We can see another problem with long unrolled gradients empirically . Consider the task of optimizing a loss surface with two local minima defined as ` ( w ) = ( w− 4 ) ( w− 3 ) w2 with initial condition w ( 0 ) = −1.2 using a momentum based optimizer with a parameterized momentum value θ ( Figure 3a ) . At low momentum values the optimizer converges in the first of the two local minima , whereas for larger momentum values the optimizer settles in the second minima . With even larger values of momentum , the iterate oscillates between the two minima before settling . We visualize both the trajectory of w ( t ) over training and the final loss value for different momentum values in Figure 3b and 3c . With increasing unrolling steps , the loss surface as a function of the momentum θ becomes less and less smooth , and develops near-discontinuities at some values of the momentum . In the case of neural network inner-problems and neural network optimizers , the outer-loss surface can grow even more complex with increasing number of unrolling steps . We illustrate this in Figure 3d and 3e for slices through the loss landscape L ( θ ) of the outer-problem for a neural network optimizer . 2.4 INCREASING BIAS WITH TRUNCATED GRADIENTS . Here , we demonstrate the bias introduced by using a short truncation window in unrolled optimization . These results are similar to those presented in Wu et al . ( 2016 ) , except that we utilize multiple truncations rather than a single , shortened unroll . First , consider outer-learning the learning rate of Adam when optimizing a small two layer neural network on MNIST ( LeCun , 1998 ) . Grid search can be used to find the optimal learning rate which is 0.004 . We initialize Adam a learning rate of 0.001 and outer-train using increasing truncation amounts ( Figure 4a ) . Despite initializing close to the optimal learning rate , when outer-training with truncated backprop the resulting learning rate decreases – that is the sum of truncated outer-gradients are anti-correlated to the true outer-gradient . We visualize the per-truncation gradients of 500 steps in Figure 4b and cumulative truncated gradients in Figure 4c . Early in inner-training there is a large negative outer-gradient which increases the learning rate . Later in inner-training , the outer-gradients are positive , decreasing the learning rate . Thus , the optimizer parameter is pulled in opposite directions by truncated gradients early versus late in inner-training , revealing an inherent tension with truncated unrolled optimization . 3 TOWARDS STABLE TRAINING OF LEARNED OPTIMIZERS . To perform outer-optimization of a loss landscape with high frequency structure like that in Figure 3 , one might intuitively want to smooth the outer-objective loss surface . To do this , instead of optimizing L ( θ ) directly , we instead optimize a smoothed outer-loss L ( θ ) , L ( θ ) = Eθ̃∼N ( θ , σ2I ) [ L ( θ̃ ) ] , where σ2 is a fixed variance ( set to 0.01 ) which determines the degree of smoothing . This is the same approach taken in variational optimization ( Staines & Barber , 2012 ) . We can construct two different unbiased gradient estimators for L ( θ ) : one via the reparameterization trick ( Kingma & Welling , 2013 ) , and one via the log-derivative trick similar to what is done in evolutionary strategies ( ES ) ( Wierstra et al. , 2008 ) . We denote the two estimates as grp and ges respectively , grp = Ee∼N ( 0 , I ) [ ∇θL ( θ + eσ ) ] , ges = Eθ̂∼N ( θ , σ2I ) [ L ( θ̃ ) ∇θ [ log ( N ( θ̃ ; θ , σ2 ) ) ] ] . Following the insight from ( Parmas et al. , 2018 ) in the context of reinforcement learning1 , we combine these estimates using inverse variance weighting ( Fleiss , 1993 ) , gmerged = grp/σ −2 rp + ges/σ −2 es σ−2rp + σ −2 es , ( 1 ) where σ2rp and σ 2 es are the empirical variances of grp and ges respectively . When outer-training learned optimizes we find the variances of ges and grp can differ by as many as 20 orders of magnitude ( Figure 5a ) . This merged estimator addresses this by having at most the lowest of the two variances . To further reduce variance , we use antithetic sampling when estimating ges . The cost of computing a single sample of ges and grp is thus 2 forward and backward passes of an unrolled optimization . To compute the empirical variance , we leverage data parallelism to compute multiple samples of ges and grp . In theory , the samples used to evaluate σ2rp , σ 2 es must be independent of ges and grp , but in practice we found good performance using the same samples for both . Finally , an increasing curriculum over steps per truncation is used over the course of outer-training . This introduces bias early in training , but also allows for far more frequent outer-weight updates , resulting in much faster outer-training in terms of wall-clock time . The full outer-training algorithm is described in Appendix B .

Summary: This paper tackles the problem of learning an optimizer, like "learning to learn by gradient descent by gradient descent" and its follow-up papers. Specifically, the authors focus on obtaining cleaner gradients from the unrolled training procedure. To do this, they use a variational optimization formulation and two different gradient estimates: one based on the reparameterization trick and one based on evolutionary strategies. The paper then uses a method from the recent RL literature to combine these two gradient estimates to obtain a variance that is upper-bounded by the minimum of the two gradients' variances.