Research map / Graph learning / Graph self-supervised learning
Enhancing Size Generalization in Graph Neural Networks through Disentangled Representation Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DisGen, a model-agnostic disentangled representation framework that uses size- and task-invariant graph augmentations, a contrastive size loss, and a linear-projection decoupling loss to improve size generalization of GNN graph classifiers.
Problem
GNN graph classifiers trained on small graphs often degrade on substantially larger test graphs because size information remains entangled in graph representations.
Approach
Disentangle size- and task-related representations by generating size-invariant (edge-removed) and task-invariant (node-removed) augmented views guided by a pre-trained GNN and explainer, applying a contrastive loss on size representations, and minimizing shared information between hidden size and task representations via a linear-projection residual loss (1/D^2) with theoretical analysis.
Key results
- GT+DisGen reaches 88.01 F1 on BBBP large test vs 83.37 for GT alone, and 79.22 vs 71.62 on PROTEINS large test.
- Average large-test improvement over the original backbone is up to +6.64 F1 (GCN+DisGen), with small-test gains up to +2.50.
- Removing the decoupling loss reduces large-test F1 by 5.41 on average, and a standard decorrelation loss does not compensate for this.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add DISGEN to GCN | GCN · BBBP large test | F1: 76.01 -> 85.63 | Helped |
| add DISGEN to GCN | GCN · PROTEINS large test | F1: 72.71 -> 83.41 | Helped |
| add DISGEN to GCN | GCN · GraphSST2 large test | F1: 83.55 -> 85.18 | Helped |
| add DISGEN to GCN | GCN · NCI1 large test | F1: 38.86 -> 43.46 | Helped |
| add DISGEN to GCN | GCN · REDDIT-BINARY large test | F1: 41.40 -> 45.16 | Helped |
| add DISGEN to GCN | GCN · FRANKENSTEIN large test | F1: 34.14 -> 38.00 | Helped |
| add augmentations (ablation) | DisGen with GCN · BBBP large test | F1: 80.29 -> 85.63 | Helped |
| add decoupling loss (ablation) | DisGen with GCN · BBBP large test | F1: 78.72 -> 85.63 | Helped |
| use decoupling loss instead of decorrelation loss (ablation) | DisGen with GCN · BBBP large test | F1: 73.28 -> 85.63 | Helped |
| add augmentations and decoupling loss rather than use decorrelation loss without augmentations (ablation) | DisGen with GCN · BBBP large test | F1: 72.54 -> 85.63 | Helped |
| add size-invariant branch (ablation) | DisGen with GCN · BBBP large test | F1: 82.45 -> 85.63 | Helped |
| add task-invariant branch (ablation) | DisGen with GCN · BBBP large test | F1: 81.01 -> 85.63 | Helped |
Surprising findings
- Standard baselines SSR and CIGAv2 fail to improve size generalization for the graph transformer backbone.
- Several baselines (e.g., RPGNN with GCN on PROTEINS large, IRM on GT NCI1 large) degrade performance below the backbone.
Limitations
- Theoretical guarantees rely on strong assumptions (twice differentiability, decoupling in terms of ground-truth size and label vectors) that are hard to verify in practice.
- Experiments are limited to graph classification on relatively small datasets with heuristic size splits and single-run-level variance reporting over five runs.
- Augmentation depends on a pre-trained GNN and explainer, and hyperparameters (beta weights, k1, k2) vary per dataset and backbone.
Experimental setup
Datasets: BBBP PROTEINS GraphSST2 NCI1 REDDIT-BINARY FRANKENSTEIN
Benchmarks: BBBP PROTEINS GraphSST2 NCI1 REDDIT-BINARY FRANKENSTEIN
Metrics: F1 score
Baselines: SizeShiftReg CIGAv2 RPGNN IRM GCN GIN Graph Transformer
Builds on: PGExplainer IRM contrastive learning SizeShiftReg CIGA
Closest papers by meaning
- Disentangled Graph Self-supervised Learning for Out-of-Distribution Generalization (ICML 2024 · similarity 0.89)
- Diffusing to the Top: Boost Graph Neural Networks with Minimal Hyperparameter Tuning (ICLR 2025 · similarity 0.88)
- Towards Bridging Generalization and Expressivity of Graph Neural Networks (ICLR 2025 · similarity 0.87)
- A Manifold Perspective on the Statistical Generalization of Graph Neural Networks (ICML 2025 · similarity 0.87)
- Subgraph Invariant Learning Towards Large-Scale Graph Node Classification (AAAI 2025 · similarity 0.87)
- On the Rademacher Complexity of Graph Neural Networks: Unifying Expressivity and Geometry (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.