How do I separate out the positive edge predictions from the negative edge predictions for computing MRR and HITS?

[devanshamin]

I'm using Link Prediction on MovieLens dataset notebook.

Metrics Definition

Modified version of the original code from OGB.

def compute_metrics( 
    y_pred_pos: torch.Tensor, 
    y_pred_neg: torch.Tensor, 
    type_info: str = 'torch', 
    topk: Tuple[int] = (5, 10, 20) 
) -> Dict[str, float]: 
    """Compute HITS and Mean Reciprocal Rank (MRR) for a batch of predictions.

    Args:
        y_pred_pos (torch.Tensor): shape (batch_size, )
        y_pred_neg (torch.Tensor): shape (batch_size, num_neg_entities)
        type_info (str): 'torch' or 'numpy'
    Returns:
        metrics (dict): dictionary with metrics
    """

    metrics = {}

    if type_info == 'torch': 
        # calculate ranks
        y_pred_pos = y_pred_pos.view(-1, 1) 
        # optimistic rank: "how many negatives have a larger score than the positive?"
        # ~> the positive is ranked first among those with equal score
        optimistic_rank = (y_pred_neg > y_pred_pos).sum(dim=1) 
        # pessimistic rank: "how many negatives have at least the positive score?"
        # ~> the positive is ranked last among those with equal score
        pessimistic_rank = (y_pred_neg >= y_pred_pos).sum(dim=1) 
        ranking_list = 0.5 * (optimistic_rank + pessimistic_rank) + 1
        mrr_list = 1. / ranking_list.to(torch.float) 
    else: 
        y_pred_pos = y_pred_pos.reshape(-1, 1) 
        optimistic_rank = (y_pred_neg >= y_pred_pos).sum(dim=1) 
        pessimistic_rank = (y_pred_neg > y_pred_pos).sum(dim=1) 
        ranking_list = 0.5 * (optimistic_rank + pessimistic_rank) + 1
        mrr_list = 1. / ranking_list.astype(np.float32)

    metrics["mrr_list"] = mrr_list.tolist() 
    for k in topk: 
        value = (ranking_list <= k) 
        if type_info == 'torch': 
            value = value.to(torch.float) 
        else: 
            value = value.astype(np.float32) 
        metrics[f'hits@{k}_list'] = value.tolist() 
    
    return metrics

Modified training script,

import tqdm
import numpy as np
import torch.nn.functional as F

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 
print(f"Device: '{device}'", end="\n\n") 
prediction_triple = ("user", "rates", "movie") 
model = model.to(device) 
optimizer = torch.optim.Adam(model.parameters(), lr=0.001) 
for epoch in range(1, 6): 
    print(f">> Epoch: {epoch:03d}") 
    total_loss = total_examples = 0
    running_metrics = {} 
    for sampled_data in tqdm.tqdm(train_loader, desc="Training"): 
        optimizer.zero_grad() 
        sampled_data.to(device) 
        pred = model(sampled_data) 
        ground_truth = sampled_data[prediction_triple].edge_label
        loss = F.binary_cross_entropy_with_logits(pred, ground_truth) 
        loss.backward() 
        optimizer.step() 
        
        total_loss += float(loss) * pred.numel() 
        total_examples += pred.numel() 
        
        batch_size = sampled_data[prediction_triple].input_id.numel() 
        pred = pred.sigmoid().view(-1) 
        y_pred_pos = pred[:batch_size] 
        y_pred_neg = pred[batch_size:].view(batch_size, -1) 
        metrics = compute_metrics(y_pred_pos, y_pred_neg) 
        for metric, values in metrics.items(): 
            running_metrics.setdefault(metric, []).extend(values)

    epoch_metrics = { 
        metric.split("_")[0].upper(): round(np.mean(values), 4) 
        for metric, values in running_metrics.items() 
    } 
    print(f"Loss: {total_loss / total_examples:.4f}\nMetrics: {epoch_metrics}", end="\n\n")

Output

Device: 'cuda' 
 
>> Epoch: 001
Training: 100%|██████████| 190/190 [00:12<00:00, 14.72it/s]
Loss: 0.4415 
Metrics: {'MRR': 0.8659, 'HITS@5': 1.0, 'HITS@10': 1.0, 'HITS@20': 1.0} 
 
>> Epoch: 002
Training: 100%|██████████| 190/190 [00:13<00:00, 14.29it/s]
Loss: 0.3521 
Metrics: {'MRR': 0.9197, 'HITS@5': 1.0, 'HITS@10': 1.0, 'HITS@20': 1.0} 
 
>> Epoch: 003
Training: 100%|██████████| 190/190 [00:13<00:00, 14.34it/s]
Loss: 0.3296 
Metrics: {'MRR': 0.9291, 'HITS@5': 1.0, 'HITS@10': 1.0, 'HITS@20': 1.0} 
 
>> Epoch: 004
Training: 100%|██████████| 190/190 [00:15<00:00, 12.55it/s]
Loss: 0.3137 
Metrics: {'MRR': 0.9364, 'HITS@5': 1.0, 'HITS@10': 1.0, 'HITS@20': 1.0} 
 
>> Epoch: 005
Training: 100%|██████████| 190/190 [00:15<00:00, 12.60it/s]
Loss: 0.2985 
Metrics: {'MRR': 0.9429, 'HITS@5': 1.0, 'HITS@10': 1.0, 'HITS@20': 1.0}

Is this the correct way to create y_pred_pos and y_pred_neg?. Assumption here is that the negative edges are created from the same head as the positive edge and the tail is corrupted.

I couldn't find any examples in the OGB examples where the LinkNeighborLoader is used and the MRR/HITS are computed from it.

[rusty1s]

Can you share some more information about how you use LinkNeighborLoader? IMO, you need three different link neighbor loaders: One for training, and two for iterating over validation and test edges defined by split_idx in the OGB datasets.

[devanshamin]

Thank you @rusty1s for the response.

I'm using the same code in the Link Prediction on MovieLens dataset provided by PyG team. I provided above the modifications that I did to the training script to calculate MRR and HITS metric.

Below is the LinkNeighborLoader defined in the "Training a Heterogeneous Link-level GNN" section of the notebook,

# In the first hop, we sample at most 20 neighbors.
# In the second hop, we sample at most 10 neighbors.
# In addition, during training, we want to sample negative edges on-the-fly with
# a ratio of 2:1.
# We can make use of the `loader.LinkNeighborLoader` from PyG:
from torch_geometric.loader import LinkNeighborLoader

# Define seed edges:
edge_label_index = train_data["user", "rates", "movie"].edge_label_index
edge_label = train_data["user", "rates", "movie"].edge_label

train_loader = LinkNeighborLoader(
    data=train_data,
    num_neighbors=[20, 10],
    neg_sampling_ratio=2.0,
    edge_label_index=(("user", "rates", "movie"), edge_label_index),
    edge_label=edge_label,
    batch_size=128,
    shuffle=True,
)

If we just focus on computing the metrics during training, where the edge_index is used for message passing and edge_label_index for supervision and computing loss,

pred = model(sampled_data) 
ground_truth = sampled_data[prediction_triple].edge_label
loss = F.binary_cross_entropy_with_logits(pred, ground_truth) 

"""
batch_size = 4
neg_sampling_ratio = 2.0
edge_label = [
  1, 1, 1, 1,              # Positive label
  0, 0, 0, 0, 0, 0, 0, 0   # Negative label
]
pred = tensor([ 0.0953,  0.0216, -1.9576,  0.6710, -1.0623, -1.5792,  0.1249, -0.9376, -1.2606, -1.4963, -0.0501, -1.4178])
"""

batch_size = sampled_data[prediction_triple].input_id.numel()  # batch_size = 4 * 3 = 12
pred = pred.sigmoid().view(-1)
y_pred_pos = pred[:batch_size] 
y_pred_neg = pred[batch_size:].view(batch_size, -1) 
"""
pred = tensor([0.5238, 0.5054, 0.1237, 0.6617, 0.2569, 0.1709, 0.5312, 0.2814, 0.2209, 0.1830, 0.4875, 0.1950])
y_pred_pos = tensor([0.5238, 0.5054, 0.1237, 0.6617])
y_pred_neg = tensor([
  [0.2569, 0.1709],  # Is this correct? Does the negative edges correspond to `0.5238` in y_pred_pos?
  [0.5312, 0.2814], 
  [0.2209, 0.1830], 
  [0.4875, 0.1950]
])
"""
metrics = compute_metrics(y_pred_pos, y_pred_neg) 

[rusty1s]

This looks good so far, although negative edges do not necessarily have the same source node (it's just a random edge in your graph). You would need to specify neg_sampling="triplet" if you want to share the same source node across positive/negative examples.