transformer block local window attention

[vgrau98]

Fixes #7349
Towards #6357

Description

Allow basic local window attention as used in Segment Anything network (image encoder) https://arxiv.org/abs/2304.02643, https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/image_encoder.py

For now window attention is only supported for 2D inputs

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[KumoLiu]

Thanks for the PR. Could you please also add unit tests to cover the code?
We may also need to consider integrate all transformer block related components into one.
c61c6ac#diff-003ce93240016c5323702f64ea7b9f24d401ef98a8caf03763765cc958d176f6

[vgrau98]

Thanks for the PR. Could you please also add unit tests to cover the code? We may also need to consider integrate all transformer block related components into one. c61c6ac#diff-003ce93240016c5323702f64ea7b9f24d401ef98a8caf03763765cc958d176f6

Unit tests added for 2d and 3d

[KumoLiu]

I would suggest moving this rearrange part to window_partition and 2d/3d into one window_partition func to make the code cleaner. After this PR, we may also consider integrating this one, which may make this TransformerBlock too complicated. What do you think?

[marksgraham]

The main integration in that PR will be added the option of cross attention for conditioning. I'm not sure how the use of local window attention + cross attention should interact. Perhaps we would use the local windowing in the self-attention block, then not use it with the second, cross attention block?

I think it should be OK to integrate the two, but I think the integration will be much easier/cleaner if we have a single window_partition func as @KumoLiu suggests.

[vgrau98]

Thx for you comments @KumoLiu and @marksgraham

1. Yes you're right, i will clean the code following you suggestion
2. I also agree that the TransformerBlock can quickly become complex if we want to provide an exhaustive block to implement the networks presented in the literature (and those to come too). Regarding this what immediately comes to mind is defining different transformer blocks according to their subtleties (cross attention, self attention local window, and others...) rather than implement every variations in one block
3. Now considering the specifing issue of merging cross attention i don't see real issue.. local attention window was designed to be used in self-attention only (from my understanding). So cross attention should not interfere with the use of local windows for self attention. Also, i see that from c61c6ac#diff-003ce93240016c5323702f64ea7b9f24d401ef98a8caf03763765cc958d176f6 cross attention is defined using a self attention block, i would suggest to implement a cross attention block instead

[marksgraham]

Agreed with point 2, at some point we should start splitting out into different transformer blocks that inherit from a base TransformerBlock class rather than trying to implement every possible variation in one.

And also agreed with 3 - I think this specific merge should be OK, either into one block or creating a CrossAttention block for the one above linked to.

[vgrau98]

• move local window utility function to utils file
• move the logic from TransformerBlock to  SABlock