ML First Principles, Part II:
From Equations to Working Code

We will trace every single computation through a small network — forward and backward — with exact numbers. No hand-waving. When you finish this section, you should be able to compute gradients for any network on paper.

Network Setup

Step 1: Forward Pass — Hidden Layer

Compute pre-activation \(\mathbf{z}^{[1]} = \mathbf{W}^{[1]}\mathbf{x} + \mathbf{b}^{[1]}\):

Apply ReLU activation \(\mathbf{a}^{[1]} = \text{ReLU}(\mathbf{z}^{[1]})\):

Step 2: Forward Pass — Output Layer

Compute pre-activation and sigmoid:

Step 3: Compute Loss

Binary cross-entropy with \(y = 1\):

Step 4: Backward Pass — Output Layer Error

The key identity for sigmoid + BCE: \(\delta^{[2]} = \hat{y} - y\):

Step 5: Gradients for Output Weights

Step 6: Propagate Error to Hidden Layer

Apply ReLU derivative: since both \(z_1^{[1]} = 0.60 > 0\) and \(z_2^{[1]} = 0.75 > 0\), the ReLU derivative is 1 for both:

Step 7: Gradients for Hidden Weights

Step 8: Update All Weights (\(\eta = 0.5\))

Since \(\delta\) was negative (model under-predicted — \(\hat{y} = 0.76 < 1\)), all gradients are negative, so all weights increase. This makes the output larger, pushing \(\hat{y}\) closer to 1.

Let's trace a minimal transformer decoder through one forward pass, computing every matrix multiplication. We'll use a vocabulary of 5 tokens, embedding dimension \(d = 4\), one attention head, and a sequence of 3 tokens.

Setup

Vocabulary: {0: "the", 1: "cat", 2: "sat", 3: "on", 4: "mat"}.
Input sequence: ["the", "cat", "sat"] → token IDs [0, 1, 2].
Target: predict the next token at each position.

Step 1: Token Embedding

Embedding matrix \(\mathbf{E} \in \mathbb{R}^{5 \times 4}\) (lookup table):

Look up rows 0, 1, 2 to get \(\mathbf{X} \in \mathbb{R}^{3 \times 4}\):

Step 2: Positional Encoding (simplified)

Add position vectors (using simple learned encodings for clarity):

Step 3: Compute Q, K, V

Using projection matrices \(\mathbf{W}^Q, \mathbf{W}^K, \mathbf{W}^V \in \mathbb{R}^{4 \times 4}\). For simplicity, let them be identity-like with small perturbations (in practice they're learned):

Suppose after projection we get (rounded):

Step 4: Attention Scores

Compute \(\mathbf{S} = \mathbf{Q}\mathbf{K}^T / \sqrt{d_k}\) where \(d_k = 4\), so divide by 2:

For example, \(s_{00} = \mathbf{q}_0^T \mathbf{k}_0\) measures how much "the" (as query) attends to "the" (as key).

Step 5: Causal Mask

Apply the causal mask — token \(i\) can only attend to positions \(\leq i\):

After masking, the first row has only one valid score (position 0); second row has two valid scores; third row has all three. Then we apply softmax row-wise:

(Example values.) Row 1: "the" can only see itself → weight = 1.0. Row 2: "cat" attends roughly equally to "the" and "cat". Row 3: "sat" attends to all three tokens, with slightly more weight on itself.

Step 6: Weighted Sum of Values

Each row of the output is a weighted combination of value vectors, where the weights came from the attention computation. This is the "new representation" of each token, enriched by information from tokens it attended to.

Step 7: Residual + LayerNorm + FFN + Residual + LayerNorm

Step 8: Final Prediction

Project last layer output to vocabulary size and apply softmax:

For position 2 ("sat"), the model outputs a probability distribution over 5 tokens. If it assigns high probability to token 3 ("on"), that's correct — "the cat sat on."

From-Scratch in NumPy: 2-Layer Network

import numpy as np

# ── Data ──
X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])  # XOR inputs
y = np.array([[0], [1], [1], [0]])              # XOR outputs

# ── Initialize weights (He initialization) ──
np.random.seed(42)
W1 = np.random.randn(2, 4) * np.sqrt(2.0 / 2)   # (2 inputs → 4 hidden)
b1 = np.zeros((1, 4))
W2 = np.random.randn(4, 1) * np.sqrt(2.0 / 4)   # (4 hidden → 1 output)
b2 = np.zeros((1, 1))

def sigmoid(z):  return 1 / (1 + np.exp(-z))
def relu(z):     return np.maximum(0, z)

# ── Training loop ──
lr = 0.5
for epoch in range(10000):
    # Forward
    z1 = X @ W1 + b1            # (4, 4)
    a1 = relu(z1)               # (4, 4)
    z2 = a1 @ W2 + b2           # (4, 1)
    a2 = sigmoid(z2)            # (4, 1) — predictions

    # Loss (BCE)
    loss = -np.mean(y * np.log(a2 + 1e-8) + (1-y) * np.log(1-a2 + 1e-8))

    # Backward
    dz2 = a2 - y                         # (4, 1) — output error
    dW2 = a1.T @ dz2 / 4                # (4, 1) — avg over batch
    db2 = np.mean(dz2, axis=0, keepdims=True)
    da1 = dz2 @ W2.T                     # propagate back
    dz1 = da1 * (z1 > 0).astype(float)  # ReLU derivative
    dW1 = X.T @ dz1 / 4
    db1 = np.mean(dz1, axis=0, keepdims=True)

    # Update
    W2 -= lr * dW2;  b2 -= lr * db2
    W1 -= lr * dW1;  b1 -= lr * db1

    if epoch % 2000 == 0:
        print(f"Epoch {epoch}: loss = {loss:.4f}")

print("Predictions:", a2.round(2).flatten())
# Should output approximately: [0.02, 0.98, 0.98, 0.03]

Self-Attention from Scratch

import numpy as np

def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))  # numerical stability
    return e / e.sum(axis=axis, keepdims=True)

def self_attention(X, Wq, Wk, Wv, mask=None):
    """
    X:    (seq_len, d_model)
    Wq/k/v: (d_model, d_k)
    mask: (seq_len, seq_len) — causal mask
    """
    Q = X @ Wq                          # (seq_len, d_k)
    K = X @ Wk                          # (seq_len, d_k)
    V = X @ Wv                          # (seq_len, d_k)

    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)     # (seq_len, seq_len)

    if mask is not None:
        scores = scores + mask           # -inf for future positions

    weights = softmax(scores)            # (seq_len, seq_len)
    output  = weights @ V                # (seq_len, d_k)

    return output, weights

# ── Create causal mask ──
seq_len = 4
mask = np.triu(np.full((seq_len, seq_len), -np.inf), k=1)
#  [[  0, -inf, -inf, -inf],
#   [  0,    0, -inf, -inf],
#   [  0,    0,    0, -inf],
#   [  0,    0,    0,    0]]

Common Bugs & Fixes

Data Preparation

Scale: LLaMA-3 was trained on ~15 trillion tokens. GPT-4 is estimated at 10T+. For context, the English Wikipedia is ~4 billion tokens — a tiny fraction of training data.

Quality filtering matters enormously: The pipeline typically includes language detection, deduplication (exact and fuzzy via MinHash), removal of toxic/low-quality content, URL filtering, and classifier-based quality scoring. Chinchilla showed that data quality affects the scaling law constants — better data means fewer tokens needed.

Tokenization in Detail

Byte Pair Encoding (BPE) builds a vocabulary bottom-up:

Typical vocab sizes: GPT-2: 50,257. LLaMA-2: 32,000. GPT-4: ~100,000. Larger vocabularies mean fewer tokens per text (more efficient) but a larger embedding matrix.

Distributed Training

No single GPU has enough memory or compute for modern LLMs. Training is distributed across hundreds or thousands of GPUs using multiple parallelism strategies:

Mixed Precision Training

Modern training uses BF16 (bfloat16) for forward/backward computation and FP32 for the master copy of weights. BF16 has the same exponent range as FP32 (so no overflow) but lower precision — this is fine for gradients, and it halves memory and doubles throughput.

For a 7B parameter model: ~126 GB just for parameters + optimizer states — requiring at least 2× A100-80GB GPUs.

Autoregressive Generation

At inference time, a decoder-only LLM generates one token at a time. At each step:

The new token is appended to the context, and the process repeats until an end-of-sequence token is generated or a maximum length is reached.

Sampling Strategies

Temperature

\(T = 1\): standard softmax. \(T \to 0\): becomes argmax (greedy, deterministic). \(T > 1\): flattens the distribution (more random, more "creative"). \(T < 1\): sharpens the distribution (more focused on high-probability tokens).

Top-k Sampling

Only consider the \(k\) most probable tokens, redistribute probability among them:

Top-p (Nucleus) Sampling

Dynamically choose the smallest set of tokens whose cumulative probability exceeds \(p\):

Where \(x_{(i)}\) are tokens sorted by decreasing probability. Top-p adapts — for confident predictions it considers few tokens; for uncertain ones, many.

The KV Cache — Why Inference Is Fast

This reduces per-token computation from \(O(T \cdot d^2)\) to approximately \(O(d^2)\) (plus the attention computation which is \(O(T \cdot d)\)).

Memory cost: For each layer, the KV cache stores \(2 \times T \times d_k \times h\) values. For a 70B model with 128K context: this can be 10+ GB of memory just for the cache.

BERT's Masked Language Modeling

Where \(\mathcal{M}\) is the set of masked positions (randomly chosen 15% of tokens). For each masked position, the model sees all other tokens (bidirectionally) and predicts the original token.

Why GPT Won (for Now)

BERT excels at understanding but struggles with generation (you'd need to iteratively unmask tokens). GPT's autoregressive approach naturally generates text and, with sufficient scale, can also understand through in-context learning. The discovery that scaling decoder-only models leads to emergent capabilities (GPT-3) shifted the field toward the GPT paradigm. However, encoder models remain superior for many specialized tasks (search, classification, retrieval).

The Encoder-Decoder Hybrid: T5

The Core Problem: Overfitting

L2 Regularization (Weight Decay)

The gradient update becomes: \(w \leftarrow w - \eta(\nabla_w \mathcal{L} + \lambda w) = w(1 - \eta\lambda) - \eta\nabla_w\mathcal{L}\). The term \((1 - \eta\lambda)\) shrinks the weight by a small fraction each step — hence "weight decay." This penalizes large weights, preferring simpler models.

Dropout

During training, randomly set each neuron's output to 0 with probability \(p\) (typically 0.1–0.5):

The \(\frac{1}{1-p}\) scaling (inverted dropout) ensures expected values don't change at test time. Why it works: Forces the network to learn redundant representations — no neuron can rely on any other specific neuron being present. This effectively trains an ensemble of \(2^n\) different sub-networks.

Layer Normalization

Normalizes across the feature dimension for each individual example. Unlike batch normalization (which normalizes across the batch), LayerNorm doesn't depend on batch size and works well for transformers and variable-length sequences.

Pre-Norm vs Post-Norm

Post-Norm (original transformer): \(\text{LayerNorm}(x + \text{sublayer}(x))\) — applies norm after the residual. Requires careful learning rate warmup.
Pre-Norm (modern standard): \(x + \text{sublayer}(\text{LayerNorm}(x))\) — applies norm before the sublayer. More stable training, dominant in modern LLMs.

The Evolution of Gradient Descent

SGD with Momentum

\(\beta \approx 0.9\) means the velocity is an exponential moving average of recent gradients. If the gradient consistently points in one direction, momentum accelerates; if it oscillates, momentum dampens.

RMSProp — Adaptive Learning Rates

Parameters with large gradients get smaller effective learning rates; parameters with small gradients get larger ones. This is crucial when different parameters have very different gradient magnitudes.

Adam — The Standard

Combines momentum (first moment) with RMSProp (second moment), plus bias correction:

Default hyperparameters: \(\beta_1 = 0.9\), \(\beta_2 = 0.999\), \(\epsilon = 10^{-8}\), \(\eta = 3 \times 10^{-4}\).

AdamW — Decoupled Weight Decay

Standard Adam applies L2 regularization to the gradient before the adaptive scaling, which weakens its effect. AdamW applies weight decay directly to the parameters, which is mathematically cleaner and empirically better. AdamW is the standard optimizer for modern LLM training.

Learning Rate Scheduling

Most LLM training uses a combination of warmup and cosine decay:

Warmup (typically 1–5% of training) gradually increases the learning rate from near-zero to prevent early instability. Cosine decay smoothly reduces it to near-zero by the end of training.

How Word Arithmetic Works

The famous equation: \(\vec{\text{king}} - \vec{\text{man}} + \vec{\text{woman}} \approx \vec{\text{queen}}\)

Mathematically, this works because the training objective (predict context words) forces the dot product \(\mathbf{u}_o^T \mathbf{v}_c\) to capture co-occurrence statistics, and co-occurrence patterns reflect semantic and syntactic relationships. The result: semantic relationships become geometric relationships.

Static vs Contextual Embeddings

Contextual embeddings are what make modern NLP so powerful. Each token's representation is computed by the transformer as a function of the entire surrounding context. The same word gets different representations depending on meaning, part of speech, and role in the sentence.

Sentence and Document Embeddings

To get a single vector for an entire sentence, common approaches include: mean-pooling all token embeddings (simple but effective), using the [CLS] token embedding from BERT, or using specialized models trained with contrastive objectives (like Sentence-BERT, E5, or BGE). These are used for semantic search, clustering, and retrieval-augmented generation (RAG).

Attention Visualization

By examining attention weight matrices, we can see what each token "looks at." Early layers tend to show local attention patterns (nearby tokens). Middle layers often show syntactic patterns (subject attending to its verb). Late layers show more diffuse, task-specific patterns.

However, attention weights alone don't tell the full story — what matters is the combination of attention weights and value vectors. A head can have high attention weight on a token but project it through a near-zero value, effectively ignoring it.

Probing Classifiers

To test whether a model has learned a concept (e.g., part of speech, syntactic tree depth, entity type), we freeze the model and train a simple linear classifier on its hidden states. If a linear probe achieves high accuracy, the information is linearly encoded in the representation — meaning the model has learned to organize its internal representations so that this information is easily extractable.

Superposition — The Key Challenge

This is why individual neurons are often "polysemantic" — a single neuron might activate for both "academic citations" and "references in code." The true "features" are not individual neurons but directions in activation space, and there may be millions of them encoded in thousands of dimensions.

Circuits — How Computation Happens

A circuit is a subgraph of the model that implements a specific behavior. For example, "induction heads" are a two-layer circuit discovered in transformers:

Layer 1 head: Copies previous token information into the current position ("previous-token head").
Layer 2 head: Attends to positions where the previous token matches the current token ("induction head").

Together, they implement the pattern: "if the sequence [A][B]....[A] has occurred, predict [B] will come next." This is a fundamental mechanism for in-context learning.

Sparse Autoencoders (SAEs) — Finding Features

Recent work has shown that training sparse autoencoders on model activations can decompose polysemantic neurons into interpretable, monosemantic features. The autoencoder learns a dictionary of directions in activation space, where each direction corresponds to an interpretable concept (a specific topic, a syntactic role, a type of reasoning).

The Grand Picture

Every layer of this stack builds on the one below. The math of dot products enables attention. Attention enables contextual representations. Contextual representations enable next-token prediction. Next-token prediction at scale produces emergent understanding. Alignment steers that understanding toward being helpful.

Your Path to Reading Research Papers

Now that you have the mathematical foundations, here's how to approach a new research paper:

1. Read the abstract and introduction — understand the problem being solved and the claim being made.

2. Look at figures first — architecture diagrams and result tables give you the high-level picture before you dive into equations.

3. Map new notation to what you know — Most papers use slight variations of the same constructs: attention, FFN, residual connections, loss functions. Identify which building blocks they use.

4. Focus on what's new — Skip the background sections. Look for the one or two novel equations or architectural choices that distinguish this paper.

5. Check the ablations — The ablation study tells you which parts of the method actually matter. This is often more informative than the main results.

6. Ask "what would I do differently?" — This is how you develop research taste. Every paper makes tradeoffs — understanding them is the sign of deep comprehension.