Local Latent Space Bayesian Optimization over Structured Inputs
Abstract
Bayesian optimization over the latent spaces of deep autoencoder models (DAEs) has recently emerged as a promising new approach for optimizing challenging black-box functions over structured, discrete, hard-to-enumerate search spaces (e.g., molecules).
Here the DAE dramatically simplifies the search space by mapping inputs into a continuous latent space where familiar Bayesian optimization tools can be more readily applied.
Despite this simplification, the latent space typically remains high-dimensional.
Thus, even with a well-suited latent space, these approaches do not necessarily provide a complete solution, but may rather shift the structured optimization problem to a high-dimensional one.
In this paper, we propose LOL-BO , which adapts the notion of trust regions explored in recent work on high-dimensional Bayesian optimization to the structured setting.
By reformulating the encoder to function as both an encoder for the DAE globally and as a deep kernel for the surrogate model within a trust region, we better align the notion of local optimization in the latent space with local optimization in the input space.
LOL-BO achieves as much as 20 times improvement over state-of-the-art latent space Bayesian optimization methods across six real-world benchmarks, demonstrating that improvement in optimization strategies is as important as developing better DAE models.
Introduction
Many challenges across the physical sciences and engineering require the optimization of structured objects .
For example, in the drug design setting, it is common to seek a molecule that maximizes some property, e.g., binding affinity for a target protein [17], or similarity to a known drug molecule under some metric [3].
These structured optimization problems are particularly challenging, as the objective function is now typically over a large combinatorial space of discrete objects.
These challenging optimization problems have motivated a large amount of very recent interest in methods that perform Bayesian optimization (BO) over latent spaces .
Using drug design as a running example, a deep auto-encoder (DAE) is first trained on a large set of unsupervised molecules.
Bayesian optimization is then run as normal in the continuous latent space of the autoencoder to produce candidate latent codes.
These candidate latent codes are then decoded back to molecules, and the objective function is evaluated.
Bayesian optimization then proceeds as normal, with the surrogate model trained on a dataset of latent codes and their resulting objective values after decoding.
Rapid progress has been made in this emerging topic.
In large part, this has been to improve the
DAE model.
The early work of Kusner et al.
[28] and Jin et al.
[21] investigated domain specific autoencoder architectures that much more reliably produced valid molecules.
The latest work of Tripp et al.
[43] and Grosnit et al.
[14] investigate methods to inject supervised signal from objective function evaluations into the latent space, enabling better surrogate modeling.
However, the optimization component of latent space Bayesian optimization remains under-explored: indeed, it is common practice to directly apply standard Bayesian optimization with expected improvement once a latent space has been constructed (all of the above cited methods do this).
We conjecture that, despite the ability of a deep encoder to dramatically simplify a structured input space, optimization remains challenging due to the high-dimensionality of the extracted latent space.
In this paper, we propose to jointly co-design the latent space model and a local Bayesian optimization strategy.
Concretely, we make the following contributions:
We develop local latent Bayesian optimization, LOL-BO , for the structured optimization setting, addressing the inherent mismatch between the notion of a trust region in the latent space and a trust region in the original structured input space.
We derive a semisupervised DAE model that better satisfies the implicit assumptions of trust region methods by performing joint variational inference over the DAE and sparse GP surrogate model (this model is similar in spirit to Kingma et al.
[25], Eissman et al.
We propose a deep autoencoder architecture using transformer layers over the recently proposed SELFIES representation for molecules [27].
We find that the SELFIES string representation provides significant improvements to optimization performance.
We show the potential of co-adapting latent space modeling with local Bayesian optimization in an empirical evaluation of LOL-BO against top-performing baselines, achieving up to 22 × improvements in objective value on commonly used benchmark tasks over prior latent space BO approaches.
Background
In black-box optimization, we seek to find the minimizer of a function, arg min x ∈X f ( x ) .
Commonly, f ( x ) is assumed to be expensive to evaluate and unknown (i.e., a “black box”).
It is therefore desirable to use fewer function evaluations to achieve a given objective value.
Most classically, this problem has been considered in the setting where the search space X is continuous (e.g., X ⊆ R ).
However, many applications across the natural sciences and engineering require optimizing over discrete and structured input spaces.
De novo molecule design using machine learning has received particular attention, where X is a search space of small organic molecules and f ( x ) is some desirable property to optimize [31].
Bayesian optimization [37, 30, 41] is a framework for solving these black-box optimization problems, most commonly when X is a continuous, real-valued search space.
A surrogate model is trained on a set of inputs and corresponding objective values D ℓ =
[ x i , y i ] ni =1 .
This surrogate model is used to suggest candidate inputs ˆ x to query next.
The candidates are then labeled, added to the dataset, and the process repeats.
Gaussian process models f ( x ) ∼ GP ( m ( x ) , k
( x , x ′ )) are commonly used as surrogate models, as access to calibrated predictive (co-)variances enables search strategies that trade off exploration and exploitation.
Deep autoencoder models.
A deep autoencoder model consists of an encoder Φ : X → Z that maps from a structured and discrete space (e.g., the space of molecules)
X to a continuous, real-valued latent space Z , and a decoder that does the reverse Γ( · ) : Z → X .
A variational autoencoder [24] uses a probabilistic form of these two functions, a variational posterior Φ( Z | X ) and a likelihood Γ( X | Z ) .
The parameters θ Φ , θ Γ of Φ and Γ are trained by maximizing the evidence lower bound (ELBO):
[Equation]
Commonly, the (amortized) variational posterior Φ( Z | X ) and the prior p ( Z ) are chosen to be
Gaussian, which leads to a convenient analytic expression for the KL divergence.
Note that, with respect to specific optimization tasks we might want to perform over the structured space X , the
ELBO as defined above is unsupervised.
Latent space Bayesian optimization.
Latent space
Bayesian optimization seeks to leverage the representation learning capabilities of deep autoencoders models (DAEs) to reduce a discrete, structured search space X to a continuous one Z [43, 5, 12, 14, 6, 22].
Because the ELBO above is unsupervised with respect to an optimization task, a large set of unlabeled inputs D u ⊆ X is first used to train the DAE.
A Gaussian process prior is again placed over the objective function, but here as a function of the latent codes , f ( z ) ∼ GP ( m ( z ) , k ( z , z ′ )) .
To train the GP initially, we collect a small set of n inputs D ℓ ⊆ X , n ≪ |D u | with known labels y ℓ .
The inputs are passed through the encoder to obtain latent representations Z =
[ z i ] ni =1 , and the surrogate model is fit to the dataset [( z i , y i )]
ni =1 .
Standard BO strategies may then be applied using this surrogate model to select a candidate latent vector ˆ z , which is then decoded back to an input ˆ x and evaluated.
Methods
Recent latent space Bayesian optimization (LS-BO) approaches have focused on improving the architecture of the VAE used to compress the search space [28, 21], and on reorganizing the subsequent latent space so that it is more amenable to BO [43, 14, 6].
However, it is commonplace to apply standard BO with expected improvement once the latent space has been constructed, despite the fact that this latent space often remains high-dimensional [21], on the order of tens or hundreds of dimensions.
In this paper, we explore whether lessons learned in the high-dimensional BO setting can be adapted to improve optimization in high-dimensional latent spaces.
3.1 Local Bayesian optimization – Fixed latent space
One approach to adapting high-dimensional BO strategies is to apply them in the extracted latent space Z .
Because Z is a synthetic latent space, it’s unlikely that the objective function decomposes additively over this space without explicitly encouraging this [23, 9, 47].
Likewise, Z is already a compression of the much larger X , so methods that exploit low-dimensional substructure may not be effective [48, 32, 7, 11].
Local Bayesian optimization is a strategy that avoids over-exploration in these search spaces without simplifying assumptions about the objective.
Commonly this is done by restricting the search in each iteration to a small region.
For example, Eriksson et al.
[8] proposes TuRBO-M, an optimization algorithm which maintains M local optimization runs, each of which is limited to search within its own hyper-rectangular trust region . Each trust region is centered at the current best input found by a given run.
Since Eriksson et al.
[8] found that TuRBO-1 is competitive or even better than TuRBO-m where m > 1 on most tasks, in this paper we focus only on TuRBO-1.
Following Eriksson et al.
[8], we will henceforth use TuRBO as short-hand for TuRBO-1.
Adapting trust regions to LS-BO is straightforward.
Suppose we have a VAE with encoder Φ( · ) and decoder Γ( · ) trained on an unsupervised dataset D u of structured inputs, and a supervised set
Z of latent codes corresponding to labeled inputs D ℓ and their objective values y .
Denote by ˆ z the incumbent latent code corresponding to the input ˆ x with best observed objective value ˆ y so far.
The trust region is an axis aligned hyper-rectangular subset of the latent space T ˆ z ⊆ Z , centered at the incumbent latent vector ˆ z .
It has a base side length L , with the actual side length for each dimension rescaled according to its lengthscale λ i in the GP’s ARD covariance function, L i = λ
[Equation] dj [Equation]
Candidate latent vectors to evaluate next are restricted to be selected from T ˆ z .
Following Nelder and
Mead [33], Eriksson et al.
[8], the base trust region length is halved after τ fail iterations where no progress is made and doubled after τ succ iterations in a row where progress is made.
When a new incumbent ˆ z ′ is found with ˆ y ′ < ˆ y , the trust region is recentered at ˆ z ′ .
3.2 Local Bayesian optimization – Adaptive latent space
By defining “locality” in terms of a small rectangular region centered at the incumbent, trust region methods implicitly assume smoothness, that points close to the center have highly correlated objective values with it.
In the standard BO setting with GP surrogates, this assumption matches that made by the GP, as many covariance functions measure covariance as a function of distance.
However, the latent space Z of a DAE is only trained so that a latent vector z =
Φ( x ) reconstructs the original input when passed through the decoder, Γ( z ) ≈ x .
Therefore, it is possible that small changes to the current incumbent latent vector ˆ z may result in large changes to the corresponding input ˆ x , violating the assumptions of both the GP surrogate model and the trust region method.
Adapting the latent space to the GP prior.
We propose to solve this mismatch by performing inference jointly over both the Gaussian process surrogate defined over the latent space p ( f | Z ) and the Bayesian VAE model p ( X , Z ) =
p ( X | Z ) p ( Z ) , while assuming in the prior conditional independence between f and X given Z :
[Equation]
By performing inference over f and Z jointly in the above model end-to-end, we encourage the latent space to organize in a way that matches the assumptions of the GP prior p ( f | Z ) .
Relation to prior work.
Learning the VAE and supervised model jointly in a semisupervised fashion is related to prior work on semisupervised VAE modelling as first proposed in Kingma et al.
This was adapted for BO in Eissman et al.
[6], who use an auxiliary supervised neural network to inject supervision before fitting the GP.
Here we consider joint variational inference over the VAE and a sparse GP, directly encouraging the latent space to satisfy the assumptions of the GP prior.
In the global BO setting, Siivola et al.
[40] finds semisupervised adaptation of VAEs to not be critical; however, in the local setting, we experimentally find this results in significant improvement due to the assumptions underlying trust region methods.
Sparse GP augmentation.
Following Hensman et al.
[15], the above prior p ( f | Z ) can be augmented with a global set of learned inducing variables u representing latent function values at a set of inducing locations V , p ( f | Z ) → p ( f | u , Z , V ) p ( u | V ) .
This results in the following full joint density over the latent function f , structured inputs X , latent representation Z , objective values y , and inducing variables u (ommitting dependencies on inducing locations V for compactness):
[Equation]
ELBO derivation.
Our goal is to derive an evidence lower bound (ELBO) for the above model to perform variational inference over both the inducing variables u and VAE latent variables Z .
The log evidence p ( y , X ) for this model is log p ( y , X )
= log (cid:82) p ( y , X , f , u , Z ) d f d u d Z .
To aid in lower bounding the log evidence, we introduce a variational distribution q ( u , Z )
=
q ( u )Φ( Z | X ) .
This is a factored joint variational distribution over the inducing values u as in [15, 42] and over the latent space of the encoder Z .
[Equation]
Because the joint distribution p ( y , X , f , u , Z ) factors as in Equation [Equation] we can further rewrite the integral of the joint distribution over f as:
[Equation]
Plugging this expected value form of the joint distribution integral into Equation 2 and repeatedly applying Jensen’s inequality, we derive an evidence lower bound
[Equation]
The ELBO in practice.
Rephrasing Equation 3 in terms of the standard SVGP (sparse variational GP; 15) and VAE ELBOs, we see that:
[Equation]
The ELBO can be estimated as follows.
Select a minibatch of inputs (e.g., molecules) X , some of which have known labels y , and some of which are unlabeled.
Pass the minibatch X through the encoder Φ and generate latent codes Z .
For latent codes Z ℓ corresponding to labeled inputs, compute
L SVGP on the supervised dataset ( Z ℓ , y ℓ ) .
For all latent codes Z , pass through the decoder and add the standard VAE loss.
Because the variational posteriors Φ( Z | X ) and q ( u ) support reparameterization, computing derivatives of the ELBO with respect to θ Φ , θ Γ , θ GP is straightforward.
Because the encoder parameters θ Φ are updated through the SVGP loss, Φ can be viewed as simultaneously acting as an encoder for the VAE and as a (stochastic) deep kernel [52, 51] for the GP model.
Recentering Under this joint model, local BO can still be applied using the same setup and hyperparameters discussed in subsection 3.1.
However, when using the above loss, the VAE–and therefore the latent space–is now modified after each iteration of BO.
Therefore, the center of the trust region ˆ z and latent codes of inputs within or near the trust region may change.
To remedy this, we pass observed data X back through the updated encoder to determine the new location of the corresponding latent points Z before selecting candidates, and update the trust region accordingly.
Computational considerations.
Updating the full encoder Φ( · ) can be significantly more expensive than training the surrogate model only.
A practical extension to the above idealized process is to only train in a joint fashion periodically.
Thus, the variational GP parameters θ GP are updated in each iteration of optimization, but not necessarily θ Φ , θ Γ .
We propose to tie updates to θ Φ , θ Γ to the sequential successes and failures already tracked by the trust region method.
If the optimization procedure is making rapid progress, updating θ Φ , θ Γ may be an unnecessary expense.
Therefore, we update the full set of parameters θ Φ , θ Γ , θ GP after τ retrain successive failures.
Experimentally we use τ retrain ≪ τ fail as this allows for multiple joint model updates for each time a trust region shrinks.
Related Work
Bayesian optimization for molecule discovery can be performed without DAEs over fixed lists of molecules [50, 16, 13].
However, the need to pre-define the molecules that can be queried upfront puts a limit to the chemical space such methods can consider: the largest virtual libraries are order of ≈ 10 10 in size [45], which is tiny compared to the space of all possible compounds [26].
Latent-space Bayesian optimization [12, 6, 43, 14, 40, 21] ties the optimization algorithm to a DAE so that, in theory, it can generate any molecule.
We broadly group many recent advancements in this approach into two categories: (a) those that investigate new training losses to improve surrogate learning of the objective [14, 43, 6, 5] and (b) those that develop new decoder architectures, for instance improving the validity of the generated molecules by building in grammars or working explicitly in the graph domain [28, 21, 38, 22, 4].
Decoder architectures have also been designed to apply these ideas to different discrete domains such as neural network architectures [55] or Gaussian process kernels [29].
Components of our method including the joint training procedure described in subsection 3.2 are similar to Eissman et al.
[6] which considers injecting semisupervision into the VAE via an auxiliary neural network and Grosnit et al.
[14] which does so using a triplet loss.
In particular, the triplet loss of Grosnit et al.
[14] likely encourages similar latent space organization to the method described in subsection 3.2 due to the use of distances in both the triplet loss and the GP covariance function.
However, these works use off-the-shelf global Bayesian optimization algorithms, such as expected improvement, as the latent optimizer.
Alternative optimization algorithms have also been proposed for molecule discovery, such as those based on genetic algorithms [20, 34], Markov chain Monte Carlo [53], or ideas from reinforcement learning (RL) [36, 54, 56].
Segler et al.
[39] use the cross entropy method to guide the generation of complete SMILES strings from an autoregressive recurrent neural network.
You et al.
[54] train a policy to build molecules up from atoms and bonds using an RL algorithm.
While such methods are often able to find molecules with better property scores than early LS-BO approaches, these methods frequently neglect sample efficiency, impeding their use on practical problems.
Experiments
We apply LOL-BO to six high-dimensional, structured BO tasks over molecules [18, 19] and arithmetic expressions.
We additionally investigate the impact of various components of our method in subsection 5.4 and subsection 5.4.
log P and Expressions (subsection 5.1) are tasks that are typically evaluated in low budget settings.
The Guacamol benchmarks (subsection 5.2) afford larger function evaluation budgets, and the protein docking task (subsection 5.3) is both low budget and the oracle is expensive.
Error bars in all plots show standard error over runs.
We implement LOL-BO leveraging
BoTorch [2] and GPyTorch [10] 1 .
Code and model weights are available at https://anonymized/ .
Datasets All methods have access to the same amount of supervised and unsupervised data for each task.
To initialize all molecular optimization tasks, we generate labels for a random subset of 10 , 000 molecules from the standardized unlabeled dataset of 1.27M molecules from the Guacamol benchmark software [3].
For the expressions task, we use the same labeled dataset of 40 ,expressions from [14] to initialize all methods.
Furthermore, wherever different VAEs are used for a given task, we use the same unsupervised dataset to train the VAEs so that the amount of unsupervised data also remains consistent.
In particular, all molecular VAEs are pre-trained on the same aforementioned dataset of 1.27M molecules from Brown et al.
Deep autoencoder models.
A variety of VAE models have been used for the benchmarks tasks we consider below.
For Expressions , we follow recent work [28, 14, 43, 35] in using a GrammarVAE model [28].
For the remaining de novo molecule design tasks, we consider two models.
Junction Tree VAE (JTVAE).
Molecules are commonly represented initially using the Simplified
Molecular-Input Line-Entry System (SMILES) [49] string representation.
SMILES representations are well-known to be brittle, in that single character mutations can easily result in strings that no longer encode a valid molecule.
The JTVAE model Jin et al.
[21] addresses the problem of generating valid strings by using a molecular graph structure to build each molecule out of valid sub-components so that the generated molecules are almost always valid.
As a result, the JTVAE has been widely used as a VAE model for latent space optimization due to its much higher rate of decoding to valid molecules [14, 43, 5, 35].
We follow prior work, using a latent dimensionality of 56.
SELFIES VAE
The SELFIES string representation for molecules is a recently proposed alternative to the SMILES string representation [27].
Notably, nearly every permutation of SELFIES tokens encodes a valid molecule.
As a result, purely sequence based models using the SELFIES representation are capable of generating valid molecules at a rate competitive with the JT-VAE while avoiding expensive computation involving graph structures.
In our molecule design experiments below, we therefore evaluate a VAE architecture based on the
SELFIES string representation in addition to the commonly used JTVAE.
We train a VAE model with 6 transformer encoder and transformer decoder layers and a latent space dimensionality of 256 [46].
We use our SELFIES VAE for all molecular optimization tasks below.
We also use the JTVAE model to enable direct comparison to prior work.
Hyperparameters.
For the trust region dynamics, all hyperparameters including the initial base and minimum trust region lengths L init , L min , and success and failure thresholds τ succ , τ fail are set to the
TuRBO defaults as used in Eriksson et al.
Our method introduces only one new hyperparameter, τ retrain , which we set to 10 in all experiments.
5.1 log P and Expressions
In this section, we consider two tasks commonly used in the latent space BO literature–optimizing the penalised water-octanol partition coefficient (log P) over molecules, and generating single variable arithmetic expressions (e.g., x × sin ( x × x ) )
Because of the small evaluation budget typically considered for these benchmarks, we use a batch size of 1 .
Baselines.
For these tasks, we compare to standard latent space BO, as well as a variety of recent prior work in this area.
We compare to LS-BO with weighted retraining ( W-LBO ) as described in
Tripp et al.
[43] and LS-BO with triplet loss retraining ( T-LBO ) from Grosnit et al.
A number of techniques outside of the BO literature consider this task as well.
These papers report final objective values ≥ 30 , but do not consider sample efficiency.
For example, GEGL [1], which we do compare to in subsection 5.2, uses more than 500 evaluations per iteration.
See Ahn et al.
[1] for an evaluation on log P.
Results.
In Figure 1, we plot the cumulative best objective value for the penalized log P and expressions tasks, averaged over three runs.
For log P, LOL-BO using the SELFIES transformer
VAE achieves an average log P score of 545.17 after 500 iterations.
Using a JTVAE, which is more directly comparable to prior work, we still achieve an average score of 123.54–higher than previous state-of-the-art scores using latent space BO of 26.11 [14].
For the expressions task, LOL-BO achieves a final regret of 0.079 on average, versus 0.194 achieved by T-LBO .
Molecule quality.
While molecules found using LOL-BO for the log P task are “valid” according to software commonly used to compute these scores, the molecules produced by this search clearly abandon any notion of reality.
A number of methods have been proposed to evaluate molecule quality [3], and indeed even some recent work in LS-BO has explored encouraging higher quality solutions [35].
However, our primary goal in this paper is to develop stronger optimization routines in this space, and indeed the baselines we consider also treat log P as an unconstrained optimization problem.
Therefore, this result should be taken primarily as evidence that the log P objective can be exploited by a sufficiently strong optimizer, rather than as evidence of novel interesting molecules.
5.2 Guacamol Benchmark Tasks
The Guacamol 2 benchmark suite [3] contains scoring oracles for a variety of molecule design tasks, with scores ranging between 0 and 1.
We select three tasks among the most challenging in terms of the best objective value achieved: Perindopril MPO , Ranolazine MPO , and Zaleplon MPO .
The goal of each of these tasks is to design a molecule with high fingerprint similarity to a known target drug but that differs in a specified target way.
The task definitions are discussed in Brown et al.
Baselines.
We compare to T-LBO , the most competitive Bayesian optimization method overall from subsection 5.1.
Additionally, a variety of techniques from beyond the BO literature using reinforcement learning and genetic programming have been evaluated on these tasks.
We compare to
GEGL and GraphGA, with GraphGA [20] among the top methods on the Guacamol leaderboard.
We also report the best score achieved over the ≈ 1.27 million molecules in the full Guacamol dataset, which serves as a performance cap for virtual screening (see section 4).
Results.
The results of optimization on these three benchmarks are shown in Figure 2.
We again report results for LOL-BO with both a JTVAE and SELFIES VAE deep autoencoder.
LOL-BO with a
SELFIES VAE optimizes significantly faster than other methods in the evaluation budget used here.
5.3 DRD3 Receptor Docking Affinity
The Therapeutics Data Commons (TDC) benchmark suite [17] contains oracles that evaluate the binding propensity of a ligand (small molecule) to a target protein.
These docking oracles incur nontrivial computation cost, requiring on the order of minutes to evaluate, making them ideal [Table] benchmark tasks for Bayesian optimization.
In this paper, we focus on designing ligands that bind to dopamine receptor D3 (DRD3), as the TDC website maintains a leaderboard of results for
DRD3 3 .
We compare LOL-BO on the DRD3 docking task to the TDC leaderboard and three BO methods– T-LBO , standard LS-BO, and fixed latent space TuRBO on this task.
Results of optimization are in Table 1 (right).
LOL-BO achieves the best performance among the BO baselines, with only TuRBO making significant progress on the task by 5000 evaluations.
Compared to the baselines from the TDC leaderboard, LOL-BO is significantly more sample efficient, achieving results comparable to the best method with a factor of 10 × fewer oracle calls.
5.4 Ablation studies
In this section, we evaluate the components of LOL-BO described in subsection 3.2 beyond the direct application of TuRBO in the latent space described in subsection 3.1.
We run TuRBO on the three
Guacamol tasks considered in subsection 5.2 using a fixed latent space extracted by the pretrained
SELFIES VAE and JTVAE.
To run TuRBO, we use the same experimental setup as for LOL-BO , but Φ( · ) and Γ( · ) are fixed at the start of optimization.
We compare LS-BO, TuRBO, and LOL-BO on all three tasks in Figure 3.
The use of SVGP is controlled across methods.
LOL-BO outperforms TuRBO across all DAE models and tasks.
While these results demonstrate the impact of joint training, the excellent performance of TuRBO in a static, pre-trained latent space highlights the value of high dimensional Bayesopt in this setting.
To further analyze whether the jointly trained latent space better-matches the distance based assumptions of the GP model, we include an additional analysis directly comparing scores of molecules sampled from a local region around the top-scoring molecule in the latent space of the pretrained
SELFIES-VAE before optimization has begun and at the end of optimization (see subsection A.3).
Since LOL-BO achieves better performance with the SELFIES VAE than with the JTVAE, we evaluate a series of metrics for both DAE models.
In Table 1 left, we compare reconstruction error and validity, as well as the standard deviation in Perindopril MPO score from decoding a fixed latent vector ztimes, averaged over 3500 distinct z .
We report supervised learning performance for a GP trained on 10 , 000 molecules taken from the initial Guacamol dataset, using the pretrained latent spaces.
Across standard metrics, the SELFIES VAE achieves excellent performance.
Beyond these, the improvement in objective variance for a single z is notable: because Guacamol objective values are in the range [0 , 1] , a standard deviation of 0 .
1 represents significant additional noise.
Thus, the DAE can be a significant source of noise even for an otherwise deterministic objective function.
Conclusion
Designing latent spaces for Bayesian optimization is challenging because the VAE and the optimizer often compete: too small of a latent space renders the VAE incapable of capturing the structure present in the input space, while too large a latent space makes optimization challenging.
The empirical results we present here demonstrate the importance of viewing these latent spaces as high-dimensional and tackling optimization with high-dimensional BO strategies, even though these latent spaces are often significantly smaller than the input space. Furthermore, we show that high-dimensional BO methods can be significantly improved when translating them to the latent space setting, demonstrating the potential of further research at the intersection of these areas.  Acknowledgements JRG and NM were supported by NSF award IIS-2145644.
JM and HJ were supported by the
Laboratory Directed Research and Development program of Los Alamos National Laboratory (LANL) under project number 20210043DR Checklist For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contribu- tions and scope? [Yes] (b) Did you describe the limitations of your work?
[Yes] (see subsection A.7) (c) Did you discuss any potential negative societal impacts of your work?
[Yes] (see subsection A.7) (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (See section 5 where we will provide a link to a github repository with our code) (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (See section 5) (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (See section 5) (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] (See subsection A.4) 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (See section 5) (b) Did you mention the license of the assets?
[Yes] (See subsection A.5)
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] (See section 5 where we will provide a link to a github repository with our code) (d) Did you discuss whether and how consent was obtained from people whose data you’re us- ing/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable informa- tion or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Appendix A.1 Additional implementation details
As the data collected by LOL-BO grows, it becomes increasingly computationally expensive to update the models in and end-to-end fashion on all of the data.
Thus, instead of updating on all data, we only update end-to-end on a subset of the data consisting of the most recently collected batch of data points, plus the top- k scoring points from the remaining data.
For the experiments in subsection 5.1 and subsection 5.3 where we have a small evaluation budget, we use k = 10 .
For the experiments in subsection 5.2 with a larger total evaluation budget, we use k = 1000 .
A.2 JTVAE vs. SELFIES VAE additional analysis
In Figure 4, we provide histograms to visualize a few of the statistics summarized in Table 1 in the main text 1.
A.3 Additional ablation analysis
In this section, we analyze whether the jointly trained latent space matches the distance based assumptions of the GP model.
On the Ranolazine MPO task, we first encode the best molecule x ∗ (score=0.9382) from a single run of LOL-BO to get the latent codes z ∗ pre and z ∗ joint from both the pretrained VAE before optimization has begun and at the end of optimization.
We then analyze the average score of molecules decoded from latent vectors sampled uniformly from an L -hypercube centered at both z ∗ joint and z ∗ pre .
By L = 5 , latent codes in the joint
VAE near z ∗ joint still achieve an average score of 0.9106, while for vectors near z ∗ pre in the pretrained model the average score has dropped to only 0.3382 Figure 5.
In this section, we additionally analyze whether the jointly trained VAE is better able to reconstruct the highest- scoring molecules found by the optimizer.
Empirically, the optimizer consistently produces higher-scoring molecules than any found in the entire initial unsupervised dataset used to train the initial VAE (see Best in Dataset comparison in Figure 2).
This raises the question of whether these high-scoring molecules found by the optimizer may often be out-of-sample for the initial VAE.
To investigate this, we took the top 128 highest-scoring molecules found by LOL-BO on the Ranolazine MPO task and computed the reconstruction accuracy of the initial
VAE on these molecules.
We repeated this experiment with the final VAE obtained at the end of the optimization run.
We found that the average reconstruction accuracy of the initial VAE was 0 .
01 while the average accuracy of the final VAE was 0 .
The extremely low reconstruction accuracy of the initial VAE indicates that many of the high-scoring molecules found by the optimizer are indeed out-of-sample.
It would therefore be very difficult to find these high-scoring molecules by simply searching in the latent space of the initial VAE.
In contrast, finding these high-scoring molecules by searching in the latent space of the final VAE would likely be much easier.
Therefore, this result provides further motivation for our procedure of jointly updating the VAE during optimization, so that as the optimization proceeds, the VAE is made to adapt so that it can better construct the types of moleucles we are interested in, even if they are unlike any molecules in the initial unsupervised dataset.
L
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
A.4 Compute
The SELFIES VAE was trained using a single server with eight Quadro RTX 8000 GPUs.
All optimization experiments were run on a single RTX A5000 GPU.
We used an internal cluster for compute.
A.5 Software licensing
All major software packages used (ie BoTorch, GPyTorch, Guacamol) are released under MIT license.
A.6 Run-time considerations
LOL-BO scales to large evaluation settings readily via minibatch training of the ELBO in Equation 4 due to the variational GP approximation used (Hensman et al., 2013).
TuRBO scales similarly by using the same variational
GP approximation.
Although we’ve shown that LOL-BO significantly improves performance over TuRBO (see Figure 5), it is important to consider the additional time complexity of LOL-BO .
In particular, updating the VAE jointly can add significant time complexity, particularly with very deep transformers.
However, this cost is generally manageable since relatively little VAE updating is required with the small amount of data acquired at each step (compared to the initial training on the large unsupervised dataset).
When the evaluation budget is small (Section 5.1), or the objective function expensive (Section 5.3), LOL-BO incurs as little as a 1.01 × slowdown over running TuRBO .
However, larger evaluation budgets require retraining the encoder on growing amounts of data.
Thus, in Section 5.2, we do incur as much as a 5 × slowdown in the worst case.
However, the total running time in Section 5.2 is still typically on the order of half a day, and we believe the improved outer-loop optimization performance makes the additional time complexity worthwhile.
Finally, the added cost of jointly updating the VAE is negligible compared to the cost of certain oracle calls typical in molecule design.
For instance, in Section 5.3 the docking oracle takes minutes to evaluate per molecule, leading to run-times of up to a few days, and even this would be considered cheap compared to wet-lab experiments.
To see specific wall-clock times for all methods, see 2.
The wall clock times given for W-LBO [43] were obtained with our modifications to parallelize the W-LBO code-base.
A.7 Potential limitations and future work
This work focuses on molecular optimization tasks centered around drug development.
While the development of new drugs has the potential to help many people, it is also important to consider the dual use nature of such methods.
In particular, AI technologies for drug discovery have the potential to be misused for the de novo design of biochemical weapons [44].
It is also important to consider the potential for negative societal impact if any new drugs discovered are not property tested before being administered.
While we hope that our work can eventually speed-up the drug development process by providing useful candidate molecules, it is essential that
[Table]
experts remain heavily involved and that FDA guidelines for drug development and approval are strictly adhered to.
Additionally
[Table]
the oracles and tasks we use here
[Table]
even ones with the potential to design candidate drug molecules like protein docking, are completely oblivious to critical aspects like human safety and efficacy, stability, manufacturing cost
[Table]
and so on.
These computational approaches to molecule design should therefore be considered first pass screening procedures, not procedures to design final end products.
In order to optimize the Guacamol molecular optimization tasks
[Table]
LOL-BO requires tens of thousands of function evaluations Figure 2.
While the fairly cheap Guacamol oracles have made it reasonable to do this many evaluations, there are many molecule design applications where each function evaluation is much more expensive.
In particular, tens of thousands of function evaluations is likely not efficient enough for many applications that would be run in vitro with actual wet-lab experiments.
Another potential limitation of LOL-BO is that is is designed to find a single best optima for a given task.
However, in the molecular space in particular, it is often desirable to find multiple good optima since a single solution might fail for unknown reasons downstream.
In future work, we plan to explore methods that can be used to extend Bayesian optimization to obtain a set of diverse solutions rather than a single solution.