Transformer block › Attention
Keys by head
×tensork
The feature axis is partitioned into 12 heads of width 64; positions remain separate.
- shape
- head i [B batch, 1, 64 d_h] all heads [B batch, 12 heads, 1, 64 d_h]
In decode this is the new token's key only; the derivation reads K for all 1,001 positions from the cache, the past followed by this one (model.py 105–108).
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 KV cache
Each call stacks only its own K and V into present (104); with past, the cached K and V are concatenated before the new ones for this call’s attention (105–108).
present · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 102–108
Read code excerpt
c = conv1d(x, 'c_attn', n_state*3)
q, k, v = map(split_heads, tf.split(c, 3, axis=2))
present = tf.stack([k, v], axis=1)
if past is not None:
pk, pv = tf.unstack(past, axis=1)
k = tf.concat([pk, k], axis=-2)
v = tf.concat([pv, v], axis=-2)Contextual evidence from KV cache
The sampling loop keeps past between calls: each step concatenates the call’s presents after it along the position axis (69); nothing is allocated ahead.
loop · openai/gpt-2 / src/sample.py
Pinned revision 9b63575ef427 ↗ · lines 62–72
Read code excerpt
def body(past, prev, output):
next_outputs = step(hparams, prev, past=past)
logits = next_outputs['logits'][:, -1, :] / tf.to_float(temperature)
logits = top_k_logits(logits, k=top_k)
logits = top_p_logits(logits, p=top_p)
samples = tf.multinomial(logits, num_samples=1, output_dtype=tf.int32)
return [
next_outputs['presents'] if past is None else tf.concat([past, next_outputs['presents']], axis=-2),
samples,
tf.concat([output, samples], axis=1)
]Contextual evidence from KV cache
attention_mask(nd, ns) keeps the lower triangle counted from the lower right; with one query (nd = 1) every position is kept.
mask · openai/gpt-2 / src/model.py
Pinned revision 9b63575ef427 ↗ · lines 58–66
Read code excerpt
def attention_mask(nd, ns, *, dtype):
"""1's in the lower triangle, counting from the lower right corner.
Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
"""
i = tf.range(nd)[:,None]
j = tf.range(ns)
m = i >= j - ns + nd
return tf.cast(m, dtype)