Learning to Pool in Graph Neural Networks for Extrapolation
This addresses a critical bottleneck in GNNs for researchers and practitioners needing robust generalization in graph-structured data, though it is incremental as it builds on existing pooling methods.
The paper tackles the problem of graph neural networks (GNNs) failing to generalize to out-of-distribution data due to poor pooling function choices, and presents GNP, a trainable pooling function that enables GNNs to extrapolate well across various tasks, sometimes outperforming existing best choices.
Graph neural networks (GNNs) are one of the most popular approaches to using deep learning on graph-structured data, and they have shown state-of-the-art performances on a variety of tasks. However, according to a recent study, a careful choice of pooling functions, which are used for the aggregation and readout operations in GNNs, is crucial for enabling GNNs to extrapolate. Without proper choices of pooling functions, which varies across tasks, GNNs completely fail to generalize to out-of-distribution data, while the number of possible choices grows exponentially with the number of layers. In this paper, we present GNP, a $L^p$ norm-like pooling function that is trainable end-to-end for any given task. Notably, GNP generalizes most of the widely-used pooling functions. We verify experimentally that simply using GNP for every aggregation and readout operation enables GNNs to extrapolate well on many node-level, graph-level, and set-related tasks; and GNP sometimes performs even better than the best-performing choices among existing pooling functions.