About transpose processing in `MultiHeadedAttention` class. about annotated-transformer

[BillyChen123]

I have the same confusion with this code.

[zh-jp]

$h*d_k$ represents the embedding dimension $d_{model}$. And $h$ represents the number of head, embedding was divided into $h$ parts, when multiheads calculate Attention separately, $h$ should be localed Index 1.

[gitfourteen]

Note that -1 represents the length $N_{token}$ (#token) of the current input or time steps of a sequence in a batch and the shape of attention scores for each head is the same, $N_{token} \times N_{token}$.

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