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.
- in
- [B batch, 12 heads, T positions, 64 d_h]
- out
- [B batch, 12 heads, T queries, T keys]
Codemodel.py 91, 93–94
def multihead_attn(q, k, v): 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))