Answers.Body: <p>IMO <code>xavier/glorot</code> is the correct way to initialize the <span class="math-container">$W_Q$</span> and <span class="math-container">$W_K$</span> matrices.</p>
  <p>In section <code>3.2.1</code> of the <a href="https://arxiv.org/pdf/1706.03762" rel="nofollow noreferrer">transformer paper</a> the authors explain why they would want the attention logits to be with unit standard deviation. So assuming the input <span class="math-container">$x$</span> is with unit std (which it probably is) you want your <code>queries</code> and <code>keys</code> to also have unit std, which is ensured by the <code>xavier</code> initialization.</p>
  <p>For the <span class="math-container">$W_V$</span> and <span class="math-container">$W_O$</span> matrices I am not really sure what is the correct approach. You are applying layer norm to the output <code>z</code> 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 <code>xavier</code> initialiaztion because of the backward pass.</p>
  <p>Usually I've seen people add bias only to the <span class="math-container">$W_O$</span> weights and leave the <span class="math-container">$W_Q$</span>, <span class="math-container">$W_K$</span> and <span class="math-container">$W_V$</span> with <code>bias=False</code>. Cannot comment on why they do it this way, but I think that it is ok all the bias to be concentrated in the <span class="math-container">$W_O$</span> layer.</p>
  <p>You may also consider some specific initialization of the <span class="math-container">$W_O$</span> weights as mentioned by Andrej Karpathy <a href="https://github.com/karpathy/nanoGPT/blob/a82b33b525ca9855d705656387698e13eb8e8d4b/model.py#L149" rel="nofollow noreferrer">here</a>, although I could not find this referenced anywhere.</p>
  <p>I have also seen batching the <span class="math-container">$W_Q$</span>, <span class="math-container">$W_K$</span> and <span class="math-container">$W_V$</span> matrices together in a single forward pass, but I guess in this case you should maunally set the variance for the initialization as <code>xavier</code> would be incorrect:</p>
  <pre><code>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)
  </code></pre>
  <p>You can also read a more detailed <a href="https://pi-tau.github.io/posts/transformer/" rel="nofollow noreferrer">blog post</a> that I wrote about the Transformer model.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-03-25T13:02:32.960Z
Answers.LastActivityDate: 2023-07-23T14:21:09.087Z
Answers.LastEditDate: 2023-07-23T14:21:09.087Z
Answers.UserDisplayName: pi-tau
Body: <p>Self-attention layers have 4 learnable tensors (in the vanilla formulation):</p>
  <ul>
  <li>Query matrix <span class="math-container">$W_Q$</span></li>
  <li>Key matrix <span class="math-container">$W_K$</span></li>
  <li>Value matrix <span class="math-container">$W_V$</span></li>
  <li>Output matrix <span class="math-container">$W_O$</span></li>
  </ul>
  <p>Nice illustration from  <a href="https://jalammar.github.io/illustrated-transformer/" rel="noreferrer">https://jalammar.github.io/illustrated-transformer/</a></p>
  <p><a href="https://i.sstatic.net/X6STQ.png" rel="noreferrer"><img src="https://i.sstatic.net/X6STQ.png" alt="enter image description here" /></a></p>
  <p>However, I do not know how should one choose the default initialization for these parameters.</p>
  <p>In the works, devoted to MLP and CNNs, one chooses <code>xavier/glorot</code> or <code>he</code> initialization by default, as they can be shown to approximately preserve the magnitude in the forward and backward pass, <a href="https://arxiv.org/abs/2012.05760" rel="noreferrer">as shown in these notes</a>.</p>
  <p>However, I wonder, whether there is some study of good initialization for Transformers. The default implementation in <code>Tensorflow</code> and <code>PyTorch</code> use <code>xavier/glorot</code>.</p>
  <p>Probably, any reasonable choice will work fine.</p>
  
Tags: <transformer><attention><weights><weights-initialization>
Title: Is there a proper initialization technique for the weight matrices in multi-head attention?
UserDisplayName: spiridon_the_sun_rotator
ThreadText: # Is there a proper initialization technique for the weight matrices in multi-head attention?
  
  (Asked by: username_0 on 2021-09-01)
  
  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/>
  
  [![enter image description here](https://i.sstatic.net/X6STQ.png)](https://i.sstatic.net/X6STQ.png)
  
  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](https://arxiv.org/abs/2012.05760).
  
  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.
  
  **Tags:** transformer, attention, weights, weights-initialization
  
  ---
  
  ## 1 answers
  
  ### Answer by username_1 on 2023-03-25. Score: 5
  
  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.
  
  
  ---
  
  