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 this paper, the authors introduce a technique for multitask Bayesian optimization based on meta learning. In general, there are two broad approaches to multi task learning in Bayesian optimization: sharing information between tasks by modelling the correlation between tasks (e.g., Swersky et al 2013), and sharing information between tasks via global weight sharing across over some embedding (e.g., Perrone et al., 2018). This paper falls in to the latter approach, learning a meta learning style task embedding (eqn. 1) to be able to make predictions over a new task given auxiliary task data (eqn. 2).