Question about Bigbird Random Attention Mechanism/possible bug
Author: SteeleRobert
Pull request: no

I was looking through the code for random attention in Big Bird and I found something that confused me and might be a bug.

In the bigbird_block_sparse_attention function in BigBirdSelfAttention class the selected blocks only seem to be from 1-9. You can see this if you look at the max value of the "rand_attn" variable after initializing it with _bigbird_block_rand_mask_with_head. To me this seems to imply that only block 1-9 are being selected for random attention when it should be that any blocks besides those in the global or windowed attention are selected. This behavior seems to happen because of the following code block taken from _bigbird_block_rand_mask_with_head:

```
# Total number of blocks in the mmask
num_blocks = from_seq_length // from_block_size
# Number of blocks per plan
plan_block_length = np.array(plan_from_length) // from_block_size
# till when to follow plan
max_plan_idx = plan_from_length.index(from_seq_length)
# Random Attention adjacency list
rand_attn = [
    np.zeros((num_blocks, np.sum(plan_num_rand_blocks[: max_plan_idx + 1])), dtype=np.int32)
    for i in range(num_heads)
]
```
In cases where input size > 704, plan_block_length should always be [11, num_blocks] and max_plan_idx will be equal to 1. plan_num_rand_blocks which is an input to this function will be [3,0]. You can find the reason for this in _get_rand_attn_plan.

This brings us to the for loop in _bigbird_block_rand_mask_with_head. since `max_plan_idx = 1` this loop will cover `range(0,2)`. On the first pass `plan_idx = 0` so it does not enter the first if statement `if plan_idx > 0`. `plan_num_rand_blocks[0] = 3` so it will not continue. `plan_idx = 0` so `from_start_block_id = 1`, `to_start_block_id = 0` and `curr_r_cnt = 3`.  This makes blk_rw_idx loop range(0, 11). This is then passed into _get_single_block_row_attention which is where my misunderstandings start. In this function to_end_block_id - global_block_right is always used for global attention, however if this is only covering blocks from 0-11 there should be no right global attention since that would be around num_block which could be significantly higher than 11. This also is why all the tokens are 1-9 since they are all chosen from this range. 

This brings us to the second pass through the plan_idx loop with a value of 1. This time we enter the first if statement. We then skip `if plan_num_rand_blocks[plan_idx] > 0` since `plan_num_rand_blocks[1] = 0`. This makes us enter `for pl_id in range(1):` which should only loop once with value of 0.  `if plan_num_rand_blocks[pl_id] == 0` is not entered since `plan_num_rand_blocks[0] = 3`. This means we loop blk_rw_idx with range(11, num_blocks). Since `pl_id = 0` then `to_start_block_id = 0` and we pass 11 to to_end_block_id in `to_end_block_id=plan_block_length[pl_id]`. This means that this section is also selecting token from 1-9 as was shown in the previous section. 
```
for plan_idx in range(max_plan_idx + 1):
    rnd_r_cnt = 0
    if plan_idx > 0:
        # set the row for all from_blocks starting from 0 to
        # plan_block_length[plan_idx-1]
        # column indx start fromm plan_block_length[plan_idx-1] and ends at
        # plan_block_length[plan_idx]
        if plan_num_rand_blocks[plan_idx] > 0:
            rnd_r_cnt = int(np.sum(plan_num_rand_blocks[:plan_idx]))
            curr_r_cnt = int(np.sum(plan_num_rand_blocks[: plan_idx + 1]))
            for blk_rw_idx in range(global_block_top, plan_block_length[plan_idx - 1]):
                for h in range(num_heads):
                    rand_attn[h][blk_rw_idx, rnd_r_cnt:curr_r_cnt] = self._get_single_block_row_attention(
                        block_id=blk_rw_idx,
                        to_start_block_id=plan_block_length[plan_idx - 1],
                        to_end_block_id=plan_block_length[plan_idx],
                        num_rand_blocks=plan_num_rand_blocks[plan_idx],
                        window_block_left=window_block_left,
                        window_block_right=window_block_right,
                        global_block_left=global_block_left,
                        global_block_right=global_block_right,
                    )

        for pl_id in range(plan_idx):
            if plan_num_rand_blocks[pl_id] == 0:
                continue
            for blk_rw_idx in range(plan_block_length[plan_idx - 1], plan_block_length[plan_idx]):
                rnd_r_cnt = 0
                to_start_block_id = 0
                if pl_id > 0:
                    rnd_r_cnt = int(np.sum(plan_num_rand_blocks[:pl_id]))
                    to_start_block_id = plan_block_length[pl_id - 1]
                curr_r_cnt = int(np.sum(plan_num_rand_blocks[: pl_id + 1]))
                for h in range(num_heads):
                    rand_attn[h][blk_rw_idx, rnd_r_cnt:curr_r_cnt] = self._get_single_block_row_attention(
                        block_id=blk_rw_idx,
                        to_start_block_id=to_start_block_id,
                        to_end_block_id=plan_block_length[pl_id],
                        num_rand_blocks=plan_num_rand_blocks[pl_id],
                        window_block_left=window_block_left,
                        window_block_right=window_block_right,
                        global_block_left=global_block_left,
                        global_block_right=global_block_right,
                    )

    if plan_num_rand_blocks[plan_idx] == 0:
        continue
    curr_r_cnt = int(np.sum(plan_num_rand_blocks[: plan_idx + 1]))
    from_start_block_id = global_block_top
    to_start_block_id = 0
    if plan_idx > 0:
        rnd_r_cnt = int(np.sum(plan_num_rand_blocks[:plan_idx]))
        from_start_block_id = plan_block_length[plan_idx - 1]
        to_start_block_id = plan_block_length[plan_idx - 1]

    for blk_rw_idx in range(from_start_block_id, plan_block_length[plan_idx]):
        for h in range(num_heads):
            rand_attn[h][blk_rw_idx, rnd_r_cnt:curr_r_cnt] = self._get_single_block_row_attention(
                block_id=blk_rw_idx,
                to_start_block_id=to_start_block_id,
                to_end_block_id=plan_block_length[plan_idx],
                num_rand_blocks=plan_num_rand_blocks[plan_idx],
                window_block_left=window_block_left,
                window_block_right=window_block_right,
                global_block_left=global_block_left,
                global_block_right=global_block_right,
            )
```

To me this seem like it would be a bug since random attention should be selecting from all blocks that are not in the windowed attention or the global attention. However, I double checked and this implementation seems to line up with the implementation that was published here https://github.com/google-research/bigbird. Could someone please explain what is happening here? Apologies for the long post. Thank you

TLDR; Random attention doesn't seem to select from all attention blocks and I am not sure why.


Comments:
@vasudevgupta7, who has worked on the model, might have more information!

Just a follow up on this post, I have since spoke to the first author on the paper, they limit random attention to the first 1024 tokens based on better empirical results from cross validation. Something to keep in mind if you are ever using this model since you might want to change this behavior depending on the task.