search

pyg-team / pytorch_geometric

Graph Neural Network Library for PyTorch
https://pyg.org
MIT License
20.97k stars 3.62k forks source link

GPSConv structural encodings. #9511

Open rballeba opened 2 months ago

rballeba commented 2 months ago

🚀 The feature, motivation and pitch

In the original implementation of the GPSLayer (found in graphgps/layer/gps_layer.py), structural encodings for edges are used in the multi-head attention block. Specifically, structural encodings are encoded in batch.attn_bias as a "real-valued" adjacency matrix. However, the PyTorch Geometric (PyG) implementation of GPS does not currently support edge structural encodings. Implementing this option should be straightforward, requiring only a few modifications:

  1. Check if the parameter attn_bias has been provided by the user in the forward call.
  2. If attn_bias is provided and multi-head attention has been selected as the transformer layer, add attn_mask=attn_bias when calling self.attn.

This enhancement would align the PyG implementation more closely with the original GPS Layer, allowing for the use of structural encodings in edge attention calculations. The only problem I see is how to implement batching of real adjacency matrices (I am a relatively new user of PyTorch geometric and I do not know how to extend properly the data class).

Alternatives

No response

Additional context

No response

rusty1s commented 1 month ago

Happy to include this. Do you want to contribute it?

rballeba commented 1 month ago

@rusty1s I'd love to, but I need some time to finish some urgent projects. If, for some reason, this is not urgent, I can help in some months.