Which Multihead Attention Implementation is Correct?

Sayed_Shaun: <pre><code class="lang-auto">class MultiHeadAttention(nn.Module):
    def __init__(self, embed_size:int, n_heads:int) -&gt; None:
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.head_dim = embed_size//n_heads
        assert(self.head_dim*n_heads==embed_size), "embed_size must be disvisible by n_heads"
        
        self.query = nn.Linear(embed_size, embed_size, bias=False)
        self.key = nn.Linear(embed_size, embed_size, bias=False)
        self.value = nn.Linear(embed_size, embed_size, bias=False)
        self.linear = nn.Linear(embed_size, embed_size)
    
    def forward(self, seq_embed:Tensor, mask:Tensor=None)-&gt;Tensor:
        batch_size, seq_len, embed_size = seq_embed.shape
        
        query = self.query(seq_embed)
        key = self.key(seq_embed)
        value = self.value(seq_embed)

        scores = torch.matmul(query, key.transpose(-2,-1))/math.sqrt(embed_size)
        if mask is not None:
            scores.masked_fill(mask==0, float("-inf"))

        weights = torch.softmax(scores, dim=-1)
        attention = torch.matmul(weights, value)
        return self.linear(attention)


attention = MultiHeadAttention(embed_size=768, n_heads=12)
seq = torch.zeros([32, 25, 768])
print(attention(seq).shape) #-&gt; output shape [32, 25, 768]
</code></pre>
<pre><code class="lang-auto">class MultiHeadAttention(nn.Module):
    def __init__(self, embed_size:int, n_heads:int) -&gt; None:
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.head_dim = embed_size//n_heads
        assert(self.head_dim*n_heads==embed_size), "embed_size must be disvisible by n_heads"
        
        self.query = nn.Linear(embed_size, embed_size, bias=False)
        self.key = nn.Linear(embed_size, embed_size, bias=False)
        self.value = nn.Linear(embed_size, embed_size, bias=False)
        self.linear = nn.Linear(embed_size, embed_size)
    
    def forward(self, seq_embed:Tensor, mask:Tensor=None)-&gt;Tensor:
        batch_size, seq_len, embed_size = seq_embed.shape
        
        query = self.query(seq_embed).reshape(batch_size, seq_len, self.n_heads, self.head_dim)
        key = self.key(seq_embed).reshape(batch_size, seq_len, self.n_heads, self.head_dim)
        value = self.value(seq_embed).reshape(batch_size, seq_len, self.n_heads, self.head_dim)

        scores = torch.matmul(query, key.transpose(-2,-1))/math.sqrt(embed_size)
        if mask is not None:
            scores.masked_fill(mask==0, float("-inf"))

        weights = torch.softmax(scores, dim=-1)
        attention = torch.matmul(weights, value)
        attention = attention.reshape(batch_size, seq_len, self.n_heads*self.head_dim)
        return self.linear(attention)


attention = MultiHeadAttention(embed_size=768, n_heads=12)
seq = torch.zeros([32, 25, 768])
print(attention(seq).shape) #-&gt; output shape [32, 25, 768]
</code></pre>

konan009: <p>Hello,</p>
<p>The idea of Multi-head Attention is that we have multiple self-attention (attention heads) that is computed in parallel, and then later, those heads gets concatenated after.</p>
<p>This one is just self-attention with <strong>1 head attention.</strong></p>
<pre><code class="lang-auto">class MultiHeadAttention(nn.Module):
    def __init__(self, embed_size:int, n_heads:int) -&gt; None:
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.head_dim = embed_size//n_heads
        assert(self.head_dim*n_heads==embed_size), "embed_size must be disvisible by n_heads"
        
        self.query = nn.Linear(embed_size, embed_size, bias=False)
        self.key = nn.Linear(embed_size, embed_size, bias=False)
        self.value = nn.Linear(embed_size, embed_size, bias=False)
        self.linear = nn.Linear(embed_size, embed_size)
    
    def forward(self, seq_embed:Tensor, mask:Tensor=None)-&gt;Tensor:
        batch_size, seq_len, embed_size = seq_embed.shape
        
        query = self.query(seq_embed)
        key = self.key(seq_embed)
        value = self.value(seq_embed)

        scores = torch.matmul(query, key.transpose(-2,-1))/math.sqrt(embed_size)
        if mask is not None:
            scores.masked_fill(mask==0, float("-inf"))

        weights = torch.softmax(scores, dim=-1)
        attention = torch.matmul(weights, value)
        return self.linear(attention)

</code></pre>
<p>This one is closer to the correct version, but you forgot to transpose the <strong>1st dim (seq_len)</strong> and <strong>2nd dim (n_heads)</strong> — So that each of our heads have the same <strong>seq_len x seq_len</strong> dimension of attention matrix. Then of course after the <strong>attention @ values</strong>, we must concatenate our heads together so that it will match the dimension of our input.</p>
<pre><code class="lang-auto">class MultiHeadAttention(nn.Module):
    def __init__(self, embed_size:int, n_heads:int) -&gt; None:
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.head_dim = embed_size//n_heads
        assert(self.head_dim*n_heads==embed_size), "embed_size must be disvisible by n_heads"
        
        self.query = nn.Linear(embed_size, embed_size, bias=False)
        self.key = nn.Linear(embed_size, embed_size, bias=False)
        self.value = nn.Linear(embed_size, embed_size, bias=False)
        self.linear = nn.Linear(embed_size, embed_size)
    
    def forward(self, seq_embed:Tensor, mask:Tensor=None)-&gt;Tensor:
        batch_size, seq_len, embed_size = seq_embed.shape
        
        query = self.query(seq_embed).reshape(batch_size, seq_len, self.n_heads, self.head_dim)
        key = self.key(seq_embed).reshape(batch_size, seq_len, self.n_heads, self.head_dim)
        value = self.value(seq_embed).reshape(batch_size, seq_len, self.n_heads, self.head_dim)

        scores = torch.matmul(query, key.transpose(-2,-1))/math.sqrt(embed_size)
        if mask is not None:
            scores.masked_fill(mask==0, float("-inf"))

        weights = torch.softmax(scores, dim=-1)
        attention = torch.matmul(weights, value)
        attention = attention.reshape(batch_size, seq_len, self.n_heads*self.head_dim)
        return self.linear(attention)
</code></pre>
<p>Here I wrote the <strong>correct version</strong>:</p>
<pre><code class="lang-auto">class MultiHeadAttention(nn.Module):
    def __init__(self, embed_size:int, n_heads:int) -&gt; None:
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.head_dim = embed_size//n_heads
        assert(self.head_dim*n_heads==embed_size), "embed_size must be disvisible by n_heads"
        
        self.query  = nn.Linear(embed_size, embed_size, bias=False)
        self.key    = nn.Linear(embed_size, embed_size, bias=False)
        self.value  = nn.Linear(embed_size, embed_size, bias=False)
        self.linear = nn.Linear(embed_size, embed_size)
    
    def forward(self, seq_embed:Tensor, mask:Tensor=None)-&gt;Tensor:
        batch_size, seq_len, embed_size = seq_embed.shape
        
        query   =   self.query(seq_embed).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1,2)
        key     =   self.key(seq_embed).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1,2)
        value   =   self.value(seq_embed).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1,2)

        scores = torch.matmul(query, key.transpose(-2,-1))/math.sqrt(query.size(-1))
        if mask is not None:
            scores.masked_fill(mask==0, float("-inf"))
        weights = torch.softmax(scores, dim=-1)
        
        attention = torch.matmul(weights, value)
        attention = attention.transpose(1,2).contiguous().view((batch_size, seq_len, embed_size))
        return self.linear(attention)

</code></pre>
<p>You can check these implementations:</p>
<ol>
<li><a href="https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L86-L106" rel="noopener nofollow ugc">Vision Transformer (ViT) - <em>Attention</em></a></li>
<li><a href="https://github.com/karpathy/minGPT/blob/master/mingpt/model.py#L29-L71" rel="noopener nofollow ugc">minGPT - <em>CausalSelfAttention</em></a></li>
<li><a href="https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py#L254-L353" rel="noopener nofollow ugc">BERT - <em>BertSelfAttention</em></a></li>
</ol>
<p>… since the solution I show you from above are just exact implementation from those links.</p>
<p>Let me know if this helps</p>
<p>Edit:<br>
I update the scaler</p>

Sayed_Shaun: <p>Hi, thanks for the reply. Could you tell me if this is correct or wrong?</p>
<pre><code class="lang-auto">class SelfAttention(nn.Module):
    def __init__(self, hidden_size:int, head_dim:int):
        super(SelfAttention, self).__init__()
        self.query = nn.Linear(hidden_size, head_dim, bias=False)
        self.key = nn.Linear(hidden_size, head_dim, bias=False)
        self.value = nn.Linear(hidden_size, head_dim, bias=False)

    def forward(self, Q:Tensor, K:Tensor, V:Tensor, mask:Tensor=None)-&gt;Tensor:
        Q, K, V = self.query(Q), self.query(K), self.value(V)
        scores = torch.bmm(Q, K.transpose(1, 2))/math.sqrt(K.shape[-1])
        if mask is not None:
            scores = scores.masked_fill(mask==0, float("-inf"))

        weights = torch.softmax(scores, dim=-1)
        attention = torch.bmm(weights, V)
        return attention

class MultiHeadAttention(nn.Module):
    def __init__(self, hidden_size:int, n_heads:int):
        super(MultiHeadAttention, self).__init__()
        self.head_dim = hidden_size//n_heads
        assert(self.head_dim * n_heads==hidden_size),"head_dim * n_heads != embed_size"
        self.heads = nn.ModuleList([SelfAttention(hidden_size, self.head_dim) for _ in range (n_heads)])
        self.linear = nn.Linear(hidden_size, hidden_size)

    def forward(self, Q:Tensor, K:Tensor, V:Tensor, mask:Tensor=None)-&gt;Tensor:
        concat = torch.cat([head(Q, K, V, mask) for head in self.heads], dim=-1)
        output = self.linear(concat)
        return output
</code></pre>

konan009: <p>Yes, this code is a more intuitive version that I sent you from above.Though this version executes sequentially because of the loop  <strong>for _ in range (n_heads)</strong>, you can use the one that I sent you since that one is faster.</p>
<aside class="quote no-group" data-username="Sayed_Shaun" data-post="3" data-topic="198996" data-full="true">
<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/sayed_shaun/48/67546_2.png" class="avatar"> Sayed_Shaun:</div>
<blockquote>
<p>Hi, thanks for the reply. Could you tell me if this is correct or wrong?</p>
<pre><code class="lang-auto">class SelfAttention(nn.Module):
    def __init__(self, hidden_size:int, head_dim:int):
        super(SelfAttention, self).__init__()
        self.query = nn.Linear(hidden_size, head_dim, bias=False)
        self.key = nn.Linear(hidden_size, head_dim, bias=False)
        self.value = nn.Linear(hidden_size, head_dim, bias=False)

    def forward(self, Q:Tensor, K:Tensor, V:Tensor, mask:Tensor=None)-&gt;Tensor:
        Q, K, V = self.query(Q), self.query(K), self.value(V)
        scores = torch.bmm(Q, K.transpose(1, 2))/math.sqrt(K.shape[-1])
        if mask is not None:
            scores = scores.masked_fill(mask==0, float("-inf"))

        weights = torch.softmax(scores, dim=-1)
        attention = torch.bmm(weights, V)
        return attention

class MultiHeadAttention(nn.Module):
    def __init__(self, hidden_size:int, n_heads:int):
        super(MultiHeadAttention, self).__init__()
        self.head_dim = hidden_size//n_heads
        assert(self.head_dim * n_heads==hidden_size),"head_dim * n_heads != embed_size"
        self.heads = nn.ModuleList([SelfAttention(hidden_size, self.head_dim) for _ in range (n_heads)])
        self.linear = nn.Linear(hidden_size, hidden_size)

    def forward(self, Q:Tensor, K:Tensor, V:Tensor, mask:Tensor=None)-&gt;Tensor:
        concat = torch.cat([head(Q, K, V, mask) for head in self.heads], dim=-1)
        output = self.linear(concat)
        return output
</code></pre>
</blockquote>
</aside>
<p>You can debug or compare them by just using a smaller input / hyperparameter size like this :</p>
<blockquote>
<p>num_heads = 2<br>
batch_no = 2<br>
window_size = 5<br>
hidden_size = 4</p>
</blockquote>
<p>Then access the weights of each <em>Q,K and V</em> by accessing <em>.weight</em> to manually compute each of the matrix multiply.</p>
<p>Regards</p>