Original: S PANNING T REE - BASED G RAPH G ENERATION FOR M OLECULES
Revised: S PANNING T REE - BASED G RAPH G ENERATION FOR M OLECULES

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: In this paper, we explore the problem of generating molecules using deep neural networks, which has recently gained much interest in chemistry.
Revised: In this paper, we explore the problem of generating molecules using deep neural networks, which has recently gained much interest in chemistry.

Original: To this end, we propose a spanning tree-based graph generation (STGG) framework based on formulating molecular graph generation as a construction of a spanning tree and the residual edges.
Revised: To this end, we propose a spanning tree-based graph generation (STGG) framework based on formulating molecular graph generation as a construction of a spanning tree and the residual edges.

Original: Such a formulation exploits the sparsity of molecular graphs and allows using compact tree-constructive operations to deﬁne the molecular graph connectivity.
Revised: Such a formulation exploits the sparsity of molecular graphs and allows using compact tree-constructive operations to deﬁne the molecular graph connectivity.

Original: Based on the intermediate graph structure of the construction process, our framework can constrain its generation to molecular graphs that satisfy the chemical valence rules.
Revised: Based on the intermediate graph structure of the construction process, our framework can constrain its generation to molecular graphs that satisfy the chemical valence rules.

Original: We also newly design a Transformer architecture with tree-based relative positional encodings for realizing the tree construction procedure.
Revised: We also newly design a Transformer architecture with tree-based relative positional encodings for realizing the tree construction procedure.

Original: Experiments on QM9, ZINC250k, and MOSES benchmarks verify the effectiveness of the proposed framework in metrics such as validity, Fr ´ echet
Revised: Experiments on QM9, ZINC250k, and MOSES benchmarks verify the effectiveness of the proposed framework in metrics such as validity, Fr ´ echet

Original: ChemNet distance, and fragment similarity.
Revised: ChemNet distance, and fragment similarity.

Original: We also demonstrate the usefulness of
Revised: We also demonstrate the usefulness of

Original: STGG in maximizing penalized LogP value of molecules.
Revised: STGG in maximizing penalized LogP value of molecules.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Researchers have extensively studied graph generative models, dating back to the early works of Erd ¨ os
Revised: Researchers have extensively studied graph generative models, dating back to the early works of Erd ¨ os

Original: R ´ enyi (Erd ¨ os et al., 1959).
Revised: R ´ enyi (Erd ¨ os et al., 1959).

Original: Recently, models based on deep neural networks (DNNs) have gained much attraction due to their expressive power in learning a graph dataset.
Revised: Recently, models based on deep neural networks (DNNs) have gained much attraction due to their expressive power in learning a graph dataset.

Original: The molecule-generating
Revised: The molecule-generating

Original: DNNs stand out among them for their success in the task of drug discovery.
Revised: DNNs stand out among them for their success in the task of drug discovery.

Original: Recent works have proposed molecule-generating DNNs based on string-based and graph-based representations (Segler et al., 2018; Jin et al., 2018; You et al., 2018; Shi et al., 2020; Jin et al., 2020).
Revised: Recent works have proposed molecule-generating DNNs based on string-based and graph-based representations (Segler et al., 2018; Jin et al., 2018; You et al., 2018; Shi et al., 2020; Jin et al., 2020).

Original: For example, Segler et al.
Revised: For example, Segler et al.

Original: (2018) proposed to train language models on the domain-speciﬁc linear string representation of molecules, i.e., simpliﬁed molecular-input line-entry system (SMILES,
Revised: (2018) proposed to train language models on the domain-speciﬁc linear string representation of molecules, i.e., simpliﬁed molecular-input line-entry system (SMILES,

Original: Weininger 1988)
Revised: Weininger 1988)

Original: Since the string-based models ignore the inherent graph structure, recent works explore the graph-based generation that use (a) atom-by-atom (You et al., 2018; Shi et al., 2020; Luo et al., 2021) or (b) substructure-based (Jin et al., 2018; 2019; 2020) operations.
Revised: Since the string-based models ignore the inherent graph structure, recent works explore the graph-based generation that use (a) atom-by-atom (You et al., 2018; Shi et al., 2020; Luo et al., 2021) or (b) substructure-based (Jin et al., 2018; 2019; 2020) operations.

Original: Notably, the substructure-based generative models (Jin et al., 2018; 2019; 2020) successfully exploit the molecular prior knowledge: the graphs are sparsely connected and can be represented as a junction tree with molecular substructure as building blocks.
Revised: Notably, the substructure-based generative models (Jin et al., 2018; 2019; 2020) successfully exploit the molecular prior knowledge: the graphs are sparsely connected and can be represented as a junction tree with molecular substructure as building blocks.

Original: Based on such knowledge, the models use the junction tree construction operators which (a) require a fewer number of steps to generate the whole molecular graph and (b) guarantee generating molecules that satisfy the chemical valence rules.
Revised: Based on such knowledge, the models use the junction tree construction operators which (a) require a fewer number of steps to generate the whole molecular graph and (b) guarantee generating molecules that satisfy the chemical valence rules.

Original: However, despite such advantages, a recent benchmark (Polykovskiy et al., 2020) suggests that they do not outperform the existing methods in terms of learning the data distribution, even when compared with the simple SMILES-based language models.
Revised: However, despite such advantages, a recent benchmark (Polykovskiy et al., 2020) suggests that they do not outperform the existing methods in terms of learning the data distribution, even when compared with the simple SMILES-based language models.

Original: We hypothesize that this is due to the models using a coarse-grained representation of the molecule and they may lack the ability to learn the inner semantics of each substructure-based building block.
Revised: We hypothesize that this is due to the models using a coarse-grained representation of the molecule and they may lack the ability to learn the inner semantics of each substructure-based building block.

Original: Contribution.
Revised: Contribution.

Original: In this work, we propose a novel framework, coined spanning tree-based graph generation (STGG), for ﬁne-grained generation of molecules while exploiting their sparsity.
Revised: In this work, we propose a novel framework, coined spanning tree-based graph generation (STGG), for ﬁne-grained generation of molecules while exploiting their sparsity.

Original: 1 Mainly inspired from the SMILES representation of molecules, our idea is to generate the molecular graph as a composition of a spanning tree and the corresponding residual edges with atoms and bonds as building blocks.
Revised: 1 Mainly inspired from the SMILES representation of molecules, our idea is to generate the molecular graph as a composition of a spanning tree and the corresponding residual edges with atoms and bonds as building blocks.

Original: Such a formulation allows our framework to utilize compact tree-constructive operations to deﬁne the molecular graph connectivity.
Revised: Such a formulation allows our framework to utilize compact tree-constructive operations to deﬁne the molecular graph connectivity.

Original: See Figure 1 for an illustration of how we formulate the generation of a molecular graph as a sequence of tree-constructive operations.
Revised: See Figure 1 for an illustration of how we formulate the generation of a molecular graph as a sequence of tree-constructive operations.

Original: Since our framework maintains the molecular graph structure during construction, it can pre-determine decisions that (a) violate the graph construction rule and (b) lead to molecules that violate the chemical valence rule.
Revised: Since our framework maintains the molecular graph structure during construction, it can pre-determine decisions that (a) violate the graph construction rule and (b) lead to molecules that violate the chemical valence rule.

Original: Such criteria allow control over the generative model to guarantee generating valid molecular graphs by forbidding invalid actions.
Revised: Such criteria allow control over the generative model to guarantee generating valid molecular graphs by forbidding invalid actions.

Original: This is in contrast to prior works (Shi et al., 2020;
Revised: This is in contrast to prior works (Shi et al., 2020;

Original: Luo et al., 2021) that generate the molecular graph atom-by-atom but determines the validity of construction operations through a sample-rejection scheme.
Revised: Luo et al., 2021) that generate the molecular graph atom-by-atom but determines the validity of construction operations through a sample-rejection scheme.

Original: To recognize the spanning tree-based representation used in our STGG framework, we propose a
Revised: To recognize the spanning tree-based representation used in our STGG framework, we propose a

Original: Transformer architecture (Vaswani et al., 2017) with tree-based relative encoding.
Revised: Transformer architecture (Vaswani et al., 2017) with tree-based relative encoding.

Original: Inspired by recent works (Villmow et al., 2021; Lukovnikov & Fischer, 2021; Ying et al., 2021) on tree-based and graph-based Transformers, our framework expresses the relative position between two vertices as the number of forward and reverse edges in the shortest path between them.
Revised: Inspired by recent works (Villmow et al., 2021; Lukovnikov & Fischer, 2021; Ying et al., 2021) on tree-based and graph-based Transformers, our framework expresses the relative position between two vertices as the number of forward and reverse edges in the shortest path between them.

Original: We also introduce an attention-based mechanism for constructing residual edges.
Revised: We also introduce an attention-based mechanism for constructing residual edges.

Original: We experiment on popular graph generation benchmarks of QM9, ZINC250 K , and MOSES to validate the effectiveness of our algorithm.
Revised: We experiment on popular graph generation benchmarks of QM9, ZINC250 K , and MOSES to validate the effectiveness of our algorithm.

Original: In the experiments on QM9 and ZINC, our STGG framework outperforms the existing graph-based generative models by a large margin.
Revised: In the experiments on QM9 and ZINC, our STGG framework outperforms the existing graph-based generative models by a large margin.

Original: In the MOSES benchmark, our algorithm achieves superior performance compared to both string-based and graphbased methods for majority of the metrics, e.g., Fr ´ echet ChemNet distance (Preuer et al., 2018) and fragment-based similarity.
Revised: In the MOSES benchmark, our algorithm achieves superior performance compared to both string-based and graphbased methods for majority of the metrics, e.g., Fr ´ echet ChemNet distance (Preuer et al., 2018) and fragment-based similarity.

Original: We also conduct experiments on the ofﬂine optimization task for high penalized octanol-water partition coefﬁcient and achieve competitive results.
Revised: We also conduct experiments on the ofﬂine optimization task for high penalized octanol-water partition coefﬁcient and achieve competitive results.

Original:  S PANNING T REE - BASED G ENERATION OF G RAPHS (STGG) 2. O VERVIEW
Revised:  S PANNING T REE - BASED G ENERATION OF G RAPHS (STGG) 2. O VERVIEW

Original: In this section, we introduce our spanning tree-based graph generation (STGG) framework to sequentially generate a molecule as a composition of a spanning tree and residual edges.
Revised: In this section, we introduce our spanning tree-based graph generation (STGG) framework to sequentially generate a molecule as a composition of a spanning tree and residual edges.

Original: To this end, we propose compact tree-constructive operations inspired by the simpliﬁed molecular-input line-entry system (SMILES, Weininger, 1988).
Revised: To this end, we propose compact tree-constructive operations inspired by the simpliﬁed molecular-input line-entry system (SMILES, Weininger, 1988).

Original: In contrast to the existing SMILES-based molecular generative methods, our framework (a) allows inferring the intermediate graph structure and (b) is generally applicable to graph types other than molecules.
Revised: In contrast to the existing SMILES-based molecular generative methods, our framework (a) allows inferring the intermediate graph structure and (b) is generally applicable to graph types other than molecules.

Original: In particular, (a) further enables our framework to control the construction process such that the sequential operations comply with tree-constructive grammar and only generate molecules satisfying the chemical valence rule.
Revised: In particular, (a) further enables our framework to control the construction process such that the sequential operations comply with tree-constructive grammar and only generate molecules satisfying the chemical valence rule.

Original: Molecular graph representation.
Revised: Molecular graph representation.

Original: To apply our framework, we represent a molecule as a bipartite graph G =
Revised: To apply our framework, we represent a molecule as a bipartite graph G =

Original: ( A , B , E ) where A and B are the set of vertices associated with atoms and bonds of the molecule, respectively.
Revised: ( A , B , E ) where A and B are the set of vertices associated with atoms and bonds of the molecule, respectively.

Original: Each edge { a, b } ∈ E is assigned for each adjacent pair of atom and bond.
Revised: Each edge { a, b } ∈ E is assigned for each adjacent pair of atom and bond.

Original: We assign attributes x a ∈ X atom and x b ∈ X bond for vertices a ∈ A and b ∈ B to indicate the corresponding atom type and bond order, respectively.
Revised: We assign attributes x a ∈ X atom and x b ∈ X bond for vertices a ∈ A and b ∈ B to indicate the corresponding atom type and bond order, respectively.

Original: For example, { "C" , "N" , "O" } ⊆
Revised: For example, { "C" , "N" , "O" } ⊆

Original: X atom and { "-" , "=" } ⊆ X bond .
Revised: X atom and { "-" , "=" } ⊆ X bond .

Original: See Figure 1 for an example of such a molecular graph representation.
Revised: See Figure 1 for an example of such a molecular graph representation.

Original: Molecular graph from sequence of decisions.
Revised: Molecular graph from sequence of decisions.

Original: To generate the molecular graph G =
Revised: To generate the molecular graph G =

Original: ( A , B , E ) , our framework makes a sequence of decisions d 1 , . . .
Revised: ( A , B , E ) , our framework makes a sequence of decisions d 1 , . . .

Original: , d T to generate a spanning tree T
Revised: , d T to generate a spanning tree T

Original: = ( A T , B T , E T ) and a set of residual edges E R =
Revised: = ( A T , B T , E T ) and a set of residual edges E R =

Original: E \E T .
Revised: E \E T .

Original: At each iteration, seven types of decisions are applicable, i.e., attach atom , attach bond , branch start , branch end , res atom , res bond , and terminate .
Revised: At each iteration, seven types of decisions are applicable, i.e., attach atom , attach bond , branch start , branch end , res atom , res bond , and terminate .

Original: [Table]
Revised: [Table]

Original: See Table 1 for examples of decisions and the corresponding operations.
Revised: See Table 1 for examples of decisions and the corresponding operations.

Original: We provide a detailed description of the graph construction process in Section 2.2.
Revised: We provide a detailed description of the graph construction process in Section 2.2.

Original: Generating valid molecular graphs.
Revised: Generating valid molecular graphs.

Original: Without any control, a model may generate decisions that (a) do not comply with the grammar of STGG or (b) leads to a molecule violating the chemical valence rule.
Revised: Without any control, a model may generate decisions that (a) do not comply with the grammar of STGG or (b) leads to a molecule violating the chemical valence rule.

Original: To prevent this scenario, we conduct two criteria for determining validity of the given decision for (a) and (b).
Revised: To prevent this scenario, we conduct two criteria for determining validity of the given decision for (a) and (b).

Original: We further elaborate this in Section 2.3.
Revised: We further elaborate this in Section 2.3.

Original: 2. D ECISION PROCESS FOR SPANNING TREE - BASED GRAPH GENERATION
Revised: 2. D ECISION PROCESS FOR SPANNING TREE - BASED GRAPH GENERATION

Original: We now explain how our STGG framework incorporates the decisions d 1 , . . .
Revised: We now explain how our STGG framework incorporates the decisions d 1 , . . .

Original: , d T to build the spanning tree T =
Revised: , d T to build the spanning tree T =

Original: ( A T , B T , E T ) and residual edges E R from scratch.
Revised: ( A T , B T , E T ) and residual edges E R from scratch.

Original: To this end, our framework introduces the state information of (a) a pointer vertex i point ∈
Revised: To this end, our framework introduces the state information of (a) a pointer vertex i point ∈

Original: A T ∪ B T for specifying the target of the next operation, (b) a stack S branch that stores vertices to use later as a starting point of a “branch” in the spanning tree and (c) a list L res that stores vertices to use later for constructing residual edges.
Revised: A T ∪ B T for specifying the target of the next operation, (b) a stack S branch that stores vertices to use later as a starting point of a “branch” in the spanning tree and (c) a list L res that stores vertices to use later for constructing residual edges.

Original: In what follows, we describe the seven types of operations, i.e., attach atom , attach bond , branch start , branch end , res atom , res bond , and terminate , corresponding to decision values d ∈ X atom ∪ X bond ∪ L res ∪ { "(" , ")" , "*" , "[eos]" } in detail.
Revised: In what follows, we describe the seven types of operations, i.e., attach atom , attach bond , branch start , branch end , res atom , res bond , and terminate , corresponding to decision values d ∈ X atom ∪ X bond ∪ L res ∪ { "(" , ")" , "*" , "[eos]" } in detail.

Original: See Table 1 for the pairs of operations and the corresponding decisions.
Revised: See Table 1 for the pairs of operations and the corresponding decisions.

Original: We also provide an example of the graph construction process in Figure 2.
Revised: We also provide an example of the graph construction process in Figure 2.

Original: Attaching atom and bond vertices to the spanning tree.
Revised: Attaching atom and bond vertices to the spanning tree.

Original: If the decision d speciﬁes one of the atom or bond attributes, i.e., d ∈ X atom or d ∈ X bond , it applies the corresponding attach atom and attach bond operations, respectively.
Revised: If the decision d speciﬁes one of the atom or bond attributes, i.e., d ∈ X atom or d ∈ X bond , it applies the corresponding attach atom and attach bond operations, respectively.

Original: To be speciﬁc, the attach atom operation adds a new atom vertex a into the spanning tree T as a neighbor of the pointer vertex i pointer , i.e., A T ←
Revised: To be speciﬁc, the attach atom operation adds a new atom vertex a into the spanning tree T as a neighbor of the pointer vertex i pointer , i.e., A T ←

Original: A T ∪ { a } , E T ← E T ∪ { a, i pointer } .
Revised: A T ∪ { a } , E T ← E T ∪ { a, i pointer } .

Original: The value d is set as the new atom attribute, i.e., x a ← d .
Revised: The value d is set as the new atom attribute, i.e., x a ← d .

Original: The newly added vertex is set as the next pointer vertex, i.e., i pointer ← a .
Revised: The newly added vertex is set as the next pointer vertex, i.e., i pointer ← a .

Original: The attach bond operation similarly adds a new bond vertex.
Revised: The attach bond operation similarly adds a new bond vertex.

Original: For example, a line graph can be expressed as a sequence of attach atom and attach bond operations, e.g., C-C-C where "C" ∈ X A and "-" ∈ X B .
Revised: For example, a line graph can be expressed as a sequence of attach atom and attach bond operations, e.g., C-C-C where "C" ∈ X A and "-" ∈ X B .

Original: Branching out the spanning tree.
Revised: Branching out the spanning tree.

Original: To express graph structures with vertices of degree larger than two, our framework utilizes pairs of the branch start and the branch end operations with decision values of "(" and ")" , respectively.
Revised: To express graph structures with vertices of degree larger than two, our framework utilizes pairs of the branch start and the branch end operations with decision values of "(" and ")" , respectively.

Original: To be speciﬁc, the branch start operation inserts the current pointer vertex into a stack S branch of vertices.
Revised: To be speciﬁc, the branch start operation inserts the current pointer vertex into a stack S branch of vertices.

Original: Then the branch end operation pops a vertex from the stack S branch and sets it as the new pointer vertex.
Revised: Then the branch end operation pops a vertex from the stack S branch and sets it as the new pointer vertex.

Original: For an example, a graph with one atom vertex of degree three is constructed from a sequence of decisions C-C(-C)(-C) .
Revised: For an example, a graph with one atom vertex of degree three is constructed from a sequence of decisions C-C(-C)(-C) .

Original: Adding residual edges.
Revised: Adding residual edges.

Original: To construct cyclic molecular graphs, our framework generates residual edges based on pairs of res atom operation and res bond operation, corresponding to decision values of "*" and d ∈ L res , respectively.
Revised: To construct cyclic molecular graphs, our framework generates residual edges based on pairs of res atom operation and res bond operation, corresponding to decision values of "*" and d ∈ L res , respectively.

Original: To be speciﬁc, the res atom operation inserts the current (atom) pointer vertex into a list L res .
Revised: To be speciﬁc, the res atom operation inserts the current (atom) pointer vertex into a list L res .

Original: Next, when the decision value d ∈ L res is received for the res bond operation, the corresponding vertex d is popped from the list L res and forms a new residual edge with the current (bond) pointer vertex, i.e., E R ← E R ∪ {{ d, i pointer }} .
Revised: Next, when the decision value d ∈ L res is received for the res bond operation, the corresponding vertex d is popped from the list L res and forms a new residual edge with the current (bond) pointer vertex, i.e., E R ← E R ∪ {{ d, i pointer }} .

Original: For an example, a cyclic molecular graph is constructed from a sequence of decisions C*-O-O-1 , where "1" indicates res bond operation with decision of the ﬁrst atom vertex with attribute "C" .
Revised: For an example, a cyclic molecular graph is constructed from a sequence of decisions C*-O-O-1 , where "1" indicates res bond operation with decision of the ﬁrst atom vertex with attribute "C" .

Original: Termination.
Revised: Termination.

Original: The decision "[eos]" applies the terminate operation to ﬁnish the construction.
Revised: The decision "[eos]" applies the terminate operation to ﬁnish the construction.

Original: We provide the full algorithm in Algorithm 1.
Revised: We provide the full algorithm in Algorithm 1.

Original: We also provide an algorithm to extract a sequence of decisions for constructing a given graph in Appendix A.
Revised: We also provide an algorithm to extract a sequence of decisions for constructing a given graph in Appendix A.

Original: Such an algorithm is used to obtain sequence of decisions as targets for training the generative model under the STGG framework.
Revised: Such an algorithm is used to obtain sequence of decisions as targets for training the generative model under the STGG framework.

Original: [Figure]
Revised: [Figure]

Original: 2. M ASKING OUT INVALID DECISIONS FOR A VALID MOLECULAR GRAPH
Revised: 2. M ASKING OUT INVALID DECISIONS FOR A VALID MOLECULAR GRAPH

Original: Based on Algorithm 1, we develop two criteria for determining whether if a sequence of decisions leads to (a) valid generation of a molecular graph and (b) generation of a molecule satisfying the valence rule.
Revised: Based on Algorithm 1, we develop two criteria for determining whether if a sequence of decisions leads to (a) valid generation of a molecular graph and (b) generation of a molecule satisfying the valence rule.

Original: Such criteria are used to mask out invalid decisions to guarantee generating a valid molecular graph.
Revised: Such criteria are used to mask out invalid decisions to guarantee generating a valid molecular graph.

Original: Validity of graph generation.
Revised: Validity of graph generation.

Original: To determine whether if a sequence of decisions lead to a valid generation of a molecular graph, we propose an algorithm that outputs a set of valid decisions given the previous decision d , stack of pointer vertices S branch , and list of atom vertices L res during execution of Algorithhm 1.
Revised: To determine whether if a sequence of decisions lead to a valid generation of a molecular graph, we propose an algorithm that outputs a set of valid decisions given the previous decision d , stack of pointer vertices S branch , and list of atom vertices L res during execution of Algorithhm 1.

Original: In what follows, we provide a brief description of grammars enforced by the algorithm.
Revised: In what follows, we provide a brief description of grammars enforced by the algorithm.

Original: We also provide the detailed algorithm in Appendix B.
Revised: We also provide the detailed algorithm in Appendix B.

Original: • The branch end operation only appears when the stack of pointer vertex S branch is non-empty.
Revised: • The branch end operation only appears when the stack of pointer vertex S branch is non-empty.

Original: • The operations res atom and res bond are atom-speciﬁc and bond-speciﬁc, hence they only appear when the pointer vertex is located at an atom vertex and a bond vertex, respectively.
Revised: • The operations res atom and res bond are atom-speciﬁc and bond-speciﬁc, hence they only appear when the pointer vertex is located at an atom vertex and a bond vertex, respectively.

Original: • All the bond vertices have degree of two, hence branch start and branch end operations only appear when the pointer vertex is located at an atom vertex.
Revised: • All the bond vertices have degree of two, hence branch start and branch end operations only appear when the pointer vertex is located at an atom vertex.

Original: • The stack do not contain duplicates of the a pointer vertex at the same time.
Revised: • The stack do not contain duplicates of the a pointer vertex at the same time.

Original: • The terminate operation is only allowed when the stack S branch and the list L res is empty.
Revised: • The terminate operation is only allowed when the stack S branch and the list L res is empty.

Original: Validity of satisfying the valence rule.
Revised: Validity of satisfying the valence rule.

Original: To consider chemical validity of molecules, our framework offers the ability to constrain the its generation on molecules that satisfy the valence rule for each atom.
Revised: To consider chemical validity of molecules, our framework offers the ability to constrain the its generation on molecules that satisfy the valence rule for each atom.

Original: That is, the generated graph G =
Revised: That is, the generated graph G =

Original: ( A , B , E ) satisﬁes the constraint v ( x a ) ≥ (cid:80) b ∈N ( a ) o ( x b ) for every atom vertex a ∈ A where v ( x a ) denotes the valence of an atom type x a and o ( x b ) denotes the bond order.
Revised: ( A , B , E ) satisﬁes the constraint v ( x a ) ≥ (cid:80) b ∈N ( a ) o ( x b ) for every atom vertex a ∈ A where v ( x a ) denotes the valence of an atom type x a and o ( x b ) denotes the bond order.

Original: To this end, we keep a record r ( a ) of available valence for each atom a ∈ A and update them for each decision.
Revised: To this end, we keep a record r ( a ) of available valence for each atom a ∈ A and update them for each decision.

Original: For example, when a bond vertex b is newly added, record of the neighboring atom vertex a is updated by r ( a ) ← v ( a ) − o ( x b ) .
Revised: For example, when a bond vertex b is newly added, record of the neighboring atom vertex a is updated by r ( a ) ← v ( a ) − o ( x b ) .

Original: The main idea is to forbid actions that lead to negative values of r ( a ) .
Revised: The main idea is to forbid actions that lead to negative values of r ( a ) .

Original: We provide a detailed algorithm in Appendix C.
Revised: We provide a detailed algorithm in Appendix C.

Original:  T RANSFORMER A RCHITECTURE FOR T REE - BASED G ENERATION
Revised:  T RANSFORMER A RCHITECTURE FOR T REE - BASED G ENERATION

Original: In this section, we describe our deep neural network architecture for generating sequence of decisions d 1 , . . .
Revised: In this section, we describe our deep neural network architecture for generating sequence of decisions d 1 , . . .

Original: , d T under the STGG framework.
Revised: , d T under the STGG framework.

Original: To accurately recognize the decision process, we employ the tree-based relative positional encodings on the intermediate spanning tree T .
Revised: To accurately recognize the decision process, we employ the tree-based relative positional encodings on the intermediate spanning tree T .

Original: We also introduce an attention mechanism to express a probability distribution over L res which depends on intermediate state of the algorithm.
Revised: We also introduce an attention mechanism to express a probability distribution over L res which depends on intermediate state of the algorithm.

Original: 3. T REE - BASED POSITIONAL ENCODING FOR MULTI - HEAD ATTENTION LAYERS
Revised: 3. T REE - BASED POSITIONAL ENCODING FOR MULTI - HEAD ATTENTION LAYERS

Original: Each intermediate layer in our model is a combination of a multi-head self-attention module and a position-wise feed-forward neural network similar to that of Vaswani et al.
Revised: Each intermediate layer in our model is a combination of a multi-head self-attention module and a position-wise feed-forward neural network similar to that of Vaswani et al.

Original: The main difference is on how we modify the architecture to incorporate tree-based positional encodings.
Revised: The main difference is on how we modify the architecture to incorporate tree-based positional encodings.

Original: To be speciﬁc, let H =
Revised: To be speciﬁc, let H =

Original: [ h (cid:62) 1 , . . .
Revised: [ h (cid:62) 1 , . . .

Original: , h (cid:62) T ]
Revised: , h (cid:62) T ]

Original: ∈ R T × (cid:96) denote the input of a self-attention module where d is the hidden dimension and h t ∈ R 1 × (cid:96) is the hidden representation at position t .
Revised: ∈ R T × (cid:96) denote the input of a self-attention module where d is the hidden dimension and h t ∈ R 1 × (cid:96) is the hidden representation at position t .

Original: The input H is projected by three matrices W Q ∈ R (cid:96) × (cid:96) K , W K ∈ R (cid:96) × (cid:96) K and W V ∈ R (cid:96) × (cid:96) V to the corresponding representations Q, K and V , respectively.
Revised: The input H is projected by three matrices W Q ∈ R (cid:96) × (cid:96) K , W K ∈ R (cid:96) × (cid:96) K and W V ∈ R (cid:96) × (cid:96) V to the corresponding representations Q, K and V , respectively.

Original: A single self-attention head is then calculated as
Revised: A single self-attention head is then calculated as

Original: [Equation]
Revised: [Equation]

Original: where Attention ( H ) is output of the attention head, M is the triangular mask to forbid the model from accessing future information while making a prediction, and ◦ denotes the element-wise multiplication between matrices.
Revised: where Attention ( H ) is output of the attention head, M is the triangular mask to forbid the model from accessing future information while making a prediction, and ◦ denotes the element-wise multiplication between matrices.

Original: Furthermore, P is the newly introduced relative positional encoding.
Revised: Furthermore, P is the newly introduced relative positional encoding.

Original: It is a summation over the trainable embedding vectors z (1) , z (2) , z (3) indexed by relative position values of φ forward ( t 1 , t 2 ) , φ backward ( t 1 , t 2 ) , and φ seq ( t 1 , t 2 ) .
Revised: It is a summation over the trainable embedding vectors z (1) , z (2) , z (3) indexed by relative position values of φ forward ( t 1 , t 2 ) , φ backward ( t 1 , t 2 ) , and φ seq ( t 1 , t 2 ) .

Original: To be speciﬁc, the tree-based relative positions φ forward ( t 1 , t 2 ) and φ backward ( t 1 , t 2 ) denotes the number of forward and backward edges in the spanning tree path between pointer vertices at the t 1 -th and t 2 -th time step.
Revised: To be speciﬁc, the tree-based relative positions φ forward ( t 1 , t 2 ) and φ backward ( t 1 , t 2 ) denotes the number of forward and backward edges in the spanning tree path between pointer vertices at the t 1 -th and t 2 -th time step.

Original: The direction of edge is decided by order of generation in the STGG framework.
Revised: The direction of edge is decided by order of generation in the STGG framework.

Original: Such an encoding was inspired from recent works (Villmow et al., 2021;
Revised: Such an encoding was inspired from recent works (Villmow et al., 2021;

Original: Lukovnikov & Fischer, 2021;
Revised: Lukovnikov & Fischer, 2021;

Original: Ying et al., 2021) using Transformers to recognize graphs and trees.
Revised: Ying et al., 2021) using Transformers to recognize graphs and trees.

Original: Finally, the sequence-based relative position φ seq ( t 1 , t 2 ) =
Revised: Finally, the sequence-based relative position φ seq ( t 1 , t 2 ) =

Original: t 1 − tdenotes the relative difference of time-steps for the decisions.
Revised: t 1 − tdenotes the relative difference of time-steps for the decisions.

Original: 3. A TTENTION FOR UPDATING RESIDUAL EDGES .
Revised: 3. A TTENTION FOR UPDATING RESIDUAL EDGES .

Original: Our model generates a categorical distribution over the space of X = X atom ∪ X bond ∪ L res ∪ { "(" , ")" , "*" , "[eos]" } .
Revised: Our model generates a categorical distribution over the space of X = X atom ∪ X bond ∪ L res ∪ { "(" , ")" , "*" , "[eos]" } .

Original: It is relatively straight-forward to output an unnormalized probability over values of X atom ∪ X bond ∪ { "(" , ")" , "*" } using a linear classiﬁer on top of the Transformer model.
Revised: It is relatively straight-forward to output an unnormalized probability over values of X atom ∪ X bond ∪ { "(" , ")" , "*" } using a linear classiﬁer on top of the Transformer model.

Original: However, it is non-trivial to assign probability values for res bond operation, i.e., decisions values of d ∈ L res , since L res varies between different time-steps.
Revised: However, it is non-trivial to assign probability values for res bond operation, i.e., decisions values of d ∈ L res , since L res varies between different time-steps.

Original: To handle this case, we use an attention-based mechanism for assigning unnormalized probability to decision values in the list L res .
Revised: To handle this case, we use an attention-based mechanism for assigning unnormalized probability to decision values in the list L res .

Original: To be speciﬁc, at the ﬁnal layer of our model, we obtain the following probability distribution p ( d ) .
Revised: To be speciﬁc, at the ﬁnal layer of our model, we obtain the following probability distribution p ( d ) .

Original: [Equation]
Revised: [Equation]

Original: where w d ∈ R 1 × (cid:96) is a decision-speciﬁc vector, W 1 , W 2 ∈ R (cid:96) ×
Revised: where w d ∈ R 1 × (cid:96) is a decision-speciﬁc vector, W 1 , W 2 ∈ R (cid:96) ×

Original: ˜ (cid:96) are weight matrices, and h is the decision embedding, i.e., output of the Transformer layer corresponding to the previously made decision.
Revised: ˜ (cid:96) are weight matrices, and h is the decision embedding, i.e., output of the Transformer layer corresponding to the previously made decision.

Original: Furthermore, h d is the embedding corresponding to a past decision d ∈ L res .
Revised: Furthermore, h d is the embedding corresponding to a past decision d ∈ L res .

Original: Finally, m v ( d ) , m g ( d ) are the mask for excluding invalid decisions that violate the validity of graph generation and valence rule, respectively.
Revised: Finally, m v ( d ) , m g ( d ) are the mask for excluding invalid decisions that violate the validity of graph generation and valence rule, respectively.

Original: The masks are obtained using the criteria explained in Section 2.3.
Revised: The masks are obtained using the criteria explained in Section 2.3.

Original: We use the mask during both training and evaluation of the model; this differs from existing graphgenerative models which forbid invalid decisions only at evaluation using a sample-rejection scheme.
Revised: We use the mask during both training and evaluation of the model; this differs from existing graphgenerative models which forbid invalid decisions only at evaluation using a sample-rejection scheme.

Original:  R ELATED W ORKS
Revised:  R ELATED W ORKS

Original: SMILES-based molecular generative models.
Revised: SMILES-based molecular generative models.

Original: The sequence of decisions in STGG for molecule generation is largely inspired from how a sequence of SMILES characters deﬁne the molecular structure.

Original: The main difference lies in how our sequence (a) does not use any abbreviations, 3 and (b) includes the decision of selecting a vertex to perform res bond operations.

Original: 4 Such modiﬁcations allow the STGG framework to interpret the sequence of decisions as sequential graph modiﬁcation of the spanning tree and residual edges.

Original: Several studieson molecule-generating deep neural networks proposed  string-based (G´omezBombarelli et al., 2016; Segler et al., 2018;
Revised: Several studies proposed to generate a SMILES representation of molecules using string-based (G´omez-Bombarelli et al., 2016; Segler et al., 2018;

Original: Kim et al., 2021) or grammar-based (Kusner et al., 2017; Dai et al., 2018) generation of the SMILES.
Revised: Kim et al., 2021) or grammar-based (Kusner et al., 2017; Dai et al., 2018) models.

Original: However, they do not utilize the underlying graph 
Revised: While our STGG is largely inspired from such works, our STGG allows realizing the intermediate graph structure of the molecule being constructed while the SMILES-based models cannot. This difference allows the adoption of structure-aware deep neural networks to STGG.

Original: [Table]
Revised: To be speciﬁc, the difference between STGG and the SMILES-based models appears from our newly introduced graph construction procedure using a pointer vertex i point , a vertex-list L , and a vertex-stack S . They allow recognizing an incomplete sequence of decisions as a graph and assigning positions to each decision.

Original: structure of the SMILES representation. Given the similarity of our decision process to the SMILES representation
Revised: In contrast, an incomplete SMILES string does not deﬁne a graph structure and assigning positions to each character is non-trivial.

Original: [Table]
Revised: [Table]

Original: our STGG algorithm could be seen as a modiﬁed SMILES-based generative model which is augmented with the information of the molecular graph structure.

Original: Graph-based molecular generative models.
Revised: Graph-based molecular generative models.

Original: Researchers have developed a large variety of molecular graph generation frameworks based on atom-wise and bond-wise operations (You et al., 2018; Popova et al., 2019; Madhawa et al., 2019; Honda et al., 2019; Shi et al., 2020; Zang & Wang, 2020; 
Revised: Researchers have developed a large variety of molecular graph generation frameworks based on atom-wise and bond-wise operations (You et al., 2018; Popova et al., 2019; Madhawa et al., 2019; Honda et al., 2019; Shi et al., 2020; Zang & Wang, 2020; Luo et al., 2021).

Original: Lippe & Gavves, 2021;
Revised: Our STGG framework simpliﬁes the decision space of such models by exploiting the tree-like graph structures of molecules.

Original: Luo et al., 2021; Chen et al., 2021a). The frameworks mainly differ in the choice of model for generating the full adjacency of the molecular graph, e.g., using recurrent neural network (Popova et al., 2019), continuous normalizing ﬂow model (Shi et al., 2020), and discrete normalizing ﬂow model (Luo et al., 2021). Our STGG framework differs from this line of research since it proposes a new type of graph-based operations for generating the molecular graph.
Revised: To be speciﬁc, STGG requires O ( |A| + |B| ) decisions for constructing a molecule while the existing atom-by-atom graph generative models typically require O ( |A| 2 ) decisions. This implies that our generative model requires a smaller number of decisions for sparse graphs like molecules, i.e., when |B| is small.

Original: Furthermore, our work is the ﬁrst to successfully train a Transformer architecture (Vaswani et al., 2017) for molecular graph generation.
Revised: Furthermore, our work is the ﬁrst to successfully train a Transformer architecture (Vaswani et al., 2017) for graph-based molecule generation.

Original: In another line of research, several works (Jin et al., 2018; 2019; 2020) proposed generative models based on using the junction-tree representation with molecular substructures as building blocks.
Revised: In another line of research, several works (Jin et al., 2018; 2019; 2020) proposed generative models based on using the junction-tree representation with molecular substructures as building blocks.

Original: Based on such a representation, such works utilize tree-constructive operations to generate the full graph.
Revised: Based on such a representation, such works utilize tree-constructive operations to generate the full graph.

Original: Since they operate on such a coarse-grained molecular representation, they typically require a fewer number of building blocks to generate the whole molecule.
Revised: Since they operate on such a coarse-grained molecular representation, they typically require a fewer number of building blocks to generate the whole molecule.

Original: In comparison, our STGG framework utilizes a more ﬁne-grained molecular representation and may additionally learn the inner semantics of substructures that are used as building blocks for the junction tree.
Revised: In comparison, our STGG framework utilizes a more ﬁne-grained molecular representation and may additionally learn the inner semantics of substructures that are used as building blocks for the junction tree.

Original:  E XPERIMENT
Revised:  E XPERIMENT

Original: In this section, we report the experimental results of the proposed spanning tree-based graph generation (STGG) framework.
Revised: In this section, we report the experimental results of the proposed spanning tree-based graph generation (STGG) framework.

Original: In Section 5.1 and 5.2, we compare with the existing graph generative models in the ZINC250 K (Irwin et al., 2012) and QM9 (Ramakrishnan et al., 2014).
Revised: In Section 5.1 and 5.2, we compare with the existing graph generative models in the ZINC250 K (Irwin et al., 2012) and QM9 (Ramakrishnan et al., 2014).

Original: We provide ablation studies on each component of our method using the ZINC250 K dataset.
Revised: We provide ablation studies on each component of our method using the ZINC250 K dataset.

Original: In Section 5.2, we compare with the existing molecule generative models using the MOSES benchmark (Polykovskiy et al., 2020).
Revised: In Section 5.2, we compare with the existing molecule generative models using the MOSES benchmark (Polykovskiy et al., 2020).

Original: Finally, in Section 5.3, we provide our results on the molecular optimization task with respect to the penalized octanol-water partition coefﬁcient function (PL OG P).
Revised: Finally, in Section 5.3, we provide our results on the molecular optimization task with respect to the penalized octanol-water partition coefﬁcient function (PL OG P).

Original: We provide the implementation details and illustration of the generated molecules in Appendix D and E, respectively.
Revised: We provide the implementation details and illustration of the generated molecules in Appendix D and E, respectively.

Original: 5. M OLECULE GENERATION ON ZINC250 K AND QM9 DATASETS
Revised: 5. M OLECULE GENERATION ON ZINC250 K AND QM9 DATASETS

Original: We ﬁrst compare to the literature standard for the molecular generation task in the ZINC250 K and the QM9 datasets.
Revised: We ﬁrst compare to the literature standard for the molecular generation task in the ZINC250 K and the QM9 datasets.

Original: To this end, we train our generative model on the respective datasets and sample 10,000 molecules to measure (a) the ratio of valid molecules (V ALID ), (b) the ratio of unique molecules (U NIQUE ), and (c) the ratio of novel molecules with respect to the training dataset (N OVEL ).
Revised: To this end, we train our generative model on the respective datasets and sample 10,000 molecules to measure (a) the ratio of valid molecules (V ALID ), (b) the ratio of unique molecules (U NIQUE ), and (c) the ratio of novel molecules with respect to the training dataset (N OVEL ).

Original: We compare with the numbers reported by recently proposed graph generative models (Shi et al., 2020; Luo et al., 2021).
Revised: We compare with the numbers reported by recently proposed graph generative models (Shi et al., 2020; Luo et al., 2021).

Original: We also provide an additional baseline of a transformer architecture trained to generate the SMILES representation for the molecule (SMILES-T RANSFORMER ).
Revised: We also provide an additional baseline of a transformer architecture trained to generate the SMILES representation for the molecule (SMILES-T RANSFORMER ).

Original: We mark C ORRECTABLE for methods which can optionally use a sample-rejection scheme to forbid decisions that violate the chemical rules at evaluation.
Revised: We mark C ORRECTABLE for methods which can optionally use a sample-rejection scheme to forbid decisions that violate the chemical rules at evaluation.

Original: Note that our framework can train the[Table] [Table] generative model under the valence correction mask during training , while existing graph generative models use the valence correction only at evaluation .
Revised: Note that our framework can train the generative model under the valence correction mask during training , while existing graph generative models use the valence correction only at evaluation .

Original: However, for comparison, we do not use the valence correction mask during training in this experiment
Revised: However, for comparison, we do not use the valence correction mask during training in this experiment.

Original: [Table]
Revised: [Table]

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: We report the experimental results in Table 2.
Revised: We report the experimental results in Table 2.

Original: In the table, we observe that our STGG framework outperforms all the existing molecular graph generative models for high V ALID at the cost of relatively lower N OVEL .
Revised: In the table, we observe that our STGG framework outperforms all the existing molecular graph generative models for high V ALID at the cost of relatively lower N OVEL .

Original: In particular, our generative model can achieve a 100% ratio of valid molecules in the QM9 dataset even without any correction procedure.
Revised: In particular, our generative model can achieve a 100% ratio of valid molecules in the QM9 dataset even without any correction procedure.

Original: Such a result highlights the how our model can effectively learn the chemical rules and model the underlying distribution.
Revised: Such a result highlights the how our model can effectively learn the chemical rules and model the underlying distribution.

Original: Finally, our STGG framework performing better than the SMILES-based transformer implies how the performance of our generative model stems from the STGG framework, rather than using the
Revised: Finally, our STGG framework performing better than the SMILES-based transformer implies how the performance of our generative model stems from the STGG framework, rather than using the

Original: Transformer architecture.
Revised: Transformer architecture.

Original: Ablation studies.
Revised: Ablation studies.

Original: We also conduct ablation studies on the
Revised: We also conduct ablation studies on the

Original: ZINC250k dataset to verify the effectiveness of our method.
Revised: ZINC250k dataset to verify the effectiveness of our method.

Original: To this end, we report the experimental results of our method without speciﬁc components.
Revised: To this end, we report the experimental results of our method without speciﬁc components.

Original: To be speciﬁc, we ablate the effects of using sequential relative positional encoding (S), treebased relative positional encoding (T), graph-construction mask (G), and valence rule mask (V).
Revised: To be speciﬁc, we ablate the effects of using sequential relative positional encoding (S), treebased relative positional encoding (T), graph-construction mask (G), and valence rule mask (V).

Original: We also consider an additional baseline of using the absolute positional encoding (A) as in the original Transformer architecture (Vaswani et al., 2017).
Revised: We also consider an additional baseline of using the absolute positional encoding (A) as in the original Transformer architecture (Vaswani et al., 2017).

Original: In
Revised: In

Original: Table 3 and Figure 4, one can observe how each component of our algorithm is crucial for achieving high V ALID .
Revised: Table 3 and Figure 4, one can observe how each component of our algorithm is crucial for achieving high V ALID .

Original: In particular, the tree encoding is essential for the performance, showing the importance of tree-based representation that we use in our model.
Revised: In particular, the tree encoding is essential for the performance, showing the importance of tree-based representation that we use in our model.

Original: 5. M OLECULE GENERATION ON THE MOSES BENCHMARK
Revised: 5. M OLECULE GENERATION ON THE MOSES BENCHMARK

Original: We also compare our method on the MOSES benchmark with the existing models.
Revised: We also compare our method on the MOSES benchmark with the existing models.

Original: The MOSES benchmark offers a large collection of metrics to access the overall quality of generated molecules.
Revised: The MOSES benchmark offers a large collection of metrics to access the overall quality of generated molecules.

Original: To be speciﬁc, in addition to V ALID , U NIQUE , N OVEL , we consider internal diversity of molecules (I NT D IV ), ratio of samples being accepted to chemical ﬁlters (F ILTERS ), Fr ´ etchet ChemNet Distance 
Revised: To be speciﬁc, in addition to V ALID , U NIQUE , N OVEL , we consider internal diversity of molecules (I NT D IV ), ratio of samples being accepted to chemical ﬁlters (F ILTERS ), Fr´etchet ChemNet Distance (FCD), nearest neighborhood similarity (SNN), frament similarity (F RAG ), and Scaffold similarity (S CAF ).

Revised: The similarity metrics of FCD, SNN, F RAG , S CAF are measured with respect to the test dataset of molecules and the scaffolds extracted from them.

Original: [Figure]
Revised: [Figure]

Original: In Table 4 and [Figure] we provide our experimental result.
Revised: In Table 4 and 5, we provide our experimental result.

Original: Here [Figure] one can observe how our algorithm outperforms the existing works for 10 out of 15 metrics including F ILTERS , FCD-T EST , FCD-T EST SF,
Revised: Here, one can observe how our algorithm outperforms the existing works for 10 out of 15 metrics including F ILTERS , FCD-T EST , FCD-T EST SF,

Original: SNN-T EST , SNN-T EST SF, F RAG -T EST , F RAG -T EST SF [Figure] and S CAF -T EST .
Revised: SNN-T EST , SNN-T EST SF, F RAG -T EST , F RAG -T EST SF, and S CAF -T EST .

Original: This highlights the ability of our STGG framework to successfully learn the training distribution [Figure]
Revised: This highlights the ability of our STGG framework to successfully learn the training distribution.

Original: 5. M OLECULAR OPTIMIZATION FOR PENALIZED OCTANOL - WATER PARTITION COEFFICIENT
Revised: 5. M OLECULAR OPTIMIZATION FOR PENALIZED OCTANOL - WATER PARTITION COEFFICIENT

Original: Finally, we compare with existing graph generative models for the task of molecular optimization.
Revised: Finally, we compare with existing graph generative models for the task of molecular optimization.

Original: To this end, we consider the literature standard of maximizing the penalized octanol-water partition coefﬁcient (PL OG P).
Revised: To this end, we consider the literature standard of maximizing the penalized octanol-water partition coefﬁcient (PL OG P).

Original: However, several works (Gao & Coley, 2020; Coley, 2020) have noted how the existing algorithms on this benchmark may not be practical, since PL OG P is ill-deﬁned as a scoring function for molecules; this scoring function may assign high values to “unrealistic” molecules that are unstable and hard to synthesize in practice.
Revised: However, several works (Gao & Coley, 2020; Coley, 2020) have noted how the existing algorithms on this benchmark may not be practical, since PL OG P is ill-deﬁned as a scoring function for molecules; this scoring function may assign high values to “unrealistic” molecules that are unstable and hard to synthesize in practice.

Original: As a result, any interaction between the scoring function and the generative model may bias the model towards generating unrealistic molecules.
Revised: As a result, any interaction between the scoring function and the generative model may bias the model towards generating unrealistic molecules.

Original: To consider this aspect, we propose an ofﬂine optimization algorithm which does not require interaction between the scoring function and the generative model, i.e., we train our model only on the dataset of real-world molecules.
Revised: To consider this aspect, we propose an ofﬂine optimization algorithm which does not require interaction between the scoring function and the generative model, i.e., we train our model only on the dataset of real-world molecules.

Original: At a high-level, we train a conditional generative model p θ ( m | γ ) under the STGG framework with PL OG P as the condition γ .
Revised: At a high-level, we train a conditional generative model p θ ( m | γ ) under the STGG framework with PL OG P as the condition γ .

Original: At the test time, we sample from a high value γ to obtain high-scoring molecules.
Revised: At the test time, we sample from a high value γ to obtain high-scoring molecules.

Original: Such an algorithm is inspired from the recent ofﬂine reinforcement learning algorithms (Schmidhuber, 2019; Kumar & Levine, 2020;
Revised: Such an algorithm is inspired from the recent ofﬂine reinforcement learning algorithms (Schmidhuber, 2019; Kumar & Levine, 2020;

Original: Chen et al., 2021b;
Revised: Chen et al., 2021;

Original: Janner et al., 2021).
Revised: Janner et al., 2021).

Original: We fully describe our molecular optimization algorithm in Appendix F.
Revised: We fully describe our molecular optimization algorithm in Appendix F.

Original: In Table 5 and Figure 6, we report the result of our molecular optimization experiment.
Revised: In Table 5 and Figure 6, we report the result of our molecular optimization experiment.

Original: We provide additional illustrations of the generated molecules in Appendix G. One can observe how our STGG model outperforms the existing generative models by a large margin.
Revised: We provide additional illustrations of the generated molecules in Appendix G. One can observe how our STGG model outperforms the existing generative models by a large margin.

Original: This demonstrates how our model can extrapolate well outside the training distribution.
Revised: This demonstrates how our model can extrapolate well outside the training distribution.

Original: Intriguingly, in Figure 6 and Appendix G, one can observe how increasing γ gradually changes the optimized molecule from realistic structures to large, chain-like, and unrealistic structures.
Revised: Intriguingly, in Figure 6 and Appendix G, one can observe how increasing γ gradually changes the optimized molecule from realistic structures to large, chain-like, and unrealistic structures.

Original: This observation is in agreement with prior works (Winter et al., 2019; Nigam et al., 2020; Shi et al., 2020;
Revised: This observation is in agreement with prior works (Winter et al., 2019; Nigam et al., 2020; Shi et al., 2020;

Original: Ahn et al., 2020; Luo et al., 2021) which report molecules with high PL OG P that have chain-like and unrealistic structures.
Revised: Ahn et al., 2020; Luo et al., 2021) which report molecules with high PL OG P that have chain-like and unrealistic structures.

Original: However, our ofﬂine optimization algorithm offers an additional freedom of the condition γ to control the trade-off between high PL OG P and realistic-ness of the generated molecules.
Revised: However, our ofﬂine optimization algorithm offers an additional freedom of the condition γ to control the trade-off between high PL OG P and realistic-ness of the generated molecules.

Original: Investigating the use of our newly proposed optimization algorithm is an interesting future research direction.
Revised: Investigating the use of our newly proposed optimization algorithm is an interesting future research direction.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper, we propose STGG which is the ﬁrst spanning tree-based framework for the generation of molecules using the Transformer architecture.
Revised: In this paper, we propose STGG which is the ﬁrst spanning tree-based framework for the generation of molecules using the Transformer architecture.

Original: The key idea of using the spanning tree for graph generation applies to any graph type outside the molecules; we believe such an extension of our work to be both promising and interesting.
Revised: The key idea of using the spanning tree for graph generation applies to any graph type outside the molecules; we believe such an extension of our work to be both promising and interesting.

Original: We also propose an ofﬂine algorithm for molecular optimization which allows the trade-off between the high score and the realistic-ness of molecules.
Revised: We also propose an ofﬂine algorithm for molecular optimization which allows the trade-off between the high score and the realistic-ness of molecules.

Original: We leave more investigation of the newly proposed optimization algorithm as future work.
Revised: We leave more investigation of the newly proposed optimization algorithm as future work.

Original:  R EPRODUCIBILITY S TATEMENT
Revised:  R EPRODUCIBILITY S TATEMENT

Original: We provide explicit description of our algorithm in Algorithm 1, Appendix 2,  and 4.
Revised: We provide explicit description of our algorithm in Algorithm 1, Appendix 2, 3, and 4.

Original: We list the hyper-parameters, the hardware used for the experiments and the data-processing information in
Revised: We list the hyper-parameters, the hardware used for the experiments and the data-processing information in

Original: Appendix D. We provide illustrations of the molecules generated for the experiments in Figure  and
Revised: Appendix D. We provide illustrations of the molecules generated for the experiments in Figure 6, and

Original: Appendix E, G.
Revised: Appendix E, G.

Original: We submit the full implementation of our STGG framework and the baselines used in our experiments as a supplementary material
Revised: We submit the full implementation of our STGG framework and the baselines used in our experiments as a supplementary material

Original: A E XTRACTING SEQUENCE OF DECISIONS FROM A MOLECULAR GRAPH
Revised: A E XTRACTING SEQUENCE OF DECISIONS FROM A MOLECULAR GRAPH

Original: In this section, we explain our algorithm for ﬁnding a sequence of decisions to construct a given molecular graph G =
Revised: In this section, we explain our algorithm for ﬁnding a sequence of decisions to construct a given molecular graph G =

Original: ( A , B , E ) .
Revised: ( A , B , E ) .

Original: The high-level idea is to ﬁrst perform a depth-ﬁrst search on G to ﬁnd a spanning tree T =
Revised: The high-level idea is to ﬁrst perform a depth-ﬁrst search on G to ﬁnd a spanning tree T =

Original: ( A , B , E T ) and the corresponding set of residual edges E R = E\E T .
Revised: ( A , B , E T ) and the corresponding set of residual edges E R = E\E T .

Original: Then the algorithm traverses the spanning tree T according to the depth-ﬁrst search tree while (a) allocating branch start and branch end for vertices with degree higher than two and (b) adding res atom and res bond operations for any vertex covered by a residual edge { a, b } ∈ E R .
Revised: Then the algorithm traverses the spanning tree T according to the depth-ﬁrst search tree while (a) allocating branch start and branch end for vertices with degree higher than two and (b) adding res atom and res bond operations for any vertex covered by a residual edge { a, b } ∈ E R .

Original: To this end, we utilize a stack S dfs that stores the list of vertices in G and branching tokens { "(" , ")" } to visit.
Revised: To this end, we utilize a stack S dfs that stores the list of vertices in G and branching tokens { "(" , ")" } to visit.

Original: At each iteration, an element i of the stack S dfs is popped.
Revised: At each iteration, an element i of the stack S dfs is popped.

Original: If i is a vertex, the algorithm adds the corresponding decision for attach atom and attach bond operations.
Revised: If i is a vertex, the algorithm adds the corresponding decision for attach atom and attach bond operations.

Original: If the vertex has more than two successors with respect to the spanning tree T , the successors are inserted into the stack S dfs with surrounding "(" and ")" tokens.
Revised: If the vertex has more than two successors with respect to the spanning tree T , the successors are inserted into the stack S dfs with surrounding "(" and ")" tokens.

Original: If the vertex has only one successor, the successor is inserted into the stack without an additional operation.
Revised: If the vertex has only one successor, the successor is inserted into the stack without an additional operation.

Original: When the branching tokens { "(" , ")" } are popped from the stack, the algorithm adds the corresponding decision value to the sequence of decisions.
Revised: When the branching tokens { "(" , ")" } are popped from the stack, the algorithm adds the corresponding decision value to the sequence of decisions.

Original: We describe the full scheme in Algorithm 2.
Revised: We describe the full scheme in Algorithm 2.

Original: Decision to attach the atom vertex.
Revised: Decision to attach the atom vertex.

Original: 9: for j ∈ { j | j ∈ N ( i ) , { i, j } ∈ E R } do (cid:46) Decisions for residual edges.
Revised: 9: for j ∈ { j | j ∈ N ( i ) , { i, j } ∈ E R } do (cid:46) Decisions for residual edges.

Original: If i ∈ A , append "*" to D .
Revised: If i ∈ A , append "*" to D .

Original: If i ∈ B , append j to D .
Revised: If i ∈ B , append j to D .

Original: Let V denote { j | j ∈ N ( i ) , j (cid:54)∈ A T ∪ B T } .
Revised: Let V denote { j | j ∈ N ( i ) , j (cid:54)∈ A T ∪ B T } .

Original: (cid:46) Successors of i in depth-ﬁrst order.
Revised: (cid:46) Successors of i in depth-ﬁrst order.

Original: If |V| > 1 , insert "(" to S branch .
Revised: If |V| > 1 , insert "(" to S branch .

Original: (cid:46) Allocate decision to record pointer vertex.
Revised: (cid:46) Allocate decision to record pointer vertex.

Original: Insert vertices in V to S branch .
Revised: Insert vertices in V to S branch .

Original: (cid:46)
Revised: (cid:46)

Original: Allocate successors to visit later.
Revised: Allocate successors to visit later.

Original: If |V| > 1 , insert ")" to S branch .
Revised: If |V| > 1 , insert ")" to S branch .

Original: (cid:46) Allocate decision to return to the pointer vertex.
Revised: (cid:46) Allocate decision to return to the pointer vertex.

Original: 16: if i ∈ { "(" , ")" } then 17:
Revised: 16: if i ∈ { "(" , ")" } then 17:

Original: Append i to D .
Revised: Append i to D .

Original: 18: while |S branch | >19: Output: sequence of decisions D = d 1 , . . .
Revised: 18: while |S branch | >19: Output: sequence of decisions D = d 1 , . . .

Original: , d to reconstruct G .
Revised: , d to reconstruct G .

Original: B A LGORITHMS FOR GRAPH MASKING
Revised: B A LGORITHMS FOR GRAPH MASKING

Original: To determine whether if a sequence of decisions lead to a valid generation of a molecular graph, we propose an algorithm that outputs a set of valid decisions given the current decision d , stack of pointer vertices S , and list of atom vertices L during execution of Algorithhm 1.
Revised: To determine whether if a sequence of decisions lead to a valid generation of a molecular graph, we propose an algorithm that outputs a set of valid decisions given the current decision d , stack of pointer vertices S , and list of atom vertices L during execution of Algorithhm 1.

Original: We provide the full description in Algorithm 3. 
Revised: We provide the full description in Algorithm 3. [Table] [Table] [Table]

Revised: [Table]

Revised: We also establish the following theoretical result on how the sequence of decisions generated from

Original: C A LGORITHM FOR VALENCE MASKING
Revised: Algorithm 3 always deﬁnes a valid graph structure. Theorem 1. Let G = ( A , B , E ) , S , and L be a graph, a stack of vertices, and a list of vertices being updated by Algorithm 1 and a sequence of decisions d 1 , . . .

Revised: , d T .

Revised: If the sequence of decisions satisﬁes the criteria deﬁned by Algorithm 3, the following properties are satisﬁed

Revised: P 1 At the t -th step of Algorithm 1, |S| > 0 if d t = ")" .

Revised: P 2 At the t -th step of Algorithm 1, d t ∈ L if d t ∈ A ∪ B .

Original: {+P 
+}
Revised: {+P 
+}

Revised: At any step of Algorithm 1, the corresponding graph G is bipartite and connected.

Revised: P 4 If d T = [eos] , the graph G satisﬁes |N ( b ) | = 2 for all b ∈ B .

Revised: Namely, P 1 and P 2 implies how the operations in Algorithm 1 are well-deﬁned for d 1 , . . .

Revised: , d T .

Revised: Next,

Revised: P 3 and P 4 implies how the generated graph represents a molecular graph structure.

Revised: Proof.

Revised: First, P 1 is enforced by Algorithm 3 which forbids the decision value of ")" when the stack S is empty.

Revised: Next, P 2 is enforced by Algorithm 3 selecting decision values from the current list of vertices L .

Revised: To enforce P 3 , Algorithm 3 allows the decision of d t ∈ X atom ∪ L only when the pointer vertex is a bond vertex, i.e., d ∈ X bond .

Revised: Similarly, d t ∈ X bond is allowed only when d ∈ X atom ∪ { "*" , "(" } .

Revised: Finally, to enforce P 4 , the algorithm does not allow adding a bond

Revised: C A LGORITHM FOR VALENCE MASKING

Original: To consider chemical validity of molecules, our framework offers the ability to constrain its generation on molecules that satisfy the valence rule for each atom.
Revised: To consider chemical validity of molecules, our framework offers the ability to constrain its generation on molecules that satisfy the valence rule for each atom.

Original: That is, the generated graph G =
Revised: That is, the generated graph G =

Original: ( A , B , E ) satisﬁes the constraint v ( x a ) ≥ (cid:80) b ∈N ( a ) o ( x b ) for every atom vertex a ∈ A where v ( x a ) denotes the valence of an atom type x a and o ( x b ) denotes the bond order.
Revised: ( A , B , E ) satisﬁes the constraint v ( x a ) ≥ (cid:80) b ∈N ( a ) o ( x b ) for every atom vertex a ∈ A where v ( x a ) denotes the valence of an atom type x a and o ( x b ) denotes the bond order.

Original: To this end, we propose an algorithm which iteratively updates a record r ( a ) of available valence for each atom vertex a ∈ A .
Revised: To this end, we propose an algorithm which iteratively updates a record r ( a ) of available valence for each atom vertex a ∈ A .

Original: The key idea is to (a) update the record accordingly for each addition of atom and bond orders and (b) pre-allocate valence for the branch start and res atom operations by the amount of minimum bond order min x ∈X bond o ( x ) .
Revised: The key idea is to (a) update the record accordingly for each addition of atom and bond orders and (b) pre-allocate valence for the branch start and res atom operations by the amount of minimum bond order min x ∈X bond o ( x ) .

Original: The second part (b) is required since the branch start and res atom operations indicate future bond vertices to be added as a neighbor of the current atom vertex.
Revised: The second part (b) is required since the branch start and res atom operations indicate future bond vertices to be added as a neighbor of the current atom vertex.

Original: We provide the full description in Algorithm 4.
Revised: We provide the full description in Algorithm 4.

Revised: We note that Algorithm 4 allows guaranteeing the generation of a valid molecule based on forbidding any operation that violate the valence rule.

Revised: Furthermore, pre-allocating the minimum amount of bond order for the branch start and res atom prevents the case where every decision value is invalid after the respective branch start and res atom decisions.

Revised: Therefore, combined with Algorithm 3,

Revised: Algorithm 4 guarantees generation of a molecular graph satisfying the valence rule.

Revised: Update record using previously added bond vertex.

Original:  (cid:46) Reject decisions requiring minimal amount of valence.
Revised: Set D ← { x | x ∈ X bond , o ( x ) > r ( i point ) } . (cid:46) Reject bond orders higher than the record.

Revised: 7: if r ( i point ) < min x ∈X bond o ( x ) then 8:

Revised: Set D ← D ∪ { "(" , "*" } .

Revised: (cid:46) Reject decisions requiring minimal amount of valence.

Original: 9: if d ∈ X bond then 10: if ˜ d (cid:54) = "(" then 11:
Revised: 9: if d ∈ X bond then 10: if ˜ d (cid:54) = "(" then 11:

Original: Set r
Revised: Set r

Original: (˜ i point ) ← r (˜ i point ) − o ( d ) .
Revised: (˜ i point ) ← r (˜ i point ) − o ( d ) .

Original: (cid:46) Update record of previously added atom vertex.
Revised: (cid:46) Update record of previously added atom vertex.

Original: 12: else 13:
Revised: 12: else 13:

Original: Set r
Revised: Set r

Original: (˜ i point ) ← r (˜ i point )
Revised: (˜ i point ) ← r (˜ i point )

Original: − o ( d ) + min x ∈X bond o ( x ) .
Revised: − o ( d ) + min x ∈X bond o ( x ) .

Original: (cid:46) Update previously added atom vertex considering pre-allocated valence.
Revised: (cid:46) Update previously added atom vertex considering pre-allocated valence.

Original: Set D ← { x | x ∈ X atom , v ( x ) < r
Revised: Set D ← { x | x ∈ X atom , v ( x ) < r

Original: ( i point )) } .
Revised: ( i point )) } .

Original: (cid:46) Reject atom valence lower than bond order.
Revised: (cid:46) Reject atom valence lower than bond order.

Original: Set D ← D ∪ { x | x ∈ L res , r ( x ) < o ( d ) − min x ∈X bond o ( x ) } .
Revised: Set D ← D ∪ { x | x ∈ L res , r ( x ) < o ( d ) − min x ∈X bond o ( x ) } .

Original: (cid:46) Reject residual edge candidates with valence lower than bond order.
Revised: (cid:46) Reject residual edge candidates with valence lower than bond order.

Original: 16: if d = "(" then 17:
Revised: 16: if d = "(" then 17:

Original: Set r
Revised: Set r

Original: ( i point )
Revised: ( i point )

Original: ← r ( i point ) − min x ∈X bond o ( x ) .
Revised: ← r ( i point ) − min x ∈X bond o ( x ) .

Original: (cid:46) Pre-allocate minimum bond order.
Revised: (cid:46) Pre-allocate minimum bond order.

Original: Set D ← { x | x ∈ X bond , o ( x ) >
Revised: Set D ← { x | x ∈ X bond , o ( x ) >

Original: r ( i point ) } .
Revised: r ( i point ) } .

Original: (cid:46) Reject bond orders higher than the record.
Revised: (cid:46) Reject bond orders higher than the record.

Original: 19: if d = ")" then 20:
Revised: 19: if d = ")" then 20:

Original: Set D ← ∅ 21: if d = "*" then 22:
Revised: Set D ← ∅ 21: if d = "*" then 22:

Original: Set r
Revised: Set r

Original: ( i point )
Revised: ( i point )

Original: ← r ( i point ) − min x ∈X bond o ( x ) .
Revised: ← r ( i point ) − min x ∈X bond o ( x ) .

Original: (cid:46) Pre-allocate minimum bond order.
Revised: (cid:46) Pre-allocate minimum bond order.

Original: Set D ← { x | x ∈ X bond , o ( x ) >
Revised: Set D ← { x | x ∈ X bond , o ( x ) >

Original: r ( i point ) } .

Original: (cid:46) Reject bond orders higher than the record.

Original: 24: if d ∈ L res then 25:

Original: Set r ( d ) ← r ( d ) + min x ∈X bond o ( x ) − o ( x i point ) .

Original: (cid:46) Update record of previously added atom vertex using current bond order.

Original: Set D ← ∅ .

Original: D I MPLEMENTATION DETAILS
Revised: D I MPLEMENTATION DETAILS

Original: In this section, we provide speciﬁc details on how we implement the STGG framework for our experiments.
Revised: In this section, we provide speciﬁc details on how we implement the STGG framework for our experiments.

Original: Training detail.
Revised: Training detail.

Original: For all the experiments, we train the Transformer under STGG framework forepochs with batch size of 128 for all the dataset.
Revised: For all the experiments, we train the Transformer under STGG framework forepochs with batch size of 128 for all the dataset.

Original: We use the AdamW (Loshchilov & Hutter, 2019) optimizer with constant learning rate of 10 − 4 .
Revised: We use the AdamW (Loshchilov & Hutter, 2019) optimizer with constant learning rate of 10 − 4 .

Original: We use three and six Transformer layers for { QM9,
Revised: We use three and six Transformer layers for { QM9,

Original: ZINC250 K } and MOSES, respectively.
Revised: ZINC250 K } and MOSES, respectively.

Original: The rest of Transformer-related conﬁgurations follow that of the original work (Vaswani et al., 2017); we use the attention module with embedding size ofwith eight heads, MLP with dimension of 2048, and dropout with probability of 0 . 1 .
Revised: The rest of Transformer-related conﬁgurations follow that of the original work (Vaswani et al., 2017); we use the attention module with embedding size ofwith eight heads, MLP with dimension of 2048, and dropout with probability of 0 . 1 .

Original: Using a single
Revised: Using a single

Original: Quadro RTX 6000 GPU, it takes approximately three, ten, and 96 hours to fully train the models on
Revised: Quadro RTX 6000 GPU, it takes approximately three, ten, and 96 hours to fully train the models on

Original: QM9, ZINC250 K , and MOSES datasets, respectively.
Revised: QM9, ZINC250 K , and MOSES datasets, respectively.

Original: Pre-processing.
Revised: Pre-processing.

Original: For the whole dataset, we use the following set of atom vocabularies X atom : { "CH" , "CH2" , "CH-" , "CH2-" , "C" , "N-" , "NH-" , "N" , "NH" , "N+" , "NH+" , "NH2+" , "NH3+" , "O-" , "O" , "O+" , "OH+" , "F" , "P" , "PH" , "PH2" , "P+" , "PH+" , "S-" , "S" , "S+" , "SH" , "SH+" , "Cl" , "Br" , "I" , } .
Revised: For the whole dataset, we use the following set of atom vocabularies X atom : { "CH" , "CH2" , "CH-" , "CH2-" , "C" , "N-" , "NH-" , "N" , "NH" , "N+" , "NH+" , "NH2+" , "NH3+" , "O-" , "O" , "O+" , "OH+" , "F" , "P" , "PH" , "PH2" , "P+" , "PH+" , "S-" , "S" , "S+" , "SH" , "SH+" , "Cl" , "Br" , "I" , } .

Original: Note that we assign different features for the same atom numbers with different number of explicit hydrogens and formal charges.
Revised: Note that we assign different features for the same atom numbers with different number of explicit hydrogens and formal charges.

Original: This allows our algorithm to properly allocate maximum valence for each atom feature.
Revised: This allows our algorithm to properly allocate maximum valence for each atom feature.

Original: Next, we use the bond vocabulary X bond = { "-" , "=" , "#" } , corresponding to bond orders of single, double, and triple, respectively.
Revised: Next, we use the bond vocabulary X bond = { "-" , "=" , "#" } , corresponding to bond orders of single, double, and triple, respectively.

Original: For explicit calculation of the atom valence during molecular construction, we train our models on kekulized molecules, i.e., aromatic bonds are ﬁxed to single or double bonds.
Revised: For explicit calculation of the atom valence during molecular construction, we train our models on kekulized molecules, i.e., aromatic bonds are ﬁxed to single or double bonds.

Original: E E XAMPLE OF GENERATED MOLECULES
Revised: E E XAMPLE OF GENERATED MOLECULES

Original: F O FFLINE OPTIMIZATION OF MOLECULES
Revised: F O FFLINE OPTIMIZATION OF MOLECULES

Original: In this section, we describe our ofﬂine molecular optimization algorithm mainly inspired by existing works in ofﬂine reinforcement learning (Schmidhuber, 2019; Chen et al., 2021b; Janner et al., 2021) and ofﬂine model-based optimization (Kumar & Levine, 2020).
Revised: In this section, we describe our ofﬂine molecular optimization algorithm mainly inspired by existing works in ofﬂine reinforcement learning (Schmidhuber, 2019; Chen et al., 2021; Janner et al., 2021) and ofﬂine model-based optimization (Kumar & Levine, 2020).

Original: For maximizing a reward function deﬁned on a molecule, our algorithm consists of two simple steps.
Revised: For maximizing a reward function deﬁned on a molecule, our algorithm consists of two simple steps.

Original: First, our ofﬂine optimization algorithm trains a conditional generative model p θ ( m | γ ) where m is the molecule and γ is the reward function evaluated on the ofﬂine dataset of molecules.
Revised: First, our ofﬂine optimization algorithm trains a conditional generative model p θ ( m | γ ) where m is the molecule and γ is the reward function evaluated on the ofﬂine dataset of molecules.

Original: Next, the reward-conditional generative model samples highly-rewarding molecules by generation conditioned on high values of γ .
Revised: Next, the reward-conditional generative model samples highly-rewarding molecules by generation conditioned on high values of γ .

Original: In particular, we set the value of γ to extrapolate outside the training dataset.
Revised: In particular, we set the value of γ to extrapolate outside the training dataset.

Original: Based on the highly expressive power of Transformer architecture, our algorithm can successfully generate highly-rewarding molecules.
Revised: Based on the highly expressive power of Transformer architecture, our algorithm can successfully generate highly-rewarding molecules.

Original: G A DDITIONAL EXPERIMENTAL RESULTS ON MOLECULAR OPTIMIZATION
Revised: G A DDITIONAL EXPERIMENTAL RESULTS ON MOLECULAR OPTIMIZATION

Revised: H C OMPARISON WITH CG-VAE

Revised: [Table]

Revised: [Table]

Revised: In this section

Revised: [Table]

Revised: we additionally compare our STGG framework with the CG-VAE model (Liu et al., 2018), which is another atom-by-atom graph generative model that allows masking out the action space to generate molecules satisfying the valence rules.

Revised: Compared to Table 2, we additionally use the FCD, SNN, Frag, Scaf metrics to measure the faithfulness of the generative models for learning the underlying distribution of molecules.

Revised: Note that the V ALID metric used in Table 2 is insufﬁcient to compare the faithfulness of STGG and CG-VAE since they both enjoy the guarantee to generate molecules satisfying the valence rule.

Revised: In Table 6 and 7, one can observe that our algorithm highly outperforms the CG-VAE in terms of faithfully learning the underlying distribution of molecules at the cost of relatively lower U NIQUE .

Revised: For example, FCD score of our STGG for learning the ZINC dataset is 0.2775 while that of the

Revised: CG-VAE is 11.33.

Revised: This highlights the expressive power of our STGG framework.