Adding support for BigBird in GraphGPS

[Rahul1epic]

This PR implements BigBird in accordance with the paper https://arxiv.org/abs/2007.14062 and provides support for the same in GraphGPS.

This PR closes issue #8320

Details

• Adds BigBird sparse attention to torch_geometric/nn/attention in bigbird.py
• Adds respective test in test/nn/attention. (the test uses sequence length of 4096)
• Adds support for BigBird in gps_conv.py
• Adds BigBird to gps_conv tests.

Note: Since sparse attention depends on the block_size and number of random blocks, BigBird may fail if sequence length is too short. When such cases arise, BigBird falls back to regular dense attention. Such case is occurring in the gps_conv_test at the moment. Further, since BigBird is designed for long sequences, I have set the deafult value of block_size to 64 and number of random blocks to 3.

[xnuohz]

Can you compare bigbird vs performer (or other attention type) on a dataset (like zinc)? To see if the metrics get better or not.

[xnuohz]

can we move this under GPSConv? when should the model use large or small block_size

[Rahul1epic]

I had set the default block size to 64 assuming the sequence length of the input is very large (1024 or 4096) for example. A sequence length of less than the block size would mean there are no blocks which is incorrect I assume.
I had added assertions for the same as done here.

Further for example lets say the sequence length of the input to the transformer is 4 , block_size is 1 and rand_blocks I want to insert is also 1 (like the above case). This means there are 4 blocks say B0, B1, B2, B3
Note that random blocks cannot be allocated for blocks to local blocks and global blocks, in this case for example for block B1 ( B0 and B3 are the global blocks, those anyways attend to all other blocks)

• local window blocks are B0, B1 (itself), and B2
• globals are B0 and B3

Thus, choosing 1 random block is not possible so the random block list becomes empty thereby throwing an error.
Thus, to handle this for consistency and to allow random blocks to be inserted (and to ensure random blocks have both incoming and outgoing edge not just an outgoing edge in the blocked attention adjacency matrix), I fallback to dense attention if the sequence length is less than or equal to (2*num_random_blocks+5)*blocksize as shown here.

I am still quite unsure about these assertions and fallbacks maybe there's a better way.....
Also, Is the above explanation in the context of BigBird sparse attention right or am I missing something ?

[Rahul1epic]

Hi,
I just made a few changes to make it more concise.

I simplified handling low sequence lengths to one constraint. Its automatically falls back to dense attention if sequence length is less than or equal to (2*num_random_blocks+5)*blocksize.

Further, I added padding if the sequence length is not divisible by the block size :)

I removed the above block from test_gps_conv.py and the tests still passes (it fall back to dense attention, however since the sequence length doesn't satisfy the above constraint..)

[Rahul1epic]

Can you compare bigbird vs performer (or other attention type) on a dataset (like zinc)? To see if the metrics get better or not.

Sure, will try to do this

[Rahul1epic]

Hi @xnuohz,
I was trying to use the zinc dataset, but i cannot see the effect of bigbird because the number of nodes in each graph is too less causing the sequence length constraint to fail resulting in it falling back to dense attention.

Instead, I created a random input graphs with varying number of nodes as 1024,2048,... so on. Add created random edges for the graphs as well. The nodes have an embedding dimension of 1024.
I then observe the time it takes to do a forward pass along a single GPSconv layer for each of the attention types. I obtained the following graph:

Is this expected?

[xnuohz]

@Rahul1epic can you get the results like this?

[Rahul1epic]

Hi @xnuohz,
I followed the training example here https://github.com/pyg-team/pytorch_geometric/blob/master/examples/graph_gps.py for multihead, performer, and bigbird on the zinc subset. I trained for 50 epochs in each case. The best test MAEs I got at the end of the 50 epochs in each case were as follows:
Multihead: 0.2006
Performer: 0.1922
BigBird: 0.2251

Also, since in the zinc dataset the number of graph nodes of each sample are less (around 20-40), bigbird was always falling back to original attention as the constraint on sequence length was not satisfied for any sample.

Are the results reasonable?