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)
p_attn = scores.softmax(dim=-1)
if dropout is not None:
p_attn = dropout(p_attn)

wq = torch.rand(512, 512)
wk = torch.rand(512, 512)
wv = torch.rand(512, 512)
def __init__(self, h, d_model, dropout=0.1):
"Take in model size and number of heads."
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"
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(
)

# 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

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.

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.
• 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. Oct 8 at 10:23