ChemTS: An Efficient Python Library for de novo Molecular Generation

Xiufeng Yang, Jinzhe Zhang, Kazuki Yoshizoe, Kei Terayama, Koji Tsuda

Introduction

In modern society, a variety of organic molecules are used as important materials such as solar cells , organic light-emitting diodes (OLEDs) , conductors , sensors and ferroelectrics . At the highest level of abstraction, design of organic molecules is formulated as a combinatorial optimization problem to find the best solutions in a vast chemical space. Most computer-aided methods for molecular design build a molecule by a combination of predefined fragments (e.g.,). Recently, Ikebata et al succeeded de novo molecular design using an engineered language model of SMILES representation of molecules . It is increasingly evident, however, that engineered models often perform worse than neural networks in text and image generation . Gomez-Bombarelli et al. were the first to employ a neural network called variational autoencoder (VAE) to generate molecules. Later Kusner et al. enhanced it to grammar variational autoencoder (GVAE) . SMILES strings created by VAEs are mostly invalid (i.e., they do not translate to chemical structures), so generation steps have to be repeated many times to obtain a molecule. Segler et al. showed that a recurrent neural network (RNN) using long short term memory (LSTM) achieves a high probablity of valid SMILES generation. In their algorithm, a large number of candidates are generated randomly and a black-box optimization algorithm is employed to choose high-scoring molecules. It is required to generate a very large number of candidates to ensure that desirable molecules are included in the candidate set. Optimization in a too large candidate space can be inhibitively slow.

In this paper, we present a novel python library ChemTS to offer material scientists a versatile tool of de novo molecular design. The space of SMILES strings is represented as a search tree where the $i$ -th level corresponds to the $i$ -th symbol. A path from the root to a terminal node corresponds to a complete SMILES string. Initially, only the root node exists and the search tree is gradually generated by Monte Carlo tree search (MCTS) . MCTS is a randomized best-first search method that showed exceptional performance in computer Go . Recently, it has been successfully applied to alloy design . MCTS constructs only a shallow tree and downstream paths are generated by a rollout procedure. In ChemTS, an RNN trained by a large database of SMILES strings is used as the rollout procedure. In a benchmarking experiment, ChemTS showed better efficiency in comparison to VAEs, creating about 40 molecules per minute. As a result, high scoring molecules were generated within several hours.

Method

ChemTS requires a database of SMILES strings and a reward function $r(S)$ where $S=\{s_{1},\ldots,s_{T}\}$ is an input SMILES string. Our definition of SMILES strings contains the following symbols representing atoms, bonds, ring numbers and branches: $s_{t}\in$ {C, c, o, O, N, F, [C@@H], n, -, S,Cl, [O-],[C@H], [NH+],[C@], s, Br, [nH], [NH3+], [NH2+], [C@@], [N+], [nH+], [S@], [N-], [n+],[S@@], [S-], I, [n-], P, [OH+],[NH-], [P@@H], [P@@], [PH2], [P@], [P+], [S+],[o+], [CH2-], [CH-], [SH+], [O+], [s+], [PH+], [PH], [S@@+], /,=, #, 1,2,3,4,5,6,7,8,9,(, ),}. In addition, we have a terminal symbol $. The reward function involves first principle or semi-empirical calculations and describes the quality of the molecule described by$ S $. If$ S $does not correspond to a valid molecule,$ r(S) $is set to an exceptionally small value. We employ rdkit (www.rdkit.org) to check if$ S $is valid or not. Before starting the search, an RNN is trained by the database and we obtain the conditional probability$ P(s_{t}|s_{1},...,s_{t-1})$ as a result. The architecture of our RNN is similar to that in and will be detailed in Section 2.1.

MCTS creates a search tree, where each node corresponds to one symbol. Nodes with the terminal symbol are called terminal nodes. Starting with the root node, the search tree grows gradually by repeating the four steps, Selection, Expansion, Simulation and Backpropagation (Figure 1). Each intermediate node has a UCB score that evaluates the merit of the node . The distinct feature of MCTS is the use of rollout in the simulation step. Whenever a new node is added, paths from the node to terminal nodes are built by a random process. In computer games, it is known that uniformly random rollout does not perform well, and designing a better rollout procedure based on available knowledge is essential in achieving high performance . Our idea is to employ a trained RNN for rollout. A node at level $t-1$ has a partial SMILES string $s_{1},\ldots,s_{t-1}$ corresponding to the path from the root to the node. Given the partial string, RNN allows us to compute the distribution of the next letter $s_{t}$ . Sampling from the distribution, the string is elongated by one. Elongation by RNN is repeated until the terminal symbol occurs. After elongation is done, the reward of the generated string is computed. In the backpropagation step, the reward is propagated backwards and the UCB scores of traversed nodes are updated. See for details about MCTS.

Our recurrent neural network (RNN) has a non-deterministic output: an input string $s_{1},\ldots,s_{T}$ is mapped to probability distributions of output symbols $P(y_{1}),\ldots,P(y_{T})$ . The RNN represents the function ${\boldsymbol{h}}_{t}=f({\boldsymbol{h}}_{t-1},{\boldsymbol{x}}_{t})$ , where ${\boldsymbol{h}}_{t}\in\Re^{512}$ is a hidden state at position $t$ and ${\boldsymbol{x}}_{t}\in\Re^{64}$ is the one-hot coded vector of input symbol $s_{t}$ . The function $f$ is implemented by two stacked gated recurrent units (GRUs) , each with 256 dimensional hidden states. The input vector ${\boldsymbol{x}}_{t}$ is fed to the lower GRU, and the hidden state of the lower GRU is fed to the upper GRU. The distribution of output symbol is computed as $P(y_{t}=j)=g_{j}({\boldsymbol{h}}_{t})$ , where $g_{j}$ is a softmax activation function depending only on the hidden state of the upper GRU.

Given $N$ strings in the training set, we train the network such that it outputs a right-shifted version of the input. Denote by ${\boldsymbol{x}}_{it}$ the one-hot coded vector of the $t$ -th symbol in the $i$ -th training string. The parameters in the network ${\boldsymbol{\theta}}$ is trained to minimize the following loss function,

where $D$ denotes the relative entropy. Our RNN was implemented using Keras library (github.com/fchollet/keras), and trained with ADAM using a batch size of 256. After the training is finished, one can compute $P(y_{t})$ from $s_{1},\ldots,s_{t-1}$ . It allows us to perform rollout by sampling the next symbol repeatedly.

Experiments

Following , we generate molecules that jointly optimize the octanol-water partition coefficient logP and the other two properties: synthetic accessibility and ring penalty that penalizes unrealistically large rings. The score of molecule $S$ is described as

The reward function of ChemTS is defined as

ChemTS was compared with two existing methods CVAE and GVAE based on variational autoencoders. Their implementation is available at https://github.com/mkusner/grammarVAE. Both methods perform molecular generation by Bayesian optimization (BO) in a latent space of VAE. RNN, CVAE and GVAE were trained with approximately 250,000 molecules in ZINC database . All methods were trained for 100 epochs. Training took 3.8, 9.4 and 33.5 hours respectively, on a CentOS 6.7 server with a GeForce GTX Titan X GPU. To evaluate the efficiency of MCTS, we prepared two alternative methods using RNN. One is simple random sampling using RNN, where the first symbol is made randomly and it is elongated until the terminal symbol occurs. The other is the combination of RNN and Bayesian optimization , where 4,000 molecules are made a priori and Bayesian optimization is applied to find the best scoring molecule.

As shown in Table 1, effectiveness of each method is quantified by the maximum score $J$ among all generated molecules at 2,4,6 and 8 hours and the speed of molecules generation (i.e., the number of generated molecules per minute). VAE methods performed substantially slower than RNN-based methods, which reflects the low probability of generating valid SMILES strings. ChemTS performed best in finding high scoring molecules, while the speed of molecular generation (40.89 molecules per minute) was only slightly worse than random generation by RNN (41.33 molecules per minute). The combination of RNN and BO could not find high scoring molecules. Preparing more candidate molecules may improve the best score, but it would further slow down the molecular generation. In general, it is difficult to design a correct reward function when there are multiple objectives. So, it is important to generate many good molecules in a given time frame to allow the user to browse and select favorite molecules afterwards. See Figure 2 for the best molecules generated by ChemTS.

Conclusion

In this paper, we presented a new python package for molecular generation. It will be further extended to include more sophisticated tree search methods and neural networks. Use of additional packages for computational physics such as pymatgen allows the users to implement their own reward function easily. We look forward to see ChemTS as a part of the open-source ecosystem for organic materials development.

Acknowledgement(s)

We would like to thank Hou Zhufeng, Diptesh Das, Masato Sumita and Thaer M. Dieb for their fruitful discussions.

Disclosure statement

Funding

This work was supported by the “Materials research by Information Integration” Initiative (MI2I) project and CREST Grant No. JPMJCR1502 from Japan Science and Technology Agency (JST). It was also supported by Grant-in-Aid for Scientific Research on Innovative Areas “Nano Informatics” (Grant No. 25106005) from the Japan Society for the Promotion of Science (JSPS). In addition, it was supported by MEXT as “Priority Issue on Post-K computer” (Building Innovative Drug Discovery Infrastructure Through Functional Control of Biomolecular Systems).

Notes on contributors

K. Tsuda proposed the research idea, supported experiments design and help draft the manuscript. K. Yoshizoe supported ChemTS design. J. Zhang and K. Terayama helped analyze the experimental data. X. Yang designed and implemented ChemTS, analyzed the data, and compiled the manuscript. All of the authors have read and approved the final manuscript.