Link Prediction-based Molecule Generation

[fork123aniket]

This question is an extension of #4341. The train, test, and val functions can be seen below:-

def train():
    model.train()
    loss_all = 0

    for data in train_loader:
        optimizer.zero_grad()
        z = model.encode(data.x, data.edge_index)
        loss = model.recon_loss(z, data.pos_edge_label_index, batched_negative_sampling(data.edge_index, data.batch, 
method='dense', force_undirected=True))
        loss = loss + (1 / data.num_nodes) * model.kl_loss()
        loss.backward()
        loss_all += data.y.size(0) * float(loss)
        optimizer.step()
    return loss_all / len(train_loader.dataset)


@torch.no_grad()
def val(loader):
    model.eval()
    auc_all, ap_all = 0, 0

    for data in loader:
        z = model.encode(data.x, data.edge_index)
        auc, ap = model.test(z, data.pos_edge_label_index, data.neg_edge_label_index)
        auc_all += data.y.size(0) * float(auc)
        ap_all += data.y.size(0) * float(ap)
    return auc_all / len(val_loader.dataset), ap_all / len(val_loader.dataset)


@torch.no_grad()
def test(loader):
    model.eval()
    graph_adj = []
    
    for graph, data in enumerate(loader):
        z = model.encode(data.x, data.edge_index)
        graph_adj.append(model.decoder.forward_all(z))
        if graph == gen_graphs - 1:
          break
    return graph_adj

As advised in #4341, I'd applied batched_negative_sampling to randomly sample negative edges for multiple graphs. However, I ended up having bad VGAE model with AUC: 0.5162 and AP: 0.0683. Interestingly (enough), when didn't put batched_negative_sampling to use while training the VGAE model (using loss = model.recon_loss(z, data.pos_edge_label_index, data.neg_edge_label_index) inside train()), got AP as high as 0.7014 (which is quite good). While employing batched_negative_sampling, why it results in worst performance of the model?? Have I applied batched_negative_sampling in a wrong manner?? Do I need to use batched_negative_sampling at the validation time as well, i.e. auc, ap = model.test(z, data.pos_edge_label_index, batched_negative_sampling(data.edge_index, data.batch, method='dense', force_undirected=True)) inside val()??

Note that in order to get negative edges (neg_edge_label and neg_edge_label_index) for all the graphs present in the training set of ZINC dataset, I've set add_negative_train_samples=True, as can be seen below:-

transform = T.Compose([
    T.RandomLinkSplit(num_val=0.05, num_test=0.1, is_undirected=True,
                      split_labels=True, add_negative_train_samples=True)
])

This can also help superseding the need for using batched_negative_sampling if I ain't mistaken.

[rusty1s]

Is that training AP or validation AP? In general, it is expected that you achieve a better training loss using fixed-size negative samples (since the model can just easily overfit on these samples). I don't think you are using batched_negative_sampling wrong here. Persoanlly, I think there is no real need to apply RandomLinkSplit at all, since IMO it is better split the data across graphs (rather than within each graph). WDYT?

[fork123aniket]

Yes, all the scores I've mentioned are on the validation set and not on the training set. Ohh, do you mean I just have to compute AP scores only on the training set?? Besides, if you recommend above not to use RandomLinkSplit, do you think batched_negative_sampling can do the trick of splitting the data across graphs just by passing data.edge_index and data.batch as input arguments to batched_negative_sampling??

Furthermore, WDYT about the three functions (train, val, and test) I've written above to generating graphs for ZINC molecules (by leveraging batch processing to add support for handling multiple graphs at once at the training time), are they implemented in a right manner or do I need to make changes to them somewhere??

One interesting question I wanna ask is why do we interpret Graph generation as Link Prediction task here?? Can just predicting missing links lead us to the better molecular graph formation??

[rusty1s]

Yes, you can use data.edge_index as positive examples, and batched_negative_sampling(data.edge_index, data.batch) for negative ones.

The test functions look good to me.

The idea here would be to not infer missing links, but to infer the complete graph from the latent node embedding space. As such, this is a little bit different from the traditional link prediction case in which edges used for message passing and edges used supervision/loss computation vary.