Explainability of Link prediction in Heterogeneous Graph

[claudiostanzione]

Hello everyone, I need your help. I'm trying to apply Captum explainer to link prediction on a heterogeneous graph to predict the existence or not of an edge. I started from the example provided in MovieLens and modified it because the objective in the example is regression. I have two main concerns: 1. Is the model I adopted suitable for link prediction to predict the existence or not of an edge? 2. Am I creating negative edges? Otherwise, I think it won't work well. I'm attaching model, training and evaluation and error I receive when apply explainer to model. The link to the Colab notebook is as follows: https://colab.research.google.com/drive/14DrEMHGgNCEc1xOVTKvkdjF01IitmZG6?usp=sharing . I appreciate any suggestions and advice on what to do.

Model:

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["user"][row]
        z_dst = z_dict["movie"][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['user', 'rates', 'movie'].edge_label_index)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

Training and evaluation:

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['user', 'rates', 'movie'].edge_label_index)
    target = train_data['user', 'rates', 'movie'].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['user', 'rates', 'movie'].edge_label_index)
    target = data['user', 'rates', 'movie'].edge_label
    
    l = loss_fn(pred, target)
    return float(l)

for epoch in range(1, 30):
    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}')

Explain:

explainer = Explainer(
    model=model,
    algorithm=CaptumExplainer('IntegratedGradients'),
    explanation_type='model',
    model_config=dict(
        mode='binary_classification',
        task_level='edge',
        return_type='raw',
    ),
    node_mask_type='attributes',
    edge_mask_type='object',
    threshold_config=dict(
        threshold_type='topk',
        value=200,
    ),
)

index = torch.tensor([2, 100])  # Explain edge labels with index 2 and 10.
explanation = explainer(
    train_data.x_dict,
    train_data.edge_index_dict,
    index=index,
    edge_label_index=train_data['user', 'rates', 'movie'].edge_label_index,
)
print(f'Generated explanations in {explanation.available_explanations}')

visualize_feature_importance(explanation, top_k=10)

Error:
AssertionError: Output must be 2D to select tensor of targets.

[mzamini92]

I guess you are creating negative edges by setting the edge_label of some edges to 0. This is necessary for the IntegratedGradients algorithm to work, as it needs to be able to explain both positive and negative predictions. The error you are getting is because the IntegratedGradients algorithm expects the output of the model to be a 2D tensor. In your case, the output of the model is a 1D tensor, which is not supported by the algorithm. you can add a final linear layer to the model that maps the output of the 1D tensor to a 2D tensor.

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()
        self.linear = torch.nn.Linear(hidden_channels, 2)

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

[claudiostanzione]

Thank you very much for your prompt reply @mzamini92.
Changing the template with the code you suggested I get the following error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[<ipython-input-13-6ebe11d4951d>](https://liixqcgw7d-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20230711-060203-RC00_547140704#) in <cell line: 51>()
     50 
     51 with torch.no_grad():
---> 52     model(train_data.x_dict, train_data.edge_index_dict, train_data['user', 'rates', 'movie'].edge_label_index)
     53 
     54 optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

3 frames
[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/linear.py](https://liixqcgw7d-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20230711-060203-RC00_547140704#) in forward(self, input)
    112 
    113     def forward(self, input: Tensor) -> Tensor:
--> 114         return F.linear(input, self.weight, self.bias)
    115 
    116     def extra_repr(self) -> str:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x70586 and 128x2)

[mzamini92]

you need to reshape the input tensor before passing it to the linear layer. Try modify it like this:

def forward(self, x_dict, edge_index_dict, edge_label_index):
    z_dict = self.encoder(x_dict, edge_index_dict)
    decoder_output = self.decoder(z_dict, edge_label_index)
    reshaped_output = decoder_output.unsqueeze(-1)  # Reshape to (batch_size, hidden_channels, 1)
    linear_output = self.linear(reshaped_output)
    return linear_output.squeeze(-1)  # Remove the extra dimension

[claudiostanzione]

Thank you @mzamini92 for you reply. The error now is the following:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (70586x1 and 128x2)

[claudiostanzione]

Also regarding negative edges, I am not sure how to create and where

[rusty1s]

This seems to be the same issue as #7702. Closing as duplicate. Feel free to re-open if that is not the case. We will try to fix this ASAP.