MSA Column attention's softmax axis is not the same as the padding_mask about esm

[tomsercu]

Shape of attn_weights: Correct.
The softmax is correct in the implementation: in normal self-attention softmax goes over length of the sequence, in this case over the rows. See how it gets summed away in L222 -- that's the dimension you want to have normalized to sum to 1.

[andrewcchang]

Hi Tom, thank you so much for your quick response. Sorry that my original title was a bit misleading. I was initially a bit confused about the axis mismatch between padding_mask, attn_weights, and softmax.

If I understand the data preprocessing correctly, we pad sequences to a max_length. Therefore, the padding_mask will be TRUE at those 'C' (seq_len) dimension, which is dim=1. And in L216, we replace those padded positions in attn_weights with -10000. So I thought we should also take the softmax at dim=1. But thank you for pointing out that the einsum in L222 is summing over the rows; that's why we take softmax(-1).

Thank you!

[tomsercu]

yw! And thanks for looking into the details here!