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 presents novel theoretical results on learning a step size for vanilla GD and SGD by unrolling the optimization steps and back-propagating, taking into account the simple problem minimizing quadratic functions and mean-square errors. The authors could demonstrate the occurrence of already-detected phenomena for learned optimizers, such as gradient explosion/vanishing and over-fitting, in the particular studied case. A few experiments illustrate what the developed theory predicts.