A graph autoencoder (GAE) learns a compact representation of graph data and uses it to reconstruct something about the graph. For link prediction, its practical job is to infer missing connections: given node features and known edges, it estimates whether pairs of nodes should be linked. Unlike a node classifier, it is trained to reconstruct graph structure rather than assign a class to each node. section 5.2 section 5.0
The link-prediction pipeline has four parts:
- Prepare node features and an edge index, then reserve some edges for validation and testing. The held-out links provide a check on predictions the model was not trained to make. section 5.2.3
- Encode the graph with GNN layers, such as graph convolutions, that progressively reduce the feature dimension. The result is a latent vector for each node. section 5.2.2
- Decode candidate node pairs from their latent vectors. An inner-product decoder scores each pair; applying a sigmoid turns its score into an edge probability. section 5.2.2 section 5.5.2
- Train against known edges and sampled non-edges, then assess predictions on held-out edges. This teaches the model to distinguish plausible links from absent ones. section 5.2.3
Listing 5.4 Graph encoder
from torch_geometric.nn import GCNConv #1
class GCNEncoder(torch.nn.Module): #2
def __init__(self, input_size, layers, latent_dim): #2
super().__init__() #2
self.conv0 = GCNConv(input_size, layers[0]) #3
self.conv1 = GCNConv(layers[0], layers[1]) #3
self.conv2 = GCNConv(layers[1], latent_dim) #3
def forward(self, x, edge_index): #4
x = self.conv0(x, edge_index).relu() #4
x = self.conv1(x, edge_index).relu() #4
return self.conv2(x, edge_index) #4
The graph encoder implementation shows three PyTorch Geometric GCNConv layers taking node features and edge_index to a smaller latent dimension. section 5.2.2
Predicting links in a product graph
For an Amazon product graph, convert the adjacency matrix to an edge index, normalize node features, and split edges into training, validation, and test sets. Encode the training graph, score positive edges alongside sampled non-edges, and compare predictions with their edge labels. Validation and test ROC AUC measure how well the scores rank held-out links. Keep the decoder output and loss compatible: the illustrated decoder applies a sigmoid, while the illustrated training setup specifies BCEWithLogitsLoss. section 5.2.2 section 5.2.3
Choose a GAE when reconstruction is the main task, especially predicting missing links in a graph with clear structure. A variational graph autoencoder (VGAE) adds latent-space regularization and is often preferable for noisy data or for generating a more diverse range of new graphs; a GAE is faster. The encoder and decoder are replaceable components, so the architecture can be adapted to other graph reconstruction tasks. section 5.3.2 section 5.2
