message function documentation description

[mdanb]

The documentation mentions that the following formula for message passing is adopted in Pytorch geometric:

Then, it mentions about the function MessagePassing.message(...) that it "Constructs messages to node i in analogy to ϕ". I looked at the implementation for GCNConv (specifically, forward) just so that I can understand the implementation under the hood better, and I noticed that x is linearly transformed before being passed into propagate (and hence message). However, based on my understanding, ϕ (and hence message) is the one "supposed" to perform this operation based on the description in the docs. I understand that it doesn't really make a difference in terms of results, but I was wondering if someone could explain the rationale behind doing the transformation in forward (rather than message). Sorry for being too pedantic! Just want to understand the design choices better, I love the library!

[rusty1s]

Yes, that does not make a real difference. The reason we do this here is due to the fact that it is more memory-friendly to perform the transformation node-wise instead of edge_wise, i.e.:

x = x @ self.weight  # [num_nodes, out_channels]
self.propagate(edge_index, x=x)

is faster than doing

self.propagate(edge_index, x=x)

def message(self, x_j):
    return x_j @ self.weight  # [num_edges, out_channels]

since usually, |E| >> |N|.

[mdanb]

I think I get it. The problem that you're referring to is that x_j and the other edge_level matrix that is produced from x_j @ self.weight will temporarily coexist, which is not memory efficient. Is this correct, or am I missing something? Sorry if I'm a bit slow, I'm not used to thinking about memory much, until now!

[rusty1s]

That is correct :)

[mdanb]

This is somewhat related, but based on our discussion, this (which is in the docs under "Memory-Efficient Aggregations"):

should maybe have the MLP part in forward as x = MLP(x) and message should maybe be x_j - x_i? I know this is just a quick cooked up example so it's not really relevant but just thought I'd mention it (I know, I'm too picky... haha)

[rusty1s]

Actually, the MLP needs to be applied in the message function as its both dependent on source and target node representations. There is no way to out-source that as a pre-processing step since MLP(x_i) - MLP(x_j) is different from MLP(x_j - x_i), with the latter being more powerful as it can learn different messages dependent on the source node.

[rusty1s]

This holds only true for non-linear transformations though, not for linear ones.