Attention, as long as gradient calculations care, is two nested tensor multiplications and a softmax. I thought that, then, multi-head attention with $h=8$ and $d_k=64$ results in the same tensor with single-head attention with $d_k = 512$, when same projection at the end of multi-head also applied. Below is my justification in Python:

import torch
import torch.nn as nn
import math

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

wq = torch.rand(512, 512)
wk = torch.rand(512, 512)
wv = torch.rand(512, 512)
class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h
        self.linears = [wq, wk, wv]
        # self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k
        query, key, value = [
            (x @ lin).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears, (query, key, value))

        # 2) Apply attention on all the projected vectors in batch.
        x, self.attn = attention(
            query, key, value, mask=mask, dropout=None

        # 3) "Concat" using a view and apply a final linear.
        x = (
            x.transpose(1, 2)
            .view(nbatches, -1, self.h * self.d_k)
        del query
        del key
        del value
        return x

multi = MultiHeadedAttention(8, 512)
sin = MultiHeadedAttention(1, 512)

x = torch.rand(2, 10, 512)

# returns tensor(True)

If I'm not mistaken and up to this point multi and single head attentions are equivalent, then where do they differ? I think they differ in the seperate optimization of heads but I can't work out the gradient calculations.


1 Answer 1


I'm not sure what you mean by "equivalent", but the output is certainly different. In the attention function, when computing the value of the scores variable the reduction happening in the matmul operation is over a subset of the components of the query and key vectors. For a single head this is the whole set of components, but for a multihead, these subsets form a a nontrivial partition of the set of components. From there on, all the results you'd get would be different.

You can see this difference when checking the shape of the attention weights.

  • $\begingroup$ MultiHeadedAttention as implemented above computes the matmul of input and same weight matrix for any combination of $h$ and $d_k$, even $h = 1$ and $d_k = d_{model}$. In the implementation, the .view after this operation puts these smaller vectors in the trailing dimension. I think when we reduce them, not the $d_{model}$ all at once, to $(L, L)$ score tables, and then concat them, the learning (backpropagation) is a different formula and hence using multiple smaller vector spaces make sense. $\endgroup$
    – tolgarecep
    Oct 8 at 10:23

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.