SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xiuyuan Hu $^{1,2*}$ , Guoqing Liu $^{2\dagger}$ , Yang Zhao $^{1}$ , Hao Zhang $^{1\dagger}$
$^{1}$ Department of Electronic Engineering, Tsinghua University
$^{2}$ Microsoft Research AI4Science
huxy22@mails.tsinghua.edu.cn, guoqingliu@microsoft.com,
zhao-yang@tsinghua.edu.cn, haozhang@tsinghua.edu.cn
De novo drug design is a pivotal issue in pharmacology and a new area of focus in AI for science research. A central challenge in this field is to generate molecules with specific properties while also producing a wide range of diverse candidates. Although advanced technologies such as transformer models and reinforcement learning have been applied in drug design, their potential has not been fully realized. Therefore, we propose MolRL-MGPT, a reinforcement learning algorithm with multiple GPT agents for drug molecular generation. To promote molecular diversity, we encourage the agents to collaborate in searching for desirable molecules in diverse directions. Our algorithm has shown promising results on the GuacaMol benchmark and exhibits efficacy in designing inhibitors against SARS-CoV-2 protein targets. The codes are available at: https://github.com/HXYfighter/MolRL-MGPT.
In recent years, significant strides have been made in computer-aided drug discovery (CADD) thanks to the development of various fields, including proteomics, genomics and deep learning [13, 66]. The conventional drug discovery process is typically time-consuming and financially demanding, with a low success rate. However, advanced machine learning techniques have the potential to reverse this predicament, greatly benefiting the economy and society's development [50, 58]. Currently, a major issue in the field of pharmacology is goal-directed de novo drug design, which involves generating new drug molecules with specific biochemical properties, such as designing compounds with high binding affinity to a designated protein target [15, 38]. Despite the numerous proposed machine learning algorithms for molecular generation, the chemical space is vast, and the relationship between molecular properties and structures is intricate, making it challenging to obtain satisfactory results in practical applications [42].
Molecular diversity is a critical concern in drug design because a diverse set of candidates can provide more choices for downstream screening and avoid drug resistance and unknown side effects [12, 41]. However, existing algorithms encounter challenges in designing diverse drug molecules. Many algorithms tend to generate sets of highly similar compounds, which is of little value for the subsequent drug development process.
Over the past several years, generative language models have made remarkable strides in natural language processing, vision, and audio. Among these models, the generative pre-trained transformer (GPT) is particularly notable, which has demonstrated impressive capabilities of language under
*Work was done while Xiuyuan Hu was a research intern at Microsoft Research.
†Corresponding author: Hao Zhang (1st), Guoqing Liu.
standing and generation [45, 46, 11]. In the field of chemistry, transformer-based language models pre-trained on the SMILES (simplified molecular input line entry system) representation of molecules have emerged. Through transfer learning, these models can be adapted to a range of tasks, including molecular property prediction, reaction prediction, and molecular optimization [26, 27, 4].
Reinforcement learning (RL) has emerged as a promising approach for de novo drug design [39, 68, 29]. The basic idea is to consider molecular property predictors (scoring functions) as rewards and train an RL agent to iteratively generate candidate compounds with increasingly high scores. These RL algorithms can explore the vast chemical space remarkably faster than human chemists, but they may have limitations in molecular diversity. Although multi-agent reinforcement learning (MARL) is a commonly used technique for promoting diversity in searching problems [34, 14], it has not yet been effectively utilized in the field of drug design.
To address the limitations of current approaches, we propose MolRL-MGPT (Molecular design using Reinforcement Learning with Multiple GPT agents), a novel MARL framework for de novo drug molecular design that utilizes GPT models as agents. Our approach treats molecular design as a cooperative Markov game, where multiple lightweight GPT agents collaborate to generate high-scoring molecules during the RL process. These agents share identical pre-trained parameters on a molecular SMILES dataset for initialization and have a common optimization objective for specific properties. To enhance the diversity of generated candidates, we incorporate an auxiliary loss function that encourages agents to explore in diverse directions. Our algorithm has demonstrated superior performance compared to various baselines on the GuacaMol benchmark. We also apply it to resolve the real-world problem of designing candidates against two SARS-CoV-2 protein targets, resulting in potentially desirable candidates with good binding affinity, drug-likeness, and synthetic accessibility. Moreover, we further validate the effectiveness of our design by comparative and ablation experiments on  $\mathrm{GNK3}\beta$ , JNK3 and QED maximization tasks.
Machine learning has become a formidable instrument in molecular generation with applications to de novo drug design, aiding scientists in identifying novel molecules that possess the desired properties for drug discovery. The molecular representation is the basis of molecular generation and optimization algorithms, which can be roughly divided into three categories: 1D string, 2D image, and 3D geometry [15].
SMILES is the most commonly used 1D representation of molecules, which employs strings of characters to encode a molecule. Although the SMILES of each molecule is not unique, there is only one canonical SMILES, and each SMILES corresponds to a maximum of one molecule. It is noteworthy that most SMILES strings are invalid; that is, their corresponding structures cannot exist in the real world. Some algorithms and techniques applied in natural language processing (NLP) have been adopted to generate molecules using SMILES representations, such as variational autoencoder (VAE) [22, 16], recurrent neural network (RNN) [49], generative adversarial network (GAN) [23] and Bayesian optimization (BO) [37].
On the other hand, 2D and 3D molecular representations are more intuitive and have become popular in molecular design in recent years. For 2D molecular graphs with vertices corresponding to atoms and edges corresponding to chemical bonds between atoms, techniques such as graph neural network (GNN), genetic algorithm (GA) and flow network for graph data have already been applied to drug development [1, 28, 30, 62, 35, 21]. For 3D geometries of compounds, which theoretically contain the most structural information of the molecules, although models including diffusion have been introduced to the chemistry field [20, 33], they cannot currently design candidates with desired biochemical properties as 1D/2D de novo drug design approaches.
RL-based drug design algorithms Reinforcement learning is the most popular technique for molecule generation. In molecular design, deep neural networks are usually employed as agents in studies that use RL for 1D string generation, which are then fine-tuned via customized reward functions [39, 9, 41, 59]. Likewise, when RL is applied to generating 2D molecular graphs, the states correspond to incomplete representations, and the actions involve adding substructures (such as atoms, bonds, and rings) to specific positions [65, 68, 29, 1, 63, 19]. What is more, recent studies
have incorporated 3D geometries into the RL process to consider the spatial properties of molecules during their generation [51, 52, 18].
Transformers for chemical language Transformer is a type of deep neural network architecture entirely based on attention mechanism and has been widely applied in the field of NLP [57]. It has demonstrated a better ability to process long text sequences and parallel computing capabilities than traditional architectures such as RNN. Some works have already applied transformers in the field of chemistry. For instance, ChemFormer [26], MolGPT [4] and [24] have focused on SMILES-based pre-trained transformer models. Meanwhile, PROTAC-RL [67], TamGent [61] and [44] have concentrated on transformer-based drug design against protein targets, while MCMG [59] has used transformer as a component to enhance the learning capability of the algorithm.
Diversity in drug design Previous research has proposed using reinforcement learning to generate diverse molecules for drug development [8, 59, 63, 41]. However, these studies have not adequately addressed this concern, and multi-agent reinforcement learning is yet to be effectively applied in de novo drug design.
3.1 Problem definition
The fundamental aspect of the molecular generation problem formulation is a scoring function  $s(x)$  of molecular properties, also known as an oracle, with the input of  $x$  being a molecule and the output being a real number. Molecular properties can include physical properties such as molecular weight and the number of aromatic rings, as well as chemical properties such as logP and drug-likeness. Additionally, the molecular properties that we are most concerned with in real-world drug design are often related to biological activity, with binding affinity to protein targets being the most common, which can be estimated using docking software.
Standardly, the scoring function  $s(\cdot)$  is typically constrained within the interval [0, 1] when applied to a valid molecular input, where a higher score corresponds to a better molecular property. Invalid molecules consistently receive a score of -1. Researchers may implement a transformation function  $t(\cdot)$  for each molecule property predictor  $p(x)$  to achieve uniformity with the standard form, and create a multi-objective scoring function through a weighted combination of various oracles.
$$
s (x) \in [ 0, 1 ] \cup \{- 1 \},
$$
$$
s (x) = t (p (x)) \text {o r} s (x) = \sum_ {i} w _ {i} \cdot t _ {i} \left(p _ {i} (x)\right), \quad \sum_ {i} w _ {i} = 1 \tag {1}
$$
The evaluation of a de novo molecule generation algorithm usually requires the assessment of a set of generated high-scoring molecules, and a common approach is to report the average score and diversity of the top-  $k$  scoring generated molecules, where  $k \in \mathbb{N}^{+}$  is given. A widely used metric of molecular diversity is internal diversity (IntDiv) [6]:
$$
\operatorname {I n t D i v} (A) := \frac {1}{| A | (| A | - 1)} \sum_ {(x, y) \in A \times A, x \neq y} d _ {T} (\mathcal {F} (x), \mathcal {F} (y)) \tag {2}
$$
where  $A$  is a set of compounds,  $d_{T}$  represents the Tanimoto distance [55], and  $\mathcal{F}(x)$  refers to the extended-connectivity fingerprint (ECFP) [47] of a molecule  $x$ .
In drug molecular design tasks, we primarily focus on evaluating the final set of generated molecular candidates. Generally, we pay little attention to the generation process, including factors like time and computing resource consumption, as these costs are relatively insignificant compared to the conventional drug discovery process.
Hence, we can represent the problem of designing novel drug candidates as a cooperative Markov game consisting of multiple generative model agents. At every iteration  $i$  ( $i = 1,2,\dots,s$ ), each agent  $k$  ( $k = 1,2,\dots,n$ ) generates  $m$  molecules (actions), and the scoring function acts as an environment providing rewards in the form of scores, which in turn are used to update the agents. The game's objective is primarily to maximize the average of the highest  $k$  scores of generated molecules and secondarily to improve the diversity of the candidates.
3.2 MolRL-MGPT algorithm
(b)
Figure 1: (a) Overview of our MolRL-MGPT algorithm. (b) The model architecture of the GPT prior model and agents.
As illustrated in Fig 1 (a), the MolRL-MGPT algorithm consists of iterations in which multiple agents are updated in a reinforcement learning way. Specifically, in each iteration, each agent first samples a set of SMILES strings with high likelihood and retrieves several high-scoring molecules in the past iterations from the molecular memory. Then, the scoring function is used to obtain the reward of each generated SMILES string. The GPT agents conduct this process in order. The loss functions are designed as follows.
Loss function In MolRL-MGPT, the reward function is defined as the scores predicted by the scoring function. Our primary objective is to increase the average scores of SMILES strings generated by each agent. Furthermore, to prevent the agents from disregarding the knowledge acquired during the pre-training phase, we impose a penalty on the deviation between the new policies and the prior model, which is also employed to initialize all the agents. The loss function of the 1st GPT agent is designed as follows:
$$
L _ {1} (x; \Theta_ {1}) = \left[ \log P (x) _ {\text {P r i o r}} - \log P (x) _ {\text {A g e n t} _ {1}} + \sigma_ {1} \cdot s (x) \right] ^ {2} \tag {3}
$$
where  $\Theta_{1}$  is the parameters of the 1st agent,  $x$  is a generated molecule,  $\sigma_{1}$  is a coefficient for controlling the term of scores, and  $P(x)_{\mathrm{model}}$  refers to the likelihood of generating  $x$  by model. It should be noted that typically  $P(x)_{\mathrm{Prior}} < P(x)_{\mathrm{Agent}}$ .
Additionally, we encourage the agents to explore different directions in the chemical space instead of conducting repetitive searches by introducing a term indicating the deviation between the current agent and previous ones. The loss function of the  $k$ th GPT agent is designed as follows:
$$
\begin{array}{l} L _ {k} (x; \Theta_ {k}) = L _ {1} (x; \Theta_ {k}) - \sigma_ {2} \sum_ {j = 1} ^ {k - 1} s (x) \cdot | \log P (x) _ {\mathrm {A g e n t} _ {k}} - \log P (x) _ {\mathrm {A g e n t} _ {j}} | \\ = \left[ \log P (x) _ {\text {P r i o r}} - \log P (x) _ {\text {A g e n t} _ {k}} + \sigma_ {1} \cdot s (x) \right] ^ {2} \tag {4} \\ - \sigma_ {2} \sum_ {j = 1} ^ {k - 1} s (x) \cdot | \log P (x) _ {\mathrm {A g e n} t _ {k}} - \log P (x) _ {\mathrm {A g e n} t _ {j}} | \\ \end{array}
$$
where  $k = 1,2,\ldots n$ ,  $\Theta_{k}$  is the parameters of the  $k$ th agent, and  $\sigma_{2}$  is a coefficient for the encouragement term.
The selection of coefficients  $\sigma_{1},\sigma_{2}$  can be adapted according to the specific task at hand. Moreover, we propose to implement a decreasing schedule for  $\sigma_{1}$  such that  $\sigma_{1}$  decreases as  $s(x)$  increases during the RL process. This is because we expect the increase in scores to correspond to a decrease in loss.
Molecular Memory MolRL-MGPT utilizes a memory consisting of high-scoring molecules with a maximum size of 1000. The memory is updated with every SMILES string sampled, and compounds stored in the memory are sorted based on their scores.
3.3 Implementations
Pre-training on chemical language Instead of the commonly used recurrent neural network (RNN) architecture for generating chemical language in SMILES strings, we adopt a transformer-based architecture. As shown in Fig 1 (b), our model is a mini version of the GPT model with 8 layers of transformer blocks, each consisting of 8 attention heads. The embedding size is 256, and the maximum length of SMILES strings is set to 128. The pre-training approach of our model involved learning molecular structure patterns. Our prior model has only around 6.4M parameters, much less than the GPT-2 model [46].
Two versions of the prior model are trained using unsupervised learning on the ChEMBL [36] and ZINC-100M [53] datasets, respectively. The data are preprocessed by removing molecules with ionized structures and SMILES strings longer than 100 characters, which are not typical for small molecule drugs. Additionally, we use SMILES randomization for data augmentation, which has been proven helpful in enhancing generating capacity [3]. Ten training epochs are conducted with a batch size of 2048 and a maximum learning rate of 0.001, using a learning rate schedule featuring warm-up and cosine decay. The training on ChEMBL (roughly 2 million SMILES) takes around 5 hours using a single NVIDIA A100 GPU.
The pre-trained SMILES models should be objectively evaluated by the valid ratio of the generated molecules, representing the proportion of valid molecules present in the generated SMILES. Generally, this value can exceed  $90\%$  [43]. Our pre-trained transformer model has achieved a valid ratio of  $98\%$  in generating molecules, indicating its success in capturing the inherent grammar and rules of molecular SMILES strings.
Experience replay Experience replay is a widely employed technique in deep reinforcement learning. The agent or learner records its experience in a memory buffer and randomly samples past experiences, enabling it to learn from previous experiences with reduced correlation between consecutive training samples. This technique helps agents avoid overfitting, generalize to unseen situations, and achieve a more stable and efficient learning process [32]. In MolRL-MGPT, we replay the "successful" experiences by randomly sampling 5 molecules from the 25 highest-scoring molecules in the molecular memory and computing the loss on these molecules together with SMILES strings generated in the current iteration.
*Similarity penalization Some previous works for de novo drug design add a penalizing term for high similarity of identical skeletons between new molecules and previously found high-scoring molecules in order to encourage the agent to search in unexplored space rather than repeatedly finding the same or similar compounds [9, 60]. However, our experiments indicate that this trick does not work with our algorithm.
To demonstrate the performance of MolRL-MGPT, we conduct three groups of experiments. Firstly, we run our algorithm on the public de novo molecular design benchmark, GuacaMol, and compare it with existing advanced methods. Secondly, we apply MolRL-MGPT to the design of inhibitors against SARS-CoV-2 protein targets, which is a current real-world challenge for human beings. Thirdly, we conduct comparative and ablation experiments on GNK3 $\beta$ , JNK3 and QED maximization tasks to validate the effectiveness of modules in our design.
4.1 GuacaMol benchmark
GuacaMol [10] is a widely recognized open-source evaluation framework for de novo molecular design algorithms. It contains 20 goal-directed tasks mimicking the drug discovery objectives, corresponding to 20 standard scoring functions described in Sec 3.1. These tasks cover commonly used objectives in drug design, such as structural, physicochemical, and biochemical properties.
For comparison, we select the following baselines: (1) SMILES GA [64], a genetic algorithm based on SMILES; (2) SMILES LSTM [49], an LSTM network generating SMILES strings autoregressively, combined with a hill-climb algorithm for optimization; (3) Graph GA [28], a genetic algorithm with crossovers and mutations performed on molecular graphs; (4) Reinvent [39], a deep reinforcement learning framework for training an RNN model generating SMILES; (5) GEGL [1], a genetic expert-guided learning framework for training an RNN for molecular generation. Results of some other baselines are shown in the Appendix.
The prior model for MolRL-MGPT was pre-trained on the official GuacaMol dataset, which is a subset of ChEMBL. The hyper-parameters are set such that we run each tasks for 5000 tasks (break if the score has achieved 1.000) with 4 GPT agents, and the batch size (number of sampled SMILES strings) of each agent is 256. The values of coefficients are set to  $\sigma_{1} = 1000$  with a linear decreasing schedule, and  $\sigma_{2} = 0.1$ . The entire set of 20 tasks tasks less than 400 hours to complete when run on a single NVIDIA A100 GPU.
Table 1: Scores of MolRL-MGPT and other baselines on the GuacaMol benchmark.
Tasks	SMILES GA	SMILES LSTM	Graph GA	Reinvent	GEGL	MolRL-MGPT
1. Celecoxib rediscovery	0.732	1.000	1.000	1.000	1.000	1.000
2. Troglitazone rediscovery	0.515	1.000	1.000	1.000	0.552	1.000
3. Thiothixene rediscovery	0.598	1.000	1.000	1.000	1.000	1.000
4. Aripiprazole similarity	0.834	1.000	1.000	1.000	1.000	1.000
5. Albuterol similarity	0.907	1.000	1.000	1.000	1.000	1.000
6. Mestranol similarity	0.790	1.000	1.000	1.000	1.000	1.000
7. C11H24	0.829	0.993	0.971	0.999	1.000	1.000
8. C9H10N2O2PF2Cl	0.889	0.879	0.982	0.877	1.000	0.939
9. Median molecules 1	0.334	0.438	0.406	0.434	0.455	0.449
10. Median molecules 2	0.380	0.422	0.432	0.395	0.437	0.422
11. Osimertinib MPO	0.886	0.907	0.953	0.889	1.000	0.977
12. Fexofenadine MPO	0.931	0.959	0.998	1.000	1.000	1.000
13. Ranolazine MPO	0.881	0.855	0.920	0.895	0.933	0.939
14. Perindopril MPO	0.661	0.808	0.792	0.764	0.833	0.810
15. Amlodipine MPO	0.722	0.894	0.894	0.888	0.905	0.906
16. Sitagliptin MPO	0.689	0.545	0.891	0.539	0.749	0.823
17. Zaleplon MPO	0.413	0.669	0.754	0.590	0.763	0.790
18. Valsartan SMARTS	0.552	0.978	0.990	0.095	1.000	0.997
19. deco hop	0.970	0.996	1.000	0.994	1.000	1.000
20. scaffold hop	0.885	0.998	1.000	0.990	1.000	1.000
Total	14.396	17.340	17.983	16.350	17.627	18.052
As shown in Table 1, the MolRL-MGPT algorithm outperforms baselines in 13 molecular design tasks in the GuacaMol benchmark, and its total score of 20 tasks also ranks first. These results indicate that MolRL-MGPT excels in general scenarios of de novo drug design.
4.2 Designing inhibitors against SARS-CoV-2 targets
Molecular docking is a computational method used to predict the binding modes of small molecules to a protein target. It involves predicting the spatial orientation and binding affinity of the small molecule in the active site of the protein. This information is useful in drug discovery since it enables identifying potential drug candidates and understanding how they interact with their targets. Autodock Vina [56] is currently the most widely used molecular docking software. However, we opt to use Quick Vina 2 [2], a novel and more efficient alternative.
The quantitative binding affinity is named docking score, which is calculated based on the energies of the interaction between the ligand and the receptor, and a lower docking score indicates a more stable and, therefore, more likely binding pose. Typically, docking scores are negative, and desirable docking scores range from -10 to -14 kcal/mol. Therefore, we use a reverse sigmoid function as the transformation function of the docking score:
$$
t _ {\text {d o c k i n g}} (p) = \frac {1}{1 + 1 0 ^ {0 . 6 2 5 \cdot (p + 1 0)}} \tag {5}
$$
SARS-CoV-2 (Severe Acute Respiratory Syndrome Coronavirus 2), commonly referred to as the novel coronavirus, is a respiratory virus which, in late 2019, emerged as a global pandemic leading to the COVID-19 disease that mainly targets the respiratory system. This disease is a severe public health concern that continues to pose a crisis worldwide and requires a comprehensive response to mitigate its spread and negative effects. Therefore, we apply our algorithm to the design of inhibitors against protein targets of SARS-CoV-2, which is a significant real-world issue. Following [48], we select two targets: PLPro_7JIR $^3$  and RdRp_6YYT $^4$ .
Besides docking scores, we also consider two additional oracles often employed in practical drug design: (1) QED (Quantitative Estimate of Drug-likeness), which quantifies the drug-likeness of a molecule based on the concept of desirability of eight molecular properties [7], ranging in [0, 1]; (2) SA (Synthetic Accessibility), which incorporates fragment contributions and a complexity penalty [17], ranging in [1, 10]. Their transformation functions are linear:
$$
t _ {\mathrm {Q E D}} (p) = p, \quad t _ {\mathrm {S A}} (p) = \frac {1 0 - p}{9} \tag {6}
$$
Moreover, the scoring function for designing inhibitors against SARS-CoV-2 protein targets is a linear combination of docking scores, QED scores and SA scores:
$$
s _ {\text {t o t a l}} (x) = 0. 8 \cdot s _ {\text {d o c k i n g}} (x) + 0. 1 \cdot s _ {\text {Q E D}} (x) + 0. 1 \cdot s _ {\text {S A}} (x) \tag {7}
$$