Discovering the Representation Bottleneck of Graph Neural Networks

TL;DR

This work identifies a representation bottleneck in graph neural networks where the learned distribution of interaction orders $J^{(m)}$ fails to match the task-specific optimum ${J^{(m)}}^*$, largely due to improper inductive bias from graph construction. It introduces a theory of multi-order interactions and a bias-aware graph rewiring method, ISGR, which dynamically adjusts node receptive fields to better capture informative interaction orders. Across molecular, dynamical, and MD datasets, ISGR consistently improves performance over strong rewiring baselines and adapts to different graph constructions (KNN vs FC). The paper also situates the GNN bottleneck within a broader context of DNN inductive biases, connects multi-order interactions to Shapley-based explanations, and demonstrates the framework’s relevance to visual architectures, highlighting practical implications for AI-for-science applications.

Abstract

Graph neural networks (GNNs) rely mainly on the message-passing paradigm to propagate node features and build interactions, and different graph learning problems require different ranges of node interactions. In this work, we explore the capacity of GNNs to capture node interactions under contexts of different complexities. We discover that GNNs usually fail to capture the most informative kinds of interaction styles for diverse graph learning tasks, and thus name this phenomenon GNNs' representation bottleneck. As a response, we demonstrate that the inductive bias introduced by existing graph construction mechanisms can result in this representation bottleneck, \emph{i.e.}, preventing GNNs from learning interactions of the most appropriate complexity. To address that limitation, we propose a novel graph rewiring approach based on interaction patterns learned by GNNs to dynamically adjust each node's receptive fields. Extensive experiments on both real-world and synthetic datasets prove the effectiveness of our algorithm in alleviating the representation bottleneck and its superiority in enhancing the performance of GNNs over state-of-the-art graph rewiring baselines.

Figures (9)

• Figure 1: Distributions of interaction strengths of EGNN and Molformer in graph-level and node-level prediction tasks. We use double-x axes to represent the order $m$ (bottom) and the ratio $m/n$ (top).
• Figure 2: Different graph constructions of the compound $\ch{C_6H_6O}$.
• Figure 3: (a) The change of $m^*$ over epochs for EGNN on QM7, QM8, Hamiltonian, and Newtonian datasets. (b) Change in interaction strengths over epochs for EGNN on these four datasets.
• Figure 4: The theoretical distributions of $F^{(m)}$ under different $n$.
• Figure 5: (a) The change of interaction strengths during training for ResNet on CIFAR-10 and ImageNet. (b) Change in interaction strengths during training for MLP-Mixer on CIFAR-10 and ImageNet.