Hetero GNN Link Prediction with LinkNeighborLoader

[fenglangsss]

I would like to use heterogeneous graph neural networks for the link prediction task. Since my graph is rather large, I want to use LinkNeighborLoader to train the edges in batches. However, I found that LinkNeighborLoader can only load on a single type of edge when instantiated on a heterogeneous graph. What can I do if I want to train on multiple types of edges?

[rusty1s]

This is currently not possible as it heavily complicates the underlying computation, e.g., we will now have edge_label_index and edge_label across multiple edge types, and need to sample negatives accordingly. During training, we need to compute heads for all potential edge types, while some of them may not even be present in the current mini-batch. I definitely see the use-cases for this, but I am not yet sure on how to best integrate this feature. cc @Padarn for any additional thoughts.

[Padarn]

I guess a simple workaround would be to create multiple loaders and put the graphs together in a batch, but it's not a very efficient solution.

[rusty1s]

I remember long time ago that you wanted to bring support for multiple node types to NeighborLoader. Do you think it is worth to come up with a roadmap for this (and potential implications)?

[Padarn]

It wasn't that long ago :-) But yeah thats a good idea - I've got a couple of items to close up but I think after that it would be a nice thing to pick up!

[Padarn]

I've started a draft of this here #4765 - would appreciate your input @fenglangsss!