Bayesian Optimization with a Neural Network Meta-learned on Synthetic Data Only
Abstract
Bayesian Optimization (BO) is an effective approach to optimize black-box functions, relying on a probabilistic surrogate to model the response surface.
In this work, we propose to use a Prior-data Fitted Network (PFN) as a cheap and flexible surrogate.
PFNs are neural networks that approximate the Posterior Predictive Distribution (PPD) in a single forward-pass.
Most importantly, they can approximate the PPD for any prior distribution that we can sample from efficiently.
Additionally, we show what is required for PFNs to be used in a standard BO setting with common acquisition functions.
We evaluated the performance of a PFN surrogate for Hyperparameter optimization (HPO), a major application of BO.
While the method can still fail for some search spaces, we fare comparable or better than the state-of-the-art on the HPO-B and PD1 benchmark.
Introduction
Gaussian processes (GPs) are today's de facto standard surrogate model in Bayesian Optimization (BO) [7].
This dominance can be attributed to both their strong performance and their mathematical convenience.
The set of priors that GPs can model is restricted, though.
A GP can only model priors that can be encoded as a valid kernel function and, as the name says, only normally-distributed priors.
Moreover, kernel hyperparameters are typically not treated in a Bayesian manner, likely due to the high cost of running Markov Chain Monte Carlo for this full Bayesian view, which was already shown to yield strong results a decade ago [1, 16].
Instead, GPs are usually fitted with maximum likelihood in most BO packages; this procedure, often referred to as “Empirical Bayes” makes Bayesian Optimization less principled from a Bayesian point of view than it could be.
The recently proposed Prior-data Fitted Networks (PFNs, 14) show that fast approximate Bayesian inference is possible by training a neural network to mimic the posterior predictive distribution (PPD).
This is a powerful approach, as it makes approximate Bayesian inference readily usable in novel applications and allows using any prior that we can sample from.
E.g., it has been applied to perform
Bayesian inference on tabular data [10] with a prior over different neural architectures and their weights, generalizing the usual notion of Bayesian deep learning to not only quantify uncertainty over the weights of a fixed architecture but over the joint space of architectures and weights.
In this work, we show how to use PFNs as a more flexible, Bayesian, drop-in replacement for Gaussian
Processes for real-world Bayesian optimization problems.
We show that PFNs can be applied to
Bayesian optimization and find important pre-processing steps to make them a strong BO surrogate.
In our experiments, we focus on BO for hyperparameter optimization (HPO 6).
This is an important task to bring machine learning algorithms to top performance and a major BO application [3, 16, 9].
Bayesian Optimization with PFNs
Prior-data Fitted Networks (PFNs, 14) are neural networks (meta-)trained to perform approximate
Bayesian prediction.
That is, PFNs are trained to predict some output y ∈ R conditioned on an input x ∈ R k and a training set of given input-output examples D
= { ( x i , y i ) } Ni =1 .
The PFN is trained on this task with samples obtained from a prior distribution p , i.e., { ( x 1 , y 1 ) , ..., ( x N , y N ) , ( x , y ) } ∼ p ( D ) .
The algorithm thus is solely trained on artificial data.
The loss function for training a PFN q θ with parameters θ is the cross entropy for predicting the hold-out example's label
[Equation]
After training on artificial datasets drawn from the prior p ( D ) , we use the trained PFN to approximate the PPD of this prior, based on which we compute the acquisition values.
Last, we maximize the acquisition values for new inputs to find the configuration to be queried next.
We contrast Bayesian optimization with PFNs and GPs (with Empirical Bayes) in Algorithm 1 in the
Appendix.
The major difference is that PFNs have an up-front cost of fitting the PFN once, while
Empirical Bayes incurs the online cost of fitting the hyperparameters in each iteration.
We emphasize that PFNs incur the training cost exactly once per prior, and that a single trained PFN can be used for BO on different tasks with different dimensions, just like the code base for GP regression is a one-time investment that is then generally applicable.
2.1 A HEBO-inspired Prior for BO
HEBO [4] is a state-of-the-art BO method that won the NeurIPS blackbox competition [20] and demonstrated excellent performance in a recent comparison [5].
Thus, we build a prior based on the one used by HEBO.
Our HEBO-inspired prior considers parameters ϕ
= { outputscale , lengthscale , noise , Kumaraswamy parameters } modelled like random variables subject to some distribution p ( ϕ ; ψ ) that depend on hyperparameters ψ .
To sample functions, we perform the following four steps per dataset: i)
Sample x from a uniform distribution, ii) sample the parameters ϕ ∼ p ( ϕ ; ψ ) , iii) transform x using the Kumaraswamy CDF, iv) draw the outputs for our dataset y ∼ N (0 , K ( x ; ϕ )) , where K is the covariance matrix defined by the 3/2 Mat ´ ern kernel, which uses the sampled outputscale , lengthscale and noise .
For additional details on ψ see
Appendix B.
2.2 PFN Adaptions for BO
Acquisition Functions When we condition a PFN on new observations, we get an approximation of the PPD in a forward-pass p ( y | x , D )
≈ q θ ( y | x , D ) in the form of a Riemann distribution, i.e., a distribution over bins spanning a reasonable output range.
The Riemann distribution [14] allows us to calculate the utility exactly for different acquisition functions, e.g., EI, PI or UCB.
To exemplify the general approach, we outline how to compute the PI ( f ∗ ) =
(cid:82) ∞
−∞ [ y > f ∗ ] p ( y ) dy for the unbounded
Riemann distribution in Appendix D. We experimented with EI, PI, UCB and the HEBO acquisition function ensemble [4] and found that simple EI works well enough and is stable.
Inference Strategies We found that two more tricks improve the performance of our PFNs considerably.
i)
We use a power transform to transform the observed outputs to a distribution more similar to a standard-normal, like proposed by HEBO [4].
ii)
We found that simple min-max input normalization between 0 and 1 works well.
In many search spaces, dimensions that naturally lie on a logarithmic scale are not correctly specified, though; this e.g., inspired input warping [18].
We found that we can improve performance in some scenarios by computing the PPD for an ensemble of d + 1 members, where one is normalized like above and for all others we apply a logarithm to a single dimension first.
We also found that ensembling over different permutations of the features can improve performance, thus we do it when we do not otherwise ensemble.
PFN Training Procedure We only train a single PFN for one particular prior, even though the search spaces can have different numbers of features.
To do this, we sample datasets with different numbers of dimensions during PFN training.
Following [10], we zero-pad and linearly scale the features when the number of features k is smaller than the maximum number of features K by Kk to make sure the magnitude of the inputs is similar across different numbers of features.
We used the original
PFNs settings for training [14]: an embedding size of 512 and six layers, using Adam [12] and
[Figure]
# Features # T r a i n i ng
E x a m p l e s
[Figure]
# Features  10 Approx.
Mean Regret
Emp. Bay.
[Figure]
110 PFN 0.
0.124 # Wins
Emp. Bay.
169 PFN  171 Mean Rank
Emp. Bay.
1.501 PFN 1.
cosine-annealing [13] without any special tricks.
The learning rate was chosen based on a simple grid search for minimal training loss.
Results
We present results on three problems: optimizing functions sampled from the HEBO prior, HPO on the HPO-B [15] benchmark and HPO on the PD1 [22] benchmark.
We only consider the discrete setting, but plan to scale the results to continuous spaces in the future.
We consider a budget ofevaluations and repeat experiments for five seeds per task for the real-world experiments.
3.1 Optimizing Prior Functions
We first study whether PFNs can do BO out of the box with as few confounders as possible, by evaluating them on functions sampled from the prior that they were trained on.
We first verified that the PFN predictions match the true GP posterior without Empirical Bayes, where it is computable; as shown in Figure 5 in the appendix.
This is not the case when training the PFN on too few (100k) prior samples but the approximation becomes very tight with enough samples (we evaluated training on up to 80M samples, but performance was already stable with 20M).
Next, we consider the HEBO prior, following their hyper-prior setup as close as possible.
We do not tune the hyper-prior hyperparameters, e.g. the distribution over the lengthscale, in this experiment to have a fair comparison.
We evaluate the HEBO posterior based on our reimplementation of the
HEBO model in GPyTorch with Empirical Bayes [8].
We stayed as close as possible to HEBO, but had to make some adaptions as one cannot sample from the HEBO model as it is.
Specifically, we had to introduce simple priors for the hyperparameters of both the kernels used in HEBO (Mat´ern and linear kernel), as they did not have a prior attached.
We introduce uniform priors ( U (0 , 1) ) for the lengthscale of both the Mat´ern and the Linear kernel and the variance of the linear kernel.
Next, we trained a PFN using the HEBO prior with distributions for the hyperparameters (see Section 2.1) in order to enable PFNs to perform approximate MCMC over the hyperpriors.
We compare the resulting PFN’s fit vs. the Empirical Bayes approximation, as used in HEBO.
Figure 1 shows that the likelihood assigned to held-out outputs is higher for the PFN in many cases; PFNs also become better with more features.
We hypothesize that this is because HEBO’s empirical Bayes approximation becomes too greedy in high dimensions.
Next, we perform BO on datasets sampled from the prior.
Here, we have to introduce the acquisition function as a confounding factor, which might favor one or the other method.
We chose expected improvement, since it is a default choice in BO, but are aware that this might have an impact on our results.
For every feature size we sample 1 000 different functions, and for each functionuniform random points in [0 . 0 , 1 . 0] .
Across dimensions, the GP with empirical Bayes and our PFN approximation worked very similar for Bayesian optimization.
To assess our method qualitatively, we provide example query points of optimizations in Figures 4 and 5 in the appendix.
Number
[Figure]
A v e r a g e
R a nk     10 − 1 .10 −
Number of trials
A v e r a g e
R e g r e t
DGP
BOHAMIANN
DNGO
Random
GP
HEBO
HEBO3 (
[Figure]
no
[Figure]
design)
PFN
PFN (
[Figure]
ensemble)
3.2 Hyperparameter Optimization
Finally, we conduct a large-scale evaluation of PFNs as a surrogate for HPO.
We perform experiments on the large-scale benchmark HPO-B [15], which contains 16 search spaces on 85 datasets.
On
HPO-B, we compare our method to the following baselines: 1) Random Search (RS) [2], 2)
Gaussian Processes (GP) [11, 16] 3) DNGO [17] 4) BOHAMIANN [19], 5)
Deep Kernel Gaussian
Processes [23], 6) HEBO [4] and 7) a variant of HEBO that only uses EI and does not have initial design sampling, like our method.
For this experiment, we used a model with the prior described in Section 2.1.
To decide on the prior hyperparameters ψ , we ran a random search over 50 configurations, picked the model that yielded the lowest regret in a subset of validation search spaces and report results on a separate set of test search spaces.
We specify details on this split in Appendix A. Figure 2 reports results for the average rank and average regret across all the datasets from the test search spaces, showing the PFN to attain competitive results compared to the baselines.
Moreover, an ensemble of PFNs with logged dimensions outperforms HEBO, the state-the-of-the-art in HPO.
However, we would like to mention that the performance rankings are very inconsistent across search spaces, and that all methods, including the PFN, had some validation search spaces for which they performed very poorly; see Figure 6 in the appendix for examples.
We also evaluated our PFN for the very different image classification tuning benchmarks in PD1 [22]; the results in Appendix E show that in this case it performs very similarly to HEBO.
Conclusion & Limitations
We showed that PFNs can be adapted to perform as efficient BO surrogates.
Still, we believe that there is a lot of room to improve upon our work.
So far, this work only considers two real-world benchmarks and no continuous benchmark, as we did not manage to obtain strong results in a continuous optimization setting yet.
This will therefore be the subject of our focus next.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] or [N/A] .
You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data are proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] See our experiments in Section 3.
(b) Did you describe the limitations of your work?
[Yes] We do have limitations.
We only evaluate on discrete benchmarks and we see that our method can utterly fail.
These are described in the experiments in Section 3 and the conclusion in Section (c) Did you discuss any potential negative societal impacts of your work?
[No] We do not see any likely bad impact on society (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] We are releasing the code once we are advancing this to a full conference submission.
So far, it is still heavily under development.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] For the main results we provide 95% confidence intervals.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] The total amount of compute required to train each of the two PFNs used in this paper was 8 GPU-days on RTX 2080 Ti GPUs.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] Both benchmarks we use are cited, as well as the GP framework, we use.
(b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A]
Input : A prior distribution over datasets p ( D ) , hyperparameter prior settings , initial observations D
= { ( x 1 , y 1 ) , ..., ( x K , y K ) } , search space X , number of BO iterations
K , black-box function to optimize f , acquisition function α
Output : Best observed configuration x ∗ and response value y ∗
Train neural network q θ by minimizing objective function in Equation 1 (once, offline) for i ← 1 to N do
Fit GP model ˆ f to data D
Suggest next candidate x ∈ arg max ˆ x ∈X α (ˆ x , D, q θ or ˆ f )
Update history with response D ← D ∪ { ( x , f ( x )) } end
Return top performing configuration: arg max ( x i ,y i ) ∈
D y
i
A Search spaces split
We split search spaces so that we have similar algorithms within the same group, therefore procuring minimal overlap between search spaces.
For choosing the prior with the validation dataset, we use 5527, 5891, 5906, 5971, 6767, 6766, 5860.
The test search spaces are 5860, 5906, 5965, 5970, 5971, 6766, 6767, 6794.
B Details on the HEBO-inspired Prior
For our final model, not the one used for the direct comparison with HEBO on artificial data, we made some adaptions to the HEBO prior.
We do not use a linear kernel and use the following priors, found by through BO on the validation search spaces of HPO-B, as described in
Section 3.2.
We use Gamma ( concentration, rate ) distributions.
For the outputscale we use
Gamma (0 . 8452 , 0 . 3993) and for the lengthscale we use Gamma (1 . 2107 , 1 . 5212) .
For the noise we follow the original kernel exactly with log( ϵ ) ∼ N ( − 4 .
The warping was performed with log normal distributions for its concentrations, s.t. log( concentration0 ) ∼ N (0 , 0 . 5939) and log( concentration1 ) ∼ N (0 , 0 . 9722) .
C Training Details
We only train a single PFN for one particular prior, even though the search spaces can have different numbers of features.
To do this, we sample datasets with different numbers of dimensions during PFN training.
Following [10], we zero-pad and linearly scale the features when the number of features k is smaller than the maximum number of features K by Kk to make sure the magnitude of the inputs is similar across different numbers of features.
Our PFNs were trained in the standard PFN settings.
We use an embedding size of 512 and six layers.
Our models were trained with Adam [12] and cosine-annealing [13] without any special tricks.
The lr was chosen based on simple grid searches for minimal training loss.
D Riemman Distribution and Acqusition Function
We outline how to compute the PI ( f ∗ ) =
(cid:82) ∞
−∞ [ y > f ∗ ] p ( y ) dy for the unbounded Riemann distribution.
[Figure]
Random
PFN
HEBO
[Figure]
[Equation]
(3) =
[Equation]
F
[Equation]
( y
[Equation]
f
[Equation]
+
M (cid:88) i =
[Equation]
(cid:40)
[Equation] [Equation]
i , if y i < f ∗ <
y i
+
[ y i
[Equation]
f
[Equation]
] p ( b i ) , else (cid:41) + F
[Equation]
( f
[Equation]
y M
[Equation]
1 ) (4) =
(1 − F
[Equation]
( y
[Equation]
f
[Equation]
(cid:88)
[Equation]
( y
[Equation]
min ( y
[Equation]
+
[Equation]
max ( f
[Equation]
))) p ( b i )
[Equation]
(5) where F l ( F r ) is the CDF of the half-normal distribution used for the left (right) side.
The acquisition function is divided into three terms: (1) a term that governs the probability mass for values lower than interval b 1 and that goes up to values of y 1 , (2) a term that summarizes over all intervals b 1 , . . .
, b M , and (3) a term that governs the probability mass for values larger than the interval b M and that start from values of y M + 1 .
E Evaluation on PD
We evaluated on all vision tasks of PD1, which amounts to a total of 17 tasks in the same search space.
We make this choice, as these share their evaluation metric: accuracy.
For each task we ranseeds with a different (shared) single initial given evaluation.
In Figure 3, we can see that even on a completely different benchmark the PFN is very competetive with HEBO.
[Figure]
− 0 .
[Figure] [Figure]
[Figure] [Figure]
[Figure]