Graph Barlow Twins: A self-supervised representation learning framework for graphs

In plain words

A graph is a set of items joined by links, such as papers that cite each other. Learning useful item descriptions without labels is hard. The popular approaches compare pairs of items and need "negative" pairs that should be far apart, and for graphs it is unclear what a good negative is. A newer method avoids negatives but uses two networks that are deliberately kept unequal.

This paper keeps one network and makes two damaged copies of the same graph. One copy loses some links, and the other loses some of the input-feature columns. Both copies go through the same network. The training goal is that the two sets of item descriptions agree with each other feature by feature, while different features stay unrelated. This is measured with a table of pairwise correlations, and the target is the identity matrix. The result is a set of item descriptions that can be fed to a simple classifier. The authors report accuracy close to the strongest earlier method on standard benchmarks, with much less training time.

TL;DR

• Self-supervised graph encoder with a symmetric, negative-free loss: the Barlow Twins cross-correlation objective applied to two augmented views of one graph §3.
• Augmentations are random edge dropping and random feature-column masking, with one shared probability pair for both views §3 §5.1.
• Author-reported: on the five small datasets G-BT (at most 1 000 epochs) is within 2.7 points of BGRL (10 000 epochs) and within 0.5 points on three of them tab. 2. On ogb-arxiv it is 1.5 points below BGRL on test tab. 3.
• Total training time is much lower than BGRL's on the transductive datasets; the "around 30 times" claim in the abstract is a ratio of total times, not per-epoch times tab. 6.
• Code is MIT-licensed and matches the paper on the loss, augmentation, encoders and training schedule; the discrepancies that matter are listed under Paper vs code.

Problem

Contrastive graph methods such as DGI, GRACE and GraphCL need negative examples, and choosing them is hard for graphs §1 §2. BGRL avoids negatives but uses an asymmetric design: an online and a target encoder, a predictor network, stop-gradient and a moving-average target §1 §2. The authors report that BGRL needs 10 000 epochs to converge in their setup §4.3.1 §4.3.2. The paper targets a symmetric, negative-free objective that converges faster and has fewer hyperparameters §1 [abstract].

Method

Graph and views. The graph is a node-feature matrix X (|V| × k) and a binary adjacency matrix A §3. Two views are produced by the same random augmentation, applied twice, each with its own random masks §3 eq. 1:

• Feature masking keeps or zeroes each of the k feature columns, and the decision is shared by all nodes of the view. Keep probability is 1 − p_X §3 augment.py.
• Edge dropping keeps each edge with probability 1 − p_A, independently per edge §3 augment.py.

Encoder. One encoder f_θ maps each view to |V| × d embeddings, and both views pass through the same weights §3. The default encoder is a two-layer GCN: GCN(k, 2d), batch normalisation with momentum 0.01, PReLU, then GCN(2d, d), with no normalisation or activation on the output §4.3.1 app. C encoders.py. The paper uses a three-layer GCN for ogb-arxiv §4.3.2 app. C and a three-layer GAT with skip connections for PPI §4.4.1 app. C encoders.py.

Loss. Each embedding dimension of Z_1 and Z_2 is normalised to zero mean and unit standard deviation over the nodes §3. The empirical cross-correlation matrix C is d × d §3 eq. 2. The loss sums the squared deviation of the diagonal from 1 (invariance term) and λ times the squared off-diagonal entries (redundancy term), with λ = 1/d §3 eq. 3. The choice λ = 1/d follows a previous note the authors cite, and §5.3 reports that it is the best of the values they tried on WikiCS §3 §5.3 fig. 4. Gradients flow through both branches. The method uses no predictor, no stop-gradient and no momentum encoder §3.

Optimisation. AdamW with weight decay 1e-5. The learning rate follows cosine annealing after a linear warmup §4.2 app. B model.py. Training runs for a fixed number of epochs, and the encoder is evaluated at a fixed interval app. B.

Algorithm. The paper's training loop is given as Algorithm 1 alg. 1:

# Graph Barlow Twins training  [alg. 1]
# Input:  graph (X, A); keep-probabilities p_x, p_a;
#         encoder f_theta; number of epochs N
# Output: trained encoder f_theta

def train(X, A, p_x, p_a, encoder, num_epochs):
    # Repeat for N epochs  [alg. 1]
    for epoch in range(num_epochs):
        # Two augmented views, each with its own random masks  [§3] [eq. 1]
        X_1, A_1 = augment(X, A, p_x, p_a)
        X_2, A_2 = augment(X, A, p_x, p_a)

        # Node embeddings from the same encoder  [§3]
        Z_1 = encoder(X_1, A_1)                       # (|V|, d)
        Z_2 = encoder(X_2, A_2)                       # (|V|, d)

        # Barlow Twins loss on the cross-correlation matrix  [§3] [eq. 2-3]
        loss = barlow_twins_loss(Z_1, Z_2)

        # Update encoder weights; the code uses AdamW  [§4.2] [app. B]
        optimizer_step(encoder, loss)

# Two-view augmentation  [§3] [eq. 1]
# Feature masking: one decision per column, shared by all nodes  [§3]

def augment(X, A, p_x, p_a):
    # Keep each feature column with probability 1 - p_x  [§3]
    feature_keep = bernoulli(1 - p_x, size=(1, k))     # (1, k)
    X_view = X * feature_keep
    # Keep each edge with probability 1 - p_a, one decision per edge  [§3]
    edge_keep = bernoulli(1 - p_a, size=num_edges)     # (E,)
    A_view = A[:, edge_keep]
    return X_view, A_view

# Barlow Twins loss on two embedding matrices  [§3] [eq. 2-3]
# Z_1, Z_2: (N, d) embeddings of the nodes of the two views

def barlow_twins_loss(Z_1, Z_2):
    # Standardise each dimension over the nodes  [§3] [eq. 2]
    Z_1_norm = (Z_1 - mean(Z_1, axis=0)) / std(Z_1, axis=0)    # (N, d)
    Z_2_norm = (Z_2 - mean(Z_2, axis=0)) / std(Z_2, axis=0)    # (N, d)
    # Cross-correlation matrix between the two views  [§3] [eq. 2]
    C = transpose(Z_1_norm) @ Z_2_norm / N                     # (d, d)
    # Invariance: diagonal entries should be 1  [§3] [eq. 3]
    invariance = sum((1 - diagonal(C)) ** 2)
    # Redundancy reduction: off-diagonal entries should be 0, weight 1/d  [§3]
    redundancy = sum(offdiag(C) ** 2) / d
    return invariance + redundancy

Evaluation. The encoder is frozen and embeddings are computed on the unaugmented graph §4.2. A classifier is then trained on the training nodes. For the five small datasets the classifier is scikit-learn's L2-regularised logistic regression, with a grid search over its regularisation strength §4.2 app. B. For ogb-arxiv, ogb-products and PPI it is a PyTorch version of the same model, because the scikit-learn version was too slow to converge §4.2 app. B. Reported numbers are mean and standard deviation over 20 model initialisations or data splits, except ogb-products (5) §4.2. Details of the classifier are in the Paper vs code section.

Augmentation search. The two augmentation probabilities are chosen by a grid over {0, 0.1, …, 0.5} for each, 500 epochs with 50 warmup epochs §4.2 app. A. The paper uses one shared pair for both views. A search over separate pairs for each view would be 630 combinations after symmetry reduction, and §5.1 reports no substantial test-accuracy difference between the two designs §5.1 fig. 2.

Hyperparameters

Per-dataset values. Columns p_A (edge drop) and p_X (feature mask) match the code's p_e and p_x tab. 8 train.yaml. Epoch, warmup and evaluation settings are from tab. 9; the feature count k and embedding size d are from tab. 10. "Best at" is the epoch the paper reports as the best model tab. 9.

Other settings:

Results

Author-reported. The paper re-uses the BGRL and GRACE numbers from the BGRL paper, and marks its own reruns with an asterisk ("*") tab. 2 tab. 3 tab. 4 tab. 5.

Transductive, small datasets, accuracy mean ± std over 20 splits tab. 2:

The paper says G-BT converges at about 500 to 900 epochs on these datasets and outperforms BGRL at 1 000 epochs §4.3.1 tab. 9.

Transductive, ogb-arxiv, accuracy over 20 re-initialisations tab. 3:

The paper says G-BT is about 1.5 points below BGRL on test §4.3.2.

Inductive, PPI, Micro-F1 over 20 initialisations tab. 4:

Inductive, ogb-products, accuracy over 5 initialisations tab. 5:

The paper states that the ogb-products table includes leaderboard numbers from semi-supervised methods, so those are not directly comparable §4.4.2.

Training time, seconds per epoch and total tab. 6. Single-epoch time over 10 epochs; "Speedup" is the ratio of total training times §4.5:

Batched training, accuracy tab. 7. The full-batch row is copied from tab. 2. Neighbour sampling is used for the batched rows §4.6:

The paper reports an expected drop in batched mode and does not tune the batched runs §4.6.

Ablations (WikiCS, one-layer and other variants trained 1000 epochs with 100 warmup; all other settings from the best G-BT model) §5.2 fig. 3:

• Both augmentations: 76.65 (the main result).
• Feature masking only: about 75.9.
• Edge dropping only: about 72.
• No augmentation: about 67; the authors expect collapse to a constant embedding with enough training.
• MLP encoder instead of GCN: about 68.
• One-layer GCN: 1.1 points below the two-layer baseline.
• Three-layer GCN: lower than two-layer; the authors suggest over-smoothing but do not test it §5.2.

Loss weight λ (WikiCS, two-layer GCN) §5.3 fig. 4: best at λ = 1/256 (the default for d = 256); λ = 1 settles near 55%; values below 0.01 give about 75%.

Projector (WikiCS). Adding a projector of widths 128 to 16 348 (as printed) made no significant difference compared with having none, by a Friedman test §5.4 fig. 5.

Paper vs code

Code: https://github.com/pbielak/graph-barlow-twins at commit ec62580aa89bf3f0d20c92e7549031deedc105ab. The owner is the first author, and the licence is MIT (Copyright 2021 Piotr Bielak) LICENSE. The repository was checked against the paper's loss, augmentation, encoders, optimiser and training table. These match: λ = 1/d loss.py; Bernoulli feature masking and Bernoulli edge dropping augment.py; 2-layer, 3-layer and GAT encoders with the stated widths and skips encoders.py; AdamW, weight decay 1e-5 and warmup/cosine schedule model.py; and the epoch, warmup and learning-rate values for the full-batch datasets train.yaml. The discrepancies below were found:

1. Normalisation and the diagonal. Paper eq. 2 divides the un-centred product by the root sums of squares §3 eq. 2. The text says the embeddings are first normalised to zero mean and unit standard deviation §3. The code centres each column, divides by the column standard deviation with PyTorch's default unbiased estimator, and divides the product by N loss.py. Two identical, perfectly correlated views therefore give a diagonal of (N−1)/N, not 1 loss.py. This is negligible for full graphs but larger for small batches: 255/256 for a 256-node batch.
2. Embedding normalisation before the probe. The code L2-normalises embeddings before every classifier in the accuracy path tasks.py. The paper does not say this.
3. Probe grids. The paper lists regularisation grids that end at 2^10 §4.2. The code's arange excludes the end value, so the largest sklearn C is 2^9 tasks.py and the largest PyTorch weight decay is 2^8 tasks.py. The paper's probe selection is not described in more detail than "grid search".
4. Sklearn probe selection. The sklearn path chooses C by 5-fold cross-validation on the training nodes, wrapped in a one-vs-rest classifier, with max_iter=100 tasks.py. The PyTorch path chooses by validation accuracy tasks.py. The paper says only "grid search" §4.2.
5. PyTorch probe settings. The PyTorch probe uses AdamW with learning rate 0.01 and 1000 full-batch steps tasks.py tasks.py. The paper gives the step count but not the learning rate §4.2.
6. PPI scaling. For the multi-label F1 path, a StandardScaler is fitted separately on train, validation and test embeddings tasks.py. The paper does not mention scaling. Fitting on each split separately means the test features are scaled with their own mean and variance.
7. Augmentation selection uses the test split. The search script ranks augmentation pairs by test accuracy (or F1 for PPI) and writes the best pair to best.json hps.py. The paper says only that the best-performing pair was used §4.2 app. A. This is test-set selection of a hyperparameter (see Pitfalls).
8. Search learning rate for ogb-arxiv. The augmentation search uses base learning rate 1e-2 for ogb-arxiv hps.yaml, while the final training uses 1e-3 train.yaml and the paper's Table 9 lists 1e-3 tab. 9. The pair chosen in search is therefore not tuned at the final learning rate.
9. Batched retrain count. The paper says batched runs are retrained five times per batch size §4.6. The batched config uses 20 splits for WikiCS, Amazon-Computers and Amazon-Photos train.yaml and 5 for the others train.yaml.
10. Augmentation values are not in the repository. The training script reads data/full_batch/hps/<dataset>/best.json train.py and the data/ directory in the repository is empty. The values in tab. 8 therefore cannot be checked against code output. Only the search grid and the search settings can be read from the configs.
11. Stale configs. experiments/configs/train_transductive.yaml and train_inductive_ppi.yaml describe other settings, for example 4000 epochs and warmup 400 for WikiCS with p_x 0.1 and p_e 0.5 train_transductive.yaml. They also contain an HSIC variant that the paper says gave no gain §3. The scripts that read them open experiments_ssl/configs/, which is not in the repository, so these files appear to be unused leftovers. The full-batch configs are the ones that match the paper.
12. Scheduler granularity. Full-batch training steps the scheduler once per epoch, with one optimiser step per epoch model.py. PPI takes one optimiser step per graph (20 per epoch), and the scheduler steps once per epoch multiple_graph_model.py. The paper does not state the step granularity.
13. Algorithm 1 versus AdamW. Algorithm 1 writes a plain gradient update alg. 1. The experiments use AdamW with cosine warmup §4.2 app. B.

Two TeX sources are in the arXiv bundle: main.tex (used for this card, with the compiled main.bbl) and kbs_formatted_final.tex, the journal-formatted version. The hyperparameter tables in both (tab. 8, tab. 9, tab. 10 and the augmentation and training tables for ogb-products) agree. The results tables were not compared line by line.

Pitfalls

• Test-set selection of augmentation probabilities. The search picks the pair with the best test accuracy hps.py. Reported test numbers therefore carry a selection bias. Reproduce with a validation-based search if the numbers matter.
• Choice of reporting epoch. Training logs validation and test metrics every 100 epochs and writes mean and standard deviation for every logged epoch train.py. The paper reports the "best model" epoch in tab. 9. How that epoch was picked, and on which split, is not stated.
• Speedup is a total-time ratio. The abstract's "around 30 times faster" compares total training times, and BGRL is trained for 10 000 epochs while G-BT is trained for 500 to 900 tab. 6. Per epoch, the table's times give about 2× to 2.6× on WikiCS, Amazon-Computers and Amazon-Photos tab. 6 (ratios computed here from the per-epoch column).
• Dropped edges may be directed. drop_edges masks each column of edge_index independently augment.py. If a dataset stores both directions of an undirected edge, the two directions can be dropped separately, which makes the augmented graph directed. The paper does not say how this should be handled.
• Feature mask is global. One Bernoulli decision per feature column is shared by every node of a view augment.py. A masked column is zero everywhere in that view, so the encoder sees a different input for each view on each epoch.
• Division by near-zero standard deviation. The normalisation adds only EPS = 1e-15 to each column's standard deviation loss.py loss.py. If an embedding dimension collapses to a constant, its normalised values become very large and dominate the loss. The paper does not discuss collapse guards.
• Batched mode costs accuracy on some datasets. Against the full-batch row, WikiCS, Amazon-Photos and Coauthor-CS lose about 1 to 1.5 points, while Amazon-Computers and Coauthor-Physics change by 0.3 or less tab. 7. Batched runs reuse the full-batch augmentation values and epoch counts §4.6, so they are not tuned for batch size.
• Ablations use one dataset. All ablations and the λ and projector studies use WikiCS only §5.2 §5.3 §5.4.
• Probe cost. The sklearn path runs a 20-value grid with 5-fold cross-validation per split tasks.py. That cost is not in the paper's timing table, which measures training epochs only tab. 6.
• Learning-rate sensitivity across datasets. The base learning rate spans 1e-5 to 5e-3 across datasets tab. 9. The paper gives no rule for choosing it.

Lineage

• Builds on: Barlow Twins, a redundancy-reduction objective for images that this paper adapts to graphs §1 §3. No card in this wiki.
• Builds on: BGRL, the negative-free graph baseline that this paper compares against §1 §2. No card in this wiki.
• Compared against: DGI, GRACE, GCA, GraphCL, MVGRL, DeepWalk and node2vec, using numbers taken from the BGRL paper tab. 2 tab. 3. None of these has a card in this wiki.
• Superseded by: not known from content in this wiki.