Link Prediction Baselines for Graphs with actual edge features

[bits-glitch]

Hi PyTorch Geometric Team, thank you very much for the fantastic library!

As most implementations of the ogb-dataset, especially the examples provided, are implemented in PyG. I decided to ask here. If this is not the right place, please let me know.

I would like to benchmark different link prediction algorithms for my custom dataset, which I have implemented in both, the PyG format and as an OGB dataset (that I store locally). For this, I have decided to utilize the reference implementations of ogb-ppa / ogb-ddi as a starting point:

https://github.com/snap-stanford/ogb/tree/master/examples/linkproppred/ppa

My graph is an undirected graph with 5-10 edge features, and 5-10 node features. As I am not that experienced with link prediction tasks (yet), I was wondering if there are baseline implementations that utilize real edge features (not aggregated node embeddings) for the link prediction task. From what I understand, the OGB graphs do not include edge features (except one graph including edge weights). Can I modify the existing baselines to utilize edge features as well or is this not easily feasible?

Many thanks!

[rusty1s]

This should be easily feasible by simply passing edge_attr as an additional argument to the GNN layers, e.g.:

def forward(self, x, edge_index, edge_attr):
    ...
    x = conv(x, edge_index, edge_attr)
    ...

given that you make use of a GNN layer that can handle multi-dimensional edge features, see here.

[bits-glitch]

Thank you very much for the answer, really helpful!

[bits-glitch]

@rusty1s would it be possible to extend the Matrix Factorization algorithm to use edge attributes as well? If so, how would you generate the embeddings?

[rusty1s]

You can also add edge_attr to the LinkPredictor, e.g.:

def forward(self, x_i, edge_attr, x_j):
    return self.MLP(torch.cat([x_i, edge_attr, x_j], dim=-1))