What Happens to the Initial Target Node Embedding in HeteroConv During Node Updates?

[faseelmo]

I’m using the following graph, modeled using HeteroData:

data = HeteroData()

num_users = 3
num_user_features = 2

num_images = 1
num_image_features = 4

data['user'].x = torch.randn(num_users, num_user_features, dtype=torch.float)
data['image'].x = torch.randn(num_images, num_image_features, dtype=torch.float)

data['user', 'follows', 'user'].edge_index = torch.tensor([[1], [2]], dtype=torch.long)
data['user', 'post', 'image'].edge_index = torch.tensor([[0], [0]], dtype=torch.long)
data['user', 'likes', 'image'].edge_index = torch.tensor([[1, 2], [0, 0]], dtype=torch.long)

class HeteroGNN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super(HeteroGNN, self).__init__()
        aggr = ["sum"]
        self.conv1 = HeteroConv({
            ('user', 'follows', 'user'): GraphConv(-1, hidden_channels, aggr=aggr),
            ('user', 'post', 'image'): GraphConv(-1, hidden_channels, aggr=aggr),
            ('user', 'likes', 'image'): GraphConv(-1, hidden_channels, aggr=aggr),
        }, aggr="cat")

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

model = HeteroGNN(hidden_channels=5)
output = model(data.x_dict, data.edge_index_dict)
print(f"User {output['user'].shape} | Image {output['image'].shape}")

The output is

User torch.Size([3, 5]) | Image torch.Size([1, 10])

My questions is what happens to the initial embedding of the target node (e.g., Image) during the node update process in HeteroConv? Is the initial embedding discarded entirely? Or is it somehow included in the output (e.g., implicitly added to the aggregated result)?

[Romain-Nicolle]

Hi,

I am not an expert but I will try to answer your question.
According to the GraphSAGE documentation here, the layer used for message-passing is the SAGEConv one here with the parameters root_weight = True by default. Meaning that the initial embedding of each node will be used (jointly with the neighbors' features) to update their embeddings. So, yes, the initial embedding of the target node (e.g. image) is implicitly added to the aggregated result, by default. You can change this behavior by setting the root_weight = False.

So the final embeddings are dependent on the initial node embeddings. You can always access your initial node_embeddings in your data object:

print(data['user'].x)
print(data['image'].x)

I hope it helps.

Best,

[faseelmo]

Thanks for the reply but my question was specifically about how partial updates from each edge type in HeteroConv gets combined and whether the previous node embedding is explicitly added during that aggregation step.

My understanding now is that the aggregator inside HeteroConv combines the partial updates that come out of each GraphConv. The original node embedding is not explicitly added again in the final aggregation. Instead, GraphConv already includes the node's own features by default. So by the time those partial updates reach the aggregator, the node's original embedding is already baked into them.