Add Performer to GPSConv

[zechengz]

Description

• Most code is adopted and simplified from https://github.com/lucidrains/performer-pytorch/blob/main/performer_pytorch/performer_pytorch.py and https://github.com/rampasek/GraphGPS/blob/main/graphgps/layer/gps_layer.py Should we add some comments to cite? Checked that both of them have the MIT License.
• There is a projection matrix redrawer that redraws the projection matrix in the performer. It seems to be not required so I just put it into the examples/graph_gps.py. Not sure where should I put it.
• There are also some other params and kernels to be added to the performer, such as softmax kernel etc. For now I added some TODOs.

ZINC Example

python3 examples/graph_gps.py --attn_type performer

I directly replace the multihead attention with performer fast attention with the ReLU kernel. The performance is similar (there are some randomness for each run and for the graph classification task) Following is the test_mae plot for each epoch.

[codecov]

Codecov Report

Merging #7465 (1481cdc) into master (0520bbb) will decrease coverage by 0.30%.
The diff coverage is 99.03%.

❗ Current head 1481cdc differs from pull request most recent head 4f98778. Consider uploading reports for the commit 4f98778 to get more accurate results

@@            Coverage Diff             @@
##           master    #7465      +/-   ##
==========================================
- Coverage   91.70%   91.40%   -0.30%     
==========================================
  Files         447      449       +2     
  Lines       24925    25024      +99     
==========================================
+ Hits        22857    22874      +17     
- Misses       2068     2150      +82     

Impacted Files	Coverage Δ
torch_geometric/nn/attention/performer.py	98.90% <98.90%> (ø)
torch_geometric/nn/attention/__init__.py	100.00% <100.00%> (ø)
torch_geometric/nn/conv/gps_conv.py	100.00% <100.00%> (ø)

... and 17 files with indirect coverage changes

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[wsad1]

Left some initial comments. Will take a look at it again.
Also, what sort of speedups do you see using Performer, it'll be nice to benchmark this.

[wsad1]

@zechengz , can PerformerAttention be used by any PyG model that needs multihead attention. If so can we move this to some common place like torch_geometric/nn/models or something?

[zechengz]

It seems that moving to torch_geometric/nn/models will cause some circular import. Any suggestions?

[wsad1]

Ah I see. Since we plan to add more attention modules. Lets create a folder called torch/nn/attention, and put all the attention modules there.

[zechengz]

Ok

[wsad1]

@zechengz I am not sure how many attention models we plan to add. But if there are potentially several out there and multiple pyg modules could use them, we should define a Attention base class and have all the attention modules we add use the same interface, similar to torch_geometric/nn/aggr. With this the user can simply pass any attention module as an argument to the GNN modules that support attention. WDYT? cc: @rusty1s

I don't want to gate this PR, with the above suggestion though. I've left some more comments and it looks good.

[wsad1]

Ah I see. Since we plan to add more attention modules. Lets create a folder called torch/nn/attention, and put all the attention modules there.

[zechengz]

@wsad1 Current plan is to add this performer and I am also working on another one called bigbird. So, with normal multihead attention we at least have three different global attentions. I also observed that there are also relative a lot of attention methods other than performer and bigbird that can be useful to add to PyG. So creating the folder torch/nn/attention makes sense to me. Thanks!

[wsad1]

Thanks @zechengz !
Good to merge after the comments are addressed.