GPT-2

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 sequence Training Prefill Decode

x[B, T, 768]

Prepare Q, K, Vprojectsplit heads Open →

Q_i[B, T, 64]K_i[B, T, 64]V_i[B, T, 64]

{\htmlData{light=a-head softmax}{A_i}}

=

\underset{\text{over keys}}{\mathrm{softmax}}

\Big(

{\htmlData{light=q-head split-heads}{{\htmlClass{vl-stream-q}{Q_i}}}}{\htmlData{light=k-head split-heads}{{\htmlClass{vl-stream-k}{K_i^{\top}}}}}

/\sqrt{64} + {\htmlClass{vl-stream-context}{M}}

\Big)

{\htmlData{light=head-mix attn-mix}{H_i}}

=

training only · identity at evaluation

\mathrm{dropout}_{p_{\mathrm{attn}}}

\big(

{\htmlData{light=a-head softmax}{A_i}}

\big)

\cdot

{\htmlData{light=v-head split-heads}{{\htmlClass{vl-stream-v}{V_i}}}}

H_i[B, T, 64]

Merge and projectmerge headsoutput projection Open →

Δ_attn[B, T, 768]

GPT-2

GPT-2 attention computation flow