# Enhancing Generative Molecular Design via Uncertainty-guided Fine-tuning of Variational Autoencoders

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 47 references that are in the map.
- First posted on arXiv 31 May 2024.

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

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [Molecular De Novo Design through Deep Reinforcement Learning](https://mlsanity.com/paper/1704.07555/) (2017)
- [Therapeutics Data Commons: Machine Learning Datasets and Tasks for Drug Discovery and Development](https://mlsanity.com/paper/2102.09548/) (2021)
- [Hierarchical Generation of Molecular Graphs using Structural Motifs](https://mlsanity.com/paper/2002.03230/) (2020)
- [Sample Efficiency Matters: A Benchmark for Practical Molecular Optimization](https://mlsanity.com/paper/2206.12411/) (2022)
- [MARS: Markov Molecular Sampling for Multi-objective Drug Discovery](https://mlsanity.com/paper/2103.10432/) (2021)
- [Sample-Efficient Optimization in the Latent Space of Deep Generative Models via Weighted Retraining](https://mlsanity.com/paper/2006.09191/) (2020)
- [Bayesian Compression for Deep Learning](https://mlsanity.com/paper/1705.08665/) (2017)
- [Multi-Objective De Novo Drug Design with Conditional Graph Generative Model](https://mlsanity.com/paper/1801.07299/) (2018)
- [Subspace Inference for Bayesian Deep Learning](https://mlsanity.com/paper/1907.07504/) (2019)
- [Rethinking Parameter Counting in Deep Models: Effective Dimensionality Revisited](https://mlsanity.com/paper/2003.02139/) (2020)
- [Improving black-box optimization in VAE latent space using decoder uncertainty](https://mlsanity.com/paper/2107.00096/) (2021)
- [Conditional molecular design with deep generative models](https://mlsanity.com/paper/1805.00108/) (2018)
- [Improving Molecular Design by Stochastic Iterative Target Augmentation](https://mlsanity.com/paper/2002.04720/) (2020)
- [Layer Adaptive Node Selection in Bayesian Neural Networks: Statistical Guarantees and Implementation Details](https://mlsanity.com/paper/2108.11000/) (2021)
- [Multi-Objective Latent Space Optimization of Generative Molecular Design Models](https://mlsanity.com/paper/2203.00526/) (2022)
- [Learning Active Subspaces for Effective and Scalable Uncertainty Quantification in Deep Neural Networks](https://mlsanity.com/paper/2309.03061/) (2023)
- [Spike-and-slab shrinkage priors for structurally sparse Bayesian neural networks](https://mlsanity.com/paper/2308.09104/) (2023)
- [Fine-Tuning Generative Models as an Inference Method for Robotic Tasks](https://mlsanity.com/paper/2310.12862/) (2023)
- [Leveraging Active Subspaces to Capture Epistemic Model Uncertainty in Deep Generative Models for Molecular Design](https://mlsanity.com/paper/2405.00202/) (2024)

## Built on by (most cited first, 2 of 2)

- [Inverse Design of Inorganic Compounds with Generative AI](https://mlsanity.com/paper/2604.11827/) (2026)
- [Compressed Active Subspaces for Scalable Bayesian Inference](https://mlsanity.com/paper/2609.19539/) (2026)

## Abstract (arXiv)

In recent years, deep generative models have been successfully adopted for various molecular design tasks, particularly in the life and material sciences. A critical challenge for pre-trained generative molecular design (GMD) models is to fine-tune them to be better suited for downstream design tasks aimed at optimizing specific molecular properties. However, redesigning and training an existing effective generative model from scratch for each new design task is impractical. Furthermore, the black-box nature of typical downstream tasksunicodex2013such as property predictionunicodex2013makes it nontrivial to optimize the generative model in a task-specific manner. In this work, we propose a novel approach for a model uncertainty-guided fine-tuning of a pre-trained variational autoencoder (VAE)-based GMD model through performance feedback in an active learning setting. The main idea is to quantify model uncertainty in the generative model, which is made efficient by working within a low-dimensional active subspace of the high-dimensional VAE parameters explaining most of the variability in the model's output. The inclusion of model uncertainty expands the space of viable molecules through decoder diversity. We then explore the resulting model uncertainty class via black-box optimization made tractable by low-dimensionality of the active subspace. This enables us to identify and leverage a diverse set of high-performing models to generate enhanced molecules. Empirical results across six target molecular properties, using multiple VAE-based generative models, demonstrate that our uncertainty-guided fine-tuning approach consistently outperforms the original pre-trained models.
