GPT-2 Small · Scale and causal mask

Transformer block › Attention

Scale and causal mask

no parametersw = w * tf.rsqrt(tf.cast(v.shape[-1].value, w.dtype)); w = mask_attn_weights(w)

Scales raw similarities by the square root of head width and masks future key positions.

{\htmlData{light=attn-scores}{\widetilde S_i}}=\frac{{\htmlData{light=qk-dot}{S_i}}}{\sqrt{64}}+{\htmlClass{vl-stream-context}{M}}

in: $S$ [B batch, 12 heads, T queries, T keys]
out: $\widetilde S$ [B batch, 12 heads, T queries, T keys]

$M_{t,s}$ is 0 for $s\le t$ and the most negative value of the score dtype for $s>t$ , with $t$ and $s$ counted from the first cached position (masking_utils.py 600–603).

Explanation

For independent, zero-mean, unit-variance query and key features, a 64-term dot product has variance 64; dividing by $\sqrt{64}$ returns it to 1, so the scores keep the same scale at any head width. Learned features need not meet these assumptions.

The formula is the maintained eager path, which adds $M$ (modeling_gpt2.py 56–61). The released code overwrites each masked score with $-10^{10}$ : it computes w*b - 1e10*(1-b), where the code's b is 1 at permitted positions and 0 elsewhere (model.py 83–89). Softmax gives the masked entries zero weight in both, to floating-point precision.

Every position reads itself and everything before it, never anything after. Across 8 illustrative positions that is 36 of 64 pairs. The mask prevents future-token leakage while attention is computed across sequence positions in parallel. Drawn at 8 positions for legibility; this model's context is 1,024.

Codemodel.py 58–97 · 2 more files

model.py 58–66

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)

model.py 83–89

    def mask_attn_weights(w):
        # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
        _, _, nd, ns = shape_list(w)
        b = attention_mask(nd, ns, dtype=w.dtype)
        b = tf.reshape(b, [1, 1, nd, ns])
        w = w*b - tf.cast(1e10, w.dtype)*(1-b)
        return w

model.py 93–97

        w = tf.matmul(q, k, transpose_b=True)
        w = w * tf.rsqrt(tf.cast(v.shape[-1].value, w.dtype))

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

modeling_gpt2.py 56–61

        scaling = query.size(-1) ** -0.5

    attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling

    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask

masking_utils.py 600–603

    if mask is not None:
        min_dtype = torch.finfo(dtype).min
        # we need 0s where the tokens should be taken into account, and -inf otherwise (mask is already of boolean type)
        mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), min_dtype)

Sources3

Direct evidence

GPT-2 scales the per-head scores, then applies the lower-triangular mask with a large negative value at excluded positions.

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

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)

… lines 67–82 omitted …

    def mask_attn_weights(w):
        # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
        _, _, nd, ns = shape_list(w)
        b = attention_mask(nd, ns, dtype=w.dtype)
        b = tf.reshape(b, [1, 1, nd, ns])
        w = w*b - tf.cast(1e10, w.dtype)*(1-b)
        return w

… lines 90–92 omitted …

        w = tf.matmul(q, k, transpose_b=True)
        w = w * tf.rsqrt(tf.cast(v.shape[-1].value, w.dtype))

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

Direct evidence

The maintained eager path scales the scores and adds the precomputed attention mask.

transformers · huggingface/transformers / src/transformers/models/gpt2/modeling_gpt2.py
Pinned revision 71c6f699ac9b ↗ · lines 56–61

Read code excerpt

        scaling = query.size(-1) ** -0.5

    attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling

    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask

Direct evidence

The maintained path builds the eager causal mask as 0 at permitted positions and the dtype minimum elsewhere.

transformers mask · huggingface/transformers / src/transformers/masking_utils.py
Pinned revision 71c6f699ac9b ↗ · lines 600–603

Read code excerpt

    if mask is not None:
        min_dtype = torch.finfo(dtype).min
        # we need 0s where the tokens should be taken into account, and -inf otherwise (mask is already of boolean type)
        mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), min_dtype)

GPT-2

GPT-2 attention computation flow