Recently I realized that you can implement multihead attention without doing any concatenation if you combined all the matrices in a particular way. It involves some linear algebra tricks like stacking marices A,B,C like
[A 0 0
0 B 0
0 0 C]
Where those 0s are matrices of the appropriate dimension.
I know pytorch's implementation of multi-head attention is probably better than what I'll write but assuming I have the option to either 1) use a for loop to compute all the head attentions before concatenation vs 2) getting rid of concatenation and the for loop by making some large block matrices with a number of zeros, which one would be expected to run faster?
I think it comes down to if pytorch doesn't worry about computing gradients after there's multiplication by zero...