Masked multi-head self-attention

Reference code openai/gpt-2 @ 9b63575, huggingface/transformers @ 71c6f69, configuration openai-community/gpt2 @ 607a30d · checkpoint tensors not independently counted · How to read

Transformer block › Attention

Query-key dot product

×

no parameterstf.matmul(q, k, transpose_b=True)

Each query position is compared with every key position in its head, including future positions. Causal masking follows this raw similarity calculation.

Si=QiKiT{\htmlData{light=qk-dot}{S_i}}={\htmlData{light=q-head split-heads}{{\htmlClass{vl-stream-q}{Q_i}}}}{\htmlData{light=k-head split-heads}{{\htmlClass{vl-stream-k}{K_i^{\mathsf T}}}}}
in
QQ [B batch, 12 heads, T positions, 64 d_h]KK [B batch, 12 heads, T positions, 64 d_h]
out
SS [B batch, 12 heads, T queries, T keys]
Codemodel.py 91, 93–94
model.py 91
    def multihead_attn(q, k, v):
model.py 93–94
        w = tf.matmul(q, k, transpose_b=True)
        w = w * tf.rsqrt(tf.cast(v.shape[-1].value, w.dtype))
Sources1

Direct evidence

The attention kernel forms the per-head query–key matrix before scaling.

openai · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 91–94

Read code excerpt
    def multihead_attn(q, k, v):

… lines 92–92 omitted …

        w = tf.matmul(q, k, transpose_b=True)
        w = w * tf.rsqrt(tf.cast(v.shape[-1].value, w.dtype))