A Unified Invariant Learning Framework for Graph Classification

TL;DR

This work addresses the challenge of out-of-distribution generalization in graph classification by arguing that existing invariant-learning approaches, which focus primarily on semantic invariance, may fail to identify truly minimal stable features. It proposes UIL, a unified framework that enforces both structural invariance in the graph space via graphon-based constraints and semantic invariance in the representation space, thereby guiding the model toward minimal stable features $G_{st}$. The method introduces a Stable & Environmental Feature Separator, Stable Graphon Estimation, and a combined loss $\,\mathcal{L}=\mathcal{L}_{suf}+\alpha\mathcal{L}_{str}+\beta\mathcal{L}_{sem}$, with unsupervised environment discovery via $K$-means. Theoretical results link the structural objective to convergence toward minimal stable features, and extensive experiments on synthetic and real graph OOD benchmarks show UIL consistently improves generalization over strong baselines, supported by ablations and analyses of graphon distances. Overall, UIL provides a principled, practical pathway to robust graph classification under distribution shifts by unifying graph-space and semantic invariances.

Abstract

Invariant learning demonstrates substantial potential for enhancing the generalization of graph neural networks (GNNs) with out-of-distribution (OOD) data. It aims to recognize stable features in graph data for classification, based on the premise that these features causally determine the target label, and their influence is invariant to changes in distribution. Along this line, most studies have attempted to pinpoint these stable features by emphasizing explicit substructures in the graph, such as masked or attentive subgraphs, and primarily enforcing the invariance principle in the semantic space, i.e., graph representations. However, we argue that focusing only on the semantic space may not accurately identify these stable features. To address this, we introduce the Unified Invariant Learning (UIL) framework for graph classification. It provides a unified perspective on invariant graph learning, emphasizing both structural and semantic invariance principles to identify more robust stable features. In the graph space, UIL adheres to the structural invariance principle by reducing the distance between graphons over a set of stable features across different environments. Simultaneously, to confirm semantic invariance, UIL underscores that the acquired graph representations should demonstrate exemplary performance across diverse environments. We present both theoretical and empirical evidence to confirm our method's ability to recognize superior stable features. Moreover, through a series of comprehensive experiments complemented by in-depth analyses, we demonstrate that UIL considerably enhances OOD generalization, surpassing the performance of leading baseline methods. Our codes are available at https://github.com/yongduosui/UIL.

Figures (7)

• Figure 1: The workflow of several prevalent methods grounded in invariant graph learning.
• Figure 2: The overview of the proposed Unified Invariant Learning (UIL) framework.
• Figure 3: Precision of extracting the minimal stable features (AUC) and classification accuracy (%).
• Figure 4: Visualization results of learned stable graphons. $e_1, e_2, e_3$ refer to different environments in Motif dataset.
• Figure 5: Visualization results of the captured stable features.

Theorems & Definitions (12)

• Definition 1: Semantic Invariance Principle
• Definition 2: Estimated Stable Feature
• Definition 3: Minimal Stable Feature Set
• Definition 4: Stable Graphon
• Lemma 1: Weak Regularity
• Definition 5: Structural Invariance Principle
• Proposition 1
• Lemma 2
• Proposition 2
• Definition 6: Graphon