Practical Massively Parallel Monte-Carlo Tree Search Applied to Molecular Design

1 INTRODUCTION . A survey paper on MCTS , published in 2012 , has cited 240 papers , including many game and non-game applications ( Browne et al. , 2012 ) . Since the invention of Upper Confidence bound applied to Trees ( UCT ) ( Kocsis & Szepesvári , 2006 ) ( the most representative MCTS algorithm ) in 2006 , MCTS has shown remarkable performance in various problems . Recently , the successful combination with Deep Neural Networks ( DNN ) in computer Go by AlphaGo ( Silver et al. , 2016 ) has brought MCTS into the spotlight . Combining MCTS and DNN is becoming one of the standard tools for solving decision making or combinatorial optimization problems . Therefore , there is a significant demand for parallel MCTS . However , in contrast to the enormous computing resources invested in training DNN models in many recent studies , MCTS is rarely parallelized at large scale . Parallelizing MCTS/UCT is notoriously challenging . For example , in UCT , the algorithm follows four steps , selection-expansion-rollout-backpropagation . Non-parallel vanilla UCT updates ( backpropagates ) the values in the tree nodes after each rollout . The behavior of the subsequent selection steps depends on the results of the previous rollouts-backpropagation . Therefore , there is no apparent parallelism in the algorithm . ∗This work was done while all the authors were at RIKEN . 1Code is available at https : //github.com/yoshizoe/mp-chemts Using virtual-loss technique ( explained in section 2.3 ) , MCTS has been efficiently parallelized on shared-memory single machine environment , where the number of CPU cores are limited ( Chaslot et al. , 2008 ; Enzenberger & Müller , 2010 ; Segal , 2010 ) . However , there is limited research on large-scale parallel MCTS using distributed memory environments . Only two approaches scale efficiently on distributed memory environment , but these were only validated in terms of the number of rollouts and the actual improvement is not validated ( Yoshizoe et al. , 2011 ; Graf et al. , 2011 ) . Recently , the combination of ( non-parallel ) MCTS and DNN has been applied to molecular design problems , which aims to find new chemical compounds with desired properties ( Yang et al. , 2017 ; Sumita et al. , 2018 ) , utilizing the ability of MCTS to solve single-agent problems . In general , designing novel molecules can be formulated as a combinatorial optimization or planning problem to find the optimal solutions in vast chemical space ( of 1023 to 1060 , Polishchuk et al . ( 2013 ) ) and can be tackled with the combinations of deep generative models and search ( Kusner et al. , 2017 ; GómezBombarelli et al. , 2018 ; Jin et al. , 2018 ; Popova et al. , 2018 ; 2019 ; Yang et al. , 2017 ) . However , there are no previous studies about massively parallel MCTS for molecular design . In this paper , we propose a novel distributed parallel MCTS and apply it to the molecule design problem . This is the first work to explore viability of distributed parallel MCTS in molecular design . Our experimental results show that a simple RNN model combined with massively parallel MCTS outperforms existing work using more complex models combined with Bayesian Optimization or Reinforcement Learning ( other than UCT ) . 2 BACKGROUND . 2.1 ( NON-PARALLEL ) MCTS . In 2006 , Kocsis and Szepesvári proposed UCT based on a Multi-Armed Bandit algorithm UCB1 ( Auer et al. , 2002 ) , which is the first MCTS algorithm having a proof of convergence to the optimal solution . It has shown good performance for many problems , including the game of Go ( Gelly et al. , 2006 ) . One round of UCT consists of four steps , as shown in Fig . 1 . It repeats the rounds for a given number of times or until a specified time has elapsed . Selection : The algorithm starts from the root node and selects the child with the highest UCB1 value ( Fig . 2 left ) until it reaches a leaf node . For each child i , vi is the number of visits , wi is the cumulative reward , and V is the number of visits at the parent node . Exploration constant C controls the behavior of UCT : the smaller , the more selective ; the greater , the more explorative . Expansion : If the number of visits exceeds a given threshold , expand the leaf node ( add the children of the leaf to the tree ) and select one of the new children for simulation . Do nothing otherwise . Simulation : UCT then evaluates the node by a simulation . This step is often called playout or rollout . A simple example of rollout is to go down the tree by selecting a random child at each node until it reaches a terminal node and returns the value at the terminal node as the reward r ( win or loss for games ) . Replacing rollout with a DNN based evaluation is becoming more popular following the success of AlphaGo . Backpropagation : UCT finally traverses the path all the way back to the root , and updates the values of the nodes in the path ( wi = wi + r , vi = vi + 1 ) . 2.2 SOLVING MOLECULAR DESIGN USING MCTS . The Simplified Molecular-Input Line-Entry System ( SMILES ) ( Weininger , 1988 ) is a standard notation in chemistry for describing molecules using ASCII strings , defined by a grammar . SMILES uses ASCII symbols to denote atoms , bonds , or structural information such as rings . For example , SMILES for Carbon dioxide is `` O=C=O '' where `` = '' means a double bond , and for Benzene , it is `` C1=CC=CC=C1 '' or `` c1ccccc1 '' which forms a ring by connecting the two `` C1 '' s. The search space for molecular design can be defined as a tree based on the SMILES , where the root node is the starting symbol ( denoted as `` & '' , usually not shown ) and each path from the root to a depth d node represents a left-to-right enumeration of the first d symbols in the ( partial ) SMILES string ( see Fig . 1 . The subscripts shows the depth. ) . Yang et al . ( 2017 ) were the first to apply non-parallel MCTS to this problem using the AlphaGo-like approach , which combines MCTS with DNN and a computational chemistry simulator . They trained an RNN model with a chemical compound database , and the RNN predicts the next symbol of a partial SMILES string which is given as the input . The model is used for expansion and rollout . Expansion for Molecular Design : The node ( =2 ) in Fig . 1 denotes SMILES strings starting with `` O= '' . The RNN receives `` O= '' as input and outputs the probability of the next SMILES symbols . Instead of simply adding all symbols as child nodes , the low probability branches are pruned based on the output of the model . Simulation for Molecular Design : Fig . 1 illustrates the simulation step for molecular design . Node ( C3 ) denotes SMILES strings starting with `` O=C '' . Firstly , a rollout generates a complete SMILES string , a new candidate molecule , by repeatedly sampling the next symbol using the model , until the model outputs the terminal symbol . Then , a computational chemistry simulator receives the SMILES string and calculates the target chemical property for the molecule , to determine the reward . 2.3 PARALLEL MCTS . Non-parallel MCTS finds the most promising leaf one at a time in the selection step . In parallel MCTS , multiple workers must find the promising leaves ( to launch simulations from ) in parallel , hence , should find leaves speculatively without knowing the latest results . Fig . 2 shows the example of a three workers case ( a worker means a thread or a process , which runs on either a CPU or GPU ) . If the workers ( shown in green , red , blue ) follow the UCB1 formula , all end up at the same leaf node ( Fig . 2- ( a ) ) , and the parallelization fails . The Virtual loss ( Chaslot et al. , 2008 ) let the workers find the promising but different leaves . The original UCB1 is modified to UCBvl shown in Fig . 2- ( b ) , where ti is the number of workers currently searching in the subtree of the child node i , and T is the sum of ti of the children . It is called virtual loss because it assumes that the current ongoing simulations will obtain 0 reward . With this modification , UCBvl value is penalized based on the number of workers , which makes the subsequent workers avoid the subtrees already being explored ( see Fig . 2- ( b ) ) . Parallel UCT with virtual loss was proved to improve the strength of shared-memory ( using up to 16 cores ) Go programs ( Chaslot et al. , 2008 ; Enzenberger & Müller , 2010 ) . Segal ’ s experiment on an emulated multithreaded machine that assumes no communication overhead shows speedup on up to 512 threads . However , his experiments on real distributed machines did not provide significant speedup beyond 32 cores ( Segal , 2010 ) . 2.4 HASH-DRIVEN PARALLEL SEARCH AND UNIFORM LOAD BALANCING . When many workers perform a search in parallel , it is essential to distribute the search space as evenly as possible . If there is an imbalance in the workload , overall performance will be less efficient because only some workers will continue the computation while the rest remain idle . If all the search space is explicitly given in advance , this is a trivial matter . However , in case of problems such as games and combinatorial optimization , where the search space is generated while the search is performed , it is difficult to distribute the search space evenly . Hash-driven ( HD ) parallel search is one of the methods for resolving this difficulty ; Evett et al . ( 1995 ) ( PRA * ) and Romein et al . ( 1999 ) ( TDS : Transposition-table Driven Scheduling ) developed this method independently and applied to parallel Iterative Deepening A * search . Kishimoto et al . ( 2013 ) later applied it to parallel A * search . HD parallel search requires a hash function that defines the hash key of each node . Each time the algorithm creates a node during the search , it assigns the node to a specific worker , the home worker ( or the home processor ) , based on the hash function , and achieves a near-equal load balancing if the hash function is sufficiently random ( such as Zobrist Hashing ( Zobrist , 1970 ) ) . Figure 3 illustrates this method . The hash function randomly divides the tree into four partitions ( shown in four colors ) , and each worker holds the nodes in the assigned partition in its hash table . A distinct drawback of this method is that it requires communication between workers almost every time the algorithm traverses a branch ( unless the two nodes are on the same worker ) . The key to efficient parallel speedup is the trade-off between uniform load balancing and frequent communication . The experimental results for IDA * and A * ( Evett et al. , 1995 ; Romein et al. , 1999 ; Kishimoto et al. , 2013 ) prove its efficiency for these search algorithms . However , a straightforward application of HD parallel search to UCT is not efficient enough as described below .

Summary: This paper aims to provide a scalable and competitive molecular design technique based on MCTS. Machine learning guided molecular design is a popular field and many solutions have been developed in the recent years. There is appeal, due to the high dimensionality of search space, to resort to MCTS based approaches that leverage stochastic pruning of the search space as an effective way of making the problem tractable.