fareedkhan-dev--train-llm-from-scratch/docs/foundations/transformer.md

Decoder-Only Transformer

This repo implements a GPT-style decoder-only Transformer. "Decoder-only" means:

• the model reads a single token sequence;
• each position can attend only to previous positions;
• the output at every position is a distribution over the next token.

It is the right architecture for autoregressive language modeling:

[ p_\theta(x_t \mid x_{<t}) ]

The forward pass

flowchart TD
    IDS["token ids (B, T)"] --> TOK["token embedding (B, T, C)"]
    IDS --> POS["position ids 0..T-1"]
    POS --> PE["position embedding (T, C)"]
    TOK --> SUM["token + position"]
    PE --> SUM
    SUM --> B1["Block 1"]
    B1 --> B2["Block 2"]
    B2 --> BN["... N blocks"]
    BN --> LN["final LayerNorm"]
    LN --> HEAD["lm_head Linear(C -> vocab)"]
    HEAD --> LOGITS["logits (B, T, V)"]

The implementation is in src/models/transformer.py:

self.token_embed = nn.Embedding(vocab_size, n_embed)
self.position_embed = nn.Embedding(context_length, n_embed)
self.attn_blocks = nn.ModuleList([
    Block(n_head, n_embed, context_length) for _ in range(N_BLOCKS)
])
self.layer_norm = nn.LayerNorm(n_embed)
self.lm_head = nn.Linear(n_embed, vocab_size)

Shape symbols used throughout the docs:

Symbol	Meaning
B	batch size
T	sequence length / context length
C	embedding width, n_embed
H	number of attention heads
D	head width, usually C / H
V	vocabulary size

Embeddings

Token ids are categorical. The embedding table is a learned lookup:

[ E_{\text{tok}} \in \mathbb{R}^{V \times C} ]

For token id (x_t), the token vector is:

[ e_t = E_{\text{tok}}[x_t] ]

The model also learns absolute position embeddings:

[ E_{\text{pos}} \in \mathbb{R}^{T_{\max} \times C} ]

The input to the first block is:

[ h_t^{(0)} = E_{\text{tok}}[x_t] + E_{\text{pos}}[t] ]

In code:

tok_embedding = self.token_embed(idx)
pos_embedding = self.position_embed(self.pos_idxs[:T])
return tok_embedding + pos_embedding

The position embedding is necessary because attention alone is permutation-equivariant: without position information, the model would not know whether a token occurred first, last, or in the middle.

Transformer block

Each block in src/models/transformer_block.py uses pre-norm residual structure:

x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))

Mathematically:

[ u = x + \text{MHA}(\text{LN}(x)) ]

[ y = u + \text{MLP}(\text{LN}(u)) ]

This gives each block two jobs:

• attention moves information across token positions;
• the MLP transforms each position independently.

Why residual connections matter

A residual block learns an update, not a full replacement:

[ y = x + f(x) ]

If a layer is not useful yet, it can learn a small update and let information pass through. This makes deep stacks trainable because gradients have a direct path backward through the addition.

Why LayerNorm appears before sublayers

LayerNorm normalizes each token vector across its feature dimension:

[ \text{LN}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta ]

For a token vector (x \in \mathbb{R}^{C}):

[ \mu = \frac{1}{C} \sum_{i=1}^{C} x_i ]

[ \sigma^2 = \frac{1}{C} \sum_{i=1}^{C} (x_i - \mu)^2 ]

This repo uses pre-norm (LN -> sublayer -> residual) instead of post-norm (sublayer -> residual -> LN). Pre-norm is common in GPT-like models because it tends to make deeper stacks easier to optimize.

MLP / feed-forward network

The block MLP in src/models/mlp.py is:

self.hidden = nn.Linear(n_embed, 4 * n_embed)
self.relu = nn.ReLU()
self.proj = nn.Linear(4 * n_embed, n_embed)

For each position independently:

[ \text{MLP}(x) = W_2 , \text{ReLU}(W_1 x + b_1) + b_2 ]

where:

• (W_1) expands from (C) to (4C);
• (W_2) projects from (4C) back to (C).

Attention lets tokens communicate. The MLP gives each token vector nonlinear compute after that communication.

Logits and the language-model head

After the final block and final norm:

[ z_t = W_{\text{lm}} h_t + b_{\text{lm}} ]

The result (z_t \in \mathbb{R}^{V}) is a vector of logits: one unnormalized score per token in the vocabulary.

The probability distribution comes from softmax:

[ p_\theta(x_{t+1}=i \mid x_{\leq t}) = \frac{\exp(z_{t,i})}{\sum_{j=1}^{V}\exp(z_{t,j})} ]

Parameter count intuition

Ignoring biases and norms, a rough per-block parameter estimate is:

[ \text{attention} \approx 4C^2 ]

because Q, K, V, and the output projection are each (C \times C).

[ \text{MLP} \approx 8C^2 ]

because (C \to 4C) and (4C \to C).

So each block is roughly:

[ 12C^2 ]

The embedding and LM head add:

[ V C + C V ]

This repo does not tie token embeddings and output embeddings, so the input embedding and lm_head are separate parameter matrices.

Repo-specific architecture notes

Choice	Repo implementation	Consequence
Absolute learned positions	nn.Embedding(context_length, n_embed)	Simple and readable; fixed max context.
Causal attention mask	lower-triangular buffer in each head	Prevents future-token leakage.
MLP activation	ReLU	Educationally simple; many production GPT models use GELU/SwiGLU variants.
Dropout	not present in base modules	Less code noise; regularization comes mostly from data and optimizer choices.
Weight tying	not used	Easier to read; more parameters than tied embeddings.
Post-training heads	use forward_hidden	Reward/value heads reuse the same backbone.

Next

The most important sublayer is attention. Continue to Attention, Masks & Heads.