Original: F RAGMENT - BASED S EQUENTIAL T RANSLATION FOR M OLECULAR O PTIMIZATION
Revised: F RAGMENT - BASED S EQUENTIAL T RANSLATION FOR M OLECULAR O PTIMIZATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Searching for novel molecular compounds with desired properties is an important problem in drug discovery.
Revised: Searching for novel molecular compounds with desired properties is an important problem in drug discovery.

Original: Many existing frameworks generate molecules one atom at a time.
Revised: Many existing frameworks generate molecules one atom at a time.

Original: We instead propose a ﬂexible editing paradigm that generates molecules using learned molecular fragments—meaningful substructures of molecules.
Revised: We instead propose a ﬂexible editing paradigm that generates molecules using learned molecular fragments—meaningful substructures of molecules.

Original: To do so, we train a variational autoencoder (VAE) to encode molecular fragments in a coherent latent space, which we then utilize as a vocabulary for editing molecules to explore the complex chemical property space.
Revised: To do so, we train a variational autoencoder (VAE) to encode molecular fragments in a coherent latent space, which we then utilize as a vocabulary for editing molecules to explore the complex chemical property space.

Original: Equipped with the learned fragment vocabulary, we propose F r a gment-based S equential
Revised: Equipped with the learned fragment vocabulary, we propose F r a gment-based S equential

Original: T ranslation (FaST), which learns a reinforcement learning (RL) policy to iteratively translate model-discovered molecules into increasingly novel molecules while satisfying desired properties.
Revised: T ranslation (FaST), which learns a reinforcement learning (RL) policy to iteratively translate model-discovered molecules into increasingly novel molecules while satisfying desired properties.

Original: Empirical evaluation shows that FaST signiﬁcantly improves over state-of-the-art methods on benchmark single/multiobjective molecular optimization tasks.
Revised: Empirical evaluation shows that FaST signiﬁcantly improves over state-of-the-art methods on benchmark single/multiobjective molecular optimization tasks.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Molecular optimization is a challenging task that is pivotal to drug discovery applications.
Revised: Molecular optimization is a challenging task that is pivotal to drug discovery applications.

Original: Part of the challenge stems from the difﬁculty of exploration in the molecular space: not only are there physical constraints on molecules (molecular strings/graphs have to obey speciﬁc chemical principles), molecular property landscapes are also very complex and difﬁcult to characterize: small changes in the molecular space can lead to large deviations in the property space.
Revised: Part of the challenge stems from the difﬁculty of exploration in the molecular space: not only are there physical constraints on molecules (molecular strings/graphs have to obey speciﬁc chemical principles), molecular property landscapes are also very complex and difﬁcult to characterize: small changes in the molecular space can lead to large deviations in the property space.

Original: Recent fragment-based molecular generative models have shown signiﬁcant empirical advantages (Jin et al., 2019a; Podda et al., 2020; Xie et al., 2021) over atom-by-atom generative models in molecular optimization.
Revised: Recent fragment-based molecular generative models have shown signiﬁcant empirical advantages (Jin et al., 2019a; Podda et al., 2020; Xie et al., 2021) over atom-by-atom generative models in molecular optimization.

Original: However, they generally operate over a ﬁxed set of fragments which limit the generative capabilities of these models.
Revised: However, they generally operate over a ﬁxed set of fragments which limit the generative capabilities of these models.

Original: 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).
Revised: 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).

Original: 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 ﬁxed fragment vocabulary.
Revised: 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 ﬁxed fragment vocabulary.

Original: 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.
Revised: 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.

Original: We further introduce a novel sequential translation scheme designed for fragment-based molecular optimization.
Revised: We further introduce a novel sequential translation scheme designed for fragment-based molecular optimization.

Original: 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.
Revised: 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.

Original: Previous works optimize molecules either by generating from scratch or a single translation from known molecules, which is inefﬁcient in ﬁnding high-quality molecules and often discovering molecules lacking novelty/diversity.
Revised: Previous works optimize molecules either by generating from scratch or a single translation from known molecules, which is inefﬁcient in ﬁnding high-quality molecules and often discovering molecules lacking novelty/diversity.

Original: Our proposed framework addresses these deﬁciencies since our method is (1) very efﬁcient in ﬁnding 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 ﬂexible learned fragment vocabulary and a sequence of fragment-based editing.
Revised: Our proposed framework addresses these deﬁciencies since our method is (1) very efﬁcient in ﬁnding 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 ﬂexible learned fragment vocabulary and a sequence of fragment-based editing.

Original: Combining the advantage of a distributional fragment vocabulary and the sequential translation scheme, we propose F r a gment-based S equential T ranslation (FaST), which is realized by an RL policy that proposes fragment addition/deletion to a given molecule.
Revised: Combining the advantage of a distributional fragment vocabulary and the sequential translation scheme, we propose F r a gment-based S equential T ranslation (FaST), which is realized by an RL policy that proposes fragment addition/deletion to a given molecule.

Original: Our proposed method can generate molecules under various objectives such as property constraints, novelty constraints, and diversity constraints.
Revised: Our proposed method can generate molecules under various objectives such as property constraints, novelty constraints, and diversity constraints.

Original: The main contribution of this paper includes:
Revised: The main contribution of this paper includes:

Original: We demonstrate a way to learn distributional molecular fragment vocabulary through a VQ-VAE and the effectiveness of the learned vocabulary in molecular optimization.
Revised: We demonstrate a way to learn distributional molecular fragment vocabulary through a VQ-VAE and the effectiveness of the learned vocabulary in molecular optimization.

Original: 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.
Revised: 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.

Original: We implement a RL policy combining the fragment vocabulary and the sequential translation scheme that signiﬁcantly outperforms state-of-the-art methods in benchmark single/multi-objective molecular optimization tasks.
Revised: We implement a RL policy combining the fragment vocabulary and the sequential translation scheme that signiﬁcantly outperforms state-of-the-art methods in benchmark single/multi-objective molecular optimization tasks.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Molecular generation and optimization.
Revised: Molecular generation and optimization.

Original: Early works on molecular optimization build on generative models on both SMILES/SELFIES string (G´omez-Bombarelli et al., 2018; Segler et al., 2018;
Revised: Early works on molecular optimization build on generative models on both SMILES/SELFIES string (G´omez-Bombarelli et al., 2018; Segler et al., 2018;

Original: 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.
Revised: 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.

Original: Jin et al.
Revised: Jin et al.

Original: (2018) generates graphs as junction trees by considering the vocabulary as the set of atoms or predeﬁned rings from the data; Jin et al.
Revised: (2018) generates graphs as junction trees by considering the vocabulary as the set of atoms or predeﬁned rings from the data; Jin et al.

Original: (2020) use the same atom+ring vocabulary to generate molecules by augmenting extracted rationales of molecules.
Revised: (2020) use the same atom+ring vocabulary to generate molecules by augmenting extracted rationales of molecules.

Original: Generating molecules with a ﬁxed 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;
Revised: Generating molecules with a ﬁxed 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;

Original: Fu et al., 2021; 2020; Leguy et al., 2020), outperforming previous atom-level models.
Revised: Fu et al., 2021; 2020; Leguy et al., 2020), outperforming previous atom-level models.

Original: 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 ﬁxed reactant pool.
Revised: 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 ﬁxed reactant pool.

Original: However, the ﬁxed fragment vocabularies use by these models, which are typically small and predeﬁned a priori, limits the chemical space spanned by the models.
Revised: However, the ﬁxed fragment vocabularies used by these models, which are typically small and predeﬁned a priori, limit the chemical space spanned by the models.

Original: 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).
Revised: 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).

Original: By sampling fragments from the learned distribution, our model can span a much larger chemical space than methods using a ﬁxed vocabulary (visualized in Figure 3a, Figure 3b).
Revised: By sampling fragments from the learned distribution, our model can span a much larger chemical space than methods using a ﬁxed vocabulary (visualized in Figure 3a, Figure 3b).

Original: Sequential generation of molecules.
Revised: Sequential generation of molecules.

Original: Guimaraes et al.
Revised: Guimaraes et al.

Original: ; Olivecrona et al.
Revised: ; Olivecrona et al.

Original: ; You et al.
Revised: ; You et al.

Original: ; Zhou et al.
Revised: ; Zhou et al.

Original: (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 efﬁciency of the search algorithm.
Revised: (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 efﬁciency of the search algorithm.

Original: Jin et al.
Revised: Jin et al.

Original: (2019b) uses a graph-to-graph translation model for property optimization.
Revised: (2019b) uses a graph-to-graph translation model for property optimization.

Original: However, it requires a large number of translation pairs to train, which often involves expert human knowledge and is expensive to obtain.
Revised: However, it requires a large number of translation pairs to train, which often involves expert human knowledge and is expensive to obtain.

Original: Others have used genetic/evolutionary algorithms to tackle this problem (Nigam et al., 2020; 2021a), which performs random mutations on chemical strings.
Revised: Others have used genetic/evolutionary algorithms to tackle this problem (Nigam et al., 2020; 2021a), which performs random mutations on chemical strings.

Original: Although these methods use learned discriminators to prune sub-optimal molecules, the random mutation process can become inefﬁcient in searching for good molecules under complex property constraints.
Revised: Although these methods use learned discriminators to prune sub-optimal molecules, the random mutation process can become inefﬁcient in searching for good molecules under complex property constraints.

Original: Xie et al.
Revised: Xie et al.

Original: Fu et al.
Revised: Fu et al.

Original: (2021) applies Markov Chain Monte
Revised: (2021) applies Markov Chain Monte

Original: Carlo (MCMC) sampling through editing molecules, while Kong et al.
Revised: Carlo (MCMC) sampling through editing molecules, while Kong et al.

Original: (2021) uses Bayesian optimization on the latent space.
Revised: (2021) uses Bayesian optimization on the latent space.

Original: 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.
Revised: 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.

Original: We train a novelty/diversity-aware RL policy to search for novel, diverse molecules that retain desired properties.
Revised: We train a novelty/diversity-aware RL policy to search for novel, diverse molecules that retain desired properties.

Original: Our method initializes searches from model-discovered molecules, which greatly improves the efﬁciency and diversity of the generated molecules.
Revised: Our method initializes searches from model-discovered molecules, which greatly improves the efﬁciency and diversity of the generated molecules.

Original: 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.
Revised: 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.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: 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.
Revised: 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.

Original: More formally, let x = ( V, E ) denote a directed graph where v i ∈ V are the atoms, and e ij ∈
Revised: More formally, let x = ( V, E ) denote a directed graph where v i ∈ V are the atoms, and e ij ∈

Original: E are the edges of the graph.
Revised: E are the edges of the graph.

Original: The network maintains hidden states h te ij for each directed edge, where t is the layer index.
Revised: The network maintains hidden states h te ij for each directed edge, where t is the layer index.

Original: At each step, the hidden representations aggregate information from neighboring nodes and edges, and captures a larger neighborhood of atoms.
Revised: At each step, the hidden representations aggregate information from neighboring nodes and edges, and captures a larger neighborhood of atoms.

Original: Iteratively, the hidden states are updated as:
Revised: Iteratively, the hidden states are updated as:

Original: [Equation]
Revised: [Equation]

Original: Here, f is parameterized by RNN cells (e.g. LSTM cells (Hochreiter & Schmidhuber, 1997) or
Revised: Here, f is parameterized by RNN cells (e.g. LSTM cells (Hochreiter & Schmidhuber, 1997) or

Original: GRU cells (Chung et al., 2014)), and N ( v i ) is the set of neighbors of v i .
Revised: GRU cells (Chung et al., 2014)), and N ( v i ) is the set of neighbors of v i .

Original: After T steps of messagepassing, the ﬁnal node embeddings h v i are obtained by summing their respective incoming edge embeddings:
Revised: After T steps of messagepassing, the ﬁnal node embeddings h v i are obtained by summing their respective incoming edge embeddings:

Original: [Equation]
Revised: [Equation]

Original: The ﬁnal node embeddings are then summed to get a graph embedding representation h x = (cid:80) v i h v i .
Revised: The ﬁnal node embeddings are then summed to get a graph embedding representation h x = (cid:80) v i h v i .

Original: Vector-Quantised Variational Autoencoders (VQ-VAE)
Revised: Vector-Quantised Variational Autoencoders (VQ-VAE)

Original: 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.
Revised: 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.

Original: The VQ-VAE deﬁnes a dictionary of k embedding elements, [ s 1 , s 2 , ...
Revised: The VQ-VAE deﬁnes a dictionary of k embedding elements, [ s 1 , s 2 , ...

Original: s k ] ∈ R k × l .
Revised: s k ] ∈ R k × l .

Original: Given an input x (here the graph for a molecular fragment), let z e ( x ) ∈ R d × l be the output of the encoder (a MPNN in our case).
Revised: Given an input x (here the graph for a molecular fragment), let z e ( x ) ∈ R d × l be the output of the encoder (a MPNN in our case).

Original: We deﬁne l to be the same dimension for both encoder output embeddings z e ( x ) and dictionary embeddings
Revised: We deﬁne l to be the same dimension for both encoder output embeddings z e ( x ) and dictionary embeddings

Original: s i , because input z q ( x ) is computed by ﬁnding the l 2 nearest neighbor dictionary elements for each row of z e ( x ) :
Revised: s i , because input z q ( x ) is computed by ﬁnding the l 2 nearest neighbor dictionary elements for each row of z e ( x ) :

Original: [Equation]
Revised: [Equation]

Original: , where k = arg
Revised: , where k = arg

Original: [Equation]
Revised: [Equation]

Original: This embedding scheme allows us to represent each molecular fragment using a length d vector, where each entry takes value from { 1 , . .
Revised: This embedding scheme allows us to represent each molecular fragment using a length d vector, where each entry takes value from { 1 , . .

Original: . , k } that corresponds to the dictionary embedding index for that row.
Revised: . , k } that corresponds to the dictionary embedding index for that row.

Original: The combinatorial vocabulary deﬁned by the VQ-VAE has the capacity to represent d k distinct molecular fragments, which lifts the constraints of a limited generative span under a ﬁxed fragment vocabulary.
Revised: The combinatorial vocabulary deﬁned by the VQ-VAE has the capacity to represent d k distinct molecular fragments, which lifts the constraints of a limited generative span under a ﬁxed fragment vocabulary.

Original: Since the discretization step does not allow for gradient ﬂow, gradients are passed through the network through approximating the gradient from the dictionary embeddings to the encoder embeddings.
Revised: Since the discretization step does not allow for gradient ﬂow, gradients are passed through the network through approximating the gradient from the dictionary embeddings to the encoder embeddings.

Original: Additionally, there is a commitment loss that encourages the encoder to output embeddings that are similar to those in the dictionary (hence commitment).
Revised: Additionally, there is a commitment loss that encourages the encoder to output embeddings that are similar to those in the dictionary (hence commitment).

Original: The total loss of the VAE is the following:
Revised: The total loss of the VAE is the following:

Original: [Equation]
Revised: [Equation]

Original: Where s ij is the closest dictionary element s j for the z e ( x ) i .
Revised: Where s ij is the closest dictionary element s j for the z e ( x ) i .

Original: Additionally, β is a hyperparameter that controls for contribution of the commitment term, and sg represents the stop-gradient operator.
Revised: Additionally, β is a hyperparameter that controls for contribution of the commitment term, and sg represents the stop-gradient operator.

Original:  M ETHODS
Revised:  M ETHODS

Original: Molecular Optimization.
Revised: Molecular Optimization.

Original: The goal of molecular optimization is to generate a set of high-quality molecules C ( C onstrained set) which satisfy or optimize a set of properties P .
Revised: The goal of molecular optimization is to generate a set of high-quality molecules C ( C onstrained set) which satisfy or optimize a set of properties P .

Original: High novelty and diversity (detailed in Section 5) are also desired for de novo generation applications.
Revised: High novelty and diversity (detailed in Section 5) are also desired for de novo generation applications.

Original: We model the molecular optimization problem as a Markov decision process (MDP), deﬁned by the 5-tuple {S , A , p, r, ρ 0 } , where the state space S is the set of all possible molecular graphs.
Revised: We model the molecular optimization problem as a Markov decision process (MDP), deﬁned by the 5-tuple {S , A , p, r, ρ 0 } , where the state space S is the set of all possible molecular graphs.

Original: As an overview, our method introduces novel designs over the action space A and the transition model p (Section 4.1) by utilizing a distributional fragment vocabulary, learned by a VQ-VAE.
Revised: As an overview, our method introduces novel designs over the action space A and the transition model p (Section 4.1) by utilizing a distributional fragment vocabulary, learned by a VQ-VAE.

Original: We deﬁne the reward and initial state distribution, r and ρ 

Revised: We deﬁne the reward and initial state distribution, r and ρ

Original: (Section 4.2) accordingly for speciﬁed tasks and to implement the proposed sequential translation generation scheme.
Revised: (Section 4.2) accordingly for speciﬁed tasks and to implement the proposed sequential translation generation scheme.

Original: An illustration of our model is in Figure 1.
Revised: An illustration of our model is in Figure 1.

Original: 4. L EARNING D ISTRIBUTIONAL F RAGMENT V OCABULARY
Revised: 4. L EARNING D ISTRIBUTIONAL F RAGMENT V OCABULARY

Original: Molecular Fragments are extracted from molecules in the ChEMBL database (Gaulton et al., 2012).
Revised: Molecular Fragments are extracted from molecules in the ChEMBL database (Gaulton et al., 2012).

Original: For each molecule, we randomly sample fragments by extracting subgraphs that contain ten or fewer atoms that have a single bond attachment to the rest of the molecule.
Revised: For each molecule, we randomly sample fragments by extracting subgraphs that contain ten or fewer atoms that have a single bond attachment to the rest of the molecule.

Original: We then use a VQ-VAE to encode these fragments into a meaningful latent space.
Revised: We then use a VQ-VAE to encode these fragments into a meaningful latent space.

Original: The use of molecular fragments simpliﬁes the search problem, while the variable-sized fragment distribution maintains the reachability of most molecular compounds.
Revised: The use of molecular fragments simpliﬁes the search problem, while the variable-sized fragment distribution maintains the reachability of most molecular compounds.

Original: Because our search algorithm ultimately uses the latent representation of the molecules as the action space, we ﬁnd that using a VQ-VAE with a categorical prior instead of the typical Gaussian prior makes RL training stable and provides good performance gains (Tang & Agrawal, 2020;
Revised: Because our search algorithm ultimately uses the latent representation of the molecules as the action space, we ﬁnd that using a VQ-VAE with a categorical prior instead of the typical Gaussian prior makes RL training stable and provides good performance gains (Tang & Agrawal, 2020;

Original: Grill et al., 2020).
Revised: Grill et al., 2020).

Original: The training instability under a normal VAE with Gaussian prior and continuous latents causes failture of the RL training.
Revised: The training instability under a normal VAE with Gaussian prior and continuous latents causes failture of the RL training.

Original: Our ablation study also shows that the fragment samples from a VQ-VAE are more diverse than the samples from a continuous VAE (Section 6).
Revised: Our ablation study also shows that the fragment samples from a VQ-VAE are more diverse than the samples from a continuous VAE (Section 6).

Original: Encoder/Decoder We use MPNN encoders for any graph inputs, which include both fragments for the VQ-VAE, as well as molecular states during policy learning.
Revised: Encoder/Decoder We use MPNN encoders for any graph inputs, which include both fragments for the VQ-VAE, as well as molecular states during policy learning.

Original: The graph models are especially suitable for describing actions on the molecular state, as they explicitly parametrize the representations of each atom and bond.
Revised: The graph models are especially suitable for describing actions on the molecular state, as they explicitly parametrize the representations of each atom and bond.

Original: Meanwhile, the decoder architecture is a recurrent network that decodes a SELFIES representation of a molecule.
Revised: Meanwhile, the decoder architecture is a recurrent network that decodes a SELFIES representation of a molecule.

Original: We choose a recurrent network for the decoder because we do not need the full complexity of a graph decoder.
Revised: We choose a recurrent network for the decoder because we do not need the full complexity of a graph decoder.

Original: Due to the construction scheme, the fragments are rooted trees, and all have a single attachment point.
Revised: Due to the construction scheme, the fragments are rooted trees, and all have a single attachment point.

Original: As our fragments are small in molecular size ( ≤ 10 atoms), the string grammar is simple to learn, and we ﬁnd the SELFIES decoder works well empirically (see Appendix F for more details).
Revised: As our fragments are small in molecular size ( ≤ 10 atoms), the string grammar is simple to learn, and we ﬁnd the SELFIES decoder works well empirically (see Appendix F for more details).

Original: Adding and deleting fragments as actions.
Revised: Adding and deleting fragments as actions.

Original: At each step of the MDP, the policy network ﬁrst takes the current molecular graph as input and produces a Bernoulli distribution on whether to add or delete a fragment.
Revised: At each step of the MDP, the policy network ﬁrst takes the current molecular graph as input and produces a Bernoulli distribution on whether to add or delete a fragment.

Original: Equipped with the fragment VQ-VAE, we deﬁne the Add and Delete actions at the fragment-level:
Revised: Equipped with the fragment VQ-VAE, we deﬁne the Add and Delete actions at the fragment-level:

Original: • Fragment Addition.
Revised: • Fragment Addition.

Original: The addition action is characterized by (1) a probability distribution over the atoms of the molecule: p add ( v i ) =
Revised: The addition action is characterized by (1) a probability distribution over the atoms of the molecule: p add ( v i ) =

Original: σ [ MLP ( h v )] , where σ is the softmax operator.
Revised: σ [ MLP ( h v )] , where σ is the softmax operator.

Original: (2) Conditioned on the graph embedding h x and the attachment point atom v add sampled from p add , we predict a d -channel categorical distribution p fragment =
Revised: (2) Conditioned on the graph embedding h x and the attachment point atom v add sampled from p add , we predict a d -channel categorical distribution p fragment =

Original: σ [ MLP ([ h v add ; h x ])]
Revised: σ [ MLP ([ h v add ; h x ])]

Original: ∈ R d × k , where each row of p fragment sums to 1.
Revised: ∈ R d × k , where each row of p fragment sums to 1.

Original: We can then sample the discrete categorical latent z add ∈ { 1 , ..., k } d from p fragment .
Revised: We can then sample the discrete categorical latent z add ∈ { 1 , ..., k } d from p fragment .

Original: The fragment to add
Revised: The fragment to add

Original: is then obtained by deocoding z add through the learned frozen fragment decoder.
Revised: is then obtained by deocoding z add through the learned frozen fragment decoder.

Original: We then assemble the decoded fragment with the current molecular graph by attaching the fragment to the predicted attachment point v add .
Revised: We then assemble the decoded fragment with the current molecular graph by attaching the fragment to the predicted attachment point v add .

Original: Note that the attachment point over the fragment is indicated through the generated SELFIES string.
Revised: Note that the attachment point over the fragment is indicated through the generated SELFIES string.

Original: • Fragment Deletion.
Revised: • Fragment Deletion.

Original: The deletion action acts over the directed edges of the molecule.
Revised: The deletion action acts over the directed edges of the molecule.

Original: A probability distribution over deletable edges is computed with a MLP: p del ( e ij )
Revised: A probability distribution over deletable edges is computed with a MLP: p del ( e ij )

Original: =
Revised: =

Original: σ [ MLP ( h e ij )] .
Revised: σ [ MLP ( h e ij )] .

Original: One edge is then sampled and deleted; since the edges are directed, the directionality specify the the molecule to keep and the fragment to be deleted.
Revised: One edge is then sampled and deleted; since the edges are directed, the directionality specify the the molecule to keep and the fragment to be deleted.

Original: With the action space A deﬁned as above, the transition model for the MDP is simply p ( s (cid:48) | s, a ) = 1 if applying the addition/deletion action a to the molecule s results in the molecule s (cid:48) , and p ( s (cid:48) | s, a ) = 0 otherwise.
Revised: With the action space A deﬁned as above, the transition model for the MDP is simply p ( s (cid:48) | s, a ) = 1 if applying the addition/deletion action a to the molecule s results in the molecule s (cid:48) , and p ( s (cid:48) | s, a ) = 0 otherwise.

Original: The fragment-based action space is powerful and suitable for policy learning as it (1) is powered by the enormous distributional vocabulary learned by the fragment VQ-VAE, thus spans a diverse set of editing operations over molecular graphs; (2) exploits the meaningful latent representation of fragments since the representation of similar fragments are grouped together.
Revised: The fragment-based action space is powerful and suitable for policy learning as it (1) is powered by the enormous distributional vocabulary learned by the fragment VQ-VAE, thus spans a diverse set of editing operations over molecular graphs; (2) exploits the meaningful latent representation of fragments since the representation of similar fragments are grouped together.

Original: These advantages greatly simplify the molecular search problem.
Revised: These advantages greatly simplify the molecular search problem.

Original: We terminate an episode when the molecule fails to satisfy the desired property or when the episode exceeds ten steps.
Revised: We terminate an episode when the molecule fails to satisfy the desired property or when the episode exceeds ten steps.

Original: 4. D ISCOVER N OVEL M OLECULES THROUGH S EQUENTIAL T RANSLATION
Revised: 4. D ISCOVER N OVEL M OLECULES THROUGH S EQUENTIAL T RANSLATION

Original: We propose sequential translation that incrementally grows the set of discovered novel molecules and use the model-discovered molecules as starting points for further search episodes.
Revised: We propose sequential translation that incrementally grows the set of discovered novel molecules and use the model-discovered molecules as starting points for further search episodes.

Original: This regime of starting exploration from states reached in previous episodes was also explored under the setting of RL from image inputs (Ecoffet et al., 2021).
Revised: This regime of starting exploration from states reached in previous episodes was also explored under the setting of RL from image inputs (Ecoffet et al., 2021).

Original: More concretely, we implement sequential translation with a reinforcement learning policy that operates under the fragment-based action space deﬁned in
Revised: More concretely, we implement sequential translation with a reinforcement learning policy that operates under the fragment-based action space deﬁned in

Original: Section 4.1, while using a moving initial state distribution ρ 0 , which is a distribution over molecules in the frontier set F = I ∪C .
Revised: Section 4.1, while using a moving initial state distribution ρ 0 , which is a distribution over molecules in the frontier set F = I ∪C .

Original: By starting new search episodes from the frontier set – the union of the initial set and good molecules that are discovered by the RL policy, we achieve efﬁcient search in the chemical space by staying close to the high-quality subspace and achieve novel molecule generation through a sequence of fragment-based editing operations to the known molecules.
Revised: By starting new search episodes from the frontier set – the union of the initial set and good molecules that are discovered by the RL policy, we achieve efﬁcient search in the chemical space by staying close to the high-quality subspace and achieve novel molecule generation through a sequence of fragment-based editing operations to the known molecules.

Original: Our proposed search algorithm is detailed in Algorithm 1.
Revised: Our proposed search algorithm is detailed in Algorithm 1.

Original: Discover novel molecules and expand the frontier.
Revised: Discover novel molecules and expand the frontier.

Original: Our method explores the chemical space with a property-aware and novelty/diversity-aware reinforcement learning policy that proposes addition/deletion modiﬁcations to the molecular state at every environment step to optimize for the reward r .
Revised: Our method explores the chemical space with a property-aware and novelty/diversity-aware reinforcement learning policy that proposes addition/deletion modiﬁcations to the molecular state at every environment step to optimize for the reward r .

Original: We gradually expand the discovered set C by adding qualiﬁed molecules found in the
Revised: We gradually expand the discovered set C by adding qualiﬁed molecules found in the

Original: RL exploration within the MDP.
Revised: RL exploration within the MDP.

Original: A molecule x is qualiﬁed
Revised: A molecule x is qualiﬁed

Original: if: (1) x satisﬁes the desired properties measured by property scores
Revised: if: (1) x satisﬁes the desired properties measured by property scores

Original: [Equation]
Revised: [Equation]

Original: where P is the set of desired properties and threshold p is the score threshold for satisfying property p .
Revised: where P is the set of desired properties and threshold p is the score threshold for satisfying property p .

Original: A molecule x satisfying all desired properties has C P ( x ) = 1 and C P ( x )
Revised: A molecule x satisfying all desired properties has C P ( x ) = 1 and C P ( x )

Original: = 

Revised: =

Original: otherwise.
Revised: otherwise.

Original: (2) x is novel/diverse compared to molecules currently in the frontier F , measured by ﬁngerprint similarity (detailed in Section 5):
Revised: (2) x is novel/diverse compared to molecules currently in the frontier F , measured by ﬁngerprint similarity (detailed in Section 5):

Original: [Equation]
Revised: [Equation]

Original: Where sim denotes ﬁngerprint similarity, threshold nov and threshold div are predeﬁned similarity thresholds for novelty and diversity, I and C are the initial set of good molecules and model discovered molecules as deﬁned in previous sections.
Revised: Where sim denotes ﬁngerprint similarity, threshold nov and threshold div are predeﬁned similarity thresholds for novelty and diversity, I and C are the initial set of good molecules and model discovered molecules as deﬁned in previous sections.

Original: A molecule that satisﬁes novelty/diversity criterion has C ND ( x ) = 1 and C ND ( x ) = 0 otherwise.
Revised: A molecule that satisﬁes novelty/diversity criterion has C ND ( x ) = 1 and C ND ( x ) = 0 otherwise.

Original: We use a reward of +1 for a transition that results in a molecule qualiﬁed for the set C , and discourage the model from producing invalid molecules by adding a reward of − 0 .
Revised: We use a reward of +1 for a transition that results in a molecule qualiﬁed for the set C , and discourage the model from producing invalid molecules by adding a reward of − 0 .

Original: 1 for a transition that produces an invalid molecular graph 1 :
Revised: 1 for a transition that produces an invalid molecular graph 1 :

Original: [Equation]
Revised: [Equation]

Original: where [ x ← a ] denotes the molecule resulting from editing x with the fragment addition/deletion action a . add encoding
Revised: where [ x ← a ] denotes the molecule resulting from editing x with the fragment addition/deletion action a . add encoding

Original: Decode fragment y = D ( z add ) 21:
Revised: Decode fragment y = D ( z add ) 21:

Original: Add fragment y to molecule: x ← x + y 22: else 23:
Revised: Add fragment y to molecule: x ← x + y 22: else 23:

Original: Sample
Revised: Sample

Original: ] ∀
Revised: ] ∀

Original: ∈
Revised: ∈

Original: Section
Revised: Section

Original: Let y be the fragment designated by e , delete fragment x ← x − y 25: if C P ( x ) = 1 & C ND ( x ) = 1 then 26:
Revised: Let y be the fragment designated by e , delete fragment x ← x − y 25: if C P ( x ) = 1 & C ND ( x ) = 1 then 26:

Original: C ← C ∪ { x } 27:
Revised: C ← C ∪ { x } 27:

Original: F ← I ∪ C 28:
Revised: F ← I ∪ C 28:

Original: Let step ← step +
Revised: Let step ← step +

Original: Initialize search episodes from promising candidates.
Revised: Initialize search episodes from promising candidates.

Original: To bias the initial state distribution ρ 0 to favor molecules that can derive more novel high-quality molecules, we keep an upper-conﬁdencebound (UCB) score for each initial molecule in the frontier F .
Revised: To bias the initial state distribution ρ 0 to favor molecules that can derive more novel high-quality molecules, we keep an upper-conﬁdencebound (UCB) score for each initial molecule in the frontier F .

Original: We record the number of times we initiate a search N ( x, t ) from a molecule x ∈ F , and the number of molecules qualiﬁed for adding to C that is found in an episode strating from x : R ( x, t ) .
Revised: We record the number of times we initiate a search N ( x, t ) from a molecule x ∈ F , and the number of molecules qualiﬁed for adding to C that is found in an episode strating from x : R ( x, t ) .

Original: Here t = (cid:80) x ∈ ρ 0 N ( x ) is the total number of search episodes.
Revised: Here t = (cid:80) x ∈ ρ 0 N ( x ) is the total number of search episodes.

Original: The UCB score of the initial molecule m is calculated by:
Revised: The UCB score of the initial molecule m is calculated by:

Original: [Equation]
Revised: [Equation]

Original: The probability of a molecule in the initialization set being sampled as the starting point of a new episode is then computed by a softmax over the UCB scores: p init ( x, t + 1) =
Revised: The probability of a molecule in the initialization set being sampled as the starting point of a new episode is then computed by a softmax over the UCB scores: p init ( x, t + 1) =

Original: [Equation]
Revised: [Equation]

Original: To summarize, FaST learns a policy that (1) choose good initial molecules to start search episodes; (2) choose to add a fragment to or delete a subgraph from a given state (a molecule in our case); (3) choose what to add through predicting a fragment latent embedding, or what to delete through predicting a directed edge, and remove part of the molecular graph accordingly.
Revised: To summarize, FaST learns a policy that (1) choose good initial molecules to start search episodes; (2) choose to add a fragment to or delete a subgraph from a given state (a molecule in our case); (3) choose what to add through predicting a fragment latent embedding, or what to delete through predicting a directed edge, and remove part of the molecular graph accordingly.

Original: Although we present our method in this section under the most realistic multi-objective optimization task settings (with experiments in Section 5), our method is easily extendable to other problem settings by modifying the constraints C P , C ND , and the reward function r accordingly.
Revised: Although we present our method in this section under the most realistic multi-objective optimization task settings (with experiments in Section 5), our method is easily extendable to other problem settings by modifying the constraints C P , C ND , and the reward function r accordingly.

Original: For example, see Appendix B for the application of our method to the standard constrained penalized log P task and Section 5 for multi-objective molecular optimization under different novelty/diversity metrics.
Revised: For example, see Appendix B for the application of our method to the standard constrained penalized log P task and Section 5 for multi-objective molecular optimization under different novelty/diversity metrics.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Datasets.
Revised: Datasets.

Original: We use benchmark datasets for molecular optimization, which aims to generate ligand molecules for inhibition of two proteins: glycogen synthase kinase-3 beta (GSK3 β ) and c-Jun Nterminal kinase 3 (JNK3).
Revised: We use benchmark datasets for molecular optimization, which aims to generate ligand molecules for inhibition of two proteins: glycogen synthase kinase-3 beta (GSK3 β ) and c-Jun Nterminal kinase 3 (JNK3).

Original: Following previous work (Jin et al., 2020; Xie et al., 2021; Nigam et al., 2021a), we adopt the same strategy of using a random forest trained on these datasets as the oracle property predictor, and incorporate the additional factors, quantitative estimate of drug-likeliness (QED) (Bickerton et al., 2012) and synthetic accessibility (SA) (Ertl & Schuffenhauer, 2009) as our
Revised: Following previous work (Jin et al., 2020; Xie et al., 2021; Nigam et al., 2021a), we adopt the same strategy of using a random forest trained on these datasets as the oracle property predictor, and incorporate the additional factors, quantitative estimate of drug-likeliness (QED) (Bickerton et al., 2012) and synthetic accessibility (SA) (Ertl & Schuffenhauer, 2009) as our

Original: [Figure]
Revised: [Figure]

Original: optimization objectives.
Revised: optimization objectives.

Original: Single property optimization is often a ﬂawed task, because the generator can overﬁt to the pretrained predictor and generate unsynthesizable compounds.
Revised: Single property optimization is often a ﬂawed task, because the generator can overﬁt to the pretrained predictor and generate unsynthesizable compounds.

Original: Evaluation metrics.
Revised: Evaluation metrics.

Original: Following previous works, we evaluate our generative model on three target metrics, success, novelty and diversity.
Revised: Following previous works, we evaluate our generative model on three target metrics, success, novelty and diversity.

Original: 5,000 molecules are generated by the model, and the metric scores are computed as follows:
Revised: 5,000 molecules are generated by the model, and the metric scores are computed as follows:

Original: Success rate (SR) measures the proportion of generated molecules that ﬁt the desired properties.
Revised: Success rate (SR) measures the proportion of generated molecules that ﬁt the desired properties.

Original: Novelty (Nov) measures how different the generated molecules are compared to the set of actives in the dataset (range [0 , 1] ), and Diversity (Div) measures how different the generated molecules are compared to each other (range [0 , 1] ).
Revised: Novelty (Nov) measures how different the generated molecules are compared to the set of actives in the dataset (range [0 , 1] ), and Diversity (Div) measures how different the generated molecules are compared to each other (range [0 , 1] ).

Original: PM is the product of the three metrics above ( PM = SR · Nov · Div ).
Revised: PM is the product of the three metrics above ( PM = SR · Nov · Div ).

Original: Implementation details.
Revised: Implementation details.

Original: We construct the initial set of molecules for our search algorithm from the rationales extracted from Jin et al.
Revised: We construct the initial set of molecules for our search algorithm from the rationales extracted from Jin et al.

Original: These rationales are obtained through a sampling process on the active molecules that tries to minimize the size of the rationale subgraph, while maintaining their inhibitory properties.
Revised: These rationales are obtained through a sampling process on the active molecules that tries to minimize the size of the rationale subgraph, while maintaining their inhibitory properties.

Original: Rationales for multi-property tasks (GSK3 β +JNK3) are extracted by combining the rationales for single-property tasks.
Revised: Rationales for multi-property tasks (GSK3 β +JNK3) are extracted by combining the rationales for single-property tasks.

Original: Initializing generation with subgraphs is commonly done in molecular generative models such as Shi et al.
Revised: Initializing generation with subgraphs is commonly done in molecular generative models such as Shi et al.

Original: (2020) and Kong et al.
Revised: (2020) and Kong et al.

Original: We train the RL policy using the Proximal Policy Optimization (PPO, Schulman et al. 2017) algorithm.
Revised: We train the RL policy using the Proximal Policy Optimization (PPO, Schulman et al. 2017) algorithm.

Original: We ﬁnd the RL training robust despite both the reward function r and the initial state distribution ρ 0 are non-stationary (i.e., changing during RL training).
Revised: We ﬁnd the RL training robust despite both the reward function r and the initial state distribution ρ 0 are non-stationary (i.e., changing during RL training).

Original: Hyperparameters used for producing the results in Section 5 and molecule samples from FaST are included in Appendix D and Appendix E.
Revised: Hyperparameters used for producing the results in Section 5 and molecule samples from FaST are included in Appendix D and Appendix E.

Original: Baseline methods.
Revised: Baseline methods.

Original: Rationale-RL (Jin et al., 2020) extracts rationales of the active molecules and then uses RL to train a completion model that add atoms to the rationale in a sequential manner to generate molecules satisfying the desired properties.
Revised: Rationale-RL (Jin et al., 2020) extracts rationales of the active molecules and then uses RL to train a completion model that add atoms to the rationale in a sequential manner to generate molecules satisfying the desired properties.

Original: GA+D & JANUS (Nigam et al., 2020; 2021a) are two genetic algorithms that use random mutations of SELFIES strings to generate promising molecular candidates; JANUS leverages a two-pronged approach, accounting for mutations towards both exploration and exploitation.
Revised: GA+D & JANUS (Nigam et al., 2020; 2021a) are two genetic algorithms that use random mutations of SELFIES strings to generate promising molecular candidates; JANUS leverages a two-pronged approach, accounting for mutations towards both exploration and exploitation.

Original: MARS (Xie et al., 2021) uses Markov Chain Monte
Revised: MARS (Xie et al., 2021) uses Markov Chain Monte

Original: Carlo (MCMC) sampling to iterative build new molecules by adding or removing fragments, and the model is trained to ﬁt the distribution of the active molecules.
Revised: Carlo (MCMC) sampling to iterative build new molecules by adding or removing fragments, and the model is trained to ﬁt the distribution of the active molecules.

Original: To provide a fair comparison against baselines that do not use rationales, we additionally include a baseline MARS+Rationale that initialize the MARS algorithm with the same starting initial rationale set used in Rationale-RL and our method.
Revised: To provide a fair comparison against baselines that do not use rationales, we additionally include a baseline MARS+Rationale that initialize the MARS algorithm with the same starting initial rationale set used in Rationale-RL and our method.

Original: where possible, we use the numbers from the original corresponding paper.
Revised: where possible, we use the numbers from the original corresponding paper.

Original: Performance.
Revised: Performance.

Original: The evaluation metrics are shown in Table 1; FaST signiﬁcantly outperforms all baselines on all tasks including both single-property and multi-property optimization.
Revised: The evaluation metrics are shown in Table 1; FaST signiﬁcantly outperforms all baselines on all tasks including both single-property and multi-property optimization.

Original: On the most challenging task, GSK3 β +JNK3+QED+SA, FaST improves upon the previous best model by over 30% in the product of the three evaluation metrics.
Revised: On the most challenging task, GSK3 β +JNK3+QED+SA, FaST improves upon the previous best model by over 30% in the product of the three evaluation metrics.

Original: Our model is able to efﬁciently search for molecules that stay within the constrained property space, and discover novel and diverse molecules by sequentially translating known and discovered active molecules.
Revised: Our model is able to efﬁciently search for molecules that stay within the constrained property space, and discover novel and diverse molecules by sequentially translating known and discovered active molecules.

Original: The MARS+Rationale model, which uses the same rationales as the initialization for their search algorithm, does not perform
Revised: The MARS+Rationale model, which uses the same rationales as the initialization for their search algorithm, does not perform

Original: [Table]
Revised: [Table]

Original: well compared to the original implementation
Revised: well compared to the original implementation

Original: [Table]
Revised: [Table]

Original: which initializes each search with a simple “C-C” molecule
Revised: which initializes each search with a simple “C-C” molecule

Original: [Table]
Revised: [Table]

Original: Sample complexity comparison given performance thresholds.
Revised: Sample complexity comparison given performance thresholds.

Original: Another comparison scheme is to let a model keep generating molecules until it achieves a good candidate set under certain performance thresholds.
Revised: Another comparison scheme is to let a model keep generating molecules until it achieves a good candidate set under certain performance thresholds.

Original: Under this evaluation protocol, all models will have the same or very similar performance in SR, Nov, and Div.
Revised: Under this evaluation protocol, all models will have the same or very similar performance in SR, Nov, and Div.

Original: The metric of interest will then be the sample complexity of the algorithm: how many molecules it requires to visit/generate to obtain a good candidate set.
Revised: The metric of interest will then be the sample complexity of the algorithm: how many molecules it requires to visit/generate to obtain a good candidate set.

Original: This places every model under the same regime, allowing each model to generate molecules in a novelty/diversity-aware setting.
Revised: This places every model under the same regime, allowing each model to generate molecules in a novelty/diversity-aware setting.

Original: We compare FaST to Rationale-RL and MARS under this setting in
Revised: We compare FaST to Rationale-RL and MARS under this setting in

Original: Figure 2, where we impose SR=1, Nov=1, Div=.7 for the candidate set.
Revised: Figure 2, where we impose SR=1, Nov=1, Div=.7 for the candidate set.

Original: FaST on average searched through 71k molecules in total to gather the 5k proposal set, while Rational-RL and MARS need to search through 205k and 759k molecules to obtain their corresponding proposal sets.
Revised: FaST on average searched through 71k molecules in total to gather the 5k proposal set, while Rational-RL and MARS need to search through 205k and 759k molecules to obtain their corresponding proposal sets.

Original: Being a pretrained generative model, Rationale-RL has a steeper slope initially but then slows down to ﬁnd more good molecules.
Revised: Being a pretrained generative model, Rationale-RL has a steeper slope initially but then slows down to ﬁnd more good molecules.

Original: The ﬂexibility of our learned vocabulary and the RL search strategy lead to the superior performance of FaST, which we further verify through ablation study in Section 6.
Revised: The ﬂexibility of our learned vocabulary and the RL search strategy lead to the superior performance of FaST, which we further verify through ablation study in Section 6.

Original: Optimize for different novelty/diversity metrics.
Revised: Optimize for different novelty/diversity metrics.

Original: The Morgan ﬁngerprints used for similarity comparison contain certain inductive biases.
Revised: The Morgan ﬁngerprints used for similarity comparison contain certain inductive biases.

Original: Under different applications, different novelty/diversity metrics may be of interest.
Revised: Under different applications, different novelty/diversity metrics may be of interest.

Original: To demonstrate the viability of our model under any metrics, we train
Revised: To demonstrate the viability of our model under any metrics, we train

Original: FaST using Atom Pairs (AP) ﬁngerprints (Carhart et al., 1985) on the GSK3 β +JNK3+QED+SA task.
Revised: FaST using Atom Pairs (AP) ﬁngerprints (Carhart et al., 1985) on the GSK3 β +JNK3+QED+SA task.

Original: The results, and discussion of the different ﬁngerprint methods, are reported in Appendix C.
Revised: The results, and discussion of the different ﬁngerprint methods, are reported in Appendix C.

Original: We ﬁnd that (1) FaST can still ﬁnd high-quality molecules that are novel and diverse, while the baseline methods suffer a low novelty; (2) FaST trained with AP ﬁngerprints still attains good performance when evaluated under Morgan ﬁngerprints but the reverse is not true.
Revised: We ﬁnd that (1) FaST can still ﬁnd high-quality molecules that are novel and diverse, while the baseline methods suffer a low novelty; (2) FaST trained with AP ﬁngerprints still attains good performance when evaluated under Morgan ﬁngerprints but the reverse is not true.

Original: This result shows that Novelty/diversity under AP ﬁngerprints is a stricter criterion to satisfy and necessitates novelty/diversity-awareness during optimization.
Revised: This result shows that Novelty/diversity under AP ﬁngerprints is a stricter criterion to satisfy and necessitates novelty/diversity-awareness during optimization.

Original: Penalized logP Maximization.
Revised: Penalized logP Maximization.

Original: To demonstrate the wide applicability of our method to any molecular optimization task, we also include our results on the standard constrained penalized logP optimization task in Appendix B.
Revised: To demonstrate the wide applicability of our method to any molecular optimization task, we also include our results on the standard constrained penalized logP optimization task in Appendix B.

Original: We show that our model signiﬁcantly outperforms all baselines in this task under different constraint levels.
Revised: We show that our model signiﬁcantly outperforms all baselines in this task under different constraint levels.

Original: We also provide insight on the task itself: while this task has
Revised: We also provide insight on the task itself: while this task has

Original: [Figure]
Revised: [Figure]

Original: been studied in many previous works, the task, as it is currently deﬁned, is not entirely chemically meaningful.
Revised: been studied in many previous works, the task, as it is currently deﬁned, is not entirely chemically meaningful.

Original: Additionally, one drawback of this task is that a model can achieve high performance by simply generating large molecules.
Revised: Additionally, one drawback of this task is that a model can achieve high performance by simply generating large molecules.

Original: A detailed discussion can be found in Appendix B.
Revised: A detailed discussion can be found in Appendix B.

Original:  A BLATION AND A NALYSIS
Revised:  A BLATION AND A NALYSIS

Original: Diversity of generation.
Revised: Diversity of generation.

Original: In addition to the ﬁngerprint diversity metrics presented in Section 5, we also examine functional group diversity.
Revised: In addition to the ﬁngerprint diversity metrics presented in Section 5, we also examine functional group diversity.

Original: We extract all unique molecular fragments of the 5, molecules generated for GSK3 β +JNK3+QED+SA task for each model, and produce t-SNE visualization of these fragments in Figure 3a and Figure 3b.
Revised: We extract all unique molecular fragments of the 5, molecules generated for GSK3 β +JNK3+QED+SA task for each model, and produce t-SNE visualization of these fragments in Figure 3a and Figure 3b.

Original: In total, we extracted 1.7k unique fragments from our model outputs vs only 1.1k unique fragments for Rational-RL and 500 unique fragments from MARS.
Revised: In total, we extracted 1.7k unique fragments from our model outputs vs only 1.1k unique fragments for Rational-RL and 500 unique fragments from MARS.

Original: The visualization shows that the fragments in the molecules generated by our model spans a much larger chemical space.
Revised: The visualization shows that the fragments in the molecules generated by our model spans a much larger chemical space.

Original: This conﬁrms the advantages of using a learned vocabulary, compared to using a ﬁxed set of fragments, as we are able to utilize a much more diverse set of chemical subgraphs.
Revised: This conﬁrms the advantages of using a learned vocabulary, compared to using a ﬁxed set of fragments, as we are able to utilize a much more diverse set of chemical subgraphs.

Original: Sampled trajectories (Figure 6) and molecules (Figure 7) are included in Appendix E.
Revised: Sampled trajectories (Figure 6) and molecules (Figure 7) are included in Appendix E.

Original: Beneﬁt of distributional vocabulary.
Revised: Beneﬁt of distributional vocabulary.

Original: To investigate the beneﬁt of using a distributional vocabulary, instead of using the pretrained VQ-VAE, we also train our model using a ﬁxed vocabulary of fragments, which consists of roughly 50k unique fragments (the same set used to pretrain the VQ-VAE).
Revised: To investigate the beneﬁt of using a distributional vocabulary, instead of using the pretrained VQ-VAE, we also train our model using a ﬁxed vocabulary of fragments, which consists of roughly 50k unique fragments (the same set used to pretrain the VQ-VAE).

Original: Figure 2 compares the performance of the two models.
Revised: Figure 2 compares the performance of the two models.

Original: On average, the model with ﬁxed fragments took 122k steps, while with VQ-VAE it only took 71k steps to ﬁnd a set of 5,000 good molecules (72% improvement).
Revised: On average, the model with ﬁxed fragments took 122k steps, while with VQ-VAE it only took 71k steps to ﬁnd a set of 5,000 good molecules (72% improvement).

Original: We further analyze the beneﬁt of using discrete latents with a VQ-VAE rather than continuous latents with a Gaussian prior VAE in Appendix A.
Revised: We further analyze the beneﬁt of using discrete latents with a VQ-VAE rather than continuous latents with a Gaussian prior VAE in Appendix A.

Original: Importance of RL search.
Revised: Importance of RL search.

Original: We also demonstrate the importance of our RL search policy compared to previous sampling methods.
Revised: We also demonstrate the importance of our RL search policy compared to previous sampling methods.

Original: To do so, we run MARS with both the 50k ﬁxed fragments and our
Revised: To do so, we run MARS with both the 50k ﬁxed fragments and our

Original: VQ-VAE.
Revised: VQ-VAE.

Original: Figure 2 shows the performance of using the 50k ﬁxed fragment vocabulary compared to the original MARS model which uses a small 1k vocabulary.
Revised: Figure 2 shows the performance of using the 50k ﬁxed fragment vocabulary compared to the original MARS model which uses a small 1k vocabulary.

Original: When the vocabulary is large, MARS exhibits very poor sample complexity.
Revised: When the vocabulary is large, MARS exhibits very poor sample complexity.

Original: Additionally, we also implemented our VQ-VAE with the sampling strategy proposed in MARS, but this model was altogether unable to successfully generate good candidate molecules.
Revised: Additionally, we also implemented our VQ-VAE with the sampling strategy proposed in MARS, but this model was altogether unable to successfully generate good candidate molecules.

Original: Therefore, we see that when the vocabulary is more complex, we need a better search strategy, highlighting the importance of our RL algorithm.
Revised: Therefore, we see that when the vocabulary is more complex, we need a better search strategy, highlighting the importance of our RL algorithm.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We propose a new framework for molecular optimization, which leverages a learned vocabulary of molecular fragments to search the chemical space efﬁciently.
Revised: We propose a new framework for molecular optimization, which leverages a learned vocabulary of molecular fragments to search the chemical space efﬁciently.

Original: We demonstrate that F r a gment-based
Revised: We demonstrate that F r a gment-based

Original: S equential T ranslation (FaST), which adaptively grows a set of promising molecular candidates, can generate high-quality, novel, and diverse molecules on single-property and multi-property optimization tasks.
Revised: S equential T ranslation (FaST), which adaptively grows a set of promising molecular candidates, can generate high-quality, novel, and diverse molecules on single-property and multi-property optimization tasks.

Original: Ablation study shows that all components of our proposed method contribute to its superior performance.
Revised: Ablation study shows that all components of our proposed method contribute to its superior performance.

Original: The learning of a ﬂexible vocabulary is a complementary module to other research in fragment-based drug design.
Revised: The learning of a ﬂexible vocabulary is a complementary module to other research in fragment-based drug design.

Original: Incorporating FaST to more practical drug discovery pipelines while taking synthesis paths in mind is an exciting avenue for future work.
Revised: Incorporating FaST to more practical drug discovery pipelines while taking synthesis paths in mind is an exciting avenue for future work.

Original: A V OCABULARY LEARNING THROUGH VQ-VAE
Revised: A V OCABULARY LEARNING THROUGH VQ-VAE

Original: To evaluate the beneﬁts of VQ-VAE over a typical VAE trained with Gaussian priors, we train both models, and look at the distribution of fragments.
Revised: To evaluate the beneﬁts of VQ-VAE over a typical VAE trained with Gaussian priors, we train both models, and look at the distribution of fragments.

Original: Figure 4a compares the t-SNE distributions of the two models, where we sample 2,000 fragments from each model.
Revised: Figure 4a compares the t-SNE distributions of the two models, where we sample 2,000 fragments from each model.

Original: The VAE model has tight clusters, while the VQ-VAE model exhibits a much more diverse set of fragments.
Revised: The VAE model has tight clusters, while the VQ-VAE model exhibits a much more diverse set of fragments.

Original: We visualize random samples from VAE and VQ-VAE Figure 4, where we see that the samples from VAE are relatively simple and generic fragments, while samples from the VQ-VAE demonstrate diverse patterns.
Revised: We visualize random samples from VAE and VQ-VAE Figure 4, where we see that the samples from VAE are relatively simple and generic fragments, while samples from the VQ-VAE demonstrate diverse patterns.

Original: This is because the more generic fragments appear more frequently in real molecules, and a Gaussian prior over the fragment latent space would favor these fragments.
Revised: This is because the more generic fragments appear more frequently in real molecules, and a Gaussian prior over the fragment latent space would favor these fragments.

Original: B C ONSTRAINED PENALIZED LOG P T ASK
Revised: B C ONSTRAINED PENALIZED LOG P T ASK

Original: To demonstrate the general applicability of our model for any molecular optimization task, we also run our model on another constrained optimization task, here optimizing for penalized octanol-water partition coefﬁcients (logP) scores of ZINC (Irwin et al., 2012) molecules.
Revised: To demonstrate the general applicability of our model for any molecular optimization task, we also run our model on another constrained optimization task, here optimizing for penalized octanol-water partition coefﬁcients (logP) scores of ZINC (Irwin et al., 2012) molecules.

Original: The penalized logP score is the logP score penalized by synthetic accessibility and ring size.
Revised: The penalized logP score is the logP score penalized by synthetic accessibility and ring size.

Original: We use the exact computation in You et al.
Revised: We use the exact computation in You et al.

Original: (2018), where the components of the penalized logP score are normalized across the entire 250k ZINC training set.
Revised: (2018), where the components of the penalized logP score are normalized across the entire 250k ZINC training set.

Original: The generated molecules are constrained to have similar Morgan ﬁngerprints (Rogers & Hahn, 2010) as the original molecules.
Revised: The generated molecules are constrained to have similar Morgan ﬁngerprints (Rogers & Hahn, 2010) as the original molecules.

Original: Following the same setup as previous work (Jin et al., 2019b;
Revised: Following the same setup as previous work (Jin et al., 2019b;

Original: You et al., 2018; Shi et al., 2020;
Revised: You et al., 2018; Shi et al., 2020;

Original: Nigam et al., 2020; Kong et al., 2021), we try to optimize the 800 test molecules from ZINC with the lowest penalized logP scores (the initial set I ).
Revised: Nigam et al., 2020; Kong et al., 2021), we try to optimize the 800 test molecules from ZINC with the lowest penalized logP scores (the initial set I ).

Original: Speciﬁcally, the task is to translate these molecules into new molecules with the Tanimoto similarity of the ﬁngerprints constrained within δ ∈ { . 4 , .
Revised: Speciﬁcally, the task is to translate these molecules into new molecules with the Tanimoto similarity of the ﬁngerprints constrained within δ ∈ { . 4 , .

Original: This task aims for optimizing a certain quantity (instead of satisfying property constraints) and is a translation task (need to stay close to original molecules rather than ﬁnding novel ones).
Revised: This task aims for optimizing a certain quantity (instead of satisfying property constraints) and is a translation task (need to stay close to original molecules rather than ﬁnding novel ones).

Original: To run FaST on this task, we apply the following changes to the reward function, the qualiﬁcation criterion, and the episode termination criterion, of FaST.
Revised: To run FaST on this task, we apply the following changes to the reward function, the qualiﬁcation criterion, and the episode termination criterion, of FaST.

Original: We denote score ( x ) to be the penalized logP scoring function, and sim ( · , · ) to be the Tanimoto similarity between two molecules:
Revised: We denote score ( x ) to be the penalized logP scoring function, and sim ( · , · ) to be the Tanimoto similarity between two molecules:

Original: • reward r = score ( x j ) − score ( x i ) for any transition from molecule x i → x j • C , the discovered set contains all explored molecules that satisfy Equation (5), where the threshold is given by the input parameter δ • We terminate an episode when the number of steps exceeds 10.
Revised: • reward r = score ( x j ) − score ( x i ) for any transition from molecule x i → x j • C , the discovered set contains all explored molecules that satisfy Equation (5), where the threshold is given by the input parameter δ • We terminate an episode when the number of steps exceeds 10.

Original: For each molecule we add to G , we keep track of its original parent (the molecule from the 800 test molecules).
Revised: For each molecule we add to G , we keep track of its original parent (the molecule from the 800 test molecules).

Original: After training, for each of the 800 test molecules, we take the set of translated molecules in G , and select the one with the highest property score.
Revised: After training, for each of the 800 test molecules, we take the set of translated molecules in G , and select the one with the highest property score.

Original: [Table]
Revised: [Table]

Original: The model generates a molecule with repeating aromatic rings; though not realistic, this molecule achieves a high score, while having close Tanimoto similarity using Morgan ﬁngerprints.
Revised: The model generates a molecule with repeating aromatic rings; though not realistic, this molecule achieves a high score, while having close Tanimoto similarity using Morgan ﬁngerprints.

Original: Results are shown in Table 2; our method greatly outperforms the other baselines, but we point out a few ﬂaws intrinsic to the task.
Revised: Results are shown in Table 2; our method greatly outperforms the other baselines, but we point out a few ﬂaws intrinsic to the task.

Original: Because the similarity is computed through Morgan ﬁngerprint, which are hashes of substructures, repeatedly adding aromatic rings can often not change the ﬁngerprint by a lot.
Revised: Because the similarity is computed through Morgan ﬁngerprint, which are hashes of substructures, repeatedly adding aromatic rings can often not change the ﬁngerprint by a lot.

Original: Nevertheless, adding aromatic rings will linearly increase penalized logP score, which allows trivial solutions to produce high scores for this task (see Figure 5).
Revised: Nevertheless, adding aromatic rings will linearly increase penalized logP score, which allows trivial solutions to produce high scores for this task (see Figure 5).

Original: This phenomenon is noted by
Revised: This phenomenon is noted by

Original: Nigam et al.
Revised: Nigam et al.

Original: but they add a regularizer to constrain the generated compounds to look similar to the reference molecules.
Revised: but they add a regularizer to constrain the generated compounds to look similar to the reference molecules.

Original: Due to the mentioned issues, we believe this task can be reformulated.
Revised: Due to the mentioned issues, we believe this task can be reformulated.

Original: For instance, one could use a different ﬁngerprint method so that the ﬁngerprint similarity is not so easily exploited (see AP (Carhart et al., 1985), MACCS (Durant et al., 2002), or ROCS (Hawkins et al., 2010)), or size constraints should be incorporated.
Revised: For instance, one could use a different ﬁngerprint method so that the ﬁngerprint similarity is not so easily exploited (see AP (Carhart et al., 1985), MACCS (Durant et al., 2002), or ROCS (Hawkins et al., 2010)), or size constraints should be incorporated.

Original: Nevertheless, we provide our results for comparison to other molecular generation methods.
Revised: Nevertheless, we provide our results for comparison to other molecular generation methods.

Original: In general, the task of optimizing (increasing) the penalized logP scores is not entirely meaningful.
Revised: In general, the task of optimizing (increasing) the penalized logP scores is not entirely meaningful.

Original: According to Lipinski’s rule of ﬁve (Lipinski et al., 1997), which are widely established rules to evaluate the druglikeness of molecules, the logP score should be lower than 5 .
Revised: According to Lipinski’s rule of ﬁve (Lipinski et al., 1997), which are widely established rules to evaluate the druglikeness of molecules, the logP score should be lower than 5 .

Original: So an unbounded optimization of logP has little practical usability.
Revised: So an unbounded optimization of logP has little practical usability.

Original: Perhaps a better task would be to optimize for all 5 rules in Linpinski’s rule of ﬁve which includes constraints involving the number of hydrogen bond donors/acceptors and molecular mass.
Revised: Perhaps a better task would be to optimize for all 5 rules in Linpinski’s rule of ﬁve which includes constraints involving the number of hydrogen bond donors/acceptors and molecular mass.

Original: C D IFFERENT N OVELTY /D IVERSITY M ETRICS
Revised: C D IFFERENT N OVELTY /D IVERSITY M ETRICS

Original: FaST is capable of optimizing for different novelty/diversity metrics.
Revised: FaST is capable of optimizing for different novelty/diversity metrics.

Original: In this section, we compute the novelty/diversity metrics using atom-pair (AP) ﬁngerprints (Carhart et al., 1985).
Revised: In this section, we compute the novelty/diversity metrics using atom-pair (AP) ﬁngerprints (Carhart et al., 1985).

Original: While Morgan ﬁngerprints have successfully been applied to many molecular tasks such as drug screening, it has some failure modes (Capecchi et al., 2020).
Revised: While Morgan ﬁngerprints have successfully been applied to many molecular tasks such as drug screening, it has some failure modes (Capecchi et al., 2020).

Original: Namely, Morgan ﬁngerprints is often not informative about the size or the shape of the molecules.
Revised: Namely, Morgan ﬁngerprints is often not informative about the size or the shape of the molecules.

Original: These properties are better captured in AP ﬁngerprints, as AP ﬁngerprints account for all atom pairs, including their pairwise distances.
Revised: These properties are better captured in AP ﬁngerprints, as AP ﬁngerprints account for all atom pairs, including their pairwise distances.

Original: We run the same experiment on the GSK3 β +JNK3+QED+SA task described in Section 5, but change the ﬁngerprint from Morgan to AP for the novelty/diversity metrics.
Revised: We run the same experiment on the GSK3 β +JNK3+QED+SA task described in Section 5, but change the ﬁngerprint from Morgan to AP for the novelty/diversity metrics.

Original: The results are shown in Table 3 with comparison to baselines.
Revised: The results are shown in Table 3 with comparison to baselines.

Original: We observe that our method outperform baselines by a greater margin, especially in the novelty metric.
Revised: We observe that our method outperform baselines by a greater margin, especially in the novelty metric.

Original: This is not surprising because our model can explicitly optimize for any similarity metric, while the baseline methods are not novelty/diversity-aware during training.
Revised: This is not surprising because our model can explicitly optimize for any similarity metric, while the baseline methods are not novelty/diversity-aware during training.

Original: Interestingly, we ﬁnd that optimizing for AP ﬁngerprints also yields molecules that score high under
Revised: Interestingly, we ﬁnd that optimizing for AP ﬁngerprints also yields molecules that score high under

Original: Morgan ﬁngerprints for this task (but the converse is not true)
Revised: Morgan ﬁngerprints for this task (but the converse is not true)

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: D I MPLEMENTATION D ETAILS
Revised: D I MPLEMENTATION D ETAILS

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: E S AMPLE TRAJECTORIES AND OPTIMIZED MOLECULES FROM F A ST
Revised: E S AMPLE TRAJECTORIES AND OPTIMIZED MOLECULES FROM F A ST

Original: We provide more example molecular optimization trajectories of our model on the
Revised: We provide more example molecular optimization trajectories of our model on the

Original: GSK3 β +JNK3+QED+SA task in Figure 6.
Revised: GSK3 β +JNK3+QED+SA task in Figure 6.

Original: Additional generated moles for the GSK3 β +JNK3+QED+SA task are shown in Figure
Revised: Additional generated moles for the GSK3 β +JNK3+QED+SA task are shown in Figure

Original: F D ECODER A RCHITECTURE : SELFIES VS SMILES
Revised: F D ECODER A RCHITECTURE : SELFIES VS SMILES

Original: For our VAE, we decode into string representations of fragments for several reasons.
Revised: For our VAE, we decode into string representations of fragments for several reasons.

Original: (1) Graph decoders are complex and requires a lot of hand engineered rules to work.
Revised: (1) Graph decoders are complex and requires a lot of hand engineered rules to work.

Original: In fact, many graph decoders require their own vocabularies, which makes the decoder architecture even more complex.
Revised: In fact, many graph decoders require their own vocabularies, which makes the decoder architecture even more complex.

Original: (2) Our fragments are small, less than 10 atoms, so a simpler representation of the fragments in string form is appropriate.
Revised: (2) Our fragments are small, less than 10 atoms, so a simpler representation of the fragments in string form is appropriate.

Original: Additionally, we choose to use SELFIES strings over SMILES strings, because SELFIES strings have a more robust grammar compared to SMILES strings, and has been shown through many past experiments (Krenn et al., 2020).
Revised: Additionally, we choose to use SELFIES strings over SMILES strings, because SELFIES strings have a more robust grammar compared to SMILES strings, and has been shown through many past experiments (Krenn et al., 2020).

Original: Empirically, we see in Figure 8 that
Revised: Empirically, we see in Figure 8 that

Original: SELFIES strings gives better reconstruction accuracy compared to using SMILES strings for the decoder in the VAE architecture.
Revised: SELFIES strings gives better reconstruction accuracy compared to using SMILES strings for the decoder in the VAE architecture.

Original: [Figure]
Revised: [Figure]