GPT‑2 (Tiny) From Scratch — End‑to‑End Explainer

1) Raw Text Dataset

• What: A plain .txt corpus (e.g., resumes, docs, notes).
• Why: Language modeling learns token sequences; more diverse/clean text helps the model learn syntax, style, and patterns.

2) Byte Pair Encoding (BPE) with GPT‑2 vocab + <|endoftext|>

• What: Convert raw text → token IDs using GPT‑2’s BPE (via tiktoken).
• Why BPE: It merges frequent byte/character pairs, balancing small vocabulary size with coverage across languages/symbols. Robust on noisy text.
• <|endoftext|>: Teaches sequence boundaries, improving coherence during generation and enabling clean stopping.

3) Sliding‑Window Dataset → (Inputs X, Targets Y)

• What: Create many training examples by slicing token IDs into windows of length T.
  • X = tokens [i : i+T]
  • Y = tokens [i+1 : i+T+1] (one‑step shifted)
• Why: Next‑token prediction is the training objective. Stride controls overlap → more overlap yields more samples (at the cost of redundancy).

4) Embeddings: Token + Position → Final Encoding X₀

• Token Embedding: Lookup table mapping token IDs → vectors of size d_model.
• Positional Embedding: Learned vectors indexed by position 0..T-1.
• Final Encoding: X₀ = token_emb + pos_emb, shape (B, T, d_model).
• Why positions: Self‑attention is permutation‑invariant; positions inject order so the model knows “who came first.”

5) Stack of N Decoder Blocks (GPT‑style)

Each block preserves shape (B, T, d_model) and contains:

5.1) LayerNorm (Pre‑Norm before Attention)

• Why: Normalizes features per token to stabilize gradients and improve training of deep stacks.

5.2) Masked Multi‑Head Self‑Attention

• Projections:
  Q = XW_Q, K = XW_K, V = XW_V, then split into h heads with per‑head dim d_k = d_model / h.
• Scaled Dot‑Product:
  Compute attention scores S = (Q Kᵀ) / √d_k. The √d_k keeps the variance of dot products in a healthy range so softmax isn’t saturated.
• Causal Mask: Add −∞ above the main diagonal so position t can only attend to positions ≤ t (no future peeking).
• Weights: A = softmax(S + mask)
• Context: C = A V (mix values per position according to attention weights)
• Multi‑Head Merge: Concat heads and apply output projection back to d_model.

5.3) Residual (“Shortcut”) Add

• Why: X₁ = X + AttnOut lets gradients flow through identity paths, stabilizes optimization, and preserves prior representations while adding refinements.

5.4) LayerNorm (Pre‑Norm before MLP)

• Why again: Normalize before the MLP for the same stability reasons (this is “pre‑norm” design used by GPT‑2+).

5.5) Feed‑Forward MLP (Channel Mixer)

• What: Linear(d → 4d) → GELU → Linear(4d → d) (4× is common; 8× in larger models).
• Why: Adds nonlinearity and feature‑space expansion so tokens can be enriched, enabling the network to model complex transformations beyond attention’s weighted averaging.

5.6) Residual Add

• What: X₂ = X₁ + MLPOut.
• Why: Same residual rationale; stack many blocks without vanishing/exploding gradients.

Shape invariance: Input and output of each block remain (B, T, d_model) so the number of tokens T stays the same before and after attention/MLP. We alter representations, not sequence length.

6) Final LayerNorm

• What: Normalize once more after the last block.
• Why: Empirically improves stability and final perplexity before projecting to vocabulary space.

7) LM Head → Logits (B, T, vocab)

• What: A linear map from hidden size d_model to vocab_size for each position.
• Why: Produces unnormalized scores (logits) for every next‑token candidate.

8) Training vs. Inference

• Training: Apply softmax to logits and compute cross‑entropy loss against targets Y (the shifted tokens). Backprop minimizes next‑token prediction error.
• Inference (Decoding):
  • Greedy: pick argmax at each step (deterministic, can be repetitive).
  • Temperature: divide logits by T (T>1 = more random, T<1 = more focused).
  • Top‑k: keep only the top‑k candidates before softmax, then sample.
  • Loop: Append sampled token to the sequence, crop to context length T, and repeat until <|endoftext|> or max tokens.

9) Why two LayerNorms per block (and another at the end)?

• Pre‑Norm before Attention: stabilizes attention computation.
• Pre‑Norm before MLP: stabilizes the wide, nonlinear transformation.
• Final LayerNorm: a last normalization pass before the LM head improves output calibration.

10) Why residuals (“shortcuts”)?

• Enable deep networks to train by providing identity gradient paths.
• Let each block refine rather than replace earlier features.
• Improve convergence speed and final quality.

11) Why scale by √d_k?

• Without scaling, Q·Kᵀ grows with d_k, pushing softmax into saturation → tiny gradients.
• Dividing by √d_k keeps logits in a numerically stable range so softmax is sensitive and trainable.

12) Why tokens “stay the same” through attention?

• The sequence length T is unchanged; attention mixes features within each position using weighted sums of others but does not insert/delete positions. Output remains (B, T, d_model).

13) Decoding and <|endoftext|>

• During generation, stop when <|endoftext|> is produced to finish a sample cleanly.
• Including it during training teaches the model where sequences naturally end.