Multi-Objective Latent Space Optimization of Generative Molecular Design Models

TL;DR

MO-LSO addresses multi-objective molecular design by integrating Pareto-based weighting into latent-space retraining of a JT-VAE, enabling efficient exploration of trade-offs among multiple properties without ad hoc scalarization. By ranking training molecules via non-dominated sorting and updating weights, the method iteratively shifts the latent space toward regions yielding higher multi-property scores, augmenting the dataset with top Pareto candidates, and repeating. Across six property pairs and a three-property case, MO-LSO achieves larger Pareto-front hypervolumes than scalarized baselines and demonstrates robustness to incomplete data, including recovering DRD2-inhibitor candidates when actives are scarce. The approach offers a scalable, data-efficient pathway for joint optimization of multiple drug-design attributes and can complement advanced search strategies such as Bayesian optimization. It shows potential for practical drug discovery workflows where balancing efficacy, synthesizability, and drug-like properties is essential.

Abstract

Molecular design based on generative models, such as variational autoencoders (VAEs), has become increasingly popular in recent years due to its efficiency for exploring high-dimensional molecular space to identify molecules with desired properties. While the efficacy of the initial model strongly depends on the training data, the sampling efficiency of the model for suggesting novel molecules with enhanced properties can be further enhanced via latent space optimization. In this paper, we propose a multi-objective latent space optimization (LSO) method that can significantly enhance the performance of generative molecular design (GMD). The proposed method adopts an iterative weighted retraining approach, where the respective weights of the molecules in the training data are determined by their Pareto efficiency. We demonstrate that our multi-objective GMD LSO method can significantly improve the performance of GMD for jointly optimizing multiple molecular properties.

Figures (13)

• Figure 1: Overview of the proposed multi-objective latent space optimization scheme. The initial JT-VAE model is trained based on the original training dataset (Step-1). The weights of the molecules in the dataset are adjusted according to their Pareto front ranking based on the properties of interest. Desirable molecules with a higher ranking are assigned larger weights, while molecules with a lower ranking are assigned smaller weights. The JT-VAE is retrained based on the re-weighted dataset (Step-2). The retrained model is used to suggest novel molecules with enhanced properties by sampling or optimization in the latent space (Step-3). Top molecules are selected and used to augment the current training dataset (Step-4). Steps 2--4 may be repeated for iterative retraining of the generative model.
• Figure 2: Evolution of the property distribution of the generated molecules due to latent space optimization via iterative weighted retraining. The plots show how the property distribution changes as a result of weighted retraining of the JT-VAE based on the proposed multi-objective latent space optimization scheme. The latent space of the JT-VAE was optimized to suggest molecules with larger logP and smaller SAS. Results are shown for different values of $k$, which determines the sensitivity of the weight to ranking. In each subfigure, the first violin plot (labeled "train") shows the property distribution of all molecules in the initial training dataset. The subsequent violin plots show the property distribution of 1,000 randomly sampled molecules after $i$-th iterative retraining ($i=0$ corresponds to the original JT-VAE without any retraining). The results clearly show that the distribution of logP is shifted upward while that of SAS is shifted downward during the iterative retraining process as desired.
• Figure 3: Evolution of the Pareto front via latent space optimization of the generative model. The latent space of the JT-VAE has been jointly optimized to maximize logP and minimize SAS of the molecules suggested by the generative model. The scatter plots show the (logP, SAS) distribution of the molecules in the initial training dataset (column 1), molecules sampled in the latent space of the baseline model (column 2), and molecules suggested by the optimized model after iteration-1 (column 3), iteration-5 (column 4), iteration-10 (column 5). The plots in the top row show the trends for the case when the complete training dataset was used, while the bottom row shows the trend when a reduced dataset was used. The results show that the Pareto front gradually shifts towards the desired direction (i.e., bottom right for larger logP and smaller SAS) resulting in a larger hypervolume (HV) of the Pareto front dominated property space in both cases.
• Figure 4: Transition of DRD2 and SAS towards optimum direction while starting with no DRD2 active training samples. The first two scatter plots represent the training data and the molecules generated from the baseline model respectively. In rest of the jointplots, the molecules from learned model after $1^{st}$, $5^{th}$ and $10^{th}$ weighted retraining for $k=10^{-6}$ are shown in the objective space. The hypervolume (HV) in each plot indicates the volume of the property space that is dominated by the Pareto front. As we iterate the weighted retraining, the resulting hypervolume tends to increase.
• Figure 6: Properties of a top molecule predicted by the optimized generative model. The properties of a top molecule (compound A) suggested by the JT-VAE, whose latent space was optimized by the proposed method. Six parameters -- POLAR (Polarity), INSOLU (Insolubility), INSATU (Instauration), FLEX (Rotable bond flexibility), LIPO (Lipophilicity), and SIZE (Molecular Weight) -- are shown. We can see that the suggested compound is within the colored zone, which corresponds to the physio-chemical space suitable for oral bioavailability SwissADME_paper.