Online Hyperparameter Meta-Learning with Hypergradient Distillation

1 INTRODUCTION . Meta-learning ( Schmidhuber , 1987 ; Thrun & Pratt , 1998 ) aims to learn a learning process itself over a task distribution . Many gradient-based meta-learning approaches assume a set of parameters that do not participate in inner-optimization ( Lee & Choi , 2018 ; Flennerhag et al. , 2020 ; Raghu et al. , 2020 ) which can be seen as hyperparameters . Those hyperparameters are important in helping the inner-learner converge faster and generalize better . As they are usually very high-dimensional such as element-wise learning rates ( Li et al. , 2017 ) , we can not meta-learn them with simple hyperparameter optimization ( HO ) techniques such as random search ( Bergstra & Bengio , 2012 ) or Bayesian optimization ( Snoek et al. , 2012 ) due to the too extensive search space . In this case , we can use gradient-based HO methods that can directly optimize the highdimensional hyperparameters by minimizing the validation loss w.r.t . the hyperparameters ( Bengio , 2000 ) . Due to the expensive computational cost of evaluating the hypergradients ( i.e . the gradient w.r.t . the hyperparameters ) , there has been a lot of efforts to improve the effectiveness and the efficiency of the algorithms . However , unfortunately , none of the existing algorithms satisfy the following criteria at the same time that should be met for their practical use : 1 ) scalable to hyperparameter dimension , 2 ) online optimization , 3 ) memory-efficient , 4 ) avoid short-horizon bias . Please See Table 1 for the comparison of existing gradient-based HO algorithms in the above four criteria . Forward-Mode Differentiation ( FMD ) ( Franceschi et al. , 2017 ) in Table 1 is an algorithm that forward-propagates Jacobians ( i.e . derivatives of the update function ) from the first to the last step , which is analogous to real-time recurrent learning ( RTRL ) ( Williams & Zipser , 1989 ) in recurrent neural networks . FMD allows online optimization ( i.e . update hyperparameters every inner-step ) with the intermediate Jacobians and also computes the hypergradients over the entire horizon . However , a critical limitation is that the time and space complexity linearly increases w.r.t . the hyperparameter dimension . Thus , we can not use FMD for solving many practical meta-learning problems that come with millions of hyperparameters , which is the main problem we tackle in this paper . Secondly , Reverse-Mode Differentiation ( RMD ) ( Maclaurin et al. , 2015 ) back-propagates the Jacobian-vector products ( JVPs ) from the last to the initial step , which is structurally identical to backprop through time ( BPTT ) ( Werbos , 1990 ) . RMD is scalable to the hyperparameter dimension , but the space complexity linearly increases w.r.t . the horizon length ( i.e. , the number of innergradient steps used to compute the hypergradient ) . It is possible to reduce the memory burden by checkpointing some of the previous weights and further interpolating between the weights to approximate the trajectory ( Fu et al. , 2016 ) . However , RMD and its variants are not scalable for online optimization . This is because they do not retain the intermediate Jacobians unlike FMD and thus need to recompute the whole second-order term for every online HO step . Thirdly , algorithms based on Implicit Function Theorem ( IFT ) are applicable to high-dimensional HO ( Bengio , 2000 ; Pedregosa , 2016 ) . Under the assumption that the main model parameters have arrived at convergence , the best-response Jacobian , i.e . how the converged model parameters change w.r.t . the hyperparameters , can be expressed by only the information available at the last step , such as the inverse of Hessian at convergence . Thus , we do not have to explicitly unroll the previous update steps . Due to the heavy cost of computing inverse-Hessian-vector product , Lorraine et al . ( 2020a ) propose to approximate it by an iterative method , which works well for high-dimensional HO problems . However , still it is not straightforward to use the method for online optimization because of the convergence assumption . That is , computing hypergradients before convergence does not guarantee the quality of the hypergradients . To our knowledge , the short horizon approximation such as one-step lookahead ( 1-step in Table 1 ) ( Luketina et al. , 2016 ) is the only existing method that fully supports online optimization , while being scalable to the hyperparameter dimension at the same time . It computes hypergradients only over a single update step and ignores the past learning trajectory , which is computationally efficient as only a single JVP is computed per each online HO step . However , this approximation suffers from the short horizon bias ( Wu et al. , 2018 ) by definition . In this paper , we propose a novel HO algorithm that can simultaneously satisfy all the aforementioned criteria for practical HO . The key idea is to distill the entire second-order term into a single JVP . As a result , we only need to compute the single JVP for each online HO step , and at the same time the distilled JVP can consider longer horizons than short horizon approximations such as one-step lookahead or first-order method . We summarize the contribution of this paper as follows : • We propose HyperDistill , a novel HO algorithm that satisfies the aforementioned four criteria for practical HO , each of which is crucial for a HO algorithm to be applied to the current meta-learning frameworks . • We show how to efficiently distill the hypergradient second-order term into a single JVP . • We empirically demonstrate that our algorithm converges faster and provides better gener- alization performance at convergence , with three recent meta-learning models and on two benchmark image datasets . 2 RELATED WORK . Hyperparameter optimization When the hyperparameter dimension is small ( e.g . less than 100 ) , random search ( Bergstra & Bengio , 2012 ) or Bayesian optimization ( Snoek et al. , 2012 ) works well . However , when the hyperparameter is high-dimensional , gradient-based HO is often preferred since random or Bayesian search could become infeasible . One of the most well known methods for gradient-based HO are based on Implicit Function Theorem which compute or approximate the inverse Hessian only at convergence . Bengio ( 2000 ) computes the exact inverse Hessian , and Luketina et al . ( 2016 ) approximate the inverse Hessian with the identity matrix , which is identical to the onestep lookahead approximation . Pedregosa ( 2016 ) approximates the inverse Hessian with conjugate gradients ( CG ) method . Lorraine et al . ( 2020b ) propose Neumann approximation , which is numerically more stable than CG approximation . On the other hand , Domke ( 2012 ) proposes unrolled differentiation for solving bi-level optimization , and Shaban et al . ( 2019 ) analyzes the truncated unrolled differentiation , which is computationally more efficient . Unrolled diffrentiation can be further categorized into forward ( FMD ) and reverse mode ( RMD ) ( Franceschi et al. , 2017 ) . FMD is more suitable for optimizing low-dimensional hyperparamters ( Im et al. , 2021 ; Micaelli & Storkey , 2020 ) , but RMD is more scalable to the hyperparameter dimension . Maclaurin et al . ( 2015 ) proposes a more memory-efficient RMD , which reverses the SGD trajectory with momentum . Fu et al . ( 2016 ) further reduce memory burden of RMD by approximating the learning trajectory with linear interpolation . Luketina et al . ( 2016 ) can also be understood as a short horizon approximation of RMD for online optimization . Our method also supports online optimization , but the critical difference is that our algorithm can alleviate the short horizon bias ( Wu et al. , 2018 ) . RMD is basically a type of backpropagation and it is available in deep learning libraries ( Grefenstette et al. , 2020 ) . Meta-learning Meta-learning ( Schmidhuber , 1987 ; Thrun & Pratt , 1998 ) aims to learn a model that generalizes over a distribution of tasks ( Vinyals et al. , 2016 ; Ravi & Larochelle , 2017 ) . While there exists a variety of approaches , in this paper we focus on gradient-based meta-learning ( Finn et al. , 2017b ) , especially the methods with high-dimensional hyperparameters that do not participate in inner-optimization . For instance , there have been many attempts to precondition the innergradients for faster inner-optimization , either by warping the parameter space with every pair of consecutive layers interleaved with a warp layer ( Lee & Choi , 2018 ; Flennerhag et al. , 2020 ) or directly modulating the inner-gradients with diagonal ( Li et al. , 2017 ) or block-diagonal matrix ( Park & Oliva , 2019 ) . Perturbation function is another form of hyperparameters that help the inner-learner generalize better ( Lee et al. , 2020 ; Ryu et al. , 2020 ; Tseng et al. , 2020 ) . It is also possible to let the whole feature extractor be hyperparameters and only adapt the last fully-connected layer ( Raghu et al. , 2020 ) . On the other hand , some of the meta-learning literatures do not assume a task distribution , but tune their hyperparameters with a holdouot validation set , similarly to the conventional HO setting . In this case , the one-step lookahead method ( Luketina et al. , 2016 ) is mostly used for scalable online HO , in context of domain generalization ( Li et al. , 2018 ) , handling class imbalance ( Ren et al. , 2018 ; Shu et al. , 2019 ) , gradient-based neural architecture search ( Liu et al. , 2019 ) , and coefficient of norm-based regularizer ( Balaji et al. , 2018 ) . Although we mainly focus on metalearning setting in this work , whose goal is to transfer knowledge through a task distribution , it is straightforward to apply our method to conventional HO problems . 3 BACKGROUND . In this section , we first introduce RMD and its approximations for efficient computation . We then introduce our novel algorithm that supports high-dimensional online HO over the entire horizon . 3.1 HYPERPARAMETER UNROLLED DIFFERENTIATION . We first introduce notations . Throughout this paper , we will specifiy w as weight and λ as hyperparameter . The series of weights w0 , w1 , w2 . . . , wT evolve with the update function wt = Φ ( wt−1 , λ ; Dt ) over steps t = 1 , . . . , T . The function Φ takes the previous weight wt−1 and the hyperparameter λ as inputs and its form depends on the current mini-batch Dt . Note that w1 , w2 , . . . , wT are functions w.r.t . the hyperparameter λ . The question is how to find a good hyperparameter λ that yields a good response wT at the last step . In gradient-based HO , we find the optimal λ by minimizing the validation loss Lval as a function of λ. min λ Lval ( wT ( λ ) , λ ) ( 1 ) Note that we let the loss function Lval ( · , · ) itself be modulated by λ for generality . According to the chain rule , the hypergradient is decomposed into dLval ( wT , λ ) dλ = ∂Lval ( wT , λ ) ∂λ︸ ︷︷ ︸ gFOT : First-order term + ∂Lval ( wT , λ ) ∂wT dwT dλ︸ ︷︷ ︸ gSOT : Second-order term ( 2 ) On the right hand side , the first-order ( FO ) term gFOT directly computes the gradient w.r.t λ by fixing wT . The second-order ( SO ) term gSOT computes the indirect effect of λ through the response wT . αT = ∂Lval ( wT , λ ) ∂wT can be easily computed similarly to gFOT , but the response Jacobian dwT dλ is more computationally challenging as it is unrolled into the following form . dwT dλ = T∑ t=1 ( T∏ s=t+1 As ) Bt , where As = ∂Φ ( ws−1 , λ ; Ds ) ∂ws−1 , Bt = ∂Φ ( wt−1 , λ ; Dt ) ∂λ ( 3 ) Algorithm 1 Reverse-HG ( RMD ) 1 : Input : The last weight wT and all the previous weights w0 , . . . , wT−1 . 2 : Output : Hypergradient gFO + gSO . 3 : α← ∂L ( wT , λ ) ∂wT , gFO ← ∂L ( wT , λ ) ∂λ , gSO ← 0 4 : for t = T downto 1 do 5 : gSO ← gSO + αBt 6 : α← αAt 7 : end for 8 : return gFO + gSO Algorithm 2 DrMAD ( Fu et al. , 2016 ) 1 : Input : The last weight wT and the initial weight w0 . 2 : Output : Approximated hypergradient gFO + gSO . 3 : α← ∂L ( wT , λ ) ∂wT , gFO ← ∂L ( wT , λ ) ∂λ , gSO ← 0 4 : for t = T downto 1 do 5 : ŵt−1 ← ( 1− t−1 T ) w0 + t−1 T wT 6 : gSO ← gSO + αB̂t 7 : α← αÂt 8 : end for 9 : return gFO + gSO Eq . ( 3 ) involves the Jacobians { A } and { B } at the intermediate steps . Evaluating them or their vector products are computationally expensive in terms of either time ( FMD ) or space ( FMD , RMD ) ( Franceschi et al. , 2017 ) . Therefore , how to approximate Eq . ( 3 ) is the key to developing an efficient and effective HO algorithm .

Summary: This paper develops a practical gradient-based hyperparameter optimization method, HyperDistill, that meets the following criteria   a) scalability in hyperparameter dimension and memory constraints,    b) accuracy (hyper-gradient update terms do not depend on only the last step of gradient updates)   c) applicability to the online setting.  The main difficulty lies in estimating the gradient of the weights with respect to the hyperparameters. The authors do so by approximating it as a single Jacobian-vector product, using a "distilled" weight and dataset pair.  Experimentally, HyperDistill achieves better validation losses and higher quality true hypergradient estimates than a variety of recent, relevant baselines, and does in an efficient fashion.