Computational Graph of Heterogenous GNN model and Implementing Reverse Message Passing

[Linnore]

Hi, I am trying to reproduce the results of the paper Provably Powerful Graph Neural Networks for Directed Multigraphs. In this paper, it proposed the following update scheme which uses reverse message passing in the directed graph settings:

The codes provided by the author can be found here.

My question is about forward compuataional procedures of a Hetero-model that at each layer, the same type of nodes are updated by different edges. The author of this paper use heterograph to represent the directed (transaction) network. The reverse message passing is done by propogating through the edges node rev_to node. I notice that after defining a homogeneous GNN model, a heterogeneous model is created by

    # training.py line 210
    if args.reverse_mp:
        model = to_hetero(model, te_data.metadata(), aggr='mean')

Then, the output will be obtained by

    # training.py line 89
    out = model(batch.x_dict, batch.edge_index_dict, batch.edge_attr_dict)

where batch contains node features x, edges node to node, reversed edges node rev_to node, and edge attributes.

From PyG's documentation, the forward computational graph of Hetero-model is unclear. To follow the update scheme proposed by the paper, the forward output x_in using only the incoming neighbourhood should be used to update x together with x_out using only the outgoing neighbourhood after each Conv layer. Therefore, I wonder whether PyG's to_hetero(model, te_data.metadata(), aggr='mean') behaves as I expected, or instead, it only aggregates x_in and x_out by mean at the very end.

Besides the computational graph of PyG's Hetero-model, I wonder whether using a heterograph to enable such a separated reverse message passing is a good way. Is it better to use homogenous settings and do reverse message passing manually like:

# pseudo codes
out = update(x, SAGEConv(x, edge_index), SAGEConv(x, flipped_edge_index))

Thanks for any help and ideas in advance!

[rusty1s]

I think combining this scheme with to_hetero is tricky, and I feel this is better implemented from scratch. I also don't see yet how one can apply the paper on heterogenuous graphs in PyG (as you cannot simply combine/aggregate from outgoing messages).