Is there a proper initialization technique for the weight matrices in multi-head attention?

Question

Self-attention layers have 4 learnable tensors (in the vanilla formulation):

• Query matrix $W_Q$
• Key matrix $W_K$
• Value matrix $W_V$
• Output matrix $W_O$

Nice illustration from https://jalammar.github.io/illustrated-transformer/

However, I do not know how should one choose the default initialization for these parameters.

In the works, devoted to MLP and CNNs, one chooses xavier/glorot or he initialization by default, as they can be shown to approximately preserve the magnitude in the forward and backward pass, as shown in these notes.

However, I wonder, whether there is some study of good initialization for Transformers. The default implementation in Tensorflow and PyTorch use xavier/glorot.

Probably, any reasonable choice will work fine.

Answer 1

IMO xavier/glorot is the correct way to initialize the $W_Q$ and $W_K$ matrices.

In section 3.2.1 of the transformer paper the authors explain why they would want the attention logits to be with unit standard deviation. So assuming the input $x$ is with unit std (which it probably is) you want your queries and keys to also have unit std, which is ensured by the xavier initialization.

For the $W_V$ and $W_O$ matrices I am not really sure what is the correct approach. You are applying layer norm to the output z to scale it to unit std (getting ready for the next layer) so as far as the forward pass is concerned the initialization probably doesn't matter. I suppose it is a good idea to have again the xavier initialiaztion because of the backward pass.

Usually I've seen people add bias only to the $W_O$ weights and leave the $W_Q$, $W_K$ and $W_V$ with bias=False. Cannot comment on why they do it this way, but I think that it is ok all the bias to be concentrated in the $W_O$ layer.

You may also consider some specific initialization of the $W_O$ weights as mentioned by Andrej Karpathy here, although I could not find this referenced anywhere.

I have also seen batching the $W_Q$, $W_K$ and $W_V$ matrices together in a single forward pass, but I guess in this case you should maunally set the variance for the initialization as xavier would be incorrect:

qkv = nn.Linear(in_dim, 3 * embed_dim, bias=False)
# nn.init.xavier_normal_(qkv) # imo is incorrect
nn.init.normal_(qkv, mean=0., std=np.sqrt(2 / (in_dim+embed_dim)))
# ...
queries, keys, values = qkv(x).chunk(chunks=3, dim=-1)

You can also read a more detailed blog post that I wrote about the Transformer model.