The Annotated Transformer

Attention

Attention maps a query and key-value pairs to an output. The model compares the query with each key, then uses the resulting weights to average the values.

We call our particular attention Scaled Dot-Product Attention. Queries and keys have dimension d_k; values have dimension d_v.

MatMul Scale Mask SoftMax MatMul QKV

In implementation, the attention scores are produced by multiplying the query matrix by the transposed key matrix, scaling by sqrt(d_k), and normalizing with softmax.

scores = torch.matmul(query, key.transpose(-2, -1))
scores = scores / math.sqrt(query.size(-1))
p_attn = scores.softmax(dim=-1)
return torch.matmul(p_attn, value)

Softmax turns the scaled scores into attention weights, so the final output is a weighted mixture of the value vectors.