How aggregation functions work for hetero data?

[shahinghasemi]

Hello, I'm wondering how aggregation functions work for heterogeneous data. To make it more clear please consider code below:

from torch_geometric.data import HeteroData
import torch
from torch_geometric.nn import SAGEConv, to_hetero, aggr

data = HeteroData()
data['man'].x = torch.eye(3, dtype=torch.float)
data['woman'].x = torch.eye(2, dtype=torch.float)
data['man', 'edge', 'woman'].edge_index = torch.tensor([[0, 1, 2],[1, 0, 0]], dtype=torch.long)
data['woman', 'edge', 'man'].edge_index = torch.tensor([[0, 0, 1],[1, 2, 0]], dtype=torch.long)

class GNNEncoder(torch.nn.Module):
    def __init__(self, neurons, aggregator ):
        super().__init__()
        if aggregator == 'sum':
            self.conv = SAGEConv((-1, -1), neurons, normalize=True, aggr=aggr.SumAggregation())
        elif aggregator == 'var':
            self.conv = SAGEConv((-1, -1), neurons, normalize=True, aggr=aggr.VarAggregation())

    def forward(self, x, edge_index):
        x = self.conv(x, edge_index).relu()
        return x

class Model(torch.nn.Module):
    def __init__(self, data, neurons, aggregator_conv, aggregator_hetero ):
        super().__init__()
        self.encoder = GNNEncoder(neurons, aggregator_conv)
        self.encoder = to_hetero(self.encoder, data.metadata(), aggr=aggregator_hetero)

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

model = Model(data=data, neurons=2, layers=1,  aggregator_conv='sum', aggregator_hetero='sum')

with torch.no_grad():
    model.encoder(data.x_dict, data.edge_index_dict)

model.train()
output = model(data.x_dict, data.edge_index_dict)
print('output: ', output)

the output is as below:

output:  {'man': tensor([[ 0.1006, -0.9949],
        [ 0.4866, -0.8736],
        [-0.3980, -0.9174]], grad_fn=<DivBackward0>), 'woman': tensor([[0.6570, 0.7539],
        [0.8027, 0.5964]], grad_fn=<DivBackward0>)}

For simplicity let's say that all the weight values are just 1s. Would you please elaborate on man[0] and woman[0] embedding calculation? It's confusing for me that for example how Variance aggregation function is applied to which features considering the relations in this hetero data.

Thanks in advance.

[EdisonLeeeee]

Hi, @shahinghasemi

According to your code example, let me introduce how the aggregation function (particularly to_hetero) works for hetero data.
First, to_hetero converts your homogeneous GNN encoder into its heterogeneous equivalent. Actually, it defines different encoders for different graph views/relations respectively, i.e.:

# We can treat a hetero graph as a group of homogenous graphs
# view 1 (man, edge, woman)
encoder1 = GNNEncoder(neurons, aggregator_conv)
# view 2 (woman, edge, man)
encoder2 = GNNEncoder(neurons, aggregator_conv)

Then, it performs aggregation separately on different graph views and gets the representation of destination node types:

woman_embedding = encoder1(...) # input (man, edge, woman) graph
man_encoder = encoder2(...) # input (woman, edge, man) graph

if you got more graph views (relations), your might have different node embeddings derived from different graph views:

woman_embedding1 = ...
woman_embedding2 = ...

Finally, aggregator_hetero (here VarAggregation) is adopted for grouping node embeddings generated by different relations.

woman_embedding = aggregator_hetero[woman_embedding1, woman_embedding2, ...]
...

Since you got only one relation for a node type, the aggregator_hetero is omitted.

For more details, you can refer to this paper and the PyG documentation. Hope this helps :)

[shahinghasemi]

Thanks for the reply. My confusion is for aggregator_conv in this case SAGE's aggr function. Let's be specific. For example we want to calculate man[0] and woman[0] embeddings which we name them man'[0] and woman'[0] respectively. For simplicity we consider the weights all 1s.

This is my assumption about the calculation which potentially can be wrong.

man'[0]
First we check what are the neighbors of man[0]. In this case woman[0] is its only neighbor. According to SAGE equation

man'[0] = data['man'][x][0] * W1_m + W2_m *VAR(data['woman']['x'][1])

woman'[0]
First we check what are the neighbors of woman[0]. In this case man[1] and man[2] is its neighbors. According to SAGE equation

woman'[0] = data['woman'][x][0] * W1_w + W2_w * (VAR{data['man']['x'][1], data['man']['x'][2]})

Is the equations above are correct? The Variance aggregator function in this case is confusing for me.

[EdisonLeeeee]

Yes, this looks correct to me.