Guarantees for Tuning the Step Size using a Learning-to-Learn Approach

Learning-to-learn—using optimization algorithms to learn a new optimizer—has successfully trained efficient optimizers in practice . This approach relies on metagradient descent on a meta-objective based on the trajectory that the optimizer generates . However , there were few theoretical guarantees on how to avoid metagradient explosion/vanishing problems , or how to train an optimizer with good generalization performance . In this paper we study the learning-to-learn approach on a simple problem of tuning the step size for quadratic loss . Our results show that although there is a way to design the meta-objective so that the metagradient remain polynomially bounded , computing the meta-gradient directly using backpropagation leads to numerical issues that look similar to gradient explosion/vanishing problems . We also characterize when it is necessary to compute the meta-objective on a separate validation set instead of the original training set . Finally , we verify our results empirically and show that a similar phenomenon appears even for more complicated learned optimizers parametrized by neural networks . 1 INTRODUCTION . Choosing the right optimization algorithm and related hyper-parameters is important for training a deep neural network . Recently , a series of works ( e.g. , Andrychowicz et al . ( 2016 ) ; Wichrowska et al . ( 2017 ) ) proposed to use learning algorithms to find a better optimizer . These papers use a learning-to-learn approach : they design a class of possible optimizers ( often parametrized by a neural network ) , and then optimize the parameters of the optimizer ( later referred to as metaparameters ) to achieve better performance . We refer to the optimization of the optimizer as the meta optimization problem , and the application of the learned optimizer as the inner optimization problem . The learning-to-learn approach solves the meta optimization problem by defining a metaobjective function based on the trajectory that the inner-optimizer generates , and then using backpropagation to compute the meta-gradient ( Franceschi et al. , 2017 ) . Although the learning-to-learn approach has shown empirical success , there are very few theoretical guarantees for learned optimizers . In particular , since the optimization for meta-parameters is usually a nonconvex problem , does it have bad local optimal solutions ? Current ways of optimizing meta-parameters rely on unrolling the trajectory of the inner-optimizer , which is very expensive and often lead to exploding/vanishing gradient problems . Is there a way to alleviate these problems ? Can we have a provable way of designing meta-objective to make sure that the inner optimizers can achieve good generalization performance ? In this paper we answer some of these problems in a simple setting , where we use the learningto-learn approach to tune the step size of the standard gradient descent/stochastic gradient descent algorithm . We will see that even in this simple setting , many of the challenges still remain and we can get better learned optimizers by choosing the right meta-objective function . Though our results are proved only in the simple setting , we empirically verify the results using complicated learned optimizers with neural network parametrizations . 1.1 CHALLENGES OF LEARNING-TO-LEARN APPROACH AND OUR RESULTS . Metz et al . ( 2019 ) highlighted several challenges in the meta-optimization for learning-to-learn approach . First , they observed that the optimal parameters for the learned optimizer ( or even just the step size for gradient descent ) can depend on the number of training steps t of the inner-optimization problem , which is also observed by Wu et al . ( 2018 ) . Ge et al . ( 2019 ) theoretically proved this in a least-squares setting . Because of this , one needs to ensure that the inner training has enough number of steps ( similar to the number of steps that it would take when we apply the learned optimizer ) . However , when the number of steps is large , the meta-gradient can often explode or vanish , which makes it difficult to solve the meta-optimization problem . Our first result shows that this is still true in the case of tuning step size for gradient descent on a simple quadratic objective . In this setting , we show that there is a unique local and global minimizer for the step size , and we also give a simple way to get rid of the gradient explosion/vanishing problem . Theorem 1 ( Informal ) . For tuning the step size of gradient descent on a quadratic objective , if the meta-objective is the loss of the last iteration , then the meta-gradient can explode/vanish . If the meta-objective is the log of the loss of the last iteration , then the meta-gradient is polynomially bounded . Further , doing meta-gradient descent with a meta step size of 1/ √ k ( where k is the number of meta-gradient steps ) provably converges to the optimal step size for the inner-optimizer . Surprisingly , even though taking the log of the objective solves the gradient explosion/vanishing problem , one can not simply implement such an algorithm using auto-differentiation tools such as those used in TensorFlow ( Abadi et al. , 2016 ) . The reason is that even though the meta-gradient is polynomially bounded , if we compute the meta-gradient using the standard back-propagation algorithm , the meta-gradient will be the ratio of two exponentially large/small numbers , which causes numerical issues . Detailed discussion for the first result appears in Section 3 ( Theorem 3 and Theorem 4 ) . The generalization performance of the learned optimizer is another challenge . If one just tries to optimize the performance of the learned optimizer on the training set ( we refer to this as the trainby-train approach ) , then the learned optimizer might overfit . Metz et al . ( 2019 ) proposed to use a train-by-validation approach instead , where the meta-objective is defined to be the performance of the learned optimizer on a separate validation set . Our second result considers a simple least squares setting where y = 〈w∗ , x〉+ ξ and ξ ∼ N ( 0 , σ2 ) . We show that when the number of samples is small and the noise is large , it is important to use train-by-validation ; while when the number of samples is much larger train-by-train can also learn a good optimizer . Theorem 2 ( Informal ) . For a simple least squares problem in d dimensions , if the number of samples n is a constant fraction of d ( e.g. , d/2 ) , and the samples have large noise , then the train-by-train approach performs much worse than train-by-validation . On the other hand , when number of samples n is large , train-by-train can get close to error dσ2/n , which is optimal . We discuss the details in Section 4 ( Theorem 5 and Theorem 6 ) . In Section 5 we show that such observations also hold empirically for more complicated learned optimizers—an optimizer parametrized by neural network . 1.2 RELATED WORK . Learning-to-learn for supervised learning Hochreiter et al . ( 2001 ) introduced the application of gradient descent method to meta-learning . The idea of using a neural network to parametrize an optimizer started in Andrychowicz et al . ( 2016 ) , which used an LSTM to directly learn the update rule . Before that , the idea of using optimization to tune parameters for optimzers also appeared in Maclaurin et al . ( 2015 ) . Later , Li & Malik ( 2016 ) ; Bello et al . ( 2017 ) applied techniques from reinforcement learning to learn an optimizer . Wichrowska et al . ( 2017 ) used a hierarchical RNN as the optimizer . Metz et al . ( 2019 ) adopted a small MLP as the optimizer and used dynamic weighting of two gradient estimators to stabilize and speedup the meta-training process . Learning-to-learn in other settings Ravi & Larochelle ( 2016 ) used LSTM as a meta-learner to learn the update rule for training neural networks in the few-shot learning setting , Wang et al . ( 2016 ) learned an RL algorithm by another meta-learning RL algorithm , and Duan et al . ( 2016 ) learned a general-purpose RNN that can adapt to different RL tasks . Gradient-based meta-learning Finn et al . ( 2017 ) proposed Model-Agnostic Meta-Learning ( MAML ) where they parameterize the update rule for network parameters and learn a shared initialization for the optimizer using the tasks sampled from some distribution . Subsequent works generalized or improved MAML , e.g. , Rusu et al . ( 2018 ) learned a low-dimensional latent representation for gradient-based meta-learning , and Li et al . ( 2017 ) enabled the concurrent learning of learning rate and update direction . Chen et al . ( 2020 ) studied a model with an optimization solver stacked on another neural component . They computed Rademacher complexity of the model , but didn ’ t give any optimization guarantee or study train-by-train versus train-by-validation . Learning assisted algorithms design Similar ideas can also be extended to develop a metaalgorithm selecting an algorithm from a family of parametrized algorithms . Gupta & Roughgarden ( 2017 ) first modeled the algorithm-selection process as a statistical learning problem and bounded the number of tasks it takes to tune a step size for gradient descent . However , they didn ’ t consider the meta-optimization problem . Based on Gupta & Roughgarden ( 2017 ) , people have developed and analyzed the meta-algorithms in many problems ( Balcan et al. , 2016 ; 2018a ; c ; b ; Denevi et al. , 2018 ; Alabi et al. , 2019 ; Denevi et al. , 2019 ) Tuning step size/step size schedule for SGD Shamir & Zhang ( 2013 ) showed that SGD with polynomial step size scheduling can almost match the minimax rate in convex non-smooth settings , which was later tightened by Harvey et al . ( 2018 ) for standard step size scheduling . Assuming that the horizon T is known to the algorithm , the information-theoretically optimal bound in convex nonsmooth setting was later achieved by Jain et al . ( 2019 ) which used another step size schedule , and Ge et al . ( 2019 ) showed that exponentially decaying step size scheduling can achieve near optimal rate for least squares regression . There are also a line of works that investigate methods which adapt a vector of step sizes ( Sutton , 1992 ; Schraudolph , 1999 ; Kearney et al. , 2018 ; Günther et al. , 2019 ; Jacobsen et al. , 2019 ) . 2 PRELIMINARIES . In this section , we first introduce some notations , then formulate the learning-to-learn framework . 2.1 NOTATIONS . For any integer n , we use [ n ] to denote { 1 , 2 , · · · , n } . We use ‖·‖ to denote the ` 2 norm for a vector and the spectral norm for a matrix . We use 〈· , ·〉 to denote the inner product of two vectors . For a symmetric matrix A ∈ Rd×d , we denote its eigenvalues as λ1 ( A ) ≥ · · · ≥ λd ( A ) . We denote the d-dimensional identity matrix as Id . We also denote the identity matrix simply as I when the dimension is clear from the context . We use O ( · ) , Ω ( · ) , Θ ( · ) to hide constant factor dependencies . We use poly ( · ) to represent a polynomial on the relevant parameters with constant degree . We say an event happens with high probability if it happens with probability 1− c for small constant c . 2.2 LEARNING-TO-LEARN FRAMEWORK . We consider the learning-to-learn approach applied to training a distribution of learning tasks . Each task is specified by a tuple ( D , Strain , Svalid , ` ) . Here D is a distribution of samples in X × Y , where X is the domain for the sample and Y is the domain for the label/value . The sets Strain and Svalid are samples generated independently from D , which serve as the training and validation set ( the validation set is optional ) . The learning task looks to find a parameter w ∈ W that minimizes the loss function ` ( w , x , y ) : W × X × Y → R , which gives the loss of the parameter w for sample ( x , y ) . The training loss for this task is f̂ ( w ) : = 1|Strain| ∑ ( x , y ) ∈Strain ` ( w , x , y ) , while the population loss is f ( w ) : = E ( x , y ) ∼D [ ` ( w , x , y ) ] . The goal of inner-optimization is to minimize the population loss f ( w ) . For the learned optimizer , we consider it as an update rule u ( · ) on weight w. The update rule is a parameterized function that maps the weight at step τ and its history to the step τ + 1 : wτ+1 = u ( wτ , ∇f̂ ( wτ ) , ∇f̂ ( wτ−1 ) , · · · ; θ ) . In most parts of this paper , we consider the update rule u as gradient descent mapping with step size as the trainable parameter ( here θ = η which is the step size for gradient descent ) . That is , uη ( w ) = w − η∇f̂ ( w ) for gradient descent and uη ( w ) = w − η∇w ` ( w , x , y ) for stochastic gradient descent where ( x , y ) is a sample randomly chosen from the training set Strain . In the outer ( meta ) level , we consider a distribution T of tasks . For each task P ∼ T , we can define a meta-loss function ∆ ( θ , P ) . The meta-loss function measures the performance of the optimizer on this learning task . The meta objective , for example , can be chosen as the target training loss f̂ at the last iteration ( train-by-train ) , or the loss on the validation set ( train-by-validation ) . The training loss for the meta-level is the average of the meta-loss across m different specific tasks P1 , P2 , ... , Pm , that is , F̂ ( θ ) = 1m ∑m i=1 ∆ ( θ , Pk ) . The population loss for the meta-level is the expectation over all the possible specific tasks F ( θ ) = EP∼T [ ∆ ( θ , P ) ] . In order to train an optimizer by gradient descent , we need to compute the gradient of meta-objective F̂ in terms of meta parameters θ . The meta parameter is updated once after applying the optimizer on the inner objective t times to generate the trajectory w0 , w1 , ... , wt . The meta-gradient is then computed by unrolling the optimization process and back-propagating through the t applications of the optimizer . As we will see later , this unroll procedure is costly and can introduce meta-gradient explosion/vanishing problems .

Summary: This paper considers algorithms that attempt to learn learning rates for gradient descent by gradient descent. Analysis is provided for a few specific quadratic losses showing that the gradient with respect to the learning rate may explode or vanish, and taking the logarithm is suggested to mitigate this. Further results suggest that implementing the gradient of the log comes with interesting numerical difficulties as *intermediate results* might explode or vanish even if the final answer does not.