GPMO: Gradient Perturbation-Based Contrastive Learning for Molecule Optimization

1 Introduction

Molecule Optimization (MO) aims to identify a target molecule, that exhibits improved pharmacological properties while retaining its similarity to the original molecule. Traditionally, chemists rely on their knowledge, experience, and intuition to apply chemical transformations to promising molecules in order to achieve a balance of desirable properties. Additionally, this manual approach can be time-consuming, labor-intensive, and may not explore the full range of possible modifications that could enhance a molecule's pharmacological properties. High throughput screening tests a large library of compounds to identify those with desirable pharmacological properties, but it can be an expensive and inefficient approach to drug discovery.

* Corresponding authors.

In contrast to the traditional method, in silico MO through automated processes to improve drug properties, which can accelerate the drug design process [Zeng et al., 2022; Ma et al., 2022; Lin et al., 2020]. Deep learning technologies, such as reinforcement learning (RL) [Jin et al., 2018a; You et al., 2018; Zhou et al., 2019] and deep generative model (DGM) [Zang and Wang, 2020; Chen et al., 2021; Jin et al., 2020; Liu et al., 2020], have been applied in MO. RL methods view MO as the Markov decision process, in which the current state of the molecule is modified through actions guided by a reward function that reflects desired property. While RL methods have shown promise in improving property, they can significantly alter structure of molecule [Gao et al., 2021]. DGMs optimize molecules by generating new structures through translation between molecule pairs while maintaining the scaffold of the original molecule through similarity constraints. The translation-based DGMs are trained using the teacher-forcing training strategy to predict gold-standard labels of the previous tokens, without being exposed to incorrectly generated tokens. During the testing phase, translation-based DGMs rely on the model's prediction of the previous ones, which can differ from the ground-truth contexts. This discrepancy between training and testing phases is known as the "exposure bias" problem [Lee et al., 2020].

To address the aforementioned issues, we propose GPMO that utilizes contrastive learning [Shen et al., 2022] to overcome the "exposure bias" problem in the teacher-forcing guided model. To generate negative and positive molecules for contrastive learning, we incorporate gradient perturbation. After pre-training on a large-scale unlabeled dataset, GPMO decodes the output of the encoder to generate a target molecule with multiple property constraints. The input of GPMO is the condition tokens and the SMILES of start molecules. During the encoder-decoder process, GPMO generates the negative and positive molecules without domainspecific knowledge, and these molecules are used to train the model with adversarial loss. GPMO generates negative molecules by adding a small gradient perturbation to the hidden state of the target molecule, and the model is trained to minimize the conditional likelihood of the negative molecule to ensure it is classified as a negative example during contrastive learning. In GPMO, the positive molecule is generated using a two-step gradient perturbation on the hidden state of the target molecule. The Kullback-Leibler (KL) divergence ensures that the distributions of the positive and real molecules are similar, guaranteeing that the positive molecule is indeed a positive example. The two-step gradient perturbation keeps the positive molecule far away from the real molecule in the embedding space. This will produce challenging molecules that the model is unable to accurately differentiate, providing it with more useful pairs to learn from. It is worth noting that GPMO is different from the one proposed in [Lee et al., 2020]. On one hand, GPMO focuses on optimizing molecules based on conditional constraints, rather than simply generating text without specific criteria. In this way, GPMO can generate molecules with specific property values. On the other hand, GPMO provides a comprehensive framework for molecule optimization applications, making it highly suitable for large-scale optimization tasks. Experimental results show that GPMO outperforms all baseline methods in identifying molecules that satisfy specific property constraints under various conditions. Our contributions can be summarized as follows:

• GPMO is the first model to address the "exposure bias" problem in MO, and it does so by generating positive and negative molecules for effective contrastive learning.
• GPMO applies positive and negative molecules generation strategy which is free from domain-specific knowledge and easily accessible in the back-propagation process.
• Experimental results show that GPMO achieves SOTA in properties optimization problems, and can optimize one molecule into another molecule with better properties while maintaining the scaffold and exhibiting more desired properties.

2 RelatedWork

2.1 Reinforcement LearningMethod

The objective of RL is to teach intelligent agents to take action in an environment to maximize the cumulative reward. The environment of the RL method is defined by chemical rules, and the improvement of a molecule's property is defined as the reward in RL. The RL method demonstrates effective property improvement and has the capability to operate in a significantly larger chemical space compared to deep learning methods [Wang et al., 2021]. However, the RL method often overcomplicates the molecular structure when optimizing multiple properties [Nigam et al., 2019]. A newly developed RL method [Gao et al., 2021] introduces a synthesis tree to alleviate the synthesis problem in MO that arises from complex structures. The synthesis of an optimized molecule is ensured by following a path in the synthesis tree.

2.2 Deep Generative Model

DGM are widely employed for molecule optimization, utilizing various approaches such as VAE-based [GómezBombarelli et al., 2018], GAN-based [Guimaraes et al., 2017], encoder-decoder based [He et al., 2021], and flowbased models [Zang and Wang, 2020; Shi et al., 2020]. DGMs for molecule optimization treat the task as a generative problem, where the starting molecules serve as input, and the target molecules are the output. Typically, the target molecules are aimed to be more drug-like compared to the starting molecules. Modof [Chen et al., 2021] is a DGM designed for molecule optimization. It predicts the disconnection site and modification between the starting and target molecules.

2.3 HybridMethod

Hybrid methods include genetic algorithms (GA) method [Nigam et al., 2019], Markov Chain Monte Carlo (MCMC) method [Fu et al., 2021], and Bayesian optimization (BO) method [Moss et al., 2020]. The GA method [Nigam et al., 2019] incorporates an adaptive penalty mechanism based on a neural network to promote the exploratory behavior of GA, thus improving the diversity of generated molecules. The MCMC method [Fu et al., 2021] first pre-trains two graph neural networks for substructure-type and molecule topology predictions. Subsequently, the optimized molecules are sampled from the pre-trained molecule distribution using MCMC sampling with multiple constraints. The BO method [Moss et al., 2020] utilizes BO loops on string kernels to encode molecules into feature-rich spaces and employs GA to obtain the targeted molecules.

In this work, we present a novel molecule optimization method that treats molecule optimization as a machine translation problem. This approach offers contrastive learning without the requirement of domain-specific knowledge for generating negative and positive molecules, which helps address the "exposure bias" problem encountered in translationbased methods for molecule optimization.

3 Method

This section outlines the specifics of GPMO, which comprises two main steps: pre-training of encoder-decoder and conditional molecule optimization with contrastive learning. GPMO learns the grammar of SMILES and the relationships between molecules in the pre-training stage by masked language model (MLM) and SMILES translation task. Condition tokens are used for conditional molecule optimization. During conditional molecule optimization, positive and negative molecules are generated by gradient perturbation for contrastive learning.

Figure 1a illustrates that conditional molecule optimization involves utilizing both the condition tokens and the SMILES of the start molecule as inputs for the encoderdecoder model. The outputs of the encoder-decoder model are the SMILES of the target molecule. By adopting this approach, the model is capable of generating a novel molecule

̅

that adheres to the specified conditions while preserving the essential structural characteristics of the start molecule. The encoder-decoder model is trained using additional negative and positive molecules, which are produced through negative (Figure 1b) and positive perturbations (Figure 1c). Negative molecules share a similar SMILES string with real molecules but differ in their context, as modifying a single token can result in an invalid molecule (the negative molecule in Figure 1a). Compared to the real molecule, the positive molecules in Figure 1a vary with SMILES but share a similar chemical structure.

3.1 ProblemDefinition

MO can be formulated as a seq2seq problem, where the optimization from the start molecule to the target molecule is the translation between SMILES of molecules. This task is typically accomplished using an encoder-decoder architecture. For a given input molecule xi , the encoder converts it into a high-level representation. Given the latent representation of xi , the decoder autoregressively generates one token of prediction molecule at a time. The encoder-decoder process can be defined as Hi = f(xi ) and Si = g(Hi ), where Hi = [hi1 , hi2 , ..., hiL] and Si = [si1 , si2 , ..., siT ] represent the hidden states of the input molecule with a length of L and the output molecule with a length of T , respectively. f and g denote the encoder and decoder respectively. The auto-regressive decoding process during training with teacher forcing, denoted as pθ (yt ∣∣yi<t , xi ), can be gained by the decoder hidden state sit with a softmax layer as pθ(yi t |yi<t , xi ) = softmax(Wsit+b), where W and b are the weight and bias of the linear layer, respectively. Given a dataset (xi , yi )Ni=1 with N molecules, the seq2seq model is trained using a maximized likelihood estimate (MLE), which can be formulated as follows:

where log pθ (yi |xi ) represents the log likelihood of yi conditional on xi . LMLE is maximized to estimate the parameter of the encoder-decoder.

3.2 Condition Token

GPMO is an approach used to optimize a given start molecule to a target molecule that exhibits structural similarity while satisfying conditional constraints. In order to ensure the similarity of molecular structures, the start and target molecules used in GPMO are constrained to comply with MMP relation [Dalke et al., 2018]. Molecules with MMP relations share the same scaffold but differ in a single chemical transformation. Conditional molecule optimization is achieved by using condition tokens, which encode the changes of desired properties between the start and target molecules. By incorporating condition tokens, the optimization process can be guided to search for molecules that satisfy the desired properties encoded in the condition tokens. Three ADMET properties-LogD, Solubility, and Clint-are used as the properties to be optimized. LogD refers to the partition coefficient of all forms of a compound at a specific pH. Solubility is the degree to which a substance dissolves in water to make a solution. Clint is the overall hepatic intrinsic clearance.

Specifically, as the condition shown in Figure 1a, the condition tokens are included as part of the GPMO input. The property changes of LogD can be defined as: ∆(LogD) = LogD(yi ) - LogD(xi ). Then, ∆(LogD) is map into a code table: (-inf,-6.9], (-6.9,-6.7], ..., (6.9, inf). -6.9 and 6.9 are the minimum and maximum values of ∆(LogD). ∆(LogD) is mapped into the corresponding code. For example, the condition token of ∆(LogD) = 0.82 can be denoted as Delta_LogD(0.7,0.9]. In terms of Solubility, the code table of ∆(Solubility) can be defined as:

The condition token of solubility can be defined as Delta_Solubilityhigh_low, Delta_Solubilityno_change, and Delta_Solubilitylow_high. The condition token of Clint is similar to Solubility, but the threshold is 0.35. Besides property improvement where higher values are desirable, GPMO can optimize molecules by selectively improving one property while keeping other properties unchanged or lower. This is accomplished by incorporating the condition tokens that specify the desired property changes.

3.3 Positive and Negative Molecules Generation

To generate negative and positive molecules for contrastive learning without domain-specific knowledge, two kinds of gradient perturbation are used. The first kind is a small negative perturbation, which generates a negative molecule that is close to the real molecule. The second kind is a large positive perturbation, which involves two steps to generate a positive molecule that is far away from the real molecule.

Small Negative Perturbation for Negative Molecule

For negative perturbation d, the log likelihood with respect to d is pθ(yi t |yi<t , xi ; sit +d) = softmax(W(sit +d)+b). As d is a negative perturbation, the log-likelihood with respect to d should be minimized. Generally, a linear approximation can be defined as:

where g = ∇SiLMLE is the negative perturbation, which is the gradient of Si in the back-propagation of LMLE. ∥.∥2 is the l2 norm function and λ controls the weight of negative perturbation. As g is the gradient of Si , thus ˇSi is near Si in the embedding space and different from Si . The contrastive learning with negative molecule gained by negative perturbation can be defined as:

where zix = φ(Hi ) and ziy = φ(Si ). Ui = {Sj ∪ ˇSi }, Sj = {Sj : j ̸= i}, and z-y = φ(Ui ). φ(.) is a projection function with a linear layer and an average pooling operation. Sj is the random molecule in the batch which acts as the negative molecule. Minimizing of Lcont- (θ) will force the negative molecule ˇziy away from the real molecule.

Large Positive Perturbation for Positive Molecule

Two-step perturbation is used to generate a positive molecule that is far away from the real molecule. The first perturbation is the gradient with respect to contrastive loss, which maximizes the similarity between paired molecules and minimizes the similarity between unpaired molecules in the same batch. Then, the intermediate positive molecule is obtained based on the first step of positive perturbation. After that, KL loss is used to constrain distribution similarity between the intermediate positive molecule and the real molecule. Next, the second positive perturbation is the gradient based on KL loss. Finally, another positive molecule is obtained based on the second step of positive perturbation.

To gain the first gradient perturbation, the contrastive loss is used to train the model to learn the representations of the real molecule by contrasting the molecule pairs with the unpaired molecules. The unpaired molecules are the randomly sampled non-target output molecule in the same batch. As shown in Figure 1c, the start and target molecules are projected onto the latent embedding space. Then, contrastive learning is used to contrast the paired molecules with the unpaired molecules. We define this as:

where Lcont (θ) maximizes the representation similarity between molecule pairs and minimizes the similarity between unpaired molecules within the same batch. Lcont (θ) trains the model to learn the representations of the real molecule. The gradient perturbation of Si in the back-propagation of Lcont (θ) is a positive perturbation. Thus the first positive perturbation can be defined as f + = ∇SiLcont.

The first positive perturbation is used to generate the intermediate positive molecule. The hidden state of the intermediate positive molecule can be defined as:

where µ controls the weight of the first positive perturbation. As ¯Si is the intermediate positive molecule, the distribution of ¯Si after the softmax layer should be similar to the real molecule. The KL loss is used to minimize the distribution between intermediate positive molecule and the real molecule which can be defined as:

where pθ(¯yi , xi ) = softmax(W¯Sit + b). The distribution similarity of the positive molecule and the real molecule is achieved by minimizing KL loss.

As the KL loss maximizes the similarity between the intermediate positive and real molecules, the gradient perturbation obtained from KL loss can also be used as the second positive perturbation to generate the final positive molecule. The hidden state of the positive molecule gained by the second positive perturbation can be defined as:

where f ++ = ∇¯SiLKL. The molecule representation ˆSi is obtained through a two-step positive perturbation, resulting in a positive molecule that is significant far away from the real molecule in the embedding space. Then, the contrastive loss with the positive molecule can be written as:

where ˆziy = φ( ˆ Si ). Lcont+ is the contrastive loss, which trains the model to learn the representation of the real molecule by contrasting the hard distinguished positive molecule with a negative molecule.

Lcont- (θ) and Lcont+ (θ) contrast real molecule with generated negative and positive molecules, aiming to differentiate between them and align the generated positive molecules with the real ones. The real molecule acts as an anchor and its similarity with positive molecule is maximized, while its similarity with negative molecule is minimized (Figure 1a). Finally, the parameters of the encoder-decoder can be estimated by minimizing the following objective:

The minimize of L trains the seq2seq model using diverse negative and positive examples through contrastive learning, thereby mitigating the "exposure bias" problem.

3.4 Self-Supervised Pre-Training

To autoregressively decode molecules from input molecules and conditions, GPMO is pre-trained using self-supervised tasks such as MLM and SMILES translation. This pretraining enables GPMO to learn the grammar of SMILES and the relationship between molecules.

• MLM. Canonical SMILES are used in MLM. For the input of the MLM, 15% of the input tokens are randomly masked. However, due to the large number of carbon atoms in molecules, reconstructing only the masked tokens may lead to overfitting on carbon atoms. Therefore, the reconstruction of the entire Canonical SMILES is predicted instead. The MLM is employed to learn the grammar of SMILES.
• SMILES translation. Amajor challenge inMO is understanding the relationship between two molecules, which is not directly captured by MLM task. To address this,

the SMILES translation pre-training task is employed to train the GPMO model to learn the relationship between two molecules. In the SMILES translation, a many-toone translation approach is employed. It involves using SMILES with a randomly chosen root atom as the input and generating the corresponding Canonical SMILES as the output, where multiple random SMILES representations of a molecule are encoded into similar embeddings. This encourages the decoder to focus on improving its ability to accurately decode Canonical SMILES.

In the pre-training of GPMO, the MLM and SMILES translation tasks are selected randomly with a 50% chance.

4 Experiment

To showcase the effectiveness of GPMO, extensive experiments were conducted on benchmark tasks. GPMO outperforms deep learning techniques, showcasing its potential as a powerful tool for molecule optimization.

4.1 Experiment Setup

GPMO is trained in the pre-training stage using the dataset fromMOSES [Polykovskiy et al., 2020]. In the molecule optimization stage, GPMO utilizes a dataset from previous work [He et al., 2021], and the statistical information is presented in Table 1. The table displays the number of molecule pairs and the property ranges in the train, validation, and test sets. For property prediction, an ensemble random forest model trained on [He et al., 2021] is employed.

Baseline Methods

The following MO baselines are compared with GPMO:

• JTNN [Jin et al., 2018b] approaches MO by learning to translate molecule pairs.
• HierG2G [Jin et al., 2020] employs a hierarchical approach to encode molecule graphs and creates compounds by connecting structural motifs.
• Modof [Chen et al., 2021] conducts MO by predicting the removal, addition, or replacement at the disconnection site.
• Transformer [Vaswani et al., 2017] utilizes a seq2seq approach for MO, where the starting and target molecules serve as the input and output, respectively.
• GPMO-w/o-per is a variant of GPMO that does not use gradient perturbation-based contrastive learning.

• GPMO-w/o-con is a variant of GPMO that does not include the use of condition tokens in the input of the encoder.

Metric

We consider the following metrics for the evaluation of generated molecules.

• LogD: the proportion of generated molecules that satisfy the LogD constraint, where the absolute difference in LogD between the generated molecule and target molecule is no more than 0.4.
• Solubility: the proportion of code(Solubility(y′)- Solubility(x)) equal to code(Solubility(y)- Solubility(x)), where x is the start molecule, y′ is the generated molecule, and y is the target molecule. Code(.) refers to the code table defined in Eq. (2).
• Clint: the definition of Clint is similar to Solubility.
• All: the proportion of generated molecules that satisfy the constraints for LogD, Solubility, and Clint constraints simultaneously.
• Validity: the proportion of chemically valid molecules among all generated molecules.
• MMP: the proportion of generated molecules that exhibit an MMP relation with the starting molecules. The MMP relation is defined using the MMP query tool mmpdb [Dalke et al., 2018].
• Similarity: the Tanimoto similarity between the starting and generated molecules with Morgan fingerprint.

4.2 Multi-Property Optimization

As shown in Table 2, GPMO outperforms the state-of-the-art models on all seven metrics. From the results shown in Table 2, GPMO outperforms baselines on single-property performance metrics such as LogD, Solubility, and Clint. In terms of multi-property optimization, GPMO outperforms the best baseline method by a significant margin of 9.14% on the All metric. GPMO achieves the best performance on Validity due to it is pre-trained on unlabeled data, which enables it to learn the grammar on SMILES. Our result for the MMP and Similarity metrics indicate that the molecules generated from GPMO maintain the MMP relation and chemical similarity with the input molecules. Based on our analysis, we can conclude that GPMO has the ability to generate valid molecule with improved properties while retaining the starting molecule.

One key factor contributing to GPMO's impressive performance is the use of the condition tokens, which allow for the differentiation of distinct optimization objectives among pairs that share the same starting molecule. In contrast, JTNN, HierG2G, and Modof are unable to distinguish between different molecule pairs that have the same starting molecule. Consequently, the optimization of these molecule pairs will be an average of all molecule pairs that share the same starting molecule in the training set. Compared with Transformer, which also uses condition tokens, GPMO achieves better performance on Validity metric due to its pre-trained on unlabeled data, which enables it to learn the grammar of SMILES.

When compared with GPMO-w/o-per, GPMO lags behind on Similarity metric since it learns a broader range of valid or incorrect samples, resulting in less chemical similarity with a starting molecule. Our findings reveal that GPMO outperforms GPMO-w/o-per by an additional 3.38% on the All score, demonstrating that our model generates a greater number of desired molecules when compared to its counterpart. Given that the variance of GPMO-w/o-per is higher than GPMO, GPMO is more robust than GPMO-w/o-per. This is likely due to the contrastive learning approach employed by GPMO. When condition tokens are removed from GPMO, as in the case of GPMO-w/o-con shown in Table 2, the Validity rate remains high. However, the All rate drops down to 0.9356 due to the loss of ability to distinguish molecule pairs with the same starting molecule.

4.3 Model Performance on Smaller Datasets

Since experimentally obtained datasets are both expensive and rare, it is crucial to employ deep learning models that can handle small-scale datasets. In our study, we conducted experiments to evaluate the performance of GPMO in comparison to baseline methods on reduced versions of the train, validation, and test datasets, with only 50%, 20%, and 10% of the original dataset sizes, respectively.

The performance of each method is measured in terms of the All and Validity metrics, which respectively measure the method's ability to perform property optimization and generate valid molecules. From the results shown in Table 3, we observe that the All and Validity metrics decayed with the reduction of samples. Despite the performance decay observed with the reduction of training samples, our model still shows competitive results compared to the baseline methods on the small-scale dataset. More specifically, in the case of the 10% dataset, GPMO achieves an 88.77% All metric on propriety optimization, which is 9.03% higher than the best baseline method. Moreover, GPMO exhibits a 6.77% improvement on the Validity metric. These results highlight the benefiting from pre-training tasks, which enable GPMO to exhibit reduced reliance on labeled data.

4.4 Sensitivity Analysis

Table 4 provides an overview of the overall performance of GPMO under across various combinations hyper-parameters λ and µ. We set the weight of negative perturbation λ = [1, 3, 5] and the weight of positive perturbation µ = [1, 3, 5]. We observe that the variance in performance across different gradient perturbations combinations was no more than 0.14, indicating the robustness of GPMO. Our results indicate that the hyper-parameter combination of λ = 3 and µ = 3 achieved the best performance, suggesting that overly large or small perturbation are inefficient. This is likely due to the fact that small perturbations produce negative and positive samples that show less contrast with real samples, while large perturbation can generate artificial samples that make it difficult for the model to convergence.

4.5 Multi-StepMolecule Optimization

In multi-step MO, molecules are continuously optimized based on prior experience, much like the way chemists optimize molecules. To achieve this, we use the currently optimized molecule as the starting molecule in the next round of optimization with the same property constraint. We repeat this process five times to generate a set of molecules with the same property constraint.

Figure 2 presents the results of multi-step MO with five steps, under the property constraints {(0.3,0.5], no_change, no_change}, which require improving the LogD of the target molecule while maintaining unchanged Solubility and Clint properties. The results of multi-step molecule optimization show an increase in LogD from 2.7 to 3.9, while the Solubility and Clint remain unchanged. Moreover, the scaffold is preserved throughout the optimization process. In the multistep optimization setting, our model is capable of optimizing the molecule with desired properties while constraining the same scaffold with a starting molecule.

5 Conclusion

In this work, we propose GPMO, a gradient perturbationbased contrastive learning network for molecule optimization. Extensive experiments demonstrate that with proper design, GPMO is able to find molecules that satisfy multiple properties and its performance is comparable to advanced methods using deep learning. The reasons for GPMO's strong performance include the use of gradient perturbation contrastive learning, the pre-trained LMwhich reduces the need for labeled molecule pairs, the condition tokens on properties distinguish different molecule pairs with the same starting molecules, and the use of the MMP dataset to ensures scaffold similarity between the starting and generated molecules. GPMO can be readily adapted to real drug discovery projects, such as large-scale MMP analysis processes without chemist guidance. Furthermore, properties other than LogD, Solubility, and Clint, such as QED/SA and other ADMET properties, can also be included to search for more promising candidates.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China (62122025, U22A2037, 62102140), Hunan Provincial Natural Science Foundation of China (2021JJ10020), and the Science and Technology Innovation Program of Hunan Province (2022RC1100).