[Paper Review] Improving black-box optimization in VAE latent space using decoder uncertainty

Pascal Notin, José Miguel Hernández-Lobato+1 more|arXiv (Cornell University)|Jun 30, 2021

Advanced Neural Network Applications26 references24 citations

TL;DR

The paper introduces decoder-uncertainty guided black-box optimization in VAE latent space, using an importance-sampling MI estimator to improve validity and objective across digits, arithmetic expressions, and molecule design.

ABSTRACT

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.

Motivation & Objective

• Motivate robust black-box optimization in high-dimensional discrete spaces via VAEs without altering model architecture or training.
• Quantify decoder epistemic uncertainty and its impact on decoded sample validity.
• Develop a robust estimator for decoder uncertainty using importance sampling in high-dimensional structured data.
• Demonstrate uncertainty-guided optimization across multiple data modalities (digits, arithmetic expressions, molecules).
• Show compatibility with Bayesian optimization and gradient-based latent-space optimization methods.

Proposed method

• Define and estimate decoder epistemic uncertainty using an importance sampling-based Mutual Information (MI) estimator (IS-MI).
• Incorporate decoder uncertainty into optimization frameworks (Bayesian optimization with uncertainty-aware surrogates or censoring; uncertainty-constrained gradient ascent).
• Provide a practical estimator that samples from an approximate posterior over model parameters and decodes under those samples to compute p(y|x,θ).
• Apply the approach in three settings (digit generation, arithmetic expression approximation, molecule generation) with varied architectures (Conv-Deconv, CVAE, JT-VAE).
• Compare IS-MI to token-independence MI (TI-MI) and baselines, assessing validity of decodings and objective performance.

Experimental results

Research questions

• RQ1 Can decoder epistemic uncertainty be reliably estimated in high-dimensional structured outputs?
• RQ2 Does censoring or penalizing high-uncertainty latent points improve the validity of decoded objects without sacrificing objective performance?
• RQ3 How does decoder-uncertainty guided optimization perform across digits, arithmetic expressions, and molecular generation tasks with different architectures?
• RQ4 Is the importance sampling MI estimator (IS-MI) more robust than TI-MI in high-dimensional settings?
• RQ5 Can uncertainty-guided optimization be integrated with Bayesian optimization and gradient-based latent space methods without altering model training?

Key findings

• Decoder uncertainty correlates with decoding validity: lower uncertainty regions yield more valid decodings across tasks.
• IS-MI provides lower-variance MI estimates and better identification of high-risk latent points than TI-MI.
• Uncertainty-guided optimization yields higher penalized objective values and/or higher validity compared with baselines, across digits, arithmetic expressions, and molecule design.
• Censoring high-uncertainty points during Bayesian optimization improves validity and enables discovering high-quality molecules with favorable properties.
• Using NLLP (prior likelihood) constraints can hinder optimization performance compared with IS-MI.
• JT-VAE experiments show uncertainty-guided optimization can achieve high penalized logP values while maintaining quality filters, outperforming baselines that ignore decoder uncertainty.