Neural Trees for Learning on Graphs

TL;DR

This paper introduces Neural Trees, a novel GNN architecture that constructs a hierarchical tree-structured graph to overcome local message-passing limitations, enabling better approximation of functions over graphs and improving prediction accuracy.

Contribution

The paper proposes Neural Trees, a new GNN architecture that operates on a hierarchical tree-structured graph, enhancing expressive power and approximation capabilities over traditional message-passing GNNs.

Findings

Neural Trees can approximate any smooth distribution over graphs.

The number of parameters scales exponentially with treewidth but linearly with graph size.

Neural Trees outperform traditional GNNs in semi-supervised node classification.

Abstract

Graph Neural Networks (GNNs) have emerged as a flexible and powerful approach for learning over graphs. Despite this success, existing GNNs are constrained by their local message-passing architecture and are provably limited in their expressive power. In this work, we propose a new GNN architecture -- the Neural Tree. The neural tree architecture does 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. Nodes in the H-tree correspond to subgraphs in the input graph, and they are reorganized in a hierarchical manner such that the parent of a node in the H-tree always corresponds to a larger subgraph in the input graph. We show that the neural tree architecture can approximate any smooth probability distribution function over an undirected graph. We also prove that the number of parameters needed to achieve an $\epsilon$-approximation of the distribution function is exponential in the treewidth of the input graph, but linear in its size. We prove that any continuous $\mathcal{G}$-invariant/equivariant function can be approximated by a nonlinear combination of such probability distribution functions over $\mathcal{G}$. We apply the neural tree to semi-supervised node classification in 3D scene graphs, and show that these theoretical properties translate into significant gains in prediction accuracy, over the more traditional GNN architectures. We also show the applicability of the neural tree architecture to citation networks with large treewidth, by using a graph sub-sampling technique.