Token + Position Embedding

GPT-2

Easy

Combined token and learned positional embedding

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

Self-attention has no built-in order. GPT-2 injects order the same way GPT did: a learned positional embedding added to a learned token embedding. Section 2.3 says the context is raised from GPT’s 512 tokens to 1024 and the vocabulary to 50,257. It does not write the addition formula or mention sinusoids. The official model() function is unambiguous: gather token rows, gather position rows, add.

Positions are absolute indices into a table of length n_ctx, not Vaswani-style (\sin/\cos) features. When decoding with a KV cache, the position index is offset by the cached length so the new token does not reuse position 0.

How it works

Let (X\in{0,\ldots,V-1}^{B\times T}) be token IDs. Two parameter matrices:

W_{te}\in\mathbb{R}^{V\times d},\qquad W_{pe}\in\mathbb{R}^{C\times d}

with (V=) n_vocab (50,257), (C=) n_ctx (1024), (d=) n_embd (768 for the 117M row of Table 2; 1024 / 1280 / 1600 for the larger models). In code they are wte and wpe.

Let (p_t = t_0 + t) for (t=0\ldots T-1), where (t_0) is past_length (0 on a full forward). Then

h_{b,t} = (W_{te})_{X_{b,t}} + (W_{pe})_{p_t}.

That is the entire embedding step. No scale by (\sqrt{d}), no extra dropout in the public graph, no token-type embedding.

positions_for implements (p_t): past_length + range(T), tiled over batch. past_length is 0 if past is None, else the cached sequence axis of past. Out-of-range positions ((p_t \ge C)) are a user-level error; the released sampler is expected to keep context (\le 1024).

Initialization, from model(): wpe ~ (\mathcal{N}(0, 0.01^2)), wte ~ (\mathcal{N}(0, 0.02^2)). The paper does not state these stddevs; they are official-code details. After the Transformer stack, the same wte is reused as the softmax output projection (weight tying). That tying is part of the full forward, not of this add.

The paper specifies context 1024 and vocab 50,257 and that the backbone is a Transformer language model. Learned (vs sinusoidal) positions and the exact add-after-gather are inferred from the GPT lineage plus src/model.py.

Official code

src/model.py: wpe / wte allocation and

h = tf.gather(wte, X) + tf.gather(wpe, positions_for(X, past_length))

inside model. Helpers: positions_for, expand_tile. Defaults in default_hparams(): n_ctx=1024, n_embd=768. Tokenizer IDs come from src/encoder.py; embedding does not call the encoder itself.

Watch-outs

• Positions must be absolute in the document, not always 0..T-1. Forgetting past_length when using a cache assigns the new token position 0 and corrupts every later layer.
• (W_{pe}) has (C=1024) rows, not an unlimited sinusoid. Indices (\ge 1024) are invalid for released GPT-2.
• Adding after an extra (\sqrt{d}) scale (a Transformer-base habit from some libraries) disagrees with this checkpoint.
• Token and position tables have different init std (0.02 vs 0.01). Sharing one initializer is a training-time mismatch with the official graph.

Sources

• Paper: Language Models are Unsupervised Multitask Learners (Radford et al., 2019), §2.3, Table 2
• Code: openai/gpt-2 src/model.py (model, positions_for)