Return of the Latent Space COWBOYS: Re-thinking the use of VAEs for Bayesian Optimisation of Structured Spaces
Authors: Henry Moss
Sebastian W. Ober
Tom Diethe
Venue: ICML 2025 spotlightposter (2025)
Keywords: Bayesian Optimisation
Primary area: probabilistic_methods->gaussian_processes
TL;DR: Don't do Bayesian optimisation in the latent space of a VAE ....

Abstract: Bayesian optimisation in the latent space of a VAE is a powerful framework for optimisation tasks over complex structured domains, such as the space of valid molecules. However, existing approaches tightly couple the surrogate and generative models, which can lead to suboptimal performance when the latent space is not tailored to specific tasks, which in turn has led to the proposal of increasingly sophisticated algorithms. In this work, we explore a new direction, instead proposing a decoupled approach that trains a generative model and a GP surrogate separately, then combines them via a simple yet principled Bayesian update rule. This separation allows each component to focus on its strengths— structure generation from the VAE and predictive modelling by the GP. We show that our decoupled approach improves our ability to identify high-potential candidates in molecular optimisation problems under constrained evaluation budgets.

Return of the Latent Space COWBOYS:
Re-thinking the use of VAEs for Bayesian Optimisation of Structured Spaces
Henry B. Moss Sebastian W. Ober
Tom Diethe
Abstract
Latent GP
Structured
VAE
VAE
GP
Bayesian optimisation in the latent space of a Variational AutoEncoder (VAE) powerful framework for optimisation tasks over complex structured domains, such the space of scien- tifically interesting molecules. However, exist- ing approaches tightly couple the surrogate and generative models, which can lead to suboptimal performance when the latent space is not tailored to specific tasks, which in turn has led to the pro- posal of increasingly sophisticated algorithms. In this work, explore new direction, instead proposing a decoupled approach that trains a gen- erative model and a Gaussian Process (GP) surro- gate separately, then combines them via simple yet principled Bayesian update rule. This sep- aration allows each component to focus on its strengths- structure generation from the VAE and predictive modelling by the GP. We show that our decoupled approach improves our ability to identify high-potential candidates in molecular optimisation problems under constrained evalua- tion budgets.
Latent Space
COWBOYS
BO
Figure I. Unlike LSBO where GPs are fit in a VAE's latent space, COWBOYS's GP is fit structure space, decoupled from VAE. acquisition routines can be employed. Candidate points selected by the optimiser then decoded back into the original structured domain to yield new query points.
Learning a complex mapping from variable-size structured inputs to fixed-size Euclidean representations requires more data than is typically gathered in LSBO itself. Therefore, LSSO relies on embeddings pre-trained on related task with abundant unlabelled data (e.g., using VAE latent space trained it is well known (Chu et al., 2024) that fitting surrogates in the latent space of VAE can lead to suboptimal mod- elling; therefore, the recent LSBO literature has focused on introducing increasingly sophisticated heuristics to fine-tune the VAE on newly acquired experimental data. These ap- proaches remain challenged by the persistent risk of overfit- ting when training neural network components of Bayesian models limited datasets (Ober et al., 2021). In this work, we also make new argument that the prevalent approach of limiting the search to a fixed box subset of the latent space necessary to apply traditional BO methods - also limits the effectiveness LSBO.
large set of valid molecules). However,
1. Introduction
First introduced by Gómez-Bombarelli et al. (2018), Bayesian Optimisation (BO) over latent spaces has emerged as powerful technique for optimising over structures. Rather than performing challenging combinatorial or high- dimensional optimisation directly on discrete structures such as molecules or proteins - Latent Space BO (LSBO) first maps inputs into a Exeedddimensiona Euclidean latent space, where standard surrogate models and gradient-
While the VAE is designed purely generative model, LSBO attempts repurpose it for discrimination funda- mental mismatch that gives rise the pathologies discussed above. We propose an alternative approach that preserves the original motivation of the VAE as a generative model by coupling with a separately trained surrogate GP model via a novel Bayesian formulation. Our primary goal straightforward: develop a sampling strategy that increas- ingly favours structures with desirable objective values.
'School of Mathematical Sciences, Lancaster University, UK Department of Applied Maths and Theoretical Physies, University of Cambridge, UK Oncology R&D, AstraZeneca. Gaithersburg, USA Biopharma R&D, AstraZeneca, Cambridge, UK. Corre- spondence to: Henry B. Moss
We propose Categorical Optimisation With Belief Of un- derl Structure (COWBOYS), novel principled frame- work that unifies GPs and VAEs within a BO loop. Although we focus on molecular search- the most popular applica- tion of 1eh0-where we demonstrate that COWBOYS
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Return of the Latent Space COWBOYS
excels under heavily constrained optimisation budgets, our approach is general and applicable to any structured prob- lem for which suitable structural GP kernel can be defined (see our discussion in Section 8).
2012; Moss et al., 2021), e.g..
arr(a; D) = P(f(x)> D*)
(I)
where f is the best objective value observed so far. Covariance (kernel) functions. In this work, where we must build a surrogate model over the space of molecules, we use the Tanimoto kernel (Tripp et al., 2023):
2. Background
2.1. Bayesian Optimisation
Bayesian Optimisation (BO) (Mockus, 2005) a powerful framework for the efficient optimisation of costly black-box functions/ f X' + R. For a successful application of BO, we typically need three components: search space, surrogate model, and an acquisition function. Given some initial data, BO constructs a surrogate model of the data, which is then used in tandem with an acquisition function to determine which point the search space will be most valuable to acquire. Once this point chosen, we query the black-box function at this point and update the surrogate model, repeating until the evaluation budget is exhausted.
m. m
(m, m') o² m 1|2
|m m m'
where m and m' are molecular fingerprints representation vectors of molecules and a E R is a tunable scaling factor. We follow the advice of Tripp & Hernández-Lobal (2024) and use count-based vectors that count specific molecular features (Landrum, 2013). Structural kernels such as the Tanimoto kernel, but also string kernels (Moss et al., 2020a) and graph kernels (Kriege et al., 2020), can often outperform deep learning alternatives, particularly in low- to medium- data regimes (Moss & Griffiths, 2020; Griffiths et al., 2022).
Search spaces. In order for optimisation to succeed, the space over which we attempt to find the optimal value must be suitably restricted. For instance, when dealing with a Euclidean space R&, is typical to restrict the search space to box region| 6- 6] where is typically chosen using some prior knowledge about the expected location of the optimum (Hvarfner et al., 2022). In this example, choosing too large a d will lead to slow convergence due to the need for more data so that the surrogate can be effective, whereas too small a 8 will risk missing the optimum. Later, we will discuss how choosing appropriate search spaces is a key bottleneck for LSBO and how COWBOYS addresses this.
2.2. Variational AutoEncoders
Variational AutoEncoders (VAEs) (Kingma & Welling, 2014; Rezende et al., 2014) are deep generative models that model data through a two-step process: sampling a latent variable (with dimension smaller than X') from a prior p(z), followed by sampling æ: through a distribution po po(xlz) ()). resulting in generative model æ po(z). The latter of these steps is performed by neural network, known as the decoder, with parameters O. The decoder outputs the parameters of the conditional distribution po(=) for in- stance the mean and variance of a normal distribution, or class probabilities for a categorical distribution. In order to train a VAE, we require the posterior po (z x)) However, as this is intractable, parameters to approximate it through amortised inference: q (x xx) N (xxx), o² [z) where 6 parameterises μu()) and o²(.) The encoder and decoder are trained jointly through a minibatch-frient -fienndyy lower bound to the marginal likelihood of the data.
Surrogate models. BO requires surrogate models that ac- curately quantify uncertainty while retaining flexibility: as such, Gaussian Processes (GPs) (Rasmussen, 2003) are a standard choice. A GP can be defined as an infinite collec- tion of random variables, any finite number of which are Gaussian distributed, and is entirely defined by a mean func- tion v()) and kernel k(,,)) Under GP
neural network encoder with
GP (P( k), making a prediction f at test point & amounts to condition- ing on the observed data D pfflD) = _N'x*, ", Σ²), where and E. can be computed in closed form via the properties of Gaussians with ((N3) computational complexity and O(N)) memory.
3. The Pitfalls of Latent Space BO
In what follows, we assume that we have access to a pre- trained VAE with a d-dimensional latent space Z = Rd, resulting in a decoding distribution px(z),), and imply a data distribution p=(x).
Acquisition functions. Acquisition functions measure the utility of acquiring an unseen point æ according to the pre- dictions of our surrogate model. Common acquisition fune- tions measure the Expected Improvement (EI), Probability of Improvement (PI) (Bergstra et al., 2011) Upper Confi- dence Bounds (UCB) (Auer, 2000) of candidate points, or use information theoretical arguments (Hennig & Schuler,
3.1. Latent Space BO
Latent Space BO (LSBO) (66mez-Bombarelll et al., 2018) can be viewed as standard BO conducted over an alternative search space - the latent space 2 of a VAE - and using
2

---

Return of the Latent Space COWBOYS
Algorithm 1 Latent Space Bayesian Optimisation
alignment," generally attributed to two VAE properties:
Input: budget N, init size Nini, search bounds 8
1) The VAE is trained offline. Lack of alignment in LSBO is typically attributed to the purely unsupervised training of VAEs, which focuses solely on reconstructing the (unla- belled) training data of the VAE. In particular, as the highly expressive neural network underlying the VAE can distort local neighbourhoods, any smoothness assumptions that mo- tivate Bayesian optimisation in the original space X'-for instance, the continuity of molecule's 's performance un- der small structural perturbations are unlikely to transfer to smoothness in the latent space of the VAE. The vast majority of recent LSBO work (see Section 6) attempts to fine-tune the VAE during optimisation to improve align- ment for the current optimisation objective. However, such approaches are plagued by the fundamental challenge of re- liably fine-tuning neural networks on small datasets without over-fitting behaviour. Indeed, as the VAE has large num- ber of hyperparameters from a Bayesian point of view (i.e., the parameters of the decoder), the VAE's decoder is prone to doing all the heavy lifting, which can lead to overfitting analogous to that described for deep kernel learning in Ober et al. (2021).
Clip search space Z5 < [- C Z
for n E (1)...)} do if n < Ninit then
4 initial design
2n ~ SpacceFillnggessign 2z))
else
4 sequential optimisation
2n + t-armmax222 a(z;
<e.gg Eq. (1)
end if
Decode chosen latent xn pz(
Evaluate new molecular structure Un + f(x,))
Update dataset D < D U((zn:3n)
Fit latent space GP on D end for
return Believed optimum across {z1,.., kn
a different surrogate model, one modelling the mapping g 2 + R from latent codes to objective function values with a surrogate ğ. At the nhhh optimisation step, LSBO proceeds analogously to standard BO. For example, under the PI acquisition function, LSBO selects new latent code En with highest utility
2) The VAE has stochastic decoder. The second chal- lenge for GP models in LSBO, one largely ignored until the recent works of Chu et al. (2024); Lee et al. (2025), is the stochasticity of the decoder. Since the decoder is stochas- tic, a single latent code = can decode to multiple different molecular structures æ, and thereby range of objective function values (see Figure 2(a)). standard GP can only explain such discrepancies additional observation noise or by over-fitting, limiting the utility of predictions across the rest of the space well-known problem that has moti- vated substantial work on providing GP models that support noisy inputs (e.g., McHutchon & Rasmussen, 2011: 1; Qing et al., 2023a).
Zn f argmax P(ğ(z) > D* ZE2
where the surrogate ğ is trained on the dataset D {Zz, 3=) of latent code-evaluation pairs. The result 2n is then decoded and the resulting structure is evaluated the objective function (sec Algorithm 1).
We dedicate the remainder of this section to precisely state the two primary pathological issues caused by this current way of coupling VAES and GPs in LSBO: i) fitting GPs in VAE latent spaces can lead to poor surrogate models, and ii) it challenging to define search space within the latent space so that it reliably targets the most performant regions of the decoder and so leads to the selection and generation of high-quality and useful molecular structures (i.e., avoiding chemically invalid or biologically irrelevant molecules),
In COWBOYS, we avoid alignment issues entirely by depart- ing from using a GP in the latent space, instead only fitting surrogate models directly in the data-space.
3.2. LSBO Surrogate Models can be Poor Predictors
3.3. LSBO Search Spaces are not Uniform Boxes
The premise behind the efficiency of BO is the ability to build an effective surrogate model by exploiting the smooth- ness of our objective f over the search space X. However, as demonstrated by many (e.g., Griffiths& & Hernández-Lobato, 2020; Moss et al., 2020a; Tripp et al., 2020), the mapping g from the latent space to objective function values can be significantly more challenging to learn with GP than f- Consequently, LSBO often ends up with poor surrogate model that under-represents variations, over-emphasises sub-optimal areas of the space, and fails to extrapolate the target property to eeveyat-e structures - a notion typically referred to by the catch-all term "poor latent space
The latent spaces of VAES are unbounded, therefore LSBO typically restricts the search to pre-specified bounded re- gion [-6, 87d C Z (see Algorithm 1). This simplification serves two main purposes. First, limiting the search space facilitates the use of established acquisition function op- timisation methods. Second, by focusing on a central re- gion hypothesized to be more well-behaved, one can limit the presence so-called "dead regions" in the latent space: areas known to yield unrealistic decodings (Griffiths & Hernández-Lobato, 2020). However, for the VAEs typi- cally considered in molecular design problems, we now show that this intuition is entirely incorrect.
3

---

Return of the Latent Space COWBOYS
4. Return of the Latent Space COWBOYS
f(x)
Z
As LSBO, COWBOYS relies on two probabilistic models: the VAE a ~ po(x) for generating likely valid molecules, and a GP-based predictive distribution p(f(x) (x)) for es- timating objective values. However, rather than predicting these values directly from latent codes-an often challeng- ing task (see above) - we leverage GP's proven effec- tiveness in the original structured space using a dataset n- {( y+]R=1 of structure-evaluation pairs and a Tanimoto kernel (as introduced in Section 2.1).
Vd.
X'
Pg
2
(a)
(b)
Figure 2. (a) In LSBO, the same latent input (blue dot) will. via the stochastic decoder (grey box), map to different values in structure space (black dots) and so corresponds to multiple objective func- tion values (red dots) - discrepancy that hinders the learning of problems,
We now summarise Categorical Optimisation With Belief Of underl underlYing Structures (COWBOYSs), Algorithm 2. Notably it avoids any explicit reference to the latent space 2. Despite its simplicity. COWBOYS's formulation allows us to exploit the power of the VAE whilst avoiding LSBO's two core challenges: (i) we do not fit models in the latent space and (ii) we do not use the latent space as a search space. To "focus" the generating distribution po towards promising regions during optimisation, we apply a Bayesian update informed by the GP surrogate. The algorithm proceeds via two main steps, described below for the n* iteration under an initial design of size Ninit-
accurate surrogate models. (b) In higher-dimensions the area of the latent space supported by the prior the VAE (bluc) concentrates in thin circular shell.
"Dead areas" are regions in generative model's 's latent space that lie far from the training data's s encoodigg which themselves populate high- boab-rrbbbbiitt zones under the prior. This concept is often referenced when manipulat- ing latent vectors across wide range of generative models (White, 2016) (for VAEs and Generative Adversarial Net- works (GANs)) and (Song et al., 2020a;b) (for diffusion- based approaches). Outside of the LSBO setting, common strategies for latent space manipulation draw on a classical result which shows that most of high-dimensional Gaus- sian's 's probability mass resides in thin annulus rather than near the centre:
1. Initial design (n S. Ninit) To generate our initial design of size Ninit we use the VAE exactly as it was designed, i.e. by decoding Gaussian latent samples. Therefore for n E {l...., Ninit) } we sample
En ~ po(æ),
(2)
Theorem 3.1 (Gaussian Annulus Theorem. Section 3.3.3 of Vershynin (2018)). Nearly all of the probability mass of standardised Gaussian is concentrated in thin annulus of width ((1) at radius Vd.
resulting in an initial dataset of structure-evaluation pairs over which we can initialise our GP model, with- out the need to specify a clipped search space.
2. Optimisation steps (n > Ninnt) After the initial de- sign, COWBOYS refines its search by only sampling molecules that the surrogate GP predicts will exceed the best observed objective value so far. More con- cretely,
Although the Gaussian Annulus Theorem is well known in machine learning llading to "radius-presern" preserving" methods for latent value manipulation like spherical linear interpola- tion (SLERP) (Shoemake, 1985) and sub-space extraction (Bodin et al., 2024)-its its implications pose a serious chal- lenge to the practice of clipping the latent search space in LSBO when considering VAES with even just moder- ate dimensionality. For example, the encoder of the 128- dimensional VAE used for our second set of experiments (Section 7.2) maps 95% a 5, 000 random subset of its train- ing data to be within a do-dimmensionnl spherical shell of thickness 0.06. Even advanced LSBO approaches, such as adaptive centring of the search space (Maus et al., 2022a),, still rely on box-shaped regions, making it impossible to target only this high-probability "shell" where the VAE en- coder produces its most useful structures.
Iy, po(x=fff f", D.").
(3)
where fæ D is a random variable following the GP's posterior predictive distribution for the objective value x, and f" denotes the highest observed objective value far.
To build intuition for COWBOYS, one might consider a simple, though computationally prohibitive, approach via rejection sampling: generate large number of candidate structures from the VAE and retain only those satisfying fn > [* (under realisation of fal DD
In COWBOYS we can avoid the need to define a search space entirely, instead proposing a crmpling-based strategy rather than one requiring acquisition optimisation.
Batch BO. Note that an arbitrary number of samples can be drawn in parallel from (3), making COWBOYS well-suited for the larch-batch evaluations often tackled in Bayesian
4

---

Return of the Latent Space COWBOYS
Algorithm 2 COWBOYS
argmax, argmax,PP(<<1))) In other words, we approximate the VAE's decoder distribution with the delta function:
Input: Budget N, init size Mint
for (1,... J do
po(x))) 5(x he (z)),,
if < Nint then
< initial design
In ~ (z)
4 vanilla VAE
i.e. deterministic decoding strategy that selects only the most-likely decoded molecule from each latent location z.
else
<I sequential optimisation
Note that we are not the first to replace the VAE's probabiiss tic decoding with the most-likely mapping (e.g.. González- Duque et al. (2024)). Removing this source of variation can, at the risk of reducing the diversity of candidate sam- ples, mitigate the alignment issues discussed in Section 3.2. In our experiments, however, we found that COWBOYS did not suffer empirically from this reduced variation (see Appendix B). We also stress that the VAE is still initially trained with a stochastic decoder, and that it is only when performing BO that we take the most likely mapping.
Un ~ papalal (a 8", Dx-1 1 GP-conditioned VAE
end if
Evaluate new structure yn + f(x.) Update dataset D + D Uf(zn,
Fit structured space GP on D end for return Believed optimum across { Fnl
optimisation (Vakili et al., 2021) and active learning (Ober et al., 2025).
Once we have defined the mapping hg : Z + X, each latent variable z E Z now defines corresponding structure E X. Consequently, rather than sampling new structure directly, we can reframe the sampling task in (3) as sampling a new latent value z from an appropriate distribution and then mapping = through hg to obtain the new structure. Specifically, under the deterministic approximation of the decoder distribution, sampling æ from (3) is equivalent to the following two step procedure:
5. Practical Sampling for COWBOYS
Although the GP-conditioned VAE is conceptually ap- pealing way to focus on increasingly narrower, relevant proportions of molecular space, sampling from (3) remains non-trivial. Indeed, the naive rejection sampler that uses the VAE (x) as a proposal becomes prohibitively ineffi- cient once the region of the search space likely to satisfy f f% narrows, which typically occurs in the early stages of the optimisation process. Consequently, we recommend leveraging more advanced sampling algorithms from the computational statistics literature.
z p(z|g0,= > f (.D)
(4)
x hg(2),
where D is the random variable following the GP posterior predictive distribution for the objective value at the decoded structure æ he h=()), i.c. g0.x = (ho(z)).
In the remainder of this section, we illustrate how a popular Markov Chain Monte Carlo (Brooks et al., 2011, MCMC) approach can be used to approximate COWBOYS. Specifi- cally, we demonstrate that sampling from COWBOYS' GP- conditioned Variational AutoEncoder ( (VAE) can be refor- mulated as sampling from a posterior distribution induced by Gaussian prior and a corresponding likelihood - a well-established setting for which effective MCMC algo- rithms already exist and can be readily applied.
By decomposing (4) via Bayes" rule
p( g8,= f" D,") x p(ge,x f D p(z), (5)
we can now see that (4) corresponds to sampling from the posterior distribution induced by the VAF's Gaussian prior p(z) and a likelihood p(ga, x > f* D ), which resembles the PI acquisition function (and so can be calculated easily in closed-form under the GP). Hence, our approximation of COWBOYS reduces to a classical form for which many well-established sampling methods are available.
5.1. COWBOYS Inference under a Gaussian Prior
Our practical implementation of COWBOYS approximates the GP-conditioned VAE (3) as the posterior distributions induced by a specific likelihood and (potentially high- dimensional) Gaussian prior. However, this approxima- tion strategy requires deterministic mapping from the la- tent space to the structure space, rather than the stochastic mapping provided by the VAE's decoder. Consequently, COWBOYS only ever considers the most likely decod- ing of the latent variable, inducing a deterministic map- ping between latent codes and structure given by ho(z)-
5.2. Preconditioned Crank-Nicolson MCMC
For our experiments, we sample from (4) using a modifi- cation of the Preconditioned Crcos-Nicolson (PCN) algo- rithm (Cotter et al., 2013). PCN is well-suited to sampling from medium- to high-dimensional Gaussian priors, as it concentrates on the annulus of high-prior probability (as discussed in Section 3.3) rather than attempting to explore the entire parameter space. This contrasts with more con- ventional Markov Chain Monte-Carlo (MCMC) algorithms
5

---

Return of the Latent Space COWBOYS
(e.g., random-walk methods), which rapidly deteriorate in ef- fectiveness dimensionality increases (Hairer et al., 2014) an observation closely aligned with our critique of existing LSBO strategies.
Encouraging valid decodings. Custom VAE decoders have been proposed that can ensure the satisfaction of string- based (Kusner et al., 2017) or graph-based constraints (Jin et al., 2018). Alternatively, the suggestion of valid structures, can be encouraged by using an additional model that predicts the validity of proposed structures to help triage the query points suggested by BO (Griffiths & Hernández- Lobato, 2020) or by avoiding areas of the latent space where the decoder's epistemic uncertainty is high (Notin et al., 2021).
Exact details of our implementation of PCN are provided in Appendix A. Note that, in order to leverage parallel hard- ware, we introduced minor alteration that transforms PCN into a non-reversible MCMC algorithm, thereby forfeiting the rigorous sample quality guarantees derived by the com- putational statistics community. However, in the context of COWBOYS, where our ultimate goal is Bayesian optimisa- tion, our empirical findings suggest that producing a set of approximate samples is sufficient.
Additional improvements for LSBO. Extensions support high- nigh-dimensional latent spaces (Maus et al., 2022a), multi- objective (Stanton et al., 2022) and batch optimisation (Maus et al., 2022b), as well as a method to include structure- level knowledge into the surrogate model alongside latent representations (Deshwal & Doppa, 2021). Recent work of (Ramchandran et al., 2025) proposed the use of GP VAEs that allows the inclusion of auxiliary variables to help learn useful latent spaces.
5.3. Computational Complexity
We now analyse the computational overhead incurred by COWBOYS and standard LSBO in non-batch setting. Ig- noring the shared computational cost of fitting GP surrogate models, the th step of COWBOYS algorithm using PCN with C MCMC chains and S MCMC steps has complexity O(CS(N2 + V)). corresponding to the need to make GP prediction O 2 (assuming cached Gram matrix inverse) and VAE decoding ( at cost V) for each point considered by PCN. In contrast, standard LSBO incurs (AN? + V). where A are the number of evaluations required to max- imise its acquisition function. We set up our experiments via choices of C and S so that GP costs are equivalent (C X S A). We stress that although COWBOYS re- quires more VAE evaluations than standard LSBO, advanced LSBO approaches that fine-tune the VAE on incoming data also require large numbers of additional VAE evaluations.
7. Results
We compare the performance of COWBOYS over multiple VAES, against 24 unique algorithms, and on 16 open-source benchmark problems for molecular design. Our experiments replicate the exact experimental setups of Chu et al. (2024), González- -uquee et al. (2024), and Maus et al. (2024); re- sults are shown in Figure 3, Table 1, and Figure 4, respec- tively. See (Brown et al., 2019) for the scientific motivation and clear practical details behind these benchmarks.
General details. Our code is integrated into the bench- marking suite of González-Duque et al. (2024) and is avail- able at on BoTorch (Balandat et al., 2020), GPyTorch (Gard- ner et al., 2018), and lastly CAUCHE (Griffiths et al., 2024) for its Tanimoto and molecule utilities, adopting all default settings for model optimisation and kernel hyperparame- ters. Baseline methods are drawn from their respective open-source implementations. For COWBOYS's MCMC- based sampling strategy (our modified PCN, described in Appendix A), we run 100 MCMC steps over single chain when the batch size is one (B- = 1), and 50 chains when B > 5. This matches the 5, 000 acquisition function evalua- tion budget typically used by TURBO (Eriksson al,, 2019), key component of all other high-performin tines. Appendix B examines COWBOYS' robustness to these MCMC settings and highlights the combined impact of its two primary components by applying COWBOYS style sampling but with latent-space GP model.
We build
6. Baselines & Related Works
Improving "alignment". The primary recent focus of the LSBO literature has been on allowing VAEs to be fine- tuned during optimisation, encouraging the alignment of latent spaces with respect to the optimisation objective, with Eissman et al. (2018), Tripp et al. (2020), Moriconi et al. (2020). Grosnit et al. (2021), Maus et al. (2022a), and Chen et al. (2024) proposing methods for fine-tuning the latent space mappings- using additional supervised neural net- works, re-weighted losses, performing joint inference over the surrogate model/mapping. Most recently, Lee et al. (2023) fine-tune VAEs with the explicit goal of ensuring the optimisation objective is smooth (in a Lipschitz sense) across the latent space- - an approach extended further in Chu et al. (2024) to also address the additional alignment issue arising from stochasticity in the VAE's decoder. Here, LSBO is recast as an inverse problem, where they use opti- misation identify the latent codes most likely to decode to any suggested structures before fitting surrogate models.
LSBO rou-
7.1. Comparison with State of the Art LSBO Methods
Baselines. Following Chu et al. (2024), we consider LOL- BO (Maus et al., 2022a), CoBO (Lee et al., 2024), W-LBO
6

---

Return of the Latent Space COWBOYS
(Tripp al., 2020), InvBO (Chu et al., 2024), and PG-LBO (Chen al., 2024), all of which require fine-tuning the VAE during optimisation. We also evaluate the standard LSBO method (Algorithm 1) and its trust region variant TURBO- I. (Eriksson et al., 2019). As a reference point, we include Graph-GA (Jensen, 2019), a BO-free baseline that uses a widely adopted graph- genetic algorithm. For more detailed discussion of these baselines, see Chu et al. (2024), Note that we used a pre-trained VAE provided via direct correspondence from the authors of (Maus et al., 2022a), as the one provided in their open-source implementation was not sufficient to recreate their results.
0.22
0.600
Mean Reward
0.575
0.550
0.20
Results. As shown in Figure 3, COWBOYS achieves marked improvements in sample efficiency over existing LSBO methods the low-data regime, outperforming even those permitted to fine-tune their VAE. We replicate the experimental setup from Maus et al. (2022a;b); Lee et al. (2024); Chu et al. (2024), focusing on lower-budget prob- lems (COWBOYS cannot currently handle evaluation bud- gets of 80,000 due to a current lack of availability of seal- able GP models for discrete data; see Section 8). Across six tasks introduced by Brown et al. (2019), we run 100 BO steps with batch size B = 5, starting from 100 rep- resentative molecules as defined by Lee et al. (2024), scenario designed to mimic real-world lead optimisation. These experiments employ the SELFIES VAE of Maus et al. (2022a), which encodes and decodes molecules into 256- dimensional latent space using six transformer layers (see Maus et al., 2022a, for details).
0.525
0.19
0.500
200
400
600
200
400
600
(a) Osimertinib MPO
(b) Median Molecules 2
0.84
0.575
Mean Reward
0.550
0.82
0.525
0.80
0.500
0.78
0.475
200
400
600
200
400
600
(c)
Amlodipine MPO
(d) Perindopril MPO
0.80
0.50
Meen Reward
0.75
0.45
0.70
0.40
7.2. High-dimensional BO Benchmark
0.35
0.65
600
200
Baselines. Following the setup in González-Duque et al. (2024), we compare COWBOYS to three categories of meth- ods: (i) non-BO evolutionary optimisers applied to the VAE latent space: hill-climbing, CMAES, and a Genetic Algo- rithm (GA), (ii) high-dimensional BO methods also operat- ing in the VAE latent space: RandomLine BO (Kirschner et al., 2019) and TURBO (Friksson et al., 2019), and (iii) BO methods that optimise directly in the discrete (molecu- lar) space: Bounce (Papemmeier et al., 2023) and ProbRep (Daulton al., 2022). For more detailed discussion of these baselines, see the online documentation in González- Duque et al. (2024).
Number of Gracle Calls
Number Oracle Calls
(c) Ranolazine MPO
(f) Zaleplon MPO
COWBOYS
Graph-GA
LOL-BO -
PG-LBO
W-LBO
CoBO
InvBO
LS DO TURBO-L
Figure 3. Average performance (#: standard error) of COWBOYS over 10 runs on problems considered by (Chu et al., 2024). COW- BOYS achieves substantial improvement sample efficiency over all existing LSBO methods (including those able to fine-tune VAES to incoming data) in this low data regime, with only the recently proposed InvBO sometimes matching its performance. We stress that COWBOYS does not fine-tune its VAE during opti- misation, rather just uses it more efficiently.
Results. Table 1 shows that COWBOYS consistently out- performs all compared methods on the high-dimensional discrete sequence optimisation tasks provided by González- Duque al. (2024), (based on the PMO benchmark of Gao et al. (2022)). All algorithms start from an initial design of 10 molecules and run for 300 BO steps with batch size of B = 1. Because the original VAE used by Gonzâlez-Duque et al. (2024) had relatively poor reconstruction accuracy, we reran the entire benchmark (and all baselines) with an

---

Return of the Latent Space COWBOYS
updated, more expressive, 22-dimmenioonl fully-connected VAE using the latest release of the benchmarking suite.
0.80
0.80
Mean Reward
0.75
7.3. Comparison with Traditional LSBB
0.70
0.65
Baselines. In our final set of experiments, we restrict our focus to methods that do not fine-tune their VAE. Specifi- cally, we run COWBOYS and LSBO/TURBO-L with exact Gaussian Processes for 100 steps of B = 20 evaluations, a limit imposed by the computational demands of exact GP EI
0.65
0.60
0.80
10
10a
10a
0.33
10
10a
10*
(a) Osimertinib MPO
(b) Fexofenadine MPO
0.200
inference. We also include scalable variants-EUUI and EULBO KG (Maus et al., 2024), Vanilla BO (Hvarfner et al., 2024), and IPA (Moss al., 2022: 2023)-thhee rely on Sparse Variational Gaussian Processes (SVGP) (Hensman et al., 2013) to extend LSBO to 4,000 steps. This allows us to examine whether significantly larger optimisation budgets enable these methods to match COWBOYS' performance. Extending COWBOYS to accommodate misation is an intriguing avenue for future work (Section 8). Unlike above, here we used the exact pre-trained VAE provided the open-source implementation of (Maus al.,
0.30
Mean Reward
0.175
0.25
0.150
0.20
0.125
0.15
0.100
10a
10a
10
10a
102
10
Number of Oracle Calls
Number Oracle
(c) Median Molecules 1
(d) Median Molecules 2
opti-
COWBOYS
TURBO-L (IPA El)
LSBO [Vanilla EI)
TURBO-L (BULBO
TURBO-L (Exact
LSBD El)
TURBO-L (EULBO KG)
LSBO (EULBO EI)
LSBO (Exact EI)
TURBO-L (Vanilla EI)
LS80 (EULBO KG)
2022a) for all approaches.
Figure 4. Average performance (+: standard error) over 20 repeti- tions with an log-scaled x-axis, demonstrating that, among LSBO methods that cannot fine-tune their latent space, COWBOYS pro- vides significant improvemnt in efficiency.
Results. As shown in Figure 4, COWBOYS achieves sub- stantially better performance than other LSBO methods that also do not fine-tune the VAE. Even when given a two orders-of-magnine marrnntuu- increase in optimisation budget, these baselines do not catch up COWBOYS, highlighting the value of our proposed method of combining VAES and GPs for Bayesian optimisation. These experiments replicate those in Maus et al. (2024), who perform LSBO using the latent space of a SELFIES VAE (see Section 7.1), focusing on a subset of the molecular design benchmarks discussed carlier.
where there is a rich, yet underused, literature. For instance, the Tanimoto kernel, originally proposed by chemists, en- codes prior knowledge of what sort of molecular attributes are important in a way that can be especially powerful in low-data regimes, Similar structural kernels exist for vari- ous structured objects and, via COWBOYS, could now be used for molecular graph optimisation with graph kernels (Schraudolph et al., 2010). enginecring design with 3d mesh kernels (Perez et al., 2024), optimising computer code via tree kernels (Beck et al., 2015), or protein design with new protein kernels (Groth et al., 2024). By enabling the direct use of such kernels, COWBOYS allows practitioners to in- corporate wealth of comain-specifte prior knowledge. We hope our work helps revitalize interest in hamnessing these powerful kernels across a range of complex design problems and spurs further synergy between rich generative models and carefully structured discriminative models.
8. Discussion and Limitations
COWBOYS leverages a simple Bayesian updating mecha- nism that links separately trained VAEs and GPs through a sampling strategy progressively biased toward promising structures. Despite its simplicity, this sampling- baaeid ap- proach directly addresses LSBO's two primary challenges. First, COWBOYS does not model the objective in the VAE latent space; instead, it retains the GP in the original, struc- tured representation of the data, allowing the use of spe- cialized kernels. Second, COWBOYS does not employ the latent space for optimisation; instead, it naturally restricts its focus to regions that are both plausible under the VAE prior and potentially high-performing under the GP. We demon- strate that COWBOYS offers a substantial boost in sample efficiency low-budget molecular design benchmarks.
Methodological extensions. Extending COWBOYS to the important task of large-scale molecular optimisation raises several exciting avenues for future research. First, we need to develop methods for fitting sparse GPs discrete molec- ular data- inherently more challenging task- where inducing points are more difficult to optimise, perhaps by applying techniques from Burt et al. (2019) or Chang et al. (2022). Second, integrating COWBOYS with LSBO tech- niques that fine-tune their VAEs during optimisation may become essential for performance once experimental bud-
Extensions into new application areas. By working in the raw structure space, our approach naturally supports specialized kernels tailored to particular domains-an an area
8

---

Return of the Latent Space COWBOYS
Objective function
COWBOYS
CMAES
GA
RandorLineR0
TURS0
Bounce
ProbRep
albuterol similarity
0.472=0.08
0.487#0.06
0.453=00004
0.356=0.05
0.454+0.06
0,456#0.04
0.16 0.21 # 0.03
amlodipine mpo
0.477=0.04
0.449=0.02
0.458=0.02
0,440=0.01
0.453=0.02
0.444+0.02
0.00 0.00 0.00
0.00
palcooxib_rediscove
0.217=0.02
0.202=0.01
0.213=0.02
0.202+0.01
0.204+0.01
0,202+0.01
0.02:0000
0.02 # 0.00
deco_hop
0.570=0.01
0.562=#001
0.563=00.0
0.5622+0.01
0.5644#001
0.562#0.01
0.50 0.51 # 0.00
drd2.docking
0.342=0.11 17
0.097=0.09
0.087=0.09
0.076=0.10
0,34640.37
0.170=0.10
0.01 0.00 0.01 0.00
fexofenadine gsk3 beta
0.68220.00
0.644 04
0.652=0.05
0.578+0.06
0.66940.03
0.63240.06
0.13 0.20 = 0.08
0.368 03
0.252=0.05
0.321#0.11
0.241 0.44
0.170=#0.05
0,302+0.05
0 0,12 0.02
isomer c7h8n2o2
1.000=0.00
0.788=0.10
0.872=0.07
0.812+0.10
0.864=0.15
0.11 0,880=0.07
0.24 0.11
isomer c9h 1 On202pf2el
0.690=0.08
0.600=0.12 12
0.567=0.11
0.639=0.07
0.65420.12
0.61840.24 24 0.01 0.06 0.03
jnk3
0. 0.122=0.044 0.138=0.05
0.136=0.02
0.130##0.04
0.120.0.00
0.05 0.092+0.03
0.06 0.01
median
0. 0.1
0.162=00.01
0.187=0.03
031
0. 0..4000.03
0.154+0.04
0.03 + 0.01 =
median_2
0.162=0..1
0. 0.146=0.01
0. 0.149
0.145=0.01
0.145500.01
0.145+0.01
0.01 0.00 0.01 0.00
mestranol_similarin
0.38420.01
0.348=0.02
0.362#0.02
0.266+0.03
0.353+00001
0.378+0.07
0:01 0.00 0.00
osimetrinib_mpp
0.722=0.03
0.718=0.02
0.711=0.05
0.707+0.04
0.714
0,30 0.711=0.03
0,59 0.04
perindopril_mpg mpo
0.382=0.04 04
0,326=0.11 11
0.335=0.11
0.334+0.12
0.337 11
0.362+0.06
0.00 0.00 0,00
ranoluzine mpo
0.648=0.00 04
0.609=0.02
0.625=0.08
0,484+0.06
0.632+0.03
0.597+0.04
0.00 0.11 0.02
rdkit logp
10.678+0.99
11.38336.12
14.474+7.41
9.307+1.67
19.48445.92
9.691+0.71
3.12 1.20 5.49 3.01
ged sa .dc
0.909=0.03
0.902=0003
0.912=0.03
0.902+0.03
0,906*0.04
0,903+0.04
0.52 + 0.09 0.60 + 0.05
7.909=0.07
7.920=0.06
8.87840. 16 7.516
7.877
7.952+0.05
8.36 0.46 8.59 0.13
scoffold_hop
0,43520.01
0.427=0.01
0.430=0.01
0.427+0.01
0.427
0.429#0.01
0.34 a: 0.01 0.34 0.00
sitagliptin mpo
0.259 007
0.1
0.237 12
0.15440.00 09
0.233+0.12
0.163+0. 12
0.00 4: 0.00 0.00 0.00
thiothixene thiothixene_redssovvee rediscovery
0.247=0.03
0.224=0.02
0.224+0.02
0.224
0.226*0.02
0.224 02
0.02 + 0,01 0.03 + 0.01
toglitazone
0.205=0.02
0.187=0.02
0.203=0.03
0.187+0.02
0.188+0.02
0.195#0.01
0.02 + 0.01 0.02 + 0.00
valsartan smarts
0.000=0.00
000000=0000
0.000=0.00
0.000=0.00
0.000=0.00
0.000+0.00
0.00 = 0.00
0.00 = 0.00
zaleplon_mpe mpo
0.379=0.05
0.337=0.08
0.293=0.12
0.297+0.11
0.344
0.297=0.11
0.00 = 0.00 0.00 # 0.00
Table Average performance (+: s.d.) over repetitions of COWBOYS on the discrete BO benehmarking suite of (González- -uzzze et al., 2024). We stress best average scores achieved after 300 evaluations (dark) and scores within single standard deviation of best (light).
gets grows and VAE adaptation becomes more feasible. Finally, more sophisticated sampling algorithms could fur- ther improve MCMC efficiency (as is important for control- ling the cost of COWBOYS under expensive VAE archi- tectures) or allow us to exploit the full stochastic nature of the decoder, rather than restricting the search to the most likely decoding. Also, note that the condition in 3 does resemble the well- Probability of Improvement (PI) acquisition function. While PI is intuitive, it is known to be suboptimal and can exhibit pathological behaviour and so adapting more sophisticated acquisition schemes into a COWBOYS framework is an exciting direction for future research, e-g.
machine learning. Real-world benefits in molecule design include significantly reducing wet-lab resources. However, the same method could be misused in non-therapeutic ap- plications. Thus, while optimisation is automatic, humans must decide how these results are applied, following appro- priate ethical guidelines.
References
Andrieu, C. and Thoms, J. A tutorial on adaptive meme.
Statistics and computing, 18:343-373, 2008.
Auer, P. Using upper confidence bounds for online learning. In Proceedings 41 annual symposium on foundations of computer science, pp. 270-279, IEEE, 2000.
schemes for multi multi--idelity (Moss et al., 2020b), multi-objective (Qing et al., 2023b) and tackling more ambitious problems like quantile optimisation (Picheny et al., 2022).
Balandat, M., Karrer, B., Jiang, D.,, Daulton, S., Letham, B., Wilson, A. G., and Bakshy, E. Botorch: A framework for efficient monte-carlo bayesian optimization. Advances in neural information processing systems, 33:21524-21533 2020.
Acknowledgments
This research was supported by AstraZeneca, Schmidt Sci- ences, and Lancaster University's 's Mathematics for AI in Real-world Systems E3 grant.
Beck, D., Cohn, T., Hardmeier, C., and Specia, I.. Learning structural kernels for natural language processing. Trans- actions of the Association for Computational Linguistics, 2015.
Impact Statement
Our work introduces an efficient framework for optimising over discrete sequences, a problem of growing interest in
Bergstra, J., Bardenet, R., Bengio, Y., and Kégl, B. Algo- rithms for hyper-parameter optimization. Advances in neural information processing systems, 24, 2011.

---

Return of the Latent Space COWBOYS
Bodin, E., Moss, H., and Ek, C. H. Linear combinations of latents in diffusion models: interpolation and beyond. arXiv e-prints, pp. arXiv-2408, 2024.
Gao, W., Fu, T., Sun, J., and Coley, C. Sample efficieney matters: benchmark for practical molecular optimiza- tion. Advances in neural information processing systems, 2022.
Brooks, S., Gelman, A., Jones, G., and Meng, X.-L. Hand- book of markov chain monte carlo. CRC press, 2011.
Gardner, J., Pleiss, G., Weinberger, K. Q., Bindel, D., and Wilson, G. Gpytorch: Blackbox matrix- gaus- sian process inference with gpu acceleration. Advances in neural information processing systems, 31, 2018.
Brown, N., Fiscato, M., Segler, M. H., and Vaucher, A. C. Guacamol: benchmarking models for de novo molecular design. Journal of chemical information and modeling, 2019.
R., Wei, J. N., Duvenaud, D., Hernández-Lobato J. M., Sánchez-l Lengeling B., She- berla, D., Iparrauurr,, J., Hirzel, T. D., Adams, R. P. and Aspuru-Guzik, A. Automatic chemical de- sign using a data-driven continuous representation of molecules. ACS central science, 2018.
Burt, D., Rasmussen, C. E., and Van Der Wilk, M. Rates of convergence for sparse variational gaussian process re- gression. International Conference on Machine Learn- ing. 2019.
Chang, P. E., Verma, P., John, S., Picheny, V., Moss, H., and Solin, A. Fantasizing with dual gps in bayesian optimization and active learning, arXiv preprint arXiv:2211.010533 2022.
González-Duque, M., Michael, R., Bartels, S.,
Zainchkovskyy. Y. Hauberg, S., and Boomsma, W. A survey and benchmark high-dimensions bayesian optimization of discrete sequences. arXiv preprint arXiv:2406.047336 2024.
Chen, T., Duan, Y., Li, D., Qi, L., Shi, Y., and Gao, Y. Pg-lbo: Enhancing high-ding nigh-dimensional bayesian optimiza- tion with pseudo-labe label and gaussian process guidance. In Proceedings of the AAAI Conference on Artificial Intelli- gence, 2024.
Griffiths, R.-R. and Hernández-Lobate J. M. Constrained Bayesian optimization for automatic chemical design us- ing variational autoencoders. Chemical science, 2020.
Griffiths, R.-R., Greenfield, I.., Thawani, A. R., Jamash, R., Moss, B., Bourached, A., Jones, P., McCorkin- dale, W., Aldrick, A. A., Fuchter, M. J., al. Data-driven discovery of molecular photoswitches with multioutput gaussian processes. Chemical Science, 2022.
Chu, J., Park, J., Lee, S., and Kim, H. J. Inversion-basee latent Bayesian optimization. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Cotter, S., Roberts, G., Stuart, A., and White, D. Meme methods for functions: Modifying old algorithms make them faster. Statistical Science, 2013.
Griffiths, R.-R., Klarner, L., Moss, H., Ravuri, A., Truong, S., Du, Y., Stanton, S., Tom, G., Rankovic, B., Jamasb, A., et al. GAUCHE: a library for Gaussian processes in chemistry. Advances Neural Information Processing Systems, 36, 2024.
Daulton, S., Wan. X., Eriksson, D., Balandat, M., Osborne, M. A., and Bakshy, E. Bayesian optimization over dis- crete and mixed spaces via probabilistic reparameteriza- tion. Advances in Neural Information Processing Systems, 2022.
Grosnit, A., Tutunov, R., Maraval, A. M., Griffiths, R.-R., Cowen-Rivers, A. I., Yang. I.., Zhu, L.., Lyu, W., Chen, Z., Wang. J., et al. High-dimensiona Bayesian optimisation with variational autoencoders and deep metric learning. arXiv preprint 1rXiv:2106.036/// 2021.
Deshwal, A. and Doppa, J. Combining latent space and structured kernels for Bayesian optimization over combi- natorial spaces. Advances in Newral Information Process- ing Systems, 34:8185 88-0000, 2021.
Groth, P. M., Kern, M., Olsen, L., Salomon, J., and Boomsma, W. Kermut: Composite kernel regression for protein variant effects. Advances in Neural Information Processing Systems, 37, 2024.
Eissman, S., Levy, D., Shu, R., Bartzsch, S., and Ermon, S. Bayesian optimization and attribute adjustment. In Proc. 34th Conference on Uncertainty in Artificial Intelligence, 2018.
Hairer, M., Stuart, A., and Vollmer, S. Spectral gaps for Metropolis- Hastings algorithm in infinite dimensions. The Annals of Applied Probability, 2014.
Eriksson, D., Pearce, M., Gardner, J., Turner, R. D., and Poloczek, M. Scalable global optimization via local bayesian optimization. Advances in neural information processing systems, 32, 2019.
Hennig, P. and Schuler, C. Entropy search for information- efficient global optimization. Journal of Machine Learn ing Research, 2012.
10

---

Return of the Latent Space COWBOYS
Hensman, J., Fusi, N., and Lawrence, N. D. Gaussian processes for big data. In Proceedings of the Twenty- Ninth Conference on Uncertainty in Artificial Intelligence, pp. 282-290, 2013.
Maus, N., Jones, H., Moore, J., Kusner, M. J., Bradshaw, J., and Gardner, J. Local latent space Bayesian optimization over structured inputs. Advances in Neural Information Processing Systems, 2022a.
Hvarfner, C., Stoll, D., Souza, A., Lindauer, M., Hutter, F., and Nardi, L. abo: Augmenting acquisition functions with user beliefs for bayesian optimization. In Tenth International Conference of Learning Representations, ICLR 2022, 2022.
Maus, N., Wu, K., Eriksson, D., and Gardner, J. Discovering many diverse solutions with Bayesian optimization. arXiv preprint arXiv:22 20. 10953, 2022b.
Maus, N., Kim, K., Pleiss, G., Eriksson, D., Cunningham, J. P., and Gardner, R. Approximation-awarre bayesian optimization. arXiv preprint wrxiv:2406.043066 2024.
Hvarfner, C., Hellsten, E. O., and Nardi, L. Vanilla bayesian optimization performs great in high dimension. arXiv preprint ar xrXiv:2402.02229, 2024.
McHutchon, A. and Rasmussen, C. Gaussian process train- ing with input noise. Advances in neural information processing systems, 2011.
Jensen, J. H. A graph-based genetic algorithm and gener- ative model/monte carlo tree search for the exploration of chemical space. Chemical science, 10(1))3357-3357, 2019.
Mockus, J. The Bayesian approach to global optimization. In System Modeling and Optimization: Proceedings of the 10th IFIP Conference New York City, USA, August 31 31-September 4, 1981, pp. 473-481. Springer, 2005.
Jin, W., Barzilay, R., and Jaakkola, T. Junction tree vari- ational autoencoder for molecular graph generation. In International conference on machine learning, 2018.
Moriconi, R., Deisenroth, M. P., and Sesh Kumar, K. High-dimensional Bayesian optimization using low- dimensional feature spaces. Machine Learning, 2020.
Kingma, D. P. and Welling, M. Auto-encoding variational Bayes. In Proceedings of the International Conference on Learning Representations (ICLR), 2014.
Moss, H., Leslie, D., Beck, D., Gonzalez, J., and Rayson, P. Boss: Bayesian optimization over string spaces. Ad- vances in neural information processing systems, 2020a.
Kirschner, J., Mutny, M., Hiller, N., Ischebeck, R., and Krause, A. Adaptive and safe bayesian optimization in high dimensions via oie-dimeniional subspaces. In International Conference Machine Learning, pp. 3429- 3438. PMLR, 2019.
Moss, H. B. and Griffiths, R.-R. Gaussian process molecule property prediction with flowmo. arXiv preprint arXiv:2010.011 118 2020.
Kriege, N. M., Johansson, F. D., and Morris, C. survey on graph kernels. Applied Network Science, 2020.
Moss, H. B., Leslie, D. S., and Rayson, P. Mumbo: Multi- task max-value bayesian optimization. In Joint European Conference Machine Learning and Knowledge Dis- covery in Databases, 2020b.
Kusner, M. J., Paige, B., and Hemándezz-obbtoo J. M. Grammar variational autoencoder. In International con- ference on machine learning, pp. 1945-1954. PMLR, 2017.
Moss, H. B., Leslie, D. S., Gonzalez, J., and Rayson, P. Gibbon: General-purpose information-based bayesian optimisation. Journal of Machine Learning Research, 2021.
Landrum, G. Rdkit documentation. Release, 2013.
Lee, S., Chu, J., Kim, S., Ko, J., and Kim, H. J. Advanc- ing Bayesian optimization via learning correlated latent space. In Thirth-seventh Conference on Neural Informa- fion Processing Systems, 2023.
Moss, H. B., Ober, S. W., and Picheny, V. Information- theoretic inducing point placement for throuuhr bayesian optimisation. arXiv preprint arXiv:2206.00 024377 2022.
Lee, S., Chu, J., Kim, S., Ko, J., and Kim, J. Advanc- ing Bayesian optimization via learning correlated latent space. Advances in Neural Information Processing Sys- tems, 2024.
Moss, H. B., Ober, S. W., and Picheny, V. Inducing point al- location for sparse gaussian processes in high-throughput bayesian optimisation. In International Conference on Artificial Intelligence and Statistics, 2023.
Lee, S., Park, J., Chu, J., Yoon, M., and Kim, H. J. La- tent bayesian optimization via autoregressive normaliz- ing flows. In The Thirteenth International Conference on Learning Representations, 2025.
Notin, P., Hermández-Lobato, J. M., and Gal, Y. Improving black-box optimization in vae latent space using decoder uncertainty. Advances in Neural Information Processing Systems, 2021.

---

Return of the Latent Space COWBOYS
Ober, S. W., Rasmussen, C. E., and van der Wilk, M. The promises and pitfalls of deep kernel learning. In Uncer- tainty in Artificial Intelligence, 2021.
Song, Y., Sohl-Dickstein, J., Kingma, D. Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201..3356, 2020b.
Ober, S. W., Power, S., Diethe, T., and Moss, H. B. Big batch Bayesian active learning by considering predictive probabilities. arXiv preprint arXiv:2501.05225 2025.
Stanton, S., Maddox, W., Gruver, N., Maffettone, De- laney, E.,, Greenside, P., and Wilson, A. G. Accelerating Bayesian optimization for biological sequence design with denoising autoencoders. In International Confer- ence on Machine Learning. 2022.
Papenmeier, L., Nardi, L., and Poloczek, M. Bounce: reli- able -igheniinniional bayesian optimization for combi- natorial and mixed spaces. Advances in Neural Informa- tion Processing Systems, 2023.
Tripp. and Hernández-Lobato, M. Diagnosing and fixing common problems in bayesian optimization for molecule design. arXiv e-prints, pp. arXiv 2406, 2024.
Perez, R. C., Da Veiga, S., Garnier, J., and Staber, B. Gaus- sian process regression with sliced wasserstein weisfeiler- lehman graph kernels. In International Conference on Artificial Intelligence and Statistics, 2024.
Tripp. A., Daxberger, E., and Hemández-Lobaul J. M. Sample-efficient optimization in the latent space of deep generative models via weighted retraining. Advances in Neural Information Processing Systems, 2020.
Picheny, V., Moss, H., Torossian, L., and Durrande, N. Bayesian quantile and expectile optimisation. In Uncer- tainty in Artificial Intelligence, pp. 1623-1633. PMLR, 2022.
Tripp, A., Bacallado, Singh, S., and Hernández- J. M. Tanimoto random features for scalable molecular machine learning. In Advances in Neural Information Processing Systems, 2023.
Qing, J., Couckuyt, I., and Dhaene, T. A robust multi- objective Bayesian optimization framework considering imput uncertainty. Journal of Global Optimization, 2023a.
Vakili, S., Moss, H., Artemev, A., Dutordoir, V., and
Picheny, V. Scalable thompson sampling using sparse gaussian process models. Advances in neural information processing systems, 14:5631-566433 2021.
Qing, J., Moss, B., Dhaene, T., and Couckuyt, I. Pf2es: Parallel feasible pareto frontier entropy search for multi- objective bayesian optimization. In International Confer- ence on Artificial Intelligence and Statistics, 2023b.
Vershynin, R. High-dimensional probability: An introduc- tion with applications data science, volume 47. Cam- bridge university press, 2018
Ramchandran, S., Haussmann, M., and Lähdesmäki, H. High-dimensional bayesian optimisation with gaussian process prior variational autoencoders. In The Thirteenth International Conference on Learning Representations, 2025.
White, T. Sampling generative networks. arXiv preprint arXiv: /609 1609.04468, 2016.
Rasmussen, C. E. Gaussian processes in machine learn- ing. In Summer school machine learning, pp. 63-71. Springer, 2003.
Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep gen- crative models. In Proceedings of the 31st International Conference on Machine Learning. 2014.
Schraudolph, N. N., Kondor, R., and Borgwardt, K. M. Graph kernels. The Journal of Machine Learning Re- search, 2010.
Shoemake, K. Animating rotation with quaternion curves. In Proceedings of the 12th annual conference on Computer graphics and interactive techniques, pp. 245-254, 1985.
Song, J., Meng, C., and Ermon, S. Denoising diffusion im- plicit models. arXiv preprint arXiv:2010.02502,, 2020a.
12

---

Return of the Latent Space COWBOYS
A. COWBOYS' sampler details
We now provide full details for our PCN MCMC sampler (Cotter et al., 2013) as Algorithms and 4 which we use to get a set of samples from COWBOYS sampling objective (4). Note that we keep track of the latent value that decoded give the best structure so far Øbest and use this start our chains. We also employ an adaptive choice of PCN's scale parameter 3 following Andrieu & Thoms (2008). Often, especially early in the optimisation before the region satisfying likely to satisfy [a) > I narrows, our PCN sampler can return a larger number of samples than our desired batch size. In this case, we select the B samples that achieve the highest expected utility under greedy maximisation of the batch Expected Improvement (qEI) acquisition function provided in BoTorch (Balandat et al., 2020). If very large batches are required, then a subset could be chosen using Thompson following Vakili et al. (2021). Although this subset selection strategy breaks MCMC's detailed balance, it ensures that we return diverse, high-value subset of samples to COWBOYS samples that are likely be approximate anyway due relatively small number of MCMC iterations and lack of burn-in phase. Note that in practical implementations, our PCN (Algorithm 4) is trivially parallelised by running all chains present in the outer loop at once.
Algorithm 3 COWBOYS' MCMC Sampler
Input: Batch size B, number chains C, number steps S, threshold f", starting latent code Z samples t PCN(C,,,,, (bbes))
4 Algorithm Keep sampling if needed Algorithm 4 again
while 7samples B do
Z samples Zsamples SP, Zbeat) end while
if Z sampieal B then
samples E argmaxzCa aammier (EI(Z)
< choose B best according to batch utility
|
end if return Zehosen
Algorithm 4 PCN
Input: number chains C, number steps S, best score so far f", starting latent code
Zsamples 6-8 0
for c E {1,..,C) C} do
8 0.1
< Initialise step size Initialise new chain at best far Do MCMC step
Zcurrent <- Zhest
for E {1, S} do
2proposal Na
Be Zeurrent, T2
Sample from PCN proposal density
P 90 Fampuall > / 1D x pl (ppppppaa)))
a - min
1,
Calculate acceptance probability using prior and likelihood from (5)
p #cueeee > 8 1 x xp(zmmmm)
u ~ 1)
if a > u then
< Stochastic acceptance step
Zeurrent Zproposal Z zamples U ] }
Store sample
end if
3 + 3 + 0.1 - 0.243)
4 Adaptive MCMC update
end for end for return Zsamples
return all accepted samples
13

---

Return of the Latent Space COWBOYS
COWBOYS
COWBOYS
COWBOYS
COWBOYS
COWBOYS
COWBOYS
with Latent GP
Oracle
10 chains
1 chains
10 chains
chains
chains
10 chains
100 steps
100 steps
50 steps
50 steps
10 steps
100 steps
simiterol
0.454+/-0.06 0.472+/-0.08
0.450+/-0.06
0.464+/-0.07 0.435+/-0.05
0.459+/-0.09
amlodipine mpo
0.461+/-0.02
0.477+/-0.04
0.471+/-0.02
0.450+/-0.01
0.461+/-0.02
0,457+/-0.01
celecoxib. rediscovery
0.264+/-0.04
0.217+/-0.02
0.240+/-0.03
0.224+/-0.02
0.240+/-0.04
0.202+/-0.00
_ho
0.569+/-0.01
0.570+/-0.01
0.572+/-0.02
0.562+/-0.01
0.568+/-0.01
0.562+/-0.01
drd2 docking
0.214+/-0.23
0.342+/-0.17
0.325+/-0.17
0.217+/-0.21
0.266+/-0.08
0.241++-0..6
fexofenadine_mpo
0.671+/-0.02
0.682+/00.02
0.681+/00.02
0.677+/-0.01
0.698+/0..02
0.674+/-0.03
gsk3 beta
0.410+/-0.02
0.368+/-0.03
0.358+/-0.06
0.354+/-0.10
0.322+/-0.05
0,290+/-0.03
isomer_c7h8n202
0.934+/-0.05
1.000+/-0000
0.889+/-0.06
0.976+/-0.05
0.921+/-0.06 0.880+/-0.07
isomer.c9h10n222pp2d
0.634+/-0.06
0.690+/-0.08
0.653+/-0.05
0.634+/-0.08
0.649+/-0.04 0.622+/-0.05
jnk3
0.138+/-0.01
0.138+/-0.05
0.1244/-0.02
0.146+/-0.02
0. 0.144+/-0.01
0.134+/-0.03
median 1
0.211+/-0.22
0.198+/-0.02
0.218+/-0.02
0.198+/-0.03
0.175+/-0.02
0.170+/-0.03
median_2
0.156+/-0.01
0.162+/-0.01
0.150+/-0.01
0.155+/-0.01
0.158+/-0.01
0.145+/-0.01
semtratol
0,423+/-0.04
0.384+/00.01
0,458+/-0.05
0.380+/-0.03
0.388+/-0.03
0.362+/-0.04
osimetrinib mpo
0.749+/-0.02
0.722+/-0.03
0.731+/-0.01
0.733+/-0.02
0.754+/-0.02
0.721+/-0.02
perindopril.mp .mo
0.390+/-0.03
0.382+/-0.04
0,388+/-0.03
0.376+/-0.04
0,399+/-0.03
0.357+/-0.06
ranolazine mpo
0.652+/-0.03
0.648+/00.04
0.643+/-0.02
0.638+/-0.03
0.628+/-0.04
0.601+/-0.02
rdkit_logp
9.733+/-0.96
10.678+/-0.99
10.508+/-0.57
10.690+/-0.69
9.517+/-0.53
9,406+/-1.34
rdkit qed.
0.917+/-0.02
0.909+/-0.03
0.926+/-0.01
0.915+/-0.02
0.906+/-0.03
0.917+/-0.02
sa_tde
8.047+/-0.15
7.909+/-0.07
8.012+/-0.13
7.982+/-0.02
7.945+/-0.03
7.798+/-0.13
scaffold_hop
0,451+/-0.02
0.435+/-0.01
0.441+/00.02
0,444+/00.0
0,440+/-0.01
0,427+/-000
sitagliptin _mpo
0.388+/-0.05
0.259+/-0.07
0.277+/-0.06
0.246+/-0.03
0,247+/-0.06
0.276+/-0.09
thiothixene rediscovery
0.243+/-0.02
0.247+/-0.03
0.244+/-0.02
0.238+/-0.01
0.247+/-0.02
0.225+/-0.01
troglitazone rediscovery
0.197+/-0.01
0.205+/-0.02
0.208+/-0.02
0.214+/-0.02
0.203+/-0.02
0.192+/-0.01
basartaa_smarty zaleplon_mps
0.000+/-0.00
0.000+/-0.00
0.000+/-0.00
0.000+/-0.00
0.000+/-0.00
0.000+/-0.00
0.382+/-0.03 0.379+/-0.05
0.396+/-0.01
0.346+/-0.05
0.357+/-0.04 0.340+/-0.07
Table 2. Average performance (# s.d.) 5 repetitions of COWBOYS the discrete BO benchmarking suite of Cfonzález-Duque et al.,
2024). We stress best average scores achieved after 300 evaluations (dark) and scores within single standard deviation of best (light). All runs of COWBOYS achieve roughly comparable results except the two the far right, ie, using few MCMC chains/iterations modifying COWBOYS to use latent space GP, respectively.
B. COWBOYS' Ablation Study
Table 2 examines COWBOYS' robustness to different configurations of its PCN MCMC sampler across the benchmark problems of González-Duque et al. (2024), demonstrating that COWBOYS is not sensitive to the number of MCMC chains or the number of MCMC steps, unless they are set be very small. Of course, no single parameter choice can yield uniformly optimal performance across different objective functions, owing to varying degrees of model mismatch between our GP and each specific problem objective and varying degrees of difficulty/locality of the optimisation problems. The final column of the table highlights the combined impact of COWBOYS' two primary components (its MCMC sampling, and its modelling in structure space) by demonstrating that if we apply model, we lose performance. More precisely, rather than fitting a Tanimoto GP in the structure space, we instead follow an approach closer to LSBO, and a latent space GP, i.e., we replace the structured GP model f (recall ((h)(z))) with a direct GP model (=z), but otherwise proceed in COWBOYS.
BOWBB-YSsyyl sampling but with a latent-space GP
In order to diagnose if our deterministic decoder approximation is limiting the flexibility of our proposed molecules, We also ran the MCMC as described (to get chosen z) but then sampled from the full stochastic decoder to get the chosen molecule. This lead to significant drop in performance for 10 of the 25 tasks (and returning statistically similar scores for the other tasks), which we hypothesise is due to the fact that the resulting "sampled" molecules are now not the same as the ones passed to the GP to see if the considered latent code will yield an improvement in score i.c. we return to a problem similar to that of the alignment issues that we discuss in Section 3.2.
14

---

Return of the Latent Space COWBOYS
Finally, also investigated the robustness of COWBOYS against degraded surrogate mode performance (to mimic settings where it may be hard to propose structure-space kernel). Here, we deliberately impaired the Tanimoto GP surrogate by randomly swapping ones to zeros in molecular fingerprints with varying probabilities. This experiment corresponds to adding significant imput noise to the fingerprints by randomly removing the count of particular (hashed) molecular substructures. Our findings indicated that COWBOYS maintains robust performance except under the most severe perturbations. In summary, perturbing Tanimoto entries with probability 0.01, 0.1, and 0.5 lead to a statistically significant drop in the best molecule found across 0, and 8 tasks from the 25 molecular optimization tasks, respectively.
15


Reviews:
[{"review_id": "SlbgxrkjcR", "reviewer": "Reviewer_dWxR", "summary": "This paper proposes a novel approach to Bayesian Optimization (BO) over structured spaces such as molecular design. Instead of fitting a Gaussian Process (GP) surrogate model within the latent space of a Variational Autoencoder (VAE), which often leads to poor predictive performance, the authors introduce COWBOYS (Categorical Optimisation With Belief Of underlYing Structure). This method decouples the generative model (VAE) from the surrogate (GP), allowing the GP to be trained directly in the structured space rather than the latent space. By leveraging a Bayesian update rule, COWBOYS refines its search iteratively without requiring the complex fine-tuning of VAEs. Experiments show the approach outperforms traditional Latent Space BO (LSBO) methods, demonstrating significant improvements in efficiency for molecular optimization tasks under constrained evaluation budgets.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your strongly positive review and for highlighting the novelty, clarity, and thorough experimental design of our work. We have carefully considered your two main comments, both of which guided additional discussion and improvements in the final version of the paper. These insights have helped us better refine our manuscript. We look forward to seeing how this work inspires new lines of inquiry and contributes to the continued advancement of Bayesian optimization.\n\n\n1) **On Incorporating Arbitrary Acquisition Functions**. You raise an important question regarding the integration of an arbitrary acquisition function into COWBOYS. In essence, our proposed approach is fundamentally different from traditional BO settings that rely on standard utility-based acquisition functions. Here, we move away from maximizing an acquisition function over a constrained set of equally likely structures. Instead, we explore searching over an entire distribution (as modeled by the VAE) of likely structures.\nBecause of this distribution-based view, conventional acquisition functions (e.g., Expected Improvement or Probability of Improvement) cannot be directly applied without inheriting the same limitations found in current LSBO methods, most notably their disregard for the distribution provided by the VAE. One of our paper’s main contributions is precisely to handle this distribution by replacing maximization-based strategies with a sampling-based approach (Equation 3).\nWhile Equation 3 does resemble Probability of Improvement (a point now greatly elaborated in the final version), adapting more sophisticated acquisition schemes is indeed an exciting direction for future research. For instance, we could condition the VAE to generate samples for which an alternative acquisition criterion (e.g., EI) exceeds a certain threshold. Determining this threshold adaptively would open up further avenues, possibly guided by theoretical regret analyses to explain COWBOYS’ strong empirical performance. We believe these avenues will become fertile ground for future studies that aim to bridge existing BO theory with generative models in complex, high-dimensional design spaces.\n2) **Generalizability of COWBOYS Beyond Molecular Design and Drug Discovery**. Please see our response to vn27 for a discussion of the significant potential of COWBOYS in different domains.", "claims_and_evidence": "Overall yes, though with some of my concerns listed below.", "methods_and_evaluation_criteria": "Yes, though the evaluation is mostly done on molecular design and drug discovery problems. Although they are the area of application that COWBOYS is designed for, it is unclear whether the claimed superior performance can be extended to other application domains.\n\nOne of the issues of the proposed method is that it still requires fitting a GP (or other probabilistic model) on the original, high-dimensional space. In many molecular design and drug discovery problems, Tanimoto is suitable for such high-dimensional binary problems and scale well, but it may likely not be the case in many other problem — simply fitting the GP (or other probabilistic model) on the high-dimensional space is already prohibitively expensive, not to mention any efficient optimization.\n\nAlso the optimization step described in eq (3) seems can be replaced by more established method such as Expected Improvement or Thompson sampling. If it can/should not be changed, more discussion of why should be included.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Overall experiments are conducted thoroughly by evaluating a number of baselines on a good number of test problems.\n\nIdeally I’d like to see more ablation study on COWBOYS’ components (e.g., sensitivity to VAE hyperparameters, impact of optimization budget such as N_init, choice of acqf used in eq 3), as we currently only have a very limited one for the PCN MCMC sampler in Appendix B.", "relation_to_broader_scientific_literature": "This is directly related to the general Bayesian Optimization and high-dimensional/latent-space BO.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The paper is overall well-written and easy to understand.", "other_comments_or_suggestions": "Minor: extra space in L284.", "questions_for_authors": "- Instead of using the sampling strategy described in eq (3), which may intuitively make sense, but has little theoretical justification, why not use something more established such as Thompson sampling or EI to select the candidates.\n\n- I’d love to see more discussion and/or experiment on other problem domains. As I mentioned in “Methods And Evaluation Criteria\" above, the proposed method may not extend well to other domain, which may be fine, but would significantly limit its impact.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EUQKwmRaiI", "reviewer": "Reviewer_1JvV", "summary": "This manuscript proposes an alternative approach to latent space Bayesian optimization in high-dimensional/structured spaces. Contrary to previous methods that aim to efficiently couple a generative model (typically VAE) and GP surrogate, the proposed method decouples the generative model and GP surrogate.  The proposed method builds on an assumption that a surrogate model can be trained in the original high-dimensional/structured data domain or using features extracted from the data objects in the original domain (here authors consider molecules and features are molecular fingerprints). The trained generative model and GP surrogate are then combined during the optimization step by defining a conditional sampling distribution for the evaluation of next molecules. Results section contains extensive evaluation against most if not all the state-of-the-art methods, and results demonstrate highly competitive performance.\n\nReview update after rebuttal: I am satisfied with authors answers, and will keep my suggestion to accept the manuscript.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your thoughtful review and valuable suggestions. We sincerely appreciate your recognition of the novelty, clarity, rigorous empirical evaluation, comprehensive background coverage, and overall manuscript quality of our work.\nWe have addressed all your suggestions thoroughly and have implemented the requested revisions as detailed below:\n1) **Generalizability of COWBOYS Beyond Molecular Design and Drug Discovery**: We fully agree with your observation on the importance of clarifying the generalizability of our approach. Following your advice, we have expanded our discussion (now included prominently in the introduction and highlighted in the limitations section) regarding the applicability and significant potential of COWBOYS beyond molecular fingerprint-based problems. Please refer to our detailed response to reviewers vn27 and 1JvV for further context.\nFurthermore, your insightful suggestion about investigating the robustness of COWBOYS against degraded surrogate mode performance (to mimic settings where it may be hard to propose a structure-space kernel) was extremely helpful. As suggested, we have conducted additional experiments on the high-dimensional BO benchmark of Section 7.2, deliberately impairing the Tanimoto GP surrogate by randomly swapping ones to zeros in molecular fingerprints with varying probabilities. This experiment corresponds to  adding significant input noise to the fingerprints by randomly removing the count of particular (hashed) molecular substructures. Our findings indicate that COWBOYS maintains robust performance except under the most severe perturbations and have been added to our ablation studies of Appendix B. In summary, perturbing Tanimoto entries with probability 0.01, 0.1, and 0.5 lead to a statistically significant drop in the best molecule found across 0, 3, and 8 tasks from the 25 molecular optimization tasks, respectively. \n2) **Small additional Points**. We have clarified our notation by using boldface m for fingerprints to distinctly differentiate from molecules represented by boldface x. Following your recommendation, we have included a graphical model illustrating the dependencies within both COWBOYS and existing LSBO approaches. We appreciate this suggestion as it significantly enhances the manuscript's clarity. We have added the suggested recent reference, highlighting its relevance. The presence of this paper (and also [1] at the same conference in 2 months time), both using extensions of standard LSBO frameworks, underscores the timeliness and importance of our contributions.\n\nThank you once again for your constructive feedback, which has significantly strengthened our manuscript.\n\n[1] Lee, Seunghun, et al. \"Latent Bayesian Optimization via Autoregressive Normalizing Flows.\" The Thirteenth International Conference on Learning Representations. 2025.", "claims_and_evidence": "The presented result support the claims of highly competitive performance.", "methods_and_evaluation_criteria": "The proposed method can be considered as a well-designed and novel approach to the high-dimensional/structured BO.  Empirical evaluations follow the standard benchmarks in the field.", "theoretical_claims": "Manuscript does not contain any formal theoretical claims but the probabilistic model construction seems well-designed and valid.", "experimental_designs_or_analyses": "Empirical evaluations use the standard benchmarks in the field and replicates and extends the evaluations from previously published papers.", "relation_to_broader_scientific_literature": "see below", "essential_references_not_discussed": "All relevant background is presented.", "other_strengths_and_weaknesses": "As far as I can tell, the proposed method can be considered as a novel approach to BO in structured spaces. \n\nThe method is presented clearly and entire manuscript is very well written. \n\nEmpirical performance evaluation seems solid, and the results demonstrate strong performance.", "other_comments_or_suggestions": "I agree with all central parts of the manuscript. My only slightly critical comment or concerns is that the proposed method assumes that the surrogate model can be trained using the molecular fingerprints as inputs. More generally, the proposed method is applicable whenever a structured GP kernel can be defined, as authors note at the end of page 1 (this limitation could be repeated in Limitations Section at the end).  While this works very well for molecules using fingerprints and Tanimoto kernel, there certainly are other problems where it less clear how a structured kernel should be defined, or whether such exist.  I have two related suggestions: 1) authors could elaborate on when and why this assumption holds (e.g. in Discussion and Limitations Section).  2) authors could carry out an ablation where they study the role and accuracy of the structured surrogate GP.  Authors could e.g. artificially remove (or add noise to) fingerprints such that the surrogate modeling becomes more challenging and less accurate, and evaluate how that affects the overall performance. \n\nNotation: my understanding is that throughout the manuscript $x$ generally denotes a molecule, except in the context of Tanimoto kernel, where $x$ denotes the molecular fingerprint representation. I suggest to use a different symbol around Tanimoto kernel, e.g. $finger(x)$  or something else to clearly note that they are different. \n\nSince your model is build on probability of improvement (PI) acquisition function, it would make sense to use the same acquisition function also for the baseline comparison methods. Perhaps that was mentioned in the text but I do not remember seeing that authors commented on that. \n\nCan you also incorporate a probabilistic graphical model representation of your proposed model (in supplement if not enough space in the main text). \n\nThe proposed method can be considered as a conditional generative model for BO. In related works section you could also cite a recent work (Ramchandran et al, ICLR, 2025, https://openreview.net/pdf?id=SIuD7CySb4 ) that has also proposed to use conditional generative model when they attempt to align the latent space for surrogate modeling.", "questions_for_authors": "-", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s9mbCJJeal", "reviewer": "Reviewer_QRHf", "summary": "This paper looks for valid molecules with a high probability of improving a desired property of the molecule. The novelty of this method is that they decouple the two probabilities: the probability of a molecule existing is estimated by VAE which is initially trained/pre-trained, then the probability of a molecule having the desired properties is estimated from a GP which operates on the space of count based vectors of molecular features. \n\nBy not fitting a GP in the latent space of the VAE they claim to avoid two pathologies. First,  there can be a lack of smoothness of the objective function in the latent space. Second, because a single point in the VAE latent space maps to several molecules through noise added to the decoder which each have different values for the objective function, a GP in the latent space will inevitably have more noise than could be predicted from a single structure. \n\nThe authors suggest candidate molecules using an MCMC sampler using a prior defined by a deterministic decoder and a likelihood given by the probability of improvement compared to the current best observation, as calculated by the GP. \n\nThe authors benchmark this method agains various alternative methods on several datasets.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your thoughtful review and positive feedback regarding our experimental design and validations. We greatly appreciate your recognition of our method's strong performance. We have carefully addressed all of your suggestions and implemented the requested revisions, detailed below. We hope these efforts demonstrate our commitment to improving the manuscript and may justify an increased evaluation score.\n\n1) **Code Accessibility** In response to your suggestion we have now made some of our code available [here](https://anonymous.4open.science/r/cowboys-AC3E/README.md). The repository allows the recreation of the benchmarking results of Section 7.2. Indeed, incorporating COWBOYS (see [cowboys.py](https://anonymous.4open.science/r/cowboys-AC3E/cowboys.py)) required only minimal additions, emphasizing both its practical usability and potential for broad community adoption. The fully-fledged codebase, allowing COWBOYS to be applied to generic pre-trained VAEs and thus applied to a range of downstream tasks (and recreate the results of Sections 7.1 and 7.2), will be provided upon acceptance as it cannot be shared without violating ICML’s anonymity rules.\n\n2) **Clarification of Equation (Line 267)** We appreciate you pointing out the ambiguity in the equation on line 267. As requested, we have clarified this equation explicitly in the revised manuscript, clearly defining the delta function as indicating a deterministic decoding strategy that selects only the most-likely decoded molecule from the latent location z. The delta function here means that the probability is 0 unless we are taking the most likely prediction from the decoder (from the chosen latent location z).\n\n3) **Addressing concerns about discarding the stochasticity of the decoder**. We understand your concerns, as we initially were also worried. We agree that there is indeed a theoretical risk of being unable to generate as wide a range of molecules as the stochastic VAE (although not one we observed in practice). We believe the risks of allocating 100% confidence to implausible molecules, which would arise when decoding areas of the latent space where the decoder is highly uncertain, are strongly mitigated by the fact that we are doing conditional VAE sampling (i.e. staying on the well-supported areas of the latent space) rather than exploring the whole latent space like standard LSBO methods.  Nevertheless, future work will seek more sophisticated sampling strategies to perform COWBOYS-like strategies for generative models (see our Discussion Section). \nWe like your idea of running the MCMC as described but then sampling from the decoder in order to properly diagnose if the restrictions of the deterministic decoder are limiting the flexibility of our proposed molecules. We have implemented and ran this suggestion for the high-dimensional BO benchmark of Section 7.2, with the full results added to the ablation study of Appendix B. Unfortunately, this lead to a significant drop in performance for 10 of the 25 tasks (and returning statistically similar scores for the other tasks), which we hypothesise is due to the fact that the resulting “sampled” molecules are now not the same as the ones passed to the GP to see if the considered latent code will yield an improvement in score – i.e. we return to a problem similar to that of the alignment issues that we discuss in Section 3.2. Based on your comments,  we have greatly increased our discussion of our decision to use a deterministic simplification of the decoder, stressing the points above and that (1) the VAE is still initially trained with a stochastic decoder (it is only at inference/BO time that we take the most likely mapping), and (2) we have added four additional specific references of previous work that also using this deterministic simplification. \n\n4) **Clarifying Differing Comparisons in Experiments**.  The diversity in initial budgets, degrees of parallelism, and subsets of molecular optimization problems stems directly from our intention to replicate and extend the experiments used in influential prior work in LSBO precisely. We believe this breadth provides a robust and comprehensive empirical evaluation of COWBOYS' performance. For peace of mind, please see the comments of the other reviewers, e.g. “Empirical evaluations use the standard benchmarks in the field and replicates and extends the evaluations from previously published papers” and “experiments are conducted thoroughly by evaluating a number of baselines on a good number of test problems”. However, acknowledging your valid concern regarding transparency—especially for readers less familiar with molecular design benchmarks—we have now included an additional section in the appendix. This section explicitly details the structure and interpretation of these benchmark suites, clearly explaining the objectives, settings, and significance of the various optimization tasks we considered.", "claims_and_evidence": "The authors claim that their method outperforms alternatives the low data regime and existing high dimensional discrete sequence which appears to be well supported by the experiments shown here.  They also claim with good support to outperform traditional Latent Space BO even when the traditional methods have much higher experimental budgets. \n\nTwo elements which make me less confident on the quality of the method are:\n\n1. The code has been retracted for the review and I do not see any anonymised version to check their implementation. I would have appreciated seeing the code during the review. \n\n2. Use of different comparison methods for different datasets.  I am not familiar with the standard benchmarks in this area and would have appreciated either seeing the same methods applied across all datasets, or a justification for why not. Please see the question below regarding justifying this.", "methods_and_evaluation_criteria": "I thought that the proposed method and evaluation criteria mostly made sense and was well justified. \n\nMy one concern was in Section 5.1 where they discard the stochasticity of the decoder. The VAE with a stochastic decoder should provide a probabilistic manifold for the density of possible molecules. The deterministic decoder seems to have a potential issue with either assigning definite existence to impossible molecules or definite existence to optimal molecules. Additionally, as the MCMC sampler was just used to find candidate molecules (as opposed to finding exact posteriors) I do not see the issue with running the MCMC as described but then sampling from the stochastic decoder, which would have potentially avoided these concerns. \n\nI would have appreciated a discussion of potential issues from this, although I appreciate that they do not seem to have been major concerns in practice given the experimental results. Please also see the question regarding this.", "theoretical_claims": "The authors make no strong theoretical claims.", "experimental_designs_or_analyses": "I do not see any issues with the experimental design carried out here, the experiments and comparisons both seem to be valid and show good results.", "relation_to_broader_scientific_literature": "I am not familiar with the literature surrounding BO for molecules, but to me the literature review seems extensive.", "essential_references_not_discussed": "I am not aware of any references which should have been discussed here, although I do not claim extensive knowledge of BO applied to molecular design.", "other_strengths_and_weaknesses": "The strengths and weaknesses are addressed in other comments adequately.", "other_comments_or_suggestions": "1. As a non-expert in BO in molecular design I would have liked 1-2 sentences describing each of the benchmark datasets and the particular challenge/reason for including each one. \n\n2. The equation on the left hand column of line 267 was unclear if it was using a Dirac Delta function, which is typically not used for categorical variables. I believe that what is being stated in that line is that $\\hat p_\\theta(x|z)$ is a vector of all zeros except for one element with value one, where the non-zero element corresponds to the maximum likelihood prediction for $p_\\theta(x|z)$. See also the question on this.", "questions_for_authors": "1. Can you comment on possible problems of using a deterministic decoder, specifically: \n     i) is there a risk of the VAE with a deterministic decoder assigning 100% confidence to molecules which in practice do not/ cannot exist? \n     ii) Does the deterministic VAE prevent the model from exploring other similar molecules to high performing ones result in missing very good candidate molecules, especially in the final rounds of a BO run in the exploitative stage? \n\n2. Can you clarify what is meant by the equation on the left hand column at line 267?\n\n3. Can you clarify why on different benchmark datasets different comparison methods are used? Do the different approaches not work on the other benchmark datasets, or was the decision for which methods to compare based on another consideration?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gNGxR2mc8X", "reviewer": "Reviewer_vn27", "summary": "In the proposed manuscript, the authors present a novel approach for Bayesian optimization within the latent space of a VAE applied to structured domains like molecular design. They aim to address multiple drawbacks of conventional methods that often result in suboptimal performance when the latent space does not align with the specific task. To address these issues, the authors introduce a decoupled framework where the generative model and a GP surrogate are trained independently and later integrated through a principled Bayesian update rule. This approach allows the VAE to specialize in generating candidate structures while the GP excels in predictive modeling, thereby improving the identification of promising candidates under limited evaluation budgets.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful assessment of our work. We appreciate your recognition of both the descriptive and empirical strengths of our approach. It is encouraging to hear that you find our decoupled strategy promising, and that our competitive results align well with other state-of-the-art methods. \n\nWe have addressed all your suggestions thoroughly and have implemented the requested revisions as detailed below. We hope that our efforts to address the comments from you and the other reviewers will be appreciated and could justify increasing your score.\n\n1) **Generalizability of COWBOYS Beyond Molecular Design and Drug Discovery**. We appreciate your interest in the broader applicability of COWBOYS outside its current focus on molecular design and drug discovery. Our evaluation focuses on these tasks because they remain the most widespread use cases for latent space Bayesian optimization (LSBO), and well-established baseline suites exist to facilitate direct comparison.\nFitting a Gaussian Process (GP) in the original, potentially high-dimensional and highly structured space can indeed be challenging with standard kernels. However, we see an ability to model directly in structure space as a key advantage of COWBOYS, rather than a limitation. By working in the raw structure space, our approach naturally supports specialized kernels tailored to particular domains—an area where there is a rich, yet underused, literature. For instance, the Tanimoto kernel, originally proposed by chemists, encodes prior knowledge of what sort of molecular attributes are important in a way that can be especially powerful in low-data regimes [1,2]. Similar structural kernels exist for various structured objects and, via COWBOYS, could now be used for molecular graph optimisation with graph kernels [3], engineering design with 3d mesh kernels [4], optimising computer code via tree kernels [5], or protein design with new protein kernels [6].\nIn contrast, many current LSBO strategies cannot leverage these specialized kernels because they rely on Euclidean latent spaces. By enabling the direct use of such kernels, COWBOYS allows practitioners to incorporate a wealth of domain-specific prior knowledge. We hope this work helps revitalize interest in harnessing these powerful kernels across a range of complex design problems and spurs further synergy between rich generative models and carefully structured discriminative models\n2) **We are pleased to hear that our discussion clarified the robustness of COWBOYS with respect to PCN parameters**. However, In the revised manuscript, we have added details on how to interpret the corresponding table, emphasizing that no single parameter choice can yield uniformly optimal performance across different objective functions, owing to varying degrees of model mismatch between our GP and each specific problem objective and varying degrees of difficulty/locality of the optimisation problems.\n\n3) **“Is the search space problem also an issue for the initial design of COWBOYS”?** In standard LSBO, a specific region of the latent space must be chosen, over which we sample an initial design and then perform BO. We argue in Section 3.3 that this clipping can lead to pathological performance issues. In contrast, COWBOYS just samples from the VAE in the standard way, i.e. we generate a random Gaussian (which has infinite support) and then push it through the decoder. There is no need to specify a search space for COWBOYS. We have greatly increased discussion around this point in the final version.\n\n4) Also, thanks for the **minor corrections**, they have now been changed in the final version, including adding more four specific references to justify ”replacing the VAE’s probabilistic decoding with the most-likely mapping is already a common strategy”.\n\n[1] Griffiths, Ryan-Rhys, et al. \"GAUCHE: a library for Gaussian processes in chemistry.\" Advances in Neural Information Processing Systems 36 (2023): 76923-76946.\n\n[2] Moss, Henry, et al. \"Boss: Bayesian optimization over string spaces.\" Advances in neural information processing systems 33 (2020): 15476-15486.]\n\n[3] Vishwanathan, S. Vichy N., et al. \"Graph kernels.\" The Journal of Machine Learning Research 11 (2010): 1201-1242.\n\n[4] Perez, Raphaël Carpintero, et al. \"Gaussian process regression with Sliced Wasserstein Weisfeiler-Lehman graph kernels.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2024.\n\n[5] Beck, Daniel, et al. \"Learning structural kernels for natural language processing.\" Transactions of the Association for Computational Linguistics 3 (2015): 461-473.\n\n[6] Groth, Peter Mørch, et al. \"Kermut: Composite kernel regression for protein variant effects.\" Advances in Neural Information Processing Systems 37 (2024): 29514-29565.", "claims_and_evidence": "The paper claims that while BO in the latent space of a VAE is a powerful framework for tackling complex optimisation tasks,  existing methods have significant drawbacks, including suboptimal surrogate modeling in the latent space and inherent search space challenges.  The paper provides a thorough discussion of these issues, offering both descriptive analysis and empirical evidence. Moreover, the proposed decoupled approach appears to address these shortcomings effectively, potentially serving as a remedy for the identified limitations.", "methods_and_evaluation_criteria": "The proposed method decouples the search space from the latent space and further uses a Preconditioned Crank-Nicolson algorithm (which concentrates on high-prior probability annulus, better than random walk sampling) for sampling. \n\nCOWBOYS conducts a number of studies, showing competitive results with state-of-the-art comparison methods, even those that fine-tune the VAE, and with high-dimensional BO benchmarks and traditional LSBO methods. Finally, the authors discuss the need to develop/extend COWBOYS to large-scale molecule optimization tasks, as this can be challenging with the proposed structure-based GP optimization approach.  \n\nWhile the presented approach shows promising results on the current task, the manuscript's impact would be significantly enhanced by testing the method in additional domains, such as protein fitness optimization or similar applications involving other structured spaces.", "theoretical_claims": "The theoretical claims in the presented work seem plausible and well discussed.", "experimental_designs_or_analyses": "The experimental design seems in line with other literature in the field and appears to be well-integrated within a benchmark suite. The approach also sounds plausible based on the information presented.", "relation_to_broader_scientific_literature": "The authors start by motivating LSBO, which has emerged to be a prominent technique in this field in recent years. Current limitations and shortcomings of LSBO are discussed, and how the community handles some of these (e.g., fine-tuning the VAE, limiting the search space). This is followed by how COWBOYS addresses these issues.", "essential_references_not_discussed": "To the best of my knowledge, the authors have included the key literature in the field, particularly regarding LSBO approaches.", "other_strengths_and_weaknesses": "Is the search space problem also an issue for the initial design of COWBOYS? Essentially, the initial design of Alg. 2 also requires sampling from the VAE’s search space and decoding the samples.", "other_comments_or_suggestions": "line 98 left: expected improvement twice \n\nline 82 right: this should be encoder instead of decoder \n\nline 94 right: “.” missing at the end of the sentence \n\nline 227 left: In algorithm, the arrow extends into the brackets\n\nLine 269 left: “*Note that replacing the VAE’s probabilistic decoding with the most-likely mapping is already a common strategy in many LSBO implementations”* You are citing here a survey paper that does not discuss this in detail, please cite the actual works to make the connection clear for the reader.", "questions_for_authors": "See questions asked above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper revisits a popular approach to Bayesian optimisation in structured spaces, combining a probabilistic model for property prediction with a latent space of an autoencoder. The decoupled approach here is likely to be broadly useful in tasks where good GP surrogates exist in data space (including for many molecule and materials generation tasks). After the rebuttal period, all reviewers had a consensus that the paper should be accepted.
Decision: Accept (spotlight poster)