Lab: Self-Supervised Learning with GraphMAE

In this lab, we will implement GraphMAE, a state-of-the-art Masked Graph Autoencoder. We will train the model to reconstruct hidden node features, learning a robust representation of the graph without using any class labels.

---

1. The Pre-training Paradigm

In many industrial applications, you have millions of nodes but only 100 labels. Self-supervised pre-training allows you to use all that unlabeled data to build a smart model, which can then be "fine-tuned" on your few labels.

---

2. Model Architecture: Masked GNN

We will use a GAT-based encoder and decoder to perform the reconstruction.

graph_mae_lab.py

import torch
from torch_geometric.nn import GATConv

class GraphMAE(torch.nn.Module):
    def __init__(self, in_dim, hid_dim):
        super().__init__()
        self.encoder = GATConv(in_dim, hid_dim, heads=4, concat=False)
        self.decoder = GATConv(hid_dim, in_dim, heads=1, concat=False)
        
        # The [MASK] token
        self.mask_token = torch.nn.Parameter(torch.zeros(1, in_dim))

    def forward(self, x, edge_index, mask_idx):
        # 1. Mask the features
        x_masked = x.clone()
        x_masked[mask_idx] = self.mask_token
        
        # 2. Encode and Decode
        z = self.encoder(x_masked, edge_index).relu()
        return self.decoder(z, edge_index)

---

3. Training: Scaled Cosine Error (SCE)

As discussed in the GraphMAE theory, we use SCE to focus on the semantic direction of the reconstructed vectors.

sce_loss.py

def sce_loss(x, y, gamma=2):
    # Normalize vectors
    x = torch.nn.functional.normalize(x, p=2, dim=-1)
    y = torch.nn.functional.normalize(y, p=2, dim=-1)
    
    # Compute Scaled Cosine Error
    loss = (1 - (x * y).sum(dim=-1)) ** gamma
    return loss.mean()

# Training step
mask_idx = torch.randperm(data.num_nodes)[:int(0.3 * data.num_nodes)]
recon = model(data.x, data.edge_index, mask_idx)
loss = sce_loss(recon[mask_idx], data.x[mask_idx])

---

4. Why this matters?

GraphMAE is currently the "Gold Standard" for generative graph pre-training.

Advantage	Benefit
No Labels Needed	Works on raw, unlabeled data.
Feature Understanding	Deeply learns the relationships between node attributes.
Robustness	The masking process makes the model resistant to noisy inputs.

Downstream Performance

After pre-training with GraphMAE, if you freeze the encoder and train a simple Linear classifier on top of the embeddings $z$, you will often outperform a GCN trained from scratch on labels alone!

---

Hands-on Resources

• Google Colab: Full GraphMAE Pre-training Lab
• Source Code: Available in the labs/ directory.