GPT-2

All models, with previews →

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 sequenceTrainingPrefillDecode

Transformer block › Attention

Row-wise softmax

×

no parameterssoftmax(w)

Normalize each query’s masked scores across key positions into attention probabilities. Each row sums to one before attention dropout.

Ai[b,t,s]=exp⁡(S~i[b,t,s])∑u=01000exp⁡(S~i[b,t,u]){\htmlData{light=a-head softmax}{A_i}}[b,t,s]=\frac{\exp({\htmlData{light=attn-scores}{\widetilde S_i}}[b,t,s])}{\sum_{u=0}^{1000}\exp({\htmlData{light=attn-scores}{\widetilde S_i}}[b,t,u])}
in
S~\widetilde S [B batch, 12 heads, 1, 1001]
out
AA [B batch, 12 heads, 1, 1001]
Explanation

bb, ii, tt and ss index batch, head, query and key; uu runs over the keys.

The released code subtracts each row’s maximum before exponentiating (model.py 20–23). The shift cancels between numerator and denominator and keeps every exponent at or below zero, so exp⁡\exp cannot overflow.

Codemodel.py 20–23, 96–97
model.py 20–23
def softmax(x, axis=-1):
    x = x - tf.reduce_max(x, axis=axis, keepdims=True)
    ex = tf.exp(x)
    return ex / tf.reduce_sum(ex, axis=axis, keepdims=True)
model.py 96–97
        w = mask_attn_weights(w)
        w = softmax(w)
Sources1

Direct evidence

The released implementation subtracts the row maximum for stability, exponentiates, and divides by the row sum on the final axis. Attention calls it after masking.

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

Read code excerpt
def softmax(x, axis=-1):
    x = x - tf.reduce_max(x, axis=axis, keepdims=True)
    ex = tf.exp(x)
    return ex / tf.reduce_sum(ex, axis=axis, keepdims=True)

… lines 24–95 omitted …

        w = mask_attn_weights(w)
        w = softmax(w)