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:  The work presented in this paper tackles the expressivity issue with graph neural networks (GNN). Drawing from the probabilistic graphical models' literature, the authors define the notion of G-compatible functions (i.e. can be factorized over the list of maximal cliques in graph G) and show how they are linked to G-invariant functions (i.e. invariant to node permutation). Then, the authors propose their new GNN architecture called Neural Tree. In their approach, an intermediate tree-structured representation, called H-tree, of the input graph G is constructed first. Traditional message passing is then performed on that H-tree instead of the input graph. The authors go on to show how neural trees can approximate G-compatible functions, thus capable of learning graph invariant functions.  Node classification experiments were conducted on two types of graphs: 3D scene graphs, and citation networks. The neural tree architecture was tested with four different aggregation functions (used during message passing) and compared against their typical architecture implementation (GCN, GraphSAGE, GAT, and GIN). Empirically, using the proposed neural tree architecture always yields higher accuracy. Additional experiments were done to measure the impact on the accuracy when changing the size of the training data (% of labeled nodes), and the number of iterations of message passing. Neural trees appear to be better when having access to a limited amount of labeled nodes, and the optimal number of iterations of message passing seems to follow a similar trend as with the standard GCN architecture.