output gradients on mini-batch inputs for graph classification

[xuesongwang]

Hi there,

I am trying to get the gradients of some outputs over the mini-batch input node features/ edge features, this way we can figure out the more important edges / nodes. My code looks pretty much like the following.

### purpose: output gradients on mini-batch inputs for graph classification
### btw I tried Captum, it's just the gradients do not make sense
### for the domain experts.


import torch
import torch.nn as nn
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
from torch_geometric.nn import ChebConv, global_mean_pool as gmp
import matplotlib.pyplot as plt

# no. nodes: 3, no. edges: 6
n_node = 3
n_edge = 6
edge_index = torch.tensor([[0, 0, 1, 1, 2, 2],
                           [1, 2, 0, 2, 0, 1]], dtype=torch.long)

node_dim = 8
dataset_size = 100
dataset = [Data(x=torch.rand((n_node, node_dim)), edge_index=edge_index,
                edge_attr=torch.rand(n_edge)) for i in range(dataset_size)]
data_loader = DataLoader(dataset, batch_size=16)


class Model(nn.Module):
    def __init__(self, in_channel=1, out_channel=2):
        super(Model, self).__init__()
        hidden_dim = 32
        self.conv = ChebConv(in_channels=in_channel, out_channels=hidden_dim, K=3)
        self.lin = nn.Linear(hidden_dim, out_channel)

    def forward(self, batch):
        x = self.conv(batch.x, batch.edge_index, batch.edge_attr)
        x = gmp(x, batch.batch)
        x = torch.relu(self.lin(x))
        return x


model = Model(node_dim)
for batch in data_loader:
    batch.x.requires_grad_(True) # (bs*n_node, node_dim)
    batch.edge_attr.requires_grad_(True) #(bs*n_edge)

    y = model(batch)

    # did not use .backward() as I might use it for model.eval(),not quite sure

    # need gradient of class_0 over x, get (bs*n_node, node_dim)
    node_grad = torch.autograd.grad(outputs=y[:, 0], inputs=batch.x, grad_outputs=torch.ones_like(y[:, 0]),
                                    allow_unused=True, retain_graph=True, create_graph=True)[0]

    node_grad = node_grad.reshape((-1, n_node, node_dim))
    print("node gradient shape:", node_grad.shape)

    # need gradient of class_0 over edge_attr, get (bs*n_edge)
    edge_grad = torch.autograd.grad(outputs=y[:, 0], inputs=batch.edge_attr, grad_outputs=torch.ones_like(y[:, 0]),
                                    allow_unused=True, retain_graph=True, create_graph=True)[0]
    edge_grad = edge_grad.reshape((-1, n_edge))
    print("edge gradient shape:", edge_grad.shape)

    # get the dataset edge importance
    dataset_edge_importance = torch.mean(edge_grad, dim=0)
    # as this is fully-connected, we can reshape it to a square matrix:
    dataset_edge_importance = dataset_edge_importance.reshape((n_node, n_node))
    plt.imshow(dataset_edge_importance.detach().cpu().numpy())

The code will plot a 3 x 3 matrix for the gradients. However, the results do not make sense to domain experts. The following screentshots are the results on my real dataset, where I have 90 nodes. They say there shouldn't be huge gradients throughout a node (large gradients for all edges connected to a specific node (left img)), it should rather look like large gradients on scattered edges (certain important edges connecting different nodes, not one node connecting all the rest nodes (right img))

So I was wondering could this be caused by the mini-batch schedule which I shouldn't have simply reshaped the gradients like this ? or this is related to the dataset itself?
Any insights would be appreciated. Please let me know if you need further elaboration.

[Padarn]

Sorry not really an expert here, but I would try your code on a much smaller graph with some synthetic data to see if what you're doing makes sense (maybe you can compare it to using .backward() directly while you figure out if you need to do this?

[rusty1s]

I think the gradient computation is fine, but I am not sure I understand the reshaping of edge_grad and how you visualize the gradients. IMO, there cannot exists any gradients in the adjacency matrix between different examples. It should more like a block-diagonal matrix. Also note that your model does not use fully-connected adjacency matrices.

[xuesongwang]

Thank you for the reply, the reason why I reshape edge_grad like this edge_grad = edge_grad.reshape((-1, n_edge))
is to get edge importance per sample, the resulting shape of the new edge_grad would become (batch_size, n_edge) , then I will calculate the average of edge_grad to get the overall edge importance/gradients regarding this dataset: torch.mean(edge_grad, dim=0).

This way I can get a vector edge_grad, shape: (n_edge, ) that shows the most important input features contributed to the final output. If all the samples use the full-connection, then this n_edge == n_node * n_node, I can reshape edge_grad.shape((n_node, n_node)) and obtain the square matrix provided in the screenshot, axes in the image refere to the node index. And yes I did not use full connection in the sample code as that is only for demonstration purpose, but I revised the code to be fully-connected and added the plotting code to be more consistent with my actual purpose.

I guess you mentioned block-diagonal matrix since in mini-batch the adjacency matrices are arranged in this manner, but looking back at the batch.edge_attr its shape is (bs*n_edge), it is not arranged in a diagonal manner? so I guess reshaping the gradient directly to (bs, n_edge) is reasonable?

Thank you for the time.

[rusty1s]

Thanks for clarifying. Indeed, if you average/aggregate gradients across examples in a batch, then your approach should work just fine (sorry for missing that). Let me know if I can be of further help.