Implementation of Multi-Head Attention

ali_hejrati: <p>Hi. is this a right implementation of MHA that I create 3 Linear Layer for forecasting q, k, v? I think definition of 3 smaller layer is more stable than one large layer.</p>
<pre><code class="lang-auto">class MHA(nn.Module):
    def __init__(self, input_dim, embed_dim, h, dropout=0, args=None):
        super(SelfAttention, self).__init__()
        self.args = {} if args is None else args
        self.multihead_attn = nn.MultiheadAttention(embed_dim, h)
        self.dropout = nn.Dropout(dropout)
        self.q_proj = nn.Linear(input_dim, embed_dim)
        self.k_proj = nn.Linear(input_dim, embed_dim)
        self.v_proj = nn.Linear(input_dim, embed_dim)

    def forward(self, x):
        query = self.dropout(self.q_proj(x))
        key   = self.dropout(self.k_proj(x))
        value = self.dropout(self.v_proj(x))
        return self.multihead_attn(query, key, value, **self.args)

</code></pre>
<p>I have some other questions</p>
<ol>
<li>Why doesn’t nn.MultiheadAttention take ‘x’ and produce q, k, v itself?</li>
<li>If number of heads is set to 1 in MHA module, then the self-attenuation is obtained, but in the documention of nn.MultiheadAttention it is written that if all the variables take the x value, the self-attenuation is calculated.</li>
</ol>

tom: <p>What do you mean by more stable here?<br>
The three are strictly equivalent mathematically: If you concatenate the three weights and biases, and then send in the input <code>x</code> and split the output, you get the exact same result as if you keep them separate.<br>
Computationally, it is often more efficient to have one large input rather than three small ones.</p>
<p>Best regards</p>
<p>Thomas</p>

ali_hejrati: <aside class="quote no-group" data-username="tom" data-post="2" data-topic="162038">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>Computationally, it is often more efficient to have one large input rather than three small ones.</p>
</blockquote>
</aside>
<p>if we assume that the linear layer maps from 256 to 64, in this case, a large layer maps from 256 to 192 thus its convergence take too long? i don’t sure about this. Even the Matrix that is formed in the GPU to perform calculations is smaller?<br>
What do you think about questions 1 and 2?</p>

tom: <aside class="quote no-group" data-username="ali_hejrati" data-post="1" data-topic="162038">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ali_hejrati/48/49381_2.png" class="avatar"> ali_hejrati:</div>
<blockquote>
<ol>
<li>Why doesn’t nn.MultiheadAttention take ‘x’ and produce q, k, v itself?</li>
</ol>
</blockquote>
</aside>
<p>I think it does have the linear layer you use here, so you can feed x, x, x if you want.<br>
Note that for “external attention” in the decoder of encoder-decoder networks (e.g. in machine translation) you typically want to feed some inputs from the decoder and some from the encoder.</p>
<aside class="quote no-group" data-username="ali_hejrati" data-post="1" data-topic="162038">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ali_hejrati/48/49381_2.png" class="avatar"> ali_hejrati:</div>
<blockquote>
<ol start="2">
<li>If number of heads is set to 1 in MHA module, then the self-attenuation is obtained, but in the documention of nn.MultiheadAttention it is written that if all the variables take the x value, the self-attenuation is calculated.</li>
</ol>
</blockquote>
</aside>
<p>Self-Attention here is in contrast to external Attention, not in contrast to multi-head.</p>
<p>Best regards</p>
<p>Thomas</p>