Improving Hyperparameter Optimization by Planning Ahead

1 INTRODUCTION . Hyperparameter optimization ( HPO ) is a ubiquitous problem within the research community and an integral aspect of tuning machine learning algorithms to ensure generalization beyond the training data . HPO is often posed as a sequential decision-making process , however , it can be seen as a special use-case of model-based reinforcement learning ( MbRL ) ( Sutton , 1991a ; Henaff et al. , 2017 ) developed under the guise of some idiosyncratic terms . In MbRL the objective is to train a transition model to approximate an underlying transition function via interactions with an environment governed by some policy , e.g . random shooting ( Nagabandi et al. , 2018 ) , to improve sample efficiency and learn from more useful interactions whereby the learned models are used as a simulator for sampling trajectories . An agent navigates the simulated environment to optimize a pre-defined reward function , while the transition model remains unchanged . Conventionally in HPO , a surrogate model is trained to estimate some black-box function , e.g . validation loss of a machine learning algorithm under investigation ( Rasmussen , 2003 ; Snoek et al. , 2015b ; Springenberg et al. , 2016 ) . An acquisition function ( Wilson et al. , 2017b ) interacts with the surrogate model to propose potential hyperparameters that optimize the black-box response , viewed as a reward function . Effectively , the surrogate model is the only unknown component for a transition model , that prevents HPO from being framed fully as MbRL problem . In this paper , we present a novel formulation for HPO defined within the context of MbRL . Namely , we learn an ensemble of probabilistic neural network models ( Lakshminarayanan et al. , 2017 ) and show that using model predictive control ( MPC ) ( Kamthe & Deisenroth , 2018 ) and a novel lookahead variant to navigate the simulated black-box environment , we can outperform conventional Bayesian optimization techniques with heuristic acquisition functions in both transfer and nontransfer learning settings . Thus , we elaborate on the importance of explicit planning in HPO that has been largely overlooked by the community . We also formally define HPO as a Markov decision process ( MDP ) with a simple , yet novel , state representation as the set of previously evaluated hyperparameters and their corresponding responses . We argue that with a clearly defined transition model , we can replace the acquisition function with a simple policy that maximizes the reward across the simulated trajectories , and achieve better results through proper planning , as shown in Figure 1 . LookAhead MPC-5 is our look-ahead strategy with MPC that simulates 5 states ahead . We focus here on the motivation with more details in Section 6 . Our main contributions are summarized as i ) a formal definition of HPO as an MDP which does not depend on any engineered heuristics , ii ) a new transfer learning surrogate model represented by an ensemble of probabilistic neural networks , iii ) a novel acquisition function that implements a look-ahead strategy paired with model predictive control , iv ) clear motivation that highlights the impact of planning in HPO , which to the best of our knowledge , has never been addressed before . 2 RELATED WORK . Hyperparameter optimization has been extensively studied in the community within two main settings : single task ( non-transfer learning ) HPO and transfer learning HPO . For single task HPO , solutions often involve estimating the hyperparameter response surface using a ( probabilistic ) surrogate , such as a Gaussian process ( Rasmussen , 2003 ; Bergstra et al. , 2011 ) , random forests ( Hutter et al. , 2011 ) , Bayesian neural networks ( Springenberg et al. , 2016 ) , or some hybrid approach ( Snoek et al. , 2015a ) . Hyperparameter candidates are then selected via an acquisition function ( Wilson et al. , 2017a ) , e.g . expected improvement , that satisfies some well-motivated assumptions and utilizes the statistics from the posterior to score each hyperparameter . Transfer learning solutions on the other hand leverage available data from experiments on related tasks to expedite HPO . Simple solutions involve initializing the surrogate using hyperparameters that perform well on datasets that have similar meta-features ( Feurer et al. , 2014 ; Jomaa et al. , 2021a ) , i.e . dataset statistics . The response surface can also be modeled jointly in a multi-task setting ( Bardenet et al. , 2013 ; Yogatama & Mann , 2014 ; Perrone et al. , 2018 ) . For example , Salinas et al . ( 2020 ) transforms the hyperparameter response into a similar distribution and learns a shared Gaussian Copula . Another way to achieve transfer learning is through a weighted combination of the surrogates ( Wistuba et al. , 2016 ; Feurer et al. , 2018 ) . More recently , it has been shown that meta-learning an initialization of the surrogate by training to estimate the response across the training tasks , and subsequently fine-tuning to the target task improves generalization and leads to better performance ( Wistuba & Grabocka , 2021 ; Jomaa et al. , 2021b ) . Additionally one can learn a transferable acquisition function ( Wistuba et al. , 2018 ; Jomaa et al. , 2019 ; Volpp et al. , 2020 ) . Jomaa et al . ( 2019 ) propose to learn a policy based on Deep Q-learning ( Watkins & Dayan , 1992 ) where they assume a temporal dependency between the hyperparameters and propose an LSTM ( Hochreiter & Schmidhuber , 1997 ) to model it . Volpp et al . ( 2020 ) propose a transferable acquisition function , MetaBO , as a policy that is meta-trained on related tasks . As another model-free RL approach , MetaBO is trained by interacting with the environment and observes an engineered state representation that is heavily influenced by underlying surrogates . MetaBO however is very sensitive to the number of trials ( Wistuba & Grabocka , 2021 ) . We discuss the proposed MDPs and their relationship to our proposed approach in Section 4 . Model-based reinforcement learning aims to learn a model of the environment by estimating the next state given a state-action pair . It is applied with great success in areas such as robotics and video games . Existing methods similarly involve learning a ( probabilistic ) model of the dynamics ( Chua et al. , 2018 ; Nagabandi et al. , 2018 ; Ko et al. , 2007 ) . Moreover , several methods improve the policy in a Dyna style approach ( Sutton , 1991a ) , by leveraging the generated trajectories from the model to mitigate the distribution shift . 3 PRELIMINARIES . Hyperparameter Optimization : Let D denote the space of all datasets and Λ ∈ RN the space of hyperparameters associated with an unknown black-box function ` ( D ) : Λ → R where ` ( D ) represents the response ` ( D ) ( λ ) , e.g . validation loss , of a certain model under investigation trained on a dataset D ∈ D with hyperparameters λ ∈ Λ . For example , we might be interested in a neural network model where Λ : = R+ × R+0 × N represents the learning rate , dropout rate , and the number of layers , respectively . The objective of HPO is hence to find the optimal hyperparameter configuration λ∗ : = arg minλ∈Λ ` ( D ) ( λ ) given a fixed budget of T trials . HPO is commonly treated as sequential decision-making process , where a surrogate model ˆ̀ : Λ→ R is iteratively fit to the history H ( D ) t : = { λi , ` ( D ) ( λi ) } ti=1 of evaluated hyperparameters and an acquisition function ŵ : ( Λ × R ) ∗ → Λ is used to select the next candidate which minimizes the expected hyperparameter response : arg min ŵ ED∼ρD ` ( D ) ( { λ1 , . . . , λt } ) , ( 1 ) where ` ( D ) ( { λ1 , . . . , λt } ) : = mini∈ { 1 , ... , t } ` ( D ) ( λi ) , λi : = ŵ ( { ( λj , ` ( D ) ( λj ) } i−1j=1 ) and ρD is some distribution over the available datasets . Among the variety of acquisition functions , the expected improvement is widely adopted ( Mockus , 1974 ) . Although HPO algorithms can be applied on both continuous and discrete search spaces , here we focus on discrete spaces because they allow for faster training in transfer learning settings ( Schilling et al. , 2016 ; Jomaa et al. , 2021b ) . Model-based reinforcement learning : MbRL uses interactions with the environment to learn a parameterized approximation of the underlying transition function which subsequently can be used for planning . The learning task can be formulated as a discrete-time Markov decision process ( MDP ) defined by the tuple 〈S , A , τ , r〉 with state space S , action space A , unknown transition function τ : S × A → S and reward function r : S × A → R. The general approach in RL is to learn the optimal policy π maximizing the expected cumulative discounted reward given by J ( π ) : = E [ ∑T t=1 γ tr ( st , at ) | st = τ ( st−1 , at−1 ) , a ∼ π ( st ) ] where st ∈ S and at ∈ A with discount factor γ ∈ [ 0 , 1 ] . We note that the explicit reward function or an approximate reward model is necessary for planning actions under some learned transition model . During inference , the optimal sequence of actions is commonly presented as the solution to the following optimization problem : arg max a1 , ... , aT T−1∑ t=0 r̂ ( τ̂ ( st , at ) , at+1 ) , ( 2 ) which leverages the transition model τ̂ to predict the next state st+1 given a state st and action at . 4 HYPERPARAMETER OPTIMIZATION VIA MODEL-BASED RL . In the following we formulate HPO as an MDP which in turn facilitates the use of a variety of RL approaches to solve the problem . As a sequential decision-making process , hyperparameters λ ∈ Λ are iteratively selected and evaluated forming a simple MDP . Particularly , the state space is defined as S : = ( Λ× R ) ∗ with s ( D ) t : = ( ( λ1 , ` ( D ) 1 ) , . . . , ( λt , ` ( D ) t ) ) ( 3 ) where ` ( D ) i : = ` ( D ) ( λi ) and action spaceA : = Λ . The ground truth transition function τ : S×Λ→ S simply appends the new observations to the previous ones , τ ( s ( D ) t , λ ) : = ( ( λ1 , ` ( D ) 1 ) , . . . , ( λt , ` ( D ) t ) , ( λ , ` ( D ) ( λ ) ) ) ( 4 ) for dataset D ∈ D and a ground truth reward function , r : = S × Λ → R , that returns the loss reduction of the new hyperparameters over the best loss so far , r ( s ( D ) t , λ ) : = max { 0 , min i=1 : t ` ( D ) i −min { ` ( D ) 1 , . . . , ` ( D ) t , ` ( D ) ( λ ) } } ( 5 ) To estimate both , the transition function and the reward function , as done in MbRL , the only missing piece is a model for the validation loss , the so called surrogate model . Given ˆ̀ ( D ) : Λ → R , the transition and reward models are just : τ̂ : S × Λ→ S where τ̂ ( s ( D ) t , λ ) : = ( ( λ1 , ` ( D ) 1 ) , . . . , ( λt , ` ( D ) t ) , ( λ , ˆ̀ ( D ) ( λ ) ) ) ( 6 ) and r̂ : S × Λ→ R with r̂ ( s ( D ) t , λ ) : = max { 0 , min i=1 : t ` ( D ) i −min { ` ( D ) 1 , . . . , ` ( D ) t , ˆ̀ ( D ) ( λ ) } } . ( 7 ) In this sense , HPO using a surrogate model can be seen as a special case of MbRL , and thus all approaches researched for MbRL can be directly applied to HPO . When learning a joint surrogate across several datasets in D , the superscript is dropped and ˆ̀ ( D ) → ˆ̀ . To the best of our knowledge , so far there exist two prior publications which attempt to solve HPO using RL by defining distinct MDPs . Jomaa et al . ( 2019 ) originally define the state as the dynamic history of evaluated hyperparameters and their respective responses as S : = ( Λ × R ) ∗ and s ( D ) t as defined earlier . However , they suggest that the order by which the hyperparameters are selected impacts the decision on which action to take next , which is why they model this temporal aspect via an LSTM in their policy , which is conditioned on some engineered meta-features . Their transition function generates new states by appending the action and the observed response to the previous state according to τ ( s ( D ) t , λ ) : = ( ( λ1 , ` ( D ) 1 ) , . . . , ( λ ( D ) t , ` t ) , ( λ , ` ( D ) ( λ ) ) ) ) . Volpp et al . ( 2020 ) represent the state space as S : = ( RM ) ∗ with st : = ( ( µt ( λi ) , σ2t ( λi ) , λi , ψi ) ) |Λ| i=1 where µt and σ2t are the mean and variance of the posterior distribution of an underlying surrogate and ψi are some engineered attributes . The transition function simply updates the parameters of the surrogate based on the new observations . Both approaches define the action space as the discrete grid of hyperparameters , A : = Λ , whereas the reward function is computed in terms of the regret . In strong contrast to these methods that replace the standard acquisition function with a policy trained via model-free RL approaches , we adopt the powerful yet simple MDP proposed by Jomaa et al . ( 2019 ) to train a transition model that allows planning to improve HPO in the context of MbRL .

Summary: This paper proposes a method for hyperparameter optimization (HO) drawing inspiration from model-based RL (MBRL). They recast the HO problem in a markov decision process (MDP) formulation. This formulation permits the use of MBRL methods to solve the HO problem.   Concretely, having reframed HO using this MDP formulation, they propose to: (1) train an ensemble of probabilistic neural networks to model the transition function (using first order meta learning), which maps a set consisting of previous hyperparam-response value pairs and the current hyperparam setting to the associated response value (eg val loss); (2) use this transition model in Model Predictive Control (MPC) with random shootouts to obtain an effective hyperparameter setting. Given the nature of their MDP and the fact that successive states are independent of one another, they use a variant of MPC that selects the best hyperparameter setting (or action) seen (rather than use the first action, as is common in MPC).  In experiments, they demonstrate their method improves on various baselines on three different settings. They perform some ablation studies showing the importance of finetuning their model on the target domains and the importance of planning.