Qs about MetaLayer and MessagePassing

[varghele]

Hey.

I have a couple questions regarding the MetaLayer of PyTorch Geometric https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/meta.html:

1. As you have stated it is inspired by the GN paper by Battaglia et. al https://arxiv.org/abs/1806.01261, however as far as I understand the provided code, it isn't identical to any of the GraphNet Blocks that are proposed. It isn't a Full GN Block, as the edge aggregation is missing in the GlobalModel, and it isn't a MessagePassing Block, since the graph level features influence edges, nodes and themselves.
   Did I get that right? Or am I missing something and the edges do indeed influence graph level features (besides their influence on the nodes that in turn influence the graph level features)?

2. I need/want to implement a MessagePassing GNN as a "molecular feature extractor" analogous to the paper Neural Message Passing in QC (https://arxiv.org/abs/1704.01212), and this is also referenced in Battaglias paper. Removing the influence of graph level features on the nodes and edges is easy, however I also want to implement the edge aggregation and their influence on the graph level features. (Unlike the paper)

def forward(self, x, edge_index, edge_attr, u, batch):
    # x: [N, F_x], where N is the number of nodes.
    # edge_index: [2, E] with max entry N - 1.
    # edge_attr: [E, F_e]
    # u: [B, F_u]
    # batch: [N] with max entry B - 1.
    row,col=edge_index
    out = torch.cat([x[row], edge_attr], dim=1)
    out = self.global_mlp_1(out)
    out = scatter_add(out, col, dim=0, dim_size=x.size(0))
    out = torch.cat([scatter_add(x, batch, dim=0),scatter_add(out,batch,dim=0)], dim=1)
    return self.global_mlp_2(out)

Now I tried to stay close to the code of the NodeModel, however I am not entirely sure if I got this correct, and advise would be much appreciated.

3. I also substituted scatter_mean with scatter_add, as the QC paper performs elementwise summation in the edge aggregation, however I am uncertain if I understood the role of scatter_mean correctly. I have read the documentation, but I am not sure if my substitution is the correct way.

If someone could help me with this it would be greatly appreciated. Thanks in advance!
-V

[rusty1s]

1. The MetaLayer tries to follow Algorithm 1 in the given paper. Global graph features u are passed to both edge and node-level modules. For the global module, you need to take care of aggregating node and edge features by yourself, e.g, step 4-6 needs to be implemented in global_model, e.g.:

class GlobalModel(torch.nn.Module):
    ...
    def forward(self, x, edge_index, edge_attr, u, batch):
        node_aggregate = scatter_mean(x, batch, dim=0)
        edge_aggregate = scatter_mean(edge_attr, batch[edge_index[1]], dim=0)
        out = torch.cat([u, node_aggregate, edge_aggregate], dim=1)
        return self.global_mlp(out)

2. This looks correct to me, although you can also implement the two scatter_add calls as one via batch[edge_index[1]] as denoted above.
3. If the authors use summation for aggregation, then scatter_add is the correct choice. scatter_mean refers to taking the average across all features with the same index.

[varghele]

Ok thank you very much, this was immensely helpful!