PaperlayerView exact revision on arXiv

Source-linked AI summary

Bayesian Optimization with Conformal Prediction Sets

Samuel Stanton, Wesley Maddox, Andrew Gordon Wilson

arXiv:2210.12496v4cs.LGstat.ML

TL;DR

Bayesian optimization can make unreliable query choices when its surrogate is misspecified or query points differ from training data, undermining expected utility as a guide. The paper introduces conformal Bayesian optimization to direct queries toward regions with conformal coverage guarantees while retaining a principled distinction between epistemic and aleatoric uncertainty. Across synthetic black-box and real tabular ranking tasks, it improves query coverage and has superior sample-efficiency under misspecification, with comparable efficiency otherwise.

  • Problem

    Bayesian optimization can select queries with unreliable actual utility when its surrogate is misspecified, and feedback covariate shift makes heldout training accuracy unrepresentative of likely queries.

    Source: §1 · p. 1

  • Method

    Conformal BayesOpt uses a conformal Bayes posterior and generalized acquisition functions to direct queries toward regions with coverage guarantees while distinguishing epistemic from aleatoric uncertainty.

    Source: §1 · pp. 1–2

  • Results

    On synthetic black-box and real tabular ranking tasks, conformal BayesOpt improves query coverage, with superior sample-efficiency under surrogate misspecification and comparable efficiency otherwise.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The approach combines conformal coverage guarantees with Bayesian optimization while retaining a principled distinction between epistemic and aleatoric uncertainty.

    Source: §1 · p. 2

    Active online BayesOpt data are not IID, so the stated conformal coverage guarantee does not apply; the paper reports empirical coverage improvements despite this theory gap.

    Source: §2.3 · p. 4

Abstract

from arXiv · show

Bayesian optimization is a coherent, ubiquitous approach to decision-making under uncertainty, with applications including multi-arm bandits, active learning, and black-box optimization. Bayesian optimization selects decisions (i.e. objective function queries) with maximal expected utility with respect to the posterior distribution of a Bayesian model, which quantifies reducible, epistemic uncertainty about query outcomes. In practice, subjectively implausible outcomes can occur regularly for two reasons: 1) model misspecification and 2) covariate shift. Conformal prediction is an uncertainty quantification method with coverage guarantees even for misspecified models and a simple mechanism to correct for covariate shift. We propose conformal Bayesian optimization, which directs queries towards regions of search space where the model predictions have guaranteed validity, and investigate its behavior on a suite of black-box optimization tasks and tabular ranking tasks. In many cases we find that query coverage can be significantly improved without harming sample-efficiency.

1 INTRODUCTION

Bayesian optimization can select queries whose predicted utility is unreliable when the surrogate is misspecified or query points differ from training data. The paper introduces conformal BayesOpt to guide queries toward regions with guaranteed conformal coverage while retaining key Bayesian optimization capabilities.

1 INTRODUCTION

  • Feedback covariate shift makes heldout training accuracy unrepresentative of likely queries, while surrogate misspecification can make expected query utility diverge from actual utility.Both issues can leave BayesOpt with little or no objective improvement after many rounds.Source: §1 · p. 1
  • Bayes credible-set coverage depends on modeling assumptions, so assumption error can produce poor coverage rather than merely wider prediction sets.The preferred trade-off is lost efficiency through wider sets instead of poor coverage.Source: §1 · p. 1
  • Conformal prediction offers reliable coverage without requiring a well-specified surrogate and can accommodate post hoc covariate shift, but poses practical challenges for BayesOpt.Challenges include nondifferentiability, discretizing continuous outcomes, estimating unknown density ratios, and not distinguishing epistemic from aleatoric uncertainty.Source: §1 · p. 1
  • With miscoverage tolerance α < 1, conformal BayesOpt directs search toward regions where conformal predictions are guaranteed coverage of at least 1 − α.The method accounts for possible modeling error while keeping feedback covariate shift in check.Source: §1 · pp. 1–2
  • The conformal Bayes posterior and generalized acquisition functions integrate conformal prediction into BayesOpt while retaining a principled distinction between epistemic and aleatoric uncertainty.The paper also develops differentiable, efficient full conformal Bayes for GP regression and a practical density-ratio estimation procedure.Source: §1 · p. 2
  • Conformal BayesOpt has superior sample-efficiency under surrogate misspecification and comparable sample-efficiency otherwise, while significantly improving query coverage on synthetic black-box and real tabular ranking tasks.The authors emphasize that the goal is to introduce conformal prediction into BayesOpt and study its empirical behavior, not primarily to beat classical alternatives.Source: §1 · p. 2

2 PRELIMINARIES

Bayesian optimization selects queries by expected utility under a posterior model, but misspecification can make those expectations unreliable. Conformal prediction supplies coverage guarantees under stated exchangeability conditions, motivating its use in BayesOpt.

2.1 Bayesian optimization

  • BayesOpt optimizes unknown objectives by selecting noisy observations with an acquisition function that averages user utility over the posterior.The optimization may involve multiple objectives.Source: §2.1 · p. 2, §2 · p. 2

2.2 Bayesian inference and model misspecification

  • Credible-set coverage depends on modeling assumptions, which can fail through misspecified kernels or likelihoods and feedback covariate shift.The authors argue that assumption error may be preferable to accommodate rather than eliminate, since simplifying assumptions can offer practical benefits.Source: §2.2 · pp. 2–3

2.3 Conformal prediction

  • Conformal sets test candidate labels using conformity scores and importance weights, including weights that account for covariate shift.The procedure includes a candidate label when the masked-weight partial sum exceeds α; its marginal coverage is at least 1−α under pseudo-exchangeability.Source: §2.3 · p. 3, Figure 2 · p. 3
  • The coverage guarantee is marginal, not specific to each input, and does not apply when active BayesOpt data lack the required pseudo-exchangeability.The authors note that large offline datasets are not necessarily IID either.Source: §2.3 · p. 4
  • Full conformal Bayes scores candidates by their log posterior predictive density after adding each candidate observation to the data.When the surrogate is correctly specified, this score yields prediction sets with the smallest expected volume among valid sets; under pseudo-exchangeability, predictions remain valid.Source: §2.3 · p. 4

3 RELATED WORK

Prior work addresses covariate shift or model misspecification in optimization, but often relies on restrictive settings or assumptions. The paper positions conformal BayesOpt as a more general approach with distribution-free coverage and compatibility with reparameterizable acquisition functions.

Conformal prediction

  • Fannjiang et al. (2022) address feedback covariate shift, but their method is difficult to extend to most BayesOpt methods because it assumes a closed-form proposal and exact conformal prediction is not differentiable.Other approaches require a minimax subproblem or focus on fully exchangeable classification data; this work allows covariate shift and estimates importance weights for implicit query densities.Source: §3 · p. 4

Robust BayesOpt

  • Robust BayesOpt methods address specific misspecification settings but often require strong smoothness assumptions, fixed misspecification types, or offline pretraining data.The discussed methods primarily concern continuous decisions, unlike much work in the bandit setting.Source: §3 · p. 4
  • Unlike TuRBO's heuristic trust region, conformal BayesOpt provides distribution-free coverage under mild assumptions and applies to any reparameterizable acquisition function.The authors describe their approach as the first BayesOpt procedure to incorporate conformal prediction.Source: §3 · p. 4

4 METHOD

Conformal BayesOpt builds prediction sets and a conformal posterior to guide acquisition under model misspecification and feedback covariate shift. It combines differentiable full conformal prediction with density-ratio estimation for query proposals.

4.1 Full conformal Bayes with Gaussian processes

  • Full conformal Bayes uses Gaussian-process posterior updates to evaluate candidate outcomes efficiently, while sigmoid relaxation makes prediction masks differentiable for query optimization.Candidates are sampled from the predictive posterior rather than selected from a dense outcome grid; as the relaxation temperature approaches zero, it becomes exact but gradients become poorly behaved.Source: §4.1 · p. 5

4.2 Conformal acquisition functions

  • The conformal Bayes posterior replaces predictive outcome weights with conformal weights, assigning total probability 1 −α to the prediction set and retaining base-posterior weighting outside it.This reweighting preserves the distinction between epistemic and aleatoric uncertainty and supports conformal versions of acquisition functions.Source: §4.2 · pp. 5–6

4.3 Accounting for Feedback Covariate Shift

  • The method targets feedback covariate shift because optimization queries are not exchangeable with training data, and the shift can intensify as optimization proceeds.Estimating the proposal-to-training density ratio is needed to adapt conformal sets to this shift.Source: §4.3 · p. 6
  • Figure 5 reports that learned density-ratio estimates have fairly minimal effect on conformal coverage, while increasing sigmoid temperature makes prediction sets more conservative.The figure compares coverage under oracle and learned density ratios and examines temperature under shifted and unshifted covariate distributions.Source: Figure 5 · p. 7
  • The method iteratively samples query proposals with SGLD and updates a classifier-based density-ratio estimate, using averaged classifier weights to stabilize acquisition gradients.It begins with an estimate of 1 and alternates between sampling from the acquisition-induced proposal and updating the classifier.Source: §4.3 · pp. 6–7

5 EXPERIMENTS

Across synthetic optimization experiments, conformal BayesOpt improves query-outcome reliability while retaining comparable sample-efficiency, and it performs well with misspecified likelihoods. The experiments also examine coverage under approximations and preview results on drug-design ranking tasks.

5.1 Do Our Approximations Impact Coverage?

  • Empirical conformal coverage concentrates around 1 −α relative to Bayesian credible-set coverage, and learned density ratios and continuous relaxation do not compromise validity.Increasing relaxation temperature makes the prediction sets more conservative.Source: §5.1 · p. 8

5.2 Model Misspecification and Sample-Efficiency

  • With a simpler, misspecified noise model, conformal BayesOpt finds a better solution with fewer queries than specialized heteroscedastic baselines on sinc and double-knot.The baselines use heteroscedastic likelihoods and require multiple replicates of each query.Source: §5.2 · p. 8

5.3 Good Query Coverage and Good Sample-Efficiency Are Not Mutually Exclusive

  • Conformal BayesOpt consistently tracks nominal coverage on query points and improves objective value with sample-efficiency comparable to standard BayesOpt and TuRBO.On levy-20d and ackley-20d, credible-set coverage varies widely on queries despite appearing reasonable on random holdout data.Source: §5.3 · p. 8, Figure 7 · p. 8
  • A drug-design ranking experiment previews improved query coverage without harming sample-efficiency when C-UCB selects small molecules for dopamine receptor binding affinity.The full experiment is reported in Appendix C.3.Source: §5.3 · p. 8, Figure 8 · p. 9

6 DISCUSSION

The discussion frames conformal prediction as a response to unreliable Bayesian credible sets under model misspecification and optimization-induced covariate shift. The authors report improved query coverage with sample-efficiency comparable to methods without coverage guarantees, and outline directions for extending the approach.

6 DISCUSSION

  • Conformal BayesOpt improves query coverage while achieving sample-efficiency comparable to methods without coverage guarantees.The discussion attributes credible-set unreliability to the combination of model misspecification and optimization-induced covariate shift.Source: §6 · p. 9
  • The authors identify transfer to larger models and datasets, discrete optimization, and theoretical analysis of coverage and regret as future directions.They focus on GP surrogates in the low-n regime and suggest split conformal Bayes or alternative surrogate models as possible paths to scale.Source: §6 · p. 9
  • The discussion argues for methods that accommodate imperfect models while balancing internal coherence with external validity.Source: §6 · p. 9

Appendices

The appendices collect assumptions and limitations, derivations, additional experiments, and implementation details. The stated coverage guarantee relies on pseudo-exchangeability, a condition the authors say BayesOpt satisfies through the order-independent mixture of past query likelihoods.

Appendices

  • Appendix A covers assumptions, limitations, and broader impacts; Appendix B gives derivations; Appendix C adds experiments; and Appendix D details implementations.Source: Appendices · p. 13

A.1 Assumptions

  • A real-valued measurable score function can yield a valid prediction set, but a poor score can reduce efficiency by making sets too large.Source: §A.1 · p. 14
  • Conformal prediction requires pseudo-exchangeability: the joint density depends on sequence values, not their ordering.The authors say BayesOpt satisfies this condition because the training-data likelihood is a mixture of past query likelihoods that do not depend on observation order.Source: §A.1 · p. 14

A.2 Limitations

The limitations section emphasizes that conformal coverage is marginal and that implementation approximations may cause deviations from expected coverage. The broader-impact discussion connects reliability to scientific teams while noting possible dual-use risks of black-box optimization.

A.2 Limitations

  • Coverage is guaranteed only in aggregate across many independent parallel experiments; coverage for a particular input in a particular trial can vary substantially.The authors identify stronger conditional validity as future work.Source: §A.2 · p. 14
  • Discretizing continuous labels and relaxing conformal prediction sets may cause deviations from the expected coverage level.The authors report that these approximations do not appear too severe empirically, while noting that avoiding discretization remains future work.Source: §A.2 · p. 14

A.3 Broader Impacts:

  • Black-box optimization is application-agnostic and could be used to discover toxins or design weapons, but the authors say their work is not especially vulnerable to misuse relative to existing work.Source: §A.3 · p. 14
  • The authors hope to spur work on reliably interacting with current models alongside efforts to build better models.They note that misspecification and covariate shift can arise from modeling choices and active data selection.Source: §A.3 · p. 14
  • The authors present this work as a step toward interpretable reliability certificates for interdisciplinary teams in experimental science.They stress that materials science and drug discovery require coordination among scientists and engineers.Source: §A.3 · p. 14

B PROOFS AND DERIVATIONS

The appendices develop randomized conformal prediction, characterize the conformal Bayes posterior, and extend acquisition procedures to single-objective, multi-objective, and batch settings. They also describe handling out-of-distribution queries and generating synthetic targets with known RKHS norm.

B.1 Smoothed conformal prediction

  • Randomization at prediction-set boundaries makes conformal sets exactly valid, avoiding conservatively wide sets that may hinder decision-making.The algorithm probabilistically accepts boundary labels, with acceptance depending on their proximity to the threshold.Source: §B.1 · p. 15

B.2 Characterizing the conformal Bayes posterior

  • As α approaches 1, the conformal Bayes posterior converges pointwise to the ordinary Bayesian posterior.Source: §B.2 · p. 16

B.3 Monte Carlo integration of conformal acquisition functions

  • Monte Carlo integration estimates conformal acquisition functions using normalized weighted sums over sampled outcomes and function draws.Source: §B.3 · pp. 17–18

B.4 Conformalized Single Objective Acquisitions

  • Conformal UCB integrates over outcomes at least as large as the predictive mean because UCB is optimistic.The predictive mean and covariance are recomputed conditional on each candidate outcome.Source: §B.4 · p. 18

B.5 Conformalized Multi-Objective Acquisition Functions

  • Multi-objective optimization evaluates approximate Pareto fronts using hypervolume relative to a user-specified reference point.The derivations also extend to composite acquisitions such as qParEGO and qNParEGO.Source: §B.5 · p. 18

B.6 Conformalizing Batch Acquisitions

  • Batch acquisition samples outcomes jointly, and the Monte Carlo procedure handles their increased dimensionality by sampling from the joint outcome distribution.The current implementation applies pointwise conformal masks before maximizing utility across batch elements; joint masks are left for future work.Source: §B.6 · p. 19

B.7 Out-of-distribution queries

  • When an out-of-distribution query would make the prediction set include every possible label, its conformal acquisition value is set to zero.This favors in-distribution points with positive expected utility over such out-of-distribution points for conformal EI and similar acquisitions.Source: §B.7 · p. 19

B.8 Generating target functions with known RKHS norm

  • Synthetic targets with known RKHS norm are generated as GP posterior means conditioned on noisy labels sampled at random basis points.The number of basis points needed for non-flat functions grows exponentially with input dimension under default GP kernels.Source: §B.8 · p. 20

C.1 Single-Objective Black-Box Optimization

Across the black-box tasks, conformal BayesOpt generally improves prediction-set coverage while achieving solution quality comparable to standard BayesOpt. The single- and multi-objective experiments examine acquisition performance and coverage, with multi-objective coverage gains remaining modest.

C.1 Single-Objective Black-Box Optimization

  • qEI, qNEI, conformal qEI, and conformal qNEI perform similarly; conformal qUCB is best except on Ackley-20, where qUCB is best.Across acquisition functions, no clear overall ranking emerges, though UCB and conformal UCB perform well in general.Source: Figure 10 · p. 21, §C.1 · p. 21
  • Conformal prediction sets track the target coverage better than Bayesian credible sets across single-objective holdout data.The reported curves are evaluated on Levy-d over varying query batch sizes and input dimensionalities.Source: §C.1 · p. 21, Figure 11 · p. 21

C.2 Multi-Objective Black-Box Optimization

  • Conformal and standard acquisitions find multi-objective solution sets with similar hypervolumes, while conformal coverage tracks 1 −α better; both are underconfident.The coverage improvement is fairly small in this setting, and boundary localization degrades as objective dimensionality increases.Source: Figure 12 · p. 22, §C.2 · p. 22

C.3 Tabular Ranking with Real-World Drug and Antibody Data

On real-world molecular and antibody ranking tasks, conformal UCB improves query coverage overall without sacrificing sample-efficiency. The experiments restrict selection to prespecified candidates and use continuous descriptor features.

C.3 Tabular Ranking with Real-World Drug and Antibody Data

  • Restricting search to a fixed candidate set simplifies conformal acquisition computation by allowing direct use of candidate samples for ratio-estimator training.The experiments begin with 32 labeled entries and sequentially select 128 candidates with q = 1.Source: §C.3 · p. 23
  • The tasks rank small molecules for penalized logP, QED, and DRD3 binding affinity, and large antibodies for stability.Small-molecule and antibody features use RDKit chemical descriptors and BioPython sequence descriptors, respectively.Source: §C.3 · p. 23
  • Conformal UCB selects queries with better coverage overall without sacrificing sample-efficiency on the ranking tasks.Sample-efficiency is measured by cumulative regret; coverage is less consistent than in black-box optimization with unrestricted search.Source: §C.3 · p. 23, Figure 13 · p. 23

C.4 Comparing Bayesian Credible Sets and Conformal Bayes Prediction Sets in the Well-Specified Regime

In well-specified settings, conformal prediction sets can be wider than Bayesian credible sets where training data are sparse, while credible sets can be well-calibrated across the domain. The comparison emphasizes that the choice between approaches depends on prior knowledge and the cost of miscalibration.

C.4 Comparing Bayesian Credible Sets and Conformal Bayes Prediction Sets in the Well-Specified Regime

  • In sparse-data regions, conformal prediction sets are underconfident and much wider than credible sets, while the credible sets are well-calibrated across the domain.This occurs when the GP prior is highly concordant with the target function.Source: §C.4 · p. 24, Figure 14 · p. 24
  • The choice between conformal and conventional Bayesian methods depends on available data, confidence in the prior, and the cost of miscalibration.The paper notes that conformal Bayes prediction sets are efficient among rules guaranteed valid at the 1 −α level, but may be underconfident in sparse regions.Source: §C.4 · p. 24
  • With five-dimensional inputs and 128 basis points, conformal coverage is less satisfactory, partly because data are not IID and density-ratio approximation has error.In this well-specified experiment, Bayesian credible sets are fairly well-calibrated and conformal performance is comparable.Source: §C.4 · p. 24, Figure 15 · p. 25

D IMPLEMENTATION DETAILS

The implementation combines a conformal BayesOpt query-selection loop with numerical techniques for stable GP predictions and specifies model settings, computational costs, and runtime. The experiments use a single GPU at a time and list the software dependencies.

D.1 Bringing Everything Together

  • The inner loop estimates ratio-based weights, constructs an outcome mask, evaluates an acquisition, updates the query with SGLD, and retrains the classifier.The algorithm returns the updated query after its iterative selection loop.Source: §D.1 · p. 26

D.2 Stable Predictions on the Training Set

  • To avoid instability in GP posterior covariance calculations, the implementation uses stochastic diagonal estimation with 10 Bernoulli probe vectors.The authors report instability from solving and subtracting an n × n system and ill-defined gradients through eigendecomposition.Source: §D.2 · p. 26

D.3 Hyperparameters

  • The GP setup uses a scaled Matérn-5/2 kernel, constant prior mean, specified Gamma priors, and a noise range of 5 × 10^-4 to 0.5.The miscoverage schedule is α = max(0.05, 1/√n), with initialization from 10 Sobol points.Source: §D.3 · pp. 26–27

D.4 Computational Complexity

  • Exact GP training costs O(n^3), while retraining on one new example costs O(n); acquisition-gradient evaluation costs O(qn^2) with a k-fold memory increase.The paper says one sample from each conformal posterior tends to suffice because its predictions concentrate near the candidate label.Source: §D.4 · p. 27
  • Conformal BayesOpt runs around an order of magnitude slower in wall-clock time than conventional BayesOpt.Burn-in for SGLD chains and convergence of the bootstrapped ratio estimator require more gradient evaluations than L-BFGS acquisition optimization.Source: §D.4 · pp. 27–28

D.5 Compute Resources

  • The experiments used one GPU at a time and were estimated to require approximately 250 GPU hours to reproduce.The estimate covers 25 seeds per variant across 10 experiments; development and debugging likely used substantially more.Source: §D.5 · p. 28

D.6 Software Packages

  • The software list names Matplotlib, Seaborn, NumPy, PyTorch, GPyTorch, and BoTorch.Source: §D.6 · p. 28

Related papers

Loading 2210.12496v4…