Graph Autoencoders and VGAE for Link Prediction
Link prediction — inferring which edges are likely to exist in a graph — is one of the most widely studied tasks in graph machine learning. Applications span social network friend recommendation, knowledge graph completion, protein interaction prediction, and e-commerce product similarity. Graph autoencoders and their variational extensions (VGAE) offer a principled, unsupervised approach to this problem by learning latent node representations that decode into edge existence probabilities.
TGraphX implements VGAE in tgraphx.mining.vgae. This article explains the encoder-decoder architecture, the ELBO training objective, evaluation metrics (AUC and average precision), and how to train a VGAE for link prediction.
Graph Autoencoders: The Core Idea
A graph autoencoder (GAE) encodes node features into a low-dimensional embedding space using a GNN encoder, then reconstructs the adjacency matrix from those embeddings using a dot-product decoder.
The encoder maps the node feature matrix X and the adjacency matrix A to an embedding matrix Z:
Z = GNN_encoder(X, A) shape: [N, d]
The decoder reconstructs the adjacency matrix by computing pairwise inner products:
 = σ(Z Z^T) shape: [N, N]
where σ is the sigmoid function. The reconstruction loss is binary cross-entropy between the observed edges and the decoded probabilities.
During training, the model is given an incomplete graph — a subset of observed edges — and asked to predict held-out edges. At test time, nodes with high dot-product similarity in embedding space are predicted to be connected.
From GAE to VGAE: Adding Variational Inference
The Variational Graph Autoencoder (Kipf & Welling, 2016) extends GAE by treating the latent embeddings as distributions rather than point estimates. The encoder now outputs the mean μ and log-variance log σ² of a Gaussian distribution for each node:
μ, log σ² = GNN_encoder(X, A)
Z ~ N(μ, diag(σ²))
The decoder remains the same dot-product form. Training now minimizes the Evidence Lower Bound (ELBO):
L = E_q[log p(A | Z)] - KL[ q(Z | X, A) || p(Z) ]
The first term is the reconstruction likelihood (how well edges are predicted). The second term is the KL divergence between the posterior q(Z | X, A) and a standard Gaussian prior p(Z) = N(0, I). The KL term acts as a regularizer, preventing the encoder from collapsing to a deterministic mapping.
The key practical benefit of VGAE over GAE is that the latent space is smoother and better calibrated — interpolating between node embeddings in the latent space produces more semantically meaningful results.
Prerequisites
You should be familiar with:
- Variational autoencoders (VAEs) in standard settings
- GNN message-passing fundamentals (see GIN architecture for background)
- Binary classification evaluation (AUC-ROC and average precision)
Install TGraphX:
pip install tgraphx
Preparing Data for Link Prediction
Link prediction requires splitting the edge set into training, validation, and test edges. Crucially, the test edges must be held out before constructing the training graph, and the GNN encoder must not see them during training.
import torch
import numpy as np
def train_test_edge_split(edge_index, num_nodes, test_ratio=0.1, val_ratio=0.05):
"""
Split edges into train/val/test sets.
Returns edge_index for training graph, and positive/negative edges for evaluation.
"""
num_edges = edge_index.size(1)
perm = torch.randperm(num_edges)
n_test = int(num_edges * test_ratio)
n_val = int(num_edges * val_ratio)
test_edges = edge_index[:, perm[:n_test]]
val_edges = edge_index[:, perm[n_test:n_test + n_val]]
train_edges = edge_index[:, perm[n_test + n_val:]]
# Sample negative edges (non-existent edges)
edge_set = set(map(tuple, edge_index.t().tolist()))
neg_edges = []
while len(neg_edges) < n_test + n_val:
i, j = np.random.randint(0, num_nodes, 2)
if (i, j) not in edge_set and i != j:
neg_edges.append([i, j])
neg_edges = torch.tensor(neg_edges, dtype=torch.long).t()
return (
train_edges,
val_edges, neg_edges[:, :n_val],
test_edges, neg_edges[:, n_val:n_test + n_val],
)
Training a VGAE with TGraphX
TGraphX's tgraphx.mining.vgae module provides the VGAE implementation. You supply the GNN encoder; the module handles the reparameterization trick, KL loss, and dot-product decoder.
import torch
import torch.nn as nn
import torch.nn.functional as F
from tgraphx.mining.vgae import VGAE
from tgraphx.layers.sage import TensorGraphSAGELayer
class VGAEEncoder(nn.Module):
"""Two-layer GNN encoder that outputs mean and log-variance."""
def __init__(self, in_channels, hidden_channels, latent_channels):
super().__init__()
self.conv1 = TensorGraphSAGELayer(in_channels, hidden_channels)
self.conv_mu = TensorGraphSAGELayer(hidden_channels, latent_channels)
self.conv_logstd = TensorGraphSAGELayer(hidden_channels, latent_channels)
def forward(self, x, edge_index):
x = F.relu(self.conv1(x, edge_index))
mu = self.conv_mu(x, edge_index)
logstd = self.conv_logstd(x, edge_index)
return mu, logstd
# Instantiate VGAE with the encoder
in_channels, hidden_channels, latent_channels = 64, 128, 32
encoder = VGAEEncoder(in_channels, hidden_channels, latent_channels)
model = VGAE(encoder=encoder)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
# Synthetic data
num_nodes = 500
x = torch.randn(num_nodes, in_channels)
edge_index = torch.randint(0, num_nodes, (2, 2000), dtype=torch.long)
# Training loop
for epoch in range(50):
model.train()
optimizer.zero_grad()
z = model.encode(x, edge_index)
recon_loss = model.recon_loss(z, edge_index)
kl_loss = model.kl_loss()
loss = recon_loss + (1 / num_nodes) * kl_loss
loss.backward()
optimizer.step()
if epoch % 10 == 0:
print(f"Epoch {epoch:3d}: recon={recon_loss.item():.4f} kl={kl_loss.item():.4f}")
The (1 / num_nodes) weight on the KL term is the standard normalization from the original VGAE paper. This prevents the KL term from dominating on large graphs.
Evaluating with AUC and Average Precision
After training, evaluate link prediction quality on the held-out test edges:
from sklearn.metrics import roc_auc_score, average_precision_score
def evaluate_link_prediction(model, x, edge_index, pos_edge_index, neg_edge_index):
model.eval()
with torch.no_grad():
z = model.encode(x, edge_index)
# Positive edge scores
pos_src, pos_dst = pos_edge_index
pos_scores = (z[pos_src] * z[pos_dst]).sum(dim=-1).sigmoid()
# Negative edge scores
neg_src, neg_dst = neg_edge_index
neg_scores = (z[neg_src] * z[neg_dst]).sum(dim=-1).sigmoid()
scores = torch.cat([pos_scores, neg_scores]).cpu().numpy()
labels = torch.cat([
torch.ones(pos_scores.size(0)),
torch.zeros(neg_scores.size(0))
]).numpy()
auc = roc_auc_score(labels, scores)
ap = average_precision_score(labels, scores)
return auc, ap
# After training, split some edges for evaluation
pos_test = edge_index[:, :100]
neg_test = torch.stack([
torch.randint(0, num_nodes, (100,)),
torch.randint(0, num_nodes, (100,))
])
auc, ap = evaluate_link_prediction(model, x, edge_index, pos_test, neg_test)
print(f"AUC: {auc:.4f} | Average Precision: {ap:.4f}")
AUC measures ranking quality — how often the model ranks a true edge higher than a false edge. Average Precision accounts for the imbalance between positive and negative edges and is often the more informative metric when negative edges vastly outnumber positive ones.
GAE vs VGAE: When Does the Variational Component Help
For most practical link prediction tasks on citation graphs (Cora, Citeseer, DBLP), the performance difference between GAE and VGAE is small. VGAE tends to help more when:
- You need uncertainty estimates on predicted links (the latent distribution provides these)
- The training graph is very sparse and regularization matters
- You plan to use the latent space for generation or interpolation, not just prediction
GAE is faster to train, has no KL term to tune, and is a reasonable first baseline. VGAE is worth the additional complexity when the downstream task requires calibrated uncertainty or generative capabilities.
For a complementary view on generative graph models, see graph generation: classical, neural, and metrics.
Using the sklearn-style Estimator Interface
TGraphX provides an estimators module with sklearn-compatible wrappers. For VGAE:
from tgraphx.estimators.vgae import VGAEEstimator
estimator = VGAEEstimator(
in_channels=64,
hidden_channels=128,
latent_channels=32,
epochs=100,
lr=0.01,
)
# fit() expects the graph as a (x, edge_index) tuple or a Graph object
# predict() returns predicted edge scores
The estimator interface is useful for quick experimentation and hyperparameter search using scikit-learn's tools (e.g., GridSearchCV with custom scoring). For production training with custom loss schedules, use the lower-level API shown above.
Limitations and Honest Notes
The dot-product decoder is not universal. It assumes that connected nodes are similar in embedding space (a strong homophily assumption). For heterophilous graphs — where edges connect dissimilar nodes — a more expressive decoder (e.g., a small MLP) may be necessary.
Negative sampling matters enormously. Random negative sampling is the default but not ideal. Degree-biased negative sampling (sampling negatives proportional to degree) or hard negative mining can significantly change evaluation numbers without improving the model itself. Always report negative sampling strategy alongside AUC/AP numbers.
AUC on dense graphs can be misleadingly high. On dense graphs where most pairs are connected, the negative sampling space is small and random negatives are easy. AUC should be interpreted relative to the graph density.
The VGAE posterior collapse risk exists. If the KL weight is too large, the encoder may collapse to the prior and ignore node features. Anneal the KL weight from 0 to its target value during early training if you observe the KL term approaching zero early.
Shape validation. For tensor-valued node features, ensure your encoder layers use TensorGraphSAGELayer or TensorGATLayer with the appropriate spatial_rank. The shape-aware validation guide explains TGraphX's validation utilities.
Frequently Asked Questions
Can VGAE be used for directed graphs?
The standard dot-product decoder produces symmetric predictions. For directed link prediction, replace the decoder with an asymmetric one (separate projections for source and destination).
How many latent dimensions should I use?
Typical values range from 16 to 64. Larger latent spaces can overfit on small graphs. Start with 32 and tune.
Can I use GIN layers instead of GraphSAGE in the encoder?
Yes. Any GNN encoder that produces per-node embeddings can be substituted. Using TensorGINLayer may improve performance on graph-level tasks but the difference for link prediction is typically small.
Where is the full source code?
See the TGraphX GitHub and the PyPI package.
Integrating VGAE with the TGraphX Graph Object
The TGraphX Graph object carries node features, edge index, and optional metadata together. Using it with VGAE removes the need to pass tensors individually:
from tgraphx import Graph
from tgraphx.mining.vgae import VGAE
# Build a Graph object
g = Graph(
node_features=torch.randn(500, 64),
edge_index=torch.randint(0, 500, (2, 2000), dtype=torch.long),
)
# The encoder receives g.node_features and g.edge_index
encoder = VGAEEncoder(64, 128, 32)
model = VGAE(encoder=encoder)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(30):
model.train()
optimizer.zero_grad()
z = model.encode(g.node_features, g.edge_index)
loss = model.recon_loss(z, g.edge_index) + (1 / g.num_nodes) * model.kl_loss()
loss.backward()
optimizer.step()
Using the Graph object makes it easier to persist trained embeddings alongside the graph structure, and to pass graphs through the TGraphX validation utilities before training.
Visualizing Latent Space
One diagnostic advantage of VGAE over GAE is that the latent space is continuous and regularized, making it amenable to visualization. After training, reduce the [N, latent_dim] embedding matrix to 2D using PCA or UMAP and color nodes by their class label:
from sklearn.decomposition import PCA
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
model.eval()
with torch.no_grad():
z = model.encode(g.node_features, g.edge_index) # [N, latent_dim]
z_2d = PCA(n_components=2).fit_transform(z.cpu().numpy())
# Assuming you have labels [N]
labels_np = torch.randint(0, 5, (500,)).numpy()
plt.figure(figsize=(7, 5))
scatter = plt.scatter(z_2d[:, 0], z_2d[:, 1], c=labels_np, cmap="tab10", s=10, alpha=0.7)
plt.colorbar(scatter, label="Node class")
plt.title("VGAE Latent Space (PCA)")
plt.tight_layout()
plt.savefig("/tmp/vgae_latent.png")
print("Saved to /tmp/vgae_latent.png")
Well-trained VGAE embeddings should show class clusters in latent space. Overlapping clusters indicate that the encoder does not have enough information to separate classes — try increasing latent dimension or adding more GNN layers.
Reproducibility and Reporting
VGAE results depend on the random initialization of the encoder, the random train/val/test split, and the negative sampling strategy. Before reporting AUC or AP numbers in a paper or comparison, ensure:
- You fix the random seed for both PyTorch and the edge split
- You report the negative sampling strategy (random, degree-biased, hard negatives)
- You report the number of training epochs and learning rate schedule
- You run at least five random seeds and report mean ± standard deviation
from tgraphx.reproducibility import set_reproducible
with set_reproducible(seed=42):
# All training and evaluation here is deterministic
encoder = VGAEEncoder(64, 128, 32)
model = VGAE(encoder=encoder)
# ... train and evaluate ...
See the GNN research reproducibility guide for a full treatment of reproducibility practices in TGraphX.