Unleash Graph Neural Networks from Heavy Tuning

TL;DR

The paper tackles the substantial cost and overfitting risk of hyperparameter tuning in graph neural networks by introducing GNN-Diff, a graph-conditioned latent diffusion framework that learns from checkpoints saved during a light coarse search to generate high-performing GNN parameters. GNN-Diff leverages a graph autoencoder to encode X and A into a graph condition and a latent DDPM to sample latent parameter representations, which are reconstructed into full GNN weights via a parameter autoencoder. Across node classification benchmarks on both homophilic and heterophilic graphs, GNN-Diff often matches or surpasses grid search while dramatically reducing computation, and ablations confirm the value of graph conditioning for robust parameter generation. This approach shifts GNN optimization from exhaustive search toward exploring the population distribution of good parameters, enabling faster deployment of effective GNNs with improved generalization.

Abstract

Graph Neural Networks (GNNs) are deep-learning architectures designed for graph-type data, where understanding relationships among individual observations is crucial. However, achieving promising GNN performance, especially on unseen data, requires comprehensive hyperparameter tuning and meticulous training. Unfortunately, these processes come with high computational costs and significant human effort. Additionally, conventional searching algorithms such as grid search may result in overfitting on validation data, diminishing generalization accuracy. To tackle these challenges, we propose a graph conditional latent diffusion framework (GNN-Diff) to generate high-performing GNNs directly by learning from checkpoints saved during a light-tuning coarse search. Our method: (1) unleashes GNN training from heavy tuning and complex search space design; (2) produces GNN parameters that outperform those obtained through comprehensive grid search; and (3) establishes higher-quality generation for GNNs compared to diffusion frameworks designed for general neural networks.

Figures (7)

• Figure 1: GNN-Diff Overview. (1) Input graph data: input graph signals, adjacency matrix, and train/validation ground truth labels. (2) Parameter collection: a coarse search with a small search space is conducted to select an appropriate hyperparameter configuration, which is then used to collect model checkpoints. (3) Training: PAE and GAE are firstly trained to produce latent parameters and graph conditions, and then G-LDM learns how to recover latent parameters from white noise with graph conditions. (4) Inference: after sampling latent parameters from G-LDM, GNN parameters are reconstructed with a PAE decoder and returned to target GNN for prediction.
• Figure 2: (a) Scatter plot between validation and test accuracy. Black circles indicate the position of the final models selected by three methods. (b) Visualization of generated and coarse parameters after dimension reduction with isomap. Each bubble represents one parameter sample, and the bubble size shows the corresponding test accuracy. (c) The kernel density estimation plot shows the distribution of test accuracy associated with parameters learned with grid search and coarse search and generated with GNN-Diff.
• Figure 3: Distribution of parameters generated with/without graph condition for GCN2 on Cora. Last two plots on the right show the distribution of corresponding test accuracy.
• Figure 4: Ablation study on generative condition. GNN-Diff employs graph conditions provided by GAE. The model is equivalent with p-diff wang2024neural when no condition is applied (Uncon.).
• Figure 5: Time costs of comprehensive grid search, coarse search, and GNN-Diff training. The sampling time of GNN-Diff is negligible (less than 10 seconds to generate 100 GNNs).