GraphMAE2: A Decoding-Enhanced Masked Self-Supervised Graph Learner

In plain words

Many graph datasets have node features but few labels. One way to learn from the graph without labels is to hide part of each node's features and train a network to fill them back in from the rest of the graph. GraphMAE does this. GraphMAE2 keeps that idea and adds two safeguards, because the hidden features can be a noisy target when they were produced by a weak feature extractor.

The first safeguard is to hide parts of the learned representation again, several times with different random choices, and require each version to still reproduce the original features. The second is to ask the model to predict a target produced by a slowly updated twin network that sees the unhidden graph. Together these stop the model from memorising the noise in the input features.

For graphs with millions of nodes, the training data is cut into small, densely connected neighbourhoods, each built around one node. This keeps the training pieces manageable and matches how the model is used later, when it embeds one node from its own neighbourhood.

TL;DR

GraphMAE2 is a masked self-supervised pre-training method for node embeddings on graphs with a GAT encoder. It masks half of the node features and reconstructs them through a decoder that sees several random re-masked views, plus a latent prediction loss against an EMA target encoder. Author-reported linear-probe results are highest on the four large datasets; on the small ones they are highest on Cora and PubMed and within 0.1 point of the best on CiteSeer. Relevant if you need label-free node embeddings on large graphs; the paper and code disagree on several details (see Paper vs code).

Problem

Labels are scarce on real graphs. Generative masked feature reconstruction (GraphMAE) is strong but depends on how discriminative the input node features are §1 §2.1. In a pilot, reducing the features to 50 dimensions with PCA lowered GraphMAE's accuracy more than the supervised GAT's on Cora (84.2 to 82.6 versus 83.0 to 82.3) and PubMed (81.1 to 78.9 versus 78.0 to 77.0) tab. 1. The authors argue that graph features are often produced by weak extractors, so reconstructing them injects noise §2.1. Scaling to very large graphs is a second problem: GNNs over-smooth and over-squash when stacked deeper, and standard subgraph samplers give sparse subgraphs that suit masked feature prediction poorly §2.2, "Extending to large-scale graph".

Method

Notation: graph G = (V, A, X), N nodes, input features X in R^(N x d_in). The encoder f_E is a GAT, the decoder f_D is a single-layer GAT, and the code H is in R^(N x d) §2.1 eq. 1.

Masking (inherited from GraphMAE). Pick a uniform random subset of nodes, without replacement, at a 50% mask ratio, and replace their features with a learnable input mask token. Encode the corrupted graph into H. The reconstruction objective is a scaled cosine error on the masked nodes only §2.1 eq. 2.

Multi-view random re-mask decoding (new). Re-mask a fresh random subset of the code rows at rate 0.5 with a second, shared learnable decoder mask token. Decode with f_D and score the reconstruction of the input features on the originally masked nodes. Repeat K times, with K=3 in all large-graph configurations, and sum the K losses §2.2 "Multi-view random re-mask decoding" eq. 2 app. A.4, tab. 10. The paper contrasts this with GraphMAE's fixed re-mask, which re-masks only the originally masked nodes §2.2.

Latent representation prediction (new). A target generator has the same architecture as the encoder and projector but its own weights. It takes the unmasked graph and produces latent targets. The online branch runs a projector g on H and predicts those targets with a scaled cosine error §2.2 "Latent representation prediction" eq. 3 eq. 4. The target weights are an exponential moving average of the online weights eq. 5. The authors contrast this with distillation and contrastive methods, which pull two augmented views together; this method uses no augmentation §2.2.

Total loss. L = L_input + lambda * L_latent, with lambda set per dataset §2.2 eq. 6 app. A.4, tab. 10.

Large graphs. For each node, run approximate PPR (PPR-Nibble) and keep the top-k nodes as a local cluster. Train on all nodes inside each cluster. At inference, embed a node using the cluster induced by that node §2.2 "Extending to large-scale graph". The paper states that time and space are linear in N §2.2 "Training and inference". The value of k is not stated; the code uses 256 localclustering.py.

Downstream use. The decoder and target generator are discarded. Only the encoder is used, either frozen for embeddings or fine-tuned §2.2 "Training and inference".

Small-scale datasets. Cora, CiteSeer and PubMed use GAT for both encoder and decoder and an MLP projector, with the public splits §3.2.

Equation numbers in this card are counted from the LaTeX source in order of appearance; the paper contains no algorithm environment.

# Scaled cosine error  [eq. 2] [eq. 4]
# Normalise each row, then average (1 - cosine)^gamma over rows.
def scaled_cosine_error(predicted, target, gamma):
    # predicted, target: (num_rows, d); gamma >= 1 [eq. 2]
    cosine = sum(normalise(predicted) * normalise(target), axis=-1)
    return mean((1 - cosine) ** gamma)

# GraphMAE2 pre-training step  [§2.2] [eq. 1]-[eq. 6]
# Input:  graph (A), features X (N, d_in)
# Output: scalar loss

def training_step(graph, features):
    # 1. Mask 50% of nodes; replace their features with the input
    #    mask token  [§2.1]
    masked_nodes = random_subset(num_nodes, rate=0.5)
    corrupted = features.copy()
    corrupted[masked_nodes] = input_mask_token      # (d_in,)

    # 2. Encode the corrupted graph into the code H  [eq. 1]
    code = encoder(graph, corrupted)                # (N, d)

    # 3. Input reconstruction over K re-masked views  [§2.2] [eq. 2]
    input_loss = 0
    for view in range(K):                           # K = 3
        remasked_nodes = random_subset(num_nodes, rate=0.5)
        view_code = code.copy()
        view_code[remasked_nodes] = decoder_mask_token   # (d,)
        reconstructed = decoder(graph, view_code)   # (N, d_in)
        input_loss += scaled_cosine_error(
            reconstructed[masked_nodes],
            features[masked_nodes],
            gamma,                                  # not stated [eq. 2]
        )

    # 4. Latent prediction against the EMA target generator, which sees
    #    the unmasked graph and gets no gradient  [§2.2] [eq. 3] [eq. 4]
    predicted = projector(code)                     # (N, d_proj)
    with no_gradient():
        target = projector_ema(encoder_ema(graph, features))
    latent_loss = scaled_cosine_error(predicted, target, gamma)  # all N

    # 5. Combine the two losses with the mixing weight lambda  [eq. 6]
    return input_loss + lam * latent_loss

# Target update, run after each optimiser step  [eq. 5]
# Applies to the encoder and projector weights only.
def ema_update(online_weights, target_weights, tau):
    # tau is the EMA decay: 0.996 in large-graph configs  [code: configs]
    target_weights = tau * target_weights + (1 - tau) * online_weights

Hyperparameters

Large-graph pre-training (paper values, app. A.4 tab. 10; values in the code's configs agree except where noted in Paper vs code):

Small-scale datasets (code only; the paper's §3.2 gives no pre-training values). Source: cora.yaml, pubmed.yaml, citeseer.yaml:

Results

Author-reported. All numbers below are transcribed from the paper; they are not reproduced by us. Hardware and software for the paper's runs: a Linux machine with 8 NVIDIA A100 80GB GPUs, Python 3.9, PyTorch 1.12.0, OGB 1.3.3 and CUDA 11.3 app. A.4.

Datasets tab. 2:

Linear probing on large datasets, mini-batch training, accuracy % tab. 3 (§3.1; 10 trials, mean ± std):

The table note says GGD's code is not released and that its Arxiv and MAG-Scholar-F results are not in its paper tab. 3 note.

Fine-tuning with 1% and 5% labels, accuracy % tab. 4 (§3.1; 10 trials, mean ± std):

Full-label fine-tuning, accuracy % tab. 8 (appendix A.2; single values, no std reported):

Full-graph training on ogbn-Arxiv, linear probing, accuracy % tab. 9 (appendix A.3; mean ± std):

Small datasets, linear probing, accuracy % tab. 5 (§3.2; 20 random initialisations; mean ± std):

Ablation on Products, MAG and Papers100M, linear probing, accuracy % tab. 6 (§3.3; mean ± std):

Sampling strategy, linear probing, accuracy % tab. 7 (§3.3):

Pilot on feature compression, accuracy % tab. 1 (§2.1): GraphMAE on Cora 84.2 with raw features and 82.6 with 50-d PCA; on PubMed 81.1 and 78.9. Supervised GAT on Cora 83.0 and 82.3; on PubMed 78.0 and 77.0. GraphMAE2 on Cora 84.5 and 83.5; on PubMed 81.4 and 80.1.

Paper vs code

Official code: https://github.com/THUDM/GraphMAE2 (linked in the abstract). The paper itself links this repository, so the owner (THUDM) is accepted as the authors' own. Cloned at commit bc473e8b750d6e90b7cba1c6bb20d858f617d90c with --depth 1; the repository has an MIT LICENSE (copyright THUDM, 2023). Other searches were not needed. Differences found:

• Latent loss is not on all masked nodes in the large-graph path. The paper sums over all N nodes in eq. 4 while the text says it matches masked nodes §2.2 eq. 4. Code, mini-batch: the loss is taken only on the seed (root) node of each local cluster, whether masked or not edcoder.py lc_sampler.py. Code, full-batch (small datasets): the loss is taken on all nodes main_full_batch.py main_full_batch.py.
• Predictor on the online branch. The paper's latent prediction is Zbar = g(H) eq. 3. Code applies a projector (Linear d→256, PReLU, Linear 256→d) followed by a predictor (PReLU, Linear d→d) on the online branch; the target branch uses only the projector edcoder.py edcoder.py. The paper does not mention a predictor.
• Linear map before re-masking. Code applies an extra bias-free linear layer (d→d) to H before the re-mask step; it is initialised with gain 1.414 edcoder.py edcoder.py edcoder.py. The paper feeds H directly to the re-mask.
• Edge dropping is not in the paper. For the large datasets the code uses drop_edge_rate 0.5 and draws two independent random edge drops, one for the online graph and one for the target graph lc_sampler.py lc_sampler.py edcoder.py ogbn-arxiv.yaml. Cora, CiteSeer and PubMed use 0.0.
• Latent-loss exponent. Eq. 4 has an exponent gamma. Code fixes it to 1 edcoder.py edcoder.py. The input-loss exponent comes from the config key alpha_l: Arxiv 6, Products 3, MAG 2, Papers100M 2 ogbn-arxiv.yaml ogbn-products.yaml mag-scholar-f.yaml ogbn-papers100M.yaml. The paper gives no value.
• EMA starts late on Arxiv. The target is updated only from epoch delayed_ema_epoch: 40 for Arxiv, 0 for the others edcoder.py ogbn-arxiv.yaml. The paper does not mention this.
• Weight decay on Products. Table 10 gives 0.06; the config gives 0.04 app. A.4, tab. 10 ogbn-products.yaml.
• Local clusters are not top-k by PPR. The paper selects the k largest PPR elements §2.2. Code runs the ACL approximate PageRank with 1000 iterations, then keeps the first 256 entries of its output in the order the routine returns them, and only then sorts those by PPR times a clipped inverse-degree term localclustering.py localclustering.py localclustering.py. A sweep-cut conductance is computed and then discarded localclustering.py localclustering.py. Before clustering, the graph is made bidirected and given self-loops localclustering.py.
• Feature standardisation. Code standardises features with a StandardScaler for Arxiv and Papers100M, fitted on all nodes including validation and test lc_sampler.py data_proc.py. The paper does not mention this.
• Pre-training nodes. Pre-training uses local clusters seeded at train, validation and test nodes. Labels are not used main_large.py. The paper says self-supervised learning covers all nodes in a cluster §2.2 without naming splits.
• Gradient clipping. Code clips the gradient norm to 3 for large graphs main_large.py. Not in the paper.
• Batch and architecture knobs. Batch size 512 clusters, heads 8/4/8/4, attention and input dropout, and drop-edge are set in the configs ogbn-arxiv.yaml and not in the paper. Residual connections and layer norm are CLI flags only; the shipped large-graph script enables them, and the configs do not set them utils.py utils.py run_minibatch.sh.
• Script values are overridden. run_minibatch.sh passes lr 0.002, weight decay 0.04, alpha_l 4, lambda 1.0, heads 4 and 60 epochs for every dataset. With --use_cfg, the YAML config replaces these for the keys it defines utils.py run_minibatch.sh README.md. The README says --use_cfg reproduces the reported results. For Arxiv that means lr 0.0025, alpha_l 6, lambda 10 and 8 heads, not the script's values.
• Seeds. Pre-training is seeded once (seed 0) main_large.py. The "10 trials" in the paper correspond to 10 seeds of the linear probe finetune.py. The shipped large-graph script runs only seed 0 at the outer level run_minibatch.sh.
• Small-scale datasets use fixed re-masking, not multi-view random re-masking. The Cora, CiteSeer and PubMed configs set remask_method to fixed. The paper's §3.2 does not say this cora.yaml citeseer.yaml pubmed.yaml. Their EMA momentums (0, 1, 0.995) also differ from the large-graph value, and CiteSeer's momentum 1 freezes the target at its initial weights citeseer.yaml.
• replace_rate is unused. Cora's config sets replace_rate 0.15, but the masking function never reads it edcoder.py edcoder.py.
• Linear probe. Code uses a logistic regression on the seed node's embedding, trained with Adam at the config's lr_f and weight decay, for up to 1000 epochs. The batch size is fixed at 5120, overriding the argument. The reported number is the test accuracy at the best validation epoch, with ties going to the later epoch finetune.py finetune.py. The paper does not describe the probe.
• Full-batch results return two numbers. The full-batch path returns both the final accuracy and an early-stopped accuracy main_full_batch.py. The paper does not say which one it reports.
• Fine-tuning 1% and 5% has no shipped script. The code supports label_rate and a non-linear-probe path main_large.py utils.py. Both shipped scripts set label_rate 1.0 and use the linear probe run_minibatch.sh run_fullbatch.sh. Reproducing Table 4 requires setting these flags by hand.
• Full-graph Arxiv has no shipped script. The README lists full-batch runs only for Cora, CiteSeer and PubMed README.md. Table 9 therefore has no runnable recipe in the repository.
• Full-batch monitoring probe. Full-batch pre-training runs a linear probe every 200 epochs. The code's comments and loop structure show it is for monitoring; the training loop does not use its result main_full_batch.py.

Pitfalls

• Latent-loss placement. The text, eq. 4 and the code disagree on which nodes receive the latent loss (see Paper vs code). Implement the variant you intend and say which one in your write-up.
• Abstract claim does not match Table 3. The abstract puts the gain over the best baselines on Papers100M at 2.45 points or more. Table 3 gives 64.89 for GraphMAE2 versus 63.50 for GGD (1.39 points) and 62.54 for GraphMAE (2.35 points). The text also quotes 2.35 points against GraphMAE on Papers100M §3.1, which matches Table 3.
• MAG claim does not match Table 3. The text says GraphMAE2 beats GraphMAE by 1.91% on MAG §3.1. Table 3 gives 59.24 versus 58.75, a gap of 0.49.
• Ablation text does not match Table 6. The text says removing random re-mask drops accuracy by 0.45% on Products §3.3. Table 6 gives 81.59 minus 81.04, or 0.55.
• Inconsistent spread. Table 7 gives 64.89 ± 0.12 for the LC setting on Papers100M, while Table 3 gives 64.89 ± 0.04. The same mean, with different spreads.
• Hidden decoder cost. The paper says the re-mask decoding "only incurs negligible computational costs" §2.2. The code runs the decoder K=3 times per step (K decoder forward passes per batch) edcoder.py. The target branch adds a forward pass with no backward edcoder.py.
• Summed loss scales with K. The loss sums K views eq. 2. A change in K shifts the effective weight of the input loss relative to lambda; re-tune lambda when you change K.
• Scaled cosine error with large exponents. Arxiv uses an exponent of 6 ogbn-arxiv.yaml. For cosine values near 1, (1 - cos)^alpha is tiny, so gradients vanish on nodes the model already reconstructs well. Check the loss scale early.
• Transductive preprocessing. Scaling uses statistics from all nodes data_proc.py. Pre-training clusters include validation and test nodes main_large.py. Labels are not used, but the pipeline is transductive.
• Local clusters need a separate preprocessing step. Pre-computed clusters must exist at ./lc_ego_graphs/ or be generated with the external localclustering package README.md lc_sampler.py. The config's ego-graph file path is not checked against the dataset.
• Missing values in the paper. The PPR parameters alpha and epsilon, the local cluster size k, the EMA decay, gamma and the batch size are not stated in the paper. Using the code's values is the only way to match the reported setup.
• Compute. The reported runs used 8 A100 GPUs with 80 GB each app. A.4. The paper does not report wall-clock time or peak memory.

Lineage

• Builds on: GraphMAE, arXiv:2205.10803 §1 §2.1 §2.2. GraphMAE contributes masked feature reconstruction, the scaled cosine error and the fixed re-mask decoding that this paper replaces. No card for GraphMAE exists in this wiki yet.
• Related ideas it cites: PPR-Nibble local clustering §2.2 with the theorem from its appendix app. A.1; the paper relates its target generator to teacher networks in self-knowledge distillation and to BYOL-style contrastive methods §2.2.
• Compared against: GRACE, BGRL, CCA-SSG and GGD on the large datasets §3.1; DGI, MVGRL, GRACE, BGRL, InfoGCL, CCA-SSG, GGD, GAE and GraphMAE on the small datasets §3.2.
• Superseded by: not known from content in this wiki.