Trying to understand nn.MultiheadAttention coming from Keras

username: jiwidi
cooked: <p>Hi!</p>
<p>I’m trying to understand the multiheadattention function at pytorch <a href="https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html" class="inline-onebox" rel="noopener nofollow ugc">MultiheadAttention — PyTorch 1.8.1 documentation</a> and if I can use it to compute the selfattention as it can be done in the keras implementation <a href="https://keras.io/api/layers/attention_layers/multi_head_attention/" class="inline-onebox" rel="noopener nofollow ugc">MultiHeadAttention layer</a></p>
<p>With Keras implementation I’m able to run selfattention over a 1D vector the following way:</p>
<pre><code class="lang-python">import tensorflow as tf

layer = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=2)
input_tensor = tf.keras.Input(shape=[8, 16])
output_tensor = layer(input_tensor, input_tensor)
print(output_tensor.shape)
(None, 8, 16)
</code></pre>
<p>I’ve tried to do the same with the pytorch implementation but couldn’t make it work. Could be because of lacking knowledge of the Transformer implementation.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
multihead_attn = nn.MultiheadAttention(16, 8)
input_tensor = torch.zeros(8,60)
multihead_attn(input_tensor,input_tensor,input_tensor)
</code></pre>
<p>With error</p>
<pre><code class="lang-auto">-&gt; 4624     tgt_len, bsz, embed_dim = query.size()
   4625     assert embed_dim == embed_dim_to_check
   4626     # allow MHA to have different sizes for the feature dimension

ValueError: not enough values to unpack (expected 3, got 2)
</code></pre>
<p>It seems the keras function does some work under the hood to make it easier for the user and perform self-attention when both Query and Value are similar. Is there a way to do the same with pytorch?</p>

username: ptrblck
cooked: <p>Do you know, what happens under the hood for 1D inputs in Keras?<br>
I guess you could mimic this behavior e.g. by using <code>unsqueeze</code> and/or <code>expand</code> on the input tensors, once you know how 1D tensors should be handled.</p>

username: jiwidi
cooked: <p>Hi, thanks for the reply.</p>
<p>I haven’t found any particular code in the keras call function <a href="https://github.com/tensorflow/tensorflow/blob/a4dfb8d1a71385bd6d122e4f27f86dcebb96712d/tensorflow/python/keras/layers/multi_head_attention.py#L487" rel="noopener nofollow ugc">call</a>, <a href="https://github.com/tensorflow/tensorflow/blob/a4dfb8d1a71385bd6d122e4f27f86dcebb96712d/tensorflow/python/keras/layers/multi_head_attention.py#L407" rel="noopener nofollow ugc">compute_attention</a> that makes some underhood 1d transformation (at least that I can identify). This goes a bit beyond my understanding, I havent been able to find any 1D attention examples. Do you know if issues in pytorch repo can be use to ask how to do calculations like this?</p>

username: ptrblck
cooked: <p>I don’t think the GitHub repository is the right place to ask questions about the usage of modules, as it’s intended to create issues in case the framework itself encounters a bug.</p>
<p>For your use case: since you cannot see what Keras is doing in the background, you could try to compare the results between Keras and PyTorch by adding the missing dimensions.<br>
I.e. the PyTorch layer expects e.g. <code>query</code> to have the shape:</p>
<blockquote>
<p>query: (L,N,E) where L is the target sequence length, N is the batch size, E is the embedding dimension.</p>
</blockquote>
<p>Based on this one dimension is missing in your input (could it be the batch dimension?).<br>
If so, use <code>input_tensor = input_tensor.unsqueeze(1)</code> and pass it to the module.</p>

username: jiwidi
cooked: <p>I think I found the solution, I can use the <a href="https://pytorch.org/docs/stable/generated/torch.nn.TransformerEncoderLayer.html" class="inline-onebox" rel="noopener nofollow ugc">TransformerEncoderLayer — PyTorch 1.8.1 documentation</a> to perform self attention myself</p>