Building a Simple Large Language Model from Scratch

Modern generative AI is built on the Transformer architecture introduced in the 2017 paper "Attention Is All You Need". While high-level libraries make it easy to call pre-trained models, implementing the architecture from raw tensors reveals the mathematical mechanics that power autoregressive generation.

This guide breaks down how to construct a causal, decoder-only transformer language model from first principles using Python and PyTorch.

Transformer Architecture & Data Flow

A decoder-only generative transformer predicts the next token in a sequence by passing token and position embeddings through stacked self-attention and feed-forward blocks:

[Raw Text Sequence]
         │
         ▼
[Character / Subword Tokenizer]
         │
         ▼
[Token IDs: (B, T)] ──┬──► [Token Embeddings: (B, T, C)]
                      │                 │
                      └──► [Positional Encodings: (T, C)]
                                        │
                                        ▼  (Element-wise Addition)
                           [Input Matrix: X_0 = (B, T, C)]
                                        │
                                        ▼
                   ┌──────────────────────────────────────────┐
                   │        Transformer Block (x N)           │
                   │                                          │
                   │  ┌────────────────────────────────────┐  │
                   │  │ LayerNorm ──► Causal Multi-Head    │  │
                   │  │               Self-Attention       │  │
                   │  └──────────────────┬─────────────────┘  │
                   │                     ▼                    │
                   │             Residual Add (+)             │
                   │                     │                    │
                   │  ┌──────────────────┴─────────────────┐  │
                   │  │ LayerNorm ──► Feed-Forward Network │  │
                   │  │               (Linear-ReLU-Linear) │  │
                   │  └──────────────────┬─────────────────┘  │
                   │                     ▼                    │
                   │             Residual Add (+)             │
                   └────────────────────┬─────────────────────┘
                                        │
                                        ▼
                           [Final Layer Normalization]
                                        │
                                        ▼
                             [Linear LM Head: (B, T, V)]
                                        │
                                        ▼
                          [Cross-Entropy Loss / Softmax]

Core Mathematical & Architectural Components

1. Tokenization and Learned Embeddings

• Text is split into integer IDs $x \in \mathbb{R}^{B \times T}$ (where $B$ is batch size and $T$ is sequence context length).

• A lookup table maps token indices into a continuous latent space:

$$\text{Token Embeddings} = E_{\text{token}}(x) \in \mathbb{R}^{B \times T \times d_{\text{model}}}$$

2. Positional Encodings

• Transformers process all tokens simultaneously, lacking inherent sequence ordering.

• Learned or sinusoidal position vectors $E_{\text{pos}} \in \mathbb{R}^{T \times d_{\text{model}}}$ are added to token embeddings:

$$X_0 = E_{\text{token}}(x) + E_{\text{pos}}$$

3. Causal Multi-Head Self-Attention

• The input is projected into Query ($Q$), Key ($K$), and Value ($V$) matrices.

• Scaled Dot-Product Attention computes relevance scores across positions:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$$

• Causal Masking ($M$): A lower-triangular mask sets upper-diagonal entries to $-\infty$, ensuring tokens cannot attend to future positions during autoregressive generation.

• Multi-Head Splitting: Multiple independent attention heads compute distinct subspace relationships in parallel before concatenating outputs.

4. Feed-Forward Networks, LayerNorm & Residuals

• MLP Block: Two linear layers with non-linear activation (ReLU or GELU) expand and contract channel dimensions:

$$\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2$$

• Pre-Layer Normalization & Skip Connections: Normalizing activations before sub-layers stabilizes gradient flow, enabling deep stacking.

5. Training & Next-Token Optimization

• The final representation passes through a linear projection (lm_head) to produce unnormalized vocabulary logits.

• The model optimizes standard Cross-Entropy Loss between predicted logits and target tokens shifted by one time step ($t+1$).

Core Tech Stack