Conditional Latent Space Molecular Scaffold Optimization for Accelerated Molecular Design

A VAE Kingma & Welling (2014) consists of an encoder $f_{\phi}^{\text{enc}}:\mathcal{X}\to\mathcal{Z}$ and a decoder $f_{\theta}^{\text{dec}}:\mathcal{Z}\to\mathcal{X}$, where $\mathcal{X}$ represents the input space and $\mathcal{Z}$ the latent space. CVAEs extend the framework of VAEs by incorporating additional condition vector $\bm{c}$ into the latent variable model, facilitating the controlled generation of new instances. In the CVAE architecture, the encoder $q_{\phi}(\bm{z}|\bm{x},\bm{c})$ maps an input $\bm{x}$ and a condition $\bm{c}$ to a latent representation $\bm{z}$. Simultaneously, the decoder $p_{\theta}(\bm{x}|\bm{z},\bm{c})$ reconstructs $\bm{x}$ using both $\bm{z}$ and $\bm{c}$. The training of CVAEs is formulated as the minimization of the conditional variational lower bound:

where $\mathrm{KL}$ denotes the Kullback-Leibler divergence. In this model, the prior distribution $p(\bm{z})$ over the latent variables is typically assumed to be a standard normal distribution, $\mathcal{N}(0,I)$. This assumption simplifies the learning process by standardizing the latent space, ensuring that the encoder learns a distribution that closely aligns with a prior distribution, thus enhancing the generative capability of the decoder conditioned on specific contexts.

In BO, we start with numerous unlabeled instances $\{{\bm{x}_{i}}\}_{i\in[\mathcal{U}]}$ and a smaller set of labeled instances ${(\bm{x}_{i},y_{i})}_{i\in[\mathcal{L}]}$, where an input $\bm{x}_{i}\in\mathcal{X}$ represents a chemical compound, and a label $y_{i}\in\mathcal{Y}\subseteq\mathcal{R}$ indicates its properties such as docking scores. BO seeks to optimize a costly black-box function $f^{\rm BB}:\mathcal{X}\to\mathcal{Y}$, which corresponds to obtaining the physical properties of chemical compounds through experiments or time-consuming simulations in the context of molecular design problems. The goal is to maximize $f^{\rm BB}$ with minimal evaluations, using typically a Gaussian Process (GP) surrogate, trained using the labeled instances $\mathcal{L}$, to predict the function over $\mathcal{X}$. BO uses the surrogate to select an input $\bm{x}$ that may yield values surpassing the current maximum $\max_{i\in\mathcal{L}}y_{i}$. However, building a GP surrogate in high-dimensional spaces like chemical compounds is challenging. LSBO tackles this by employing a VAE/CVAE trained on the unlabelled instances $\mathcal{U}$ to reduce dimensionality, encoding instance in $\mathcal{X}$ to lower dimensional latent space $\mathcal{Z}$. This simplifies surrogate modeling and optimization because $\mathcal{Z}$ has lower-dimension than $\mathcal{X}$. During LSBO iterations, the acquisition function applied to the GP’s predictions selects new points in $\mathcal{Z}$ to evaluate. The chosen latent variable $\bm{z}_{i\textquoteright}$ is decoded into a new input $\bm{x}_{i\textquoteright}=f_{\theta}^{\rm dec}(\bm{z}_{i^{\prime}})$. This new instance is evaluated by $f^{\rm BB}$, and the results update $\mathcal{L}$ and refine the GP model. This cycle repeats until optimal results are achieved or resources are exhausted. In contexts like molecular design, LSBO aims to discover chemical compounds with optimal properties by efficiently navigating the reduced latent space.

Scaffolds Bemis & Murcko (1996) are the stable core structures within molecules that serve as the framework for chemical modifications in drug design. Scaffolds retain the essential biological activity of the molecule. They play a crucial role in molecular design by providing a foundation for chemical modifications aimed at optimizing properties like QED. Researchers often use scaffolds to systematically explore chemical variations Schreiber (2000), Welsch et al. (2010), which can lead to the discovery of new compounds with improved properties.

In this study, we follow the scaffold extraction method from Bemis & Murcko (1996), where non-essential components like side chains are removed, leaving the core structure. This extracted scaffold serves as a starting point for further modifications, allowing efficient exploration of chemical space. By focusing on scaffolds, molecular design becomes more streamlined, increasing the likelihood of identifying novel, synthesizable compounds with the desired biological activity. Examples of whole molecule and scaffold pairs are provided in Fig. 2.

We denote the scaffold, which is the base of the modification, as $\bm{S}$, and the modified molecule as $\bm{S}^{\prime}$. Our goal is to efficiently find the modification that maximizes the molecular property P which is evaluated by $f^{\text{BB}}(\bm{S}^{\prime})$, while keeping the difference between $\bm{S}$ and $\bm{S}^{\prime}$ small. Directly optimizing over all possible $\bm{S}^{\prime}$ to find the best modification that maximizes $f^{\text{BB}}(\bm{S\textquoteright})$ is impractical, as it involves high complexity and costly evaluations of $f^{\text{BB}}$.

The primary challenge in optimizing molecular scaffolds lies in i) determining the optimal bonding point on the base scaffold $\bm{S}$, and ii) selecting the appropriate substructures added to the bonding point to ensure meaningful improvements in the desired property $\mathcal{P}$. A molecular scaffold $\bm{S}$, composed of several atoms $p_{1},p_{2},\ldots,p_{k}$, may have atoms with the remaining capacity to form additional chemical bonds. These atoms serve as potential candidates for bonding with newly generated substructures. Therefore, the task involves not only selecting the right substructure but also identifying the most suitable bonding point $p_{i}$ to optimize scaffold properties. This adds complexity, as the need for precise modifications must be balanced with the challenges of high-dimensional search spaces and evaluation costs. Consequently, a more efficient approach is required to explore scaffold modifications effectively while minimizing the number of evaluations.

To address this, the problem can be reframed as an optimization task in a reduced latent space $\mathcal{Z}$, obtained through a CVAE. In this space, each point $\bm{z}\in\mathcal{Z}$ corresponds to a potential substructure that can be integrated into the scaffold. By encoding the molecular substructures into this lower-dimensional space $\mathcal{Z}\in\mathbb{R}^{d}$, the search becomes more tractable. The objective is to find the optimal latent representation $\bm{z^{*}}$ conditioned on the optimal bonding point $p_{i}^{*}$ that, together, maximize the desired property $\mathcal{P}$ when the generated substructure $\bm{s^{\prime}}\leftarrow f^{\text{dec}}(\bm{z})$ is added to the scaffold. Let us denote this modification as $\bm{S}^{{}^{\prime}}\leftarrow g(\bm{S}\oplus\bm{s^{\prime}},p_{i})$, where the function $g()$ adds substructure $\bm{s}$ to the scaffold $\bm{S}$ at $p_{i}$. The optimization problem is then formulated as:

where $B(\bm{S})$ is the set of possible bonding points on the scaffold $\bm{S}$. In Section 4.2, we demonstrate that atomic features at $p_{i}$ are used as condition vectors to generate new substructures, enabling targeted substructure generation for atom $p_{i}$.

Our proposed method requires a uniquely tailored dataset because no existing dataset in the literature fully meets the specific needs of our approach. To create this dataset, we developed a BRICS (Breaking Retrosynthetically Interesting Chemical Substructures) Degen et al. (2008) based approach. BRICS is an algorithm designed to decompose organic molecules into smaller, synthetically feasible substructures by identifying breaking points within a molecule’s structure based on chemical retrosynthetic rules. These breaking points, also known as division points, are the connections between subgraphs within the molecule that BRICS identifies. The algorithm systematically breaks down a molecule $\bm{M}$ into $k$ substructures, Fig. 3 demonstrates this procedure.

The division points found by BRICS are of particular importance because they serve as both the points where the molecule is split into substructures and the potential bonding sites where these substructures can be reattached. Therefore, they provide us a valuable information about these substructures in terms of what kind of bonds they can form. This dual role makes the division points a crucial piece of information for our generative model. By preserving the properties and features of atoms at these division points, we capture the chemical context around the atom that dictates how substructures can bond with other parts of a molecule. The atomic features we consider are atom type, hybridization, valence, formal charge, degree, and ring membership. Detailed explanations of these features are provided in the Appendix A.1. These features, when used as condition vectors in our CVAE, are essential for guiding the generation of substructures that can successfully bond with the scaffold, as they provide a detailed understanding of the chemical environment at each division point, such as their atom types and if the selected atom has the capacity for forming additional bonds. Therefore, our dataset for CVAE training includes the substructures obtained via the BRICS algorithm and six different atomic features of the atom at their breaking points. The illustration of the BRICS algorithm, atomic feature extraction points, and their integration with CVAE is provided in Fig. 3.

Our CVAE consists of an encoder, parameterized by $f^{\rm{enc}}_{\phi}$, and a decoder, parameterized by $f^{\rm{dec}}_{\theta}$. The encoder takes the substructure $\bm{s}$ and the associated condition vector $\bm{c}$, and maps this input into a latent space, producing a latent representation $\bm{z}$. The decoder then takes a point $\bm{z}$ from this latent space, along with the condition vector $\bm{c}$, and generates a substructure $\bm{s}^{\prime}$ that is conditioned on the given atomic properties. The loss function of the proposed CVAE is defined as:

where $q_{\bm{\phi}}(\bm{z}|\bm{s},\bm{c})$ is the encoder, $p_{\bm{\theta}}(\bm{s}|\bm{z},\bm{c})$ is the decoder, and $\beta$ balances the reconstruction and KL divergence terms Higgins et al. (2016). The condition vectors allow the CVAE to learn how different substructures interact with specific atomic environments, building a deeper understanding of their bonding behavior. Figure 3B demonstrates the input substructures that CVAE uses in its training, and Fig. 3C visualizes the CVAE architecture along with the condition vectors.

After training, if during the generation phase, we provide the CVAE with a condition vector that specifies the atomic environment of the target scaffold, the decoder, using a latent vector along with the condition vector, generates substructures that are tailored to bond effectively with the scaffold. This conditioning mechanism ensures that the generated substructures are not random but specifically designed to fit the scaffold’s atomic environment, improving the efficiency of scaffold modifications and the likelihood of successful bonding. In the proposed CLaSMO method, we jointly optimize the bonding position $p_{i}$ and the latent vector $\bm{z}$, conditioned on the atomic properties $\bm{c}$ of the bonding site in order to optimize both the bonding position and the substructure added to the position.

In CLaSMO, we perform a targeted search in the latent space $\mathcal{Z}\in\mathbb{R}^{d}$ of the CVAE, where decodings from each point $\bm{z}$ represent a potential substructure to be added to the scaffold. The challenge is to modify the scaffold $\bm{S}$ by selecting a substructure and integrating it at an appropriate bonding point $p_{i}$, optimizing the desired molecular property $\mathcal{P}$. Therefore, the search space $\Omega$ we consider in our optimization tasks is defined as the Cartesian product of the latent space and the set of possible bonding points on the scaffold:

Within this search space, the goal is to modify the scaffold $\bm{S}$ to maximize the BB function $f^{\rm BB}(\bm{S^{\prime}})$.

However, although our main goal is optimizing the target property, we also aim to keep the similarity between $\bm{S}$ and $\bm{S^{\prime}}$ at a certain level. Therefore, in CLaSMO, we apply a similarity constraint using the Dice Similarity Dice (1945) metric to compare the input scaffold $\bm{S}$ with the modified molecule $\bm{S^{\prime}}$ before evaluating the BB function. In order to measure the similarity between $\bm{S}$ and $\bm{S^{\prime}}$, we use Morgan Fingerprints (Morgan, 1965, Rogers & Hahn, 2010). The Morgan fingerprint of a molecule is represented as a binary vector, where each bit indicates the presence or absence of a specific substructure within the molecule, making them effective and popular for comparing molecular similarities. Using these, the similarity between $\bm{S}$ and $\bm{S^{\prime}}$ is measured by computing the Dice Similarity of their Morgan fingerprint vectors¹¹1 Bajusz et al. (2015) discusses that the dice similarity is one of the best metrics to assess similarity between molecules using Morgan fingerprints., defined as:

where $\bm{M_{S}}$ and $\bm{M_{S}^{\prime}}$ are the Morgan fingerprint vectors of $\bm{S}$ and $\bm{S^{\prime}}$, respectively. Dice Similarity, ranging from 0 (no overlap) to 1 (identical), helps us track structural changes during optimization.

The final optimization objective, incorporating both the search for optimal substructures and the similarity constraint, is given by:

where $\bm{S}^{{}^{\prime}}\leftarrow g(\bm{S}\oplus f^{\text{dec}}(\bm{s^{\prime}}|\bm{z}^{*},\bm{c}))$. By this approach, CLaSMO efficiently navigates the latent space, optimizing molecular properties while allowing modifications only if $\text{DICE}(\bm{S},\bm{S}^{\prime})\geq\tau$, effectively controlling the degree of divergence from the input scaffold.

For the joint optimization of substructures and bonding positions, we consider a GP model: $(\bm{z},p)\mapsto y^{\prime}_{\Delta}$, where $\bm{z}$ is the latent vector representing a substructure, $p$ is the bonding position of the modification, and $y^{\prime}_{\Delta}:=y-y^{\prime}$ represents the improvement in the property, with $y$ and $y^{\prime}$ indicating the properties of molecules $S$ and $S^{\prime}$, respectively. To prepare the training dataset $\mathcal{D}$ for the GP, we conducted a random sampling of substructures from random latent vectors $\bm{z}$, each paired with randomly selected bonding regions $p$ on the scaffold $\bm{S}$, and evaluated the $y^{\prime}_{\Delta}$ obtained from these additions (e.g., in the case of QED optimization, QED score differences between $\bm{S}$ and $\bm{S^{\prime}}$ are calculated). These triplets of $(\bm{z},p,y^{\prime}_{\Delta})$ are used in GP training. This setup allows the GP model to learn the relationship between the latent space representations, atoms in the input scaffold, and the resulting property changes, guiding the optimization process toward regions of the latent space that are more likely to yield beneficial modifications to the scaffold.

Among many possible choices, we employ the Upper Confidence Bound (UCB) acquisition function to guide the optimization process. To ensure LSBO samples from regions that meet the similarity constraint, we introduced a penalization mechanism. Specifically, we assign a negative improvement value $y^{\prime}_{\Delta}$ when the condition $\text{DICE}(\bm{S},\bm{S}^{\prime})\geq\tau$ is not satisfied after sampling from $f^{\text{dec}}([\bm{z}^{*},\bm{c}])$. Additionally, we also apply a penalty when a substructure cannot bond with the target molecule. We outline our approach in Algorithm 1, and provide details about the selection of hyperparameters within our framework in the Appendix A.3.

For our data preparation, we used the QM9 dataset Ruddigkeit et al. (2012), Ramakrishnan et al. (2014), which contains 130,000 small molecules. QM9 was chosen for its simplicity and suitability for our task, as it consists of smaller molecules compared to other well-known datasets like ZINC Irwin & Shoichet (2005), Irwin et al. (2020). Using our data preparation strategy, we extracted 18,706 unique substructure pairs along with their atomic environment features to train our CVAE. Most of the substructures obtained from QM9 through our method are under 100 Daltons, making them ideal for training a model focused on generating small substructures to modify the input scaffold. Since our goal is to optimize the scaffold while maintaining a high degree of similarity in the modifications, using such a dataset allows CVAE to learn to generate small substructures, which is crucial for achieving our goals. Of the 18,706 instances, 80% were used for model training, with the remaining data allocated for testing and validation. We represented the molecules using SELFIES Krenn et al. (2020), a string-based molecular representation, in the form of one-hot encoding matrices. Our CVAE architecture consists of three fully connected layers in the encoder and three GRU layers in the decoder.

Given the simplicity of the dataset, we determined that a 2-dimensional latent space was sufficient to achieve over 99% reconstruction accuracy on the test set. This lower-dimensional space also enhances the performance of BO, which typically excels in smaller-dimensional spaces. Additionally, to maintain this simplicity, as explained in Section 4.2.1, we generated a 2-dimensional embedding for the condition vectors using an Autoencoder model, which is trained with fully connected layers in both the encoder and decoder, achieved 93% reconstruction accuracy for the condition vectors. Prior to CVAE training, we obtained the condition vector embeddings via the encoder of this VAE model, and used them during CVAE training. As a result, the CVAE was trained with both a 2-dimensional latent space for the substructures and 2-dimensional condition vectors, effectively balancing simplicity and optimization performance.

In this section, we present the results of our QED optimization experiments, where we used RDKit Landrum (2010) to calculate QED values. The QED metric, defined between 0 and 1, inherently limits the range of potential improvements, making optimization within this closed range particularly challenging. To further complicate this task, we imposed similarity constraints to ensure the optimized molecules remain close to their original scaffolds, by running the CLaSMO algorithm with $\tau$ values $\tau\in[0,0.25,0.5,0.6]$. For each threshold, CLaSMO was run for 100 iterations on 100 input scaffolds, where their whole molecules are sampled randomly from the ZINC250K dataset Gómez-Bombarelli et al. (2018). The search domain for each latent dimension in $\mathcal{Z}$ is set to $[-6,6]$. The GP model is trained using 100 training instances. Such an experimental setting is designed to demonstrate the optimization capabilities of CLaSMO at changing similarity thresholds, within the limited range of QED optimization tasks.

Initially, the average QED score of the input scaffold molecules was 0.5876, with a maximum QED of 0.8902. After optimization, CLaSMO with no similarity constraint ($\tau=0$) achieved the highest QED score of 0.9480 and an average QED of 0.6778, representing a mean improvement rate of 21.43%²²2Mean improvement rates calculated as the average of the relative improvements obtained by CLaSMO for each input scaffolds. and a maximum improvement rate of 81.17%. As the similarity constraint was tightened, the performance remained competitive. With $\tau=0.25$, the maximum QED score was 0.9464, with an average QED of 0.6715, leading to a mean improvement rate of 19.14%. Further increasing the similarity threshold to $\tau=0.50$, the maximum QED score was 0.9437, and the average QED was 0.6602, yielding a mean improvement rate of 16.08%. For $\tau=0.60$, the maximum QED score was 0.9402, and the average QED was 0.6434, resulting in a mean improvement rate of 11.70%. The results clearly demonstrate that CLaSMO consistently enhances the QED values across a range of similarity thresholds, effectively balancing the trade-off between maximizing QED and preserving structural similarity. Notably, even under more stringent similarity constraints, CLaSMO achieves significant improvements over the initial QED values, highlighting its robust optimization capabilities. This consistent performance across different thresholds underscores CLaSMO’s versatility and effectiveness in optimizing QED, making it a powerful tool for scaffold-based molecular optimization tasks. Table 1 summarizes these findings. These results were achieved with only 100 iterations per input scaffold, demonstrating the sample efficiency of our approach. Additionally, Fig. 4 demonstrates two example sets of scaffold optimization results from the CLaSMO experiments, in which it can be observed that our algorithm finds a molecule with a QED score of 0.948 by very limited additions, keeping the similarity between the input and optimized molecule at the high level.

In this section, we present our docking simulation score optimization results using the KAT1 protein as the target. KAT1 is an ion channel protein, and compounds with favorable docking scores are more likely to bind effectively and influence its function, which is crucial for developing therapeutic agents targeting ion channel-related conditions⁴⁴4Details of the KAT1 protein can be found at https://pdbj.org/mine/summary/6v1x.. By optimizing docking scores, we aim to identify compounds with the potential to modulate ion channel function and develop novel therapeutic strategies. Unlike the relatively simple calculation of QED values, optimizing docking scores for KAT1 requires computationally intensive simulations. For these calculations, we used Schrödinger Software Schrödinger (2023), known for providing reliable and accurate results, with each simulation taking anywhere from a few minutes to several hours. Given the high computational cost of these simulations, CLaSMO’s sample-efficient approach is particularly valuable, enabling effective optimization while minimizing the number of expensive evaluations.

Similar to the QED optimization setting, the search domain for each latent dimension in $\mathcal{Z}$ was set to $[-6,6]$. To address the high cost of obtaining labeled data for docking scores, we leveraged the same training set used in the QED experiments, treating them as low-fidelity instances for training the GP model. This approach allows the surrogate model to learn the relationships between atoms and latent vectors, even though the labels originate from QED rather than docking scores. By reusing this shared data, we enhance the GP model’s capacity to generalize and effectively guide the optimization process, reducing the dependency on extensive docking score evaluations. Given the significant time required for each docking score calculation, we limited the experiments to 10 compounds from the ZINC250K dataset and ran CLaSMO for 100 iterations per compound using two Dice similarity thresholds, $\tau\in[0.25,0.50]$. This setup ensures that despite the computational expense, we can still obtain comprehensive and meaningful results for docking score optimization, and observe if CLaSMO is capable of optimizing docking scores under similarity constraints.

We present the distribution of docking scores obtained from both CLaSMO experiments in Fig. 7. The results clearly demonstrate that CLaSMO, at both similarity thresholds, achieves significant improvements in docking scores compared to the initial scaffolds. Figure 7 provides detailed metrics, including the maximum and average improvements across the models, as well as the initial best values. Notably, CLaSMO with a similarity threshold of $\tau=0.25$ achieved improvements of up to 96.3% over the initial scaffold docking score, while CLaSMO with $\tau=0.50$ achieved improvements of up to 75.1%, both indicating strong optimization performance despite the limited number of LSBO iterations, showcasing CLasMO’s sample-efficiency. In Fig. 8, we provide example molecules obtained from our experiment, where we observe that, even at the beginning stages of the optimization where $\text{DICE}(\bm{S},\bm{S^{\prime}})$ still above 0.70, we observe substantial improvements.

We conducted an ablation study to evaluate the impact of incorporating atomic environment features as condition vectors in the generative process. In the training of the baseline VAE model, only substructures of the molecules were used, without conditioning on atomic properties. This setup allowed us to isolate the effect of conditioning on scaffold optimization.

We repeated the QED optimization experiments using the VAE model. The QED optimization results, shown in Table 3, demonstrate that incorporating atomic environment features in the CVAE significantly improves the generation of substructures. The conditioning mechanism enhances the model’s ability to generate substructures that bond more effectively with the scaffold, leading to higher success rates in optimizing molecular properties. While the VAE without conditioning still gains from the LSBO framework’s efficient optimization, incorporating atomic properties as conditions significantly enhances the quality of the generated molecules.

In CLaSMO, the modification region of the input scaffold is typically selected during the automated optimization process. However, the framework also supports an interactive mode, where a chemical expert manually selects the region of the molecule to modify. In this mode, the expert identifies the specific atom or region for modification, rather than relying on CLaSMO’s automated selection. Once the region is chosen, the rest of the process remains unchanged—CLaSMO continues to optimize the molecule by generating substructures using the CVAE and refining them through LSBO to improve target properties.

This interactive approach offers several key advantages. It enables the integration of expert knowledge into the optimization process, allowing CLaSMO to operate in a Human-in-the-loop setting. Experts can leverage their domain-specific insights to target specific regions they find promising, ensuring that modifications are not only data-driven but also aligned with scientific understanding. Meanwhile, CLaSMO maintains its efficient optimization process, using LSBO to enhance molecular properties and preserve sample efficiency. Detailed user guidelines for this open-source web application are provided in the Appendix A.