GPT-2

All models, with previews →

OpenAI · Feb–Nov 2019Dense decoder

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

Full sequenceTrainingPrefillDecode

Transformer block › Attention

Query-key dot product

×

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

The new token's query is compared with every key read from the cache (1,001 positions), in its head.

Si=QiKiT{\htmlData{light=qk-dot}{S_i}}={\htmlData{light=q-head split-heads}{{\htmlClass{vl-stream-q}{Q_i}}}}{\htmlData{light=kv-cache}{{\htmlClass{vl-stream-k}{K_i^{\mathsf T}}}}}
in
QQ [B batch, 12 heads, 1, 64 d_h]
out
SS [B batch, 12 heads, 1 queries, 1001 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))