Source-linked AI summary

Parallel and Distributed Thompson Sampling for Large-scale Accelerated Exploration of Chemical Space

José Miguel Hernández-Lobato, James Requeima, Edward O. Pyzer-Knapp, Alán Aspuru-Guzik

arXiv:1706.01825v1stat.ML

TL;DR

The paper addresses the inability of existing Bayesian optimization methods to scale to massive molecular libraries and large batches of parallel measurements. It proposes parallel and distributed Thompson sampling, showing comparable performance to parallel expected improvement for small batches and better performance than scalable baselines when parallel EI is infeasible. The results support PDTS as a scalable approach for large-scale parallel Bayesian optimization.

  • Problem

    Existing Bayesian optimization methods do not scale to the massive molecular libraries and large numbers of parallel measurements used in high-throughput screening.

  • Method

    PDTS implements Thompson sampling in a distributed parallel Bayesian optimization procedure, generating batch evaluation locations independently and in parallel.

  • Results

    PDTS performs as well as parallel EI for small batches and outperforms greedy, ϵ-greedy, random, and Monte Carlo scalable baselines in large-scale screening applications.

  • Takeaways & Limitations

    PDTS provides a highly scalable method for parallel Bayesian optimization in large-scale chemical-space searches where competing approaches are limited by scalability.

  • Takeaways & Limitations

    Parallel EI is limited by sequential batch selection and high computational costs, including posterior updates that can be expensive with many training points.

Abstract

from arXiv · show

Chemical space is so large that brute force searches for new interesting molecules are infeasible. High-throughput virtual screening via computer cluster simulations can speed up the discovery process by collecting very large amounts of data in parallel, e.g., up to hundreds or thousands of parallel measurements. Bayesian optimization (BO) can produce additional acceleration by sequentially identifying the most useful simulations or experiments to be performed next. However, current BO methods cannot scale to the large numbers of parallel measurements and the massive libraries of molecules currently used in high-throughput screening. Here, we propose a scalable solution based on a parallel and distributed implementation of Thompson sampling (PDTS). We show that, in small scale problems, PDTS performs similarly as parallel expected improvement (EI), a batch version of the most widely used BO heuristic. Additionally, in settings where parallel EI does not scale, PDTS outperforms other scalable baselines such as a greedy search, $ε$-greedy approaches and a random search method. These results show that PDTS is a successful solution for large-scale parallel BO.

1. Introduction

Chemical space contains vast numbers of possible molecules, making trial-and-error discovery infeasible and motivating high-throughput screening and scalable Bayesian optimization. The paper introduces PDTS to address the limited scalability of existing parallel BO methods.

  • Chemical space is estimated to contain over 10^60 molecules, while fewer than 100 million compounds are available in public repositories or databases.
  • Trial-and-error discovery is often slow, tedious, and infeasible because relevant compounds must be identified among many unimportant possibilities.
  • High-throughput virtual screening analyzes large molecular libraries with many parallel computer-cluster simulations to reduce them to promising leads for expensive experiments.
  • Existing Bayesian optimization methods cannot scale to the parallel measurements and massive candidate libraries used in high-throughput screening.
  • PDTS is proposed as a distributed Thompson-sampling solution for scalable parallel Bayesian optimization, alongside an illustration of Thompson sampling and PDTS.

2. BO and Thompson Sampling

Bayesian optimization guides expensive molecule evaluations using predictive models and acquisition functions, while Thompson sampling provides a posterior-sampling alternative. Parallel expected improvement supports batches but becomes computationally costly and sequential at large scale, motivating a distributed approach.

  • Bayesian optimization: Bayesian optimization uses past evaluations to model expensive objectives and selects new molecules by maximizing an acquisition function.The acquisition function quantifies the benefit of evaluating each candidate.
  • Bayesian optimization: Expected improvement balances exploitation of high predicted values with exploration under predictive uncertainty.Its utility is U(yj|xj, DI) = max(0, yj − y⋆), where y⋆ is the best observed value.
  • Thompson sampling: Thompson sampling draws model parameters from the posterior and selects the candidate with the highest sampled expected objective value.Its Monte Carlo approximation induces exploration through posterior sampling despite using a utility equal to yj.
  • Parallel Bayesian optimization: Parallel EI selects batch locations sequentially because each proposal depends on earlier pending evaluations.It uses fantasized outcomes and repeatedly updates predictive quantities while constructing the batch.
  • Parallel Bayesian optimization: Parallel EI lacks scalability for thousands of simultaneous measurements because posterior updates and acquisition optimization become expensive over large datasets and molecule libraries.Its sequential construction also prevents full use of multiple processors.
  • Motivation for PDTS: The paper introduces a fully parallel and distributed batch Bayesian optimization method based on Thompson sampling.The method is designed to exploit multiple processors in a computer cluster.

3. Parallel and Distributed Thompson Sampling

PDTS approximates parallel Bayesian optimization by running multiple Thompson sampling executions from one shared posterior without updating it between selections. This makes batch selection distributed, while duplicate molecule choices require centralized coordination.

  • PDTS construction: PDTS uses Thompson sampling with one Monte Carlo sample for pending outcomes and one posterior draw to approximate parallel acquisition.Synthetic pending outcomes preserve the initial posterior on average.
  • PDTS construction: A batch of size S in PDTS is equivalent to running sequential Thompson sampling S times without updating the current posterior.Each execution can be assigned to a separate computer node.
  • Distributed execution: PDTS distributes the current posterior and evaluated-index set to worker nodes, which independently select locations and later evaluate the objective.The central node coordinates batch formation while workers perform selections and evaluations.
  • Discrete-library coordination: Because the candidate library is discrete, simultaneous PDTS executions may select the same molecule, requiring a central list of currently selected molecules.This coordination addresses duplicate selections during parallel batch construction.

4. Related Work

Related parallel Bayesian optimization methods use fantasized observations, joint acquisition procedures, matching objectives, confidence bounds, entropy reduction, or knowledge gradients. PDTS is distinguished by avoiding extra marginalization cost and accommodating broad model classes.

  • Fantasized observations: Fantasized-data methods augment observations for pending experiments, using either fixed values such as constant liar or predictive means such as Kriging believer.These approaches perform Bayesian optimization on the augmented dataset.
  • Parallel EI: Parallel EI averages acquisition functions over predictive samples for pending evaluations, while later methods recommend points jointly.The cited extensions modify the parallel EI procedure rather than using independent Thompson samples.
  • Batch construction: Simulated matching constructs batches intended to match samples recommended by a sequential Bayesian optimization policy.A batch is considered good when it likely contains a point near the sequential policy’s best sample value.
  • Alternative heuristics: Other approaches extend UCB, predictive entropy search, or knowledge gradient to select parallel samples using uncertainty, entropy reduction, or expected solution quality.Examples include GP-BUCP, GP-UCB-PE, PPES, and parallel knowledge gradient.
  • PDTS distinction: PDTS adds no extra computational cost for approximate marginalization and can use exact or approximate posterior samples from Gaussian processes or Bayesian neural networks.These properties make PDTS highly parallelizable across model classes.

5. Bayesian Neural Networks for High-throughput Screening

The paper uses Bayesian neural networks with probabilistic back-propagation to provide scalable uncertainty estimates for Bayesian optimization in chemical-space searches. PBP approximates otherwise intractable posteriors and supports Thompson sampling through posterior weight samples.

  • Bayesian neural networks: Bayesian neural networks provide probabilistic predictions suited to exploring chemical spaces with multiple local optima.Their uncertainty estimates support efficient exploration, while stochastic optimization enables application to large datasets.
  • Implementation choice: PBP is preferred because it is fast and avoids tuning hyperparameters such as learning rates or regularization constants.
  • Model specification: The network models evaluations as noisy neural-network outputs with weights W and noise precision γ, using rectifier activations in hidden layers.The reported experiments assume normalized targets, although the broader methods can also handle additive Gaussian noise.
  • Posterior approximation: PBP replaces the generally intractable posterior over W, γ, and λ with a tractable approximation q.Its parameters are tuned by iterating assumed density filtering over the training data.
  • Posterior approximation: PBP updates q by matching moments after incorporating each likelihood factor, using Gaussian approximations to network outputs during sampled-weight forward passes.The forward pass approximates non-Gaussian neuron outputs with Gaussians having matching means and variances.
  • Thompson sampling: After training, predictions for a new feature vector are obtained by applying the same approximate forward-pass process to f(x⋆; W).For Thompson sampling, the method samples W from q and optimizes the acquisition function using predictions from the resulting deterministic network.

6. Experiments with GPs and Parallel EI

The experiments compare PDTS with parallel EI using Gaussian-process models on benchmark and GP-sampled objectives. Across these tests, PDTS and parallel EI generally obtain very similar results, with differences depending on the objective landscape.

  • Experimental setup: PDTS is compared with parallel EI using a Gaussian-process model because parallel EI is a leading parallel Bayesian-optimization approach.The experiments specifically note that parallel EI cannot scale to the large batch sizes used in high-throughput screening.
  • Experimental setup: The simulated cluster has 11 nodes, with one controlling the optimization process and 10 nodes performing parallel evaluations.Because evaluation costs are assumed to dominate computation, the experiments update the GP in blocks of 10 data points.
  • Benchmarks: The test suite includes the two-dimensional Bohachevsky and Branin-Hoo functions, the six-dimensional Hartmann function, and GP-prior samples on a two-dimensional unit square.Performance is evaluated using immediate regret after each objective evaluation.
  • Evaluation: Across 50 repetitions, Figure 3 reports mean and standard errors for the logarithm of the best immediate regret versus the number of function evaluations.TS and EI update after each sample, whereas PDTS and parallel EI update after every 10 samples.
  • Results: PDTS and parallel EI obtain very similar results overall, although parallel EI performs better on Hartmann while PDTS performs better initially on GP-prior samples.The paper attributes the Hartmann difference to PDTS being more explorative and parallel EI more exploitative.

7. Experiments with Molecule Data Sets

The experiments evaluate PDTS on molecular libraries using fingerprint-based inputs, recall of high-value molecules, and comparisons with scalable sampling baselines. PDTS outperforms random sampling and generally improves on greedy exploitation, while reducing the sampling or compute required for discovery.

  • Molecules are represented by 512-bit Morgan circular fingerprints with bond radius 2, derived from canonical SMILES using RDkit.
  • The evaluation measures 1% recall, with CEP additionally reporting recall for molecules whose PCE exceeds 10%.Each sampling step selects a batch of previously unsampled molecules.
  • PDTS significantly outperforms Monte Carlo sampling and performs better than greedy sampling on the molecule data sets.Greedy sampling ranks molecules by the mean predictive distribution and performs no exploration.
  • In CEP, greedy sampling initially finds better molecules, but PDTS later overtakes as the greedy strategy's promising region becomes exhausted.
  • 20 times higher recall than Monte Carlo is achieved by PDTS on CEP, with estimated virtual-screening use falling from 30,000 to 1,500 CPU years.
  • Approximately 6,000 sampled molecules locate about 70% of the top 1% in both One-dose and Malaria, versus 14,000 required by Monte Carlo.

8. Conclusions

The paper presents PDTS as a highly scalable approach to parallel Bayesian optimization. It performs favorably against parallel EI for small batches and outperforms scalable alternatives in large-scale chemical-space searches.

  • PDTS is a highly scalable parallel and distributed implementation of Thompson sampling for Bayesian optimization.
  • PDTS compares favorably with parallel EI for problems with small batch sizes.
  • PDTS outperforms greedy search, ϵ-greedy approaches, and random search in large-scale chemical-space applications.
Loading 1706.01825v1…