Link prediction

[simonemarasi]

Hi,
I'm trying to perform link prediction in the simplest fashion, say predicting if there is (or not) an edge between two nodes.
I have built my custom dataset and then I performed negative sampling in my train, val and test splits, balancing negative samples with positive ones. The edge labels then are only 1s and 0s. My model is simply the following:

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

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

class EdgeDecoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
     
    def forward(self, z_dict, edge_label_index):
        row, col = edge_label_index
       
        z_src = z_dict["image"][row]
        z_dst = z_dict["attribute"][col]
        
        return torch.sum(z_src * z_dst, dim=-1)

class Model(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.encoder = GNN(hidden_channels, hidden_channels)
        self.encoder = to_hetero(self.encoder, data.metadata(), aggr='sum')
        self.decoder = EdgeDecoder()

    def forward(self, x_dict, edge_index_dict, edge_label_index):
        z_dict = self.encoder(x_dict, edge_index_dict)
        return self.decoder(z_dict, edge_label_index)

model = Model(hidden_channels=128).to(device)

with torch.no_grad():
    model(train_data.x_dict, train_data.edge_index_dict, train_data['image', 'has', 'attribute'].edge_label_index)

# I have tried also adam with several learning rates
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

Now the problem is that during training only the train loss is decreasing while val and test increase from the beginning... I thought I'm overfitting, but I cannot understand where is the problem. Here the training loop:

# It is a binary classification problem so I use Binary CrossEntropy with Sigmoid layer before...
loss_fn = torch.nn.BCEWithLogitsLoss()

def train():
    model.train()
    optimizer.zero_grad()
    pred = model(train_data.x_dict, train_data.edge_index_dict, train_data['image', 'attribute'].edge_label_index)
    target = train_data['image', 'attribute'].edge_label
    l = loss_fn (pred, target)
    l.backward()
    optimizer.step()
    return float(l)

@torch.no_grad()
def evaluate(data):
    model.eval()
    pred = model(data.x_dict, data.edge_index_dict, data['image', 'attribute'].edge_label_index)
    target = data['image', 'attribute'].edge_label
    
    l = loss_fn(pred, target)
    return float(l)

for epoch in range(1, 301):
    loss = train()
    val_loss = evaluate(val_data)
    test_loss = evaluate(test_data)
    print(f'Epoch: {epoch:03d}, Train Loss: {loss:.4f}, Val: {val_loss:.4f}, Test: {test_loss:.4f}')

Could you see something wrong in my code? Or do you have some tips? Thanks in advance!

[rusty1s]

Your example code looks pretty good to me. How do you generate val_data and test_data in the first place? What happens if you re-sample negatives for each training epoch? What happens if you replace SGD with Adam?

[simonemarasi]

Thank you for answering. Using Adam does not change in terms of loss increasing, unfortunately...
The splits are generated through RandomLinkSplit as follows:

train_data, val_data, test_data = T.RandomLinkSplit(
                      is_undirected=True, 
                      add_negative_train_samples=True,
                      neg_sampling_ratio=1.0,
                      key = "edge_label",
                      disjoint_train_ratio=0,
                      edge_types=[('image', 'has', 'attribute')],
                      rev_edge_types=[('attribute', 'rev_has', 'image')]
                      )(data)

Then, I tried to switch to Adam and to re-sample negatives for each training epoch (disabling the sampling in the RandomLinkSplit), as you suggested, adapting a function found on this forum as follows:

def neg_sampling(data):
    neg_edge_index = negative_sampling(
        edge_index = data['image', 'has', 'attribute'].edge_index, 
        num_nodes = (data['image'].x.shape[0], data['attribute'].x.shape[0]),
        num_neg_samples = data['image', 'has', 'attribute'].edge_index.size(1), method='sparse')

    edge_label_index = torch.cat([data['image', 'has', 'attribute'].edge_label_index, neg_edge_index], dim=-1,)

    edge_label = torch.cat([
        data['image', 'has', 'attribute'].edge_label,
        data['image', 'has', 'attribute'].edge_label.new_zeros(neg_edge_index.size(1))
    ], dim=0)

    return edge_label_index, edge_label

Then I call this function between model encoder and decoder both in training and in evaluate steps, on the corresponding split, e.g.

z_dict = model.encoder(data.x_dict, data.edge_index_dict)
edge_label_index, target = neg_sampling(data)
pred = model.decoder(z_dict, edge_label_index)
l = custom_loss(pred, target)

By plotting the losses now I have this situation after 100 epochs (green and orange lines are validation and test losses:

The situation seems to have improved, now validation and test losses are not exploding anymore like before... Are my corrections right and are they pointing to the direction you suggested?

[rusty1s]

I don't think you would want to resample negatives for evaluation. These should be fixed to guarantee comparable measurements across epochs. Your plot does look better though, and one can see clear signs of overfitting after epoch 15.