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 proposes a spanning tree based generative model (STGG) for molecular graphs. STGG sequentially generate a molecule's spanning tree and fill in the residual edges on the way. The spanning tree construction is similar to the standard SMILES representation, but the model operates on a molecular graph rather than a SMILES string. STGG adopts a tree-based transformer with relative positional encoding for tree generation, and a attention-based predictor for residual edge prediction. The method is evaluated on standard ZINC250K, QM9, and MOSES benchmarks and outperform existing baselines.