Reinforced Few-Shot Acquisition Function Learning for Bayesian Optimization

1 Introduction . Bayesian optimization ( BO ) has served as a powerful and popular framework for global optimization in many real-world tasks , such as hyperparameter tuning [ 1–4 ] , robot control [ 5 ] , automatic material design [ 6–8 ] , etc . To search for global optima under a small sampling budget and potentially noisy observations , BO imposes a Gaussian process ( GP ) prior on the unknown black-box function and continually updates the posterior as more samples are collected . BO relies on acquisition functions ( AFs ) to determine the sample location , i.e. , those with larger AF values are prioritized than those with smaller ones . AFs are often designed to capture the trade-off between exploration and exploitation of the global optima . The design of AFs has been extensively studied from various perspectives , such as optimism in the face of uncertainty ( e.g. , GP-UCB [ 9 ] ) , optimizing information-theoretic metrics ( e.g. , entropy search methods [ 10–12 ] ) , and maximizing one-step improvement ( e.g. , expected improvement or EI [ 13 , 14 ] ) . As a result , AFs are often handcrafted according to different perspectives of the trade-off , and the best-performing AFs can vary significantly under different types of black-box 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . functions [ 11 ] . This phenomenon is repeatedly observed in our experiments in Section 4 . Therefore , one critical issue in BO is to design an AF that can adapt to a variety of black-box functions . To achieve better adaptability to new tasks in BO , recent works propose to leverage meta-data , the data previously collected from similar tasks [ 15–17 ] . For example , in the context of hyperparameter optimization under a specific dataset , the meta-data could come from the evaluation of previous hyperparameter configurations for the same learning model over any other related dataset . In [ 15 ] , meta-data is used to fine-tune the initialization of the GP parameters and thereby achieves better GP model selection for each specific task . However , the potential benefit of using meta-data for more efficient exploration via AFs is not explored . On the other hand , in [ 16 , 17 ] , meta-data is split into multiple subsets and then used to construct a transferable acquisition function based on some off-the-shelf AF ( e.g . EI ) and an ensemble of GP models . Each of the GP models is learned over a separate subset of the meta-data . However , to achieve effective knowledge transfer , this approach would require a sufficiently large amount of meta-data , which significantly limits its practical use . As a result , there remains a critical unexplored challenge in BO : how to design an AF that can effectively adapt to a wide variety of back-box functions given only a small amount of meta-data ? To tackle this challenge , we propose to rethink the use of meta-data in BO through the lens of few-shot acquisition function ( FSAF ) learning . Specifically , our goal is to learn an initial AF model that allows few-shot fast adaptation to each specific task during evaluation . Inspired by the similarity between AFs and Q-functions , we use a deep Q-network ( DQN ) as a surrogate differentiable AF , i.e. , the Q-network would output an indicator for each candidate sample point given its posterior mean and variance as well as other related information . Given the parametric nature and the differentiability of a Q-network , it is natural to leverage optimization-based few-shot learning approaches , such as model-agnostic meta-learning ( MAML ) [ 18 ] , for the training of few-shot AFs . Despite this natural connection , we find that a direct combination of standard DQN and MAML ( DQN+MAML ) is prone to overfitting , as illustrated by the comparison of training and testing results . Note that in Figure 1 ( with detailed configuration in Appendix C ) , DQN+MAML achieves better regret than EI on the training set but suffers from much higher regret during testing . We hypothesize this is because both DQN and MAML are prone to overfitting [ 19–22 ] . This issue could be particularly critical in BO due to the uncertainty required in adaptive sampling . Based on our findings and inspired by [ 23 ] , we propose a Bayesian variant of DQN with the following three salient features : ( i ) The Bayesian DQN learns a distribution of parameters of DQN based on the Kullback-Leibler ( KL ) regularization framework . Through this , the problem of minimizing the Q-learning temporal-difference ( TD ) error in DQN is converted into a Bayesian inference problem . ( ii ) For the prior of the Bayesian DQN , we propose to use a demonstration policy induced by an off-the-shelf AF to further stabilize the training process ; ( iii ) We then use the chaser meta-loss in [ 20 ] , which serves as a natural companion to the proposed Bayesian DQN . As shown by the experimental results in Section 4 , the proposed design effectively mitigates overfitting and achieves good generalization under various black-box functions . Moreover , FSAF is trained solely on synthetic GP functions ( with details in Appendix B ) and is able to adapt to a broad variety of functions based only on a small amount of metadata ( which could be costly to generate ) . With the proper design of the Q-networks , FSAF is general-purpose in the sense that it is agnostic to both the dimension and the cardinality of the input domain . The main contributions of this paper can be summarized as follows : • We consider a novel setting of few-shot acquisition function learning for BO and present the first few-shot acquisition function that can use a small amount of meta-data to achieve better task-specific exploration and thereby effectively adapt to a wide variety of black-box functions . • Inspired by the similarity between AFs and Q-functions , we view DQN as a parametric and differentiable AF and use it as the base of our FSAF . We identify the important overfitting issue in the direct combination of DQN and MAML and thereafter present a Bayesian variant of DQN that mitigates overfitting and enjoys stable training through a demo-based prior . • We extensively evaluate the proposed FSAF in a variety of tasks , including optimization benchmark functions , real-world datasets , and synthetic GP functions . We show that the proposed FSAF can indeed effectively adapt to a variety of tasks and outperform both the conventional benchmark AFs as well as the recent state-of-the-art meta-learning BO methods . 2 Preliminaries . Our goal is to design a sampling policy to optimize a black-box function f : X→ R , where X ⊂ Rd denotes the compact domain of f . f is black-box in the sense that there are no special structural properties ( e.g. , concavity or linearity ) or derivatives ( e.g. , gradient or Hessian ) about f available to the sampling policy . In each step t , the policy selects xt ∈ X and obtains a noisy observation yt = f ( xt ) + εt , where εt are i.i.d . zero-mean Gaussian noises . To evaluate a sampling policy , we define the simple regret as Regret ( t ) : = maxx∈X f ( x∗ ) −max1≤s≤t f ( xs ) , which quantifies the best sample up to t. For convenience , we let F t : = { ( xi , yi ) } t−1i=1 denote the observations up to step t. Bayesian optimization . To optimize f in a sample-efficient manner , BO first imposes on the space of objective functions a GP prior , which is fully characterized by a mean function and a covariance function , and then determines the next sample based on the resulting posterior [ 24 , 25 ] . In each step t , given F t , the posterior predictive distribution of each x ∈ X is N ( µt ( x ) , σ2t ( x ) ) , where µt ( x ) : = E [ f ( x ) | F t ] and σt ( x ) : = ( V [ f ( x ) | F t ] ) 1 2 can be derived in closed form . In this way , BO can be viewed as a sequential decision making problem . However , it is typically difficult to obtain the exact optimal policy due to the curse of dimensionality [ 24 ] . To obtain tractable policies , BO algorithms construct AFs Ψ ( x ; F t ) , which resort to maximizing one-step look-ahead objectives based on F t and the posterior [ 24 ] . For example , EI chooses the sample location based on the improvement made by the immediate next sample in expectation , i.e. , ΨEI ( x ; F t ) = E [ f ( x ) −max1≤i≤t−1 f ( xi ) | F t ] , which enjoys a closed form in µt ( x ) , σt ( x ) , and max1≤i≤t−1 f ( xi ) . Meta-learning with few-shot fast adaptation . Meta-learning is a generic paradigm for generalizing the knowledge acquired during training to solving unseen tasks in the testing phase . In the few-shot setting , meta-learning is typically achieved via a bi-level framework : ( i ) On the upper level , the training algorithm is meant to determine a proper initial model with an aim to facilitating subsequent task-specific adaptation ; ( ii ) On the lower level , given the initial model and the task of interest , a fast adaptation subroutine is configured to fine-tune the initial model based on a small amount of task-specific data . Specifically , during training , the learner finds a model parameterized by θ based on a collection of tasks T , where each task τ ∈ T is associated with a training setDtrτ and a validation set Dvalτ . For any initial model parameters θ and training setD tr of a task τ , letM ( θ , Dtrτ ) be an algorithm that outputs the adapted model parameters by applying few-shot fast adaptation to θ based onDtrτ . The performance of the adapted model is evaluated on Dvalτ by a meta-loss function L ( M ( θ , D tr τ ) , D val τ ) . Accordingly , the overall training can be viewed as solving the following optimization problem : θ∗ : = argmin θ ∑ τ∈T L ( M ( θ , Dtrτ ) , D val τ ) . ( 1 ) By properly configuring the loss function , the formulation in ( 1 ) is readily applicable to various learning problems , including supervised learning and RL . Note that the adaptation subroutine is typically chosen as taking one or a few gradient steps with respect to L ( · , · ) or some relevant loss function . For example , under the celebrated MAML [ 18 ] , the adaptation subroutine isM ( θ , Dtrτ ) ≡ θ − η∇θ L ( θ , Dtrτ ) , where η denotes the learning rate . Reinforcement learning and Q-function . Following the conventions of RL , we use st , at , and rt to denote the state , action , and reward obtained at each step t. Let R and γ be the reward function and the discount factor . The goal is to find a stationary randomized policy that maximizes the total expected discounted reward E [ ∑∞ t=0 γ trt ] . To achieve this , given a policy π , a helper function termed Q-function is defined as Qπ ( s , a ) : = E [ ∑∞ t=0 γ tR ( st , at ) |so = s , a0 = a ; π ] . The optimal Q-function can then be defined as Q∗ ( s , a ) : = maxπ Qπ ( s , a ) , for each state-action pair . One fundamental property of the optimal Q-function is the Bellman optimality equation , i.e. , Q∗ ( s , a ) = E [ rt + γmaxa′ Q ∗ ( st+1 , a ′ ) ∣∣st = a , at = a ] . Then , the Bellman optimality operator can be defined by [ BQ ] ( s , a ) : = R ( s , a ) + γ Es′ [ maxa′∈AQ ( s′ , a′ ) ] . It is known that Q∗ is the unique fixed point of B . We leverage this fact to describe the proposed FSAF in Section 3.2 .

Summary: The paper proposes "Reinforced Few-Shot Acquisition Function Learning" (FSAF), a meta-learning approach for designing few-shot acquisition functions (AFs) for Bayesian Optimization (BO).   The authors map BO onto an reinforcement learning (RL) problem, i.e., identify the AF with an RL policy which proposes evaluation points, given the current state of the optimization in form of the Gaussian Process (GP) posterior (and further features). The authors develop a Bayesian version of Deep Q-learning (DQN), leveraging Stein Variational Gradient Descent (SVGD), and combine it with ideas from Bayesian MAML (BMAML) to meta-train an initial AF on function samples from GP priors. This initial AF can be adapted to a family of target optimization problems using a few (~1-5) context optimization problems from this family for fast (few-shot) adaptation.   The method is evaluated on various sets of optimization problems, showing favorable performance compared to existing hand-designed AFs (EI, PI, GP-UCB, MES) as well as the transfer-learning approach MetaBO [1].