[Python] GNN Link Prediction (GraphSAGE/PyG) - Validation AUC Consistently Below 0.5 Despite...

I'm working on a task dependency prediction problem using Graph Neural Networks with PyTorch Geometric. The goal is to predict directed precedence links (A -> B) between tasks within specific sets (called "gammes", typically ~50-60 tasks at inference).

Data & Features:

• I'm currently training on a subset of historical data related to one equipment type family ("ballon"). This subset has ~14k nodes (tasks) and ~15k edges (known dependencies), forming a Directed Acyclic Graph (DAG).
• Node features (data.x fed into the first GNN layer, dim ~401): Sentence Embeddings (from sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2, dim 384) for the task name (Nom de l'activite), which is semantically important. Learned categorical embeddings (via torch.nn.Embedding, dim 16) for the specific equipment type variant (3 unique types in this subset). Normalized duration (1 dim).
• The original Gamme name and Projet source were found to be uninformative and are not used as input features.
• Data Splitting: Using torch_geometric.transforms.RandomLinkSplit (num_val=0.1, num_test=0.1, is_undirected=False, add_negative_train_samples=True, neg_sampling_ratio=1.0, split_labels=True).

Model Architecture:

Encoder: 2-layer GraphSAGEEncoder (using SAGEConv) that takes node features + type embeddings and edge_index (training links) to produce node embeddings (currently dim=32). Includes ReLU and Dropout(0.5) between layers.

class GraphSAGEEncoder(nn.Module):
def init(self, input_feat_dim, hidden_dim, output_dim, num_types, type_embed_dim, num_layers=2):


""" Initializes the GraphSAGE encoder.
Args:
input_feat_dim (int): Dimension of continuous input features (e.g., 384 name embedding + 1 normalized duration = 385).
hidden_dim (int): Dimension of GraphSAGE hidden layers and learned embeddings.
output_dim (int): Dimension of the final node embedding.
num_types (int): Total number of unique 'Equipment Type'.
type_embed_dim (int): Desired dimension for the 'Equipment Type' embedding.
num_layers (int): Number of SAGEConv layers (e.g., 2 or 3).
"""
super(GraphSAGEEncoder, self).__init__()

# Embedding layer for Equipment Type
self.type_embedding = nn.Embedding(num_types, type_embed_dim)

# Input dimension for the first SAGEConv layer
# It's the sum of continuous features + type embedding
actual_input_dim = input_feat_dim + type_embed_dim

self.convs = nn.ModuleList()
# First layer
self.convs.append(SAGEConv(actual_input_dim, hidden_dim))
# Subsequent hidden layers
for _ in range(num_layers - 2):
self.convs.append(SAGEConv(hidden_dim, hidden_dim))
# Final layer to output dimension
self.convs.append(SAGEConv(hidden_dim, output_dim))

self.num_layers = num_layers

def forward(self, x, edge_index, type_equip_ids):
"""
Forward pass of the encoder.

Args:
x (Tensor): Continuous node features [num_nodes, input_feat_dim].
edge_index (LongTensor): Graph structure [2, num_edges].
type_equip_ids (LongTensor): Integer IDs of the equipment type for each node [num_nodes].

Returns:
Tensor: Final node embeddings [num_nodes, output_dim].
"""
# 1. Get embeddings for equipment types
type_embs = self.type_embedding(type_equip_ids)

# 2. Concatenate with continuous features
x_combined = torch.cat([x, type_embs], dim=-1)

# 3. Pass through SAGEConv layers
for i in range(self.num_layers):
x_combined = self.convs(x_combined, edge_index)
# Apply activation (except maybe for the last layer)
if i < self.num_layers - 1:
x_combined = F.relu(x_combined)
x_combined = F.dropout(x_combined, p=0.5, training=self.training) # Dropout for regularization

return x_combined


`

Link Predictor: Simple MLP that takes embeddings of source u and target v nodes and predicts link logits. (Initially included pooled global context, but removing it gave slightly better initial AUC, so currently removed). Input dim 2 * 32, hidden dim 32, output dim 1.

class LinkPredictor(nn.Module):
def __init__(self, embedding_dim, hidden_dim=64):
super(LinkPredictor, self).__init__()
self.layer_1 = nn.Linear(embedding_dim * 2, hidden_dim)
self.layer_2 = nn.Linear(hidden_dim, 1)


def forward(self, emb_u, emb_v):
# Concatenate only emb_u and emb_v
combined_embs = torch.cat([emb_u, emb_v], dim=-1)
x = F.relu(self.layer_1(combined_embs))
x = self.layer_2(x)
return x # Still returning the logits


`

Training Setup:

Optimizer: AdamW(lr=1e-4, weight_decay=1e-5) (also tried other LRs and weight decay values). Loss: torch.nn.BCEWithLogitsLoss. Process: Full-batch. Generate all node embeddings using the encoder, then predict logits for positive and negative edge pairs specified by train_data.pos_edge_label_index and train_data.neg_edge_label_index, combine logits and labels (1s and 0s) for loss calculation. Validation is similar using val_data.

The Problem:

The model learns the training data (training loss decreases steadily, e.g., from ~0.69 down to ~0.57). However, it fails to generalize:

Validation loss starts okay but increases epoch after epoch (overfitting). Crucially, Validation AUC consistently drops well below 0.5 (e.g., starts around 0.5-0.57 in the very first epoch, then quickly drops to ~0.25-0.45) and stays there. This happens across various hyperparameter settings (LR, weight decay, model dimensions).

What I've Tried:

Reducing model complexity (hidden/output dimensions). Adjusting learning rate (1e-3, 1e-4, 1e-5). Adding/adjusting weight_decay (0, 1e-6, 1e-5). Removing the explicit global context pooling from the link predictor. Verified input features (data.x) don't contain NaNs. Training runs without numerical stability issues (no NaN loss currently).

My Question:

What could be causing the validation AUC to consistently be significantly below 0.5 in this GNN link prediction setup ?

Continue reading...