AI Engineer Exhaustive
MediumPlusGPT-2

Causal Masked Attention

GPT-2

Medium

Attention with causal mask for autoregressive decoding

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

GPT-2 is a left-to-right language model. Equation (1) in the paper factorizes a sequence as (\prod_i p(s_i\mid s_1,\ldots,s_{i-1})). Self-attention, left unmodified, lets position (i) read keys and values from (j>i), which would leak the token being predicted and make the factorization invalid. A causal (autoregressive) mask forbids that read.

The 2019 paper never writes the mask tensor. It inherits decoder-only attention from GPT / Vaswani: each position may attend to itself and the past, never the future. The official graph is the precise source. attention_mask builds a 0/1 band; mask_attn_weights applies it to already-scaled scores with a large negative offset rather than a hard zero, so softmax underflow sends blocked mass to ~0.

The same mask is written so it still works when the key/value sequence is longer than the query sequence. That is the KV-cache case: one new query attends over cached past keys plus the new key.

How it works

After the scaled product (W = QK^\top / \sqrt{d_k}), (W) has shape [batch, heads, nd, ns] — destination length by source length. Information is defined to flow from source to destination (mask_attn_weights comment in src/model.py).

attention_mask(nd, ns) builds a boolean matrix (M\in{0,1}^{n_d\times n_s}) with ones in the lower triangle, “counting from the lower right corner”:

M_{i,j}=\mathbf{1}\!\left[i \ge j - n_s + n_d\right],\qquad i=0\ldots n_d-1,\; j=0\ldots n_s-1.

When (n_d=n_s=n) (full forward, no cache) this is the familiar (i\ge j) lower-triangular mask, including the diagonal so a token may see itself. When keys have been concatenated with a past of length (n_s-n_d), the extra columns on the left are the cached past; the inequality still allows every query to see all cached keys and the current-step prefix, and still blocks the future relative to that query.

Scores are then

W \leftarrow W\odot M - 10^{10}\,(1-M)

in the mask’s dtype (tf.cast(1e10, w.dtype) in the official code). Softmax is over the source axis. Finite but huge negatives are used instead of (-\infty) so TPU/GPU softmax stays defined.

During incremental decoding, attn concatenates past keys/values along the sequence axis before multihead_attn. Queries have length 1 (the new token); keys/values have length (\mathrm{past}+1). The mask formula with (n_d=1), (n_s=\mathrm{past}+1) yields a single row of ones: the new token may see the entire cache and itself. No extra “generation mask” is required.

The paper’s contribution here is only the modeling statement (autoregressive (p(s_i\mid s_{<i})), context 1024). The index arithmetic, the (10^{10}) offset, and the past-concat contract are official-code details.

Official code

src/model.py: attention_mask, the nested mask_attn_weights, and the past concat in attn. attention_mask’s docstring notes it matches tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd) but avoids a TPU-unfriendly path. sample.py drives generation using the present cache the attention block returns; it does not reimplement the mask.

Watch-outs

Sources