Adversarial Graph Augmentation to Improve Graph Contrastive Learning

In plain words

Graph contrastive learning trains a graph neural network without labels. It teaches the network that two altered copies of the same graph are one graph. Usually people choose the alterations by hand, for example deleting random edges. This paper lets the network learn them instead. A second network, the augmenter, learns which edges to drop, and it is pushed to make the copies as different as possible. A penalty stops it from dropping too many edges. The main encoder keeps learning to match each graph with its altered copy, so it must rely on the parts of the graph that survive the alteration.

The paper argues that encoders trained this way keep less of the redundant detail that can hurt downstream tasks. It reports gains over hand-chosen augmentation on molecule and social-network benchmarks. These numbers are the authors' own.

TL;DR

AD-GCL is a min-max contrastive method. An encoder maximises an InfoNCE agreement between each graph and a copy with soft-dropped edges, while a GNN augmenter minimises the same agreement and pays a penalty on the expected drop ratio. The paper gives a theoretical link to graph information bottleneck §3.1. It reports gains on 18 datasets under unsupervised, transfer and semi-supervised protocols §5. Use it if you want a learnable edge-drop augmentation for GIN-style molecular or graph-level pretraining, and check the pitfalls, since the code differs from the paper in several places.

Problem

Graph contrastive learning (GCL) depends on hand-designed augmentations, and their choice strongly affects downstream results. The paper says the usual options either use domain knowledge or require extensive evaluation of combinations §1. Maximising mutual information between a graph and its representation can also keep redundant features that do not help the task. It can even hurt, as a motivating experiment shows: an encoder trained on random labels still keeps one-to-one graph identity but generalises worse §3.1 fig. 2 app. G. The information bottleneck principle addresses this, but it needs downstream labels, which self-supervised training lacks §1 §3.1.

Method

The paper defines a min-max objective over an augmentation family. The encoder f maximises the mutual information between a graph and its augmented copy. The augmentation T minimises the same quantity eq. 6 §3.1:

\min_{T \in \mathcal{T}} \max_{f} I(f(G); f(t(G))), \quad t(G) \sim T(G) eq. 6

The instantiation in §3.2 uses edge dropping. Each edge e has a drop probability ω_e, computed from node embeddings of a GNN augmenter with K layers eq. 1 and a small MLP over the two endpoints eq. 7. The final objective adds a penalty λ_reg times the mean drop probability per graph eq. 8. The encoder and the augmenter are updated alternately alg. 1.

Procedure, following alg. 1 and §3.2:

1. Encode each graph with the GIN encoder f, then project it with the 2-layer projection head g §3.2.
2. Run the augmenter GNN on the graph and compute per-edge logits eq. 7. The paper states the drop probability as ω_e; the code applies a sigmoid to the logit to obtain the keep weight test_minmax_ogbg.py.
3. Relax the Bernoulli keep variable with the Gumbel-max trick, so the sample is differentiable §3.2. The relaxed keep value is p_e = Sigmoid((log δ − log(1 − δ) + ω_e) / τ) with δ ~ Uniform(0, 1), and τ → 0 makes it nearly binary. The paper gives no numeric τ.
4. Encode the augmented graph with the same encoder, using the relaxed values p_e §3.2.
5. Compute the InfoNCE estimate over the minibatch with cosine similarity eq. 9 §3.2.
6. Penalise the augmenter by the mean drop probability per graph eq. 8.
7. Update the augmenter by gradient ascent on the NCE loss minus λ_reg times the penalty. Then update the encoder and projection head by gradient descent on the NCE loss alg. 1 eq. 8.

Estimator for the mutual information in step 5 eq. 9:

\hat{I} = \frac{1}{m} \sum_{i=1}^{m} \log \frac{\exp(sim(z_{i,1}, z_{i,2}))} {\sum_{i' \neq i} \exp(sim(z_{i,1}, z_{i',2}))} eq. 9

Here z_{i,1} = g(f(G_i)) and z_{i,2} = g(f(t(G_i))), sim is cosine similarity, and m is the minibatch size §3.2. The training loss in alg. 1 is the negative of this estimate, averaged over the minibatch.

Pseudocode for one minibatch step, following alg. 1 and §3.2. The wording is ours. The paper does not give the code's exact symmetric loss, which is covered under Paper vs code.

# One AD-GCL minibatch step  [alg. 1] [§3.2]
# Input:  graphs (minibatch of G); encoder f; projection head g;
#         augmenter T; regularisation strength lam
# Output: gradient for the augmenter (ascent), gradients for f and g

def ad_gcl_step(graphs, encoder, head, augmenter, lam):
    # Drop logit omega_e for every edge  [eq. 7]
    omega = augmenter(graphs)  # (num_edges,)
    # Relaxed keep weight p_e in (0, 1), Gumbel-max  [§3.2]
    keep = relaxed_keep(omega)
    # Augmented graph: edges are weighted, not removed  [§3.2]
    anchors = head(encoder(graphs))
    views = head(encoder(graphs, edge_weight=keep))
    # Mean drop probability per graph, averaged over the batch  [eq. 8]
    reg = mean_drop_per_graph(omega, graphs)
    # InfoNCE over the minibatch  [eq. 9]
    nce = info_nce_loss(anchors, views)
    # Augmenter: ascend on nce - lam * reg  [alg. 1] [eq. 8]
    augmenter_grad = grad(nce - lam * reg, augmenter.params)
    # Encoder and head: descend on nce, keep weights detached
    # [alg. 1] [code: test_minmax_ogbg.py#L156]
    encoder_grad = grad(nce, encoder.params + head.params)
    return augmenter_grad, encoder_grad

Key equations, in the paper's numbering:

• GNN message passing and graph pooling eq. 1 eq. 2.
• InfoMax objective, the basis of GCL eq. 3.
• Information bottleneck form for graphs, GIB eq. 4.
• GDA-based GCL with two fixed augmentations eq. 5.
• AD-GCL with a learned augmentation family eq. 6.
• Edge keep weights from the augmenter eq. 7.
• Final regularised objective, with λ_reg eq. 8.
• InfoNCE estimator eq. 9.

Theorem 1 §3.1 gives two bounds on the optimal pair (f, T) under assumptions: f is as powerful as the 1-WL test, and the graph space is countable. The first bounds the information kept about G that is irrelevant to Y. The second lower bounds the information about Y. The proof is in app. B.

Hyperparameters

Values for the unsupervised setting unless noted. The first column lists what the paper states. Gaps are marked "not stated".

Optimal λ_reg per dataset for AD-GCL-OPT is listed in app. F (not repeated here).

Results

Author-reported. Setup: GIN encoder, unsupervised pretraining on each dataset, then a linear probe, with the protocol in §5.1 and app. G. Std over 10 runs. The paper marks differences of at least 0.5 std in bold and of at least 2 std with a star tab. 1; the markings are not reproduced here.

Unsupervised, TU datasets (classification accuracy %) tab. 1 top:

Unsupervised, OGB and ZINC-10K (RMSE and MAE lower is better; ROC-AUC % higher is better) tab. 1 bottom:

The paper's claims from these tables §5.1:

• Relative to InfoGraph and GraphCL on TU, AD-GCL is significantly better on 5 of 8 and 6 of 8 datasets.
• AD-GCL-FIX and AD-GCL-OPT give similar results; the paper reads this as robustness to λ_reg §5.1.
• The paper says AD-GCL beats the non-adversarial ablation NAD-GCL on every dataset §5.1. The table has one exception to the FIX-vs-FIX comparison: NAD-GCL-OPT (53.43) beats AD-GCL-FIX (53.00) on RDT-M5K tab. 1.

Regularisation analysis fig. 3: a large λ_reg gives a small expected drop ratio. For classification, λ_reg from 0.1 to 10 gives drop ratios from about 80% down to under 10%, and performance is insensitive to λ_reg §5.2. For regression, a small λ_reg, meaning a larger drop ratio, performs better app. F.

Transfer learning, chemistry and biology tab. 2 (mean ROC-AUC %, 10 runs):

The paper reports a mean rank of 2.4 across the nine datasets for AD-GCL-FIX and says it significantly beats baselines on 3 of 9 §5.3.

Semi-supervised, TU, 10% labels (10-fold accuracy %, mean ± std) tab. 3:

The paper reports a mean rank of 1.5 for AD-GCL-FIX across the six datasets §5.4.

Headline numbers from the abstract: up to 14% gain in the unsupervised setting, 6% in transfer and 3% in semi-supervised learning, against state-of-the-art GCL methods, over 18 datasets [abstract].

Paper vs code

Code: github.com/susheels/adgcl, MIT license, copyright 2021 susheel suresh, at commit 2605ef8f980934c28d545f2556af5cc6ff48ed18. The repository is linked from the paper's abstract; its owner is the first author. Discrepancies, each with both citations:

• Dropout. Paper uses 0.5 app. G. The code's default is 0.0 for molecule encoders and TU molecule_encoder.py test_minmax_ogbg.py. The repository has no per-dataset config, so which value produced the reported numbers cannot be checked here.
• OGB batch size. Paper uses 32 app. G. The OGB script defaults to 64 test_minmax_ogbg.py.
• Edge dropping is soft. Paper describes edges as removed when p_e = 0 §3.2. The code keeps every edge and multiplies each message by the relaxed keep weight gine_conv.py test_minmax_ogbg.py.
• Regulariser source. Paper penalises the augmenter's drop probability ω_e eq. 8. The code penalises 1 minus the sampled, Gumbel-noised keep value per edge, averaged per graph test_minmax_ogbg.py.
• Relaxation temperature and noise. Paper states only that τ goes to 0 §3.2. The code uses τ = 1.0 and uniform noise with bias 1e-4 test_minmax_ogbg.py.
• NCE loss. Paper's eq. 9 is one-directional and has no temperature. The code's loss uses temperature 0.2 and averages the two directions (symmetric), by default ginfominmax.py ginfominmax.py.
• Augmenter. Paper specifies a K-layer GNN, with K not stated §3.2. The code uses a separate 5-layer encoder of the same type and a 64-unit MLP test_minmax_ogbg.py test_minmax_ogbg.py.
• Augmenter step uses eval mode. The code runs the encoder in eval mode while updating the augmenter, then switches to train mode for the encoder update test_minmax_ogbg.py test_minmax_ogbg.py. Paper does not say.
• ZINC λ default. The ZINC script defaults λ_reg to 0.0 test_minmax_zinc.py, while the paper's FIX setting is 5 §5.1. Set --reg_lambda explicitly when reproducing ZINC.
• Algorithm 1 typo. The paper's Algorithm 1 assigns z_{1,n} twice, and the second assignment should be z_{2,n} alg. 1. The code's two branches are distinct test_minmax_ogbg.py.
• Encoder. Paper says GIN §5.1. Molecule experiments use GINE with bond embeddings molecule_encoder.py.
• Environment. README pins Python 3.8.8, PyTorch 1.8.1, PyTorch Geometric 1.7.0 README.md. Not tested on other versions.

Pitfalls

• Learnable augmentation is a min-max loop. Each minibatch makes an augmenter step and an encoder step. The augmenter is driven by the NCE loss and a penalty, so λ_reg controls the drop ratio. For classification the paper reports expected drop ratios from about 80% (λ_reg = 0.1) down to under 10% (λ_reg = 10) §5.2 fig. 3.
• λ_reg is a hyperparameter, and the paper's choice of it is not blind. AD-GCL-OPT tunes λ_reg on validation data, which the paper marks as using "weak information from the downstream task" §5.1. AD-GCL-FIX uses λ=5 and claims no downstream information. Check which setting a comparison uses.
• Protocol gap with original baselines. The paper evaluates all methods with linear probes. InfoGraph and GraphCL report results with a non-linear SVM. The paper says the SVM makes random untrained GIN look much better and therefore uses linear probes §5.1. Do not compare the paper's TU numbers directly with the original InfoGraph and GraphCL papers.
• Transfer and semi-supervised baselines are copied. Those baseline numbers come from prior papers, not from re-runs app. G. Only the AD-GCL rows were run here.
• Seed count is inconsistent. Unsupervised results are over 10 seeds app. G. The semi-supervised table caption says 5 runs tab. 3.
• Extra cost per step. Each minibatch runs the encoder on original and augmented graphs in both the augmenter update and the encoder update, and the augmenter has its own GNN test_minmax_ogbg.py. The paper reports about 14.4 minutes for AD-GCL on COLLAB on one GPU, against 3.2 hours for JOAO app. H. Timing for other datasets is not stated.
• Theorem assumptions. Theorem 1 is stated for the optimal pair (f, T) under a countable graph space and a 1-WL-powerful encoder §3.1 app. B. The paper uses it as motivation; it does not give guarantees for a trained GNN.
• Edge dropping may not match the task. The paper says edge dropping was chosen as a proof of concept, and that other search spaces may do better and are left for future work §3.2 §5.1. For regression, the paper says the edge-drop family may not be the most apt one app. F.
• Algorithm and loss sign. Eq. 9 is an estimate of mutual information (higher is more correlated). Algorithm 1 uses the negated form as a loss. An implementation that mixes these signs will update the augmenter in the wrong direction.

Lineage

• Builds on: InfoMax-based graph contrastive learning, including InfoGraph and GraphCL §1 §5.1. Information bottleneck for graphs, GIB, with the cited GNN variants §1 §3.1. Gumbel-max relaxation and a related edge-mask parameterisation for GNN explanations §3.2.
• Compared against: RU-GIN, InfoGraph, GraphCL, the non-adversarial ablation NAD-GCL (which the paper builds itself), and a fully supervised GIN anchor §5.1. JOAO is compared in app. H. Transfer baselines are EdgePred, AttrMasking, ContextPred and InfoGraph §5.3. Semi-supervised baselines are GAE, SS-GCN-A, InfoGraph and GraphCL §5.4.
• Superseded by: not known from content in this wiki.