PaperlayerView exact revision on arXiv

Source-linked AI summary

OOD-GNN: Out-of-Distribution Generalized Graph Neural Network

Haoyang Li, Xin Wang, Ziwei Zhang, Wenwu Zhu

arXiv:2112.03806v2cs.LG

TL;DR

Existing GNNs can fail when testing graphs differ from training graphs because distribution shifts induce spurious correlations. OOD-GNN decorrelates graph representations using random Fourier features, sample reweighting, and global-local weight estimation, and experiments report superior out-of-distribution generalization across synthetic and real-world benchmarks.

  • Problem

    Existing GNNs generally assume identical training and testing distributions, but complex shifts can substantially degrade performance through spurious correlations.

  • Method

    OOD-GNN uses nonlinear representation decorrelation with random Fourier features and sample reweighting, plus a scalable global-local estimator for graph weights.

  • Results

    OOD-GNN substantially improves out-of-distribution graph classification and regression and greatly outperforms baselines across synthetic and real-world benchmarks.

  • Takeaways & Limitations

    Decorrelating graph representations provides an approach for learning graph predictors intended to generalize under complex and heterogeneous distribution shifts.

  • Takeaways & Limitations

    The theoretical independence result assumes finite kernel expectations and a characteristic kernel.

Abstract

from arXiv · show

Graph neural networks (GNNs) have achieved impressive performance when testing and training graph data come from identical distribution. However, existing GNNs lack out-of-distribution generalization abilities so that their performance substantially degrades when there exist distribution shifts between testing and training graph data. To solve this problem, in this work, we propose an out-of-distribution generalized graph neural network (OOD-GNN) for achieving satisfactory performance on unseen testing graphs that have different distributions with training graphs. Our proposed OOD-GNN employs a novel nonlinear graph representation decorrelation method utilizing random Fourier features, which encourages the model to eliminate the statistical dependence between relevant and irrelevant graph representations through iteratively optimizing the sample graph weights and graph encoder. We further present a global weight estimator to learn weights for training graphs such that variables in graph representations are forced to be independent. The learned weights help the graph encoder to get rid of spurious correlations and, in turn, concentrate more on the true connection between learned discriminative graph representations and their ground-truth labels. We conduct extensive experiments to validate the out-of-distribution generalization abilities on two synthetic and 12 real-world datasets with distribution shifts. The results demonstrate that our proposed OOD-GNN significantly outperforms state-of-the-art baselines.

1 Introduction

Existing GNNs often rely on identical training and testing distributions, but complex shifts can induce spurious correlations and substantially degrade graph-prediction performance. OOD-GNN addresses this challenge by decorrelating graph representations through nonlinear random Fourier features, sample reweighting, and scalable global-local weight estimation.

  • 1 Introduction: Existing GNNs assume identically distributed training and testing graphs, an assumption often violated in applications such as molecular prediction.Testing graphs may exhibit unobserved or uncontrolled shifts relative to training data.
  • 1 Introduction: Complex heterogeneous shifts can cause substantial degradation because GNNs exploit spurious correlations between irrelevant and relevant representations.These shifts may affect graph structure, node features, or both.
  • 1 Introduction: OOD-GNN uses random Fourier features and sample reweighting to eliminate statistical dependence between relevant and irrelevant graph representations.The method targets nonlinear dependencies that are difficult to measure and remove in graph representations.
  • 1 Introduction: A scalable global-local weight estimator jointly maintains local mini-batch and global dataset weights while graph-encoder parameters and weights are optimized iteratively.This design addresses the computational and storage costs of consistently learning one global weight per graph.
  • 1 Introduction: OOD-GNN substantially improves graph classification and regression under distribution shifts on synthetic and real-world benchmarks.The experiments cover synthetic datasets and well-known real-world graph benchmarks.

2 Related Works

Related work covers GNN representation learning, size generalization, expressiveness, representation decorrelation, and disentangled graph representations. OOD generalization differs from ordinary training-distribution generalization and from methods that force factorized graph representations.

  • Graph Neural Network: GNNs generally combine node features and graph topology through message passing and graph pooling to obtain whole-graph representations.This foundational architecture underlies the graph representation-learning methods reviewed here.
  • Size generalization of GNNs: Size-generalization methods train GNNs on smaller graphs and evaluate them on larger graphs using mechanisms such as attention or self-supervised learning.These approaches target changed graph-size distributions rather than the full range of heterogeneous distribution shifts.
  • The expressiveness of GNNs: OOD generalization studies unseen graph topological structures and features, making it orthogonal to expressiveness and generalization within the training distribution.Expressiveness is commonly assessed with the Weisfeiler-Lehman graph isomorphism test.
  • Representation decorrelation: Representation-decorrelation methods target spurious correlations but may incur computational overhead, yield marginal improvements, or require extra supervision.The cited approaches explicitly penalize high correlation.
  • Disentangled graph representation learning: Disentangled graph methods factorize representations or graphs into explanatory components, but factorization can alter semantic meaning and degrade downstream performance.The related work identifies a trade-off between disentanglement and task performance.

3 Method

OOD-GNN learns graph representations for distribution-shifted settings by reweighting samples to reduce dependence among representation dimensions. A global-local estimator makes this decorrelation scalable while iteratively updating graph weights, the encoder, and classifier.

  • 3.1 Notations and Problem Formulation: OOD-GNN jointly optimizes graph weights, a GNN encoder, and a classifier to reduce representation dependence and improve prediction under distribution shifts.The objective combines weighted prediction loss with a decorrelation objective over graph-representation dimensions.
  • 3.2 Statistical Independence with Graph Reweighting: Because relevant and irrelevant representation parts are difficult to identify without costly supervision, the method encourages independence among all representation dimensions.This design targets spurious correlations without requiring labels that distinguish relevant from irrelevant dimensions.
  • 3.2 Statistical Independence with Graph Reweighting: Random Fourier features approximate a kernel-based independence test, enabling nonlinear dependence measurement whose accuracy increases with the number of sampled functions.The paper notes that Q = 5 is practically sufficient for independence measurement.
  • 3.2 Statistical Independence with Graph Reweighting: Graph weights minimize the squared Frobenius norm of a weighted partial cross-covariance matrix, while encoder and classifier updates use weighted prediction loss.Cross-entropy is used for graph classification and mean squared error for graph regression.
  • 3.3 Global-Local Graph Weight Estimator: The global-local estimator combines mini-batch local weights with persistent global representations and weights to balance optimization efficiency with dataset-wide consistency.Global memory size depends on mini-batch size rather than the full training-set size, and each mini-batch costs O((K + 1)|B|) instead of O(N tr).

4 Experiments

The experiments evaluate OOD-GNN on synthetic and real-world graph datasets under distribution shifts. The dataset suite includes varied prediction tasks and split schemes designed to represent complex, realistic shifts.

  • 4 Experiments: The evaluation covers synthetic and real-world datasets, with additional ablation studies and analyses of sensitivity, training dynamics, weight distributions, and time complexity.The supplied experiment passage states that both synthetic and real-world datasets are evaluated.
  • 4 Experiments: The dataset statistics report graph counts, average node and edge counts, task dimensionality, and task type across the benchmark suite.Task types include binary classification, multi-classification, and regression.
  • 4 Experiments: Training, validation, and testing splits are designed to cover complex and realistic distribution shifts.The supplied table caption identifies split methods as part of the dataset characterization.

4.1 Experimental Setup

The experiments compare OOD-GNN with representative GNN baselines on synthetic and real-world graph datasets designed with distribution shifts, using graph classification and related prediction settings.

  • Baselines: The study compares OOD-GNN with GCN, GIN, virtual-node variants, FactorGCN, PNA, TopKPool, and SAGPool.These baselines cover message passing, expressive graph encoders, virtual-node augmentation, disentangled representations, neighborhood aggregation, and graph pooling.
  • Datasets: The evaluation includes two synthetic datasets and real-world molecule, social, and Open Graph Benchmark datasets.The datasets are intended to cover graph-size, graph-feature, and graph-structure distribution shifts.
  • Synthetic Datasets: TRIANGLES trains on graphs with 4–25 nodes and tests on graphs with 4–100 nodes to evaluate size distribution shifts.The task predicts one of 10 triangle-count classes from graph data.
  • Synthetic Datasets: MNIST-75SP evaluates digit classification under noisy and colorized node features while keeping graph structures unchanged.The two shifted test settings add Gaussian noise, with colorized images receiving two additional channels.
  • Implementation Details: Implementation uses PyTorch, 100 training epochs, 20 graph-weight epochs, selected batch sizes and learning rates, and GIN as the graph encoder.Representation and hidden-layer dimensionalities are selected from dataset-dependent candidate sets.

4.2 Results on Synthetic Graphs

On both synthetic benchmarks, OOD-GNN achieves the strongest reported performance under graph-size and graph-feature distribution shifts.

  • Synthetic Results: OOD-GNN consistently achieves the best testing performance on TRIANGLES and MNIST-75SP compared with the evaluated baselines.TRIANGLES tests larger graph sizes, whereas MNIST-75SP tests noisier graph features and colorized images.
  • TRIANGLES: On TRIANGLES, OOD-GNN outperforms baselines on out-of-distribution graphs despite PNA’s impressive training-graph accuracy.The result supports the method’s reported advantage under graph-size shifts.
  • MNIST-75SP: On MNIST-75SP, OOD-GNN consistently achieves the best performance when testing graph features contain larger noise levels.The graph topology remains unchanged in the shifted test graphs, making feature noise the principal distribution change described here.

4.3 Results on Real-world Graphs

Across real-world graph benchmarks, OOD-GNN is reported as the strongest or most consistently competitive method under graph-size and scaffold-based distribution shifts.

  • Molecule and Social Datasets: OOD-GNN consistently yields the best testing performance on COLLAB, PROTEINS, and D&D under graph-size distribution shifts.It improves over the strongest baselines by 2.2% on PROTEINS25, 6.0% on D&D200, and 1.7% on D&D300.
  • Distribution Shifts: The real-world evaluations test graphs split by size or scaffold, creating shifts in graph size, topology, and features.COLLAB, PROTEINS, and D&D use graph-size splits, while OGB uses scaffold-based splits.
  • Open Graph Benchmark: On nine OGB datasets, no baseline is consistently competitive across all tasks, whereas OOD-GNN shows impressive performance across the benchmark.Table 4 reports ROC-AUC for classification and RMSE for regression, with direction indicated by the metric arrows.
  • Interpretation: The COLLAB example identifies graph topology as predictive while warning that size-label correlations can harm predictions on larger OOD graphs.This illustrates how spurious correlations can undermine generalization when graph-size distributions change.

4.4 Ablation Studies

Ablation studies show that larger random Fourier feature dimensions improve OOD performance, while removing them and retaining only linear decorrelation causes a clear drop.

  • Ablation Scope: The ablations are reported on TRIANGLES, D&D300, and OGBG-MOLBACE, with similar patterns stated for the other datasets.The study also notes a trade-off between stronger independence and the computational cost of sampling more random Fourier features.
  • Random Fourier Features: Increasing the dimensionality of random Fourier features consistently improves OOD-GNN performance on the reported datasets.The ablation varies feature dimensionality relative to graph representations, including settings such as 0.2x and 2x.
  • No-RFF Variant: Removing random Fourier features causes a clear performance drop compared with the full method.The no-RFF variant reduces the optimization to eliminating linear correlation rather than encouraging independence across representation dimensions.

4.5 Training Dynamic

The weighted prediction loss converges within 100 epochs across TRIANGLES, D&D300, and OGBG-MOLBACE, reaching dataset-specific values.

  • The weighted prediction loss converges in no more than 100 epochs on TRIANGLES, D&D300, and OGBG-MOLBACE.The final losses are about 0.67, 0.30, and 0.25, respectively.

4.6 Weights Distribution

After training, OOD-GNN assigns non-trivial graph weights whose distributions vary slightly across datasets.

  • OOD-GNN learns non-trivial graph weights, with slightly different distributions across TRIANGLES, D&D300, and OGBG-MOLBACE.

OGBG-MOLBACE

Experiments on OGBG-MOLBACE examine training loss, learned weights, efficiency, parameter count, and hyper-parameter sensitivity. Larger global-weight representations improve performance, while momentum has only a slight effect and excessive depth risks over-smoothing.

  • Efficiency: The method’s time complexity is O(|E| d + |V | d2 + K|B|d2), incorporating graph encoding and global-weight optimization costs.
  • Model Complexity: The method has about 0.9M parameters on OGBG-MOLBACE with five message-passing layers and representation dimensionality 300.
  • Hyper-parameter Analysis: A very large number of graph-encoder layers can cause over-smoothing, while the optimal representation dimensionality differs across datasets.
  • Hyper-parameter Analysis: Larger global-weight representations improve performance by supporting consistent graph sample weights across the dataset.
  • Hyper-parameter Analysis: The momentum coefficient γ has only a slight performance influence: larger values slow updates, whereas smaller values accelerate them.

5 Conclusions

OOD-GNN addresses GNN generalization under complex and heterogeneous distribution shifts through nonlinear representation decorrelation and scalable graph weighting. Experiments on synthetic and real-world datasets report superiority over state-of-the-art baselines for out-of-distribution generalization.

  • OOD-GNN combines random Fourier features and sample reweighting to decorrelate learned graph representations under distribution shifts.
  • Its scalable global-local weight estimator learns consistent graph weights for the whole dataset efficiently.
  • Experiments on synthetic and real-world datasets demonstrate superiority over state-of-the-art baselines for out-of-distribution generalization.
Loading 2112.03806v2…