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

TL;DR

This work presents an uncertainty-guided fine-tuning framework for pre-trained VAE-based generative molecular design models, leveraging a low-dimensional active subspace of decoder parameters to perform black-box optimization guided by downstream QoI feedback. By constructing an active subspace over stochastic decoder components and optimizing the subspace distribution via Bayesian optimization or REINFORCE (with KL constraints), the approach yields diverse high-performing models that decode latent points into molecules with improved properties across six tasks. Empirical results across JT-VAE, SELFIES-VAE, and SMILES-VAE show consistent QoI gains, with BO generally outperforming REINFORCE for JT-VAE and yielding competitive gains for the other models. The work also analyzes intrinsic biases in learned subspaces, demonstrates the method’s applicability to multi-task scenarios, and outlines future directions for robustness and risk-aware optimization in generative molecular design.

Abstract

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 tasks$\unicode{x2013}$such as property prediction$\unicode{x2013}$makes 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.

Figures (11)

• Figure 1: Illustration of the quantity of interest ($QoI$) enhancement problem. Using a pre-trained VAE-based generative model (PTM), an algorithm $\mathcal{A}$ finds a set of design points -- $Q$ in its latent space. As a downstream task, a property predictor is applied to the molecules corresponding to $Q$ to obtain the pre-trained model' $QoI$ ($QoI_{\text{PTM}}$). Our objective is to fine-tune the model parameters to further enhance $QoI$ for the same $Q$. We propose to leverage the active subspace of model parameters and perform black-box optimization over the subspace parameters with $QoI$ feedback.
• Figure 2: Improvement in $QoI$ relative to the pre-trained JT-VAE model for two optimization methods: BO and REINFORCE. Positive $QoI$ improvement values indicate better $QoI$ than $QoI_{\text{PTM}}$. Each boxplot includes individual $QoI$ improvements for the best fine-tuned distributions found across $10$ different $Q$ sets over 3 optimization trials per optimization method. Some individual observations are horizontally adjusted within each category to remove overlaps among them.
• Figure 3: Comparison of subspace similarity between random subspaces and active subspaces for the JT-VAE tree decoder, SELFIES-VAE decoder and SMILES-VAE decoder. Each entry of the normalized subspace similarity is obtained using (\ref{['sim_eqn']}). In each case, the subspace similarity is calculated between two subspaces generated using two different random seeds.
• Figure 4: Improvement in $QoI$ relative to the pre-trained SELFIES-VAE model for two optimization methods: BO and REINFORCE. Positive $QoI$ improvement values indicate better $QoI$ than $QoI_{\text{PTM}}$. Each boxplot includes individual $QoI$ improvements for the best fine-tuned distributions found across $10$ different $Q$ sets over 3 optimization trials per optimization method. Some individual observations are horizontally adjusted within each category to remove overlaps among them.
• Figure 5: Improvement in $QoI$ relative to the pre-trained SMILES-VAE model for two optimization methods: BO and REINFORCE. Positive $QoI$ improvement values indicate better $QoI$ than $QoI_{\text{PTM}}$. Each boxplot includes individual $QoI$ improvements for the best fine-tuned distributions found across $10$ different $Q$ sets over 3 optimization trials per optimization method. Some individual observations are horizontally adjusted within each category to remove overlaps among them.