Neural Trees for Learning on Graphs

1 Introduction . Graph-structured learning problems arise in several disciplines , including biology ( e.g. , molecule classification [ 15 ] ) , computer vision ( e.g. , action recognition [ 19 ] , image classification [ 48 ] , shape and pose estimation [ 30 ] ) , computer graphics ( e.g. , mesh and point cloud classification and segmentation [ 20 , 35 , 41 ] ) , and social networks ( e.g. , fake news detection [ 42 ] ) , among others [ 9 ] . In this landscape , Graph Neural Networks ( GNN ) have gained popularity as a flexible and effective approach for regression and classification over graphs . Despite this growing research interest , recent work has pointed out several limitations of existing GNN architectures [ 56 , 43 , 38 , 8 ] . Local message passing GNNs are no more expressive than the WeisfeilerLehman ( WL ) graph isomorphism test [ 56 ] , neither can they serve as universal approximators to all G-invariant ( equivariant ) functions , i.e. , functions defined over a graph G that remain unchanged by ( or commute with ) node permutation . The work [ 12 ] proves an equivalence between the ability to do graph isomorphism testing and the ability to approximate any G-invariant function . Various GNN architectures have been proposed , that go beyond local message passing or use tensor representations , in order to improve expressivity . Graph isomorphism testing , G-invariant/equivariant function approximation , and the generalized k-order WL ( k-WL ) tests have served as end objectives and guided recent progress of this inquiry . For example , k-order linear GNN [ 39 ] and k-order folklore GNN [ 38 ] have expressive powers equivalent to k-WL and ( k + 1 ) -WL test , respectively [ 3 ] . While 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . these architectures can theoretically approximate any G-invariant function ( as k → ∞ ) , they use k-order tensors for representations , rendering them impractical for any k > 3 . There is a need for a new way to look at constructing GNN architectures . With better end objectives to guide theoretical progress . Such an attempt can result in new and expressive GNNs that are provably tractable – if not in general , at least in reasonably constrained settings . A GNN , by its very definition , operates on graph structured data . The graph structure of the data determines inter-dependency between nodes and their features . Probabilistic graphical models present a reasonable and well-established way of articulating and working with such inter-dependencies in the data . Prior to the advent of neural networks , inference algorithms on such graphical models were successfully applied to many real-world problems . Therefore , we pose that a GNN architecture operating on a graph should have at least the expressive power of a probabilistic graphical model , i.e. , it should be able to approximate any distribution defined by a probabilistic graphical model . This is not a trivial requirement as exact inference ( akin to learning the distribution or its marginals ) on a probabilistic graphical model , without any structural constraints on the input graph , is known to be an NP-hard problem [ 13 ] . Even approximate inference on a probabilistic graphical model is known to be NP-hard in general [ 46 ] . A common trick to perform exact inference , consists in constructing a junction tree for an input graph and performing message passing on the junction tree instead . In the junction tree , each node corresponds to a subset of nodes of the input graph . The junction tree algorithm remains tractable for graphs with bounded treewidth , while [ 11 ] shows that treewidth is the only structural parameter , bounding which , allows for tractable inference on graphical models . Contribution . We first define the notion of G-compatible function and argue that approximating G-compatible functions is equivalent to approximating any probability distribution on a probabilistic graphical model ( Section 4 ) ; we also show that G-invariant/equivariant functions considered in related work can be approximated using a nonlinear combination of G-compatible functions . We then propose a novel GNN architecture – the Neural Tree – that can approximate any G-compatible function ( Section 5 ) . Neural trees do not perform message passing on the input graph , but on a tree-structured graph , called the H-tree , that is constructed from the input graph . Each node in the H-tree corresponds to a subgraph of the input graph . These subgraphs are arranged hierarchically in the H-tree such that the parent of a node in the H-tree always corresponds to a larger subgraph in the input graph . The leaf nodes in the H-tree correspond to singleton subsets ( i.e. , individual nodes ) of the input graph . The H-tree is constructed by recursively computing tree decompositions of the input graph and its subgraphs , and attaching them to one another to form a hierarchy . Neural message passing on the H-tree generates representations for all the nodes and important subgraphs of the input graph . We next prove that the neural tree architecture can approximate any smooth G-compatible function defined over a given undirected graph ( Section 6 ) . We also bound the number of parameters required by a neural tree architecture to obtain an -approximation of an arbitrary ( smooth ) G-compatible function . We show that the number of parameters increases exponentially in the treewidth of the input graph , but only linearly in the input graphs size . Thus , for graphs with bounded treewidth , the neural tree can tractably approximate any smooth distribution function . We apply the neural tree architecture for semi-supervised node classification in 3D scene graphs and citation networks ( Section 7 ) . Our experiments on 3D scene graphs demonstrate that neural trees outperform standard , local message passing GNNs , by a large margin . Citation networks on the other hand , typically have large treewidth ; therefore we make use of a recently proposed bounded treewidth graph sub-sampling algorithm [ 62 ] , that sub-samples the input graph ( i.e. , removes edges ) to reduce its treewidth to a specified number . We show that applying the neural tree architecture in conjunction with such sub-sampling algorithm makes our architecture scalable to large graphs while still preserving its advantage over traditional architectures . Our code is publically available at https : //github.com/MIT-SPARK/neural_tree 2 Related Work . Expressive Power of Graph Neural Networks . Since the seminal works [ 18 , 49 ] , various GNN architectures have been proposed including Graph Convolutional Networks ( GCN ) [ 29 , 21 , 14 , 29 , 9 ] , Message Passing Neural Networks ( MPNN ) [ 17 ] , GraphSAGE [ 33 ] , Graph Attention Networks ( GAT ) [ 53 , 34 , 10 ] , message passing GNN [ 17 ] . Limited expressive power of these standard GNNs has been a major concern . For instance , it is known that local message passing GNNs can neither distinguish between non-isomorphic graphs ( provably worse than the 1-Weisfeiler-Lehman ( WL ) test ) [ 56 , 43 ] , nor can they compute even simple graph properties [ 16 ] . Many GNN architectures have been proposed to overcome this expressivity bottleneck . Graph substructure network is proposed in [ 8 ] and is shown to be more powerful than the 1-WL test . k-order GNNs , in which message passing is performed among a subset of nodes in the input graph , is shown to have expressive power equivalent to the generalized k-WL test [ 43 , 38 ] . It is generally understood that to improve the expressivity of GNNs one has to extract features corresponding to important subgraphs , and operate on them . A hierarchical architecture that pools a representation vector from a subset of nodes , at each layer , is proposed in [ 61 ] . A junction-tree based message passing GNN is proposed for molecular graph generation in [ 23 ] . Graph neural networks have been investigated as function approximators since the beginning . [ 50 ] introduces the notion of unfolding equivalence and derives a universal approximation result for graph neural networks . Recent research in developing expressive GNN architectures has been towards approximating graph invariant/equivariant functions [ 40 , 39 , 38 , 27 , 47 ] . While , invariance and equivariance are desirable properties , the problem of designing GNNs that are universal approximators of G-invariant/equivariant functions has been difficult . For instance , the k-order GNNs [ 39 , 38 ] can provably approximate any graph invariant function , but only as k → ∞ , rendering them impractical [ 3 ] . An equivalence between designing GNN architectures to approximate graph invariant functions and graph isomorphism testing is shown in [ 12 ] . The generalization power of GNNs has also been investigated in [ 51 , 16 , 57 , 58 ] . Scene Graphs . Scene graphs are a popular model to abstract information in images or model 3D environments . 2D scene graphs have been used in image retrieval [ 24 ] , caption generation [ 26 , 1 ] , visual question answering [ 44 , 31 ] , and relationship detection [ 37 ] . GNNs are a popular tool for joint object labels and/or relationship inference on scene graphs [ 55 , 36 , 59 , 63 ] . Recently , there has been a growing interest towards 3D scene graphs , which are constructed from 3D data , such as meshes [ 2 ] , point clouds [ 54 ] , or raw sensor data [ 28 , 45 ] . GNNs have been very recently applied to 3D scene graphs for scene layout prediction [ 54 ] or object search [ 32 ] . 3 Problem Statement and Preliminaries . In this section , we state the node classification problem and review standard graph neural networks . Problem . We focus on the standard problem of semi-supervised node classification [ 29 ] . We are given a graph G = ( V , E ) along with node features X = ( xv ) v∈V ; where xv denotes the node feature of node v ∈ V . The graph is not necessarily connected . A subset of nodes A ⊂ V in G are labeled , i.e. , { zv ∈ L | v ∈ A } is given ; here zv denotes the label for node v and L the finite set of label classes . We need to design a model to predict the labels of all the unlabeled nodes v ∈ V\A . Graph Neural Networks ( GNN ) . Various GNN architectures have been successfully applied to solve the node classification problem [ 33 , 29 , 53 , 56 , 23 ] . Standard GNN architectures construct representation vectors for each node in G by iteratively aggregating representation vectors of its neighboring nodes . At iteration t , the representation vector of node v is generated as follows : htv = AGGt ( ht−1v , { ( ht−1u , κu , v ) | u ∈ NG ( v ) } ) , ( 1 ) with h0v , xv , ∀v ∈ V ; where NG ( v ) denotes the set of neighbors of node v in graph G and the aggregation function AGGt can depend on the trainable edge parameters κu , v . This process of sharing and aggregating representation vectors among neighboring nodes in G is often called message passing . This procedure runs for a fixed number of iterations T . The node labels are then generated from the representation vectors at the final iteration T . Node labels are extracted as yv = READ ( hTv ) , ( 2 ) for all v ∈ V . The functions AGGt and READ are modeled as single or multi-layer perceptrons .

Summary: This paper presents a new GNN architecture, which performs message passing on a hierarchical H-tree structure constructed from the original graph. Nodes in the H-tree represent subgraphs of the original graph, and thus message passing on H-tree can achieve higher expressive power by going beyond the node-to-node communication. Theoretical results show that the proposed neural tree architecture can approximate any smooth probability distribution function over an undirected graph. Such expressive power is achieved with number of parameters exponential with the width of the H-tree and linear with the original graph size. Experiments on a scene graph and a citation network show neural tree with existing GNN aggregation functions outperform the corresponding GNN alone.