Source-linked AI summary
Graph Contrastive Learning Automated
Yuning You, Tianlong Chen, Yang Shen, Zhangyang Wang
TL;DR
GraphCL’s performance depends on manually chosen augmentations for diverse graph datasets, motivating automated selection. JOAO formulates augmentation choice as bi-level min-max optimization and adds augmentation-aware projection heads. Across diverse datasets, JOAO is comparable to hand-tuned GraphCL, while JOAOv2 can outperform it in specific bioinformatics domains.
Problem
GraphCL requires hand-picked, dataset-specific augmentations because graph datasets are diverse, and selecting them relies on trial-and-error or loose heuristics.
Method
JOAO uses bi-level optimization, instantiated as min-max optimization, to optimize augmentation-pair sampling from self-supervised signals, with augmentation-aware projection heads for varying augmentations.
Results
Across diverse datasets, JOAO performs comparably to hand-tuned GraphCL, while JOAOv2 achieves better performance than GraphCL on specific bioinformatics datasets.
Takeaways & Limitations
JOAO automates augmentation selection without dataset-specific manual tuning, and its selections generally align with augmentation best practices from exhaustive hand tuning.
Takeaways & Limitations
The alternating gradient-descent solver is an approximation for non-convex non-concave min-max optimization and lacks a theoretical convergence guarantee.
Abstract
from arXiv · showhide
Self-supervised learning on graph-structured data has drawn recent interest for learning generalizable, transferable and robust representations from unlabeled graphs. Among many, graph contrastive learning (GraphCL) has emerged with promising representation learning performance. Unfortunately, unlike its counterpart on image data, the effectiveness of GraphCL hinges on ad-hoc data augmentations, which have to be manually picked per dataset, by either rules of thumb or trial-and-errors, owing to the diverse nature of graph data. That significantly limits the more general applicability of GraphCL. Aiming to fill in this crucial gap, this paper proposes a unified bi-level optimization framework to automatically, adaptively and dynamically select data augmentations when performing GraphCL on specific graph data. The general framework, dubbed JOint Augmentation Optimization (JOAO), is instantiated as min-max optimization. The selections of augmentations made by JOAO are shown to be in general aligned with previous "best practices" observed from handcrafted tuning: yet now being automated, more flexible and versatile. Moreover, we propose a new augmentation-aware projection head mechanism, which will route output features through different projection heads corresponding to different augmentations chosen at each training step. Extensive experiments demonstrate that JOAO performs on par with or sometimes better than the state-of-the-art competitors including GraphCL, on multiple graph datasets of various scales and types, yet without resorting to any laborious dataset-specific tuning on augmentation selection. We release the code at https://github.com/Shen-Lab/GraphCL_Automated.
1. Introduction
Graph contrastive learning can learn robust representations from unlabeled graphs, but its effectiveness depends on manually selecting dataset-specific augmentations. JOAO automates this selection through bi-level optimization and augments GraphCL with augmentation-aware projection heads.
- Graph contrastive learning extends contrastive learning to learn generalizable, transferable, and robust representations from unlabeled graph data.
- GraphCL requires hand-picked augmentations for each diverse graph dataset, typically chosen through empirical rules of thumb and trial-and-error.
- JOAO is a principled bi-level optimization framework that automatically, adaptively, and dynamically selects graph data augmentations.
- JOAO can be instantiated as min-max optimization and solved with an alternating gradient-descent algorithm.
- The augmentation-aware projection head keeps one nonlinear projection head per augmentation pair and routes each selected augmentation through its corresponding head.
- Experiments show GraphCL with JOAO performs on par with or sometimes better than state-of-the-art competitors across diverse graph datasets without dataset-specific manual tuning.
- JOAO aims to match results from exhaustive manual tuning rather than necessarily surpassing every manually tuned method.
2. Preliminaries and Notations
GraphCL encodes augmented graph views and maximizes their agreement, but its augmentation pair is fixed manually for each dataset. The framework uses graph and projection encoders with cosine-similarity-based contrastive learning.
- A graph G = {V, E} consists of nodes V, edges E, and node features X_v ∈ R^D, while a GNN maps G to a D′-dimensional vector.
- Self-supervised graph learning can exploit unlabeled data for generalizable, transferable, and robust representations, but earlier tasks often require domain knowledge.
- GraphCL maximizes agreement between two augmented graph views to enforce perturbation invariance in GNNs.
- GraphCL samples two augmentation operators from NodeDrop, Subgraph, EdgePert, AttrMask, and Identical.
- The GraphCL pipeline uses a shared-weight GNN and projection head, cosine similarity, and graph samples as the negative-sampling distribution.
- GraphCL pre-fixes its augmentation pair per dataset, making the distribution over augmentation pairs a Dirac distribution selected through loose heuristics.
3. Methodology
JOAO learns graph augmentation-pair sampling through a bi-level min-max optimization procedure, using self-supervised signals rather than labeled validation data. An augmentation-aware projection-head design addresses distribution distortion from dynamically varying augmentations, while experiments examine optimization behavior, selection alignment, and performance.
- JOAO optimization: JOAO optimizes the augmentation-pair sampling distribution jointly with GraphCL training, using the self-supervised objective without downstream labeled data.The upper level optimizes model parameters, while the lower level optimizes augmentation-pair probabilities.
- JOAO optimization: The min-max formulation uses a prior-regularized objective, with a uniform prior encouraging diverse augmentation selections rather than collapse to a few pairs.The distance term compares the learned sampling distribution with the prior distribution.
- Optimization procedure: Alternating gradient descent performs upper-level minimization and lower-level maximization, while projected gradient descent updates the sampling distribution on the probability simplex.The lower-level objective is approximated to make gradients with respect to augmentation-pair probabilities tractable.
- Optimization procedure: AGD is a computationally motivated heuristic for the non-convex, non-concave problem, whose theoretical convergence guarantee remains open; empirical convergence is examined on NCI1 and PROTEINS across γ values.More exact bi-level optimization approaches are described as having high time or space complexity.
- Sanity check: Across four datasets, JOAO selections generally align with augmentation pairs that manual GraphCL tuning found effective, including dataset-dependent preferences for EdgePert, AttrMask, NodeDrop, and Subgraph.The alignment is reported through comparisons of JOAO sampling distributions with exhaustive augmentation-pair performance.
- Augmentation-aware projection heads: JOAO uses augmentation-aware projection heads to separate feature distributions produced by different augmentation pairs, routing each sampled augmentation through its corresponding head.The design addresses the possibility that varying and aggressive augmentations distort the training distribution and reduce fit to the original distribution.
4. Experiments
The experiments compare JOAO and JOAOv2 with heuristic self-supervised methods and GraphCL across diverse graph datasets and bioinformatics domains. Table 3 summarizes the main performance results.
- The evaluation covers diverse graph sources and specific bioinformatics domains, using comparisons against heuristic methods and GraphCL with predefined augmentation rules.
- Table 3 provides a summary of JOAO's performance across the reported experimental settings.
4.1. Datasets and Experiment Settings
The study uses diverse TUDataset graphs and domain-specific bioinformatics datasets under semi-supervised, unsupervised, transfer-learning, and large-scale settings. Experiments follow established protocols and specified GNN architectures, augmentations, and optimization hyperparameters.
- Datasets: Datasets include molecules, proteins, computer-vision graphs, relation networks, and additional domain-specific bioinformatics benchmarks.
- Learning protocols: Learning experiments evaluate semi-supervised, unsupervised representation, transfer-learning, and large-scale graph settings following protocols from state-of-the-art work.
- GNN architectures & augmentations: The experiments adopt setting-specific GNN architectures, including ResGCN for semi-supervised learning and GIN for unsupervised, transfer, and large-scale experiments.
- GNN architectures & augmentations: The optimization trade-off parameter γ is tuned over {0.01, 0.1, 1}, while augmentation strength uses GraphCL's default value of 0.2.
Training from scratch (with augmentations) and graph
The paper situates its experiments against graph-kernel baselines, heuristic self-supervised methods, and GraphCL using manually specified augmentation pools. These comparisons span different graph representations and domain assumptions.
- The comparisons include random-initialization training and graph-kernel methods such as GL, WL, and DGK.
- Heuristic self-supervised baselines use domain-informed objectives including edge reconstruction, vertex-feature masking, and substructure preservation.
- GraphCL constructs augmentation pools with dataset-type-specific rules: NodeDrop and Subgraph for biochemical molecules, all augmentations for dense networks, and all except AttrMask for sparse networks.
4.3. Results
JOAO matches or improves on manually tuned GraphCL across diverse datasets, while JOAOv2 generally outperforms heuristic methods and improves generalization to unseen bioinformatics datasets. However, domain-specific heuristics can remain stronger in particular domains, whereas JOAOv2 also scales well to large datasets.
- Diverse datasets from TUDataset: JOAO matches or beats manually tuned GraphCL in 7 out of 10 semi-supervised TUDataset experiments, despite a slightly worse average rank.
- Diverse datasets from TUDataset: JOAO's automatically selected augmentation pairs are generally consistent with GraphCL's empirically derived best practices in both semi-supervised and unsupervised learning.
- Diverse datasets from TUDataset: JOAOv2 achieves best average ranks of 2.0 at 1% labels and 2.8 at 10% labels in semi-supervised learning, while ranking 2.8 in unsupervised learning versus GraphCL's 2.6.
- Diverse datasets from TUDataset: Across diverse datasets, JOAOv2 generally outperforms heuristic self-supervised methods, with reported average-ranking gaps under both 1% and 10% label rates.
- Specific bioinformatics datasets: JOAOv2 underperforms some domain-specific heuristic methods on specific bioinformatics datasets, although it improves average rank over GraphCL and JOAO.
- Specific bioinformatics datasets: On unseen bioinformatics datasets, JOAO improves average rank by 0.1 and JOAOv2 by 0.3 compared with GraphCL's empirically predefined augmentation rules.
- Large datasets: On ogbg-ppa, JOAO gains over GraphCL exceed 3.49% accuracy at 1% labels and 1.55% at 10% labels, while JOAO and JOAOv2 scale at least as well as GraphCL.
4.4. Summary of Main Findings
JOAO performs comparably to or better than established methods across diverse graph-learning settings while automating augmentation selection. JOAOv2 further improves performance, including on bioinformatics and large-scale datasets.
- Across diverse datasets, JOAO with adaptive augmentation selection performs comparably to GraphCL with exhaustively hand-tuned augmentation rules.
- Augmentation-aware projection heads further boost JOAOv2 performance, sometimes beyond the original JOAO variant.
- Transfer-learning results are reported for bioinformatics datasets, while semi-supervised results use accuracy on ogbg-ppa and F1 score on ogbg-code.
- On bioinformatics datasets, JOAOv2 outperforms GraphCL configured with empirical rules not derived from those domains.
- JOAO and JOAOv2 outperform heuristic self-supervised methods with few exceptions, while domain knowledge may further enhance them.
- JOAOv2 scales to large datasets as well as GraphCL and sometimes achieves more significant improvements than on smaller datasets.
5. Conclusions & Discussions
The paper presents JOAO as a unified framework for automatically and dynamically selecting GraphCL augmentations, with an augmentation-aware projection head addressing distribution distortion. Experiments show performance on par with or sometimes better than state-of-the-art competitors without dataset-specific manual tuning, while full automation remains limited by the augmentation pool.
- JOAO automatically and dynamically selects GraphCL augmentations through a unified bi-level optimization framework instantiated as min-max optimization.
- An augmentation-aware projection head addresses potential training-distribution distortion from JOAO’s aggressive and varying augmentations.
- JOAO and its variant perform on par with or sometimes better than state-of-the-art competitors across graph datasets of various scales and types without tedious manual tuning.
- JOAO still depends on human prior knowledge to construct and configure the augmentation pool, so full automation remains an open direction.