PaperlayerView exact revision on arXiv

Source-linked AI summary

Local Latent Space Bayesian Optimization over Structured Inputs

Natalie Maus, Haydn T. Jones, Juston S. Moore, Matt J. Kusner, John Bradshaw, Jacob R. Gardner

arXiv:2201.11872v2cs.LGstat.ML

TL;DR

Structured discrete search spaces make black-box optimization challenging, while latent-space BO can leave a high-dimensional optimization problem. LOL-BO adapts trust-region methods to structured inputs and jointly co-designs the latent model and local optimizer. It achieves up to 22× improvement over prior latent-space BO approaches on commonly used benchmark tasks.

  • Problem

    Structured optimization involves large combinatorial spaces, and latent-space BO can remain challenging because the latent space is high-dimensional.

    Source: §1 · p. 2

  • Method

    LOL-BO jointly co-designs a latent-space model and local Bayesian optimization strategy to address the mismatch between latent and structured-input trust regions.

    Source: §1 · p. 2

  • Results

    Up to 22× improvements in objective value over prior latent-space BO approaches were achieved on commonly used benchmark tasks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support treating latent spaces as high-dimensional optimization problems and adapting high-dimensional BO methods for latent-space settings.

    Source: §6 · p. 10

    For log P, the molecules found abandon any notion of reality, so the result is evidence of optimizer performance rather than novel interesting molecules.

    Source: §5.1 · p. 7

Abstract

from arXiv · show

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.

1 Introduction

Structured optimization over large discrete spaces is challenging, and latent-space Bayesian optimization may remain difficult because its latent space is high-dimensional. LOL-BO jointly co-designs the latent-space model and local Bayesian optimization strategy.

1 Introduction

  • Optimization in latent-space BO remains under-explored, and high dimensionality may keep optimization challenging despite the encoder's simplification of structured inputs.Source: §1 · p. 2
  • LOL-BO addresses the mismatch between trust regions in latent space and trust regions in the original structured input space.Source: §1 · p. 2
  • The proposed semisupervised DAE jointly performs variational inference over the DAE and sparse GP surrogate to better satisfy trust-region assumptions.Source: §1 · p. 2
  • The proposed molecular DAE uses transformer layers over SELFIES, whose string representation improves optimization performance.Source: §1 · p. 2
  • Up to 22× improvements in objective value over prior latent-space BO approaches were achieved on commonly used benchmark tasks.The evaluation compared LOL-BO against top-performing baselines.Source: §1 · p. 2

2 Background

Black-box optimization seeks good objective values with fewer evaluations, including in structured discrete spaces. The background introduces Bayesian optimization and latent-space BO, in which autoencoders map structured inputs to continuous codes for optimization.

2 Background

  • Black-box optimization seeks to minimize an unknown, expensive-to-evaluate function, making fewer evaluations desirable for a given objective value.Source: §2 · p. 2
  • Structured optimization includes discrete search spaces, such as small organic molecules, where the objective is a property to optimize.Source: §2 · p. 2
  • Bayesian optimization fits a surrogate to inputs and objective values, then uses it to suggest candidates that are evaluated and added to the data.Source: §2 · pp. 2–3
  • A deep autoencoder maps structured discrete inputs to continuous latent codes with an encoder and maps codes back with a decoder.Source: §2 · p. 3
  • A VAE trains its encoder and decoder by maximizing an ELBO, while the task-independent ELBO described here is unsupervised with respect to optimization tasks.Source: §2 · p. 3
  • Latent-space BO trains a DAE on unlabeled inputs, fits a GP to labeled inputs' latent codes and objective values, then decodes and evaluates selected codes.Source: §2 · p. 3

3 Methods

The methods adapt trust-region Bayesian optimization to latent spaces, then jointly adapt the encoder and GP so latent-space locality better matches the assumptions of local optimization. They also specify how the trust region is managed as the latent space changes.

3.1 Local Bayesian optimization – Fixed latent space

  • Fixed latent-space trust regions may not match structured-input locality because the latent space need not have additive structure or effective low-dimensional substructure.These are reasons the authors give for adapting high-dimensional BO strategies rather than assuming standard methods transfer directly.Source: §3.1 · p. 3
  • TuRBO searches an axis-aligned box centered on the incumbent, scales side lengths by GP ARD lengthscales, and adjusts its base length after sequential successes or failures.The box is recentered when a better incumbent is found.Source: §3.1 · pp. 3–4

3.2 Local Bayesian optimization – Adaptive latent space

  • The encoder is updated jointly with the sparse GP so the latent representation better matches the GP prior’s assumptions about local smoothness.The authors frame this as addressing the possibility that small latent-space changes can produce large changes in structured inputs.Source: §3.2 · p. 4
  • The joint model performs inference over the GP function and VAE latent variables, with conditional independence between the objective function and inputs given the latent representation.The model augments the GP with inducing variables and derives an ELBO for variational inference over inducing and latent variables.Source: §3.2 · p. 4
  • For ELBO estimation, minibatch inputs pass through the encoder; labeled codes contribute the sparse GP loss, while all codes contribute the VAE loss.The encoder is updated through the sparse GP loss, so it acts both as the VAE encoder and as a stochastic deep kernel for the GP.Source: §3.2 · p. 5
  • Because joint training changes the latent space, observed inputs are re-encoded after updates and the trust region is adjusted before selecting candidates.The same local BO setup and hyperparameters are retained.Source: §3.2 · p. 5
  • To reduce the cost of full encoder updates, the method can update GP parameters each iteration and retrain the full model after successive trust-region failures.The experiments use a retraining threshold much smaller than the trust-region shrinkage threshold.Source: §3.2 · p. 5

4 Related Work

Prior work includes molecule optimization over fixed libraries, latent-space BO methods that improve representations or decoders, and alternative search algorithms. The paper distinguishes its approach by combining local BO with joint VAE–GP training.

4 Related Work

  • Fixed-list molecule optimization is limited by the need to specify queryable molecules upfront; the largest virtual libraries are about 10^10 compounds, far smaller than all possible compounds.The passage contrasts fixed libraries with latent-space BO, which is described as able, in theory, to generate any molecule.Source: §4 · p. 5
  • Recent latent-space BO work improves surrogate learning with new training losses or improves molecule validity with specialized decoder architectures.The cited decoder designs include grammars and graph-domain approaches, with extensions to other discrete domains also noted.Source: §4 · p. 5
  • Related semisupervised VAE methods inject supervision through an auxiliary network or triplet loss, whereas this paper jointly performs variational inference over the VAE and sparse GP.Those prior approaches use off-the-shelf global BO optimizers, while the present work uses local BO.Source: §4 · pp. 5–6
  • Genetic algorithms, MCMC, and reinforcement-learning methods can find better molecule scores than early latent-space BO, but often neglect sample efficiency.The passage also mentions cross-entropy search over generated SMILES strings.Source: §4 · p. 6

5 Experiments

The experiments compare LOL-BO with latent-space BO baselines across molecule and expression tasks, test two molecular autoencoders, and examine local optimization and joint latent-space training. Results cover low-budget tasks, larger-budget Guacamol benchmarks, expensive protein docking, and ablations.

5.1 log P and Expressions

  • 545.17 average log P after 500 iterations with SELFIES; JTVAE reaches 123.54 versus the prior latent-space BO score of 26.11 [14], while expression regret is 0.079 versus T-LBO’s 0.194.The log P molecules are valid under the scoring software but abandon any notion of reality, so the result demonstrates optimization of the objective rather than discovery of interesting molecules.Source: Figure 1 · p. 6, §5.1 · p. 7

5 Experiments

  • Experiments use matched supervised and unsupervised data across methods, with molecular VAEs pretrained on 1.27M molecules and expression methods initialized from 40,000 labeled expressions.Molecular experiments use JTVAE for direct comparison with prior work and a six-layer transformer SELFIES VAE for the molecular optimization tasks.Source: §5 · pp. 6–7

5.2 Guacamol Benchmark Tasks

  • LOL-BO with the SELFIES VAE optimizes significantly faster than other methods on the three Guacamol benchmarks within the evaluation budget used.The selected tasks are Perindopril MPO, Ranolazine MPO, and Zaleplon MPO.Source: Figure 2 · p. 8, §5.2 · pp. 8–9

5.3 DRD3 Receptor Docking Affinity

  • LOL-BO achieves the best performance among BO baselines on DRD3 docking, while matching the best TDC leaderboard method with 10× fewer oracle calls.Only TuRBO makes significant progress on the task by 5000 evaluations.Source: Table 1 · p. 9, §5.3 · p. 9

5.4 Ablation studies

  • LOL-BO outperforms TuRBO across DAE models and tasks; strong fixed-latent TuRBO results also highlight the value of high-dimensional BO.The ablation compares LS-BO, TuRBO, and LOL-BO on three Guacamol tasks while controlling the use of SVGP.Source: §5.4 · p. 9
  • The SELFIES VAE has excellent standard-metric performance, but decoding a fixed latent vector can add substantial objective noise even for deterministic oracles.On Guacamol’s [0, 1] objective scale, a standard deviation of 0.1 represents significant additional noise.Source: §5.4 · pp. 9–10

6 Conclusion

The conclusion argues that latent spaces should be treated as high-dimensional optimization spaces and that high-dimensional BO strategies can be adapted to improve them.

6 Conclusion

  • Although latent spaces are smaller than their input spaces, the results support using high-dimensional BO strategies to optimize them.The paper also identifies potential for further research at the intersection of high-dimensional BO and latent-space optimization.Source: §6 · p. 10

Checklist

The checklist records affirmative responses for reporting experimental error bars and compute resources, citing existing assets, and describing work limitations and societal impacts.

Checklist

  • The checklist says the paper describes its limitations and potential negative societal impacts in subsection A.7.The checklist marks human-subject and participant-risk items as not applicable.Source: Checklist · p. 14
  • The checklist reports that experimental error bars and compute resources are documented in the paper or supplement.It points to Section 5 for error bars and subsection A.4 for compute details.Source: Checklist · p. 14
  • The checklist indicates that existing assets are cited and their licenses are discussed.It also reports that new assets are included with a repository link.Source: Checklist · p. 14

A.1 Additional implementation details

To limit the cost of end-to-end updates as data grows, LOL-BO updates on the latest batch plus top-scoring points from older data.

A.1 Additional implementation details

  • End-to-end updates use the latest batch plus top-k older points, with k = 10 for low-budget tasks and k = 1000 for larger-budget Guacamol experiments.The paper describes this subset strategy as a response to increasingly expensive updates on all collected data.Source: §A.1 · p. 15

A.2 JTVAE vs. SELFIES VAE additional analysis

Figure 4 visualizes variation in Perindopril MPO scores among molecules decoded from one latent code and the time needed to decode a batch.

A.2 JTVAE vs. SELFIES VAE additional analysis

  • Figure 4 visualizes selected statistics summarized in Table 1 in the main text.Source: §A.2 · p. 15, Figure 4 · p. 15
  • The left histogram shows Perindopril MPO objective-value variance across 20 molecules decoded from a single latent code.The figure visualizes statistics summarized in Table 1.Source: Figure 4 · p. 15, §A.2 · p. 15
  • The right panel reports the time taken to decode a batch of points.Source: Figure 4 · p. 15

A.3 Additional ablation analysis

The analyses compare locality and reconstruction in the jointly trained and pretrained VAEs. Nearby codes in the joint VAE retain higher Ranolazine MPO scores, and the final VAE reconstructs optimizer-found molecules much more accurately.

A.3 Additional ablation analysis

  • At L = 5, nearby joint-VAE codes average 0.9106 Ranolazine MPO, versus 0.3382 for nearby pretrained-VAE codes.The comparison samples latent vectors from hypercubes centered on the respective codes for the top-scoring molecule.Source: §A.3 · p. 15, Figure 5 · p. 16
  • For the top 128 optimizer-found Ranolazine molecules, reconstruction accuracy averages 0.89 with the final VAE and 0.01 with the initial VAE.Source: §A.3 · p. 15
  • The initial VAE’s low reconstruction accuracy indicates many optimizer-found high-scoring molecules are out-of-sample, motivating joint VAE updates during optimization.The final VAE can adapt to construct molecules unlike those in the initial unsupervised dataset.Source: §A.3 · p. 15

A.4 Compute

The appendix describes compute resources, runtime trade-offs, and boundaries of the molecular optimization setting. Joint VAE updates add cost, while some oracle and task constraints limit practical use.

A.4 Compute

  • The SELFIES VAE used one server with eight Quadro RTX 8000 GPUs; optimization experiments used one RTX A5000 GPU.Compute was provided by an internal cluster.Source: §A.4 · p. 16

A.6 Run-time considerations

  • Joint VAE updates add time, especially with very deep transformers, but relatively little updating is needed per step compared with initial training.LOL-BO and TuRBO both use variational GP approximations for minibatch ELBO training.Source: §A.6 · p. 16
  • With small evaluation budgets or expensive objectives, LOL-BO can incur as little as a 1.01× slowdown over TuRBO.For larger budgets, retraining the encoder on growing data can produce a worst-case 5× slowdown; Section 5.2 runtimes are typically about half a day.Source: §A.6 · p. 16

A.7 Potential limitations and future work

  • The molecular oracles omit safety, efficacy, stability, and manufacturing cost, so the methods are first-pass screening rather than final-product design.Experts should remain involved and FDA drug-development and approval guidelines should be followed.Source: §A.7 · p. 17
  • LOL-BO uses tens of thousands of evaluations on Guacamol tasks, likely too many for many in-vitro applications.Source: §A.7 · p. 17
  • LOL-BO is designed to find one best optimum, although molecular design may call for multiple good solutions when a single candidate fails downstream.The authors identify diverse solutions as future work.Source: §A.7 · p. 17

Related papers

Loading 2201.11872v2…