Invariant Layers for Graphs with Nodes of Different Types
This work addresses the challenge of handling heterogeneous graph data in machine learning, offering incremental improvements in invariant layer design and approximation bounds.
The paper tackles the problem of designing neural network layers invariant to permutations that preserve node types in heterogeneous graphs, and shows that implementing these layers improves learning of node interactions more effectively than existing techniques. It also provides tighter bounds on tensor sizes for function approximation on graphs and images, reducing from ≤ n(n-1)/2 to ≤ n for graphs and from d^4 to 2d-1 for images.
Neural networks that satisfy invariance with respect to input permutations have been widely studied in machine learning literature. However, in many applications, only a subset of all input permutations is of interest. For heterogeneous graph data, one can focus on permutations that preserve node types. We fully characterize linear layers invariant to such permutations. We verify experimentally that implementing these layers in graph neural network architectures allows learning important node interactions more effectively than existing techniques. We show that the dimension of space of these layers is given by a generalization of Bell numbers, extending the work (Maron et al., 2019). We further narrow the invariant network design space by addressing a question about the sizes of tensor layers necessary for function approximation on graph data. Our findings suggest that function approximation on a graph with $n$ nodes can be done with tensors of sizes $\leq n$, which is tighter than the best-known bound $\leq n(n-1)/2$. For $d \times d$ image data with translation symmetry, our methods give a tight upper bound $2d - 1$ (instead of $d^{4}$) on sizes of invariant tensor generators via a surprising connection to Davenport constants.