Captum for link prediction (edge level)

[IlyaTyagin]

Hello!
I am wondering if there is a way to explain link prediction in terms of existing edges (which were used to train a GNN model)? I am trying to find edges in the existing graph, which contribute the most to predicting new edges (new edges are only formed between existing nodes).
 
I need to get explanations for thousands of edges, so something supporting multiple predictions in batches should be better. I know you're working on batch explanations with Captum #5917, which should speed things up (and I am very looking forward to it!). As far as I can understand, methods like GNNExplainer require training a NN to explain a single sample, which, I think, is going to take way too much time for my problem.
 
I am trying to get captum work with IG method, but I am somewhat stuck (I am using this example from @RBendias as a reference: #4058 (comment)), but I replaced the attribute argument with edge mask (just like in captum_explainability example https://github.com/pyg-team/pytorch_geometric/blob/master/examples/captum_explainability.py), so the code looks like this:

output_idx = 10
 
captum_model = to_captum_model(model, mask_type='edge', output_idx=output_idx)
edge_mask = torch.ones(train_data.num_edges, requires_grad=True, device=device)
 
ig = IntegratedGradients(captum_model)
ig_attr_edge = ig.attribute(edge_mask.unsqueeze(0),
                            additional_forward_args=(test_data.x, test_data.edge_index, test_data.edge_label_index),
                            internal_batch_size=1)

 
But I'm getting the following error:
 

IndexError: The shape of the mask [9502] at index 0 does not match the shape of the indexed tensor [8976] at index 0

 
If I replace edge_mask and make it shaped like test_data.num_edges, the code works.
But then I do not quite understand what it outputs (and how it is relevant to the existing graph, which we use to train the model). Can you help?

Thanks a lot!

[RBendias]

Hi! The edge_mask is used to determine which edges in the graph that you are using for the prediction are important for the output 10. To be compatible with the edge_index you are using for the prediction, it should have the same shape, which is (1, test_data.num_edges).

[IlyaTyagin]

Hi @RBendias and thank you for your prompt reply!
I think I need to figure out something:
Graphs for training_data and test_data are different, so the edges in test_data, which I would like to explain with the GNN (which I trained with train_data), do not exist in training_data. But I would like to get the output as an edge-level learned mask of shape (1, train_data.num_edges).

I read what IG accepts as input, so my question is: what to do about test_data edges (which I would like to explain in terms of train_data edges)? Is that possible or my logic is wrong somewhere?

[RBendias]

Hi. Could you share the code of the model and also the training loop? That might help me to understand the problem better. My first assumption would be that you can only generate explanations with respect to the input you provide to the model during inference, which would be only the test_data.

[IlyaTyagin]

Oh, okay, I just checked and it seems to be the case indeed. I think I actually can format my test data in the same manner.
Do you think that it would make sense to pass the mask created in the following way:
Instead of what I originally did
edge_mask = torch.ones(test_data.num_edges, ...)
I put:

edge_mask = torch.cat(
    [
        torch.ones(train_data.num_edges),
        torch.zeros(test_data.num_edges - train_data.num_edges)
    ]
)

such that new edges (which we want to explain) are not involved in message passing?

[rusty1s]

Isn't this already handled by edge_label_index? Here, edge_index refers to the edges used for message passing (training and validation edges), and edge_label_index refers to the edges we want to explain (test edges).

[IlyaTyagin]

Hi @rusty1s, thank you for your reply!
Yes, I think that passing train_data.x and train_data.edge_index to the encoder part and test_data.edge_label_index to the decoder part should do the trick and give me the explanation of an edge from test_data.edge_label_index in terms of train_data graph.

I have another question though: is there any reason to use output_idx when we call to_captum_model?

I performed the following experiment:

Case 1: use output_idx in calling to_captum_model

output_idx = 10
edge_mask = torch.ones(train_data.num_edges, requires_grad=True, device=device)

captum_model = to_captum_model(model, mask_type='edge', output_idx=output_idx)

ig = IntegratedGradients(captum_model)
ig_attr_edge_1 = ig.attribute(
    edge_mask.unsqueeze(0),
    additional_forward_args=(train_data.x, train_data.edge_index, test_data.edge_label_index),
    internal_batch_size=1
)

Case 2: use output_idx in model forward call, but do not provide output_idx in to_captum_model

output_idx = 10
edge_mask = torch.ones(train_data.num_edges, requires_grad=True, device=device)

captum_model = to_captum_model(model, mask_type='edge', output_idx=None)
edge_mask = torch.ones(train_data.num_edges, requires_grad=True, device=device)

ig = IntegratedGradients(captum_model)
ig_attr_edge_2 = ig.attribute(
    edge_mask.unsqueeze(0),
    additional_forward_args=(train_data.x, train_data.edge_index, test_data.edge_label_index[:, output_idx]),
    internal_batch_size=1
)

And the attribution tensors are the same

torch.allclose(ig_attr_edge_1, ig_attr_edge_2)
> True

I am asking this because I am also interested to check batch captum predictions (which is a not-yet-merged pull request now), but I'd like to figure out how multiple samples are handled.

[RBendias]

@IlyaTyagin There shouldn't be any difference, also not when considering batch captum predictions.

[IlyaTyagin]

Got it, thank you very much!