GPT-1 › Worked Example: Forward Pass Through a Tiny GPT

Before diving into pre-training, let's trace a concrete forward pass through a minimal GPT-1-like model to build intuition for how information flows. We'll use a tiny 2-layer model with 4-dimensional embeddings and a vocabulary of 6 tokens, small enough to reason about numerically.

Suppose our vocabulary is {the, cat, sat, on, mat, [EOS]} with token IDs 0 through 5. Our input is the sequence [the, cat, sat], corresponding to token IDs [0, 1, 2].

Step 1: Token and Position Embedding Lookup. We look up the 4-dimensional embedding vector for each token, and the 4-dimensional position embedding for each position, then add them:

$$\mathbf{e}_i = \text{TokenEmb}[x_i] + \text{PosEmb}[i]$$

where $x_i$ is the token ID at position $i$. Say the result is:

$$\begin{aligned} \mathbf{h}_0^{(0)} &= [0.2, -0.1, 0.5, 0.3] \quad \text{("the" at position 0)} \\ \mathbf{h}_1^{(0)} &= [0.8, 0.3, -0.2, 0.6] \quad \text{("cat" at position 1)} \\ \mathbf{h}_2^{(0)} &= [-0.1, 0.7, 0.4, -0.3] \quad \text{("sat" at position 2)} \end{aligned}$$

Step 2: Causal Self-Attention (Layer 1). We compute queries $\mathbf{Q} = \mathbf{H}^{(0)} W_Q$, keys $\mathbf{K} = \mathbf{H}^{(0)} W_K$, and values $\mathbf{V} = \mathbf{H}^{(0)} W_V$ where $\mathbf{H}^{(0)}$ is the $3 \times 4$ matrix of input embeddings. The attention scores are:

$$S_{ij} = \frac{\mathbf{q}_i \cdot \mathbf{k}_j}{\sqrt{d_k}}$$

where $d_k = 4$ (the head dimension). We then apply the causal mask: positions where $j > i$ are set to $-\infty$ before the softmax, so that token $i$ only attends to positions $0$ through $i$. After softmax, the attention weights for position 2 ("sat") might look like $[\alpha_{20}, \alpha_{21}, \alpha_{22}]$ with $\alpha_{20} + \alpha_{21} + \alpha_{22} = 1$. The output for position 2 is the weighted sum $\alpha_{20} \mathbf{v}_0 + \alpha_{21} \mathbf{v}_1 + \alpha_{22} \mathbf{v}_2$, which blends information from "the", "cat", and "sat" according to their relevance.

Step 3: Feed-Forward Network. The attention output passes through a two-layer feed-forward network with GELU activation. For a hidden dimension of 16 (4x expansion), this computes:

$$\text{FFN}(\mathbf{h}) = W_2 \cdot \text{GELU}(W_1 \mathbf{h} + \mathbf{b}_1) + \mathbf{b}_2$$

Each token is processed independently by the FFN, allowing the model to apply a non-linear transformation to each position's representation after the attention aggregation step.

Step 4: Repeat for Layer 2. The output of Layer 1 becomes the input to Layer 2, which applies another round of causal self-attention and feed-forward processing. After 2 layers, we apply a final LayerNorm and multiply by the weight-tied output matrix (which is the transpose of the token embedding matrix) to get logits over the vocabulary.

Step 5: Compute Loss. The model's logit vector at position 0 is used to predict the token at position 1 ("cat"), the logit vector at position 1 predicts position 2 ("sat"), and the logit vector at position 2 predicts the next token ("on", ID=3). We apply softmax to each logit vector and compute cross-entropy against the correct targets. The total loss is the average of these three per-position losses. Backpropagation then updates all parameters to increase the probability of the correct continuations.

The key insight from this trace is that the model performs $n - 1$ next-token predictions simultaneously in a single forward pass. This efficiency, called teacher forcing, is why autoregressive language models train much faster than they might seem: every token in every training sequence contributes a training signal.