Missing matmul of ShapedAttention and Norm

[kimborgen]

Hi, first of excellent paper. Thank you for the research!

I am looking to make a standalone version of the SAS-P model, and when I went through your code, I could not find any place where you matmul the Norm with the output of Shaped Attention, as indicated in Fig 10. Red circle for convenience. I am not a math guru, but I also can't find this connection in the equations. It kinda makes sense since you experiment with fewer "skip" connections. Thus, maybe this is just a bug in the figures?

I probably missed something, but it would be amazing if you could clarify! :)

[bobby-he]

Hi, thanks for your interest in our work!

I think the confusion is that the output of shaped attention is the T-by-T attention matrix not a T-by-d representation matrix (for sequence length T and width d), which is then matmulled by the T-by-d value representations. In SAS-P (and also SAS), the value weights are simply the identity matrix (and hence the value representations are just the output of the normalisation layer).

Let me know if that helps :)