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 proposes a method to learn a neural network to perform optimization. The idea is that the neural network will receive as an input several parameters, including the weights of the network to be trained, the gradient, and so on, and will output new updated weights. The neural network that is used to compute new weights can be trained through a complicated process called un-rolled optimization. The authors of the paper show two problems with this approach. Namely, the gradients tend to explode as the number of iterations increases. Truncating the gradient computation introduces some bias. To solve these problems the authors propose a variational objective that smooths the objective surface. The proposed method is evaluated on the image net dataset showing better results than first order methods optimally optimized.