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 to perform molecular optimization by sequential translation. A reinforcement learning policy is employed to increase the diversity of the generated molecules. The proposed method is shown to be able to generate high-quality molecules on single-property and multi-property optimization tasks. The paper is well written and easy to follow.