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

Keys by head

×

tensork

The feature axis is partitioned into 12 heads of width 64; positions remain separate.

shape
KiK_{i} head i [B batch, T positions, 64 d_h]KK all heads [B batch, 12 heads, T positions, 64 d_h]
Sources

Contextual evidence from Reshape into attention heads

The released split_heads function introduces a head axis and moves it before the position axis.

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

Read code excerpt
    def split_heads(x):

… lines 76–76 omitted …

        return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3])

Contextual evidence from Query-key dot product

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))