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: The paper presents a novel transfer learning approach to hyperparameter optimization (HPO) by formulating the problem in a model-based reinforcement learning (MbRL) framework. In this setting, the transition function serves as a surrogate model for learning the validation loss of any black-box ML model which is trained via an ensemble of probabilistic neural networks. The authors further introduce a lookahead search strategy to sample trajectories for a model predictive control problem. Experiments demonstrate that the proposed method outperforms several competitors on a few meta-datasets for HPO tasks.