Lab: Link Prediction with GAE

In this lab, we will use a Graph Autoencoder (GAE) to perform Link Prediction. Instead of classifying nodes, we will hide some existing edges and train the model to "predict" them back by learning a low-dimensional structural embedding.

---

1. The Challenge: Predicting the Future

Link prediction is the task of determining whether two nodes should be connected. This has massive applications in:

• Recommendation Systems: "You might know this person."
• Drug Discovery: "This molecule might interact with this protein."
• Knowledge Graphs: Filling in missing facts.

---

2. Model Architecture: GAE

We use a simple GCN encoder and an Inner Product Decoder.

gae_lab.py

import torch
from torch_geometric.nn import GCNConv, GAE
from torch_geometric.utils import train_test_split_edges

class Encoder(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, 2 * out_channels)
        self.conv2 = GCNConv(2 * out_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        return self.conv2(x, edge_index)

# Initialize
channels = 16
model = GAE(Encoder(dataset.num_features, channels))

---

3. Training: The Reconstruction Task

Unlike node classification, we don't need labels ($y$). We use the graph structure itself as the ground truth.

train_gae.py

# 1. Split edges into train, val, and test sets
# This automatically creates 'positive' and 'negative' edge samples
data = train_test_split_edges(data)

def train():
    model.train()
    optimizer.zero_grad()
    z = model.encode(data.x, data.train_pos_edge_index)
    
    # Loss: Reconstruction loss of positive edges vs negative edges
    loss = model.recon_loss(z, data.train_pos_edge_index)
    loss.backward()
    optimizer.step()
    return loss

---

4. Evaluation: AUC-ROC & AP

For link prediction, we don't use Accuracy because the graph is sparse. We use AUC-ROC (Area Under the Curve) and Average Precision (AP).

def test(pos_edge_index, neg_edge_index):
    model.eval()
    with torch.no_grad():
        z = model.encode(data.x, data.train_pos_edge_index)
    return model.test(z, pos_edge_index, neg_edge_index)

auc, ap = test(data.test_pos_edge_index, data.test_neg_edge_index)
print(f'Test AUC: {auc:.4f}, Test AP: {ap:.4f}')

Performance Tip

A well-trained GAE on a standard citation graph should achieve an AUC-ROC > 0.90. This means the model has learned a high-fidelity internal "map" of the network topology.

---

Hands-on Resources

• Google Colab: Full Link Prediction Lab
• Source Code: Available in the labs/ directory.