Questions about return_semantic_attention_weights in nn.conv.han_conv

[xubingze]

Hi, I need some help with the interpretation of final results from a heterogeneous network.

I am working on a node classification task on a heterogeneous network using this architecture. I would like to identify the different weights of connected nodes on the final classification results of focal nodes. For example, if the focal node N0 (type A) is classified as good and is directly connected to nodes N2 (type A), N3 (type B), and N4 (type C), I would like to get attention scores from N2, N3, and N4 that sum up to 1, indicating their respective impact on node N0's final classification result.

I have read some papers that provide such implementations and even visualizations of such attention weights. However, I am wondering if PyG has a function that supports this implementation, as I am using PyG for my project. So far, I have tried "return_attention_weights" in GATConv, which does not support heterogeneous graph learning, and "return_semantic_attention_weights" in HANConv, which does not provide attention scores for each node's classification result.

Does PyG offer any functions or support for implementing this type of attention mechanism in a heterogeneous network? Thank you very much! I am looking forward to your insightful feedback, as this is of great help to my projects.

[rusty1s]

If you want to retrieve attention scores in a heterogeneous GNN, then this is possible but a bit cumbersome. I suggest to write a HeteroGATConv implementation:

class HeteroGATConv(torch.nn.Module):
    def __init__(self, node_types, edge_types, channels):
        self.convs = torch.nn.ModuleDict()
        for edge_type in edge_types:
             self.convs[edge_type] = GATConv((-1, -1), channels)
             
     def forward(self, x_dict, edge_index_dict):
         out_dict = defaultdict(list)
         attention_dict = {}
         for edge_type edge_index in edge_index_dict.items():
             x_src = x_dict[edge_type[0]]
             x_dst = x_dict[edge_type[2]]
             out, att = self.convs[edge_type]((x_src, x_dst), edge_index, return_attention_weights=True)
             attention_dict[edge_type] = att
             out_dict[edge_type[2]] = out
             
       # Group `out_dict` into a single tensor
       # ...
       
       return out_dict, attention_dict

[xubingze]

Thank you very much for your response! I tried the code, but I encountered an issue with the HeteroGATConv architecture.

Specifically, the weights from the same node type sum to one, while the weights for two other types each sum to one as well. In my case, if the focal node N0 (type A) is classified as good and is directly connected to nodes N1 (type A), N2 (type A), N3 (type B), and N4 (type C), then the sum of the weights (edge N1_N0 + edge N2_N0) is one, while the weights for edge N3_N0 and edge N4_N0 are each one respectively. How can I adjust the codes so that the sum of all four edge weights equals one?

[xubingze]

@rusty1s Hi Matthias, can you help me with my problem? Thank you very much!

[rusty1s]

In this case, you would need to merge edge types into a single edge_index representation. HeteroGATConv will apply the softmax for each edge type individually.

[xubingze]

@rusty1s Thank you Matthias for your response, can you elaborate more how to merge edge types into a single edge_index representation? For example,

data = HeteroData()
data['type1'].x = torch.rand(5,2)
data['type2'].x = torch.rand(3,2)
data['type3'].x = torch.rand(3,2)

data['type1','to','type2'].edge_index = torch.tensor([[0, 1, 2, 4],
                                                      [1, 0, 1, 0]]).long().contiguous()

data['type1','to','type2'].edge_attr = torch.rand(3,1)

data['type2','to','type1'].edge_index = torch.tensor([[0, 1, 2],
                                                      [2, 1, 1]]).long().contiguous()

data['type2','to','type1'].edge_attr = torch.rand(3,1)

data['type1','to','type3'].edge_index = torch.tensor([[0, 1],
                                                      [1, 1]]).long().contiguous()

data['type1','to','type3'].edge_attr = torch.rand(2,1)

data['type2','to','type3'].edge_index = torch.tensor([[0, 1, 2],
                                                      [0, 1, 2]]).long().contiguous()

data['type2','to','type3'].edge_attr = torch.rand(3,1)

Then HeteroGATConv will encode each edge type one by one, and then layer by layer.

GraphModule(
  (GNN): ModuleDict(
    (type1__to__type2): GATConv((-1, -1), 2, heads=1)
    (type2__to__type1): GATConv((-1, -1), 2, heads=1)
    (type1__to__type3): GATConv((-1, -1), 2, heads=1)
    (type2__to__type3): GATConv((-1, -1), 2, heads=1)
  )
)

Can you explain how to merge the edge types together given my examples? Thank you very much!

[rusty1s]

You can call data.to_homogeneous() and then use RGATConv to process the data.