Probabilistic Meta-Learning for Bayesian Optimization

Transfer and meta-learning algorithms leverage evaluations on related tasks in order to significantly speed up learning or optimization on a new problem . For applications that depend on uncertainty estimates , e.g. , in Bayesian optimization , recent probabilistic approaches have shown good performance at test time , but either scale poorly with the number of data points or under-perform with little data on the test task . In this paper , we propose a novel approach to probabilistic transfer learning that uses a generative model for the underlying data distribution and simultaneously learns a latent feature distribution to represent unknown task properties . To enable fast and accurate inference at test-time , we introduce a novel meta-loss that structures the latent space to match the prior used for inference . Together , these contributions ensure that our probabilistic model exhibits high sample-efficiency and provides well-calibrated uncertainty estimates . We evaluate the proposed approach and compare its performance to probabilistic models from the literature on a set of Bayesian optimization transfer-learning tasks . 1 INTRODUCTION . Bayesian optimization ( BO ) is arguably one of the most proven and widely used blackbox optimization frameworks for expensive functions ( Shahriari et al. , 2015 ) with applications that include materials design ( Frazier & Wang , 2016 ) , reinforcement learning ( Metzen et al. , 2015 ) , and automated machine learning ( ML ) ( Hutter et al. , 2019 ) . In practical applications , BO is repeatedly used to solve variations of similar tasks . In these cases , the sample efficiency can be further increased by not starting the optimization from scratch , but rather leveraging previous runs to inform and accelerate the latest one . Several approaches to this emerged under the name of transfer-learning ( Weiss et al. , 2016 ) and meta-learning ( Vanschoren , 2018 ) . Compared to early work by Swersky et al . ( 2013 ) ; Golovin et al . ( 2017 ) , recent publications leverage the representative flexibility of neural networks , which allows for more powerful models and impressive results ( Gordon et al. , 2019 ; Rusu et al. , 2019 ; Garnelo et al. , 2018b ; a ; Zintgraf et al. , 2019 ) . Despite these significant advances , only a small subset of algorithms offers the well-calibrated uncertainty estimates on which BO relies to guide its sampling strategy efficiently . Additionally , BO benefits greatly from a meaningful prior over tasks that quickly converges to the true function to provide the highest sample efficiency . Existing work mostly focuses on deterministic models and , for those providing uncertainty estimates , sample-efficiency at test time is often a challenge . Contributions We set out to close this gap and introduce BAyesian optimization with Neural Networks and Embedding Reasoning ( BaNNER ) , a flexible meta-learning method for BO . We go beyond previous work of Perrone et al . ( 2018 ) and introduce a generative regression model explicitly conditioned on a low-dimensional latent representation for the tasks . This allows our model to ( i ) encode a meaningful prior over tasks and ( ii ) remain highly sample-efficient , since each new task only requires inference over a low-dimensional latent representation . To ensure robust training of our model , we introduce a novel loss function to regularize the latent distribution and optimize our model ’ s hyper-parameters using the available meta-data . We evaluate BaNNER on a set of synthetic benchmarks and two meta-learning problems and compare with the state-of-the-art in the literature . 2 PROBLEM STATEMENT AND BACKGROUND . Our goal is to efficiently optimize an unknown function f ( x , τ ) over a domain x ∈ X for some unknown but fixed task parameters τ that are sampled from an unknown distribution τ ∼ p ( T ) . To this end , at each iteration n we can select function parameters xn and observe a noisy function value yn = f ( xn , τ ) + n , with n drawn i.i.d . from some distribution p . While our method can handle arbitrary noise distributions , we assume a Gaussian distribution , i.e . p = N ( 0 , σ2 ) , for the remainder of the paper . We assume that each evaluation of f is expensive either in terms of monetary cost or time , so that we want to minimize the number of evaluations of f during the optimization process . The most data-efficient class of algorithms for this setting are Bayesian optimization ( BO ) algorithms , which use the observations collected up to iteration n , Dn = { xi , yi } n−1i=1 , in order to infer a posterior belief over the function values f ( x , τ ) . To select parameters xn that are informative about the optimum maxx f ( x , τ ) , BO algorithms define an acquisition function α ( · ) that uses the posterior belief to select parameters as xn = argmaxx α ( p ( f ( x , τ ) | Dn ) . While BO algorithms have been studied extensively , their performance crucially depends on the properties of the statistical model used for f . The two key requirements for BO algorithms to be data-efficient are i ) , that the prior belief over f concentrates quickly on the true function f as we observe data in Dn and ii ) , that the posterior uncertainty estimates are calibrated , so that the model always considers the true function f to be statistically plausible . The latter requirement means that the true function f ( · , τ ) must always be contained in the model ’ s confidence intervals with high probability . Since the task parameters τ are unknown , in general this requires a conservative model that works well for all possible tasks τ . We propose to use meta-learning in order to learn an effective prior ( Fig . 1a ) that can quickly adapt to a new task τ ( Fig . 1b ) . We are given data from T previous tasks τt ∼ p ( T ) with Nt observations Dmetat = { ( xn , t , yn , t } Nt n=1 each . We show the resulting generative model on the left in Fig . 2 . Metalearning aims to distill the information in Dmeta into a model gθ by optimizing the meta-parameters θ . At test time , we then keep these parameters fixed and use the learned model gθ to speed up the optimization of the new function f ( · , τ ) . 3 RELATED WORK . There are several approaches to improve the sample efficiency of BO methods based on information from related tasks . We refer to ( Vanschoren , 2018 ) for a broad review of meta-learning in the context of automated machine learning and focus on the most relevant approaches below . One strategy to improve sample-efficiency is to initialize the BO algorithm with high-quality query points . These initial configurations can be either constructed to be complementary ( Feurer et al. , 2014 ; 2015 ; Lindauer & Hutter , 2018 ) or learned based on data set features ( Kim et al. , 2017 ) . An alternative strategy is to transfer knowledge between tasks by adapting the acquisition function . These approaches balance exploration and exploitation by weighting the usefulness of a given x not only on the current task , but also on the related tasks . This trade-off can either be heuristically motivated ( Wistuba et al. , 2015 ; Wistuba et al. , 2018 ) , or learned directly ( Volpp et al. , 2020 ) . All of these approaches change the BO algorithm , but not the underlying probabilistic model . Most related to our work are methods that focus on alternative ways to model the observations . Several methods build a global Gaussian process ( GP ) models across all tasks ( Swersky et al. , 2013 ; Golovin et al. , 2017 ; Marco et al. , 2017 ; Feurer et al. , 2018 ; Law et al. , 2019 ) , which requires approximations due to the cubic scaling of GPs . A method to avoid this scaling is to use a shared neural network to learn specialized features for Bayesian Linear Regression ( BLR ) . The method , dubbed adaptive BLR ( ABLR ) ( Perrone et al. , 2018 ) , learns a new BLR model for each task based on the shared neural network features . This allows ABLR to quickly adapt to new tasks and scale better than GPs , but limits its predictive power for a small number of test data points . We consider ABLR to be the state-of-the-art for our setting and use it as a baseline in our experiments . The problem of adapting a model to a new problem is not unique to BO , but arises in other fields too . Recent progress in transfer learning ( Weiss et al. , 2016 ) and meta-learning ( Vanschoren , 2018 ) is concerned with the adaptation of a ML model from one or multiple related tasks to a new one . This transfer often focuses on data-efficiency , which makes them similar to BO . While Frameworks for general ( probabilistic ) model adaption ( Finn et al. , 2017 ; 2018 ) exist , our contribution relates to a subset of approaches that directly model latent task distributions ( Gordon et al. , 2019 ; Rusu et al. , 2019 ; Garnelo et al. , 2018b ; a ; Zintgraf et al. , 2019 ) . All of them model the relationship between tasks using the task specific variables , latent variables , and a model that adapts to new tasks by using these latent variables ( often as an input ) . Most of these meta-learning algorithms do not consider active learning as an application and focus on deterministic models gθ . In the following , we propose a method to meta-learn a probabilistic prior model that we can use together with approximate inference in order to estimate a posterior distribution for BO . 4 PROBABILISTIC META-LEARNING FOR BAYESIAN OPTIMIZATION . In this section , we discuss how to approximate f ( x , τ ) with a generative model based on the metadata Dmeta and how to use this model to infer predictions for BO on a new task , given test-data Dn . We use a neural network gθ ( x , z ) to approximate the unknown function f ( x , τ ) . Next to the function inputs x , the network depends on trainable meta-parameters θ shared between all tasks , and an encoding z ∈ Z ⊆ Rd of the task-specific parameters τ . We refer to the latter as task embeddings in the following , since they represent unknown task parameters τ . In addition to being used as additional inputs to the network g , the latent parameter space Z can also include task-dependent network parameters , e.g. , weights of the output layer . Since we do not know τ and p ( T ) , we can not match z to τ directly . Instead , we assume a fixed and known prior distribution p ( Z ) to turn gθ into a generative network . Without loss of generality , we focus on a Gaussian prior p ( Z ) = N ( 0 , I ) in the following . The resulting approximative generative model is shown on the right in Fig . 2 . Meta training The goal of meta-learning is to select appropriate global parameters θ so that for each task τ ∈ T there exists a corresponding latent encoding z ∈ Z with f ( x , τ ) = gθ ( x , z ) . Since we have no information about τ , we train a separate task embedding zt for each task τt by optimizing them jointly with the parameters θ . We can view the resulting model as a variational autoencoder , where we directly optimize over the outputs zt of the encoder , without committing to a specific parametric form . We can optimize the model ’ s predictive performance on the meta-data by maximizing the log-likelihood L ( xn , t , yn , t ; θ , zt ) = log p ( yn , t | gθ ( xn , t , zt ) ) of the parameters and observations for each of the T meta-tasks together with the corresponding task embedding : max θ , z1 , ... , zT T∑ t=1 1 Nt ∑ ( x , y ) ∈Dmetat L ( x , y ; θ , zt ) − λR ( { z1 , . . . , zT } || p ( Z ) ) , ( 1 ) where we rescale the likelihood for each meta-task by Nt , the number of data points per task , to account for imbalanced data which skews the loss towards tasks with many evaluations . Without regularization , this model could easily overfit to each task by placing the task embeddings zt in disjoint areas of the domain . In that case , samples from the Gaussian prior p ( Z ) would lead to function samples with low probability-mass under p ( T ) . To avoid this , we introduce a regularization termR scaled by an appropriate constant λ ∈ R > 0 in order to get a more uniform loss approximation over Z . We discuss the specific choice that we use to enable reliable and efficient inference below . Inference Given the meta-trained generative model , we can make predictions about a new task f ( · , τ ) . That is , after n ≥ 0 noisy observations of function values in Dn , we can infer a posterior belief over the task embeddings , p ( z | Dn , θ ) . While this is generally intractable analytically , approximate inference methods are generally reliable enough to make this tractable . For example , we use Hamiltonian Monte Carlo ( HMC ) ( Neal , 2011 ) in our experiments . Given samples from this posterior belief , we obtain a posterior belief over function values f ( x , τ ) from a monte-carlo approximation of p ( gθ ( x ) | Dn , x , θ ) = ∫ gθ ( x , z ) p ( z | Dn , θ ) ∂z = Ez∼p ( z|Dn , θ ) [ gθ ( x , z ) ] . ( 2 ) We can view this posterior inference as a task-specific adaptation of the parameters z in order to match g to f . In contrast to deterministic meta-learning approaches , here the adaptation does not happen directly through an optimization process such as gradient descent , but through approximate Bayesian inference . This enables us to naturally capture uncertainty depending on the amount of data in Dn that we condition on . Regularizing the Latent space We can only expect the inference over the predictive distribution in ( 2 ) to work well if the two generative models in Fig . 2 encode similar distributions over functions ; that is , if f ( x , τ ) with τ ∼ p ( T ) has a similar distribution as gθ ( x , z ) with z ∼ p ( Z ) for all x ∈ X . This generally requires each sample z ∼ p ( Z ) to be associated with a meaningful function gθ ( x , z ) . In this paper , we achieve this through regularizing zt to be close to the prior distribution p ( Z ) and regularizing the network gθ to be smooth . Regularizing p ( Z ) encourages a meaningful prior over functions where each sample has probability mass under p ( T ) , while the regularization of g ensures that samples from the posterior p ( gθ ( x ) | Dn , x , θ ) vary smoothly across Z . To regularize the task embeddings one might be tempted to use the log-likelihood of the prior p ( Z ) . However , Tolstikhin et al . ( 2018 ) show that this can lead to poor inference results , since a Gaussian prior regularizes ‖zt‖2 , rather than covering the probability mass of p ( Z ) uniformly . Instead , Ghosh et al . ( 2020 ) propose to use a deterministic regularizer during training and use density estimation to determine the prior p ( Z ) . We propose a novel regularizer that is both deterministic ( does not require stochastic approximations ) and regularizes the empirical distribution of zt to a Gaussian prior p ( Z ) = N ( 0 , I ) in a tractable way . We are inspired by the Kolmogorov-Smirnov test for one-dimensional distributions , which we use to compare the empirical cumulative distribution function ( CDF ) of the elements [ zt ] i in zt with the marginal CDFs over the dimensions of the prior p ( Z ) . We show an example in Fig . 5 in the appendix . The empirical CDF is defined as F ( z , d ) = 1T ∑T t=1 I ( [ zt ] d ≤ z ) , where I ( [ zt ] i ≤ z ) is the indicator function that returns one if the ith component of zt is smaller or equal than z and zero otherwise . In addition to the marginals , we account for correlations by regularizing the empirical covariance matrix Cov ( { z1 , . . . , zT } ) to be close to that of p ( Z ) , which is the identity matrix in our Algorithm 1 Probabilistic Meta-learning for Bayesian Optimization 1 : Given meta-data Dmetat for tasks t = 1 , . . . , T 2 : θ ← meta-train by minimizing loss ( 1 ) on the metadata 3 : New , unknown task τ ∼ P ( τ ) , D1 ← ∅ 4 : for iteration n = 1 , . . . do 5 : Approximate predictive distribution p ( gθ ( x ) | D , x , θ ) in ( 2 ) 6 : xn = argmaxα ( p ( gθ ( x ) | x , θ , Dn ) ) . Optimize BO acquisition function 7 : Dn+1 ← Dn ∪ { ( xn , f ( xn , τ ) + n ) } setting . The resulting regularizer , R ( { z1 , . . . , zT } || p ( Z ) ) = d∑ i=1 ( F ( [ zt ] i ) − Φ ( [ zt ] i ) ) 2︸ ︷︷ ︸ Match marginal CDF of p ( Z ) + λc‖I− Cov ( { z1 , . . . , zT } ) ‖2F︸ ︷︷ ︸ Match second moment of p ( Z ) , ( 3 ) trades off the loss for the empirical , marginal CDF of the task embeddings relative to the one for the second moment through a scaling parameter λc . Unlike the original Kolmogorov-Smirnov test , ( 3 ) uses the average distance between the CDFs at the points zt . We found this to lead to more stable training than the original formulation which uses the maximum over Z . In practice , many different tests for assessing multivariate normality exist ( Korkmax et al. , 2014 ) and could be used instead of ( 3 ) . For example , Marida ’ s test ( Mardia , 1970 ) also considers higher-order moments . However , we found it to be too computationally expensive relative to the cheap O ( dT ( log ( T ) + d ) complexity of ( 3 ) and the latter was sufficient to achieve high performance in our experiments . While the log-likelihood term L in ( 1 ) encourages a good fit for each meta-task and the regularization R in ( 3 ) forces the meta-task embeddings conform with the prior p ( Z ) , we also have to ensure that gθ interpolates between the different task embeddings smoothly in order to obtain meaningful functions gθ ( · , z ) for z ∼ p ( z ) that are not part of the task-embeddings during training . A simple trick to enforce smoothness is to add noise to the embeddings during training , which is equivalent to regularizing the Hessian ∂2gθ ( x , z ) / ∂z∂z for Gaussian likelihoods p ( Webb , 1994 ; Bishop , 1995 ; An , 1996 ) . Bayesian optimization algorithms can directly use the predictive distribution ( 2 ) in the acquisition function α in order to select informative parameters xn . Most acquisition functions only depend on the mean and variance of the predictions or samples from the posterior , all of which can be directly computed from ( 2 ) . The overall algorithm is summarized in Algorithm 1 : We first use the meta-data to minimize the regularized meta-loss in Line 2 . In Line 3 we get a new task τ and start without any test data D1 = ∅ . After that , we proceed iteratively and approximate the predictive distribution , select new parameters to use by maximizing BO ’ s acquisition function , evaluate test-task f ( xn , τ ) , and add the data point to our test data set . Selecting hyperparameters Like most machine learning methods , BaNNER depends on hyperparameters that include the regularization constants λ , λc , parameters of the BO acquisition function α , the amount of noise to add for regularization , training-specific hyperparameters like learning-rates and batch-sizes , and inference-specific hyperparameters such as the length and number of the chains that we sample from . In practice , these have to be selected based on the meta-data only , since we typically can not generate additional meta-training data easily or cheaply . For our experiments , we split the metadata into train and validation sets and select hyperparameters by comparing the average log likelihood of the trained models . For BO tasks , we split each validation task randomly into data points that we condition on in ( 2 ) and points that we use to evaluate the likelihood . To efficiently optimize all these parameters we use BOHB ( Falkner et al. , 2018 ) , a highly parallel framework to optimize hyperparameters . To speed up computation , we use BOHB ’ s multi-fidelity capabilities and scale the number of meta-training iterations and validation tasks used to compute the validation loss with the fidelity . This allows us to quickly find promising regions of the hyperparameter search space and focus the compute resources there .

Summary: In the paper, a probabilistic method is proposed for Bayesian optimization transfer learning (or meta-learning). A latent representation is modeled separately for each task. The surrogate used in BO will be approximated by a neural network with its input being the test point and the latent representation of the task. The meta-training procedure optimizes the prediction loss over the neural network weight and the latent representations of the source tasks. Additional regularizer is added to induce smooth latent representations. For a new task, predictions are made on the new data point using the surrogate model g_\theta(x,z) with a fully Bayesian treatment. The p(z|D,\theta) is approximated by Hamiltonian Monte Carlo. In the experiments, the proposed method is compared to random search, GP-BO, and adaptive Bayesian linear regression.