Fragment-Based Sequential Translation for Molecular Optimization

1 INTRODUCTION . Molecular optimization is a challenging task that is pivotal to drug discovery applications . Part of the challenge stems from the difficulty of exploration in the molecular space : not only are there physical constraints on molecules ( molecular strings/graphs have to obey specific chemical principles ) , molecular property landscapes are also very complex and difficult to characterize : small changes in the molecular space can lead to large deviations in the property space . Recent fragment-based molecular generative models have shown significant empirical advantages ( Jin et al. , 2019a ; Podda et al. , 2020 ; Xie et al. , 2021 ) over atom-by-atom generative models in molecular optimization . However , they generally operate over a fixed set of fragments which limit the generative capabilities of these models . Shifting away from previous frameworks , we learn a distribution of molecular fragments using vector-quantized variational autoencoders ( VQ-VAE ) ( van den Oord et al. , 2017 ) . Our method builds molecular graphs through the addition and deletion of molecular fragments from the learned distributional fragment vocabulary , enabling the generative model to span a much larger chemical space than models with a fixed fragment vocabulary . Considering atomic edits as primitive actions , the idea of using fragments can be thought of as options ( Sutton et al. , 1999 ; Stolle & Precup , 2002 ) as a temporal abstraction to simplify the search problem . We further introduce a novel sequential translation scheme designed for fragment-based molecular optimization . We start the molecular search by translating from known active molecules and store the discovered molecules as new potential initialization states for subsequent searches ; we incorporate a delete action in our model , enabling our method to backtrack to good molecular states . Previous works optimize molecules either by generating from scratch or a single translation from known molecules , which is inefficient in finding high-quality molecules and often discovering molecules lacking novelty/diversity . Our proposed framework addresses these deficiencies since our method is ( 1 ) very efficient in finding molecules that satisfy property constraints as the model stay close to the high-property-score chemical manifold ; and ( 2 ) able to produce highly novel molecules with our flexible learned fragment vocabulary and a sequence of fragment-based editing . Combining the advantage of a distributional fragment vocabulary and the sequential translation scheme , we propose Fragment-based Sequential Translation ( FaST ) , which is realized by an RL policy that proposes fragment addition/deletion to a given molecule . Our proposed method can generate molecules under various objectives such as property constraints , novelty constraints , and diversity constraints . The main contribution of this paper includes : 1 . We demonstrate a way to learn distributional molecular fragment vocabulary through a VQ-VAE and the effectiveness of the learned vocabulary in molecular optimization . 2 . We propose a novel molecular search scheme of sequential translation , which gradually improves the quality and novelty of generation through backtracking and a stored frontier . 3 . We implement a RL policy combining the fragment vocabulary and the sequential translation scheme that significantly outperforms state-of-the-art methods in benchmark single/multi-objective molecular optimization tasks . 2 RELATED WORK . Molecular generation and optimization . Early works on molecular optimization build on generative models on both SMILES/SELFIES string ( Gómez-Bombarelli et al. , 2018 ; Segler et al. , 2018 ; Kang & Cho , 2018 ; Krenn et al. , 2020 ; Nigam et al. , 2021b ; Shin et al. , 2021 ) , and molecular graphs ( Simonovsky & Komodakis , 2018 ; Liu et al. , 2018 ; Ma et al. , 2018 ; De Cao & Kipf , 2018 ; Samanta et al. , 2020 ; Mercado et al. , 2021 ) and generate molecules character-by-character or node-by-node . Jin et al . ( 2018 ) generates graphs as junction trees by considering the vocabulary as the set of atoms or predefined rings from the data ; Jin et al . ( 2020 ) use the same atom+ring vocabulary to generate molecules by augmenting extracted rationales of molecules . Generating molecules with a fixed molecular fragment vocabulary is a well-established idea in traditional drug design ( Erlanson , 2011 ) , and has been recently explored through deep learning models ( Podda et al. , 2020 ; Xie et al. , 2021 ; Kong et al. , 2021 ; Fu et al. , 2021 ; 2020 ; Leguy et al. , 2020 ) , outperforming previous atom-level models . Recent synthesizability-aware models can also generate single-step reaction ( Bradshaw et al. , 2019 ) and molecule synthesis graphs ( Bradshaw et al. , 2020 ) based on a fixed reactant pool . However , the fixed fragment vocabularies used by these models , which are typically small and predefined a priori , limit the chemical space spanned by the models . In our work , we utilize a learned molecular fragment vocabulary , which is obtained by training a VQ-VAE on a large set of fragments extracted from ChEMBL ( Gaulton et al. , 2012 ) . By sampling fragments from the learned distribution , our model can span a much larger chemical space than methods using a fixed vocabulary ( visualized in Figure 3a , Figure 3b ) . Sequential generation of molecules . Guimaraes et al . ( 2017 ) ; Olivecrona et al . ( 2017 ) ; You et al . ( 2018 ) ; Zhou et al . ( 2019 ) frame the molecular optimization problem as a reinforcement learning problem , but they generate on the atom/character level and from scratch each time , reducing the efficiency of the search algorithm . Jin et al . ( 2019b ) uses a graph-to-graph translation model for property optimization . However , it requires a large number of translation pairs to train , which often involves expert human knowledge and is expensive to obtain . Others have used genetic/evolutionary algorithms to tackle this problem ( Nigam et al. , 2020 ; 2021a ) , which performs random mutations on chemical strings . Although these methods use learned discriminators to prune sub-optimal molecules , the random mutation process can become inefficient in searching for good molecules under complex property constraints . Xie et al . ( 2021 ) ; Fu et al . ( 2021 ) applies Markov Chain Monte Carlo ( MCMC ) sampling through editing molecules , while Kong et al . ( 2021 ) uses Bayesian optimization on the latent space . While there are extensive studies in exploration strategies for RL ( Pathak et al. , 2017 ; Burda et al. , 2019 ; Ecoffet et al. , 2021 ) , diversity/novelty driven molecular generation is under-explored . We train a novelty/diversity-aware RL policy to search for novel , diverse molecules that retain desired properties . Our method initializes searches from model-discovered molecules , which greatly improves the efficiency and diversity of the generated molecules . Our ablation in Section 6 show that our proposed RL framework works well with the learned fragment vocabulary , while a simpler search strategy is not able to utilize this powerful editing paradigm . 3 PRELIMINARIES . Message Passing Neural Networks ( MPNN ) Molecules are represented as directed graphs , where the atoms are the nodes and the bonds are the edges of the graph . More formally , let x = ( V , E ) denote a directed graph where vi ∈ V are the atoms , and eij ∈ E are the edges of the graph . The network maintains hidden states hteij for each directed edge , where t is the layer index . At each step , the hidden representations aggregate information from neighboring nodes and edges , and captures a larger neighborhood of atoms . Iteratively , the hidden states are updated as : ht+1eij = f ( [ h 0 eij ; ∑ k∈N ( vi ) \ { vj } hteki ] ) ( 1 ) Here , f is parameterized by RNN cells ( e.g . LSTM cells ( Hochreiter & Schmidhuber , 1997 ) or GRU cells ( Chung et al. , 2014 ) ) , and N ( vi ) is the set of neighbors of vi . After T steps of messagepassing , the final node embeddings hvi are obtained by summing their respective incoming edge embeddings : hvi = ReLU ( Wo [ h 0 vi ; ∑ vk∈N ( vi ) hTeki ] ) ( 2 ) The final node embeddings are then summed to get a graph embedding representation hx = ∑ vi hvi . Vector-Quantised Variational Autoencoders ( VQ-VAE ) To learn useful representations of fragments , we employ the VQ-VAE architecture ( van den Oord et al. , 2017 ) , which maps molecule fragment graphs to a discrete latent space through using categorical distributions for the prior and posterior . The VQ-VAE defines a dictionary of k embedding elements , [ s1 , s2 , ... sk ] ∈ Rk×l . Given an input x ( here the graph for a molecular fragment ) , let ze ( x ) ∈ Rd×l be the output of the encoder ( a MPNN in our case ) . We define l to be the same dimension for both encoder output embeddings ze ( x ) and dictionary embeddings si , because input zq ( x ) is computed by finding the l2 nearest neighbor dictionary elements for each row of ze ( x ) : zq ( x ) i = sk , where k = argmin j ||ze ( x ) i − sj ||2 for i = 1 , . . . , d ( 3 ) This embedding scheme allows us to represent each molecular fragment using a length d vector , where each entry takes value from { 1 , . . . , k } that corresponds to the dictionary embedding index for that row . The combinatorial vocabulary defined by the VQ-VAE has the capacity to represent kd distinct molecular fragments , which lifts the constraints of a limited generative span under a fixed fragment vocabulary . Since the discretization step does not allow for gradient flow , gradients are passed through the network through approximating the gradient from the dictionary embeddings to the encoder embeddings . Additionally , there is a commitment loss that encourages the encoder to output embeddings that are similar to those in the dictionary ( hence commitment ) . The total loss of the VAE is the following : L = log p ( x|zq ( x ) ) + ∑ i ||sg [ ze ( x ) i ] − sij ||22 + β ∑ i ||ze ( x ) i − sg [ sij ] ||22 ( 4 ) Where sij is the closest dictionary element sj for the ze ( x ) i. Additionally , β is a hyperparameter that controls for contribution of the commitment term , and sg represents the stop-gradient operator . 4 METHODS . Molecular Optimization . The goal of molecular optimization is to generate a set of high-quality molecules C ( Constrained set ) which satisfy or optimize a set of properties P . High novelty and diversity ( detailed in Section 5 ) are also desired for de novo generation applications . We model the molecular optimization problem as a Markov decision process ( MDP ) , defined by the 5-tuple { S , A , p , r , ρ0 } , where the state space S is the set of all possible molecular graphs . As an overview , our method introduces novel designs over the action spaceA and the transition model p ( Section 4.1 ) by utilizing a distributional fragment vocabulary , learned by a VQ-VAE . We define the reward and initial state distribution , r and ρ0 ( Section 4.2 ) accordingly for specified tasks and to implement the proposed sequential translation generation scheme . An illustration of our model is in Figure 1 .

Summary: This paper proposes a new reinforcement learning method for molecular optimization. They train VQ-VAE to learn the vocabulary of molecular fragments which enables the optimizer to efficiently search the chemical space. The optimizer consisted of general reinforcement learning and MPNN, and devised sampling schemes to improve performance. By integrating evaluation criteria as a reward function, the proposed model can additionally improve the performance in terms of those criteria. The experiments show not only its superior performance but shorter optimization time.