A Study of Bayesian Neural Network Surrogates for Bayesian Optimization

Extensive Experiments

Our paper includes a diverse set of experiments showing how many different aspects of Bayesian optimization interact with our results. For each of our surrogate models, we include experiments in the appendix where we vary their hyperparameters, such as the number of models in a deep ensemble or the kernel for a GP, and measure their impact on performance. Additionally, we also include experiments showing the impact of different BNN architectures, and we also show the effects of the batch size of the acquisition function. While it is always possible to request additional experiments, we believe that we have done an exhaustive and extensive exploration of BNNs for Bayesian optimization.

How would recent methods such as PFNs4BO rank?

Thank you for the reference! Our paper focuses on the performance of Bayesian neural networks trained from scratch, and consequently, we find methods reliant on pre-training to be outside the scope of our work. The inclusion of methods not trained from scratch makes an apples-to-apples comparison especially difficult, since many of the other neural net surrogates could also be pre-trained in some way. However, we are excited about methods that leverage auxiliary data, which we believe may be a notable advantage for using neural network surrogates for Bayesian optimization. We have included a reference to PFNs4BO in the updated manuscript and included it in our discussions.

We appreciate your review, and we respond to your questions as follows.

Importance of insights

The choice of surrogate model in Bayesian optimization is a fundamental design decision that has been largely overlooked in the literature, and many papers default to using standard Gaussian process surrogates (e.g., with Matern or RBF kernels). Our paper is the first to provide a comprehensive study of BNN surrogates, and the wide range of neural network-based surrogates explored in our submission allows us to evaluate the role of representation learning, non-stationarity, and stochasticity in modeling Bayesian optimization objectives. These findings offer pragmatic insights for BO practitioners by providing guidance on selecting appropriate surrogate models.

Our paper also has many novel and important findings, in some cases even prompting us to re-think conventional wisdom in BO. Surprising and unexpected findings include the following: (1) There are limited empirical benefits of stochasticity in the surrogate, given the competitive performance of DKL to BNN models; (2) the performance of deep ensemble surrogates is surprisingly poor, despite their success in other non-BO contexts; (3) I-BNN demonstrate compelling performance on high dimensional benchmarks; (4) Different BO objectives exhibit highly different salient structures (e.g., in contrast to many vision and NLP problems). We additionally note that we are, to the best of our knowledge, the first to consider infinite BNN and HMC surrogates, and I-BNNs have strikingly good results in the high dimensional settings. While not every result would be expected to be surprising, many of the results in our submission are both surprising and informative.

In this paper, we do not try to develop new methods for specific Bayesian optimization applications and note that in order to reach state-of-the-art performance in drug discovery, materials engineering, or some applied problem of that nature, the paper would need to be entirely about that single applied problem and the idiosyncrasies associated with it. Our focus is instead on a general foundational scientific study on the effectiveness of BNN surrogate models on a wide range of problems, covering a wide array of standard benchmarks, with various different properties. We have further outlined the contributions of the paper in the separate general post.

Number of Datasets

We respectfully disagree that we consider a "very small number of datasets". We benchmarked our results on 15 different datasets, which is more than many other related Bayesian optimization papers:

• 8 datasets: Promises and Pitfalls of the Linearized Laplace in Bayesian Optimization, Kristiadi et al. 2023
• 5 datasets: Scalable Global Optimization via Local Bayesian Optimization, Eriksson et al. 2019
• 8 datasets: Scalable Bayesian Optimization Using Deep Neural Networks, Snoek et al. 2015
• 3 datasets: Maximizing Acquisition Functions for Bayesian Optimization, Wilson et al. 2018
• 9 datasets: Bayesian Optimization with High-Dimensional Outputs, Maddox et al. 2021
• 11 datasets: Bayesian Optimization with Robust Bayesian Neural Networks, Stefan et al. 2016

What happens if I regularize the DKL?

We explored this question in the submitted manuscript. In Appendix D.7, we include results for DKL as suggested in [2], showing the performance of optimizing the parameters with the marginal likelihood (ML) compared to the conditional marginal likelihood (CML). We see that there is no clear preference for using the ML or the CML, and the behavior of DKL with ML and CML is very similar across many of the objectives.

Runtime of Bayesian Optimization

We provide wall-clock times of the surrogate models across our experiments in Appendix D.11. However, we would like to emphasize that Bayesian optimization applications often include expensive computation querying of the objective function, which may include actions such as synthesizing a new material, training a large neural network to convergence, etc. For these scenarios, the quality of the surrogate model uncertainties may be much more important than the cost of inference. That said, all of our runtimes are available, and we are not advocating for one method over another in general. Practitioners can use the runtimes we provide as a guide in making a decision for what matters most to them in the context of their problem. We also note that I-BNNs, a novel surrogate model in our paper, has both relatively fast runtime, and strong results.

Edit: See new comment below! We have provided a new experiment in Appendix D.11 to address this point.

In response to your comment on runtime, we have now updated Appendix D.11 with a new experiment comparing the performance of the surrogate models within a time budget, under the assumption of very fast function queries. We find that in this setting, GPs and I-BNNs outperform other BNNs which have slower inference times. However, we would like to re-emphasize that this scenario is often not representative of many common Bayesian optimization problems, which instead have very expensive function queries which would dominate the time budget.

We hope we have addressed your concerns, and please let us know if you have any remaining questions.

Hello, the rebuttal period will end in about 12 hours! Please let us know if you have any further questions we can address.

5 trials is too little

With five trials we do see sufficient evidence to draw many informative conclusions. However, we agree more trials would be helpful and will run all of our experiments for a minimum of ten trials in the final paper to further differentiate the performance between surrogate models.

Although Figure 6 has heavily overlapping confidence intervals, running more trials may not decrease the width of the intervals. This figure plots the distribution of the relative performance of each surrogate model across all trials and all objective functions. For example, we see that for at least 25% of the trials, the relative performance of GPs is less than 0.5, meaning the maximum reward found by GPs for that trial was closer to the reward found by the worst surrogate than the reward found by the best surrogate. We would expect this ratio to stay the same even as we increase the number of trials.

Thank you again for the thoughtful feedback, which we believe has improved the paper. We have added substantial new experiments inspired by your questions, and would appreciate if you could consider raising your score in light of our response.

We really appreciate your review and thoughtful comments. Below we respond to your questions, and include several substantial experiments inspired by your comments. We also have a general response, summarizing our contributions.

Question 1: GP vs HMC experimental details

It is a fair point that our process for optimizing HMC differed from our process for setting GP hyperparameters. Indeed we were interested in understanding the potential of BNNs when the hyperparameters are well-specified for a particular problem.

Inspired by your comments, we have now conducted additional experiments for HMC where we optimize the hyperparameters on a per-iteration basis, similar to the procedure for GPs. Due to time constraints, we have focused on a subset of the objective functions, but we will include a full set of results across all problems in the final paper. Specifically, we fix the width and depth of the network across all experiments, and we conduct a small grid-search after each iteration to optimize the prior and noise hyperparameters based on the validation likelihood, and we include more details about the experiment setup and results in Appendix D.12.

Overall, we find the performance of per-iteration optimized HMC to be competitive with per-trial optimized HMC for problems like Hartmann and Cell Coverage. Furthermore, although per-trial optimization for HMC does improve its performance on Ackley, we are still able to observe similar trends as before.

Furthermore, to make the HMC results reported in the submission more comparable to the GP results, we have also conducted additional experiments where we optimize GP hyperparameters on a per-trial basis. Specifically, we use grid search to select the parameters of the prior distribution of the length scale and output scale for a GP, and for each objective, we choose the parameters which lead to the highest maximum reward over one trial.

We see that these optimized GPs are often comparable to the default GP and do not see significant improvements. This may be due to the GPs being optimized over one specific trial of Bayesian optimization, and the optimal GP may vary depending on the set of initial points. This shows that the observations of the submission hold for both the default GP as well as the meta-optimized version.

We thank you again for the question, and we believe the added experiments, inspired by your comments, will provide significant additional value to the paper.

Question 2: Do GPs still reign supreme?

Our results show that GPs with a standard Matern kernel works well for lower-dimensional problems and GPs with the I-BNN kernel work well in higher dimensions.

Indeed, GPs encompass a broad class of models. GPs with I-BNN kernels exhibit very different behavior from standard GPs with a Matern or RBF kernel. The infinite-width BNN has a non-Euclidean and non-stationary similarity metric, which enables its relative success in high dimensions. DKL based GPs are also significantly different, with their own properties, such as representation learning.

For this reason, we have been careful to distinguish "standard GPs"---GPs with the kernels almost always used in Bayesian optimization---from other model classes, as GPs allowing for any kernel are so general it becomes difficult to make informative conclusions or comparisons. Indeed, the line can get blurry---I-BNNs are a particular type of GP derived from a BNN.

We also note that we have several nuanced take-aways, rather than a single conclusion, and are careful not to proclaim standard GPs or BNNs as overall winners or losers. There is no desire to show that BNNs are generally superior, and we are in fact careful to highlight the broad appeal of standard GP surrogates. Our goal is to understand how various properties of the surrogate (stochasticity, representation learning, etc.) interact with performance in different settings, which transcend the GP vs. BNN distinction, though many of our considered surrogates are NN-inspired (HMC, SGHMC, Deep Ensembles, I-BNN, DKL). If we define BNNs narrowly, and GPs broadly, it may be reasonable to say that GPs have an edge, but we are most concerned with the low-level property distinctions of the surrogates.

Thank you for your thoughtful and supportive comments! We really appreciate it.

Discussion of non-Bayesian neural networks as the surrogate in BO

Thank you for the references! While our paper focuses on Bayesian neural network surrogates for Bayesian optimization, for a thorough investigation, we recognize that there are other works focusing on non-Bayesian methods such as neural bandits and NTK. We have updated our related works section with the suggested papers, and we have also updated Appendix D.3 to include a discussion of the performance of deep ensembles in other works.

Deep ensembles do not do well in exploration

As you pointed out in your review, we noticed that deep ensembles often plateau at low objectives. We explore this phenomenon in Appendix D.3, and showed that smaller training sizes can often lead to deep ensembles having less uncertainty. In the low-data regime, we show that the loss landscape is relatively smooth, making it more difficult to find diverse solutions, corresponding to distinct basins of attraction, through only re-initializing optimization runs. The lack of model diversity in the ensemble leads to smaller predictive uncertainty estimates and worse exploration. This makes deep ensembles less suitable surrogate models when there is a minimal number of data points. Interestingly, this is in contrast to the success of deep ensembles in general applications. More training data (i.e., objective queries) helps deep ensembles again become competitive.

Refer to specific sections of the Appendix

Thank you for the suggestion! We have updated our draft accordingly.