Graph neural networks adapt without labels for better out-of-distribution results

Implicit Feature Reweighting for Label-Free Graph Generalization

Machine Learning

Summary

When computers analyze networks like molecules or social graphs, they often fail if the data changes in ways they haven't seen before. This paper introduces a method that helps these graph-analyzing systems work well even when the data shifts, without needing any example labels during training. The method learns to ignore misleading patterns that only show up in certain environments, improving reliability. The authors tested this on synthetic and real-world drug datasets and found it worked better than previous approaches that also did not use labels.

What this means in practice

Authors

Tianjun Yao, Chen Xi, Xinyang Li, Lixian Zhong

Abstract

Graph-level out-of-distribution (OOD) generalization is critical for deploying graph neural networks in real-world applications where distribution shifts are inevitable. Existing approaches typically require class labels during representation learning to decide which substructures are predictive, limiting their applicability. We propose SCARCE, a label-free representation-learning framework that yields distribution-shift-robust graph representations from an environment partition of the training graphs alone, whether that partition comes from dataset metadata or, when none is available, from clustering a self-supervised embedding. Our method operates in two stages: (1) a variational graph autoencoder is warm-started for graph reconstruction; and (2) an adversarial confusion mechanism forces the encoder to produce representations from which a discriminator cannot distinguish environments, while preserving graph-reconstruction fidelity. We provide theoretical analysis showing that environment confusion contracts environment-discriminative directions in the representation space, which, under regularized downstream classifiers, is equivalent to imposing a stronger directional penalty on environment-associated features. This feature-reweighting mechanism reduces reliance on spurious correlations without erasing information, offering a label-free, continuous relaxation of last-layer reweighting methods. Experiments on a controlled synthetic datasets and the DrugOOD datasets show that SCARCE achieves state-of-the-art OOD generalization among label-free methods, using no class labels during representation learning.