How to properly perform inference on link prediction task for unseen graph?

[stefanschutz]

Hi, I successfully trained a model on a set of heterogeneous graphs for the link prediction task. I used PYG x lightning for training, and the AUC scores for both val and test > 0.9. The data is split into train/val/test with 720/90/90 graphs. However, when I compute the pair-wise dot products at the inference state for new unseen graphs, it seems not to give me any useful results, as the scores for the true edges are very low once applying a sigmoid function and a threshold T= 0.9. Could you identify what I did wrong here? My graphs have one node type and two edge types (type1 and type2). During training, both edge types are available. In the inference stage, I want to predict the edge_index of type2.

Model definitions:

class HeteroGNN(pl.LightningModule):
    def __init__(self, emb_size, edge_dim, heads, num_conv_layers, num_lin_layers):
        super().__init__()
        self.convs = torch.nn.ModuleList()

        for i in range(num_conv_layers):
            in_dim = -1 if (i == 0) else -1 * heads
            out_dim = emb_size
            concat = False if (i == num_conv_layers - 1) else True

            conv = HeteroConv({
                ("node", "type1", "node"): GATv2Conv(in_dim,
                                                    out_dim,
                                                    heads=heads,
                                                    edge_dim=edge_dim,
                                                    add_self_loops=False,
                                                    concat=concat),
                ("node", "type2", "node"): GATv2Conv(in_dim,
                                                               out_dim,
                                                               heads=heads,
                                                               add_self_loops=False,
                                                               concat=concat)
            }, aggr="sum")
            self.convs.append(conv)

        lin_layers = []
        for i in range(num_lin_layers):
            if i == num_lin_layers - 1:
                lin_layers.extend([Linear(emb_size, 4)])
            else:
                lin_layers.extend([Linear(emb_size, emb_size), ELU()])

        self.linear = Sequential(*lin_layers)

    def forward(self, x_dict, edge_index_dict, edge_attr):
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict, edge_attr)
            x_dict = {key: elu(x) for key, x in x_dict.items()}
        return self.linear(x_dict["node"])


class Classifier(pl.LightningModule):
    def forward(self, x_face, edge_index_cir):
        # apply dot product
        return (x_face[edge_index_cir[0]] * x_face[edge_index_cir[1]]).sum(dim=-1)


class LinkPredictor(pl.LightningModule):
    def __init__(self, **kargs):
        super().__init__()
        self.save_hyperparameters()
        self.hparams.update(kargs)
        torch.manual_seed(kargs["seed"])
        edge_dim = kargs["edge_dim"]
        heads = kargs["heads"]
        emb_size = kargs["emb_size"]
        num_conv_layers = kargs["conv_layers"]
        num_lin_layers = kargs["lin_layers"]
        self.threshold = kargs["threshold"]
        self.bce = BCEWithLogitsLoss()
        self.gnn = HeteroGNN(emb_size, edge_dim, heads, num_conv_layers, num_lin_layers)
        self.classifier = Classifier()

    def forward(self, hetero_g, neg_edge_index, infer=False):
        # only one node type
        h = self.gnn(hetero_g.x_dict, hetero_g.edge_index_dict, hetero_g.edge_attr_dict)
        if infer:
            # inference process, does not give good results
            edge_scores = h @ h.T
        else:
            # concat edge index of true graph & negative-sampled graph
            pos_edge_index = hetero_g["node", "type2", "node"].edge_index
            assert pos_edge_index.shape == neg_edge_index.shape
            combined_edge_index = torch.cat([pos_edge_index, neg_edge_index], dim=-1)
            edge_scores = self.classifier(h, combined_edge_index)
        return edge_scores

Example of the training_step (similar implementations for validation_step and test_step):

    def training_step(self, batch, batch_idx):
        pos_batch_idx = batch["node", "type2", "node"].edge_index
        neg_batch_idx = batched_negative_sampling(pos_batch_idx, batch["node"].batch)

        labels = torch.cat([torch.ones(pos_batch_idx.shape[1]), torch.zeros(neg_batch_idx.shape[1])])
        edge_scores = self(batch, neg_batch_idx)
        loss = self.bce(edge_scores, labels)
        edge_preds = get_labels_from_predictions(edge_scores, self.threshold)
        auc_scores = roc_auc_score(labels, edge_preds.detach().cpu().numpy())
        metrics = {
            "loss": loss,
            "auc": auc_scores
        }
        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True, batch_size=len(batch))
        return loss

In the inference step, I first delete the type2 edges in the new graphs, then compute the pair-wise dot products after feeding the node features into the HGNN part. I'm looping through each graph for the convenience of testing:

def predict_step(self, batch, batch_idx):
    preds = []
    list_batch = batch.to_data_list()
    for g in list_batch:
        del g["node", "type2", "node"]
        g["node", "type2", "node"] = torch.empty(size=(2, 0), dtype=torch.long)
        scores = self(g, None, True)
        preds.append(scores)
    return preds, batch.to_data_list()

Getting results:

results = trainer.predict(model=model, datamodule=datamodule)
first_graph_pred = results[0][0][0]
first_graph_origin = results[0][1][0]
num_nodes =  first_graph_origin["node"].x.shape[0]
adj = first_graph_pred .reshape((num_nodes, num_nodes))
torch.sum(sigmoid(adj) > 0.9) # ~ 600 connections, but the original edge_index of type 2 has only 16 connections.

[wsad1]

During validation and test are the edges in the prediction set also used for message passing? That might be the reason why you might be overfitting on the train/val/test data.

[stefanschutz]

I want to sample pos:neg edges with ratio 1:2 for each graph in the batch, where the graph sizes are variable. What I have tried: (i) looping and sampling each graph individually then batch them again, (ii) set num_neg_samples=batch.edge_index.size(1) * 2 for batched_negative_sampling, which does not work as the first case.

I wonder if we can control the num_neg_samples parameter to work exactly the same for (i)?

[rusty1s]

I see. You are right that the num_neg_samples argument in batched_negative_sampling is a bit non-intuitive to use. What we could do would be to also support floating-point numbers for num_neg_samples. Would that work in your case?

[stefanschutz]

That would be super great! That would work perfectly for my use-case. Thank you :)

[rusty1s]

Do you have interest in contributing this? Otherwise, I'll try to take a look.

[stefanschutz]

I would love to! However, I need sometimes to wrap up my project first, then I can submit a pull request later! Thank you :D.

[stefanschutz]

Hi again,

I just realized that during training/validating, both edge types are included to update the node embeddings (validation_step is similar to training_step provided above, but on a validation set).

For inference, I need to predict the entire set of type2, so the node embeddings are generated based on only type1 and its edge attributes (type2 set is completely removed before inference). This leads to the predictions being almost the same everywhere for unseen graphs on the test set. Conceptually, am I doing something wrong with the training pipeline?

[rusty1s]

Sorry if I am missing something, but why do you need to complete remove type2 from your graph?

[stefanschutz]

The assumption is the new unseen graphs should not have type2 edges, and the model should predict all edges of this type. Please ignore the removal part, it just means the input graphs for inference only has type1 edges.

[DoctorDinosaur]

@stefanschutz sorry for the necro, but I'm working on a similar task. Did you resolve this? Have you got a completed repo anywhere?

[AminKeshavarzi]

@stefanschutz
@rusty1s
Hi, thanks for this fruitfull discussion. I am working on a similar task. I would appreciate it if you can provide me any references or example on the inference step of a link prediction problem.