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) .contiguous() .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) torch.all(sin(x,x,x).eq(multi(x,x,x))) # 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.