Acquisition Function-Guided Generation for Bayesian Optimization over Discrete Spaces

Tommaso Ancilli; Thomas Wedenig; Henry Moss; Robert Peharz

cs.LGOct 16, 2026

Acquisition Function-Guided Generation for Bayesian Optimization over Discrete Spaces

Authors: Tommaso Ancilli, Thomas Wedenig, Henry Moss, Robert Peharz

Organizations: Institute of Machine Learning and Neural Computation, TU Graz, Austria · Graz Center for Machine Learning (GraML), Austria · School of Mathematical Sciences, Lancaster University, UK

Abstract

Bayesian optimization (BO) over discrete spaces is challenging due to missing gradients and the sparsity of valid solutions. Latent-space BO sidesteps both issues by fitting a surrogate in the latent space of a generative model, which however substantially complicates the direct incorporation of domain knowledge that is naturally expressed in the original discrete space. In this work, we instead combine a pretrained generative model with a surrogate defined directly in data space, where inductive biases such as domain-specific kernels can be readily incorporated. Rather than maximizing the acquisition function over the discrete design space, our method samples candidates from the generative model under a distribution tilted by the acquisition function. To recover useful gradients for sampling, we embed this tilted distribution into a diffusion bridge that yields a smooth, posterior-averaged acquisition function. This leads to acquisition function guided generation (AFGG), which we instantiate with a categorical diffusion model and an autoregressive normalizing flow. Across a broad suite of molecular benchmarks, the diffusion model variant performs on par with leading BO methods, while the flow model variant outperforms state-of-the-art methods in most cases and remains highly competitive otherwise. Our implementation is available at https://anonymous.4open.science/r/afgg.

Figures & tables

Figure 1: Sequential setting. Optimization trajectories on six representative objectives, with 10 initial observations followed by 300 oracle evaluations, one per BO iteration. Curves show the median over 5 seeds.

	baselines					AFGG
Objective	CoBO	InvBO	LOL-BO	NFBO	COWBOYS	CDCD	SeqFlow
albu_sim	0.593 $\pm$ 0.111	0.629 $\pm$ 0.079	0.509 $\pm$ 0.065	0.593 $\pm$ 0.119	0.448 $\pm$ 0.037	0.691 $\pm$ 0.073	0.975 $\pm$ 0.038
amlo_mpo	0.523 $\pm$ 0.012	0.538 $\pm$ 0.050	0.516 $\pm$ 0.022	0.563 $\pm$ 0.095	0.472 $\pm$ 0.038	0.546 $\pm$ 0.039	0.565 $\pm$ 0.045
cele_red	0.262 $\pm$ 0.041	0.286 $\pm$ 0.068	0.243 $\pm$ 0.035	0.455 $\pm$ 0.035	0.243 $\pm$ 0.020	0.392 $\pm$ 0.024	0.836 $\pm$ 0.125
deco_hop	0.565 $\pm$ 0.008	0.597 $\pm$ 0.019	0.560 $\pm$ 0.009	0.673 $\pm$ 0.147	0.559 $\pm$ 0.008	0.599 $\pm$ 0.010	0.803 $\pm$ 0.175
drd2	0.510 $\pm$ 0.370	0.963 $\pm$ 0.037	0.626 $\pm$ 0.184	0.971 $\pm$ 0.066	0.796 $\pm$ 0.243	0.730 $\pm$ 0.226	1.000 $\pm$ 0.000

Table 1: Sequential benchmark (partial results). Final best score on all 25 PMO objectives after 300 oracle calls (10 initial points, 1 point per BO iteration), mean

\pm

sample standard deviation over 5 seeds. For readability the table shows five of the eleven baselines; Table 6 gives the full comparison. AFGG is shown in two instantiations: CDCD and SeqFlow . Every cell is judged against the opposing group. Bold : the cell’s mean exceeds every mean in the opposing group. Green, shaded : lead exceeds the cell’s own sample standard deviation; bold without shading denotes a smaller lead. Objective abbreviations are defined in Table 5 .

Figure 2: Batch setting, uncoupled generation. Optimization trajectories on the same six representative objectives, with 100 initial observations followed by 500 oracle evaluations in batches of five. Curves show the median over 5 seeds.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

	CDCD	SeqFlow
noised variable	embedding $x$	flow latent $g_{\theta}^{-1}(x)$
schedule $(\alpha_{t},\sigma_{t})$	$(1,\,t)$ , VE	$\big(\rho_{t},\,\sqrt{1-\rho_{t}^{2}}\big)$ , VP Eq. ( 51 )
(i) $\nabla_{z_{t}}\log p_{t}$	learned Eq. ( 6 )	$-z_{t}$ , exact
(ii) draws from $p(\cdot\mid z_{t})$	mean-field $\tilde{p}$ Eq. ( 49 ), one network pass	exact Gaussian Eq. ( 27 ), $N$ passes of $g_{\theta}$
$\nabla_{z_{t}}\log\ell_{t}$	SNIS Eq. ( 42 ) or REINFORCE Eq. ( 50 )	SNIS Eq. ( 42 )
noise end	$p^{\star}_{T}\approx\mathcal{N}(0,T^{2}I)$	$\pi^{\star}_{T}\approx\mathcal{N}(0,I)$

Appendix

Table 2: What each backbone supplies to the guided score Eq. ( 19 ). Quantities (i) and (ii) are backbone-specific; the guidance term

\nabla_{z_{t}}\log\ell_{t}

is estimated from the draws of (ii) (Section 3.3 ).

objective	$K=0$	$K=32$	$\Delta$	$\Delta/\mathrm{sd}$
med2	$0.1637\pm 0.0193$	$0.1694\pm 0.0243$	$+0.0058$	$+0.26$
osmb	$0.7094\pm 0.0578$	$0.7055\pm 0.0611$	$-0.0039$	$-0.07$
adip	$0.4335\pm 0.0297$	$0.4483\pm 0.0029$	$+0.0148$	$+0.70$
zale	$0.3811\pm 0.0935$	$0.2842\pm 0.0612$	$-0.0969$	$-1.23$
rano	$0.7503\pm 0.0361$	$0.7088\pm 0.0570$	$-0.0416$	$-0.87$
qed	$0.8992\pm 0.0370$	$0.9106\pm 0.0387$	$+0.0114$	$+0.30$

Appendix

Table 3: The mean-field approximation Eq. ( 49 ) against rollouts of the reverse process.

K

is the number of inner denoising steps used to draw each of the

N

SNIS samples:

K=0

is the factorized draw from the per-position marginals Eq. ( 7 ), i.e. Eq. ( 49 ), and is what Algorithm 1 uses;

K>0

replaces it by a

K

-step ancestral rollout, which couples the positions through the network and resolves the joint. Final objective value after

300

BO epochs on six GuacaMol/TDC objectives, mean

\pm

sd over

3

seeds; higher is better.

\Delta

K{=}32

minus

K{=}0

, and the last column expresses it in units of the pooled seed-to-seed sd. At

K=32

— beyond the point where proposal fidelity saturates, and at

\approx 24\times

the wall clock (

8.77

against

0.36

h per arm) — no objective improves beyond the seed spread and one degrades.

Architecture
Hyperparameter	Value
Embedding dimension	$16$
Hidden dimension	$512$
Query/key/value dimension	$512$
Time embedding dimension	$16$
Layers	$8$

Appendix

Table 4: Hyperparameters of the CDCD generator.

Abbreviation	PMO objective	Abbreviation	PMO objective
albu_sim	albuterol_similarity	osim_mpo	osimertinib_mpo
amlo_mpo	amlodipine_mpo	peri_mpo	perindopril_mpo
cele_red	celecoxib_rediscovery	rano_mpo	ranolazine_mpo
deco_hop	deco_hop	logp	rdkit_logp
drd2	drd2_docking	qed	rdkit_qed
fexo_mpo	fexofenadine_mpo	sa	sa_tdc

Appendix

Table 5: Abbreviations of the 25 PMO objectives, used in all tables.

Figure 3: Sequential setting : full trajectories on all 25 PMO objectives (10 initial points, 300 oracle calls, 1 point per BO iteration).

	baselines											AFGG
Objective	CoBO	InvBO	LOL-BO	NFBO	CMA-ES	COWBOYS	DE	GA	HC	RLBO	TuRBO	CDCD	SeqFlow
albu_sim	0.593 $\pm$ 0.111	0.629 $\pm$ 0.079	0.509 $\pm$ 0.065	0.593 $\pm$ 0.119	0.500 $\pm$ 0.030	0.448 $\pm$ 0.037	0.488 $\pm$ 0.035	0.385 $\pm$ 0.074	0.454 $\pm$ 0.070	0.483 $\pm$ 0.091	0.475 $\pm$ 0.047	0.691 $\pm$ 0.073	0.975 $\pm$ 0.038
amlo_mpo	0.523 $\pm$ 0.012	0.538 $\pm$ 0.050	0.516 $\pm$ 0.022	0.563 $\pm$ 0.095	0.459 $\pm$ 0.038	0.472 $\pm$ 0.038	0.438 $\pm$ 0.143	0.419 $\pm$ 0.082	0.460 $\pm$ 0.048	0.491 $\pm$ 0.030	0.475 $\pm$ 0.024	0.546 $\pm$ 0.039	0.565 $\pm$ 0.045
cele_red	0.262 $\pm$ 0.041	0.286 $\pm$ 0.068	0.243 $\pm$ 0.035	0.455 $\pm$ 0.035	0.227 $\pm$ 0.032	0.243 $\pm$ 0.020	0.216 $\pm$ 0.036	0.210 $\pm$ 0.040	0.228 $\pm$ 0.032	0.224 $\pm$ 0.031	0.221 $\pm$ 0.033	0.392 $\pm$ 0.024	0.836 $\pm$ 0.125
deco_hop	0.565 $\pm$ 0.008	0.597 $\pm$ 0.019	0.560 $\pm$ 0.009	0.673 $\pm$ 0.147	0.563 $\pm$ 0.010	0.559 $\pm$ 0.008	0.571 $\pm$ 0.018	0.549 $\pm$ 0.015	0.553 $\pm$ 0.012	0.556 $\pm$ 0.015	0.552 $\pm$ 0.012	0.599 $\pm$ 0.010	0.803 $\pm$ 0.175
drd2	0.510 $\pm$ 0.370	0.963 $\pm$ 0.037	0.626 $\pm$ 0.184	0.971 $\pm$ 0.066	0.184 $\pm$ 0.068	0.796 $\pm$ 0.243	0.862 $\pm$ 0.156	0.154 $\pm$ 0.051	0.183 $\pm$ 0.045	0.380 $\pm$ 0.246	0.556 $\pm$ 0.337	0.730 $\pm$ 0.226	1.000 $\pm$ 0.000

Appendix

Table 6: Final best score on all 25 PMO objectives after 300 oracle calls (10 initial points, 1 point per BO iteration), mean

\pm

sample standard deviation over 5 seeds. AFGG is shown in two instantiations: CDCD and SeqFlow . Every cell is judged against the opposing group: an AFGG cell against the 11 baselines, a baseline cell against the two AFGG columns. Bold : the cell’s mean exceeds every mean in the opposing group. Green : that lead is significant, exceeding the cell’s own standard deviation; bold without green therefore means ahead but within noise.

Figure 4: Batch regime, uncoupled generation : 500 acquisitions (100 initial points, then 100 epochs of 5), all 25 objectives.

	baselines					AFGG
Objective	CoBO	InvBO	LOL-BO	NFBO	COWBOYS	CDCD	SeqFlow
albu_sim	0.692 $\pm$ 0.095	0.771 $\pm$ 0.057	0.579 $\pm$ 0.043	0.609 $\pm$ 0.073	0.551 $\pm$ 0.081	0.779 $\pm$ 0.073	0.960 $\pm$ 0.050
amlo_mpo	0.540 $\pm$ 0.021	0.569 $\pm$ 0.029	0.522 $\pm$ 0.026	0.582 $\pm$ 0.024	0.501 $\pm$ 0.021	0.565 $\pm$ 0.052	0.618 $\pm$ 0.075
cele_red	0.307 $\pm$ 0.035	0.366 $\pm$ 0.064	0.304 $\pm$ 0.036	0.452 $\pm$ 0.052	0.316 $\pm$ 0.052	0.457 $\pm$ 0.098	0.861 $\pm$ 0.135
deco_hop	0.588 $\pm$ 0.005	0.610 $\pm$ 0.007	0.586 $\pm$ 0.004	0.617 $\pm$ 0.018	0.583 $\pm$ 0.004	0.614 $\pm$ 0.019	0.878 $\pm$ 0.120
drd2	0.930 $\pm$ 0.090	0.996 $\pm$ 0.006	0.728 $\pm$ 0.247	0.994 $\pm$ 0.012	0.774 $\pm$ 0.210	0.792 $\pm$ 0.288	1.000 $\pm$ 0.000

Appendix

Table 7: Batch regime, uncoupled generation. Final best score on all 25 PMO objectives after 500 acquisitions (100 initial points, then 100 epochs of 5), mean

\pm

sample standard deviation over 5 seeds. AFGG generates each batch uncoupled here: every molecule is guided by its own acquisition value, and the batch only forms at selection. Every cell is judged against the opposing group: an AFGG cell against the five baselines, a baseline cell against the two AFGG columns. Bold : the cell’s mean exceeds every mean in the opposing group. Green : that lead is significant, exceeding the cell’s own standard deviation; bold without green therefore means ahead but within noise.

Figure 5: Batch regime, coupled generation : 500 acquisitions (100 initial points, then 100 epochs of 5), all 25 objectives.

	baselines					AFGG
Objective	CoBO	InvBO	LOL-BO	NFBO	COWBOYS	CDCD	SeqFlow
albu_sim	0.692 $\pm$ 0.095	0.771 $\pm$ 0.057	0.579 $\pm$ 0.043	0.609 $\pm$ 0.073	0.551 $\pm$ 0.081	0.703 $\pm$ 0.057	0.910 $\pm$ 0.083
amlo_mpo	0.540 $\pm$ 0.021	0.569 $\pm$ 0.029	0.522 $\pm$ 0.026	0.582 $\pm$ 0.024	0.501 $\pm$ 0.021	0.606 $\pm$ 0.068	0.630 $\pm$ 0.051
cele_red	0.307 $\pm$ 0.035	0.366 $\pm$ 0.064	0.304 $\pm$ 0.036	0.452 $\pm$ 0.052	0.316 $\pm$ 0.052	0.402 $\pm$ 0.034	0.645 $\pm$ 0.153
deco_hop	0.588 $\pm$ 0.005	0.610 $\pm$ 0.007	0.586 $\pm$ 0.004	0.617 $\pm$ 0.018	0.583 $\pm$ 0.004	0.615 $\pm$ 0.018	0.866 $\pm$ 0.153
drd2	0.930 $\pm$ 0.090	0.996 $\pm$ 0.006	0.728 $\pm$ 0.247	0.994 $\pm$ 0.012	0.774 $\pm$ 0.210	1.000 $\pm$ 0.000	1.000 $\pm$ 0.000

Appendix

Table 8: Batch regime, coupled generation. Final best score on all 25 PMO objectives after 500 acquisitions (100 initial points, then 100 epochs of 5), mean

\pm

sample standard deviation over 5 seeds. AFGG generates each batch coupled here: one joint qLogEI tilt per candidate batch, so its members are generated aware of each other. Every cell is judged against the opposing group: an AFGG cell against the five baselines, a baseline cell against the two AFGG columns. Bold : the cell’s mean exceeds every mean in the opposing group. Green : that lead is significant, exceeding the cell’s own standard deviation; bold without green therefore means ahead but within noise.

Metric	CDCD	SeqFlow
Validity $\uparrow$	1.0000	0.9972
Uniqueness $\uparrow$	0.9932	0.9862
Novelty $\uparrow$	0.9956	0.9838
FCD $\downarrow$	20.99	6.02

Appendix

Table 9: Unconditional sample quality.

5000

molecules drawn from each generative model and scored against a

30

k-molecule subset of GuacaMol.

	Guidance on vs off
Objective	on	off ( $\lambda{=}0$ )
osim_mpo	0.714 $\pm$ 0.064	0.722 $\pm$ 0.048
med_2	0.199 $\pm$ 0.013	0.173 $\pm$ 0.012
amlo_mpo	0.486 $\pm$ 0.053	0.408 $\pm$ 0.018
zale_mpo	0.420 $\pm$ 0.054	0.431 $\pm$ 0.030
rano_mpo	0.788 $\pm$ 0.012	0.745 $\pm$ 0.026

Appendix

Table 10: Guidance on vs. off; off sets

\lambda=0

, flattening the SNIS weights to uniform so the tilted mean equals the untilted one.

	Guidance estimator
Objective	SNIS (closed form)	REINFORCE
osim_mpo	0.714 $\pm$ 0.064	0.671 $\pm$ 0.033
med_2	0.199 $\pm$ 0.013	0.170 $\pm$ 0.010
amlo_mpo	0.486 $\pm$ 0.053	0.410 $\pm$ 0.062
zale_mpo	0.420 $\pm$ 0.054	0.250 $\pm$ 0.085
rano_mpo	0.788 $\pm$ 0.012	0.687 $\pm$ 0.036

Appendix

Table 11: Guidance estimator: closed-form SNIS vs. REINFORCE (score-function estimator with a leave-one-out baseline) in place of the mean-swap identity.

	Acquisition function
Objective	EI	PI	UCB
osim_mpo	0.714 $\pm$ 0.064	0.708 $\pm$ 0.065	0.696 $\pm$ 0.041
med_2	0.199 $\pm$ 0.013	0.183 $\pm$ 0.015	0.170 $\pm$ 0.009
amlo_mpo	0.486 $\pm$ 0.053	0.495 $\pm$ 0.029	0.490 $\pm$ 0.013
zale_mpo	0.420 $\pm$ 0.054	0.439 $\pm$ 0.045	0.445 $\pm$ 0.055
rano_mpo	0.788 $\pm$ 0.012	0.753 $\pm$ 0.026	0.746 $\pm$ 0.037

Appendix

Table 12: Acquisition function driving both the guidance tilt and candidate selection; UCB uses

\beta=0.1

	Candidate pool size
Objective	1	5	10	20
osim_mpo	0.736 $\pm$ 0.010	0.712 $\pm$ 0.053	0.714 $\pm$ 0.064	0.711 $\pm$ 0.063
med_2	0.186 $\pm$ 0.012	0.192 $\pm$ 0.008	0.199 $\pm$ 0.013	0.204 $\pm$ 0.007
amlo_mpo	0.486 $\pm$ 0.045	0.504 $\pm$ 0.038	0.486 $\pm$ 0.053	0.496 $\pm$ 0.029
zale_mpo	0.432 $\pm$ 0.048	0.411 $\pm$ 0.084	0.420 $\pm$ 0.054	0.406 $\pm$ 0.060
rano_mpo	0.764 $\pm$ 0.008	0.768 $\pm$ 0.019	0.788 $\pm$ 0.012	0.771 $\pm$ 0.015

Appendix

Table 13: Candidate pool size: guided trajectories rolled out per epoch, of which the best by acquisition value is evaluated.

	Where the temperature acts
Objective	power tilt $\gamma{=}2$	scale drift $\lambda{=}0.5$	scale drift $\lambda{=}2$	scale drift $\lambda{=}5$
osim_mpo	0.714 $\pm$ 0.064	0.702 $\pm$ 0.034	0.694 $\pm$ 0.040	0.682 $\pm$ 0.026
med_2	0.199 $\pm$ 0.013	0.194 $\pm$ 0.010	0.191 $\pm$ 0.013	0.182 $\pm$ 0.007
amlo_mpo	0.486 $\pm$ 0.053	0.474 $\pm$ 0.029	0.423 $\pm$ 0.029	0.496 $\pm$ 0.051
zale_mpo	0.420 $\pm$ 0.054	0.449 $\pm$ 0.026	0.402 $\pm$ 0.041	0.401 $\pm$ 0.029
rano_mpo	0.788 $\pm$ 0.012	0.750 $\pm$ 0.031	0.782 $\pm$ 0.023	0.773 $\pm$ 0.023

Appendix

Table 14: Different temperature parameters and how they affect the final generation process.

Bayesian Optimization

Bayesian optimization uses probabilistic models and sequential experiments to find strong solutions when evaluating a black-box objective is costly.

See all 240 papers

Sep 16, 2026

Search at the Cost of Sampling: Nearly-Instant Latent Space Bayesian Optimization

Generative models are increasingly central to many de novo discovery pipelines, in which designs are generated at scale and filtered through virtual screens to determine a set of candidates to experimentally validate. While Bayesian optimization (BO) is a natural fit for this setting, as it uses past evaluations to guide future proposals, the computational overhead required for its sequential decision-making becomes a bottleneck when virtual screens are relatively cheap. We make BO practical in this regime by exploiting the unique combination of a linear model constrained to a spherical domain where high-dimensional latents concentrate. We build off recent work justifying the use of linear surrogates, while deriving nearly closed-form solutions to the surrogate modelling and acquisition problems that exploit spherical symmetry. The result is at least a 100x speedup over state-of-the art baselines, with matching or improved performance across molecular and image generation benchmarks. Altogether, our method makes BO a practical drop-in for de novo pipelines where it was previously too slow to consider.

Donney Fan, Colin Doumont, Aleksandra Kalisz +4

University of British Columbia · Vector Institute · Tübingen AI Center +3

Surrogate Modeling Bayesian Optimization Latent Space Optimization

Ask PDF

Oct 12, 2026

Generative Models as Adaptive Trust Regions for Discrete Bayesian Optimisation

Bayesian optimisation (BO) is a powerful framework for sample-efficient optimisation of expensive black-box functions, but its application to high-dimensional discrete scientific design remains challenging. Local trust-region BO methods such as TuRBO address high dimensionality by restricting search to adaptively sized neighbourhoods. However, a direct discrete analogue based on edit distance can be poorly suited to structured discrete objects: neighbourhoods grow combinatorially, and metric proximity does not capture the dependencies that determine whether a design is plausible or useful. Motivated by the growing availability of pretrained generative models, we introduce GenTuRBO, which replaces the geometric trust region with an adaptively scaled proposal distribution defined by a pretrained generator. The generator provides a learned prior over discrete transformations, encoding coordinated changes without explicitly enumerating edit-distance neighbourhoods. GenTuRBO performs strongly across all three domains and achieves the highest mean final scores on Ehrlich sequence optimisation and Sokoban coordination.

Michael Dodds, Erik Bodin, Sebastian W. Ober +3

University of Cambridge · Monumo · BigHat Biosciences +2

Bayesian Optimization Black-Box Optimization Trust-Region Optimization

Ask PDF

Jun 7, 2026

Improving Bayesian Optimization via Training-Aware Conditional Diffusion Models

Bayesian optimization (BO) is a widely used approach for black-box optimization that uses a Gaussian process (GP) as a surrogate and guides sequential evaluations via an acquisition function, with the ultimate goal of locating the global optimum

\mathbf{x}^{\star}

. To align with this goal, information-based acquisition functions such as Predictive Entropy Search (PES) model

\mathbf{x}^{\star}

as a random variable and reduce the entropy of its distribution, but approximating this distribution via traditional GP posterior sampling is computationally expensive. To address this limitation, we leverage Conditional Diffusion Models (CDMs) to efficiently approximate the distribution of

\mathbf{x}^{\star}

and develop BO-inherent training strategies for CDMs. Motivated by the structural properties of the CDM-learned distribution, we further develop an acquisition strategy termed Diffusion-based Mode Seeking (DMS) to guide the sequential evaluation. We establish a sub-optimality guarantee for the CDM-learned distribution and demonstrate through extensive experiments that DMS outperforms standard BO baselines.

Yilin Zheng, Haowei Wang, Szu Hui Ng +1

National University of Singapore · Singapore · Georgia Institute of Technology +1

Bayesian Optimization Black-Box Optimization Conditional Diffusion Models

Ask PDF

Abstract

Bayesian Optimization

Explore similar work

Search at the Cost of Sampling: Nearly-Instant Latent Space Bayesian Optimization

Generative Models as Adaptive Trust Regions for Discrete Bayesian Optimisation

Improving Bayesian Optimization via Training-Aware Conditional Diffusion Models