Multiresolution Equivariant Graph Variational Autoencoder

1 INTRODUCTION . Understanding graphs in a multiscale and multiresolution perspective is essential for capturing the structure of molecules , social networks , or the World Wide Web . Graph neural networks ( GNNs ) utilizing various ways of generalizing the concept of convolution to graphs ( Scarselli et al. , 2009 ) ( Niepert et al. , 2016b ) ( Li et al. , 2016 ) have been widely applied to many learning tasks , including modeling physical systems ( Battaglia et al. , 2016 ) , finding molecular representations to estimate quantum chemical computation ( Duvenaud et al. , 2015 ) ( Kearnes et al. , 2016 ) ( Gilmer et al. , 2017b ) ( Hy et al. , 2018 ) , and protein interface prediction ( Fout et al. , 2017 ) . One of the most popular types of GNNs is message passing neural nets ( MPNNs ) that are constructed based on the message passing scheme in which each node propagates and aggregates information , encoded by vectorized messages , to and from its local neighborhood . While this framework has been immensely successful in many applications , it lacks the ability to capture the multiscale and multiresolution structures that are present in complex graphs ( Rustamov & Guibas , 2013 ) ( Chen et al. , 2014 ) ( Cheng et al. , 2015 ) ( Xu et al. , 2019 ) . Ying et al . ( 2018 ) proposed a multiresolution graph neural network that employs a differential pooling operator to coarsen the graph . While this approach is effective in some settings , it is based on soft assigment matrices , which means that ( a ) the sparsity of the graph is quickly lost in higher layers ( b ) the algorithm isn ’ t able to learn an actual hard clustering of the vertices . The latter is important in applications such as learning molecular graphs , where clusters should ideally be interpretable as concrete subunits of the graphs , e.g. , functional groups . In contrast , in this paper we propose an arhictecture called Multiresolution Graph Network ( MGN ) , and its generative cousin , Multiresolution Graph Variational Autoencoder ( MGVAE ) , which explicitly learn a multilevel hard clustering of the vertices , leading to a true multiresolution hierarchy . In the decoding stage , to “ uncoarsen ” the graph , MGVAE needs to generate local adjacency matrices , which is inherently a second order task with respect to the action of permutations on the vertices , hence MGVAE needs to leverage the recently developed framework of higher order permutation equivariant message passing networks ( Hy et al. , 2018 ; Maron et al. , 2019b ) . Learning to generate graphs with deep generative models is a difficult problem because graphs are combinatorial objects that typically have high order correlations between their discrete substructures ( subgraphs ) ( You et al. , 2018a ) ( Li et al. , 2018 ) ( Liao et al. , 2019 ) ( Liu et al. , 2019 ) ( Dai et al. , 2020 ) . Graph-based molecular generation ( Gmez-Bombarelli et al. , 2018 ) ( Simonovsky & Komodakis , 2018 ) ( Cao & Kipf , 2018 ) ( Jin et al. , 2018 ) ( Thiede et al. , 2020 ) involves further challenges , including correctly recognizing chemical substructures , and importantly , ensuring that the generated molecular graphs are chemically valid . MGN allows us to extend the existing model of variational autoencoders ( VAEs ) with a hierarchy of latent distributions that can stochastically generate a graph in multiple resolution levels . Our experiments show that having a flexible clustering procedure from MGN enables MGVAE to detect , reconstruct and finally generate important graph substructures , especially chemical functional groups . 2 RELATED WORK . There have been significant advances in understanding the invariance and equivariance properties of neural networks in general ( Cohen & Welling , 2016a ) ( Cohen & Welling , 2016b ) , of graph neural networks ( Maron et al. , 2019b ) , of neural networks learning on sets ( Zaheer et al. , 2017 ) ( Serviansky et al. , 2020 ) ( Maron et al. , 2020 ) , along with their expressive power on graphs ( Maron et al. , 2019c ) ( Maron et al. , 2019a ) . Our work is in line with group equivariant networks operating on graphs and sets . Multiscale , multilevel , multiresolution and coarse-grained techniques have been widely applied to graphs and discrete domains such as diffusion wavelets ( Coifman & Maggioni , 2006 ) ; spectral wavelets on graphs ( Hammond et al. , 2011 ) ; finding graph wavelets based on partitioning/clustering ( Rustamov & Guibas , 2013 ) ; graph clustering and finding balanced cuts on large graphs ( Dhillon et al. , 2005 ) ( Dhillon et al. , 2007 ) ( Chiang et al. , 2012 ) ( Si et al. , 2014 ) ; and link prediction on social networks ( Shin et al. , 2012 ) . Prior to our work , some authors such as ( Zhou et al. , 2019 ) proposed a multiscale generative model on graphs using GAN ( Goodfellow et al. , 2014 ) , but the hierarchical structure was built by heuristics algorithm , not learnable and not flexible to new data that is also an existing limitation of the field . In general , our work exploits the powerful group equivariant networks to encode a graph and to learn to form balanced partitions via back-propagation in a datadriven manner without using any heuristics as in the existing works . In the field of deep generative models , it is generally recognized that introducing a hierarchy of latents and adding stochasticity among latents leads to more powerful models capable of learning more complicated distributions ( Blei et al. , 2003 ) ( Ranganath et al. , 2016 ) ( Ingraham & Marks , 2017 ) ( Klushyn et al. , 2019 ) ( Wu et al. , 2020 ) ( Vahdat & Kautz , 2020 ) . Our work combines the hierarchical variational autoencoder with learning to construct the hierarchy that results into a generative model able to generate graphs at many resolution levels . 3 MULTIRESOLUTION GRAPH NETWORK . 3.1 CONSTRUCTION . An undirected weighted graph G = ( V , E , A , Fv , Fe ) with node set V and edge set E is represented by an adjacency matrix A ∈ N∣V ∣×∣V ∣ , where Aij > 0 implies an edge between node vi and vj with weight Aij ( e.g. , Aij ∈ { 0,1 } in the case of unweighted graph ) ; while node features are represented by a matrix Fv ∈ R∣V ∣×dv , and edge features are represented by a tensor Fe ∈ R∣V ∣×∣V ∣×de . The second-order tensor representation of edge features is necessary for our higher-order message passing networks described in the next section . Indeed , Fv can be encoded in the diagonal of Fe . Definition 1 . A K-cluster partition of graph G is a partition of the set of nodes V into K mutually exclusive clusters V1 , .. , VK . Each cluster corresponds to an induced subgraph Gk = G [ Vk ] . Definition 2 . A coarsening of G is a graph G̃ of K nodes defined by a K-cluster partition in which node ṽk of G̃ corresponds to the induced subgraph Gk . The weighted adjacency matrix Ã ∈ NK×K of G̃ is Ãkk′ = { 1 2 ∑vi , vj∈Vk Aij , if k = k ′ , ∑vi∈Vk , vj∈Vk′ Aij , if k ≠ k ′ , where the diagonal of Ã denotes the number of edges inside each cluster , while the off-diagonal denotes the number of edges between two clusters . Fig . 3.1 shows an example of Defs . 1 and 2 : a 3-cluster partition of the Aspirin C9H8O4 molecular graph and its coarsening graph . Def . 3 defines the multiresolution of graph G in a bottom-up manner in which the bottom level is the highest resolution ( e.g. , G itself ) while the top level is the lowest resolution ( e.g. , G is coarsened into a single node ) . Definition 3 . An L-level coarsening of a graph G is a series of L graphs G ( 1 ) , .. , G ( L ) in which 1 . G ( L ) is G itself . 2 . For 1 ≤ ` ≤ L − 1 , G ( ` ) is a coarsening graph of G ( ` +1 ) as defined in Def . 2 . The number of nodes in G ( ` ) is equal to the number of clusters in G ( ` +1 ) . 3 . The top level coarsening G ( 1 ) is a graph consisting of a single node , and the corresponding adjacency matrix A ( 1 ) ∈ N1×1 . Definition 4 . An L-level Multiresolution Graph Network ( MGN ) of a graph G consists of L−1 tuples of five network components { ( c ( ` ) , e ( ` ) local , d ( ` ) local , d ( ` ) global , p ( ` ) ) } L ` =2 . The ` -th tuple encodes G ( ` ) and transforms it into a lower resolution graph G ( ` −1 ) in the higher level . Each of these network components has a separate set of learnable parameters ( θ ( ` ) 1 , θ ( ` ) 2 , θ ( ` ) 3 , θ ( ` ) 4 , θ ( ` ) 5 ) . For simplicity , we collectively denote the learnable parameters as θ , and drop the superscript . The network components are defined as follows : 1 . Clustering procedure c ( G ( ` ) ; θ ) , which partitions graph G ( ` ) into K clusters V ( ` ) 1 , . . . , V ( ` ) K . Each cluster is an induced subgraph G ( ` ) k of G ( ` ) with adjacency matrix A ( ` ) k . 2 . Local encoder elocal ( G ( ` ) k ; θ ) , which is a permutation equivariant ( see Defs . 5 , 6 ) graph neural network that takes as input the subgraph G ( ` ) k , and outputs a set of node latents Z ( ` ) k represented as a matrix of size ∣V ( ` ) k ∣ × dz . 3 . Local decoder dlocal ( Z ( ` ) k ; θ ) , which is a permutation equivariant neural network that tries to reconstruct the subgraph adjacency matrixA ( ` ) k for each cluster from the local encoder ’ s output latents . 4 . ( Optional ) Global decoder dglobal ( Z ( ` ) ; θ ) , which is a permutation equivariant neural network that reconstructs the full adjacency matrixA ( ` ) from all the node latents ofK clusters Z ( ` ) =⊕kZ ( ` ) k represented as a matrix of size ∣V ( ` ) ∣ × dz . 5 . Pooling network p ( Z ( ` ) k ; θ ) , which is a permutation invariant ( see Defs . 5 , 6 ) neural network that takes the set of node latents Z ( ` ) k and outputs a single cluster latent z̃ ( ` ) k ∈ dz . The coarsening graph G ( ` −1 ) has adjacency matrix A ( ` −1 ) built as in Def . 2 , and the corresponding node features Z ( ` −1 ) =⊕k z̃ ( ` ) k represented as a matrix of size K × dz . Algorithmically , MGN works in a bottom-up manner as a tree-like hierarchy starting from the highest resolution graph G ( L ) , going to the lowest resolution G ( 1 ) ( see Fig . 3.1 ) . Iteratively , at ` -th level , MGN partitions the current graph into K clusters by running the clustering procedure c ( ` ) . Then , the local encoder e ( ` ) local and local decoder d ( ` ) global operate on each of the K subgraphs separately , and can be executed in parallel . This encoder/decoder pair is responsible for capturing the local structures . Finally , the pooling network p ( ` ) shrinks each cluster into a single node of the next level . Optionally , the global decoder d ( ` ) global makes sure that the whole set of node latents Z ( ` ) is able to capture the inter-connection between clusters . In terms of time and space complexity , MGN is more efficient than existing methods in the field . The cost of global decoding the highest resolution graph is proportional to ∣V ∣2 . For example , while the encoder can exploit the sparsity of the graph and has complexity O ( ∣E ∣ ) , a simple dot-product decoder dglobal ( Z ) = sigmoid ( ZZT ) has both time and space complexity ofO ( ∣V ∣2 ) which is infeasible for large graphs . In contrast , the cost of running K local dot-product decoders is O ( ∣V ∣2/K ) , which is approximately K times more efficient .

Summary: This paper proposes Multiresolution Equivariant Graph Variational Autoen-coders (MGVAE) which can learn and generate graphs in a multiresolution and equivariant fashion. Higher order message passing is used to encode the graph while maintaining learning mutually exclusive clusters so that coarsening into a lower resolution, thus creating a hierarchy of latent distributions. The model also maintain an end-to-end permutation equivariant with respect to node ordering. Their experimental results show that MGVAE achieves competitive re- sults with several generative tasks and graph link prediction etc.