MessagePassing for undirected graphs

[jazuria]

Hi,

I want to use a message-passing architecture (subclass of MessagePassing) on an undirected graph. I wasn't aware that for undirected graphs both edge directions need to be added to the edge_index list. I have updated my list accordingly so that if there is an edge between 0 and 1, edge_index will containboth (0, 1) and (1, 0).

However, updating my list does not change the model output for both training and inference. After some investigation, I found that the MessagePassing class has an argument to specify the flow for message calculations, which can either be source_to_target or target_to_source. I am now confused because that keyword argument suggests that the messages can only be computed in 1 way, which would explain why I get the same output, but which is confusing because that would mean that MessagePassing doesn't work for undirected graphs.

Can someone please tell me what I got wrong and what do I need to do to get message calculations for both sides ?

Thanks

[akihironitta]

Your approach sounds correct to me as there's one conv layer in PyG that does what you mentioned:

pytorch_geometric/torch_geometric/nn/conv/supergat_conv.py

Line 268 in 9ac9aa3

edge_index = to_undirected(edge_index, num_nodes=num_nodes)

If you could share your code here, we could maybe further help with this :)

[jazuria]

Thanks for your help! I actually used to_undirected to get the new edge list. Here is the message-passing model I am using:

class MP(MessagePassing):
    def __init__(self, input_dim, latent_dim, output_dim):
        super(MP, self).__init__(aggr="mean")
        self.message_mlp = Sequential(...)
        self.update_mlp = Sequential(...)

    def forward(self, h, edge_index):
        return self.propagate(edge_index, h=h)

    def message(self, h_i, h_j):
        m = torch.cat([h_i, h_j - h_i], dim=-1)
        return self.message_mlp(m)

    def update(self, message, h):
        z = torch.cat([h, message], dim=-1)
        return self.update_mlp(z)

I am getting almost the same output across training/inference 1e-3 of relative error between both cases (directed and undirected edge index list). The code is also getting very slow for the undirected case.

[akihironitta]

I am getting almost the same output across training/inference 1e-3 of relative error between both cases (directed and undirected edge index list).

I misunderstood your question. If the output is different, I suppose it should be working correctly. I'm not sure if 1e-3 is large enough to say the output doesn't change as I don't know the scale of the output. If you're too skeptical, you could simply pass a very simple input to a very simple version of your MP and see if the output matches the expected output calculated by hand :)

The code is also getting very slow for the undirected case.

I'd recommend looking at the profile of your script with a profiler, e.g., https://pytorch.org/docs/stable/profiler.html.

[jazuria]

Thanks for your reply! I have checked the aggregated message in the update method and the scale was indeed the problem, so I assume it is just working correctly. As for the efficiency of the message-passing model, is there any way to make the above code more efficient? I have seen some docs about certain optimizations such as: memory-efficient aggregation or k-NN graph generation for EdgeConv layer, which is similar to the architecture I have now, but do these apply to my case ?

[rusty1s]

The only way you could speed this up is via torch.compile. Memory-efficient aggregations do not apply here.