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: The authors devise tools for investigating the behaviour of learned, gradient-based, optimization algorithms in order to understand better how they achieve their empirically validated high performance.   To this end, they visualize the update function, i.e., the change in parameters in dependence of the gradient input, which allows to identify mechanisms employed also by hand-designed optimizers, such as gradient clipping, momentum, and learning rate adaptation. The authors demonstrate that this visualization is particularly insightful when analyzing the update rule using the linearized optimizer state dynamics around fixed points. Furthermore, the authors identify that learned optimizers encode learning rate schedules through autonomous state trajectories.   Using their tools, the authors analyze one existing type of learned optimizer (the RNN-based approach proposed by Andrychowicz et al., "Learning to Learn by Gradient Descent by Gradient Descent") on four relatively simple problems (linear regression, Rosenbrock, two moons, MNIST). They find that (depending on the problem) the learned optimizer can exhibit well-known mechanisms like gradient clipping, momentum, learning rate adaptation, and learning rate schedules. Furthermore, they provide clues that the optimizer can learn parameter-type (bias/weights) and layer-specific learning-rate schedules in neural-network based problems.