GPT-2

All models, with their blocks →

Size
Published
OpenAI · Feb–Nov 2019
Line
Dense decoder
View
Architecture
Source

openai/gpt‑2 @ 9b63575, huggingface/transformers @ 71c6f69, configuration openai‑community/gpt2 @ 607a30d · checkpoint tensors not independently countedHow to read

Merge and project

Full sequenceTrainingPrefillDecode

Transformer block › Attention

Merge attention heads

×

no parametersmerge_heads(a)

Move the head axis beside the feature axis and concatenate 12 × 64 features into width 768.

[B,12,1,64][B,12,1,64]
→  [B,1,12,64]\to\;[B,1,12,64]
→  [B,1,768]\to\;[B,1,768]
in
HH [B batch, 12 heads, 1, 64 d_h]
out
HH [B batch, 1, 768 d]
Explanation

Heads are concatenated in order: head ii fills features 64i,…,64i+63{64i},\dots,{64i+63}.

Codemodel.py 79, 81
model.py 79
    def merge_heads(x):
model.py 81
        return merge_states(tf.transpose(x, [0, 2, 1, 3]))
Sources1

Direct evidence

The released merge_heads function reverses the head permutation and merges the final two axes.

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

Read code excerpt
    def merge_heads(x):

… lines 80–80 omitted …

        return merge_states(tf.transpose(x, [0, 2, 1, 3]))