einsum operation in Linear Attention Part

Hi, Thanks a lot for your FLASH_pytorch, which helps a lot. I found that there are some differences from the paper in the Linear Attention Part: https://github.com/lucidrains/FLASH-pytorch/blob/main/flash_pytorch/flash_pytorch.py#L342-L343

lin_kv = einsum('b g n d, b g n e -> b d e', lin_k, v) / n
lin_out = einsum('b g n d, b d e -> b g n e', lin_q, lin_kv)

the lin_kv is three-dim (bde) And the code in the paper is

lin_kv = tf.einsum('bhke,bgh→bgke', lin_kv, mask) 
linear = tf.einsum('bgnk,bgke→bgne', lin_q, lin_kv)

the lin_kv is four-dim (bgke) It seems that the two ways are not equivalent.

Looking forward to your reply. Best,

[ShomyLiu]

Hi, It indeed that there is a reduction for all groups. However, in the final page Code 8: Pseudocode for FLASH, there is no reduction for groups. So maybe both are OK. (In my opinion, if there is a sum reduction for all groups, the attention results would be quite larger than the quad_part?)

[wangleiofficial]

When I read this part of the expressions and formulas, it should be that the reduction is the group dimension.