How is add_bias_kv implemented in MultiHeadAttention?

lnair: <p>Hello,</p>
<p>I’m trying to gain more insight into how Pytorch implements the <code>MultiheadAttention</code> layer. Specifically, what does the <code>add_bias_kv</code> parameter do in the code? I see that they instantiate <code>self.bias_k</code> and <code>self.bias_v</code> in addition to <code>in_proj_bias</code> (which is determined by the <code>bias</code> param). I found the low-level implementation here: <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/transformers/attention.cpp#L321" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/transformers/attention.cpp#L321</a> - However, it just uses <code>qkv_bias</code> (which I believe is the <code>in_proj_bias</code>?). Can anyone point me to where I can find more specifics on this?</p>
<p>Thank you!</p>

lnair: <p>In case it helps someone in the future, I found the implementation here: <a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L5305" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L5305</a>. By <code>add_kv_bias</code>, it seems Pytorch actually means a <em>concatenation</em> operation, to the sequence length dimension. If you have <code>k</code> of dimensions <code>[B, L, E]</code> (batch, sequence length, embedding dim), you repeat the <code>bias_k</code> (or <code>bias_v</code>) along the batch dimension, and concatenate it along the sequence length dimension. Specifically, the <code>bias_k</code> of <code>[1, 1, E]</code> is repeated to be <code>[B, 1, E]</code>. Then it is concatenated to <code>k</code> of <code>[B, L, E]</code>, with <code>axis=1</code> (the repeat, concatenation axes will vary depending on whether batch dimension is first or not). After that operation, you will end up with a different sequence length <code>[B, L+1, E]</code>. You can see that they update the source length sequence after these adjustments, here: <a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L5351" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/torch/nn/functional.py#L5351</a>.</p>

sotovju1: <p>Thank you so much <a class="mention" href="/u/lnair">@lnair</a> for the info!</p>
<p>However, now I really wonder what’s the difference between <code>add_bias_kv</code> and <code>add_zero_attn</code> arguments: <a href="https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html#torch.nn.MultiheadAttention" class="inline-onebox" rel="noopener nofollow ugc">MultiheadAttention — PyTorch 2.0 documentation</a></p>
<p>To me, both sound very similar. The only apparent differences are that:</p>
<ol>
<li><code>add_bias_kv</code> seems to be concat alongside <code>dim=0</code> (batch dimension?) while <code>add_zero_attn</code> is done at <code>dim=1</code> (sequence dimension?)</li>
<li><code>add_bias_kv</code> adds learnable parameters, while <code>add_zero_attn</code> fills with zeros.</li>
</ol>
<p><a href="https://github.com/pytorch/pytorch/issues/27461#issuecomment-656991245" rel="noopener nofollow ugc">This is the only info I have been able to find about <code>add_zero_attn</code></a>. But then I have no idea what purpose <code>add_bias_kv</code> would serve…</p>
<p>I don’t know if anyone would have more insights on this.</p>

thinhnguyen0703: <aside class="quote no-group" data-username="sotovju1" data-post="3" data-topic="181841">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/e480ec/48.png" class="avatar"> sotovju1:</div>
<blockquote>
<p><code>add_bias_kv</code> seems to be concat alongside <code>dim=0</code> (batch dimension?) while <code>add_zero_attn</code> is done at <code>dim=1</code> (sequence dimension?)</p>
</blockquote>
</aside>
<p>add_bias_kv is the option that concats the bias along the sequence. The dim=0 is actually the sequence dim according to the implementation of mha_forward function, after that the dimentions get swapped (make batch first).</p>