# Improving black-box optimization in VAE latent space using decoder uncertainty

arXiv: https://arxiv.org/abs/2107.00096 · HTML: https://mlsanity.com/paper/2107.00096/

## Key facts

- Cited by 51 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 46 references that are in the map.
- First posted on arXiv 30 Jun 2021.

## Builds on (most cited first, 25 of 25)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Towards Evaluating the Robustness of Neural Networks](https://mlsanity.com/paper/1608.04644/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [On the Properties of Neural Machine Translation: Encoder-Decoder Approaches](https://mlsanity.com/paper/1409.1259/) (2014)
- [A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges](https://mlsanity.com/paper/2011.06225/) (2020)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [Convolutional Networks on Graphs for Learning Molecular Fingerprints](https://mlsanity.com/paper/1509.09292/) (2015)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [Generating Sentences from a Continuous Space](https://mlsanity.com/paper/1511.06349/) (2015)
- [BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning](https://mlsanity.com/paper/1906.08158/) (2019)
- [GraphAF: a Flow-based Autoregressive Model for Molecular Graph Generation](https://mlsanity.com/paper/2001.09382/) (2020)
- [GuacaMol: Benchmarking Models for De Novo Molecular Design](https://mlsanity.com/paper/1811.09621/) (2018)
- [Grammar Variational Autoencoder](https://mlsanity.com/paper/1703.01925/) (2017)
- [Detecting Adversarial Samples from Artifacts](https://mlsanity.com/paper/1703.00410/) (2017)
- [Learning Deep Generative Models of Graphs](https://mlsanity.com/paper/1803.03324/) (2018)
- [Optimization of Molecules via Deep Reinforcement Learning](https://mlsanity.com/paper/1810.08678/) (2018)
- [Understanding Measures of Uncertainty for Adversarial Example Detection](https://mlsanity.com/paper/1803.08533/) (2018)
- [Sample-Efficient Optimization in the Latent Space of Deep Generative Models via Weighted Retraining](https://mlsanity.com/paper/2006.09191/) (2020)
- [Syntax-Directed Variational Autoencoder for Structured Data](https://mlsanity.com/paper/1802.08786/) (2018)
- [ChemBO: Bayesian Optimization of Small Organic Molecules with Synthesizable Recommendations](https://mlsanity.com/paper/1908.01425/) (2019)
- [A Model to Search for Synthesizable Molecules](https://mlsanity.com/paper/1906.05221/) (2019)
- [Uncertainty Estimation in Autoregressive Structured Prediction](https://mlsanity.com/paper/2002.07650/) (2020)
- [Actively Learning what makes a Discrete Sequence Valid](https://mlsanity.com/paper/1708.04465/) (2017)

## Built on by (most cited first, 30 of 51)

- [End-to-End Diffusion Latent Optimization Improves Classifier Guidance](https://mlsanity.com/paper/2303.13703/) (2023)
- [Local Latent Space Bayesian Optimization over Structured Inputs](https://mlsanity.com/paper/2201.11872/) (2022)
- [MolGenSurvey: A Systematic Survey in Machine Learning Models for Molecule Design](https://mlsanity.com/paper/2203.14500/) (2022)
- [Feedback Efficient Online Fine-Tuning of Diffusion Models](https://mlsanity.com/paper/2402.16359/) (2024)
- [Combining Latent Space and Structured Kernels for Bayesian Optimization over Combinatorial Spaces](https://mlsanity.com/paper/2111.01186/) (2021)
- [LIMO: Latent Inceptionism for Targeted Molecule Generation](https://mlsanity.com/paper/2206.09010/) (2022)
- [HyAR: Addressing Discrete-Continuous Action Reinforcement Learning via Hybrid Action Representation](https://mlsanity.com/paper/2109.05490/) (2021)
- [Directional Connectivity-based Segmentation of Medical Images](https://mlsanity.com/paper/2304.00145/) (2023)
- [Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models](https://mlsanity.com/paper/2405.19673/) (2024)
- [Bayesian Optimization over High-Dimensional Combinatorial Spaces via Dictionary-based Embeddings](https://mlsanity.com/paper/2303.01774/) (2023)
- [A survey and benchmark of high-dimensional Bayesian optimization of discrete sequences](https://mlsanity.com/paper/2406.04739/) (2024)
- [Diffusion Model Guided Sampling with Pixel-Wise Aleatoric Uncertainty Estimation](https://mlsanity.com/paper/2412.00205/) (2024)
- [Latent Bayesian Optimization via Autoregressive Normalizing Flows](https://mlsanity.com/paper/2504.14889/) (2025)
- [Advancing Bayesian Optimization via Learning Correlated Latent Space](https://mlsanity.com/paper/2310.20258/) (2023)
- [GLSO: Grammar-guided Latent Space Optimization for Sample-efficient Robot Design Automation](https://mlsanity.com/paper/2209.11748/) (2022)
- [Return of the Latent Space COWBOYS: Re-thinking the use of VAEs for Bayesian Optimisation of Structured Spaces](https://mlsanity.com/paper/2507.03910/) (2025)
- [Multi-Objective Latent Space Optimization of Generative Molecular Design Models](https://mlsanity.com/paper/2203.00526/) (2022)
- [Computationally Efficient High-Dimensional Bayesian Optimization via Variable Selection](https://mlsanity.com/paper/2109.09264/) (2021)
- [DiscoBAX: Discovery of Optimal Intervention Sets in Genomic Experiment Design](https://mlsanity.com/paper/2312.04064/) (2023)
- [Inversion-based Latent Bayesian Optimization](https://mlsanity.com/paper/2411.05330/) (2024)
- [An Equivariant Generative Framework for Molecular Graph-Structure Co-Design](https://mlsanity.com/paper/2304.12436/) (2023)
- [A conditional latent autoregressive recurrent model for generation and forecasting of beam dynamics in particle accelerators](https://mlsanity.com/paper/2403.13858/) (2024)
- [Meta-Learning Parameterized Skills](https://mlsanity.com/paper/2206.03597/) (2022)
- [Optimizing Training Trajectories in Variational Autoencoders via Latent Bayesian Optimization Approach](https://mlsanity.com/paper/2207.00128/) (2022)
- [Biological Sequence Kernels with Guaranteed Flexibility](https://mlsanity.com/paper/2304.03775/) (2023)
- [PAnDR: Fast Adaptation to New Environments from Offline Experiences via Decoupling Policy and Environment Representations](https://mlsanity.com/paper/2204.02877/) (2022)
- [MolHF: A Hierarchical Normalizing Flow for Molecular Graph Generation](https://mlsanity.com/paper/2305.08457/) (2023)
- [Thompson Sampling via Fine-Tuning of LLMs](https://mlsanity.com/paper/2510.13328/) (2025)
- [High-Dimensional Bayesian Optimization via Semi-Supervised Learning with Optimized Unlabeled Data Sampling](https://mlsanity.com/paper/2305.02614/) (2023)
- [Logit-Based Ensemble Distribution Distillation for Robust Autoregressive Sequence Uncertainties](https://mlsanity.com/paper/2305.10384/) (2023)

## Abstract (arXiv)

Optimization in the latent space of variational autoencoders is a promising approach to generate high-dimensional discrete objects that maximize an expensive black-box property (e.g., drug-likeness in molecular generation, function approximation with arithmetic expressions). However, existing methods lack robustness as they may decide to explore areas of the latent space for which no data was available during training and where the decoder can be unreliable, leading to the generation of unrealistic or invalid objects. We propose to leverage the epistemic uncertainty of the decoder to guide the optimization process. This is not trivial though, as a naive estimation of uncertainty in the high-dimensional and structured settings we consider would result in high estimator variance. To solve this problem, we introduce an importance sampling-based estimator that provides more robust estimates of epistemic uncertainty. Our uncertainty-guided optimization approach does not require modifications of the model architecture nor the training process. It produces samples with a better trade-off between black-box objective and validity of the generated samples, sometimes improving both simultaneously. We illustrate these advantages across several experimental settings in digit generation, arithmetic expression approximation and molecule generation for drug design.