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.
