Source-linked AI summary
Sample-Efficient Optimization in the Latent Space of Deep Generative Models via Weighted Retraining
Austin Tripp, Erik Daxberger, José Miguel Hernández-Lobato
TL;DR
The paper addresses sample-inefficient optimization of expensive black-box objectives over high-dimensional, structured spaces. It performs optimization in a generative model’s latent space and uses weighted retraining to steer that space toward promising regions, improving efficiency and performance across evaluated tasks.
Problem
Expensive black-box optimization over high-dimensional, structured spaces requires evaluating the objective as few times as possible.
Method
The method combines latent space optimization with objective-weighted data and periodic generative-model retraining to refine the feasible latent region.
Results
Weighted retraining significantly benefits latent space optimization across application domains and generative models, improving over state-of-the-art methods on a chemical design benchmark.
Takeaways & Limitations
Weighted retraining is a conceptually simple extension that can improve the efficiency and performance of existing latent space optimization methods.
Takeaways & Limitations
When training data are mostly low-scoring, the initial generative model can leave insufficient feasible-region space for novel high-scoring points.
Abstract
from arXiv · showhide
Many important problems in science and engineering, such as drug design, involve optimizing an expensive black-box objective function over a complex, high-dimensional, and structured input space. Although machine learning techniques have shown promise in solving such problems, existing approaches substantially lack sample efficiency. We introduce an improved method for efficient black-box optimization, which performs the optimization in the low-dimensional, continuous latent manifold learned by a deep generative model. In contrast to previous approaches, we actively steer the generative model to maintain a latent manifold that is highly useful for efficiently optimizing the objective. We achieve this by periodically retraining the generative model on the data points queried along the optimization trajectory, as well as weighting those data points according to their objective function value. This weighted retraining can be easily implemented on top of existing methods, and is empirically shown to significantly improve their efficiency and performance on synthetic and real-world optimization problems.
1 Introduction
Many scientific and engineering tasks require optimizing expensive objectives over complex, high-dimensional inputs. LSO addresses this by optimizing in a learned low-dimensional latent space, while weighted retraining aligns that space with the optimization trajectory.
- Drug design exemplifies optimization over structured inputs with expensive wet-lab or simulation-based evaluations.
- LSO maps a low-dimensional continuous latent space onto the input manifold and optimizes a surrogate objective there.
- Standard LSO decouples generative-model training from downstream optimization, limiting its alignment with the target objective.
- The proposed method combines objective-based data weighting with periodic generative-model retraining to address two LSO decoupling problems.
- Across application domains and generative models, weighted retraining improves LSO efficiency and performance, including on a chemical-design benchmark.
2 Problem Statement and Background
The paper studies sample-efficient black-box optimization when inputs are high-dimensional or structured and evaluations are expensive. LSO makes this tractable by replacing direct optimization in input space with surrogate optimization in a continuous latent space.
- The target problem is optimizing an expensive, derivative-free objective over high-dimensional, structured input spaces using as few evaluations as possible.
- Model-based optimization approximates the objective with a surrogate, but high dimensionality and structured discreteness make optimization difficult.
- LSO uses a generative model g: Z 7→ X and optimizes a latent objective model h: Z 7→ R over a low-dimensional continuous space.
- An approximate inverse q maps observed inputs into latent points so the latent objective model can be trained from evaluated data.
3 Failure Modes of Latent Space Optimization
LSO can fail because the generative model represents the training distribution rather than regions useful for the objective, and it does not incorporate information from newly queried points. These failures can exclude promising solutions from the feasible latent region.
- The feasible latent region concentrates probability where the generative model was trained, even when the latent space itself is unbounded.
- Because training data are often low-scoring, most of the feasible region may contain sub-optimal points and leave little space for novel high-scoring solutions.
- Even unlimited objective evaluations may not find a substantially better novel point if that point is absent from the feasible region.
- The learned feasible region can mismatch the region useful for optimization, preventing fixed-model LSO from approaching a distant global optimum.
- Standard LSO also fails to propagate newly acquired information into the generative model, preventing the feasible region from shifting toward promising areas.
4 Latent Space Optimization with Weighted Retraining
Weighted retraining steers the generative model toward high-scoring data and periodically updates it with newly queried points. Rank-based weighting controls how strongly scores influence training, while retraining propagates optimization discoveries into the latent manifold.
- Weighting the data distribution: Discarding low-scoring data is inadequate; weighted training instead assigns positive weights that favor high-scoring points without removing the rest.
- Weighting the data distribution: Weighted generative-model training replaces the empirical loss mean with a weighted empirical mean, implementable through sampling points with probabilities proportional to their weights.
- Weighting the data distribution: The rank-based weighting function is positive, resistant to outliers, and uses k to control weighting strength from uniform weighting at k = ∞ to concentration on the top point at k = 0.
- Periodic retraining: Periodic retraining updates the generative model during optimization, either by training from scratch or fine-tuning the previous model on novel data.
- Combined procedure: Algorithm 1 alternates weighted model training with latent surrogate optimization, querying, evaluation, and dataset augmentation.
- Periodic retraining: Retraining only on new data risks catastrophic forgetting, whereas weighting can give a small number of high-scoring new points disproportionate influence.
- Combined procedure: The combined method makes the generative model an active participant that keeps the latent manifold occupied by updated, optimization-relevant points.
- Practical considerations: Weighted retraining has minimal weighting overhead, and fine-tuning can reduce retraining cost, although training may remain prohibitive in some applications.
5 Related Work
The paper situates weighted retraining among methods for optimizing expensive objectives in high-dimensional or structured spaces. It contrasts latent-space optimization and related generative-model methods with Bayesian optimization and reinforcement learning.
- Latent Space Optimization: Latent space optimization maps low-dimensional latent variables through a generative model, enabling optimization over structured inputs.
- Weighted Retraining: Weighted retraining relates the cross-entropy method, reward-weighted regression, feedback GAN, and DbAS/CbAS through higher weights for high-scoring points.
- Bayesian Optimization: Bayesian optimization models objective values probabilistically but usually targets continuous, low-dimensional spaces; ChemBO extends it to molecules with domain-specific knowledge.
- Reinforcement Learning: Reinforcement learning applies to structured optimization, including chemical design, but is generally extremely sample inefficient when evaluations are costly.
- Conditional Generative Models: Conditional generative models produce points conditioned on target properties, but their sample efficiency remains unclear.
6 Empirical Evaluation
The evaluation tests how weighting and retraining affect latent spaces and optimization across three tasks, then compares the method with existing approaches. Results support complementary benefits from weighting and retraining and stronger performance than the compared baselines.
- Evaluation Questions: The experiments address weighted training’s effect on latent spaces, parameter effects on optimization, and comparison with existing methods.
- Tasks and Metrics: Experiments cover 2D shape area maximization, arithmetic expression fitting, and chemical design using distinct data and model types.
- Tasks and Metrics: TopK score measures the Kth-best novel evaluated point versus objective-function evaluations, with averages and standard deviations across five random seeds.
- Latent-Space Effects: Weighted training is tested with rank weighting, where k = ∞ denotes uniform weighting, to assess whether high-scoring points become more prevalent in the latent space.
- Parameter Effects: Both weighting and retraining help individually, while their combination performs best across cases and often improves performance suddenly after retraining.
- Comparison with Other Methods: The proposed method significantly outperforms the compared baselines in sample efficiency and final performance.
- Comparison with Other Methods: 27.84 was achieved with 500 samples, compared with the best previously reported machine-learning score of 11.84 using approximately 5,000 samples.
7 Discussion and Conclusion
The paper presents weighted retraining as a simple addition to latent space optimization for black-box optimization in high-dimensional, structured spaces. It reports substantial efficiency and performance gains while identifying practical limitations for future work.
- Conclusion: Weighted retraining combines latent space optimization with dataset weighting and periodic generative-model retraining.
- Conclusion: The method significantly boosts efficiency and performance on challenging real-world optimization problems and is easy to implement on previous methods.
- Limitations: Performance depends on training a latent objective model that optimizes effectively in latent space.
- Limitations: The method requires a large labelled dataset to train a deep generative model, limiting suitability for some problems.
Broader Impact
The paper discusses potential benefits for faster discovery and broader machine-learning applications, while acknowledging possible misuse and the preliminary state of the technology.
- Potential Benefits: The authors suggest this research could enable faster discovery of medicines, energy materials, and device designs.
- Potential Risks: The technology could also support discovery of chemical or biological weapons, although substantial resources and infrastructure remain necessary.
- Potential Applications: The approach may influence conditional image and text generation, where human judgments can serve as expensive objective functions.
A.1 More Information on Rank-Based Weighting
Rank-based weighting allocates data emphasis according to objective-value rank, with its quantile-level behavior determined by k rather than dataset size N under mild assumptions. The appendix also describes practical implementations, variance reduction, and parameter effects.
- Dataset-size independence: The fraction of total rank weight assigned to a data quantile depends on k but not dataset size N under mild assumptions.This follows from approximating rank-weight sums with harmonic-series logarithms.
- Dataset-size independence: For kN > 1, the empirical distribution of rank weights is essentially independent of N and closely matches the analysis.The independence fails when kN < 1.
- Implementation: Rank weighting can be implemented by sampling points proportionally to their weights or by uniformly sampling points and weighting their loss contributions.Both approaches are described for mini-batch stochastic gradient descent.
- Implementation: Naive weighted mini-batches can have high variance, so points with weights above w_max are duplicated and split into equal smaller weights.The typical setting was w_max = 5.0 after normalization, balancing variance reduction against dataset growth.
- Implementation: The weighting procedure is exposed through weighted samplers or per-example loss weighting, while rank weights can be computed from output ordering using a single parameter k.The provided implementations use WeightedRandomSampler, weighted losses, and reciprocal rank weights.
B.2 Top10 and Top50 Optimization Results
Top10 and Top50 optimization scores show results qualitatively similar to the main experiment. This suggests weighted retraining finds many distinct high-scoring points rather than repeatedly recovering the same solution.
- Top10 and Top50 results: Top10 and Top50 scores are qualitatively similar to the main experiment, suggesting that weighted retraining finds many unique high-scoring points.The corresponding evaluations are shown in Figures 11 and 12.
B.3 Comparison of Chemical Design Results with Previous Papers
The chemical design experiments compare weighted retraining with prior methods and describe the Bayesian optimization and scoring procedures used to obtain the reported results.
- Comparison with previous papers: Weighted retraining clearly beats previous reinforcement-learning methods while using fewer samples on the chemical design task.The comparison uses the top three scores, with results reported as medians of five runs unless otherwise stated.
- Best molecules: Weighted retraining finds extremely large, high-scoring molecules outside the original training distribution, although molecular size reduces the value of this design task for reinforcement-learning comparisons.The figure reports scores for example molecules found by weighted retraining.
- Optimization procedure: The experiments use a query budget of B = 500 function evaluations, with retraining performed repeatedly at multiples of the chosen frequency.The budget is double that used in the cited prior methods.
- Optimization procedure: The Bayesian optimization baseline uses a variational sparse Gaussian process surrogate with expected improvement over the latent manifold.The sparse Gaussian process uses 500 inducing points.
- Optimization procedure: Fitting the sparse Gaussian process on the 2,000 highest-scoring points and 8,000 random points substantially improves optimization performance, including for the baseline model.The subset both reduces computational cost and focuses fitting on high-performing regions.
- Chemical scoring: The chemical score combines normalized log P, SA, and cycle properties, with a lower bound of −4 to prevent highly negative values from disrupting optimization.The lower bound changes little about maximization outcomes but substantially helps optimization.
C.7 Other Reproducibility Details
The paper reports hyperparameter selection, runtime measurement, and computing infrastructure to support reproducibility of the experiments.
- Hyperparameters: The authors selected k = 10^-3 in advance of the final experiments as an intermediate value that consistently performed well across tasks.They first examined values ranging from 10^1 to 10^-5 and observed a regime where improvement became significant below a threshold.
- Runtime reporting: Runtime results for the main experiments are reported in Table 2 as approximate runtimes for each result.The paper separately identifies the table as the source of runtime measurements.
- Computing infrastructure: All experiments used a single GPU, with parallel execution on a high-performance computing cluster used only for convenience.The authors state that the experiments could in principle be performed on one machine with one GPU.