[Paper Review] Graph Contrastive Learning Automated

Yuning You, Tianlong Chen+2 more|arXiv (Cornell University)|Jun 10, 2021

Advanced Graph Neural NetworksComputer Science93 citations

TL;DR

JOAO introduces a bi-level optimization framework to automatically select data augmentations for graph contrastive learning, making GraphCL adaptation to new graphs automatic, adaptive, and dynamic; JOAOv2 adds augmentation-aware projection heads for improved robustness.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for dataset-specific, automatic augmentation selection in graph contrastive learning due to graph heterogeneity.
• Propose JOAO, a unified bi-level optimization framework, to automatically and adaptively select augmentation pairs for GraphCL.
• Show that JOAO matches or surpasses state-of-the-art GraphCL with manual augmentation tuning across diverse datasets.
• Introduce augmentation-aware projection heads to mitigate distribution shifts caused by dynamic augmentations.
• Demonstrate that JOAO selections align with previously observed best practices while removing the need for labeled validation data.

Proposed method

• Formulate JOAO as a bi-level optimization where the upper level minimizes the GraphCL loss and the lower level optimizes the augmentation pair distribution.
• Instantiate JOAO as a min-max problem with a prior distribution to encourage diverse augmentations and a distance term controlling deviation from the prior.
• Solve via alternating gradient descent between updating network parameters and updating augmentation distributions."
• Introduce a new augmentation-aware projection head mechanism with one projection head per augmentation type to route features accordingly during training."
• Provide JOAOv2 by integrating the augmentation-aware heads into the JOAO objective, keeping JOAO’s optimization structure while improving robustness to distortion from varying augmentations.

Experimental results

Research questions

• RQ1 Can JOAO automatically discover effective augmentation policies for GraphCL without labeled validation data?
• RQ2 Do JOAO and JOAOv2 achieve comparable or better performance than hand-tuned GraphCL and heuristic methods across diverse graph datasets?
• RQ3 Do augmentation selections produced by JOAO align with known best practices observed from manual tuning?
• RQ4 Does the augmentation-aware projection head mitigate distribution shifts caused by dynamic augmentations during training?
• RQ5 Is the JOAO framework adaptable to other graph self-supervised learning frameworks beyond GraphCL?

Key findings

• JOAO achieves comparable or better performance than state-of-the-art methods across multiple graph datasets without dataset-specific augmentation tuning.
• JOAO selections generally align with prior best practices, indicating sensible automatic augmentation choices.
• JOAOv2 with augmentation-aware projection heads yields further improvements over JOAO on certain datasets and augmentation strengths.
• The framework is flexible and not tied to GraphCL, enabling integration with other graph contrastive learning approaches.
• JOAO demonstrates that automatic augmentation optimization can scale to diverse graph data without requiring labeled downstream validation.