High-Dimensional Bayesian Optimisation with Gaussian Process Prior Variational Autoencoders
Authors: Siddharth Ramchandran
Manuel Haussmann
Harri Lähdesmäki
Venue: ICLR 2025 Poster (2025)
Keywords: Variational autoencoders
Gaussian processes
Bayesian optimisation
Primary area: generative models
TL;DR: In this work, we propose a conditional generative model for efficient high-dimensional BO that uses a GP surrogate model together with GP prior VAEs.

Abstract: Bayesian optimisation (BO) using a Gaussian process (GP)-based surrogate model is a powerful tool for solving black-box optimisation problems but does not scale well to high-dimensional data. Previous works have proposed to use variational autoencoders (VAEs) to project high-dimensional data onto a low-dimensional latent space and to implement BO in the inferred latent space. In this work, we propose a conditional generative model for efficient high-dimensional BO that uses a GP surrogate model together with GP prior VAEs. A GP prior VAE extends the standard VAE by conditioning the generative and inference model on auxiliary covariates, capturing complex correlations across samples with a GP. Our model incorporates the observed target quantity values as auxiliary covariates learning a structured latent space that is better suited for the GP-based BO surrogate model. It handles partially observed auxiliary covariates using a unifying probabilistic framework and can also incorporate additional auxiliary covariates that may be available in real-world applications. We demonstrate that our method improves upon existing latent space BO methods on simulated datasets as well as on commonly used benchmarks.

Published as a conference paper at ICLR 2025
HIGH-DIMENSIONAL BAYESIAN OPTIMISATION WITH GAUSSIAN PROCESS PRIOR VARIATIONAL AUTOENCODERS
Siddharth Ramchandran
Manuel Haussmann
Department of Computer Science
Department of Mathematics and Computer Science University of Southern Denmark Odense, Denmark
Aalto University
Espoo, Finland
Harri Lähdesmäki
Department of Computer Science
Aalto University
Espoo, Finland
ABSTRACT
Bayesian optimisation (BO) using a Gaussian process (GP)-based surrogate model is powerful tool for solving black-box optimisation problems but does not scale well to high-dimensional data. Previous works have proposed use variational autoencoders (VAEs) to project high-dimensional data onto a low-dimensional latent space and to implement BO in the inferred latent space. In this work, we propose conditional generative model for efficient high-dimensional BO that uses a GP surrogate model together with GP prior VAEs. A GP prior VAE extends the standard VAE by conditioning the generative and inference model auxiliary covariates, capturing complex correlations across samples with a GP. Our model incorporates the observed target quantity values as auxiliary covariates learning a structured latent space that better suited for the GP-based BO surrogate model. It handles partially observed auxiliary covariates using unifying probabilistic frame- work and can also incorporate additional auxiliary covariates that may be available in real-world applications. We demonstrate that our method improves upon exist- ing latent space BO methods on simulated datasets as well as on commonly used benchmarks.
INTRODUCTION
Bayesian optimisation (BO) (Mockus, 1989; Shahriari al., 2015; Frazier, 2018) is a technique for complex optimisation problems, where the true functional form of a target quantity of interest is unknown. This target quantity may be expensive compute or may require time consuming experiments to obtain its value. Hence, one would like to minimise the number of evaluations that are required to optimise it. Although BO offers an approach for black-box optimisation problems, it does not efficiently scale to high- high-dimensional data settings (Shahriari et al., 2015).
Variational autoencoders (VAEs) (Kingma & Welling, 2014; Rezende et al., 2014) are a popular family of latent-variable models that are often used to learn low-dimensional representations of high-dimensional data. The low-dimensional latent space afforded by VAEs, that is representative of the high-dimensional,, potentially discrete-valued data which is trained, offers powerful scaling strategy for BO. BO is performed on the inferred low-dimensional continuous-valued manifold instead of the nigh-dimensional data space combining the benefits VAEs with BO, known as VAE BO, is black-box optimisation method with many practical applications, such molecule discovery (Gómez- Bombarelli et al., 2018; Griffiths & Hermnndez-Lobato, 2020; Jin et al., 2018), neural architecture
Boobarllo et al., 2018). This method of dinge--immsiinall
"Correspondence s sidcaa rth, ramchandran@aallt..o

---

Published as a conference paper at ICLR 2025
R Σ
GP 0.
K 10
(x
x 1
100)
Step Decode candidate
Step
Step Encode molecules
Step 2
Repeat
BO in latent space obtain candidate
Y
Step Fit GP prior VAE to data
Black Box Function
New penalised LogP
Encader
Latent space coloured by quantity interest
Decoder
Figure I: An overview of our model. Consider the example application of discovering novel drug-like molecules. Our method uses GP prior VAE with an additive kernel over various partially observed auxiliary covariates such as molecular weight, number of hydrogen bonds, total polar surface area, etc. and the partially observed quantity of interest (represented by 2((r) in this image for the y additive kernel) learn structured latent space. The black-box function evaluates the quantity of interest for the chosen molecule.
search (Kandasamy al., 2018; Ru et al., 2021) and chemical synthesis (Felton et al., 2020; Shields ct al., 2021; Korovina et al., 2020).
Sohn et al. (2015) proposed conditional VAEs (cVAEs) extension that conditions generative model on auxiliary covariates. However, similar to standard VAEs, this family of models ignores possible correlations between data samples. The Gaussian process (GP) prior VAE (Casale et al., 2018) extends the conditional VAE framework by replacing the i.ddd standard Gaussian prior on the latent variables with GP prior in order capture arbitrary, but preferably smooth, correlations between data samples. These models have been shown to compare favourably to VAEs and e cAAss as well as effectively handle missing data in the observations. Ramchandran et (2024) introduced method to impute the missing auxiliary covariates in cVAss and thereby enhance their applicability real-world datasets.
Our Contribution We propose novel conditional deep generative model for high-dimensional BO that improves upon the existing VAE BO methods. Our proposed model uses GP prior VAE to learn low-dimensional, structured latent representation of the data samples, and implements the GP surrogate model to optimise the target quantity (or quantities) of interest in the repeatedly re-trained latent space. We these partially observed target quantity values directly auxiliary covariates to condition the GP prior VAE model. The model also incorporates additional (partially or fully) observed auxiliary covariates that may be available for given application. Furthermore, it can effectively handle missing values in both the high- high-dimensional observations as well as the auxiliary covariates using principled technique that is particularly developed for learning conditional VAEs. Fig. summarises our model.
Our contributions can be summarised follows:
We introduce conditional VAE-based method for efficiently performing Bayesian optim- isation on high-dimensional datasets.
We learn structured latent representations of high-dimensional data points using GP prior VAE that handle missing values the observations, target quantity values, and other possible auxiliary covariates.
We demonstrate the efficacy of our method on a synthetic dataset and on common bench- marks.
The code available https: thub. com/St diamm http//Gpppoooooo
2

---

Published as a conference paper at ICLR 2025
2
RELATED WORKS
Bayesian optimisation is a popular black-box optimisation technique that is challenging to scale to high-dimensional data (Mockus, 1989; Shahriari et al., 2015; Frazier, 2018). Binois & Wycoff (2022) reviews the recent advancements in improving the efficiency of Bayesian Optimisation (BO) for high-dimensional problems, particularly through various structural model assumptions. To address the curse of dimensionality, Griffiths & Hernández-Lobate (2020) uses autoencoder to learn low-dimensional. non-linear manifold to scale BO to nigh-dimensional datasets. They perform a constrained BO over the latent space in order to incorporate the application-specifie idiosynerasies and thereby generate a high proportion of valid reconstructions. Stanton et al. (2022) integrate Denoising Autoencoders with discriminative multi-task Gaussian process head into BO to learn latent space that captures meaningful features of biological sequences. As autoencoders cannot be used sample novel observations from their representation space, VAEs are an approach make it possible to leverage the diw-dmensional latent representation for generative purposes (Kusner al., 2017: Geanez-Bombarelll et al., 2018). However, a vanilla VAE BO is sub-optimal as the learnt latent space is not constructed by leveraging the black-box function labels (Urtasun & Darrell, 2007; Siivola et al., 2021; Grosnit et al., 2021). Building upon this, some methods: use an automatic statistician perspective by learning the kernel combination of the surrogate GP (Lu et al., 2018 8), use manifold GPs in the encoder and manifold multi-output GPs in the decoder (Moriconi et al., 2020), reformulate the encoder to effectively act both as the encoder for the VAE well as a deep kernel for the surrogate model within a local Bayesian optimisation framework using trust region method (Maus al., 2022), and use label guidance in the latent space (Eissman al., 2018; Tripp et al., 2020; Maus et al., 2022). Furthermore, Grosnit et al. (2021) proposed a method that combines VAES with deep metric learning. They make use of label guidance from the labelled data points by incorporating various metric losses (c.g., triplet loss, contrasting loss, log ratio loss. etc.). However, this method does not incorporate additional information in the form of auxiliary covariates and the triplet loss requires an additional matching procedure as pre- rre-processing step, which can be time consuming. Other relevant works include (Notin et al., 2021; Maus et al., 2023; Lee al., 2024)
Variational autoencoders (Kingma& & Welling, 2014; Rezende et al., 2014) are popular deep learning methods that map digeesinnsionn,, complex data to a oin-eiensional space and vice-versa. Most VAE-based models assume the data to be fully observed or choose to substitute unobserved values of the encoder input with zeros (Nazabal et al., 2020; Mattei & Frellsen, 2019). Conditional variational autoencoders (Sohn et al., 2015) include information about the auxiliary covariates into both the inference and generative networks. Building upon this idea, Gaussian process prior VAEs have been well as auxiliary covariates via proposed as extension to incorporate arbitrary correlations Ramchandran et al., 2021). These Gaussian process priors (Casale et al., 2018; Fortuin et al., 2020; methods have shown competitive performance as well handle missing values in the observed data. Ramchandran et al. (2024) proposed a conditional VAE-based learning approach that can robustly handle missing values in the auxiliary covariates.
3 BACKGROUND
Throughout the paper, we use the following notation: y E y is high-dimensional observation, c E TR the target quantity that we want to optimise, x = , æ X denotes additional auxiliary covariates, and = E Z R² is L-dimensional latent variable. We define á [c, ær| IR X. A set of N observations is denoted as (....... yy.. with X, X, and Z defined analogously. The target quantity e = [cI,...,cN' is typically partially observed.
3.1 BAYESIAN OPTIMISATION
Bayesian optimisation is technique for performing efficient global optimisation of black-box functions (or unknown scoring functions) that are difficult to compute and whose functional form may not be known (Kushner, 1962; 1964; Mockus, 1989; Frazier, 2018). Given function f - + R we aim to find a point y E y that corresponds to the global optimum of f. The black-box function f is also referred to utility function as it is measure of the target quantity, trying to optimise and informs us on the quality of the chosen sample. The problem can be written as (assuming maximin maximisation), y² arg max,ey f(y). Since, the unknown function ʃ is assumed to
f that
3

---

Published as a conference paper at ICLR 2025
be difficult expensive evaluate, Bayesian optimisation requires surrogate model model the true function f as well as an acquisition function which is a function of the posterior and guides the process of choosing the next sample point until stopping criteria is met the evaluation budget B exhausted.
Gaussian Processes and the Surrogate Model We use a non-parametrie Gaussian process as the surrogate model of f as GPs define probability distribution over functions and for Gaussian likelihood models the posterior distribution analytically tractable. Moreover, they maintain smoothness and uncertainty estimates to guide the exploration of new points as well as represent prior beliefs (Schulz et al., 2018). Following Williams & Rasmussen (2006), for inputs y: y E Vo GP is defined as g
GP(((3)) k(y,y')) where puy) is the mean and k( y, y is
kernel function given by kyy y²) = (g))) For N data points Y = [y....... yy the induced prior probability density g(Y) [g(y)) (.................... yyNN )]² is multivariate Gaussian distribution: g( ~ NO, ). We assume μ = throughout this work. The elements of the covariance matrix are defined by the kernel function [Ky,ylij
U,). GPs are intractable for large datasets the time complexity scales by ((V*). Several approximate methods have been proposed to address this through sparse Gaussian processes (Smola & Bartlett, 2000; Lawrence et al., 2002; Mamnonero-Cander & Rasmussen, 2005) or via (stochastic) variational formulations (Titsias, 2009; Hensman et al., 2013) for sparse approximations.
Acquisition Functions An acquisition function is a function of the posterior that captures the trade-off between exploration and exploitation of our surrogate of the function f given the known evaluations. responsible for selecting the next candidate point in y that should be evaluated measured. We use an acquisition function azy) to choose the next sample point VN+1 arg max, axy)) A good acquisition function exploits regions around the current maximum by selecting points to query from that region while also suggesting points from unexplored regions in order to escape local maxima. There are several candidate functions such as upper confidence bound, expected improvement, probability of improvement, and Thompson sampling (Shahriari et al., 2015). Our proposed method is agnostic to the choice of acquisition function.
3.2
VARIATIONAL AUTOENCODERS
We define a latent variable generative model as Pw (w, x) Po (y (2)) which is parameterised by {t), 0), and where is unobserved. We generally interested in inferring this latent variable given y. The posterior distribution, P.. (z y) = Pv. (yy = =))p(z)/ pp (y), is usually intractable due to the lack of closed-form marginalisation
the latent space (Murphy, 2023). The standard
VAE model comprises the generative model (the probabilistic decoder) ( 2) and an inference model (the probabilistic encoder) ( y) that approximates the true posterior. VAEs use amortised variational inference that exploits the inference model go y) to obtain approximate distributions for each nn- The encoder and decoder are typically parameterised by deep neural networks. In variational inference we minimise the Kullback-Leibler (KL) divergence from (z y) to (z y). or equivalently maximise the ELBO of the marginal logelikelihood w.r.t. ø. For VAEs, approximate inference is typically conducted alongside learning the generative model's parameters, that is, w.r.t.
d, 1/, 0:
N Σ
lo Pe (Y) 2 C(o, 1b, BY A
max, Egg (o =x)] [L[gg(Zn yy))
0.0.0
It is straightforward to apply computationally efficient mini-batch based stochastic gradient descent to the above equation.
4 OUR METHOD
4.1 BAYESIAN OPTIMISATION WITH VAES
The low-dimessional nonlinear latent manifold learnt by VAE can be used to perform BO (Kusner et al., 2018; Tripp et al., 2020). The VAE first pre-trained on
et al., 2017: the high-dimensional observations without access to the utility function values. As described in
4

---

Published as a conference paper at ICLR 2025
Sec. 3.2, the encoder (=(2 y) of the learnt VAE is used to map the observations y y onto low-dimensional latent representation & E Z. The VAE-based methods then perform latent space optimisation (LSO) (Tripp et al., 2020) by fitting a surrogate model over the latent space to model the utility function of interest. The VAE BO aims to identify a z" such that the corresponding y", that is obtained from the pre-trained decoder, minimises utility function of interest, f(y"). In other words, we would like to obtain a x" such that we maximise the expectation over the utility function evaluated on y ~ (y x*), i.c., arg m gy (f))) Once we have a new y and its associated utility function value c, we append them to the training dataset and update the parameters and d either after each BO step or at a chosen frequency. Tripp et al. (2020) use this approach with
the help of a weighted retraining scheme according to their utility function values.
4.2 GAUSSIAN PROCESS PRIOR VAES FOR BO
A limitation of standard VAE BO is that it infers an unconditional latent-variable model without any guidance from the observed target quantit- ies. Departures from this limitation have been proposed, e.g.. in (Eissman al., 2018; Tripp et al., 2020; Maus et al., 2022). Recently, Gros- nit et al. (2021) built upon VAE BOs by using deep metric learning to actively steer the gener- ative model to maintain a latent manifold that is useful for the BO task. We propose use GP prior VAEs that guide the generative model by conditioning the GP prior with auxiliary covari- ates.
coveristes
N
GP
BO
x
Figure 2: Our proposed model. Solid lines refer to the generative model and dashed lines to the inference model. Empty circles are unobserved, shaded circles are observed, and partially shaded circles are partially observed. Target quantity d' and possible additional covariates x' refer to the new candidate observation that will be added to the training set.
The key distinction of GP prior VAEs is that the factorisable conditional prior defined over
the latent space pa( xxx) = I (a æ;)
replaced by GP prior. Assuming a func-
tion T X + 2, which maps auxiliary cov-
ariates to the T--dimensional latent space, we
denote z ((x) (TI (x))²) GP
prior VAEs model each latent dimension with an independent GP (xx) ~ GP(p(())) ki x' (t)).
where (x) is the mean, f x' 0:) is the covariance function, and 0g denotes the parameters of the covariance function. The GP prior for the 1th latent dimension can be written as a joint mul- tivariate Gaussian distribution for the function values ž₁ =x(X) (212) (xxN)) ))T, such
that po [2t] X) = (X)) = (z, 0, K xx), where {x }is = ka my 0.). Our joint conditional prior is po x) IT (= x)= - 0, K xx
We propose to learn low-dimensional latent embedding for BO using a GP prior VAE that is conditioned on the target quantity of interest, i.e., (Z c). We hypothesise that using the target quantity the conditioning variable will automatically guide the latent embeddings to a smooth manifold that is beneficial for the BO task. Since the target quantity e E IR, the GP prior VAE can be defined using any of the commonly used smooth kernel functions, such the squared exponential kernel. Following the same reasoning, if data points y have any additional known properties x. we can incorporate those in the GP prior VAE framework as well by conditioning the latent variable generation with both e and (we denote ã [c, æ)), i.e., po Z X). If all auxiliary covariates in ă continuous, we could incorporate ă, e.g., via a single squared exponential kernel with a shared length- scale parameter or use automatic relevance determination (ARD) kernel to define covariate-specific length-scales. In practice, however, some of the auxiliary covariates may be, e.g. binary or categorical. Ramchandran al. (2021) have shown that it is possible have flexible and expressive covariance functions depending on the nature of the auxiliary covariates. In this work, we similarly assume
Q + 1 additive covariance functions, kū W' (0) kı(c, e' + (I, 2/, 0t,r) + oïr implying that (()) k!! + D K(lr)) X,X, + o JN. where the choice of the kernels depends on the application and X, denotes the th auxiliary variable.
X
5

---

end end Published as a conference paper at ICLR 2025
Algorithm 1: An overview of our proposed algorithm
Input: Budget B, frequency v. initial dataset D, pre-trained VAE for j to J = [B/v] do // Train the GP prior VAE on D
Solve 6j, 3,,,,,0 00 <- m.. ELBBOG-.VAB-mistee....0.0
Compute D2 f- (z f(y))) by using the encoder 0j to obtain 2;;
for k = to
and (T 2 y do // Perform Bayesian optimisation
Fit surrogate GP on (U)))l))
Optimise E/ for ,,,++i
Use decoder 45 to map 2 to
// Decode point chosen by B.O.
Evaluate = f()), augment data Do. Dz; // Evaluate black-box function
Increment N
Output: 3" = arg yaxx (1)
4.3 PARTIALLY OBSERVED TARGET QUANTITY AND ADDITIONAL COVARIATES
In the BO setting, the target quantity of interest that we are optimising is typically available only for a (very) small number of data points. This is problematic for conditional generative models, such as GP prior VAEs, as they assume that covariates that are used to condition the generation are always known and observed. Moreover, in our problem setting, the additional auxiliary covariates that may be available in specific application may also have missing values. We follow a formulation similar to that of Ramchandran et al. (2024) to handle the missing values in the covariates.
We augment our generative model with a prior distribution, pa ( factorising over ă. parameterised by A. Representing the observed and unobserved parts as Y = (Yo, Y") and X (x), we approximate the true posterior distribution of the unobserved variables Z and X represen- ted as p, (Z, X yo, X²) and parameterised by 7 = {tb 0, AJ. using amortised variational inference. We make use of a conditionally independent, factorisable variational approximation: (2, X y (²) = (o Y", """X""" X = II g0 y 2 (6 (á: d:)). The
latent variables 2; are assumed to have Gaussian variational distribution and, for the discrete and continuous- valued covariates z", categorical and Gaussian distributions respectively. Following Ramchandran et al. (2024), we write the ELBO objective with missing covariates M.L.BOG--VE--iii
as logp, 89) 2 Egg [logp>
[go y X>) X", X91]
Z)] -ggs X²) IPx((>X X)),,
(1)
-KT X
where the first and the second expectations with respect the latent variables Z and missing covariates X", respectively, and can be approximated using Monte Carlo (see the Sec. A of the Appendices for details of deriving the ELBO). For each specific value of the missing covariates, the KI. divergence in Eq. 1 has computation complexity of O(N3). Earlier work by Ramchandran et al. (2021; 2024) has shown that using the low-rank inducing point approximation for the multi-output GP (Z X", {))
po(4 Z X), one can derive a scalable ELBO that provides an unbiased, mini-batch compatible lower bound for efficient learning. See the Sec. A.1 of the Appendices for the specific expression of the scalable lower bound that we use.
4.4 SIGH--DMENSIONAL. BO WITH GAUSSIAN PROCESS PRIOR VAES
We use the latent space learnt by the GP prior VAE to perform efficient BO. In particular, our method can handle missing values in both the observations and covariates æ (partially observed features denoted as y and x;",, as well as large datasets through the scalable ELBO described Sec. 4.3. Consider a dataset D- Do U DD where De represents the data points whose target quantity is observed (indexed by Io) and Dy represents the data points whose c is unobserved (indexed by Ig). Therefore, Do (v), o. Civz, with G f(y,) and Dy {D8, Jiezy where, before,

---

Published as a conference paper at ICLR 2025
D; A' refers to the additional auxiliary covariates (that may or may not be available, depending on the application), and f(-) refers to an expensive black-box function. N refers to the total number of observations, comprising the number of observations with observed quantity of interest N = [ao and number of observations with unobserved quantity of interest N = [Zul
Algorithm See Algorithm for a pseudo-code summary of our method. Budget refers to the maximum number of evaluations of the black-bon function that can be performed, refers to the number of BO steps performed before re-optimising our GP prior VAE model with the augmented dataset D, and El pertains the expected improvement acquisition function (the algorithm is agnostic to this choice).
We obtain the optimal encoder and decoder parameters (Ø) and पे respectively) by optimising the ELBO (in Sec. 4.3). The method computes the fully-ohserved, low-dimensional latent space representation z; of the observations y using the optimal encoder at the current iteration, and implements a BO step. The new chosen observation ÿ. its covariates (if known), and the obtained target quantity of interest e = fi) are appended to Do. After budget has been exhausted, our algorithm returns the best candidate acquired so far.
We periodically re-train and fit a conditional generative model using the entire dataset - comprising both the initial data and samples collected during the BO steps, where covariates may be partially observed. Unlike Tripp et al. (2020), our model fitting remains unbiased toward high objective values. Instead, periodic training guides the embeddings of high-dimensional samples toward smooth manifold, as specified by the GP prior, which conditions on both the objective values and any available auxiliary covariates. The BO algorithm operates in this learned latent space, inherently structured for the BO surrogate model. Ultimately, it is the BO and its acquisition function that drives the preference for higher objective values, as in classical BO.
EXPERIMENTS
We demonstrate the efficacy of method described in Algorithm on simulated datasets well as on a molecular discovery benchmark dataset. In all our experiments, 10% of the training data is used held-out validation set for early-stopping to ensure that the generative model does not overfit. We describe the neural network architectures in the Sec. Fo the Appendices, We use the same BO options and underlying VAE architectures for all methods in a particular experiment. We benchmark against the following methods:
LSO This latent space optimisation method VAE learn low-dimensional representation of the high-dimentn sigh-dimensional dataset. BO is performed over the low-dimensional latent space. LSO does not take into account any auxiliary covariate information and makes use of a standard Gaussian prior over the latent space.
Grosnit et al. (2021) proposed a VAE-based method that tries to construct discriminative latent spaces for VAE based BO methods by incorporating a metric loss term in the ELBO. We compare our model against the triplet loss, log-ratio loss, and contrastive loss.
Triplet Loss (T-LBO) As described in Grosnit et al. (2021), the triplet loss measures distances between input triplets. In other words, this loss tries to introduce a structured space where positive and negative pairs cluster together subject to separation by a margin. The triplet pairs are assigned as a pre-processing step.
Log-Ratio Loss (LR-LBO) This metric loss is described in Kim et al. (2019) and is continuous metric loss that is applied to triplets of inputs. This loss is used with the model described in Grosnit et al. (2021).
Contrastive Loss (C-LBO) This deep metric loss is described in Hadsell et al. (2006). The contrastive loss operates on input pairs by separating the latent encodings based class label information. This is used with the model described in Grosnit et al. (2021).
Local Latent Bayesian Optimisation (LOL-BO) As proposed in Maus al. (2022), this method is a latent space BO approach that addresses the mismatch between the notion of a trust region in the latent space and trust region in the structured input space.
7

---

Published as a conference paper at ICLR 2025
2018
te to maree
112
150
(a)
(c)
Figure 3: Results from our experiments with a synthetic dataset. Lower values are better. (a) Comparing the performance of our model with the LSO benchmark. The dataset comprises 500 instances out of which the target quantity is observed only for 100 instances. O pertains to an additive GP prior VAE over all three covariates and the partially observed quantity of interest c, O₂ to GP prior VAE over only the partially observed target and O3 to an additive GP prior VAE over the partially observed quantity of interest and shift,. (b) Similarly, we also demonstrate our model's 's performance on dataset which comprises 5000 instances out of which the quantity of interest observed only for 500 instances. (c) Effect the choice of latent dimension with the dataset comprising 500 instances. |Z| pertains to the dimensionality of the latent space. All plots depict the mean quantity of interest value with the 95% confidence interval (shaded region) obtained over 100 repetitions with regenerated training data and target images. The grey line pertains to the lowest RMSE in the training set.
5.1 DEMONSTRATION ON SYNTHETIC DATA
We demonstrate our model's ability to perform effective high-dimensional BO by modifying digits from the MNIST dataset. In particular, randomly select instance of digit and resize this digit to dimension of 52 x 52 pixels for larger image space. We perform three different manipulations (which would form our additional auxiliary covariates x)) to this digit: rotation about the centre, shift along the r-axis, and shift along the y-axis, by stochastically choosing these values 500 or 5000 times. Our final training set comprises either 500 or 5000 samples. Furthermore, we define 'targer image', nr with particular rotation and shift values. This target image is not included in the training set and we ensure that the manipulations are sufficiently different from the target values. See Suppl. Fig. 7 for a random sample of the training data.
In this experiment, the black-box function the root mean squared error (RMSE) between the unseen target digit and chosen digit (either digit from the training set or new candidate), i.c., f(y) = RMSEEy - y). Our objective is to find a new digit that minimises the value returned by the black-box function. In other words, we want the quantity of interest to be close to zero as possible. Furthermore, in our training set, we assume that the quantity of interest e = f(y) (i.e. the RMSE between the unseen target and the digit) is known for only few digits (i.e., [Te] equals 100 or 500) and unobserved for the rest.
To ensure that there is sufficient stochasticity in the choice of targets, repeat our experiments as well as the generation of data 100 times. We fit our GP prior VAE-based method by making use of the partially observed target quantity of interest as well as subset of) the auxiliary covariates æ. We make use of Algorithm with v = 10 as well as B = 200. We experiment with different choices of kernels to empirically obtain the optimal model.
Fig. 3 demonstrates our experiments on the synthetic dataset and we can see that our method finds candidate points with a quantity of interest (or RMSE) that are significantly lower than those in the training set. In Fig. 3(a) we demonstrate the performance of our model on 500 instances out of which the quantity of interest is observed for 100 instances. The LSO method trained with all 500 instances and the BO computation is performed using the 100 instances for which the quantity of interest observed. Our method outperforms the LSO model already when the GP prior VAE is fitted with the partially observed target quantity (0z), and results improves further if additional auxiliary covariates are available (O; and ²)).
Similarly, Fig. 3(b) demonstrates that our method outperforms the baseline LSO with 5000 instances out of which the quantity of interest is observed for 500 instances and Fig. 3(c) demonstrates the effect of the choice of latent dimension.
8

---

Published as a conference paper at ICLR 2025
We visualise the learnt latent space and the BO steps taken in Suppl. Fig. 10. Our method learns latent representations where the target quantity increases smoothly from the lower left corner to the upper right corner (though visualisation). The BO steps explore the region of the latent space where the target quantity has high values. In Suppl. Fig. 6, we perform ablations with different subsets of additional auxiliary covariates. Furthermore, we demonstrate the performance of Vanilla VAE BO in Suppl. Fig.
again emphasise that this is a 2-D UMAP (MeInnes et al., 2018)
5.2 EXPRESSION RECONSTRUCTION
We consider the common task of generating single-variable mathematical expressions from formal grammar (Kusner et al., 2017; Tripp et al., 2020; Grosnit et al., 2021; Maus et al., 2022). The objective mean squared error (MSE) (defined as log(1 + MSE))
to minimise distancefregret based between a generated expression and the target expression, 2' * sin * z). We followed the data preparation proposed by Grosnit et al. (2021) to obtain 40000 data points and augmented the data with 8 additional covariates (count of the elements Y2, ** '+', 'exp', 'sine", '1', '2', and '3' in the expressions) which can be easily gleaned from the expressions. In order to appropriately handle the mathematical expressions, we use the Grammar VAE (Kusner et al., 2017). To demonstrate the efficacy of our method, we make use of Algorithm 1 with kernel over the additional covariates and regret. Fig. 4(a) demonstrates that our method achieves competitive performance against the benchmark methods. In Suppl. Fig. 11, we visualise the mean regret achieved by our method together with the 95% confidence interval.
10 and B 500 as well an additive
5.3 MOLECULE OPTIMISATION
al. (2018), which consists of 250000 drug-like like commercially available molecules extracted from the ZINC database (Irwin et al., 2012) public dataset for ligand discovery. The dataset includes the molecular structures in the SMILES string representation (Weininger, 1988) and three molecular properties: the water-octanol partition coefficient (logP), the Synthetic Accessibility Score (SAS), and the Quantitative Estimation of Drug-likeness (QED) (Bickerton et al., 2012). The objective of the task is to maximise the penalised logP which is defined the logP penalised by the SAS and the number of long cycles: penalised logP(m) cycle(-) the number of long cycles.
We use the ZINC-250K molecular dataset used in
logP(m) - SAS(m) - cycle(m) where m is the molecular instance and
We augmented the ZINC-250K with five additional covariates: molecular weight, number of hydrogen donors, number of hydrogen acceptors, number of rotatable bonds, and total polar surface area. These values were computed using a popular open-source chem-informatie 2013) (see Suppl. Fig. 12 for a visualisation of the distribution of these properties in the form of histograms). Including QED and SAS, there are seven additional auxiliary covariates and the penalised logP is the quantity of interest which we are trying to maximise. For a new molecule, is possible to compute the penalised logP using RDKit (acting we assume that the penalised logP is partially observed (observed for only 1% of the data).
tool, RDKit (Landrum et al.,
black-box function). Furthermore,
We demonstrate the ability of our model optimise the structure of the molecule in order maximise a property of interest (the penalised logP). To handle the SMILES representation of the molecules, we use the Junction Tree VAE (JT-VAE) (Jin al., 2018), which introduced an encoder and decoder suitable to molecular graphs In our experiments, we extend the implementation by Grosnit et al. (2021). We note that our method is not limited JT-VAE but can be applied with any latent-variable model. Furthermore, we use JT- VAE with all the baseline methods for a fair comparison.
We use Algorithm with v 10 and B 450. In Fig. 4(b) we demonstrate that our method is able to identify candidate molecules that have higher penalised logP value than competing methods. Furthermore, in Fig. 4(c), we demonstrate the performance of our method with different choices of auxiliary covariates for the additive kernel of the GP prior. The results show that it is indeed beneficial to include the additional auxiliary covariate information, whenever they are available in an application, along with the partially observed quantity of interest (penalised logP). See Suppl. Fig. 13 for a visualisation of the marginal variance of the additive components. We visualise the mean and standard deviation of the marginal variance for each of the kernel components across the 56 latent dimensions. In Suppl. Fig. 14, we demonstrate the performance of our method with fewer instances for which the penalised logP observed. For this experiment, we assumed that the penalised log P is

---

Published as a conference paper at ICLR 2025
5L2 Creanhatal. 2022
55 Tres stat. 2024 Lnta
L5D1 2E231
LSD 7i5a R. 38211
25-
(Nearttal
Sormshes
LMAEO
Rogget
logp
183
700
503
123 X0
163
(a)
(b)
(c)
Figure 4: Results from the expression reconstruction and molecule optimisation experiments, (a) The mean regret achieved over repetitions. O₂ pertains to a kernel over only the target quantity. Lower values are better. (b) The penalised logP score achieved by our method compared to competing methods. (c) Comparison of the penalised logP score achieved by our method using different choices of additional covariates in the additive kernel of the GP prior VAE. θ1 pertains an additive kemnll over all 7 covariates and the partially observed target quantity, 2₂ to a kernel over only the partially observed target quantity, and G3 to additive kernel over only the 5 additional properties that were calculated with RDKit (and does not include the partially observed target quantity). In figures (b) and (c), the mean penalised logp score
10 repetitions is visualised together with the 95% confidence interval (shaded region). The grey line pertains to the highest penalised logP score in the training set. Higher values are better
observed for only 0.1% of the data. The overall performance of all the methods decreases because the BO has fewer points to fit the surrogate model with. However, it is interesting that our method with an additive kernel that does not include the quantity of interest performs the best. We postulate that this is because only a few instances of the quantity of interest are observed and hence the estimated values for the unobserved quantities of interest have low quality. We believe that this demonstrates that our model can learn meaningful latent representations without making use of the partially observed quantity of interest and can be applied to datasets with only few labelled instances.
In Suppl. Fig 15, visualise the latent space using t-SNE (Van der Maaten & Hinton, 2008) and colour the latent embedding by the respective molecular properties. We note that the model learns a latent embedding that changes smoothly with respect the target quantity well as with the respect to the additional covariates.
5.4 EVALUATING LATENT SPACE STRUCTURE FOR GAUSSIAN PROCESSES
We evaluate our model's s ability to construct meaningful discriminative latent spaces for Gaussian processes (GPs). Following an approach similar to Grosnit et al. (2021), we leverage the trained encoder to map data points from the original space, y E V, onto a low-dimensional latent space, E Z RL, where we fit a GP using the original labels. To assess whether structured latent representations enhance GP generalisation, we ensure a unified experimental setup across all tasks. Specifically, we use 80% of the encoded latent points (from the respective training splits) to train sparse GP with 500 inducing points and compute the predictive log- logllikelhood on the remaining 20% of held-out data. This experiment provides insights into the impact of clustered latent inputs GP regression-asd an essential factor in the Bayesian Optimisation (BO) process. In Suppl. Table 2, we show that method achieves the highest predictive log-likelihood, highlighting how the discriminative latent space enhances GP generalisation.
6 CONCLUSION
In this paper, we proposed a novel GP prior VAE-based method to perform high-dimensional BO. We demonstrated the efficacy of our method simulated datasets as well as in the discovery of novel molecules that optimise quantity of interest. Our method shows that it can be beneficial to include auxiliary covariates (even partially observed) for performing BO in the latent space. Furthermore, our approach can efficiently handle partially observed target quantities. Given the flexibility and performance of our model, expect approach be beneficial to scalable BO.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We gratefully acknowledge the computational resources provided by Aalto Science-IT, Finland. Our sincere thanks also go to Charles Gadd for his valuable discussions.
REFERENCES
G Richard Bickerton, Gaia V Paolini, Jérémy Besnard, Sorel Muresan, and Andrew L Hopkins. Quantifying the chemical beauty of drugs. Nature chemistry, (2):900-9,, 2012.
Mickael Binois and Nathan Wycoff. A survey on high- ligh-dimensional gaussian process modeling with application to bayesian optimization. ACM Transactions on Evolutionary Learning and Optimization, 2(2):1-26, 2022.
Francesco Paolo Casale, Adrian Dalca, Luca Saglietti, Jennifer Listgarten, and Nicolo Fusi. Gaussian process prior variational autoencoders. In Advances in Neural Information Processing Systems, 2018.
Stephan Eissman, Daniel Levy, Rui Shu, Stefan Bartzsch, and Stefano Ermon. Bayesian optimization and attribute adjustment. In Proceedings of the 34th Conference on Uncertainty in Artificial Intelligence, 2018.
Kobi Felton, Daniel Wigh, and Alexei Lapkin. Multi-task bayesian optimization of chemical reactions. In Machine Learning for Molecules Workshop at NeurlPS, 2020.
Vincent Fortuin, Dmitry Baranchuk, Gunnar Rätsch, and Stephan Mandt. Gp-vae: Deep probabilistic time series imputation. In International conference artificial intelligence and statistics, pp.
1651-1661. PMLR, 2020.
Peter I Frazier. Bayesian optimization. In Recent advances in optimization and modeling of contemporary problems, pp. 255-2552220 258. Informs, 2018.
Rafael Gómez--Bobarelll, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato. Timothy D Hirzel, Benjamín Sánchez- Lengeliig. Dennis Sheberla, Jorge design using data-driven continuous Ryan P Adams, and Alán Aspuru- Guzik. Automatic chemical representation of molecules. ACS central science, 4(2):268 276.. 2018.
Ryan-Rhys Griffiths and José Miguel Hermández-Lobato Constrained bayesian optimization for automatic chemical design using variational autoencoders. Chemical science, 11(2):577-586, 2020.
Antoine Grosnit, Rasul Tutunov, Alexandre Max Maraval, Ryan-Rhys Griffiths, Alexander Cowen- Rivers, Lin Yang, Lin Zhu, Wenlong Lyu, Zhitang Chen, Jun Wang, et al. High-dimensional bayesian optimisation with variational autoencoders and deep metric learning. arXiv preprint 1oXiv:2/06603609 2021.
Raia Hadsell, Sumit Chopra, and Yann LeCun. Dimensionality reduction by learning an invariant mapping. In IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 06), volume 2, pp. 1735-1742. IEEE, 2006.
James Hensman, Nicolo Fusi, and Neil D Lawrence. Gaussian processes for big data. In Proceedings of the 29th Conference on Uncertainty and Artificial Intelligence, 2013.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. 3-VAE: Learning basic visual concepts with a constrained variational framework. International conference learning representations, 2017.
John J Irwin, Teague Sterling, Michael M Mysinger, Erin S Bolstad, and Ryan G Coleman. Zine: free tool to discover chemistry for biology. Journal of chemical information and modeling, 52(7): 1757-1768, 2012.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In International Conference on Machine Learning, pp. 2323-2332. PMLR, 2018.
11

---

Published as a conference paper at ICLR 2025
Kirthevasan Kandasamy, Willie Neiswanger, Jeff Schneider, Barnabas Poczos, and Eric P Xing. Neural architecture search with bayesian optimisation and optimal transport. Advances Neural Information Processing Systems, 31, 2018.
Sungyeon Kim, Minkyo Seo, Ivan Laptev, Minsu Cho, and Suha Kwak. Deep metric learning beyond binary supervision. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2288 229722297222 2019.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In 2nd International Conference on Learning Representations, ICLR, 2014.
Ksenia Korovina, Sailun Xu, Kirthevasan Kandasamy, Willie Neiswanger, Barnabas Poczos, Jeff Schneider, and Eric Xing. Chembo: Bayesian optimization of small organic molecules with syn- thesizable recommendations. In International Conference on Artificial Intelligence and Statistics, pp. 3393-3403. PMLR, 2020.
H. J. Kushner. A method of locating the maximum point of arbitrary multipeak curve in the presence of noise. Journal of Basic Engineering, 86(1) (6(1):97-106, 03 1964.
Harold J Kushner. A versatile stochastic model a function of unknown and time varying form. Journal of Mathematical Analysis and Applications, 5(1):150-167, 1962.
Matt J Kusner, Brooks Paige, and José Miguel Hermández-Lobato. Grammar variational autoencoder. In International Conference on Machine Learning, pp. 1954. PMLR, 2017.
Greg Landrum et al. Rdkit: software suite for cheminformatics, computational chemistry, and predictive modeling. Greg Landrum, 8, 2013.
Neil Lawrence, Matthias Seeger, and Ralf Herbrich. Fast sparse gaussian process methods: The informative vector machine. Advances in Neural Information Processing Systems, 15, 2002.
Seunghun Lee, Jaewon Chu, Sihycon Kim, Juyeon Ko, and Hyunwoo Kim. Advancing bayesian optimization via learning correlated latent space. Advances in Neural Information Processing Systems, 36, 2024.
Xiaoyu Lu, Javier Gonzalez, Zhenwen Dai, and Neil Lawrence. Structured variationally auto- encoded optimization. In International Conference on Machine Learning, pp. 3267-3275. PMLR, 2018.
Pierre-A Aleee-Ale Mattei and Jes Frellsen. Miwae: Deep generative modelling and imputation of incomplete data sets. In International Conference on Machine Learning, pp. 4413-4423. PMLR, 2019.
Natalie Maus, Haydn T Jones, Juston Moore, Matt Kusner, John Bradshaw, and Jacob R Gardner. Local latent space bayesian optimization over structured inputs. Advances in Neural Information Processing Systems, 2022.
Natalie Maus, Kaiwen Wu, David Eriksson, and Jacob Gardner. Discovering many diverse solutions with bayesian optimization. In Francisco Ruiz, Jennifer Dy, and Jan-Willem van de Meent (eds.), Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pp. 1779-1798. PMLR, 25-27 Apr 2023.
Leland McInnes, John Healy, Nathaniel Saul, and Lukas Großberger. Umap: Uniform manifold approximation and projection. Journal of Open Source Software, 3(29):861, 2018, doi: 10.21 105/ 1os..00861. URL https: /to.ooooo///ooo/ooooo orgg 10. 21105 joss. 00861.
Jonas Mockus. The Bayesian approach to local optimization. Springer, 1989.
Riccardo Moriconi, Marc Peter Deisenroth, and KS Sesh Kumar. High-dimensional bayesian optimization using low-dimensional feature spaces. Machine Learning, 109(9): 1925-1943, 2020.
Kevin P Murphy. Probabilistic machine learning: Advanced topics. MIT Press, 2023.
Alfredo Nazabal, Pablo M Olmos, Zoubin Ghahramani, and Isabel Valera. Handling incomplete heterogeneous data using vaes. Pattern Recognition, 107:107501, 2020.
12

---

Published as a conference paper at ICLR 2025
Pascal Notin, José Miguel Hrnndnddz--Lbbbte and Yarin Gal. Improving black-box optimization in vae latent space using decoder uncertainty. Advances Neural Information Processing Systems, 34, 2021.
Joaquin Geinonero-Candeee and Carl Edward Rasmussen. A unifying view of sparse approximate gaussian process regression. The Journal of Machine Learning Research,
2005.
Siddharth Ramchandran, Gleb Tikhonov, Kalle Kujanpää, Miika Koskinen, and Harri Lähdesmäki. Longitudinal variational autoencoder. In International Conference Artificial Intelligence and Statistics, pp. 3898-3906. PMLR, 2021.
Siddharth Ramchandran, Gleb Tikhonov, Otto Lönnroth, Pekka Tiikkainen, and Harri Lähdesmäki. Learning conditional variational autoencoders with missing covariates. Pattern Recognition, 147: 110113, 2024.
Danilo Jimenez. Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning, pp. 1278-1286. PMLR, 2014.
Binxin Ru, Xingchen Wan, Xiaowen Dong, and Michael Osborne. Interpretable neural architecture search via bayesian optimisation with weisfeiler-lchman kernels. In International Conference on Learning Representations, 2021.
Eric Schulz, Maarten Speekenbrink, and Andreas Krause. A tutorial on gaussian process regression: Modelling, exploring, and exploiting functions. Journal of Mathematical Psychology, 85:1-16, 2018.
Bobak Shahriari. Kevin Swersky, Ziyu Wang, Ryan Adams, and Nando De Freitas. Taking the human out of the loop: A review of bayesian optimization. Proceedings of the IEEE, 104(1): 148-175, 2015.
Benjamin J Shields, Jason Stevens, Jun Li, Marvin Parasram, Farhan Damani, Jesus I Martinez Alvarado, Jacob M Janey, Ryan P Adams, and Abigail G Doyle. Bayesian reaction optimization as tool for chemical synthesis. Nature, 590(7844):89-96 2021.
Eero Siivola, Andrei Paleyes, Javier González, and Aki Vehtari. Good practices for bayesian optimization of high dimensional structured spaces. Applied AI Letters, 2(2):e24, 2021.
Alex Smola and Peter Bartlett. Sparse greedy gaussian process regression. Advances in Neural Information Processing Systems, 13, 2000.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in Neural Information Processing Systems, 28, 2015.
Samuel Stanton, Wesley Maddox, Nate Gruver, Phillip Maffettone, Emily Delaney, Peyton Greenside, and Andrew Gordon Wilson. Accelerating bayesian optimization for biological sequence design with denoising autoencoders. International Conference on Machine Learning, pp. 20459-20478. PMLR. 2022.
Michalis Titsias. Variational learning of inducing variables in sparse gaussian processes. In Artificial Intelligence and statistics, pp. 577- PMLR, 2009.
Austin Tripp, Erik Daxberger, and José Miguel Hernández-Lobato Sample-efficient optimization in the latent space of deep generative models via weighted retraining. Advances in Neural Information Processing Systems, 33: 13:1912911222. 2020.
Raquel Urtasum and Trevor Darrell. Discriminative gaussian process latent variable model for classification. In Proceedings of the 24th International Conference on Machine Learning, pp. 927-934, 2007.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
13

---

Published as a conference paper at ICLR 2025
David Weininger. SMILES, a chemical language and information system. 1. introduction to methodo- logy and encoding rules. Journal of Chemical Information and Computer Sciences, 28(1): 1988.
Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning, volume 2. MIT press Cambridge, MA, 2006.
14

---

Published as a conference paper at ICLR 2025
Appendices
A DERIVATION OF THE ELBO
Following (Kingma & Welling, 2014), the ELBO of the marginal log-likelihood for the standard VAE model can be written as:
N Σ Eng P [=,] - KL 4s (z,, (2m (pn)l +
Pw) 2 L(0,0,0;3
max, o.v.0
(2)
R=1
where Φ and vi are the encoder and decoder weights respectively, z refers the low-dimensional latent representation, y refers to the observations, and KL refers to the Kullback-Leible divergence.
Since use GP prior VAEs that assume independent priors for each of the latent dimensions, write the joint conditional prior as:
L
L N 21 O, xx
Po(Z X) = II
Po ( & X) = II 1=1
where 2, = (X) = ) <(x) ~ (P(p)( la(a, x' |0,) such that (a(()) is the mean, (2r, xc' (1) is the covariance function and 0, denotes the parameters of the covariance function, and K x is a N x N covariance matrix for the 1th latent dimension.
We summarise the observed and unobserved parts as Y = (Y", Y") and X = -Xx,, (X), and write the ELBO as:
> Po Pv ( (2, 11,
(3)
We use a conditionally independent factorisable variational approximation:
("|'0 X²) = 4o (Z|Y
"orr, X9) (X"
u ( (2)
(4)
We assume that 4s ( 2)) factorises also across the latent dimensions, which allows us write the variational approximation alternatively as
N
co(2,X'1o, Xo) (Z X X9) (X' = II (z (x))))) II (e))))) (5)
Following Ramchandran et al. (2024), simplify the KL term in Eq. 3 as:
po. (L'
X {0)]
[X)] x Egg KL(go(Z11 X (Z|X"
+ KLLga ""
Bys KL loo {-1]
KL KL (1 2
=i) (6)
<D
D I po po(2. X0)
by using the assumption of factorising latent space, po p0 Xu, 80) and mean-field normal posterior for (2,2; Xo), with a variational mean ;= (f. .i........................ (* y) Wi diag(o?, yp),. ., 2 (m) v for the 1hh latent dimension. The expectation in Eq. 6 is w.r.t. the unobserved auxiliary covariates X that are the inputs to the GP kernel and, therefore, the expectation does not have a closed form but can be approximated by Monte Carlo sampling.
and
covariance
matrix
There are several different approaches to approximate the GP prior in order for the ELBO in Eq. to scale to large datasets. We make use of mini-batch compatible approach proposed by Ramchandran et al. (2021) that uses the inducing point method (Titsias, 2009; Hensman al., 2013) and exploits the structure of the GP prior, described in the section.
15

---

Published as a conference paper at ICLR 2025
A.1 SCALABLE COMPUTATION AND MINIBATCHING
Each of the KL divergences KL q [ Y l (z:[X"".X)) in Eq. 6 has a computation com- plexity of ((N3). Below we drop the index of the latent dimension, 1, for simplicity. Relying on the derivation proposed by Ramchandran et al. (2021) obtain scalable ELBO,
use the low-rank inducing point approximation for GPs and use M inducing locations S =81 8M) in X' and the corresponding inducing function values u, (71) (s......, (sM)))
(WI,.... w(M)
) for each
latent dimension (Hensman et al., 2013). We explicitly keep track of the distribution of the Gaussian inducing values u, (m,, Hi), where mi and H, are global variational parameters. We can then derive an upper-bound for the KL divergence KL[W( WO K xxx)] I DKL well as an unbiased, batch-normalised partial sum over subset of indices, (I,..., N} of size | N such that DKI DKL where
ŌKL
1N 2
E(0=² iEZ
- as)² + o. 0=20, - 2 KKn
N log o²
N 2
+ toz 2r (K (Ksa, Ka,s)) - log loggo?) (LI[(M,)))VV// Kss)),
where na = (ā, v8) and of (ā yt) are the encoder means and variances, K₃ denotes the 7th diagonal element of K = K x Ks K sx. and Kas as well as Ka.s KTa.. Sa. are defined similarly as K Xxx- The conditional probability px(X (X)) Eq. simplifies to 0 = px (). As described in (Ramchandran et al., 2024), PA ()") can be an informative prior and D amenable to mini-batching.
Therefore, the ELBO for GP prior VAE models that marginalises missing covariates and affords efficient optimisation with stochastic gradient descent is obtained from the Eqs. 6, 7. For more detailed derivation, please refer to Ramchandran et al. (2021; 2024).
B THE EXPECTED IMPROVEMENT ACQUISITION FUNCTION
The expected improvement acquisition function estimates improvement that would be achieved when choosing a specific point z as the next point to query. It balances between exploration and exploitation of the black-box function f. The expected improvement (EI) defined as
a(z) EI(2)) = f max max( ([(=)- uu(2) df(z). ((2)
Improvement
where f is our current optimum, is instance in the data space, +- Iv- (-1)) where di pertains to the trained decoder weights and 2 is the chosen latent space location, and (p()) is the probability density function of the standard normal distribution, ((0 1). The expected improvement can be analytically evaluated under the GP surrogate model:
(u(2) f()) - (T) (Z())((T)
if >
EI(2)= =
(4)
(45)
(8)
if
where,
u(x) 1(0)=
of if of) 0.
T
((2)
In the above equation, μ²(2) and (2) are the mean and standard deviation of the surrogate GP posterior predictive at z. & and ip are the cumulative distributive function and probability density function of the standard normal distribution, respectively.
16

---

Published as a conference paper at ICLR 2025
In Eq. 8. part (i) corresponds to the exploitation term and part (ti) corresponds the exploration term. The parameter & controls the amount of trade-off between exploration and exploitation (higher values leads to more exploration). We set & to be 0.01 in all our experiments recommended default value). For a detailed review of the expected improvement and other popular acquisition functions refer the reader (Frazier, 2018; Garnett, 2023).
C OPTIMISATION AND PRACTICAL CONSIDERATIONS
To maximise the evidence lower bound, an adaptive learning rate method that maintains an exponentially decaying average of past gradients well squared gradients. The parameters that need to be optimised include the neural network weights for the encoder (Ø) as well as decoder (t)) and the GP kernel parameters (0). Moreover, we separately fit the GP surrogate model for the Bayesian optimisation. In the case of mini-bater training, the optimisation steps are conducted interchangeably with natural gradient-based updates of the variational parameters.
the Adam optimiser (Kingma & Ba, 2015), which is
We use PyTorch (Paszke et al., 2019) for the inference implementation which allows the computation of derivatives using automatic differentiation and we use BoTorch (Balandat et al., 2020) for Bayesian optimisation. For all experiments we set the frequency of retraining y 10 and the stopping criterion r = 0.1. set the number of latent dimensions expression reconstruction experiment, and 56 for the molecular discovery experiment.
for the synthetic dataset experiment, for
D EXPERIMENT WITH VANILLA VAE BO
We demonstrate the performance of Vanilla VAE BO (i.e. weighted retraining) using synthetic data. From Suppl. Fig. 5, we can clearly see that our method as well as the other baselines demonstrate better performance.
VOU E IE
2014
111
[a]
(b)
Figure 5: Results from our experiments with synthetic dataset. Lower values are better. (a) Comparing the performance of our model with the LSO benchmark. The dataset comprises 500 instances out of which the target quantity is observed only for 100 instances. O pertains to an additive GP prior VAE over all three covariates and the partially observed quantity interest c, O2 to GP prior VAE over only the partially observed target and O3 to an additive GP prior VAE over the partially observed quantity of interest e and shift,. (b) Similarly, we also demonstrate our model's performance dataset which comprises 5000 instances out of which the quantity of interest is observed only for 500 instances. All plots depict the mean quantity of interest value with the 95% confidence interval (shaded region) obtained over 100 repetitions with regenerated training data and target images. The grey line pertains to the lowest RMSE in the training set.
E ABLATION STUDY
We demonstrate how method performs with different subsets of additional auxiliary covariates. The ablations were run on the simulated data with 5000 samples and the target quantity of interest was observed for 500 samples. Fig. 6 depicts the mean target quantity of interest obtained over 100 repetitions with regenerated training data and target images.
17

---

Published as a conference paper at ICLR 2025
rradtin Whh
24
DET. EL HORTA
RMSE
26
Number of BO steps
Figure Results from experimenting with the choice of kernel for the synthetic dataset.
NEURAL NETWORK ARCHITECTURES
F.1 SYNTHETIC DATASET
Table I: Neural network architectures used in the simulated dataset.
Hyperparameter
Value
Dimensionality of input
52 x 52
Number of convolution layers Number of filters per convolution layer
3
144
Kernel size Stride
3 3
2
Inference
Pooling
Max pooling
network
Pooling keree size
2 2
Pooling stride
2
Number of feedforward layers
2
Width of feedforward layers Dimensionality of latent space
500, 50
8
Activation function of layers
RELU
Dimensionality of input
8
Number of transposed convolution layers
3
Number filters per transposed convolution layer
Generative network
Kernel size
4 4
Stride
2
Padding
2
Number of feedforward layers
2
Width feedforward layers
50,
Activation function of layers
RELU
In the synthetic data experiment, we use the neural network architecture described in Table
F.2 EXPRESSION RECONSTRUCTION
In the expression reconstruction experiment the Grammar VAE (Kusner et al., 2017) which is a computational model used in natural language processing and generative modelling, It combines principles from VAEs and context-frea grammars to learn and generate structured sequences of symbols, such as sentences, mathematical expressions, or code.
In other words, the Grammar VAE extends the VAE framework to handle structured data where the order and relationships between elements matter. Furthermore. Grammar VAE incorporates context-free grammars which define the syntax and structure of sequences. This allows the model to
18

---

Published as a conference paper at ICLR 2025
capture the hierarchical and compositional nature of sequences, making it well-suited for generating structured outputs. In addition to this, Kusner et al. (2017) proposed to represent the discrete data using a parse tree from the context-free grammar. Therefore, the model is variational autoencoder which encodes and decodes directly to and from the generated parse trees while ensuring that the generated outputs are always valid.
In our experiments, we use a latent space dimension of 25 and make use of the neural network architecture specified in (Kusner al., 2017).
F.3 MOLECULE OPTIMISATION
In the molecule generation experiment we use the Junction Tree VAB (JT-VAE) (Jin et al., 2018) which extends VAEs to molecular graphs by introducing suitable encoder and decoder. The encoder learns two latent representations: one that encodes the tree structure and high-level cluster information while the other encodes fine-gra ane-grained connectivity details. In particular, the model generates a molecular graph in two phases: first it generates tree-structured scaffold over chemical sub-structure and then combines them into molecules with a graph message passing network. The molecule is encoded into two latent representations: the clusters that are the tree without fully capturing how exactly the clusters are mutually connected. The graph capture the fine-grained connectivity encoded by zg-
[2T, =c] where 2T encodes the tree structure and the information of
The latent representation is then decoded back into a molecular graph in two stages. First, reproduce the junction using a tree decoder. Then, predict the fine-grained connectivity between the clusters in the junction tree using a graph decoder to obtain the full molecular graph. The decoder generates the molecule piece-by-pin by-rece utilising the components and how they interact instead assembling the molecule atom-by-atom and/or through chemically invalid intermediaries.
The graph encoder is graph message passing network (graph neural networks), the tree encoder is a tree message passing network (related to RNNs and tree-LSTM), the junction tree is reconstructed using a structured tree decoder, and the fine-grained cluster details are obtained using a graph decoder. Furthermore, the latent space dimension is 56 (tree and graph representations are 28 dimensions each).
The key benefit is the incremental expansion of the molecule while maintaining chemical validity at every step. Furthermore, each molccule is built from sub-graphs chosen out a vocabulary of valid components. In this work, 2018). Our proposed model is agnostic to the choice of the underlying neural network architecture.
the same neural network architecture specification in (Jin et al.,
G SUPPLEMENTARY IMAGES
19

---

Published as a conference paper at ICLR 2025
Figure random sample of digits (from the the synthetic dataset) that have been rotated and shifted along the and axis.
RMSE
50
Figure 8: Violin plot visualising the distribution of the quantity of interest in instance of the synthetic dataset.
20

---

Published as a conference paper at ICLR 2025
PANSE
15-
steps
Figure 9: A demonstration of the effect of B (as in B-VAE (Higgins al., 2017)) on the model performance in the synthetic dataset.
10
80
70
150
60
125
Quantity in
6
80 Steps
50
40
4
75
30
50
20
10
0-
2
4
6
10
Figure 10: Visualisation of the latent space in the synthetic data experiment. We performed projection of the latent space down two dimensions using UMAP (McInnes al., 2018). The inverted triangles refer to the BO steps and the blue cross refers to the latent space representation of the "optimal" instance which hope our method would find (not included in the training set). The latent embedding is coloured by the quantity of interest.
21

---

Published as a conference paper at ICLR 2025
3.5
T-LBO [Grosnit et al., 2021]
LSO [Tripp et al., 2020]
LOL-BO [Maus et al., 2022]
2.5
C-LBO
LR-LBO
Our method
Regret
2.0
0.5
100
200
300
400
500
Number BO steps
Figure 11: Results from the expression reconstruction experiment. The mean regret achieved by our method compared to competing methods over 5 repetitions is visualised together with the 95% confidence interval (shaded region). Lower values are better.
22

---

Published as a conference paper at ICLR 2025
(a)
(b)
(c)
(d)
17000
3:30
2.00
(e)
(f)
$1100
41100
2000
(g)
(h)
Figure 12: Histograms visualising the distribution of the properties in the ZINC-250K dataset.
23

---

Published as a conference paper at ICLR 2025
logo-
MolWt
No. of H donors
No. acceptors
No, rotatable bonds
TPSA
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
Marginal variance (log)
Figure 13: Marginal variance of the additive components in the molecule discovery experiment. We visualise the mean and standard deviation of the marginal variance for each of the kernel components across the 56 latent dimensions. This pertains an additive kernel over all the covariates.
[Grosnit al., 2021]
25
LSO [Tripp al., 2020]
Penalised logP-score
LOL-BO [Maus et al., 2022]
Our method (with 0,)
20
Our method (with 0,)
15
100
400
Number of BO steps
Figure 14: Results from the molecule optimisation experiment with the penalised logP observed for the
only 0.1% of the data. The mean penalised logP score over repetitions is visualised together with the 95% confidence interval (shaded region). The grey line pertains to the highest penalised logP score the training set. ④ pertains to an additive kernel over all 7 covariates and the partially observed target quantity, and O₃ pertains an additive kernel over only the 5 additional properties that were calculated with RDKit (and does not include the partially observed target quantity). The overall performance of all the methods decrease because the Bayesian optimisation has fewer number of points to fit the surrogate model with. However, is interesting to note that our method with additive kernel that does not include the quantity of interest performs the best (by small margin). The other approaches make use of the quantity of interest (penalised logP) while fitting the model. We believe that this demonstrates that our model (with +₂) able to learn meaningful latent representations without making use of the partially observed target quantity of interest. Higher values are better.
24

---

Published as a conference paper at ICLR 2025
acoeptrr
03
-299
186
102
306
108
200
-59
46
<3
%
-26
116 -56 59
-203
Figure 15: Visualisation of the latent space the molecule discovery experiment. We performed projection of the latent space from 56 dimensions down to two dimensions using t-SNE (Van der Maaten Hinton, 2008) and using the validation dataset for convenience. The latent embedding is coloured by the respective molecular properties. The proposed model learns latent embedding that changes smoothly with respect to the target quantity as well as with respect to the additional covariates (noting again that the visualisation corresponds to 2-D t-SNE embedding).
25

---

Published as a conference paper at ICLR 2025
H SUPPLEMENTARY TABLES
Table 2: This table illustrates how separation the latent space enhances GP generalisation. It reports the GP predictive log-likelihood the held-out validation sets, along with the standard deviation (#) on the validation set. Higher (less negative) values are better.
Expression reconstruction
Molecular discovery
LSO (Iripp et al., 2020)
-3.1 + 0.09
-2.01 = 0.25
T-LBO (Grosnit et al., 2021)
1.85 1 0.07
-1.49 1 0.29
LOO--BO (Maus et al., 2022)
-1.75 0.07
-1.39 0.25
Our method
-1.72 + 0.08
-1.37 + 0.27
Table 3: Average run time / wall clock time. In the synthetic dataset experiment 200 BO steps are performed and in the molecule discovery experiment 450 BO steps are performed.
Method
Exporiment
Configuration
GPU type
CPU type
Runtime (avg.)
Kerne 9
AMD M1250x
AMD EPYC "Trento"
152 mins
Synthetic data (5000 obs.)
Kernel
M1250x
EPYC "Teetoo
140 mins
Kernel 6
AMD MI250x
AMD EPYC "Trento"
163 mins
Our method
Expression reconstruction
Nvidia Tesla V100
Intel Xeen Gold 6134
1064 mins
Kernel
Nvidia Tesla V100
Intel Xeon Gold 6134
1682 mins
Molecular discovery
Kernel 6
Nvidia Tesla V100
losel Xcon Gold 6134
1641 mins
Kernel 6
Nvidia Tesla V100
Intel Xeen Gold 6134
1668 mins
Synthetic data (5000 obs.)
MI250x
EPYC "Trento"
102 mins
1.50 (Tripp et al., 2020)
Expression reconstraction
Nvidia Tesla V100
Insel Xeon Gold 6134
723 mins
Molecular discovery
Nvidia Tesla V100
Intel Xeon Gold 6134
1038 mins
Expression reconstruction
Nvidia
mins
T-LBO (Grosnit al., 2021)
Molecalar discovery
Nvidia Tesla V100
Intel Xeon Gold 61
1582 mins
LOL-BO (Maus
Molecular discovery
Expression reconstruction
Nvidia Tesla V100
Intel Xeen Gold 6134
802 mins
I LIMITATIONS
While our method proposes a novel approach to performing high-dimensional Bayesian optimisation efficiently, it shares several the limitations of standard VAEs. For example:
It can be challenging to model complex (or multi modall data.
The performance is dependent on the expressiveness of the chosen neural network architec- ture for the encoder and decoder.
The latent space is assumed to follow a Gaussian distribution which may not hold true for all datasets.
Sensitivity to the hyperparameter values such dimensionality of the latent space, weight the KL divergence (3), minibatch size, etc.
Furthermore, in GP prior VAEs, the choice the auxiliary covariates used for the GP prior needs to be done empirically. Despite these limitations, GP prior VAEs have been successful in various ap- plications and have contributed to advances in generative modelling and unsupervised representation learning.
J BROADER IMPACTS
Generative machine learning models have gained significant attention in recent times. In this work, we make use the variational autoencoder which has been primarily used for representation learning, imputation, and data generation tasks. However, VAEs (and deep generative models in general) present several socictal implications that extend beyond the scope of academia and research. Furthermore, they present ethical considerations due to their potential malicious applications including contributing
26

---

Published as a conference paper at ICLR 2025
to misinformation and possible privacy concerns. Robust frameworks as well guidelines need to be established to address these concerns and to ensure the responsible deployment of generative machine learning technologies. It is essential navigate the ethical concerns and to ensure the responsible use of deep generative models for the hetterment of society.
REFERENCES
Maximilian Balandat, Brian Karrer, Daniel R. Jiang, Samuel Daulton, Benjamin Letham, An- drew Gordon Wilson, and Eytan Bakshy. BoTorch: A Framework for Efficient Monte-Cart Bayesian Optimization. In Advances in Neural Information Processing Systems 33, 2020. URL http://arrt htt.aaaiwvv...aaaaaaaaaaaaaa aatpaa//s/9./.// 06403.
Peter Frazier. Bayesian optimization. In Recent advances in optimization and modeling of contemporary problems, pp. 255-278. Informs, 2018.
Roman Gamett. Bayesian optimization. Cambridge University Press, 2023.
Antoine Grosnit, Rasul Tutunov, Alexandre Max Maraval, Ryan-Rhys Griffiths, Alexander Cowen- Rivers, Lin Yang, Lin Zhu, Wenlong Lyu, Zhitang Chen, Jun Wang, et al. High-dimensional bayesian optimisation with variational autoencoders and deep metric learning. arXiv preprint 0rXiv:2106.03609 2021.
James Hensman, Nicolo Fusi, and Neil D Lawrence. Gaussian processes for big data. In Proceedings of the 29th Conference on Uncertainty and Artificial Intelligence, 2013.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junetion tree variational autoencoder for molecular graph generation. In International Conference Machine Learning, pp. PMLR, 2018
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR, 2015.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. 2nd International
Conference Learning Representations, ICLR, 2014.
Matt J Kusner, Brooks Paige, and José Miguel Hrrmández-Lobateee Grammar variational autoencoder.
In International Conference on Machine Learning, pp. 1945-1954, PMLR, 2017.
Natalie Maus, Haydn T Jones, Juston S Moore, Matt J Kusner, John Bradshaw, and Jacob R Gardner. Local latent space bayesian optimization over structured inputs. Advances in Neural Information Processing Systems. 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, NeurIPS, 2019.
Siddharth Ramchandran, Gleb Tikhonov, Kalle Kujanpää, Miika Koskinen, and Harri Lähdesmäki. Longitudinal variational autoencoder. In International Conference on Artificial Intelligence and Statistics, pp. 3898-3906. PMLR, 2021.
Siddharth Ramchandran, Gleb Tikhonov, Otto Lönnroth, Pekka Tikkkainen, and Harri Lähdesmäki. Learning conditional variational autoencoders with missing covariates. Pattern Recognition, 147: 110113, 2024.
Michalis Titsias. Variational learning of inducing variables in sparse gaussian processes. In Artificial Intelligence and statistics, pp. 567-574. PMLR, 2009.
Austin Tripp. Erik Daxberger, and José Miguel Hernández-Lobato. Sample-efficient optimization in the latent space of deep generative models via weighted retraining, Advances in Neural Information Processing Systems, 33: 18:91199-11222 2020.
27


Reviews:
[{"review_id": "PceWRyN0Ca", "reviewer": "Reviewer_naas", "summary": "Bayesian Optimization is a class of methods for optimizing expensive black box functions $f:\\mathbb{R}^d \\to \\mathbb{R}$ typicality for low dimensional search spaces $d<10$.\n\nIn the high dimensional case (e.g. $d>100$) such as the space of images, if we have access to an unlabelled dataset of points in the search space, in this work denoted $y_1,...,y_N \\in \\mathbb{R}^d$, one may first train a variational autoencoder to map from the high dim search space to a low dim latent space $q:\\mathbb{R}^d \\to \\mathbb{R}^L$ and back $p:\\mathbb{R}^L \\to \\mathbb{R}^d$ where $L<<d$ (simplifying  notation somewhat). Then we simply use $z=q(y)$ and $y=p(z)$ as intermediate translation layers mapping between high and low dimensional spaces. BO is performed in the low dimensional space, modelling a dataset of points $\\\\{z_i, f(p(z_i))\\\\}$ with a GP and optimizing the acquisition function over $z\\in\\mathbb{R}^d$, meanwhile the objective function is evaluated in the high dimensional space $f(p(z))$.\n\nThis work considers the case where we have even more data available for some or all of the points in search space denoted $(x_i, y_i) \\in \\mathcal{X}\\times \\mathbb{R}^D$ with $\\mathcal{X} \\subset \\mathbb{R}^k$ where $k < 10$. In such a case, we have a regression dataset with low dim inputs and high dim outputs, we may use the GP-VAE architecture. As above with BO, we use the VAE function $z=q(y)$ to convert all the $y_i$ values to low dimensional and now we have a dataset $(x_i, z_i)$ which we can use for normal GP regression, mapping from $x$ to $z$.", "questions": "- is it possible to include Trip 2021 as a baseline with and without retraining?\n- what is the intuition that means training the GP-VAE would improve the outer BO modelling? Whiles integrating more data can helpat is the justification for the main hypothesis?\n- can the autrhors", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- __impactful problem__ high dimensional BO and VAE-BO are large problems, and the case where we have extra \"meta-data\" for points in the high dim search space seems perfectly reasonable and impactful (and surprising it hasn't been considered seriously until now)\n- __nice architecture__ the combination of GP-VAE and VAE-BO seems like an intuitive and good choice for such a problem.\n- __good benchmarks__ a toy example with MNIST images, mathematical expression tuning and molecule tuning, while the MNIST example  is rather artificial, I felt the molecule example really highlighted the benefit of incorporating covariates.\n- __accounting for missing data__ the authors also integrate previous approaches that handle missing data, although this is not novel in this work it is a nice to have and demonstrates broader practicality.", "weaknesses": "## Technical Comments\n- __Trip 2021 baseline__ is this baseline with weighted retraining or not? It weould be nuice to see a vanilla VAE-BO approach as well as the method of Tripp 2021 with weighted retraining.\n- __preference for high valued $y$__ The method of Trip et. al. 2020. starts with a VAE that is a generative model of the whole search space (on a high dim manifold) and after collecting a few fitness values, gradually retrain the VAE to become a generative model of high value parts of the search space, conceptually similar to CMA-ES or a trust region approach. In my view, this method has a nice intuition. In contrast with the above point, I may have misunderstood however it appears as though the proposed ELBO for GP-VAE with missing data does not have a bias for learning high valued $y$.\n- __learning without covariates possible failure mode__ when there are no extra covariates beside fitness values $c$, there are two GPs in the latent space,\n  - the first GP within the BO algorithm maps from latent points to fitness, modelling $\\hat{f}(z): \\mathbb{R}^L\\to\\mathbb{R}$,  \n  - the second GP maps from fitness back to latent $q(z|c):\\mathbb{R}\\to\\mathbb{R}^L$.\n\n  in a normal BO setting (e.g. VAE encoder and decoder are identity functions) I find this very counter-intuitive, the inverse GP must learn to map from a scalar value $c$ to all the points $\\\\{z|\\hat{f}(z)=c\\\\}\\subset\\mathbb{R}^L$, the level set of $c$ for multi-modal function, and this is being modelled by a single uni-modal Gaussian distribution. I have not seen this in the BO literature and it is not immediately obvious why such an approach would help. With _extensive_ retraining the latent space can be remoulded so that the level sets are clustered but this is speculative.\n\n- __limitations__ I may have missed this, but I there does not seem to be much discussion of failure cases and limitations, I have mentioned one above. As with any method that allows to incorporate more data/complexity also allows for more ways to break, if the $x_i$ values are pure noise or if all the optimal $y_i$ points happen to have dramatically different $x_i$ values. The paper does not seems to expose any failure modes or warning for users.\n\n## Minor Comments\n- __background__ the proposed method is an intelligent combination of prior methods, and much of section 4 (all of 4.1, 4.2 and parts of 4.3) are describing such prior works and may arguably belong in section 3.\n- __section 4.3__ I found this section to be a little bit dense and confusing, adding \"yet another\" distribution over $x$ (which conditions $z$ which conditions $y$). Although handling missing data is nice and shows practicality, moving this to the appendix as a \"bonus feature\" and using the space to provide more intuition and justification for the benefit of the main method might be better.\n\n[Tripp et.al. 2020](https://proceedings.neurips.cc/paper_files/paper/2020/file/81e3225c6ad49623167a4309eb4b2e75-Paper.pdf)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "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": "0L9aEwHBXq", "reviewer": "Reviewer_Ea83", "summary": "&nbsp;\n\nThe authors propose a novel VAE-based Bayesian optimization (VAE-BO) scheme that makes use of a Gaussian process prior VAE to leverage auxiliary covariates in learning a latent space that is better suited for Bayesian optimization. The method is evaluated against several baselines on synthetic data as well as the penalized logP molecule generation benchmark. While the framework is novel and interesting, I have some concerns about the empirical evaluation of the method as well as the reproducibility of the results. In relation to the empirical evaluation, the use of additional covariates that are correlated with the objective would appear to be biased if the baseline methods are not supplied with these covariates. Additionally, I have questions about the discrepancy between the LOL-BO optimization trace reported in the original paper [19] for penalized logP and the trace reported in the current paper. If these issues can be addressed in the rebuttal I will be happy to increase my score.\n\n&nbsp;", "questions": "&nbsp;\n\n1. My main question relates to how the additional covariates are used for each method. I would be very grateful if the authors could expand on this aspect of the empirical evaluation.\n\n2. For future work the authors may wish to consider the inversion problem [13] namely that under the mapping x -> z -> x' there is typically a reconstruction gap meaning that x is not equal to x'. Enforcing invertibility has been shown in some recent papers to improve VAE-BO performance systematically across architectures. This being said, the contribution is somewhat orthogonal to the contribution of the current work. I believe approaches such as deep metric learning as the authors have compared against are indeed the most appropriate baselines.\n\n&nbsp;", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "&nbsp;\n\n1. The work proposes a principled probabilistic framework for VAE-BO leveraging GP prior VAEs to enable conditioning on auxiliary covariates. The scheme is highly general and may be used in conjunction with many disparate VAE-BO architectures.\n\n2. The paper is exceptionally well-written and presented and the scholarship is excellent (an example of which is tracing back the ideas of label guidance to construct discriminative latent spaces to Urtasun et al. 2007).\n\n3. The empirical results are impressive pending the clarifications below as well as the release of the code.\n\n&nbsp;", "weaknesses": "&nbsp;\n\n**MAJOR POINTS**\n\n&nbsp;\n\n1. It would be great to perform diagnostic experiments on the GP surrogate fit on the latent space as in Grosnit et al. 2021 [16] in order to validate that improved BO performance is achieved due to a better GP fit on the latent space.\n\n2. On lines 468/469, the authors state, \"We augmented the ZINC-250K with five additional covariates: molecular weight, number of hydrogen donors, number of hydrogen acceptors, number of rotatable bonds, and total polar surface area.\". Are all baseline methods given access to these covariates? These additional descriptors are highly correlated with the water-octanol partition coefficient logP and so for fair comparison, I would expect that all methods be able to make use of these features in some fashion? It would be great if the authors could clarify exactly how these additional covariates are used for each method.\n\n3. The results reported in Figure 1 of the LOL-BO paper [19] are vastly at odds with the optimization trace reported for LOL-BO in Figure 5b) of the current paper. Why is this the case?\n\n4. The authors do not appear to have released the code for the submission and hence I have some concerns over the reproducibility of the results. This could be supplied as an anonymous GitHub link during the rebuttal phase.\n\n&nbsp;\n\n**MINOR POINTS**\n\n&nbsp;\n\n1. There are some missing capitalizations in the references e.g. \"Gaussian\" and \"Bayesian\".\n\n2. When introducing Bayesian optimization, it may be worth citing the originating papers for the method [1, 2] as discussed in [3].\n\n3. The statement that, \"Although BO offers an approach for black-box optimisation problems, it does not efficiently scale to high-dimensional data settings\" should probably be expanded on in light of recent work [4] which demonstrates that a vanilla GP surrogate where the lengthscale prior is scaled with the dimensionality of the problem can perform effective Bayesian optimization in 100s of dimensions.\n\n4. In terms of the references for applications of VAE-BO, Felton et al. 2020 use a multitask GP surrogate for Bayesian optimization over chemical reaction conditions (as opposed to chemical synthesis) and hence do not make use of a VAE-BO scheme. Additionally, Shields et al. use Bayesian optimization for chemical reaction conditions (as opposed to chemical synthesis) but do not use VAE-BO. Korovina et al. 2020 similarly do not use VAE-BO but rather define an optimal transport kernel directly over molecules. They do however consider chemical synthesis.\n\n5. In Figure 1, the task is articulated as discovering novel drug-like molecules. The penalized logP objective function, however, does not optimize for drug-likeness as noted in e.g. Section 5 of [5]. The penalized logP objective introduced in [6] is misspecified as a metric for drug-likeness since it attempts to maximize logP. For drug-like molecules the logP should however lie within the range of -0.4 to 5.6 according to the commonly-used heuristic Lipinski Rule of 5. As such, I would recommend rephrasing the task to molecule optimization or something comparably generic.\n\n6. In the related work, there has been limited empirical evidence that VAE-BO is beneficial in continuous high-dimensional spaces. In particular, techniques such as random embeddings [7] or SAASBO [8] or more well-known methods. VAE-BO however, has been demonstrated to help in applying Bayesian optimization over structured input spaces such as molecules, images, and biological sequences. As such, it may be worth rephrasing the discussion to focus slightly more on the structured nature of the input spaces as opposed to the dimensionality.\n\n7. The related work does a very good job at covering the majority of VAE-BO methods. Some works that should also be mentioned are [9-12]. Additionally, [13] does not yet appear to be formally published but would be worth mentioning once it is.\n\n8. The citation to Urtasun et al. 2007 shows great scholarship in tracing back the ideas underpinning VAE-BO. Additionally, it would be worth citing Jasper Snoek's PhD thesis [14] which also contained early ideas on label guidance to construct discriminative latent spaces.\n\n9. The decision to use the variable y to represent a high-dimensional observation is somewhat confusing. It may be better to use this variable to represent (noisy) observations of the objective function f.\n\n10. On line 163, it may be worth clarifying the training time complexity is O(N^3).\n\n11. It would be worth citing UMAP [15] given that it is used.\n\n12. Reference [17] should be cited when introducing the Quantitative Estimate of Drug-Likeness (QED) metric.\n\n13. In Figure 5a) it may be worth plotting the log regret to see a clearer distinction between the methods.\n\n14. It would be worth citing t-SNE [18] given that it is used.\n\n15. Reference [20] should be cited when introducing Expected Improvement (EI) as discussed in [3].\n\n16. When discussing Adam in Section C of the appendix it may be worth mentioning that it is an amalgam of the momentum and RMSProp optimizers.\n\n&nbsp;\n\n**REFERENCES**\n\n&nbsp;\n\n[1] Kushner, HJ., [A Versatile Stochastic Model of a Function of Unknown and Time\nVarying Form](https://www.sciencedirect.com/science/article/pii/0022247X62900112). Journal of Mathematical Analysis and Applications 5(1):150–167. 1962.\n\n[2] Kushner HJ., [A New Method of Locating the Maximum Point of an Arbitrary Multipeak Curve in the Presence of Noise](https://asmedigitalcollection.asme.org/fluidsengineering/article-abstract/86/1/97/392213/A-New-Method-of-Locating-the-Maximum-Point-of-an?redirectedFrom=fulltext). Journal of Basic Engineering 86(1):97–106. 1964.\n\n[3] Garnett, R., [Bayesian optimization](https://bayesoptbook.com/). Cambridge University Press. 2023.\n\n[4] Hvarfner, C., Hellsten, E.O. and Nardi, L. [Vanilla Bayesian Optimization Performs Great in High Dimensions](https://proceedings.mlr.press/v235/hvarfner24a.html). Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:20793-20817. 2024.\n\n[5] Griffiths, Ryan-Rhys, Philippe Schwaller, and Alpha A. Lee. [Dataset bias in the natural sciences: a case study in chemical reaction prediction and synthesis design.](https://arxiv.org/abs/2105.02637) arXiv preprint arXiv:2105.02637. 2021.\n\n[6] Gómez-Bombarelli, Rafael, Jennifer N. Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez-Lengeling, Dennis Sheberla, Jorge Aguilera-Iparraguirre, Timothy D. Hirzel, Ryan P. Adams, and Alán Aspuru-Guzik. [Automatic chemical design using a data-driven continuous representation of molecules.](https://pubs.acs.org/doi/full/10.1021/acscentsci.7b00572) ACS Central Science 4, no. 2, 268-276. 2018.\n\n[7] Wang, Ziyu, Frank Hutter, Masrour Zoghi, David Matheson, and Nando De Feitas. [Bayesian optimization in a billion dimensions via random embeddings.](https://www.jair.org/index.php/jair/article/view/10983) Journal of Artificial Intelligence Research 55, 361-387. 2016.\n\n[8] Eriksson, David, and Martin Jankowiak. [High-dimensional Bayesian optimization with sparse axis-aligned subspaces.](https://proceedings.mlr.press/v161/eriksson21a.html) In Uncertainty in Artificial Intelligence, pp. 493-503. PMLR, 2021.\n\n[9] Verma, Ekansh, Souradip Chakraborty, and Ryan-Rhys Griffiths. [High dimensional Bayesian optimization with invariance.](https://realworldml.github.io/files/cr/paper53.pdf) In ICML Workshop on Adaptive Experimental Design and Active Learning. 2022.\n\n[10] Notin, P., Hernández-Lobato, J.M. and Gal, Y., 2021. [Improving black-box optimization in VAE latent space using decoder uncertainty.](https://openreview.net/pdf?id=F7LYy9FnK2x) Advances in Neural Information Processing Systems, 34, pp.802-814.\n\n[11] Maus, N., Wu, K., Eriksson, D. and Gardner, J., 2023, [Discovering Many Diverse Solutions with Bayesian Optimization.](https://proceedings.mlr.press/v206/maus23a/maus23a.pdf) In International Conference on Artificial Intelligence and Statistics (pp. 1779-1798). PMLR.\n\n[12] Lee, Seunghun, Jaewon Chu, Sihyeon Kim, Juyeon Ko, and Hyunwoo J. Kim. [Advancing Bayesian optimization via learning correlated latent space.](https://proceedings.neurips.cc/paper_files/paper/2023/hash/98e967164ae2f6811b975d686dece3eb-Abstract-Conference.html) Advances in Neural Information Processing Systems 36. 2024.\n\n[13] Chu et al. [Inversion-Based Latent Bayesian Optimization](https://nips.cc/virtual/2024/poster/95013), NeurIPS 2024.\n\n[14] Snoek, Jasper. [Bayesian Optimization and Semiparametric Models with Applications to Assistive Technology.](https://library-archives.canada.ca/eng/services/services-libraries/theses/Pages/item.aspx?idNumber=1033018520) PhD diss., University of Toronto, 2013.\n\n[15] McInnes, Leland, John Healy, Nathaniel Saul, and Lukas Großberger. [UMAP: Uniform Manifold Approximation and Projection.](https://par.nsf.gov/servlets/purl/10104557) Journal of Open Source Software 3, no. 29. 2018.\n\n[16] Grosnit, A., Tutunov, R., Maraval, A.M., Griffiths, R.R., Cowen-Rivers, A.I., Yang, L., Zhu, L., Lyu, W., Chen, Z., Wang, J. and Peters, J., [High-dimensional Bayesian optimisation with variational autoencoders and deep metric learning.](https://arxiv.org/abs/2106.03609) arXiv preprint arXiv:2106.03609. 2021.\n\n[17] Bickerton, G. Richard, Gaia V. Paolini, Jérémy Besnard, Sorel Muresan, and Andrew L. Hopkins. [Quantifying the chemical beauty of drugs.](https://www.nature.com/articles/nchem.1243) Nature Chemistry 4, no. 2 (2012): 90-98.\n\n[18] Van der Maaten, Laurens, and Geoffrey Hinton. [Visualizing data using t-SNE.](https://www.jmlr.org/papers/volume9/vandermaaten08a/vandermaaten08a.pdf?fbcl) Journal of Machine Learning Research 9, no. 11. 2008.\n\n[19] Maus, Natalie, Haydn Jones, Juston Moore, Matt J. Kusner, John Bradshaw, and Jacob Gardner. [Local latent space bayesian optimization over structured inputs.](https://proceedings.neurips.cc/paper_files/paper/2022/hash/ded98d28f82342a39f371c013dfb3058-Abstract-Conference.html) Advances in Neural Information Processing Systems 35, 34505-34518. 2022.\n\n[20] Saltines, VR., One Method of Multiextremum Optimization. Avtomatika i Vychislitel’naya Tekhnika (Automatic Control and Computer Sciences) 5(3):33–38. 1971.\n\n&nbsp;", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "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": "iIFA46AKSk", "reviewer": "Reviewer_t5E1", "summary": "This paper introduces a novel approach to Bayesian optimization (BO) that addresses high-dimensional data incorporating auxiliary covariates, even when some of the covariates contain missing values.  The method builds upon a variational autoencoder (VAE) where the latents follow a Gaussian Process (GP) prior. The approach is technically sound, and the experimental results show its effectiveness in both synthetic and real-world settings.\n\nHowever, the paper's novelty is somewhat limited, as it builds on existing GP prior VAE models with missing covariates. Furthermore, the absence of a key baseline in one of the experiments raises concerns about the completeness of the experimental evaluation. Addressing these issues—especially by adding the missing baseline—would make a stronger case for the proposed method.\n\nIf the authors include this baseline in the first experiment and their method shows significant improvements (or they provide a justified reason for its absence), I would be inclined to increase my evaluation score from a weak reject to a weak accept.", "questions": "See weaknesses.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "**1. Clarity and Readability:** The paper is well-structured and clearly written. The authors explain complex ideas in an accessible way, ensuring that the technical details are easy to follow. I am not an expert for BO but could get a good understanding of the paper within a few hours. \n\n**2. Technical Rigor:** The paper is technically solid. The authors carefully describe the underlaying probabilistic model step-by-step. The integration of auxiliary covariates, even when some values are missing, is handling in a principled manner by applying variational inference. \n\n**3. Experimental Thoroughness:** The experiments are in general carefully executed. This was the part of the paper that I enjoyed reading the most. They span multiple datasets from different domains and the results are studied in-depth.", "weaknesses": "**1. Limited Novelty:** While the application of the GP prior VAE to Bayesian optimization is novel, the underlying model itself—GP prior VAE with partial observations —has been published previously in Ramchandran et al. (2024). This implies that the core contribution of the paper lies primarily in applying the method to the BO context. This limits the overall novelty of the work.\n\n**2. Competitive/Missing Baselines:** The method LOL-BO (Maus et al., 2022) shows competitive performance on the expression reconstruction dataset, and is only marginally outperformed on the molecular optimization experiment. However, in the first experiment on synthetic data, I found this method missing. It is important for the authors to include this baseline in this experiment to ensure a fair comparison. If their approach significantly outperforms this baseline, it would strengthen the contribution. If there is a specific reason why this approach cannot be applied for this experiments, it needs to be stated more clearly in the paper.\n\n**3. Scalability Dependency:** The scalability of the method relies on leveraging sparse Gaussian Process (GP) techniques, using inducing points to approximate the GP. However, selecting meaningful inducing points becomes challenging when the dataset changes over time, as it is the case in BO applications. As new data arrives, it is unclear how to update the inducing points in a principled and computationally efficient manner. This is a critical issue for ensuring that the model scales well over BO iterations, and the paper should at least discuss this problem.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "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": {}}]

Author rebuttal:

Decision: Accept (Poster)