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

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.

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.