Graph Contrastive Learning with Adaptive Augmentation

In plain words

A graph is a set of things (nodes) joined by links (edges), such as papers and citations or products and co-purchases. Each node also has a list of properties. The method learns a short vector for every node, so that nodes which play similar roles end up with similar vectors. It learns these vectors without any labels.

It works by making two damaged copies of the same graph, each with some links and some properties hidden at random. The network is trained so that a node's vector in one copy matches its vector in the other copy, and differs from the vectors of every other node. The key idea is where the damage falls. Links and properties that matter less to the overall shape of the graph are hidden more often, so the important parts survive in both copies and the network learns to rely on them.

The learned vectors can then be used with a simple classifier. The authors report that this beats other label-free graph methods on five public benchmark graphs and competes with two methods trained on labels.

TL;DR

GCA is a self-supervised contrastive method for node embeddings. Each training step builds two views of the graph by removing edges and masking feature dimensions, with probabilities set from node centrality, then pulls the two embeddings of each node together and pushes them away from all other nodes. Its main claim is that adaptive augmentation beats uniform dropping. Results are author-reported, and the official code differs from the paper in several implementation details (see Paper vs code).

Problem

Supervised graph neural networks need many labelled nodes §1. Contrastive methods avoid labels but depend on how the two views are made. Earlier graph methods used little augmentation: DGI shuffles node features, GMI uses none, and MVGRL uses graph diffusion with uniform subgraph sampling §1 §2. The authors argue that uniform edge dropping removes influential edges as often as unimportant ones, which hurts embeddings §1.

Method

The model has one shared graph encoder and a projection head used only in the loss §3.2. Each training step draws two random augmentations, encodes both views with the same encoder, and maximises an agreement objective that is averaged over nodes §3.2 alg. 1. Figure 1 shows this pipeline fig. 1.

Encoder: a two-layer graph convolutional network (GCN) with symmetric normalisation and self-loops §4.1. The paper writes it as

\mathrm{GC}_i(\bm X, \bm A) = \sigma\left(\hat{\bm D}^{-1/2}\hat{\bm A}\hat{\bm D}^{-1/2}\bm X\bm W_i\right), \quad f(\bm X, \bm A) = \mathrm{GC}_2(\mathrm{GC}_1(\bm X, \bm A), \bm A) \quad [\S 4.1]

where \hat{\bm A} = \bm A + \bm I §4.1.

Contrastive objective. For node i, the embedding in view 1, \bm u_i, is the anchor. Its positive is \bm v_i, the same node in view 2. All other nodes in both views are negatives §3.2. The critic is cosine similarity between projected embeddings, divided by a temperature \tau, with projection g a two-layer perceptron §3.2. The pairwise loss is eq. 1:

\ell(\bm u_i, \bm v_i) = \log \frac{e^{\theta(\bm u_i, \bm v_i)/\tau}}{e^{\theta(\bm u_i, \bm v_i)/\tau} + \sum_{k \neq i} e^{\theta(\bm u_i, \bm v_k)/\tau} + \sum_{k \neq i} e^{\theta(\bm u_i, \bm u_k)/\tau}} \quad [\text{eq. 1}]

The overall objective, which the paper maximises, is eq. 2:

\mathcal{J} = \frac{1}{2N}\sum_{i=1}^{N}\left[\ell(\bm u_i, \bm v_i) + \ell(\bm v_i, \bm u_i)\right] \quad [\text{eq. 2}]

Topology augmentation §3.3.1. Each edge (u,v) is kept with probability 1 - p^e_{uv} eq. 3:

P\{(u, v) \in \widetilde{\mathcal E}\} = 1 - p^e_{uv} \quad [\text{eq. 3}]

Removal probabilities come from edge centrality. Let \varphi_c be a node centrality. Undirected edges take the average of their endpoints' scores, and the paper says directed edges take the score of one endpoint (see Pitfalls). Scores are log-transformed, normalised over edges, scaled by p_e, and capped at p_\tau eq. 4 §3.3.1:

p^e_{uv} = \min\left(\frac{s^e_{\max} - s^e_{uv}}{s^e_{\max} - \mu^e_s}\cdot p_e,\ p_\tau\right), \quad s^e_{uv} = \log w^e_{uv} \quad [\text{eq. 4}]

Here s^e_{\max} and \mu^e_s are the maximum and mean of s^e_{uv} over edges §3.3.1. The three node centralities are degree (in-degree for directed graphs), eigenvector (leading eigenvector of the adjacency matrix) and PageRank (damping \alpha = 0.85) §3.3.1. PageRank is defined as eq. 5:

\bm\sigma = \alpha\bm A\bm D^{-1}\bm\sigma + \bm 1 \quad [\text{eq. 5}]

Feature augmentation §3.3.2. One Bernoulli mask \widetilde{\bm m} with one entry per feature dimension is drawn per view, and it is shared by every node in that view eq. 6:

\widetilde{\bm X} = [\bm x_1 \circ \widetilde{\bm m}; \dots; \bm x_N \circ \widetilde{\bm m}]^\top, \quad \widetilde m_i \sim \mathrm{Bern}(1 - p^f_i) \quad [\text{eq. 6}]

The weight of dimension i sums node centrality over the nodes that have it. For binary (one-hot) features eq. 7:

w^f_i = \sum_{u \in \mathcal V} x_{ui} \cdot \varphi_c(u) \quad [\text{eq. 7}]

For dense features, the paper uses absolute values eq. 8:

w^f_i = \sum_{u \in \mathcal V} |x_{ui}| \cdot \varphi_c(u) \quad [\text{eq. 8}]

Dimension removal probabilities use the same normalisation and cap as edges, with s^f_i = \log w^f_i and strength p_f eq. 9 §3.3.2:

p^f_i = \min\left(\frac{s^f_{\max} - s^f_i}{s^f_{\max} - \mu^f_s}\cdot p_f,\ p_\tau\right) \quad [\text{eq. 9}]

Each view draws its own p_e and p_f (called p_{e,1}, p_{f,1} for view 1 and p_{e,2}, p_{f,2} for view 2) so the two views differ §3.3.2. Centralities and weights are computed once from the input graph §3.3.2.

Training. Algorithm 1 runs per epoch: sample two augmentations; build two views; encode each; compute \mathcal J; take a gradient step alg. 1. The sketch below follows that structure and the equations above.

# Node centrality on the input graph, computed once  [§3.3.1]
# Measure: degree, eigenvector or PageRank (damping 0.85)  [eq. 5]
def node_centrality(graph, measure):
    # Degree: in-degree for directed graphs  [§3.3.1]
    if measure == "degree":
        return in_degrees(graph)
    # Leading eigenvector of the adjacency matrix  [§3.3.1]
    if measure == "eigenvector":
        return leading_eigenvector(graph)
    # PageRank  [§3.3.1] [eq. 5]
    return pagerank(graph, damping=0.85)

# Edge weight from endpoint centralities  [§3.3.1]
# Undirected: mean of the two endpoints; directed: see Pitfalls
def edge_centrality(edges, node_scores):
    return (node_scores[edges.src] + node_scores[edges.dst]) / 2

# Removal probability per item from its weight  [§3.3.1] [eq. 4]
# Same form for feature dimensions with s = log(weight)  [§3.3.2] [eq. 9]
def removal_probs(weights, strength, p_tau):
    # Log compresses weights that differ by orders of magnitude  [§3.3.1]
    scores = log(weights)
    top = max(scores)
    average = mean(scores)
    # Low-scoring items get high removal probabilities  [eq. 4]
    probs = (top - scores) / (top - average) * strength
    # Cap at p_tau so no item is almost always removed  [eq. 4]
    return minimum(probs, p_tau)

# Dimension weights from node scores  [§3.3.2] [eq. 7] [eq. 8]
def dimension_weights(features, node_scores, binary):
    # Binary features: count nodes that have the dimension  [eq. 7]
    if binary:
        present = (features != 0)
        return present.T @ node_scores
    # Dense features: weight by absolute value  [eq. 8]
    return abs(features).T @ node_scores

# Keep each edge with probability 1 - removal  [§3.3.1] [eq. 3]
def drop_edges(edges, removal):
    # One Bernoulli draw per edge  [§3.3.1]
    keep = bernoulli(1 - removal)
    return edges[keep]

# Zero whole feature dimensions, shared by all nodes  [§3.3.2] [eq. 6]
def mask_features(features, removal):
    # One Bernoulli draw per dimension  [§3.3.2] [eq. 6]
    keep = bernoulli(1 - removal)
    return features * keep

# Pairwise loss for every anchor in one view  [§3.2] [eq. 1]
# cross: cosine similarities to the other view, shape (N, N)
# same: cosine similarities within this view, shape (N, N)
def pair_loss(cross, same, tau):
    # Temperature scaling before the exponential  [§3.2]
    cross_exp = exp(cross / tau)
    same_exp = exp(same / tau)
    # The positive is the diagonal of the cross-view matrix  [§3.2]
    positive = diagonal(cross_exp)
    # Denominator: positive, inter-view and intra-view negatives  [eq. 1]
    # The self-similarity on the intra-view diagonal is excluded  [eq. 1]
    denominator = (
        sum(cross_exp, axis=1)
        + sum(same_exp, axis=1)
        - diagonal(same_exp)
    )
    return -log(positive / denominator)

# Symmetric loss, averaged over nodes  [§3.2] [eq. 2]
# The paper maximises J; this returns -J, which has the same optimum
def contrastive_loss(view_1, view_2, projection, tau):
    # Projection head g, applied before the cosine critic  [§3.2]
    proj_1 = projection(view_1)
    proj_2 = projection(view_2)
    cross_12 = cosine_matrix(proj_1, proj_2)
    same_1 = cosine_matrix(proj_1, proj_1)
    same_2 = cosine_matrix(proj_2, proj_2)
    loss_1 = mean(pair_loss(cross_12, same_1, tau))
    loss_2 = mean(pair_loss(transpose(cross_12), same_2, tau))
    return (loss_1 + loss_2) / 2

# One training epoch on the full graph  [alg. 1] [§3.2]
# edge_weights, feature_weights: centrality weights, computed once  [§3.3.2]
def train_epoch(features, edges, edge_weights, feature_weights,
                encoder, projection, optimizer, settings):
    # Two views, each with its own strengths  [§3.3.2]
    edges_1 = drop_edges(edges, removal_probs(
        edge_weights, settings.p_e_1, settings.p_tau))
    edges_2 = drop_edges(edges, removal_probs(
        edge_weights, settings.p_e_2, settings.p_tau))
    features_1 = mask_features(features, removal_probs(
        feature_weights, settings.p_f_1, settings.p_tau))
    features_2 = mask_features(features, removal_probs(
        feature_weights, settings.p_f_2, settings.p_tau))
    # Encode both views with the same encoder  [§3.2]
    view_1 = encoder(features_1, edges_1)
    view_2 = encoder(features_2, edges_2)
    loss = contrastive_loss(view_1, view_2, projection, settings.tau)
    # One optimiser step per epoch  [alg. 1] [§3.2]
    optimizer.step(loss)

Theory. Theorem 1 says the objective is a lower bound on the mutual information between the input and the two views' embeddings §3.4 app. B. Theorem 2 relates the objective to a triplet loss under an inner-product critic, identity projection and a first-order approximation §3.4 app. B. The paper says it does not derive how specific augmentations relate to the bound, and leaves that to future work §3.4.

Variants. GCA-DE, GCA-EV and GCA-PR use degree, eigenvector and PageRank centrality respectively §3.3.1 tab. 3.

Hyperparameters

Per-dataset settings, complete from the paper tab. 5. p_{e,1}, p_{e,2}, p_{f,1} and p_{f,2} are the augmentation strengths of the two views. p_\tau is the cap. \tau is the temperature. Epochs are full-batch update steps.

tab. 5. The official code's values differ for Wiki-CS and Amazon-Computers (see Paper vs code).

Settings that apply to all datasets:

Results

Author-reported. Setup: five transductive node classification datasets §4.1 tab. 2. Datasets, from the paper tab. 2:

Protocol: embeddings from the encoder on the full graph train a logistic regression. Accuracy is the metric. The paper reports averages over twenty runs on different splits, with Wiki-CS on its public splits and the other four on random 10/10/80 splits §4.1. How the official code does this is in Paper vs code. Baselines are DeepWalk, node2vec, GAE, VGAE, DGI, GMI and MVGRL (unsupervised), and GCN and GAT (supervised) §4.1. Accuracy in percent with the standard deviation as reported tab. 3:

Headline reading tab. 3: the best unsupervised row is a GCA variant on all five datasets. On Wiki-CS, Amazon-Computers, Coauthor-CS and Coauthor-Physics, a GCA variant is above both GCN and GAT. On Amazon-Photo, GAT (92.56) is above all three GCA variants (92.24 to 92.53). The paper's wording that GCA "surpasses some supervised counterparts" §1 rests on margins of a few tenths of a point or less.

Ablation tab. 4, using degree centrality in all variants. GCA--T--A is uniform topology and attribute augmentation, which the paper says reduces to GRACE §4.3:

The paper reports that each adaptive part helps and that the two together help most on every dataset §4.3. Amazon-Computers gains 1.5 points from the full adaptive scheme over the base model §4.3.

Sensitivity fig. 3: on Amazon-Photo, the paper describes accuracy as stable for strengths that are not too large, and says that above about 0.5 the original graph is heavily undermined. At p_e = 0.9 almost every edge is removed §4.4.

Figures are described, not reproduced: Figure 1 is the pipeline fig. 1. Figure 2 shows the three edge-centrality schemes on the Karate Club graph. All three give more weight to edges that link the coaches inside each group and less to links between peripheral nodes across groups fig. 2 §3.3.1.

Paper vs code

Official code: CRIPAC-DIG/GCA, linked in the paper's introduction §1. It is cloned at commit 0f17497f3ae89718ed610c671a5b9f56033d4d51 under the MIT licence (copyright CRIPAC, 2021) LICENSE. Every citation below is relative to that repository. Each difference is listed with both sides.

• Undirected edge weights. The paper averages the two endpoint centralities §3.3.1. The code uses only one endpoint's score: the target node, for degree functional.py, eigenvector functional.py and PageRank (default aggr='sink') functional.py. The two directions of one undirected edge therefore get different probabilities.
• PageRank update. The paper's fixed point is \bm\sigma = \alpha\bm A\bm D^{-1}\bm\sigma + \bm 1 eq. 5. The code iterates x = (1 - damp) * x + damp * agg_msg from x = 1 with no constant term utils.py. The iteration count is 200 for edge weights train.py and 10 for feature weights train.py utils.py. The paper gives no count.
• Eigenvector weights. For edges, negative scores are set to 0 and 1e-8 is added before the log functional.py. For features, the raw scores are passed on without that step train.py. The paper says nothing about either.
• Edge sampling. The code keeps or drops each row of the edge index independently functional.py, so the two directions of an undirected edge are sampled separately. The paper's edge-set notation does not say whether undirected pairs are sampled jointly §3.3.1.
• Projection head. The code is Linear, then ELU, then Linear model.py. The paper says only "two-layer perceptron" §3.2. The width is a per-dataset setting in the code (256, 128 or 64) wikics.json amazon_computers.json amazon_photo.json. The default in the code is 32 train.py.
• Encoder width. The first GCN layer outputs twice the embedding width, and the last outputs the embedding width model.py model.py. The paper's table calls the value "hidden dimension" tab. 5 without this detail.
• RReLU. The code calls F.rrelu without training=True utils.py. In PyTorch that default gives the fixed mean slope during training as well, not a random slope. The paper does not specify this.
• Initialisation. The paper says all parameters use Glorot initialisation app. A. The GCN layers use PyTorch Geometric defaults. The projection layers are plain nn.Linear with no explicit initialiser model.py.
• Linear classifier. The paper says ℓ2-regularised logistic regression §4.1. The code uses Adam with weight_decay=0.0 eval.py. It also selects the classifier epoch with the best validation accuracy, checked every 20 epochs, and reports the test accuracy at that epoch eval.py. The paper does not state this.
• Number of runs. The paper says "twenty runs for different data splits" §4.1. The code averages 20 public splits for Wiki-CS only train.py. For the other four datasets it makes one random split in one run train.py. That split is redrawn on every call eval.py. The saved split is used only for split == 'preloaded' eval.py, which the training script never passes. So the saved split does not fix the variance.
• Feature preprocessing. All datasets are row-normalised with NormalizeFeatures dataset.py. The paper says nothing about it. Bag-of-words values therefore are not 0 or 1 in the code's input, although the weight in eq. 7 is computed on the nonzero pattern functional.py.
• Hyperparameters, Wiki-CS. The table gives p_{e,2} = 0.4 and \tau = 0.6 tab. 5. The code has p_{e,2} = 0.3 wikics.json and \tau = 0.4 wikics.json.
• Hyperparameters, Amazon-Computers. The table gives PReLU, p_{e,1} = p_{e,2} = 0.5, p_{f,2} = 0.1, \tau = 0.1 and 1,500 epochs tab. 5. The code has RReLU amazon_computers.json, p_{e,1} = 0.6 amazon_computers.json, p_{e,2} = 0.3 amazon_computers.json, p_{f,2} = 0.3 amazon_computers.json, \tau = 0.2 amazon_computers.json and 2,000 epochs amazon_computers.json.
• Matches. Amazon-Photo, Coauthor-CS and Coauthor-Physics match Table 5 on every value listed there amazon_photo.json coauthor_cs.json coauthor_phy.json. The loss matches eq. 1 and eq. 2: the denominator includes the positive, inter-view negatives and intra-view negatives, and the intra-view self-term is excluded model.py model.py. Weight decay is 1e-5 train.py. The cap is 0.7 train.py.
• Drop scheme. The code selects GCA-DE, GCA-PR or GCA-EV with --drop_scheme degree|pr|evc, defaulting to degree train.py. The parameter files do not set the scheme, so the variants share the same per-dataset settings.

Pitfalls

• Edge weight rule. The paper's undirected rule §3.3.1 differs from the code (see Paper vs code). For directed graphs the text says the "tail node" but writes \varphi_c(v) §3.3.1. The code uses the target node. The datasets here are undirected, so this does not arise in the reported results.
• The cap lowers the average removal rate. The normalised weights are scaled by p and then capped at p_\tau = 0.7 eq. 4 functional.py. Capping only lowers values, so when the cap binds the realised fraction removed is below the nominal p_e. The paper's statement that p_e controls the overall probability is approximate.
• Asymmetric views. Because each direction is sampled separately, the augmented graph can be asymmetric. Implementations that assume a symmetric adjacency will not match the code.
• log(0). Eq. 4 and eq. 7 to 9 take the log of a centrality weight §3.3. A feature that is zero for every node, or a zero eigenvector score, gives log 0. The code adds 1e-8 only on the edge side. The paper gives no guard.
• PageRank settings. The update differs from eq. 5 and the iteration count is not stated (see Paper vs code). Expect different edge and feature probabilities from a reimplementation that follows the paper's formula.
• Quadratic memory. The full-batch loss builds N×N similarity matrices for both views model.py. The code switches to a batched loss only for Coauthor-Physics (batch 1,024) train.py. At 34,493 nodes tab. 2, one float32 N×N matrix is about 4.8 GB by our arithmetic. The paper calls the objective "cheap" compared with triplet losses §3.4 and does not discuss memory.
• Text vs table on strengths. The appendix says every strength is chosen between 0.0 and 0.4, and that p_{e,1} \neq p_{e,2} app. A. Table 5 lists 0.5 for Amazon-Computers and Amazon-Photo p_{e,2} and equal values for Amazon-Computers p_{e,1} and p_{e,2} tab. 5. Use the table values, and note the code's Amazon-Computers values differ again.
• Sensitivity sweep. The sweep sets p_{e,1} = p_{e,2} and p_{f,1} = p_{f,2} on Amazon-Photo only §4.4 fig. 3. It does not test the per-view asymmetry that the method uses.
• Text inconsistencies. The experiments say "six widely-used datasets" and then list five §4.1. The method says "at each iteration" and "at each training epoch" §3.2. In the code these coincide, because training is full-batch.
• Ablation gain. The text says Amazon-Computers gains 1.5 points from the full adaptive scheme §4.3. Table 4 gives 87.85 for GCA and 86.25 for GCA--T--A, a difference of 1.60 tab. 4.
• Reported spread. The paper does not say how the ± values are computed. GCA-DE on Wiki-CS is listed as 78.30 ± 0.00 tab. 3, and Wiki-CS uses 20 splits, so check the run logs before quoting that row.
• Small margins. Many of the comparisons with supervised models differ by under half a point, and no significance test is reported tab. 3. For example, GCA-DE is 93.10 against GCN's 93.03 on Coauthor-CS tab. 3.
• Theory gap. Theorem 1 is a lower bound on mutual information §3.4. The paper does not show that any specific augmentation improves the bound, and leaves that open §3.4. Theorem 2 relies on a first-order Taylor approximation under an identity projection app. B.
• Classifier protocol. Reported accuracies depend on the validation-based epoch choice in the code and on the classifier's training budget (800 or 3,000 epochs, train.py train.py). Compare only under the same protocol.

Lineage

• Builds on: GRACE, the authors' preliminary work §2. The paper says GCA--T--A reduces to GRACE §4.3. Its arXiv id is not in this wiki.
• Compared against: DeepWalk, node2vec, GAE, VGAE, DGI, GMI and MVGRL, and the supervised GCN and GAT §4.1 tab. 3. The paper's own comparison table covers DGI, GMI and MVGRL §2 tab. 1.