Performance Heterogeneity in Graph Neural Networks: Lessons for Architecture Design and Preprocessing
Lukas Fesser, Melanie Weber
TL;DR
The paper investigates performance heterogeneity in graph-level learning across both message-passing and transformer-based GNNs, revealing that graph topology alone cannot explain per-graph differences. It introduces heterogeneity profiles and leverages the Tree Mover's Distance to connect class-distance ratios with heterogeneity, guiding practical design choices. The authors propose selective rewiring to align graph spectra and a spectral-depth heuristic based on the Fiedler value ($\lambda_2$), demonstrating improved per-graph performance on benchmarks. These insights offer concrete guidelines for preprocessing and architecture selection in heterogeneous graph datasets, with potential to automate model choices in GNN pipelines.
Abstract
Graph Neural Networks have emerged as the most popular architecture for graph-level learning, including graph classification and regression tasks, which frequently arise in areas such as biochemistry and drug discovery. Achieving good performance in practice requires careful model design. Due to gaps in our understanding of the relationship between model and data characteristics, this often requires manual architecture and hyperparameter tuning. This is particularly pronounced in graph-level tasks, due to much higher variation in the input data than in node-level tasks. To work towards closing these gaps, we begin with a systematic analysis of individual performance in graph-level tasks. Our results establish significant performance heterogeneity in both message-passing and transformer-based architectures. We then investigate the interplay of model and data characteristics as drivers of the observed heterogeneity. Our results suggest that graph topology alone cannot explain heterogeneity. Using the Tree Mover's Distance, which jointly evaluates topological and feature information, we establish a link between class-distance ratios and performance heterogeneity in graph classification. These insights motivate model and data preprocessing choices that account for heterogeneity between graphs. We propose a selective rewiring approach, which only targets graphs whose individual performance benefits from rewiring. We further show that the optimal network depth depends on the graph's spectrum, which motivates a heuristic for choosing the number of GNN layers. Our experiments demonstrate the utility of both design choices in practice.