Wrong implementation of l2 attention

[PeterL1n]

https://github.com/mlpc-ucsd/ViTGAN/blob/d57078e49a1f6f8a0588f3a53b26c95c6c12cd1f/models/gan/stylegan2/vit_common.py#LL144C54-L144C54

Here. Why is there still a separateq,k? Shouldn't they be completely tied?

Also, I think the proper way is to use sum instead of mean. Also you need to negate the l2 distance.

This is the correct implementation

AB = torch.matmul(qk, qk.transpose(-1, -2))
AA = torch.sum(qk ** 2, -1, keepdim=True)
BB = AA.transpose(-1, -2)    # Since query and key are tied.
attn = -(AA - 2 * AB + BB)
attn = attn.mul(self.scale).softmax(-1)