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: The authors propose a hyperparameter optimization algorithm in meta-learning, where parameters w/o being involved in the inner loop optimization are treated as hyperparameters. The proposed algorithms approximate the second-order hypergradients via knowledge distillation. They further evaluate the effectiveness on tinyImageNet and CIFAR100.