Spanning Tree-based Graph Generation for Molecules

1 INTRODUCTION . Researchers have extensively studied graph generative models , dating back to the early works of Erdös Rényi ( Erdös et al. , 1959 ) . Recently , models based on deep neural networks ( DNNs ) have gained much attraction due to their expressive power in learning a graph dataset . The molecule-generating DNNs stand out among them for their success in the task of drug discovery . 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 ) . For example , Segler et al . ( 2018 ) proposed to train language models on the domain-specific linear string representation of molecules , i.e. , simplified molecular-input line-entry system ( SMILES , Weininger 1988 ) . 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 . 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 . 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 . 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 . 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 . Contribution . In this work , we propose a novel framework , coined spanning tree-based graph generation ( STGG ) , for fine-grained generation of molecules while exploiting their sparsity.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 . Such a formulation allows our framework to utilize compact tree-constructive 1While our framework is designed for general sparse graphs , we focus on the molecular graphs in this paper . operations to define the molecular graph connectivity . See Figure 1 for an illustration of how we formulate the generation of a molecular graph as a sequence of tree-constructive operations . 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 . Such criteria allow control over the generative model to guarantee generating valid molecular graphs by forbidding invalid actions . This is in contrast to prior works ( Shi et al. , 2020 ; Luo et al. , 2021 ) that generate the molecular graph atom-by-atom but determines the validity of construction operations through a sample-rejection scheme . To recognize the spanning tree-based representation used in our STGG framework , we propose a Transformer architecture ( Vaswani et al. , 2017 ) with tree-based relative encoding . 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 . We also introduce an attention-based mechanism for constructing residual edges . We experiment on popular graph generation benchmarks of QM9 , ZINC250K , and MOSES to validate the effectiveness of our algorithm . In the experiments on QM9 and ZINC , our STGG framework outperforms the existing graph-based generative models by a large margin . In the MOSES benchmark , our algorithm achieves superior performance compared to both string-based and graphbased methods for majority of the metrics , e.g. , Fréchet ChemNet distance ( Preuer et al. , 2018 ) and fragment-based similarity . We also conduct experiments on the offline optimization task for high penalized octanol-water partition coefficient and achieve competitive results . 2 SPANNING TREE-BASED GENERATION OF GRAPHS ( STGG ) . 2.1 OVERVIEW . 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 . To this end , we propose compact tree-constructive operations inspired by the simplified molecular-input line-entry system ( SMILES , Weininger , 1988 ) . 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 . 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 . Molecular graph representation . To apply our framework , we represent a molecule as a bipartite graph G = ( A , B , E ) where A and B are the set of vertices associated with atoms and bonds of the molecule , respectively.2 Each edge { a , b } ∈ E is assigned for each adjacent pair of atom and bond . We assign attributes xa ∈ Xatom and xb ∈ Xbond for vertices a ∈ A and b ∈ B to indicate the corresponding atom type and bond order , respectively . For example , { `` C '' , `` N '' , `` O '' } ⊆ Xatom and { `` - '' , `` = '' } ⊆ Xbond . See Figure 1 for an example of such a molecular graph representation . Molecular graph from sequence of decisions . To generate the molecular graph G = ( A , B , E ) , our framework makes a sequence of decisions d1 , . . . , dT to generate a spanning tree T = ( AT , BT , ET ) and a set of residual edges ER = E \ET . 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 . 2Many existing works , e.g. , ( Shi et al. , 2020 ) , use a non-bipartite graph with bonds assigned to edges . See Table 1 for examples of decisions and the corresponding operations . We provide a detailed description of the graph construction process in Section 2.2 . Generating valid molecular graphs . 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 . To prevent this scenario , we conduct two criteria for determining validity of the given decision for ( a ) and ( b ) . We further elaborate this in Section 2.3 . 2.2 DECISION PROCESS FOR SPANNING TREE-BASED GRAPH GENERATION . We now explain how our STGG framework incorporates the decisions d1 , . . . , dT to build the spanning tree T = ( AT , BT , ET ) and residual edges ER from scratch . To this end , our framework introduces the state information of ( a ) a pointer vertex ipoint ∈ AT ∪ BT for specifying the target of the next operation , ( b ) a stack Sbranch that stores vertices to use later as a starting point of a “ branch ” in the spanning tree , and ( c ) a list Lres that stores vertices to use later for constructing residual edges . Algorithm 1 Tree-based generation of molecular graphs 1 : Input : sequence of decisions d1 , . . . , dT . 2 : Output : graph G = ( A , B , E ) , atom attributes { xa } a∈A , and bond attributes { xb } b∈B 3 : Set AT ← ∅ , BT ← ∅ , ET ← ∅ , ER ← ∅ , and T ← ( AT , BT , ET ) . . Initialize the empty graph . 4 : Set Lres as an empty list and Sbranch as an empty stack . 5 : for t = 1 , . . . , T do 6 : if dt ∈ Xatom then . Add a new atom vertex . 7 : Create a new atom vertex a and set A ← A∪ { b } and xa ← dt . 8 : If |BT | > 0 , set ET ← ET ∪ { { a , ipoint } } . . Edge is added when tree is not empty . 9 : Set ipoint ← a . 10 : if dt ∈ Xbond then . Add a new bond vertex . 11 : Create a new bond vertex b and set ET ← ET ∪ { { b , ipoint } } , BT ← BT ∪ { b } , and xb ← dt . 12 : Set ipoint ← b . 13 : if dt = `` * '' then Insert ipoint into Lres. . Add pointer vertex into the queue . 14 : if dt ∈ Lres then Pop dt from Lres and update ER ← ER ∪ { { ipoint , dt } } . . Add a new residual edge . 15 : if dt = `` ( `` then Insert ipoint into Sbranch. . Add pointer vertex into the stack . 16 : if dt = `` ) '' then Set ipoint ← pop ( Sbranch ) . Update pointer vertex from the stack . 17 : Set A ← AT , B ← BT , and E ← ET ∪ ER . 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 ∈ Xatom ∪ Xbond ∪ Lres ∪ { `` ( `` , `` ) '' , `` * '' , `` [ eos ] '' } in detail . See Table 1 for the pairs of operations and the corresponding decisions . We also provide an example of the graph construction process in Figure 2 . Attaching atom and bond vertices to the spanning tree . If the decision d specifies one of the atom or bond attributes , i.e. , d ∈ Xatom or d ∈ Xbond , it applies the corresponding attach atom and attach bond operations , respectively . To be specific , the attach atom operation adds a new atom vertex a into the spanning tree T as a neighbor of the pointer vertex ipointer , i.e. , AT ← AT ∪ { a } , ET ← ET ∪ { a , ipointer } . The value d is set as the new atom attribute , i.e. , xa ← d. The newly added vertex is set as the next pointer vertex , i.e. , ipointer ← a . The attach bond operation similarly adds a new bond vertex . 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 '' ∈ XA and `` - '' ∈ XB . Branching out the spanning tree . 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 . To be specific , the branch start operation inserts the current pointer vertex into a stack Sbranch of vertices . Then the branch end operation pops a vertex from the stack Sbranch and sets it as the new pointer vertex . For an example , a graph with one atom vertex of degree three is constructed from a sequence of decisions C-C ( -C ) ( -C ) . Adding residual edges . 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 ∈ Lres , respectively . To be specific , the res atom operation inserts the current ( atom ) pointer vertex into a list Lres . Next , when the decision value d ∈ Lres is received for the res bond operation , the corresponding vertex d is popped from the list Lres and forms a new residual edge with the current ( bond ) pointer vertex , i.e. , ER ← ER ∪ { { d , ipointer } } . 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 first atom vertex with attribute `` C '' . Termination . The decision `` [ eos ] '' applies the terminate operation to finish the construction . We provide the full algorithm in Algorithm 1 . We also provide an algorithm to extract a sequence of decisions for constructing a given graph in Appendix A . Such an algorithm is used to obtain sequence of decisions as targets for training the generative model under the STGG framework .

Summary: This paper presents a method to construct a molecular graph, which is inspired by a spanning tree. In specific, a molecule is constructed by a sequence of actions, each of which adds an atom, adds a bond, starts/ends a branch, adds residual edges for a circular structure, and terminates. The generative process is controlled by a transformer-based neural network, which is specialized to tree construction procedure. The empirical studies show that the proposed method can learn the distribution of molecules comparably or better than the existing methods, and can be used to molecular optimization.