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 study presents a multi-scale graph VAE with hierarchical graph coarsening. The Gumbel-max trick is employed to generate learnable hard partition for clustering, and permutation equivariant tensor operations are used (Kondor et al., 2018) to construct the group equivariance network. The authos find that the proposed framework exhibits competitive performance in graph generation, molecular generation, molecular representation learning, link prediction and graph-based image generation. Given the applicability of this framework, this proposed framework may be interesting to the graph generation community, and this paper is well-written.