GPT: Autoregressive Language Models

Autoregressive Factorization

The foundational equation of GPT is the chain rule decomposition of a joint sequence distribution. Given a sequence of tokens $x = (x_1, x_2, \ldots, x_T)$:

Explanation: The probability of the entire sequence is the product of conditional probabilities, each asking "given everything so far, how likely is this next token?"

Symbols:

• $P(x)$: joint probability of the full token sequence
• $x_t$: the token at position $t$ (an integer index into the vocabulary)
• $T$: total sequence length

The training objective is to maximize the log-likelihood across a corpus $\mathcal{D}$:

Explanation: We sum log-probabilities over every token position and every document in the corpus—every token is simultaneously a supervised training signal.

Symbols:

• $\theta$: all learnable parameters of the transformer
• $\mathcal{L}(\theta)$: log-likelihood objective (maximized, or equivalently, negative cross-entropy is minimized)
• $P_\theta$: the model's predicted distribution parameterized by $\theta$

Scaled Dot-Product Attention (Causal)

Each transformer layer applies causal self-attention. For queries $Q$, keys $K$, values $V$ derived from the same input, with a causal mask $M$:

Explanation: Each position computes a weighted average over all previous positions' values, where weights are proportional to query-key dot products scaled to prevent gradient vanishing in the softmax.

Symbols:

• $Q \in \mathbb{R}^{T \times d_k}$: query matrix (current positions ask "what do I need?")
• $K \in \mathbb{R}^{T \times d_k}$: key matrix (past positions broadcast "what I contain")
• $V \in \mathbb{R}^{T \times d_v}$: value matrix (actual information to aggregate)
• $d_k$: key/query dimension; dividing by $\sqrt{d_k}$ prevents dot products from growing too large
• $M$: causal mask matrix where $M_{ij} = -\infty$ for $j > i$ and $0$ otherwise, enforcing autoregressive order

Multi-Head Attention

GPT uses $H$ parallel attention heads to capture different relationship types simultaneously:

Explanation: Different heads learn different notions of "relevance"—one head might track syntactic agreement while another tracks coreference—and their outputs are concatenated then linearly projected.

Symbols:

• $W^Q_h, W^K_h \in \mathbb{R}^{d_\text{model} \times d_k}$: per-head projection matrices for queries and keys
• $W^V_h \in \mathbb{R}^{d_\text{model} \times d_v}$: per-head value projection
• $W^O \in \mathbb{R}^{H d_v \times d_\text{model}}$: output projection that merges all heads back to model dimension
• $H$: number of heads; typically $d_k = d_v = d_\text{model}/H$

GPT Block and Output

Each of the $N$ transformer blocks applies: LayerNorm → Causal Self-Attention → residual → LayerNorm → FFN → residual. After $N$ blocks, a linear head projects to vocabulary logits:

Explanation: The final hidden state at position $t$ is projected back into vocabulary space via the transposed embedding matrix (weight tying), giving a probability distribution over the next token.

Symbols:

• $h_t^{(N)} \in \mathbb{R}^{d_\text{model}}$: the hidden state at position $t$ after all $N$ layers
• $W_E \in \mathbb{R}^{|\mathcal{V}| \times d_\text{model}}$: token embedding matrix (shared with the output projection for parameter efficiency)
• $|\mathcal{V}|$: vocabulary size (e.g., 50,257 for GPT-2; ~100K for GPT-4)

PyTorch Implementation Sketch

```python import torch import torch.nn as nn import torch.nn.functional as F class CausalSelfAttention(nn.Module): def __init__(self, d_model: int, n_heads: int, max_seq_len: int): super().__init__() assert d_model % n_heads == 0 self.n_heads = n_heads self.d_k = d_model // n_heads

# Fused QKV projection for efficiency self.qkv = nn.Linear(d_model, 3 * d_model, bias=False) self.out_proj = nn.Linear(d_model, d_model, bias=False)

# Causal mask: upper-triangular = -inf mask = torch.triu(torch.ones(max_seq_len, max_seq_len), diagonal=1) self.register_buffer("mask", mask.masked_fill(mask.bool(), float("-inf"))) def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, C = x.shape

# Project and split into Q, K, V qkv = self.qkv(x).split(C, dim=-1) # each: (B, T, C)

# Reshape to (B, n_heads, T, d_k) Q, K, V = [t.view(B, T, self.n_heads, self.d_k).transpose(1, 2) for t in qkv]

# Scaled dot-product with causal mask scores = (Q @ K.transpose(-2, -1)) / (self.d_k ** 0.5) scores = scores + self.mask[:T, :T] # broadcast over B, n_heads attn = F.softmax(scores, dim=-1) out = attn @ V # (B, n_heads, T, d_k) out = out.transpose(1, 2).contiguous().view(B, T, C) return self.out_proj(out) class GPTBlock(nn.Module): def __init__(self, d_model: int, n_heads: int, ffn_mult: int = 4, max_seq_len: int = 1024): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.attn = CausalSelfAttention(d_model, n_heads, max_seq_len) self.ln2 = nn.LayerNorm(d_model) self.ffn = nn.Sequential( nn.Linear(d_model, ffn_mult d_model), nn.GELU(), nn.Linear(ffn_mult d_model, d_model), ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.ln1(x)) # residual + attention x = x + self.ffn(self.ln2(x)) # residual + FFN return x ```