Hi @Nat
From my understanding, Pytorch forces the embedding size to be consistent all over the computation. Hence, the embed_dim must be divisible by num_heads so later on when you “concatenate” all heads, the matrix size will be embed_dim.
The use of W0 in the documentation you showed above is not for reshaping the concatenate of heads back to embed_dim.
Here is the proof.
You can notice this code line 6
attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
attn_output is reshaped back to embed_dim already (no need W0, out_proj_weight, to help reshape)
In the 2nd line, out_proj_weight (W0) is initiated with dimension (embed_dim, embed_dim) 1
So you can see that the use of W0 is not for reshaping concatenate heads matrix.
To summary, with this implementation, Pytorch forces the matrix concatenate attention heads to have dimension is embed_dim. Hence, follow Attention is all you need paper, the purpose of using W0 to reshape concatenate heads to embed_dim is no need. Adding more evidence this statement, as shown above, out_proj_weight (W0) is initiated with size (embed_dim, embed_dim). So, W0 is not used for reshaping the “concatenate” attention heads matrix.