Reverse engineering learned optimizers reveals known and novel mechanisms

1 Introduction . Optimization algorithms underlie nearly all of modern machine learning ; thus advances in optimization have broad impact . Recent research uses meta-learning to learn new optimization algorithms , by directly parameterizing and training an optimizer on a distribution of tasks . These so-called learned optimizers have been shown to outperform baseline optimizers in restricted settings [ 1–7 ] . Despite improvements in the design , training , and performance of learned optimizers , fundamental questions remain about their behavior . We understand remarkably little about how these optimizers work . Are learned optimizers simply learning a clever combination of known techniques ? Or do they learn fundamentally new behaviors that have not yet been proposed in the optimization literature ? If they did learn a new optimization technique , how would we know ? Contrast this with existing “ hand-designed ” optimizers such as momentum [ 8 ] , AdaGrad [ 9 ] , RMSProp [ 10 ] , or Adam [ 11 ] . These algorithms are motivated and analyzed using intuitive mechanisms and theoretical principles ( such as accumulating update velocity in momentum , or rescaling updates based on gradient magnitudes in RMSProp or Adam ) . This understanding of underlying mechanisms allows future studies to build on these techniques by highlighting flaws in their operation [ 12 ] , ∗Work conducted while at Google Research . Currently at Meta Reality Labs . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . studying convergence [ 13 ] , and developing deeper knowledge about why key mechanisms work [ 14 ] . Without analogous understanding of the inner workings of a learned optimizers , it is incredibly difficult to analyze or synthesize their behavior . In this work , we develop tools for isolating and elucidating mechanisms in nonlinear , highdimensional learned optimization algorithms ( §4 ) . Using these methods we show how learned optimizers utilize both known and novel techniques , across four disparate tasks . In particular , we demonstrate that learned optimizers learn momentum ( §5.1 ) , gradient clipping ( §5.2 ) , learning rate schedules ( §5.3 ) , and methods for learning rate adaptation ( §5.4 , §5.5 ) . Taken together , our work can be seen as part of a new approach to scientifically interpret and understand learned algorithms . We provide code for training and analyzing learned optimizers , as well as the trained weights for the learned optimizers studied here , at https : //bit.ly/3eqgNrH . 2 Related Work . Our work is heavily inspired by recent work using neural networks to parameterize optimizers . Andrychowicz et al . [ 1 ] originally showed promising results on this front , with additional studies improving robustness [ 2 , 3 ] , meta-training [ 6 ] , and generalization [ 7 ] of learned optimizers . Here , we study the behavior of optimizers by treating them as dynamical systems . This perspective has yielded a number of intuitive and theoretical insights [ 15–18 ] . We also build on recent work on reverse engineering recurrent neural networks ( RNNs ) . Sussillo and Barak [ 19 ] showed how linear approximations of nonlinear RNNs can reveal the algorithms used by trained networks to solve simple tasks . These techniques have been applied to understand trained RNNs in a variety of domains , from natural langauge processing [ 20 , 21 ] to neuroscience [ 22 ] . Additional work on treating RNNs as dynamical systems has led to insights into their computational capabilities [ 23–25 ] . 3 Methods . 3.1 Preliminaries . We are interested in optimization problems that minimize a loss function ( f ) over parameters ( x ) . We focus on first-order optimizers , which at iteration k have access to the gradient gki ≡ ∇f ( xki ) and produce an update ∆xki . These are component-wise optimizers that are applied to each parameter ( xi ) of the problem in parallel . Standard optimizers used in machine learning ( e.g . momentum , Adam ) are in this category1 . Going forward , we use x for the parameter to optimize , g for its gradient , k for the current iteration , and drop the parameter index ( i ) to reduce excess notation . One can think of an optimizer as being comprised of two parts : the optimizer state ( h ) that stores information about the current problem , and readout weights ( w ) that are used to update parameters given the current state . An optimization algorithm is specified by an initial state , state transition dynamics , and readout , defined as follows : hk+1 = F ( hk , gk ) ( 1 ) xk+1 = xk + wThk+1 , ( 2 ) where h is the optimizer state , F governs the optimizer dynamics , and w are the readout weights . Learned optimizers are constructed by parameterizing the function F , and then learning those parameters along with the readout weights through meta-optimization ( detailed in App . D.2 ) . Handdesigned optimization algorithms , by distinction , specify these functions at the outset . For example , in momentum , the state variable is a single number ( known as the velocity ) , and is updated as a linear combination of the previous state and the gradient ( e.g . hk+1 = βhk + gk , where β is the momentum hyperparameter ) . For momentum and other hand-designed optimizers , the state variables are low-dimensional , and their dynamics are ( largely ) straightforward . In contrast , learned optimizers have high-dimensional state variables , and the potential for rich , nonlinear dynamics . As these systems learn complex behaviors , it has historically been difficult to extract simple , intuitive descriptions of the behavior of a learned optimizer . 1Notable exceptions include quasi-Newton methods such as L-BFGS [ 26 ] or K-FAC [ 27 ] . 3.2 Training learned optimizers . We parametrize the learned optimizer with a recurrent neural network ( RNN ) , similar to Andrychowicz et al . [ 1 ] . The only input to the optimizer is the gradient . The RNN is trained by minimizing a meta-objective , which we define as the average training loss when optimizing a target problem . See App . D.2 for details about the optimizer architecture and meta-training procedures . Below , we only analyze the final ( best ) trained optimizer , however we do analyze aspects of the meta-training dynamics in App B.2 . We trained learned optimizers on each of four tasks . These tasks were selected because they converge in a relatively small number of iterations ( particularly important for meta-optimization ) and cover a range of loss surfaces ( convex and non-convex , low- and high-dimensional , deterministic and stochastic ) : Linear regression : The first task consists of random linear regression problems f ( x ) = 12‖Ax−b‖ 2 2 , where A and b are randomly sampled . Much of our theoretical understanding of the behavior of optimization algorithms is derived using quadratic loss surfaces such as these , in part because they have a constant Hessian ( ATA ) over the entire parameter space . The choice of how to sample the problem data A and b will generate a particular distribution of Hessians and condition numbers . A histogram of condition numbers for our task distribution is shown in Figure 1a . Rosenbrock : The second task is minimizing the Rosenbrock function [ 28 ] , a commonly used test function for optimization . It is a non-convex function with a curved valley and a single global minimum . The function is defined over two parameters as f ( x , y ) = ( 1 − x ) 2 + 100 ( y − x2 ) 2 ( Figure 1b ) . The distribution of problems for this task consists of different initializations sampled uniformly over a grid . The grid used to sample initializations is the same as the grid shown in the figure ; the x- and y- coordinates are sampled from the ranges ( -2 , 2 ) and ( -1 , 3 ) , respectively . Two moons : The third task involves training a fully connected neural network to classify a toy dataset , the two moons dataset ( Figure 1c ) . As the data are not linearly separable , a nonlinear classifier is required . The optimization problem is to train the weights of a three hidden layer fully connected neural network , with 64 units per layer and tanh nonlinearities ( for a total of 8,577 parameters ) . The distribution of problems involves sampling the initial weights of the fully connected network . MNIST : The fourth task is to train a four layer convolutional network to classify digits from the MNIST dataset . We use a minibatch size of 100 examples ; thus the gradients fed to the optimizer are stochastic , unlike the previous three problems . The network consists of three convolutional layers each with 16 channels with a 3×3 kernel size and ReLU activations , followed by a final ( fully connected ) dense layer , for a total of 82,250 parameters ( Figure 1d ) . We additionally tuned three baseline optimizers ( momentum , RMSProp , and Adam ) individually for each task . We selected the hyperparameters for each problem out of 2500 samples randomly drawn from a grid . Details about the exact grid ranges used for each task are in App . D.3 . Figure 1 ( top row ) compares the performance of the learned optimizer ( blue ) to baseline optimizers ( red , yellow , and orange ) , on each of the four tasks described above . Across all tasks , the learned optimizer outperforms the baseline optimizers on the meta-objective2 ( App . Fig . 17 ) . 4 Tools for understanding optimizers . In order to analyze learned optimizers , we make extensive use of two methods . The first is a way to visualize what an optimizer is doing at a particular optimizer state . The second is a way of making sense of how the optimizer state changes , that is , understanding the optimizer state dynamics . We describe these below , and then use them to analyze optimizer behavior and mechanisms in §5 . 4.1 Update functions First , we introduce a visualization tool to get a handle on what an optimizer is doing . Any optimizer , at a particular state , can be viewed as a scalar function that takes in a gradient ( g ) and returns a change in the parameter ( ∆x ) . We refer to this as the optimizer update function . Mathematically , the update function is computed as the state update projected onto the readout , ∆x = wTF ( h , g ) , following equations ( 1 ) and ( 2 ) . In addition , the slope of this function with respect to the input gradient ( ∂∆x ∂g ) can be thought of as the effective learning rate at a particular optimizer state3 . A steeper slope means that the parameter update is more sensitive to the gradient ( as expected for a higher learning rate ) , and a shallower slope means that parameter updates are smaller for the same gradient magnitude ( i.e . a lower learning rate ) . As the optimizer state varies , the update function and effective learning rate can change . This provides a mechanism for learned optimizers to implement different types of behavior : through optimizer state dynamics that induce particular types of changes in the update function . 2As the meta-objective is the average training loss during an optimization run , it naturally penalizes the training curve earlier in training ( when loss values are large ) . This explains the discrepancy in the curves for linear regression ( Fig . 1a , top ) where momentum continues to decrease the loss late in training . Despite this , the learned optimizer has an overall smaller meta-objective due to having lower loss at earlier iterations . 3We compute this slope at g = 0 , in the middle of the update function . We find that the update function is always affine in the middle with saturation at the extremes , thus the slope at g = 0 is a natural way to summarize the effective learning rate . It is instructive to first visualize update functions for commonly used optimizers ( Figure 2 ) . For gradient descent , the update ( ∆x = −αg ) is stateless and is always a fixed linear function whose slope is the learning rate , α ( Fig . 2a ) . Gradient clipping is also stateless , but is a saturating function of the gradient ( Fig . 2b ) . For momentum , the update is ∆x = βv−αg , where v denotes the momentum state ( velocity ) and β is the momentum timescale . The velocity adds an offset to the update function ( Fig . 2c ) . As the optimizer picks up positive ( or negative ) momentum , the curve shifts downward ( or upward ) , thus incorporating a bias to reduce ( or increase ) the parameter . For adaptive optimizers such as RMSProp , the state variable changes the slope , or effective learning rate , within the linear region of the update function ( Fig . 2d ) . Now , what about learned optimizers , or optimizers with much more complicated or high-dimensional state variables ? One advantage of update functions is that , as scalar functions , they can be easily visualized and compared to the known methods in Figure 2 . Whether or not the underlying hidden states are interpretable , for a given learned optimizer , remains to be seen .

Summary: In this work, the author got trained learned optimizers on four different optimization tasks, and analyzed their behavior. They discovered that learned optimizers learn a plethora of intuitive mechanisms: momentum, gradient clipping, schedules, forms of learning rate adaptation. While the coarse behaviors are qualitatively similar across different tasks, the mechanisms are tuned for particular tasks.