Ferhat Erata’s ML/AI Study Notes

Comprehensive notes covering machine learning fundamentals, optimization theory, math foundations, deep learning, sequence models, LLM training pipelines, and distributed systems.

Quick Navigation

Part 1: Learning Through Examples

1.1 What is Machine Learning?

Machine learning is about finding patterns in data. Instead of writing explicit rules, we:

1. Define a model with adjustable parameters (weights)
2. Define a loss function that measures how wrong we are
3. Use an algorithm to adjust parameters to reduce the loss

The goal: Find parameter values that make our model’s predictions match the real data as closely as possible.

Data → [Model with parameters w] → Predictions
                    ↓
           Compare to real answers
                    ↓
               Loss (error)
                    ↓
         Adjust w to reduce loss  ← Gradient descent!

1.2 Simple Linear Regression

The Problem

You have data about houses: square footage (x) and sale price (y). You want to predict price from size.

The Model: A Line

We assume price is (roughly) linear in size:

\[ \hat{y} = w \cdot x + b \]

where:

• \(x\) = input (square footage)
• \(\hat{y}\) = predicted price
• \(w\) = weight (slope of the line — how much price increases per sq ft)
• \(b\) = bias (intercept — base price)

Our job: Find the best values of \(w\) and \(b\).

Visualizing Different Choices

Price ($1000s)
    ↑
500 |                    * (3000, 500)
400 |               * (2500, 400)
350 |          * (2000, 350)
280 |     * (1500, 280)
200 | * (1000, 200)
    |________________________→ Size (sq ft)
      1000  1500  2000  2500  3000

Bad line (\(w = 0.05\), \(b = 150\)):

\[\hat{y} = 0.05 \cdot x + 150\]

The line is too flat!

Better line (\(w = 0.15\), \(b = 50\)):

\[\hat{y} = 0.15 \cdot x + 50\]

Much better!

Figure: Comparing a bad fit (left) with large residuals (MSE = 13,730) to a good fit (right) with minimized residuals (MSE = 130). The dashed lines show the errors — gradient descent minimizes the sum of squared errors.

The Loss Function: Mean Squared Error

How do we quantify “how wrong” a line is? Use Mean Squared Error (MSE):

\[ \text{Loss} = \frac{1}{N} \sum_{i=1}^{N} (y_i - \hat{y}_i)^2 = \frac{1}{N} \sum_{i=1}^{N} (y_i - (w \cdot x_i + b))^2 \]

Why squared?

• Positive: Errors don’t cancel out
• Penalizes big errors more: An error of 100 contributes \(100^2 = 10000\), not just 100
• Smooth: Has nice derivatives for optimization

For our “bad line”: \[ \text{Loss} = \frac{1}{5}(0^2 + 55^2 + 100^2 + 125^2 + 200^2) = \frac{1}{5}(0 + 3025 + 10000 + 15625 + 40000) = 13730 \]

For our “better line”: \[ \text{Loss} = \frac{1}{5}(0^2 + 5^2 + 0^2 + 25^2 + 0^2) = \frac{1}{5}(0 + 25 + 0 + 625 + 0) = 130 \]

Lower loss = better fit!

1.3 Gradient Descent: Finding the Best Line

The Optimization Problem

We want to find \((w^*, b^*)\) that minimize the loss:

\[ (w^*, b^*) = \arg\min_{w, b} \frac{1}{N} \sum_{i=1}^{N} (y_i - (w \cdot x_i + b))^2 \]

For simple linear regression, there’s a closed-form solution. But for neural networks, there isn’t — we need gradient descent.

Figure: 3D visualization of a loss landscape showing the optimization surface and gradient descent path.

The Key Idea

The loss is a function of \(w\) and \(b\). If we plot loss vs parameters, we get a surface:

Figure: The loss function forms a landscape over parameter space. Gradient descent finds the minimum (valley) by taking steps in the direction of steepest descent.

Gradient descent: Start somewhere, repeatedly take small steps “downhill” until you reach the bottom.

What is a Gradient? (Intuitive Explanation)

Before diving into the math, let’s build intuition about what a gradient actually is.

The core idea: A gradient is a vector that points in the direction of steepest increase of a function.

Physical analogy: Imagine standing on a hilly terrain. The gradient at your location is like an arrow pointing directly uphill—the direction you’d go if you wanted to climb as steeply as possible. If you want to descend to the valley (minimize elevation), you walk in the opposite direction of this arrow.

Why is it a vector? Because we have multiple parameters to adjust! If we have two parameters (\(w\) and \(b\)), the gradient has two components:

• The first component tells us: “How much does the loss change if I nudge \(w\) a tiny bit?”
• The second component tells us: “How much does the loss change if I nudge \(b\) a tiny bit?”

Together, these form a vector that points toward the steepest uphill direction in the 2D parameter space.

Mathematical definition: The gradient of a function \(f(w, b)\) is the vector of all its partial derivatives:

\[\nabla f = \begin{bmatrix} \frac{\partial f}{\partial w} \\ \frac{\partial f}{\partial b} \end{bmatrix}\]

Each partial derivative measures the “sensitivity” of the output to changes in one input while holding others fixed.

Key insight for optimization: Since the gradient points uphill, we go the opposite direction to minimize the loss: \[\text{new parameters} = \text{old parameters} - \alpha \cdot \nabla \text{Loss}\]

where \(\alpha\) is the learning rate (step size).

The Gradient: Which Way is Downhill?

The gradient tells us the direction of steepest uphill. So we go the opposite direction!

For our linear regression:

\[ \frac{\partial \text{Loss}}{\partial w} = \frac{1}{N} \sum_{i=1}^{N} -2x_i(y_i - (w \cdot x_i + b)) = -\frac{2}{N} \sum_{i=1}^{N} x_i(y_i - \hat{y}_i) \]

\[ \frac{\partial \text{Loss}}{\partial b} = \frac{1}{N} \sum_{i=1}^{N} -2(y_i - (w \cdot x_i + b)) = -\frac{2}{N} \sum_{i=1}^{N} (y_i - \hat{y}_i) \]

Intuition:

• If predictions are too low (\(y_i - \hat{y}_i > 0\)), we need to increase \(w\) and \(b\)
• The gradient is negative → subtracting it makes \(w\) and \(b\) bigger ✓

The Update Rule

\[ w_{\text{new}} = w_{\text{old}} - \alpha \cdot \frac{\partial \text{Loss}}{\partial w} \]

\[ b_{\text{new}} = b_{\text{old}} - \alpha \cdot \frac{\partial \text{Loss}}{\partial b} \]

where \(\alpha\) is the learning rate (step size).

Concrete Example: Step by Step

Let’s use our housing data: \(x = [1000, 1500, 2000, 2500, 3000]\), \(y = [200, 280, 350, 400, 500]\)

Initialize: \(w = 0.0\), \(b = 0.0\), \(\alpha = 0.0000001\) (tiny because our \(x\) values are large)

Step 1: Compute predictions \[ \hat{y} = [0, 0, 0, 0, 0] \]

Step 2: Compute errors \[ y - \hat{y} = [200, 280, 350, 400, 500] \]

Step 3: Compute gradients \[ \frac{\partial L}{\partial w} = -\frac{2}{5}(1000 \cdot 200 + 1500 \cdot 280 + 2000 \cdot 350 + 2500 \cdot 400 + 3000 \cdot 500) \]

\[ = -\frac{2}{5}(200000 + 420000 + 700000 + 1000000 + 1500000) = -\frac{2}{5}(3820000) = -1528000 \]

\[ \frac{\partial L}{\partial b} = -\frac{2}{5}(200 + 280 + 350 + 400 + 500) = -\frac{2}{5}(1730) = -692 \]

Step 4: Update parameters \[ w = 0 - 0.0000001 \cdot (-1528000) = 0.1528 \]

\[ b = 0 - 0.0000001 \cdot (-692) = 0.0000692 \]

After just one step, \(w \approx 0.15\) — already close to the good value!

Repeat until loss stops decreasing.

Python Implementation

import numpy as np

# Data
X = np.array([1000, 1500, 2000, 2500, 3000])
y = np.array([200, 280, 350, 400, 500])

# Initialize parameters
w = 0.0
b = 0.0
learning_rate = 0.0000001
n = len(X)

# Gradient descent
for step in range(1000):
    # Predictions
    y_pred = w * X + b
    
    # Loss (MSE)
    loss = np.mean((y - y_pred) ** 2)
    
    # Gradients
    dw = -2/n * np.sum(X * (y - y_pred))
    db = -2/n * np.sum(y - y_pred)
    
    # Update
    w = w - learning_rate * dw
    b = b - learning_rate * db
    
    if step % 100 == 0:
        print(f"Step {step}: w={w:.4f}, b={b:.4f}, loss={loss:.2f}")

# Final: Step 900: w=0.1480, b=53.7000, loss=133.83

1.4 Logistic Regression: Classification

The Problem: Binary Classification

Now instead of predicting a continuous value, we want to classify: Is this email spam or not?

Why Linear Regression Fails

If we use \(\hat{y} = w_1 x_1 + w_2 x_2 + w_3 x_3 + b\), we might get outputs like \(-0.5\) or \(2.3\) — but we need probabilities between 0 and 1!

The Sigmoid Function

We “squash” the linear output through the sigmoid function:

\[ \sigma(z) = \frac{1}{1 + e^{-z}} \]

Figure: The sigmoid function σ(z) = 1/(1+e^(-z)) squashes any input to a value between 0 and 1.

Properties:

• Always outputs between 0 and 1 ✓
• \(\sigma(0) = 0.5\)
• \(\sigma(\text{large positive}) \approx 1\)
• \(\sigma(\text{large negative}) \approx 0\)

Deriving the Sigmoid Derivative (Chain Rule Example)

Understanding the sigmoid derivative is crucial for backpropagation. Let’s derive it step-by-step using the chain rule.

Starting point: \[\sigma(z) = \frac{1}{1 + e^{-z}} = (1 + e^{-z})^{-1}\]

Step 1: Apply the chain rule

Let \(u = 1 + e^{-z}\), so \(\sigma = u^{-1}\)

We need \(\frac{d\sigma}{dz} = \frac{d\sigma}{du} \cdot \frac{du}{dz}\)

Step 2: Compute each part

• Power rule: \(\frac{d\sigma}{du} = \frac{d}{du}(u^{-1}) = -u^{-2} = -\frac{1}{(1+e^{-z})^2}\)

• Chain rule on exponential: \(\frac{du}{dz} = \frac{d}{dz}(1 + e^{-z}) = -e^{-z}\)

Step 3: Multiply

\[\frac{d\sigma}{dz} = -\frac{1}{(1+e^{-z})^2} \cdot (-e^{-z}) = \frac{e^{-z}}{(1+e^{-z})^2}\]

Step 4: Simplify to the elegant form

Notice that: \[\frac{e^{-z}}{(1+e^{-z})^2} = \frac{1}{1+e^{-z}} \cdot \frac{e^{-z}}{1+e^{-z}} = \sigma(z) \cdot \frac{e^{-z}}{1+e^{-z}}\]

And \(\frac{e^{-z}}{1+e^{-z}} = \frac{1+e^{-z}-1}{1+e^{-z}} = 1 - \frac{1}{1+e^{-z}} = 1 - \sigma(z)\)

Final result: \[\boxed{\sigma'(z) = \sigma(z)(1 - \sigma(z))}\]

This is remarkably elegant — the derivative depends only on the output, not the input!

Why the Maximum Derivative is 1/4 (and Why It Matters)

The sigmoid derivative \(\sigma'(z) = \sigma(z)(1 - \sigma(z))\) has a maximum value of 0.25. Let’s prove this and understand its implications.

Figure: The sigmoid function and its derivative. The derivative reaches its maximum of 0.25 at z=0, where σ(z)=0.5.

Step 1: Reframe the problem

Let \(y = \sigma(z)\). Since sigmoid outputs are in \((0, 1)\), we need to find the maximum of: \[f(y) = y(1 - y) = y - y^2 \quad \text{for } y \in (0, 1)\]

Step 2: Find the maximum

This is a downward-opening parabola. Taking the derivative: \[\frac{df}{dy} = 1 - 2y\]

Setting to zero: \(1 - 2y = 0 \Rightarrow y = 0.5\)

Step 3: Calculate the maximum value

\[f(0.5) = 0.5 \times (1 - 0.5) = 0.5 \times 0.5 = \boxed{0.25 = \frac{1}{4}}\]

What this means: At the optimal point (\(z = 0\), where \(\sigma(z) = 0.5\)), the gradient through a sigmoid is only 0.25. In the saturation regions (where \(\sigma \to 0\) or \(\sigma \to 1\)), the derivative approaches 0.

The Vanishing Gradient Problem

This maximum of 1/4 has catastrophic implications for deep networks:

During backpropagation, gradients are multiplied at each layer: \[\frac{\partial L}{\partial W_1} = \frac{\partial L}{\partial z_n} \cdot \sigma'(z_{n-1}) \cdot \sigma'(z_{n-2}) \cdots \sigma'(z_1)\]

Even in the best case (all neurons at \(\sigma = 0.5\)):

In practice, it’s worse: Neurons rarely sit at \(\sigma = 0.5\). In saturation regions, \(\sigma' \approx 0\), making gradients vanish even faster.

Symptoms of vanishing gradients: - Early layers (near input) barely update - Loss decreases very slowly - Deep networks fail to train

Figure: Gradient magnitude decay during backpropagation for different activation functions. Sigmoid’s max derivative of 0.25 causes exponential decay — after 10 layers, gradients shrink to 10⁻⁶ (best case). ReLU maintains gradient = 1 for positive activations, enabling training of very deep networks. The bar chart shows how quickly sigmoid gradients become negligible.

Figure: Comparison of sigmoid and ReLU derivatives. ReLU has gradient = 1 for positive inputs, avoiding exponential decay.

ReLU: The Solution to Vanishing Gradients

ReLU (Rectified Linear Unit) solves this problem:

\[\text{ReLU}(z) = \max(0, z)\]

ReLU derivative: \[\text{ReLU}'(z) = \begin{cases} 1 & \text{if } z > 0 \\ 0 & \text{if } z \leq 0 \end{cases}\]

Why ReLU works:

The key insight: With ReLU, gradients pass through unchanged (multiplied by 1) for positive activations. There’s no exponential decay, so deep networks can actually train!

Trade-off: ReLU has “dead neurons” — if a neuron’s input is always negative, its gradient is always 0 and it never updates. Solutions include: - Leaky ReLU: Small gradient for negative inputs (\(0.01x\) instead of \(0\)) - ELU/GELU: Smooth alternatives with better properties

Interview Q: “Why do we use ReLU instead of sigmoid in hidden layers?”

A: Sigmoid has a maximum derivative of 0.25, causing vanishing gradients in deep networks. After just 10 layers, even in the best case, gradients shrink to \(0.25^{10} \approx 10^{-6}\). Early layers barely learn. ReLU has derivative = 1 for positive inputs, so gradients pass through without decay. This enables training of deep networks. The trade-off is “dead neurons” (gradient = 0 for negative inputs), addressed by variants like Leaky ReLU.

Logistic Regression Model

\[ P(y=1|x) = \sigma(w^\top x + b) = \frac{1}{1 + e^{-(w^\top x + b)}} \]

Interpretation: “The probability this email is spam, given its features.”

The Loss Function: Binary Cross-Entropy

For classification, we use cross-entropy loss (not MSE):

\[ \text{Loss} = -\frac{1}{N}\sum_{i=1}^{N} \left[ y_i \log(\hat{y}_i) + (1 - y_i) \log(1 - \hat{y}_i) \right] \]

Why this?

• If true label is 1 and we predict 1: \(-\log(1) = 0\) (no penalty) ✓
• If true label is 1 and we predict 0.01: \(-\log(0.01) \approx 4.6\) (big penalty!) ✓
• Works with probability outputs from sigmoid

Deriving Cross-Entropy from Maximum Likelihood

The cross-entropy loss isn’t arbitrary—it comes directly from Maximum Likelihood Estimation (MLE). Understanding this derivation reveals why cross-entropy is the natural loss for classification.

Step 1: Model the output as a probability

In logistic regression, we model the probability that \(y = 1\) given input \(x\):

\[P(y = 1 | x) = \hat{y} = \sigma(w^\top x + b)\]

This means \(P(y = 0 | x) = 1 - \hat{y}\).

Step 2: Write the likelihood of one data point

For a single example \((x_i, y_i)\) where \(y_i \in \{0, 1\}\):

\[P(y_i | x_i) = \hat{y}_i^{y_i} \cdot (1 - \hat{y}_i)^{1 - y_i}\]

This elegant formula works because:

• If \(y_i = 1\): \(P = \hat{y}_i^1 \cdot (1 - \hat{y}_i)^0 = \hat{y}_i\) ✓
• If \(y_i = 0\): \(P = \hat{y}_i^0 \cdot (1 - \hat{y}_i)^1 = 1 - \hat{y}_i\) ✓

Step 3: Write the likelihood of all data (assuming independence)

\[\mathcal{L}(w, b) = \prod_{i=1}^{N} P(y_i | x_i) = \prod_{i=1}^{N} \hat{y}_i^{y_i} \cdot (1 - \hat{y}_i)^{1 - y_i}\]

Step 4: Take the log (log-likelihood)

Products are numerically unstable and hard to optimize. Taking the log converts products to sums:

\[\log \mathcal{L} = \sum_{i=1}^{N} \left[ y_i \log(\hat{y}_i) + (1 - y_i) \log(1 - \hat{y}_i) \right]\]

Step 5: Maximize log-likelihood = Minimize negative log-likelihood

We want to maximize the likelihood, but optimizers minimize. So we minimize the negative log-likelihood (NLL):

\[\text{NLL} = -\log \mathcal{L} = -\sum_{i=1}^{N} \left[ y_i \log(\hat{y}_i) + (1 - y_i) \log(1 - \hat{y}_i) \right]\]

Dividing by \(N\) for the average gives us exactly the binary cross-entropy loss!

\[\boxed{\text{BCE} = -\frac{1}{N}\sum_{i=1}^{N} \left[ y_i \log(\hat{y}_i) + (1 - y_i) \log(1 - \hat{y}_i) \right]}\]

Key insight: Cross-entropy is the principled loss for classification because it directly maximizes the probability of the correct labels under our model. This connection to MLE also explains:

• Why it gives clean gradients: \(\frac{\partial \text{BCE}}{\partial z} = \hat{y} - y\)
• Why it’s related to information theory: \(\text{BCE}(p, q) = H(p) + D_{KL}(p || q)\)
• Why it’s the natural choice for probabilistic outputs

Gradient for Logistic Regression

The gradient has a beautiful form:

\[ \frac{\partial \text{Loss}}{\partial w} = \frac{1}{N} \sum_{i=1}^{N} (\hat{y}_i - y_i) x_i \]

Just the error times the input — same form as linear regression!

Note: This gradient computes the average over all examples in the batch: dw = X.T @ error / len(y). The division by len(y) ensures the gradient magnitude doesn’t depend on batch size.

Python Implementation

import numpy as np

def sigmoid(z):
    return 1 / (1 + np.exp(-z))

# Data: [contains_FREE, contains_!, has_link]
X = np.array([[1, 1, 1],
              [0, 0, 0],
              [1, 1, 0],
              [0, 1, 1],
              [1, 0, 1]])
y = np.array([1, 0, 1, 0, 1])

# Initialize
w = np.zeros(3)
b = 0.0
lr = 0.5

for step in range(1000):
    # Forward pass
    z = X @ w + b
    y_pred = sigmoid(z)
    
    # Loss (cross-entropy)
    eps = 1e-15  # avoid log(0)
    loss = -np.mean(y * np.log(y_pred + eps) + (1 - y) * np.log(1 - y_pred + eps))
    
    # Gradients
    error = y_pred - y
    dw = X.T @ error / len(y)
    db = np.mean(error)
    
    # Update
    w = w - lr * dw
    b = b - lr * db
    
    if step % 200 == 0:
        predictions = (y_pred > 0.5).astype(int)
        accuracy = np.mean(predictions == y)
        print(f"Step {step}: loss={loss:.4f}, accuracy={accuracy:.0%}")

# Final: w=[1.23, 0.41, 0.41], b=-0.82, accuracy=100%

1.5 Multi-Layer Perceptron (MLP)

Limitation of Linear Models

Linear and logistic regression can only learn linear decision boundaries:

Linearly separable:         NOT separable (XOR):

  x₂↑                         x₂↑
    |  + + +                    |  -     +
    | + + +                     |
----+--------→ x₁          -----+-----→ x₁
    | - - -                     |
    |  - -                      |  +     -

To learn complex patterns, we need non-linear models.

The MLP: Stacking Layers

An MLP adds hidden layers between input and output:

Figure: Multi-layer perceptron with input layer, hidden layer, and output layer. Each connection represents a weight.

Forward Pass (Step by Step)

For a network with one hidden layer:

Layer 1: Linear transformation + activation \[ z_1 = W_1 x + b_1 \]

\[ h = \text{ReLU}(z_1) = \max(0, z_1) \]

Layer 2: Another linear transformation \[ z_2 = W_2 h + b_2 \]

\[ \hat{y} = \text{softmax}(z_2) \quad \text{(for classification)} \]

Activation Functions: Why We Need Them

Without activation functions, stacking linear layers is useless:

\[ W_2(W_1 x + b_1) + b_2 = W_2 W_1 x + W_2 b_1 + b_2 = W' x + b' \]

Still linear! Activation functions introduce non-linearity.

Common Activation Functions

Figure: Common activation functions including ReLU, Sigmoid, Tanh, and their variants.

ReLU: Why It Works

• Simple: Fast to compute
• Sparse: Many neurons output 0 (efficient)
• No vanishing gradient: Gradient is 1 for positive inputs

Concrete Example: Learning XOR

XOR function:

No line can separate 0s from 1s!

MLP Solution

Hidden layer (2 neurons): \[ h_1 = \text{ReLU}(x_1 + x_2 - 0.5) \quad \text{(detects "at least one 1")} \]

\[ h_2 = \text{ReLU}(x_1 + x_2 - 1.5) \quad \text{(detects "both are 1")} \]

Output: \[ \hat{y} = h_1 - 2 \cdot h_2 \]

The hidden layer creates a new representation where the problem becomes linearly separable!

Figure: The power of hidden layers. Left: In the original input space, XOR is NOT linearly separable — no single line can separate the red class (0) from the blue class (1). Middle: The hidden layer computes new features \(h_1\) (“at least one 1”) and \(h_2\) (“both are 1”). Right: In the transformed space, the problem becomes linearly separable! This is the key insight behind deep learning: each layer transforms data into more useful representations.

Python Implementation

import numpy as np

def relu(x):
    return np.maximum(0, x)

def relu_derivative(x):
    return (x > 0).astype(float)

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

# XOR data
X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
y = np.array([[0], [1], [1], [0]])

# Network: 2 inputs → 4 hidden → 1 output
np.random.seed(42)
W1 = np.random.randn(2, 4) * 0.5
b1 = np.zeros((1, 4))
W2 = np.random.randn(4, 1) * 0.5
b2 = np.zeros((1, 1))

lr = 1.0

for step in range(10000):
    # Forward pass
    z1 = X @ W1 + b1
    h = relu(z1)
    z2 = h @ W2 + b2
    y_pred = sigmoid(z2)
    
    # Loss
    loss = np.mean((y - y_pred) ** 2)
    
    # Backward pass (will explain in next section!)
    d_z2 = (y_pred - y) * y_pred * (1 - y_pred)
    d_W2 = h.T @ d_z2
    d_b2 = np.sum(d_z2, axis=0, keepdims=True)
    
    d_h = d_z2 @ W2.T
    d_z1 = d_h * relu_derivative(z1)
    d_W1 = X.T @ d_z1
    d_b1 = np.sum(d_z1, axis=0, keepdims=True)
    
    # Update
    W2 -= lr * d_W2
    b2 -= lr * d_b2
    W1 -= lr * d_W1
    b1 -= lr * d_b1
    
    if step % 1000 == 0:
        print(f"Step {step}: loss={loss:.4f}")

# Test
print("\nPredictions:")
print(np.round(y_pred, 2))
# [[0.02], [0.98], [0.98], [0.02]] ✓

1.6 The Backpropagation Algorithm

The Problem: How to Get Gradients in Deep Networks?

For linear regression: straightforward calculus.

For deep networks with millions of parameters: we need an efficient algorithm.

Backpropagation = Backward propagation of errors

The Chain Rule: The Key Insight

If \(y = f(g(x))\), then:

\[ \frac{dy}{dx} = \frac{dy}{dg} \cdot \frac{dg}{dx} \]

In neural networks, we have a chain of operations:

\[ x \to z_1 \to h \to z_2 \to \hat{y} \to \text{Loss} \]

To find \(\frac{\partial \text{Loss}}{\partial W_1}\), we apply the chain rule through each step.

Visualizing Backpropagation: The Computational Graph

A computational graph makes backpropagation intuitive. Each node represents either a variable (data) or an operation. Edges show data flow.

Figure: A computational graph for a 2-layer MLP. Blue arrows show forward pass (data flowing left to right). Red arrows show backward pass (gradients flowing right to left). Each operation node computes a “local gradient” that gets multiplied along the path.

Key insight: During backprop, we traverse the same graph in reverse, multiplying local gradients along each path.

The algorithm:

1. Forward pass: Compute all intermediate values, storing them for later
2. Backward pass: Starting from the loss, compute gradients by:
   • For each operation node, multiply the incoming gradient by the local gradient
   • Sum gradients when paths merge (multiple outputs from one node)

Example with specific values:

Forward: x=2 → [×W₁=3] → z₁=6 → [ReLU] → h=6 → [×W₂=0.5] → z₂=3 → [σ] → ŷ=0.95

If y=1:  Loss = -log(0.95) = 0.05

Backward:
  ∂L/∂ŷ = -1/0.95 = -1.05
  ∂L/∂z₂ = -1.05 × σ'(3) = -1.05 × 0.95 × 0.05 = -0.05
  ∂L/∂W₂ = -0.05 × h = -0.05 × 6 = -0.30
  ∂L/∂h  = -0.05 × W₂ = -0.05 × 0.5 = -0.025
  ∂L/∂z₁ = -0.025 × 1 (ReLU'=1 since z₁>0) = -0.025
  ∂L/∂W₁ = -0.025 × x = -0.025 × 2 = -0.05

Why computational graphs matter: - Automatic differentiation frameworks (PyTorch, TensorFlow) build these graphs automatically - Any differentiable computation can be expressed as a graph - Backprop becomes mechanical: just apply the chain rule at each node

Backprop: Step by Step

Consider our 2-layer MLP:

Forward pass (compute and cache everything):

z₁ = W₁x + b₁
h = ReLU(z₁)
z₂ = W₂h + b₂
ŷ = sigmoid(z₂)
L = MSE(y, ŷ)

Backward pass (work backwards from loss):

1. Gradient of loss w.r.t. output: \[ \frac{\partial L}{\partial \hat{y}} = \hat{y} - y \]

2. Through sigmoid: The sigmoid derivative has an elegant form. Starting from \(\sigma(z) = \frac{1}{1+e^{-z}}\):

\[\sigma'(z) = \frac{d}{dz}\frac{1}{1+e^{-z}} = \frac{e^{-z}}{(1+e^{-z})^2} = \frac{1}{1+e^{-z}} \cdot \frac{e^{-z}}{1+e^{-z}} = \sigma(z)(1-\sigma(z))\]

Therefore: \[ \frac{\partial L}{\partial z_2} = \frac{\partial L}{\partial \hat{y}} \cdot \sigma'(z_2) = (\hat{y} - y) \cdot \hat{y}(1 - \hat{y}) \]

3. Gradient for \(W_2\): \[ \frac{\partial L}{\partial W_2} = h^\top \cdot \frac{\partial L}{\partial z_2} \]

4. Propagate to hidden layer: \[ \frac{\partial L}{\partial h} = \frac{\partial L}{\partial z_2} \cdot W_2^\top \]

5. Through ReLU: \[ \frac{\partial L}{\partial z_1} = \frac{\partial L}{\partial h} \cdot \mathbf{1}_{z_1 > 0} \]

6. Gradient for \(W_1\): \[ \frac{\partial L}{\partial W_1} = x^\top \cdot \frac{\partial L}{\partial z_1} \]

Why It’s Efficient

• Forward pass: \(O(\text{network size})\) — just matrix multiplications
• Backward pass: \(O(\text{network size})\) — same complexity as forward!
• Total: \(O(\text{network size})\) per example

Without backprop, computing each parameter’s gradient separately would be \(O((\text{network size})^2)\).

Automatic Differentiation

Modern frameworks (PyTorch, TensorFlow) implement backprop automatically:

import torch
import torch.nn as nn

# Define network
model = nn.Sequential(
    nn.Linear(2, 4),
    nn.ReLU(),
    nn.Linear(4, 1),
    nn.Sigmoid()
)

# Training
X = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=torch.float32)
y = torch.tensor([[0], [1], [1], [0]], dtype=torch.float32)

optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
loss_fn = nn.MSELoss()

for step in range(10000):
    # Forward
    y_pred = model(X)
    loss = loss_fn(y_pred, y)
    
    # Backward (automatically computes all gradients!)
    optimizer.zero_grad()
    loss.backward()
    
    # Update
    optimizer.step()
    
    if step % 1000 == 0:
        print(f"Step {step}: loss={loss.item():.4f}")

1.7 Softmax and Multi-class Classification

From Binary to Multi-class

Logistic regression handles 2 classes. For K classes, we use softmax:

\[P(y = k | x) = \frac{e^{z_k}}{\sum_{j=1}^{K} e^{z_j}}\]

where \(z_k = w_k \cdot x + b_k\) is the logit for class \(k\).

Properties of Softmax

1. Outputs sum to 1: \(\sum_k P(y = k | x) = 1\) ✓
2. All outputs positive: \(P(y = k | x) > 0\) ✓
3. Preserves ranking: Largest logit → highest probability

Why Exponential (Not Simple Normalization)?

Interview Question: “Why does softmax use \(e^{z_i}\) instead of just \(\frac{z_i}{\sum_j z_j}\)?”

Simple normalization would be: \(\text{normalize}(z_i) = \frac{z_i}{\sum_j z_j}\)

This fails for several critical reasons:

Problem 1: Negative Values Break Everything

Logits: [2, -3, 1]
Sum = 0 → Division by zero!

Logits: [2, -5, 1]  
Sum = -2
"Probabilities": [2/-2, -5/-2, 1/-2] = [-1, 2.5, -0.5]
→ Negative "probabilities"! Invalid!

Exponentials are always positive: \(e^x > 0\) for all \(x\), so softmax outputs are always valid probabilities.

Problem 2: Beautiful Gradient Properties

The softmax + cross-entropy gradient has an elegant form:

\[\frac{\partial L}{\partial z_i} = p_i - y_i\]

where \(p_i\) is predicted probability and \(y_i\) is the target (1 for correct class, 0 otherwise).

This clean gradient comes directly from the exponential. Simple normalization would give messier, harder-to-optimize gradients.

Problem 3: Amplification with Temperature Control

Exponentials amplify differences between logits:

Logits: [2.0, 1.0, 0.5]

Simple normalize: [2/3.5, 1/3.5, 0.5/3.5] = [0.57, 0.29, 0.14]
Softmax:          [0.59, 0.24, 0.17]  # Winner amplified

Temperature provides explicit control over this amplification:

\[\text{softmax}(z_i / T) = \frac{e^{z_i/T}}{\sum_j e^{z_j/T}}\]

• \(T \to 0\): Approaches argmax (one-hot, hard selection)
• \(T = 1\): Standard softmax
• \(T \to \infty\): Approaches uniform distribution

This is used in knowledge distillation (soft targets) and sampling diversity control.

Problem 4: Theoretical Grounding (Maximum Entropy)

Softmax is the maximum entropy distribution subject to linear constraints on expected features. This connects to:

• Statistical mechanics (Boltzmann distribution)
• Information theory (exponential family distributions)
• Principled probabilistic modeling

Summary: Exponentials ensure:

• ✅ Always positive outputs (valid probabilities)
• ✅ Sum to 1 (proper distribution)
  
• ✅ Clean gradients for efficient learning
• ✅ Controllable sharpness via temperature
• ✅ Theoretically principled (max entropy)

Example: 3-Class Classification

Input x → [Linear Layer] → Logits z → [Softmax] → Probabilities

z = [2.0, 1.0, 0.1]

softmax(z):
  e^2.0 = 7.39
  e^1.0 = 2.72  
  e^0.1 = 1.11
  sum = 11.22

  P = [7.39/11.22, 2.72/11.22, 1.11/11.22]
    = [0.66, 0.24, 0.10]

Figure: Effect of temperature on softmax distribution. Lower temperature makes the distribution sharper (more confident), higher temperature makes it more uniform.

Multi-class Cross-Entropy Loss

For one-hot label \(y\) (e.g., \(y = [0, 1, 0]\) for class 2):

\[\mathcal{L} = -\sum_{k=1}^{K} y_k \log(\hat{y}_k) = -\log(\hat{y}_c)\]

where \(c\) is the true class. We’re penalizing low probability on the correct class.

PyTorch Implementation

import torch.nn as nn

# Method 1: Separate softmax and loss
probs = nn.Softmax(dim=1)(logits)
loss = nn.NLLLoss()(torch.log(probs), labels)

# Method 2: Combined (numerically stable, preferred!)
loss = nn.CrossEntropyLoss()(logits, labels)  # Takes raw logits!

⚠️ Common mistake: CrossEntropyLoss expects logits, not probabilities!

1.8 Batching: Why and How

What This Means (For Beginners)

When training a neural network, three terms come up constantly: Epoch, Batch, and Iteration. Understanding how they relate is fundamental.

Epoch = One complete pass through the entire training dataset

Think of studying for an exam: - 1 epoch: You go through the entire syllabus once - 5 epochs: You study the same syllabus five times, each time understanding more

Batch = A subset of the training data processed together

Since the entire dataset is often too large to fit in GPU memory at once, we divide it into smaller chunks called batches. Each batch is processed in one forward pass and one backward pass.

Iteration = One forward + backward pass on one batch

Each iteration updates the model’s weights once.

The Key Formula

\[\text{Iterations per epoch} = \frac{\text{Total training samples}}{\text{Batch size}}\]

Worked Example

Suppose you have 1,000 training samples and set batch size = 100:

Total samples:     1,000
Batch size:        100
Iterations/epoch:  1,000 / 100 = 10

What happens in 1 epoch:
  Iteration 1:  Process samples 1-100    → Update weights
  Iteration 2:  Process samples 101-200  → Update weights
  Iteration 3:  Process samples 201-300  → Update weights
  ...
  Iteration 10: Process samples 901-1000 → Update weights

After 10 iterations, you've completed 1 epoch!

For 3 epochs with batch_size=100 on 1,000 samples: - Total iterations = 3 epochs × 10 iterations/epoch = 30 weight updates - Each sample is seen 3 times total (once per epoch)

Visual Summary

┌────────────────────────────────────────────────────────────────┐
│                     FULL DATASET (1000 samples)                │
├─────────┬─────────┬─────────┬─────────┬───────────┬────────────┤
│ Batch 1 │ Batch 2 │ Batch 3 │ Batch 4 │    ...    │  Batch 10  │
│ (100)   │ (100)   │ (100)   │ (100)   │           │   (100)    │
└─────────┴─────────┴─────────┴─────────┴───────────┴────────────┘
     ↓         ↓         ↓         ↓                      ↓
  Iter 1    Iter 2    Iter 3    Iter 4       ...      Iter 10
     
     └──────────────────────┬──────────────────────────┘
                            │
                      = 1 EPOCH

The Pizza Analogy 🍕

• The entire pizza = Full training dataset
• Each slice = One batch
• Eating one slice = One iteration (process one batch, update weights)
• Finishing the whole pizza = One epoch (processed all samples once)

If the pizza has 8 slices and you eat 1 slice at a time: - You need 8 iterations to finish 1 pizza (1 epoch) - Training for 3 epochs = eating 3 pizzas = 24 slices eaten = 24 iterations

Why Multiple Epochs?

Training for a single epoch usually isn’t enough: - The model makes one pass through each sample - Weights are updated based on each batch, but may not have converged - By repeating for multiple epochs, the model gradually refines its understanding

Epoch 1: Model learns rough patterns
Epoch 2: Model refines understanding  
Epoch 3: Model fine-tunes details
...
Epoch N: Model converges (loss stops decreasing)

Warning: Too many epochs can lead to overfitting! The model memorizes the training data instead of learning generalizable patterns. This is why we monitor validation loss.

Interview Q: “What’s the difference between epoch, batch, and iteration?”

A: An epoch is one complete pass through the entire training dataset. A batch is a subset of the data processed together in one forward/backward pass. An iteration is one weight update, which happens after processing one batch.

The relationship: Iterations per epoch = Dataset size / Batch size

For example, with 10,000 samples and batch size 100, you have 100 iterations per epoch. Training for 5 epochs means 500 total weight updates, with each sample seen 5 times.

The Problem with Single Examples

Computing gradient on one example at a time:

• Very noisy updates
• Can’t utilize GPU parallelism
• Slow convergence

Computing gradient on entire dataset:

• Very stable but slow
• One update per epoch
• Memory can’t hold all data

Mini-batch: The Sweet Spot

Dataset: 10,000 examples
Batch size: 32

→ 10,000 / 32 = 312 batches per epoch
→ 312 gradient updates per epoch

Why Batching Works

The mini-batch gradient is an unbiased estimator of the full gradient:

\[\mathbb{E}\left[\frac{1}{B}\sum_{i \in \text{batch}} \nabla \ell_i\right] = \frac{1}{N}\sum_{i=1}^{N} \nabla \ell_i\]

Batch Size Tradeoffs

The “Noise is Good” Insight

Small batch noise acts as implicit regularization:

• Helps escape sharp minima (which generalize poorly)
• Finds flatter minima (which generalize better)

Figure: Sharp minima (left) overfit because small weight changes cause large loss changes. Flat minima (right) generalize better because they’re robust to perturbations.

Practical Guidelines

Code Example: Batching in PyTorch

from torch.utils.data import DataLoader, TensorDataset

# Create dataset and loader
dataset = TensorDataset(X_tensor, y_tensor)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# Training loop
for epoch in range(num_epochs):
    for batch_X, batch_y in loader:
        # Forward pass on batch
        predictions = model(batch_X)
        loss = criterion(predictions, batch_y)
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

1.9 Data Preprocessing

Why Preprocess?

Raw features often have different scales:

• Age: 0-100
• Salary: 20,000-500,000
• Height: 1.5-2.0 meters

Without preprocessing:

• Large-scale features dominate gradients
• Small learning rates needed for some features
• Optimization is slow and unstable

Standardization (Z-score Normalization)

\[x' = \frac{x - \mu}{\sigma}\]

Result: Mean = 0, Std = 1

# Training set
mean = X_train.mean(axis=0)
std = X_train.std(axis=0)

# Apply to all sets using TRAINING statistics!
X_train_scaled = (X_train - mean) / std
X_test_scaled = (X_test - mean) / std  # Same mean/std!

⚠️ Critical: Always compute statistics on training data only, then apply to test data!

Min-Max Normalization

\[x' = \frac{x - x_{\min}}{x_{\max} - x_{\min}}\]

Result: Range [0, 1]

Figure: Why preprocessing matters. Left: Raw features at vastly different scales — salary dominates (10,000× larger than height). Middle: After standardization (z-score), all features have mean=0 and std=1. Right: After min-max normalization, all features are in [0, 1]. Without normalization, gradient descent would be dominated by large-scale features.

When to Use Which?

Image Preprocessing

# Typical ImageNet preprocessing
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),  # Converts to [0, 1]
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],  # ImageNet stats
        std=[0.229, 0.224, 0.225]
    )
])

Text Preprocessing

1. Tokenization: “Hello world” → [“Hello”, “world”] → [15496, 995]
2. Padding: Make all sequences same length
3. Embedding: Token IDs → dense vectors

Handling Missing Values

Real-world datasets often have missing values. Understanding why data is missing guides how to handle it.

Types of Missingness:

Why It Matters: MCAR is “safe” — any handling method works. MAR can be handled with imputation if you model the missingness. MNAR is problematic — you can’t fully correct for it without additional information.

Handling Strategies:

⚠️ Important: Imputation Considerations

Imputation means filling in missing values with estimated values. Key points:

1. Imputation introduces bias: The imputed values are estimates, not real data. They reduce variance and can make relationships appear stronger than they are.
2. Never impute the target variable: If your label/outcome is missing, that sample should typically be excluded, not imputed.
3. Fit imputer on training data only: Just like with normalization, compute imputation statistics (mean, median, etc.) on training data, then apply to validation/test sets. This prevents data leakage.
4. Consider multiple imputation: For statistical inference, advanced techniques like MICE (Multiple Imputation by Chained Equations) account for uncertainty in imputed values.
5. Document your approach: Always report what percentage of data was missing and how you handled it—this affects reproducibility.

Python Examples:

import pandas as pd
import numpy as np
from sklearn.impute import SimpleImputer, KNNImputer

# Sample data with missing values
df = pd.DataFrame({
    'age': [25, 30, np.nan, 45, 50],
    'income': [50000, np.nan, 70000, 80000, np.nan],
    'category': ['A', 'B', np.nan, 'A', 'B']
})

# Method 1: Drop rows with any missing values
df_dropped = df.dropna()  # 2 rows remain

# Method 2: Mean imputation for numerical columns
mean_imputer = SimpleImputer(strategy='mean')
df['age_imputed'] = mean_imputer.fit_transform(df[['age']])

# Method 3: Median imputation (robust to outliers)
median_imputer = SimpleImputer(strategy='median')
df['income_imputed'] = median_imputer.fit_transform(df[['income']])

# Method 4: Mode imputation for categorical
mode_imputer = SimpleImputer(strategy='most_frequent')
df['category_imputed'] = mode_imputer.fit_transform(df[['category']])

# Method 5: KNN imputation (considers feature relationships)
knn_imputer = KNNImputer(n_neighbors=2)
numerical_cols = df[['age', 'income']].values
df_knn = pd.DataFrame(knn_imputer.fit_transform(numerical_cols),
                      columns=['age_knn', 'income_knn'])

# Method 6: Add missingness indicator
df['income_was_missing'] = df['income'].isna().astype(int)

Interview Q: “How would you handle missing values in a dataset?”

A: First, I’d analyze the missingness pattern to determine if it’s MCAR, MAR, or MNAR — this guides the approach. For MCAR with few missing values (<5%), simple deletion may work. For numerical features, I’d use median imputation (robust to outliers) or KNN imputation if features are correlated. For categorical features, mode imputation or a separate “Unknown” category. If missingness itself might be informative (e.g., people skip income questions intentionally), I’d add an indicator variable. For MAR in large datasets, model-based imputation (like using Random Forest to predict missing values) can capture complex relationships. I’d always validate by comparing model performance with different imputation strategies.

Outlier Detection and Handling

Interview Q: “How do you detect and handle outliers in your data? How do you make models robust to outliers?”

Outliers are data points that are significantly different from other observations. They can be:

• True outliers: Rare but valid data points (e.g., a billionaire in income data)
• Errors: Data entry mistakes, sensor malfunctions, or corruption

Detection Methods

1. Interquartile Range (IQR) Method (Box-Plot/Tukey’s Fences)

The most common statistical method, based on the interquartile range:

IQR = Q3 - Q1

Lower bound = Q1 - 1.5 × IQR
Upper bound = Q3 + 1.5 × IQR

Points outside these bounds → Outliers

Figure: Anatomy of a box-plot (left) showing Q1, median, Q3, IQR, and whiskers. The IQR method (right) detects outliers as points beyond Q1 - 1.5×IQR or Q3 + 1.5×IQR.

import numpy as np

def detect_outliers_iqr(data):
    """Detect outliers using IQR method (1.5×IQR rule)"""
    Q1 = np.percentile(data, 25)
    Q3 = np.percentile(data, 75)
    IQR = Q3 - Q1
    
    lower_bound = Q1 - 1.5 * IQR
    upper_bound = Q3 + 1.5 * IQR
    
    outliers = (data < lower_bound) | (data > upper_bound)
    return outliers, lower_bound, upper_bound

# Example
data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 100])  # 100 is outlier
is_outlier, lb, ub = detect_outliers_iqr(data)
print(f"Outliers: {data[is_outlier]}")  # [100]

2. Z-Score Method

For normally distributed data, points far from the mean (typically > 3σ) are outliers:

\[z = \frac{x - \mu}{\sigma}\]

def detect_outliers_zscore(data, threshold=3):
    """Detect outliers using Z-score (points > threshold std from mean)"""
    mean = np.mean(data)
    std = np.std(data)
    z_scores = np.abs((data - mean) / std)
    return z_scores > threshold

3. ML-Based Methods (High-Dimensional Data)

from sklearn.ensemble import IsolationForest

# Isolation Forest for high-dimensional data
clf = IsolationForest(contamination=0.1, random_state=42)
outlier_labels = clf.fit_predict(X)  # -1 for outliers, 1 for inliers

Handling Strategies

Once outliers are detected, you have several options:

# Winsorization: Cap at percentiles
from scipy.stats import mstats

def winsorize(data, limits=(0.01, 0.01)):
    """Cap outliers at 1st and 99th percentiles"""
    return mstats.winsorize(data, limits=limits)

# Log transform: Reduce right skew
def log_transform(data):
    return np.log1p(data)  # log(1+x) handles zeros

Making Models Robust to Outliers

1. Use Robust Loss Functions

# Huber loss: MSE for small errors, MAE for large
import torch.nn as nn
criterion = nn.HuberLoss(delta=1.0)  # delta controls transition point

2. Use Robust Scaling

from sklearn.preprocessing import RobustScaler

# Uses median and IQR instead of mean and std
# Much less affected by outliers than StandardScaler
scaler = RobustScaler()  # (x - median) / IQR
X_scaled = scaler.fit_transform(X)

3. Choose Robust Algorithms

4. Regularization Helps

Regularization reduces variance and makes models less sensitive to individual points:

• L2 (Ridge): Shrinks weights, reduces influence of any single feature
• Dropout: Random neuron dropping prevents over-reliance on specific patterns

Interview Q: “Why are tree-based models robust to outliers?”

A: Tree-based models (Random Forest, XGBoost) split data based on thresholds, not magnitudes. A value of 100 vs 1,000,000 might fall in the same leaf node if they’re both above the split threshold. The prediction depends on which leaf the point lands in, not the exact value. This makes trees naturally robust to outliers—unlike linear models where outliers directly pull the regression line.

Quick Decision Guide

Outlier detected?
       │
       ├── Is it a data error?
       │       │
       │       ├── YES → Remove or impute
       │       │
       │       └── NO (rare but valid)
       │               │
       │               ├── Use robust methods (MAE, Huber, trees)
       │               └── Or winsorize/cap if you must use non-robust methods
       │
       └── Not sure?
               │
               └── 1. Investigate the data point
                   2. Try both with/without
                   3. Use cross-validation to decide

1.10 Loss Function Comparison

The loss function (also called cost function or objective function) is arguably the most important design choice in machine learning—it defines exactly what “good” means for your model. During training, the optimizer’s sole job is to minimize this function, so your model will learn whatever behavior the loss function rewards.

Why does loss function choice matter so much?

1. Defines the learning signal: The gradients that update your weights come from the loss. Choose the wrong loss, and your model receives misleading gradients that don’t guide it toward the right solution.

2. Must match your task:

   • Regression (predict continuous values): Use MSE, MAE, or Huber
   • Classification (predict categories): Use Cross-Entropy (binary or categorical)
   • Using MSE for classification can fail spectacularly—see below!

3. Affects optimization dynamics: Some losses have better gradient properties than others. Cross-entropy gives clean gradients for classification; MSE with sigmoid can cause vanishing gradients.

How to think about it: The loss function is a contract between you and the optimizer. You specify “minimize this number,” and the optimizer will find weights that do exactly that—even if that’s not what you actually wanted. Choosing the right loss ensures that minimizing the number also means solving your actual problem.

Quick Reference Table

MSE vs Cross-Entropy for Classification

Why not MSE for classification?

Figure: Comparison of Cross-Entropy and MSE loss for classification. Cross-entropy penalizes confident wrong predictions much more severely.

True label: y = 1
Prediction: ŷ = 0.99 (very confident, correct)

MSE: (1 - 0.99)² = 0.0001
CE:  -log(0.99) = 0.01

Prediction: ŷ = 0.01 (very confident, WRONG!)

MSE: (1 - 0.01)² = 0.98
CE:  -log(0.01) = 4.6  ← Much stronger penalty!

Cross-entropy penalizes confident wrong predictions much more heavily!

MSE Gradient Problem

For sigmoid output with MSE: \[\frac{\partial \text{MSE}}{\partial z} \propto \sigma(z)(1-\sigma(z))\]

When \(\sigma(z) \approx 0\) or \(\sigma(z) \approx 1\): gradient vanishes!

For cross-entropy: \[\frac{\partial \text{CE}}{\partial z} = \hat{y} - y\]

No vanishing gradient! Clean, constant-scale updates.

Huber Loss: Best of Both Worlds

\[L_\delta(y, \hat{y}) = \begin{cases} \frac{1}{2}(y - \hat{y})^2 & |y - \hat{y}| \leq \delta \\ \delta|y - \hat{y}| - \frac{1}{2}\delta^2 & |y - \hat{y}| > \delta \end{cases}\]

• MSE for small errors (smooth)
• MAE for large errors (robust to outliers)

1.11 Convolutional Neural Networks (CNNs)

The Problem with MLPs for Images

A 224×224 RGB image has 224 × 224 × 3 = 150,528 input features.

Fully connected layer with 1000 hidden units:

• 150,528 × 1000 = 150 million parameters in first layer alone!
• Doesn’t exploit spatial structure
• No translation invariance

The Convolution Operation

A filter (kernel) slides across the image:

Input (5×5):              Filter (3×3):           Output (3×3):
┌───┬───┬───┬───┬───┐     ┌───┬───┬───┐
│ 1 │ 2 │ 3 │ 0 │ 1 │     │ 1 │ 0 │-1 │
├───┼───┼───┼───┼───┤     ├───┼───┼───┤
│ 0 │ 1 │ 2 │ 3 │ 2 │     │ 1 │ 0 │-1 │         Slide filter,
├───┼───┼───┼───┼───┤  *  ├───┼───┼───┤    →    compute dot product
│ 1 │ 0 │ 1 │ 0 │ 1 │     │ 1 │ 0 │-1 │         at each position
├───┼───┼───┼───┼───┤     └───┴───┴───┘
│ 2 │ 1 │ 0 │ 1 │ 2 │
├───┼───┼───┼───┼───┤
│ 1 │ 0 │ 2 │ 1 │ 0 │
└───┴───┴───┴───┴───┘

Example computation (top-left):
1×1 + 2×0 + 3×(-1) + 0×1 + 1×0 + 2×(-1) + 1×1 + 0×0 + 1×(-1) = 1-3-2+1-1 = -4

Why Convolutions Work

1. Parameter sharing: Same filter used everywhere → fewer parameters
2. Local connectivity: Each output depends on small local region
3. Translation equivariance: Cat in corner detected same as cat in center

Key CNN Concepts

Stride: How many pixels to move the filter each step

• Stride 1: Output ≈ input size
• Stride 2: Output ≈ half input size

Padding: Add zeros around input to preserve dimensions

• “Same” padding: Output = Input size
• “Valid” padding: No padding, output smaller

Pooling: Downsample by taking max/average over regions

2×2 Max Pool:
┌───┬───┐     ┌───┐
│ 1 │ 3 │     │ 4 │
├───┼───┤  →  └───┘
│ 2 │ 4 │
└───┴───┘

Receptive Fields: What Each Neuron “Sees”

The receptive field of a neuron is the region of the input image that can influence its output. Understanding receptive fields is crucial for CNN design.

Figure: Receptive field growth in a CNN. Each 3×3 convolution increases the receptive field by 2 pixels per side. After 2 layers with 3×3 kernels, a single output neuron “sees” a 5×5 region of the input.

Receptive field formula (for stride-1 convolutions):

\[RF_{out} = RF_{in} + (k - 1) \times \prod_{i=1}^{l-1} s_i\]

Where: - \(RF\) = receptive field size - \(k\) = kernel size - \(s_i\) = stride at layer \(i\)

Simplified rule for 3×3 kernels with stride 1: \[RF_l = RF_{l-1} + 2\]

Starting from \(RF_0 = 1\) (single pixel): - After Conv1 (3×3): RF = 3 - After Conv2 (3×3): RF = 5 - After Conv3 (3×3): RF = 7 - After Conv4 (3×3): RF = 9

Why receptive fields matter:

1. Feature hierarchy: Early layers have small RFs → detect edges, textures. Deep layers have large RFs → detect objects, scenes
2. Network depth: Deeper networks = larger RFs = can capture more global information
3. Design tradeoff: Larger kernels (5×5, 7×7) increase RF faster but add more parameters

Interview Q: “Why do modern CNNs use stacked 3×3 convolutions instead of larger kernels?”

A: Two 3×3 convolutions have the same receptive field as one 5×5 (RF = 5), but with: - Fewer parameters: \(2 \times (3^2) = 18\) vs \(5^2 = 25\) - More non-linearity: Two ReLU activations vs one - More representational power

A Simple CNN Architecture

Input: 32×32×3 (e.g., CIFAR-10 image)
  ↓
Conv: 32 filters, 3×3 → 32×32×32
ReLU
MaxPool 2×2 → 16×16×32
  ↓
Conv: 64 filters, 3×3 → 16×16×64
ReLU
MaxPool 2×2 → 8×8×64
  ↓
Flatten → 4096
  ↓
FC → 256
ReLU
  ↓
FC → 10 (classes)
Softmax

PyTorch CNN

import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 8 * 8, 256)
        self.fc2 = nn.Linear(256, 10)
    
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # 32×32 → 16×16
        x = self.pool(F.relu(self.conv2(x)))  # 16×16 → 8×8
        x = x.view(-1, 64 * 8 * 8)            # Flatten
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

1.12 Word Embeddings

The Problem: How to Represent Words as Numbers?

Neural networks need numerical inputs. How do we convert words to numbers?

Approach 1: One-Hot Encoding

"cat"  → [1, 0, 0, 0]
"dog"  → [0, 1, 0, 0]
"bird" → [0, 0, 1, 0]
"fish" → [0, 0, 0, 1]

Problems:

1. Sparse: V = 50,000 → vectors with 49,999 zeros
2. No similarity: cat·dog = 0 (orthogonal), even though semantically related
3. Memory: Large vocabulary = huge vectors

Approach 2: Dense Embeddings

Map each word to a dense, low-dimensional vector:

"cat"  → [0.2, -0.4, 0.1, 0.8, ...]   (d = 300 dimensions)
"dog"  → [0.3, -0.3, 0.0, 0.7, ...]   (similar to cat!)
"fish" → [-0.5, 0.2, 0.9, -0.1, ...]  (different)

Why Embeddings Work

Distributional hypothesis: Words that appear in similar contexts have similar meanings.

“The cat sat on the mat” “The dog sat on the mat”

cat and dog appear in similar contexts → similar embeddings!

Embedding Layer in PyTorch

import torch.nn as nn

# Vocabulary of 10,000 words, embedding dimension 300
embedding = nn.Embedding(num_embeddings=10000, embedding_dim=300)

# Convert word indices to embeddings
word_indices = torch.tensor([42, 1337, 99])  # 3 words
word_vectors = embedding(word_indices)        # Shape: (3, 300)

Pretrained Embeddings

Word2Vec, GloVe: Trained on billions of words

• Capture semantic relationships
• “king” - “man” + “woman” ≈ “queen”

# Using pretrained GloVe
from torchtext.vocab import GloVe
glove = GloVe(name='6B', dim=300)
cat_embedding = glove['cat']  # 300-dimensional vector

How Word2Vec Learns: Skip-gram Training

Understanding how Word2Vec learns embeddings provides deep insight into self-supervised learning—the foundation of modern LLMs.

Figure: Skip-gram and CBOW architectures. Skip-gram predicts context words from the center word; CBOW predicts the center word from context.

The Skip-gram Objective

Given a center word, predict its surrounding context words:

\[\text{maximize} \quad \sum_{t=1}^{T} \sum_{-c \leq j \leq c, j \neq 0} \log P(w_{t+j} | w_t)\]

Where \(c\) is the context window size.

The probability model:

\[P(w_O | w_I) = \frac{\exp(v'_{w_O} \cdot v_{w_I})}{\sum_{w=1}^{V} \exp(v'_w \cdot v_{w_I})}\]

This is just a softmax! The dot product measures similarity between: - \(v_{w_I}\): Input embedding of the center word - \(v'_{w_O}\): Output embedding of the context word

Why this learns semantic similarity:

The training signal “predict context from word” forces words appearing in similar contexts to have similar embeddings. Consider:

"The cat sat on the mat"
"The dog sat on the mat"

Both “cat” and “dog” must predict the same context words (“the”, “sat”, “on”, “mat”), so they’re pushed to have similar embeddings!

Negative Sampling (Practical Training)

The full softmax over vocabulary is expensive (V can be 100,000+). Negative sampling approximates it:

Instead of computing the full softmax, contrast the positive (real) context word against random “negative” words:

\[\mathcal{L} = \log \sigma(v'_{w_O} \cdot v_{w_I}) + \sum_{i=1}^{k} \mathbb{E}_{w_i \sim P_n(w)} [\log \sigma(-v'_{w_i} \cdot v_{w_I})]\]

• First term: Push the real context word embedding closer to center word
• Second term: Push random noise words away from center word
• Typically \(k = 5\)-20 negative samples

This is contrastive learning! The same principle underlies modern self-supervised methods like SimCLR and CLIP.

From Word Embeddings to Transformers

Modern LLMs don’t use fixed word embeddings:

1. Subword tokenization: “unhappiness” → [“un”, “happiness”]
2. Contextual embeddings: Same word, different meaning in context
3. Learned during pretraining: Not pretrained separately

Static embedding (Word2Vec):
  "bank" → same vector always

Contextual embedding (BERT, GPT):
  "river bank" → one vector
  "bank account" → different vector!

1.13 The XOR Problem: A Complete MLP Example

Why XOR?

XOR is the classic example showing why we need hidden layers:

No single line can separate the classes! (It’s not linearly separable)

  x₂
  ↑
1 │  ●        ○      ← Can't draw one line
  │                     to separate ● from ○
0 │  ○        ●
  └──────────────→ x₁
     0        1

The Network Architecture

Input Layer      Hidden Layer (2 neurons)     Output Layer
    x₁  ─────┐
             ├───→ h₁ ─────┐
    x₂  ─────┤             ├───→ y
             ├───→ h₂ ─────┘
    1 (bias)─┘

Dimensions:

• Input: 2 features
• Hidden: 2 neurons (with ReLU)
• Output: 1 neuron (with sigmoid)

Step 1: Initialize Weights

Let’s use specific weights that solve XOR:

Hidden layer weights \(W^{(1)}\) and biases \(b^{(1)}\): \[W^{(1)} = \begin{bmatrix} 1 & 1 \\ 1 & 1 \end{bmatrix}, \quad b^{(1)} = \begin{bmatrix} 0 \\ -1 \end{bmatrix}\]

Output layer weights \(W^{(2)}\) and bias \(b^{(2)}\): \[W^{(2)} = \begin{bmatrix} 1 \\ -2 \end{bmatrix}, \quad b^{(2)} = 0\]

Step 2: Forward Pass (for input [1, 1])

Hidden layer pre-activation: \[z^{(1)} = W^{(1)} \begin{bmatrix} 1 \\ 1 \end{bmatrix} + b^{(1)} = \begin{bmatrix} 1 \cdot 1 + 1 \cdot 1 \\ 1 \cdot 1 + 1 \cdot 1 \end{bmatrix} + \begin{bmatrix} 0 \\ -1 \end{bmatrix} = \begin{bmatrix} 2 \\ 1 \end{bmatrix}\]

Hidden layer activation (ReLU): \[h = \text{ReLU}(z^{(1)}) = \begin{bmatrix} \max(0, 2) \\ \max(0, 1) \end{bmatrix} = \begin{bmatrix} 2 \\ 1 \end{bmatrix}\]

Output layer pre-activation: \[z^{(2)} = W^{(2)T} h + b^{(2)} = 1 \cdot 2 + (-2) \cdot 1 + 0 = 0\]

Output (sigmoid): \[\hat{y} = \sigma(0) = \frac{1}{1 + e^0} = 0.5\]

Step 3: All Four Inputs

Better-Tuned Weights for XOR

The weights above give outputs at 0.5 for (0,0) and (1,1) — not ideal. Here are better weights that give outputs closer to 0 and 1:

Optimal hidden layer weights \(W^{(1)}\) and biases \(b^{(1)}\): \[W^{(1)} = \begin{bmatrix} 20 & 20 \\ 20 & 20 \end{bmatrix}, \quad b^{(1)} = \begin{bmatrix} -10 \\ -30 \end{bmatrix}\]

Optimal output layer weights \(W^{(2)}\) and bias \(b^{(2)}\): \[W^{(2)} = \begin{bmatrix} 20 \\ -20 \end{bmatrix}, \quad b^{(2)} = -10\]

With these weights:

Intuition behind these weights: - Hidden neuron 1: \(h_1 = \text{ReLU}(20x_1 + 20x_2 - 10)\) fires when at least one input is 1 - Hidden neuron 2: \(h_2 = \text{ReLU}(20x_1 + 20x_2 - 30)\) fires only when both inputs are 1 - Output: \(20h_1 - 20h_2 - 10\) is large positive only when \(h_1 > 0\) and \(h_2 = 0\)

Step 4: Compute Loss (Binary Cross-Entropy)

For input (1, 1) with \(\hat{y} = 0.5\), target \(y = 0\):

\[\mathcal{L} = -[y \log(\hat{y}) + (1-y) \log(1-\hat{y})]\]

\[= -[0 \cdot \log(0.5) + 1 \cdot \log(0.5)]\]

\[= -\log(0.5) = 0.693\]

Step 5: Backward Pass

Output layer gradient: \[\frac{\partial \mathcal{L}}{\partial z^{(2)}} = \hat{y} - y = 0.5 - 0 = 0.5\]

Gradient w.r.t. output weights: \[\frac{\partial \mathcal{L}}{\partial W^{(2)}} = h \cdot \frac{\partial \mathcal{L}}{\partial z^{(2)}} = \begin{bmatrix} 2 \\ 1 \end{bmatrix} \cdot 0.5 = \begin{bmatrix} 1.0 \\ 0.5 \end{bmatrix}\]

Gradient flowing to hidden layer: \[\frac{\partial \mathcal{L}}{\partial h} = W^{(2)} \cdot \frac{\partial \mathcal{L}}{\partial z^{(2)}} = \begin{bmatrix} 1 \\ -2 \end{bmatrix} \cdot 0.5 = \begin{bmatrix} 0.5 \\ -1.0 \end{bmatrix}\]

Through ReLU (gradient is 1 where input > 0, else 0): \[\frac{\partial \mathcal{L}}{\partial z^{(1)}} = \frac{\partial \mathcal{L}}{\partial h} \odot \mathbf{1}_{z^{(1)} > 0} = \begin{bmatrix} 0.5 \\ -1.0 \end{bmatrix} \odot \begin{bmatrix} 1 \\ 1 \end{bmatrix} = \begin{bmatrix} 0.5 \\ -1.0 \end{bmatrix}\]

Gradient w.r.t. hidden weights: \[\frac{\partial \mathcal{L}}{\partial W^{(1)}} = \frac{\partial \mathcal{L}}{\partial z^{(1)}} \cdot x^T = \begin{bmatrix} 0.5 \\ -1.0 \end{bmatrix} \cdot \begin{bmatrix} 1 & 1 \end{bmatrix} = \begin{bmatrix} 0.5 & 0.5 \\ -1.0 & -1.0 \end{bmatrix}\]

PyTorch Implementation

import torch
import torch.nn as nn

# XOR data
X = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=torch.float32)
y = torch.tensor([[0], [1], [1], [0]], dtype=torch.float32)

# Network
model = nn.Sequential(
    nn.Linear(2, 2),
    nn.ReLU(),
    nn.Linear(2, 1),
    nn.Sigmoid()
)

criterion = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1.0)

# Training
for epoch in range(1000):
    y_pred = model(X)
    loss = criterion(y_pred, y)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 200 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

# Test
print("\nPredictions:")
print(model(X).detach().round())  # Should be [[0], [1], [1], [0]]

Key Insights

1. Hidden layer creates new representation: The hidden layer transforms the space so that XOR becomes linearly separable
2. Non-linearity is essential: Without ReLU, two linear layers collapse to one
3. Backprop chain rule: Gradients flow backward through each layer

1.14 Weight Initialization

Why Does Initialization Matter?

Bad initialization leads to:

• All zeros: All neurons compute the same thing, learn the same features (symmetry problem)
• Too small: Activations shrink to zero layer by layer (vanishing signals)
• Too large: Activations explode, gradients explode

The Goal

Keep activations and gradients at reasonable scale throughout the network.

Xavier/Glorot Initialization (2010)

For tanh/sigmoid activations:

\[W \sim \mathcal{N}\left(0, \frac{2}{n_{\text{in}} + n_{\text{out}}}\right) \quad \text{or} \quad W \sim \mathcal{U}\left(-\sqrt{\frac{6}{n_{\text{in}} + n_{\text{out}}}}, \sqrt{\frac{6}{n_{\text{in}} + n_{\text{out}}}}\right)\]

Intuition: Balance the variance of inputs and outputs so signal doesn’t explode or vanish.

# PyTorch
nn.init.xavier_uniform_(layer.weight)
nn.init.xavier_normal_(layer.weight)

He/Kaiming Initialization (2015)

For ReLU activations (ReLU kills half the signal, so we compensate):

\[W \sim \mathcal{N}\left(0, \frac{2}{n_{\text{in}}}\right)\]

Why factor of 2? ReLU zeroes out negative half, so variance is halved. We double the initial variance to compensate.

Deriving He Initialization: Variance Propagation

Let’s derive why He initialization uses \(\text{Var}(W) = \frac{2}{n_{in}}\) for ReLU networks.

Figure: How activation variance propagates through layers. Without proper initialization, variance either explodes or vanishes exponentially. The goal: keep Var(y) ≈ Var(x) at each layer.

Setup: Consider a single layer \(y = Wx + b\) where: - Input \(x\) has \(n_{in}\) components, each with variance \(\text{Var}(x_j)\) - Weights \(W_{ij} \sim \mathcal{N}(0, \sigma_w^2)\) are independent of inputs - We want \(\text{Var}(y_i) = \text{Var}(x_j)\) (preserve variance)

Step 1: Variance of one output neuron (before activation)

\[y_i = \sum_{j=1}^{n_{in}} W_{ij} x_j + b_i\]

For zero-mean \(x\) and \(W\), assuming independence:

\[\text{Var}(y_i) = \sum_{j=1}^{n_{in}} \text{Var}(W_{ij} x_j) = \sum_{j=1}^{n_{in}} \text{Var}(W_{ij}) \cdot \text{Var}(x_j) = n_{in} \cdot \sigma_w^2 \cdot \text{Var}(x)\]

Step 2: Preserve variance (Xavier derivation)

For \(\text{Var}(y) = \text{Var}(x)\), we need: \[n_{in} \cdot \sigma_w^2 = 1 \implies \sigma_w^2 = \frac{1}{n_{in}}\]

This is Xavier initialization — perfect for linear layers or tanh (which is approximately linear near 0).

Step 3: Account for ReLU (He derivation)

ReLU zeros out negative values. For zero-mean Gaussian input, exactly half the values are negative:

\[\text{ReLU}(y) = \begin{cases} y & \text{if } y > 0 \\ 0 & \text{if } y \leq 0 \end{cases}\]

The variance after ReLU is halved: \[\text{Var}(\text{ReLU}(y)) = \frac{1}{2} \text{Var}(y)\]

To compensate, we need twice the initial variance:

\[\sigma_w^2 = \frac{2}{n_{in}}\]

This is He/Kaiming initialization!

Why this matters: Without this factor of 2, variance shrinks by half at each layer. After 10 layers: \(0.5^{10} \approx 0.001\) — activations become tiny, gradients vanish.

# PyTorch (default for nn.Linear with ReLU)
nn.init.kaiming_uniform_(layer.weight, nonlinearity='relu')
nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')

Quick Reference

What About Biases?

Almost always initialize to zero:

nn.init.zeros_(layer.bias)

Exception: LSTM forget gate biases often initialized to 1 to encourage remembering.

Example: Why Bad Init Fails

# BAD: All zeros
for layer in model.modules():
    if hasattr(layer, 'weight'):
        layer.weight.data.fill_(0)
# Result: All neurons output the same thing, all gradients identical

# BAD: Too large
for layer in model.modules():
    if hasattr(layer, 'weight'):
        layer.weight.data.normal_(0, 10)
# Result: Activations explode, NaN losses

# GOOD: He initialization for ReLU network
for layer in model.modules():
    if isinstance(layer, nn.Linear):
        nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')

1.15 Dropout

What is Dropout?

During training, randomly set neurons to zero with probability \(p\):

Without dropout:      With dropout (p=0.5):

[0.5] → [0.3]        [0.5] → [0.0]  ← dropped!
[0.8] → [0.2]        [0.8] → [0.4]  ← scaled by 1/(1-p)
[0.1] → [0.7]        [0.0] → [0.0]  ← dropped!
[0.9] → [0.5]        [0.9] → [1.0]  ← scaled

Why Does It Work?

1. Prevents co-adaptation: Neurons can’t rely on specific other neurons being present
2. Ensemble effect: Each forward pass uses a different “sub-network”
3. Implicit regularization: Similar effect to training many models and averaging

Figure: Dropout creates an implicit ensemble of exponentially many sub-networks. During training, each forward pass uses a different random subset of neurons (shown as different colored sub-networks). At inference, we use the full network with scaled weights, which approximates the average prediction of all sub-networks.

Training vs Inference

Training: Drop neurons with probability \(p\), scale remaining by \(\frac{1}{1-p}\)

Inference: Use all neurons (no dropping)

class Dropout(nn.Module):
    def __init__(self, p=0.5):
        self.p = p
    
    def forward(self, x):
        if self.training:  # Training mode
            mask = (torch.rand_like(x) > self.p).float()
            return x * mask / (1 - self.p)  # Scale to maintain expected value
        else:  # Inference mode
            return x  # No change

PyTorch Usage

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Dropout(p=0.5),  # 50% dropout
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(128, 10)
)

# CRITICAL: Set mode correctly!
model.train()  # Enable dropout
model.eval()   # Disable dropout

Typical Dropout Rates

Dropout Variants

Interview Q: “Why use dropout instead of L2?”

A: Different mechanisms:

• L2 shrinks all weights smoothly
• Dropout forces redundancy and prevents co-adaptation
• Dropout is stochastic (different network each pass)
• Can use both together!

1.16 Overfitting & Regularization Preview

Overfitting is perhaps the most fundamental challenge in machine learning. It occurs when a model learns the training data too well — including its noise and idiosyncrasies — rather than the underlying patterns that generalize to new data. Understanding and preventing overfitting is what separates models that work in practice from those that only work on paper.

The Central Problem

The core tension in machine learning is this: we want our model to perform well on data it has never seen before, but we can only train it on data we have. This creates a fundamental dilemma.

                    Training Data
                         ↓
                    [Your Model]
                         ↓
        ┌────────────────┴────────────────┐
        ↓                                 ↓
   Fits training                    Generalizes to
   data well                        new data?
   
   EASY                             THE HARD PART

A sufficiently complex model can memorize any training set — given enough parameters, it can simply store each training example and its label. Such a model achieves zero training error but learns nothing useful. When shown a new example, it has no idea what to do because it never learned the underlying pattern; it only memorized specific instances.

The goal of regularization is to prevent this memorization by encouraging the model to learn simpler, more generalizable patterns. The intuition is that the true underlying relationship is usually simpler than the noise-contaminated training data suggests.

Visual: The Overfitting Spectrum

The figure above illustrates the spectrum from underfitting to overfitting:

• Underfitting (High Bias): The model is too simple to capture the underlying pattern. A linear model trying to fit a quadratic relationship, for example. Both training and test error are high.

• Good Fit: The model captures the true underlying pattern without memorizing noise. Training error is low, and test error is similarly low.

• Overfitting (High Variance): The model is too complex and has memorized training-specific noise. Training error is very low (often near zero), but test error is high because the “learned” noise doesn’t exist in new data.

Learning Curves: Your Diagnostic Tool

Learning curves are your most powerful tool for diagnosing training problems. They plot training and validation loss (or accuracy) against training progress (epochs or iterations).

How to read learning curves:

1. Healthy training: Both curves decrease together and converge to similar values. The gap between them is small and stable.

2. Overfitting signature: Training loss continues to decrease, but validation loss stops improving or starts increasing. The gap between them grows over time. This tells you the model is memorizing training data rather than learning generalizable patterns.

3. Underfitting signature: Both training and validation loss remain high and plateau early. Neither improves much with more training. This tells you the model lacks capacity to learn the pattern, or there’s a fundamental problem with the setup.

4. When to stop: The optimal stopping point is typically where validation loss is lowest — just before the gap starts widening. This is why early stopping is so effective.

The Regularization Toolbox

Regularization encompasses any technique that helps your model generalize better, typically by constraining its complexity in some way. Here are the main tools:

L2 Regularization (Weight Decay)

\[\text{Total Loss} = \underbrace{\text{Data Loss}}_{\text{fit the data}} + \underbrace{\lambda \sum_j w_j^2}_{\text{keep weights small}}\]

Effect: Larger \(\lambda\) → smaller weights → simpler model → less overfitting

Key property: Shrinks all weights toward zero, but rarely makes them exactly zero.

L1 Regularization (Lasso)

\[\text{Total Loss} = \underbrace{\text{Data Loss}}_{\text{fit the data}} + \underbrace{\lambda \sum_j |w_j|}_{\text{push weights to zero}}\]

Effect: Larger \(\lambda\) → more weights become exactly zero → automatic feature selection

Key property: Creates sparse models — some features are completely ignored.

L1 vs L2: When to Use Which?

Visual intuition (in 2D weight space):

L2 Constraint (circle):        L1 Constraint (diamond):
       
      ___                            /\
    /     \                         /  \
   |   ●   |   ← Solution          /    \
    \     /                       /  ●   \   ← Solution hits corner
      ‾‾‾                         \      /     (one weight = 0!)
                                   \    /
                                    \  /
                                     \/

The L1 diamond has corners on the axes — the optimal point often lands on a corner where one or more weights are exactly zero.

Quick Decision Tree

Model overfitting?
       │
       ├── YES → 1. Add/increase dropout
       │         2. Add/increase L2 (weight decay)
       │         3. Get more data / data augmentation
       │         4. Reduce model size (last resort)
       │
       └── NO (underfitting) → 1. Bigger model
                               2. More features
                               3. Train longer
                               4. Less regularization

The Bias-Variance Tradeoff (Preview)

\[\text{Expected Error} = \text{Bias}^2 + \text{Variance} + \text{Noise}\]

💡 See Part 2 for full treatment of bias-variance and Part 3 for detailed regularization math!

1.17 Putting It All Together: Complete Neural Network Setup

This section consolidates everything from Part 1 into a single, practical code example showing how to configure each component of a neural network.

Complete Keras Example

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, regularizers

# ═══════════════════════════════════════════════════════════════════
# MODEL DEFINITION
# ═══════════════════════════════════════════════════════════════════

model = keras.Sequential([
    # Input layer (optional explicit definition)
    layers.Input(shape=(784,)),
    
    # Hidden Layer 1
    layers.Dense(
        units=256,                           # Number of neurons
        activation='relu',                   # Activation function
        kernel_initializer='he_normal',      # Weight initialization  
        kernel_regularizer=regularizers.l2(0.01)  # L2 regularization (weight decay)
    ),
    layers.BatchNormalization(),             # Normalize activations
    layers.Dropout(0.5),                     # Regularization (50% dropout)
    
    # Hidden Layer 2
    layers.Dense(
        units=128,
        activation='relu',
        kernel_initializer='he_normal',
        kernel_regularizer=regularizers.l2(0.01)
    ),
    layers.BatchNormalization(),
    layers.Dropout(0.3),
    
    # Output Layer (10-class classification)
    layers.Dense(
        units=10,
        activation='softmax',                # Softmax for multi-class
        kernel_initializer='glorot_uniform'  # Xavier for output layer
    )
])

# ═══════════════════════════════════════════════════════════════════
# COMPILE: OPTIMIZER + LOSS + METRICS
# ═══════════════════════════════════════════════════════════════════

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.001),
    loss='sparse_categorical_crossentropy',  # For integer labels [0, 1, 2, ...]
    metrics=['accuracy']
)

# ═══════════════════════════════════════════════════════════════════
# TRAINING
# ═══════════════════════════════════════════════════════════════════

history = model.fit(
    X_train, y_train,
    epochs=20,
    batch_size=32,
    validation_split=0.2,
    callbacks=[
        keras.callbacks.EarlyStopping(patience=5, restore_best_weights=True),
        keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=3)
    ]
)

# ═══════════════════════════════════════════════════════════════════
# EVALUATION
# ═══════════════════════════════════════════════════════════════════

test_loss, test_acc = model.evaluate(X_test, y_test)
predictions = model.predict(X_test)

Quick Reference: Component Options

Task-Specific Configurations

Common Mistakes to Avoid

# ❌ WRONG: Softmax + sparse_categorical expects integer labels
model.compile(loss='categorical_crossentropy')  # Use with one-hot labels
model.fit(X, y)  # where y = [0, 1, 2] ← integers

# ✅ CORRECT: Match loss to label format
model.compile(loss='sparse_categorical_crossentropy')  # For integer labels
model.compile(loss='categorical_crossentropy')  # For one-hot labels

# ❌ WRONG: Using sigmoid activation with CrossEntropyLoss
layers.Dense(10, activation='sigmoid')  # DON'T do this for multi-class!

# ✅ CORRECT: Softmax for multi-class
layers.Dense(10, activation='softmax')

# ❌ WRONG: He init with sigmoid/tanh
layers.Dense(256, activation='sigmoid', kernel_initializer='he_normal')

# ✅ CORRECT: Match initializer to activation
layers.Dense(256, activation='sigmoid', kernel_initializer='glorot_uniform')
layers.Dense(256, activation='relu', kernel_initializer='he_normal')

PyTorch Equivalent (for reference)

import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(784, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 10)  # No softmax! CrossEntropyLoss includes it
        )
        self._init_weights()
    
    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                nn.init.zeros_(m.bias)
    
    def forward(self, x):
        return self.layers(x)

# Training uses:
# optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
# criterion = nn.CrossEntropyLoss()  # Combines LogSoftmax + NLLLoss

1.18 Batch Normalization

The Problem: Internal Covariate Shift

Internal Covariate Shift (ICS) refers to the phenomenon where the distribution of inputs to each layer changes during training because the weights of previous layers are constantly being updated. Let’s understand why this is problematic.

What does “distribution shift” actually mean?

Consider a hidden layer that receives activations from the previous layer. At the start of training:

• The input activations might have mean ≈ 0.5, std ≈ 0.3
• The layer learns weights optimized for this distribution

After a few gradient updates to earlier layers:

• The same inputs now produce activations with mean ≈ 0.8, std ≈ 0.1
• The layer’s weights are now suboptimal for this new distribution
• It must re-adapt, but by the time it does, the distribution shifts again!

Figure: Internal Covariate Shift visualization. At training step 1, Layer 3 receives inputs with mean=0.5, std=0.3 and learns weights for that distribution. By step 100, the distribution has shifted (mean=0.8, std=0.1), making the learned weights suboptimal.

Why is this problematic?

1. Chasing a moving target: Each layer tries to learn, but its optimal weights depend on the input distribution, which keeps changing
2. Requires small learning rates: Large updates cause dramatic distribution shifts, destabilizing training
3. Slower convergence: Layers waste capacity constantly re-adapting instead of learning useful features
4. Saturated activations: If distributions shift into saturation regions of sigmoid/tanh, gradients vanish

A concrete example:

# Without BatchNorm: distributions shift wildly
# Layer 1 output at step 0:    mean=-0.02, std=0.98  
# Layer 1 output at step 100:  mean=2.31,  std=3.45   ← 100x shift!
# Layer 1 output at step 1000: mean=-0.87, std=0.23  ← shifted again!

# With BatchNorm: distributions stay normalized
# Layer 1 output (after BN) at any step: mean≈0, std≈1 (before γ,β)

Modern Perspective: Beyond ICS

While the ICS explanation is intuitive, recent research suggests BatchNorm may help primarily through other mechanisms:

1. Smoother loss landscape: BN makes the optimization surface more well-behaved with fewer sharp cliffs
2. Gradient flow: Normalized activations have more stable gradient magnitudes
3. Implicit learning rate adaptation: BN effectively adjusts the learning rate for each layer

Regardless of the exact mechanism, the empirical benefits are clear: BatchNorm enables faster, more stable training.

The Solution: Normalize Each Layer

Batch Normalization (Ioffe & Szegedy, 2015) normalizes activations within each mini-batch:

\[\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}\]

where:

• \(\mu_B = \frac{1}{B}\sum_{i=1}^{B} x_i\) (batch mean)
• \(\sigma_B^2 = \frac{1}{B}\sum_{i=1}^{B} (x_i - \mu_B)^2\) (batch variance)
• \(\epsilon\) ≈ \(10^{-5}\) (for numerical stability)

Learnable Parameters

After normalizing, we add learnable scale and shift parameters:

\[y_i = \gamma \hat{x}_i + \beta\]

• \(\gamma\) (gamma): learned scale parameter
• \(\beta\) (beta): learned shift parameter

Why? The network might need non-zero mean or non-unit variance for some layers. These parameters let it learn the optimal distribution.

Training vs Inference

Training: Use batch statistics (\(\mu_B\), \(\sigma_B^2\))

Inference: Use running averages of statistics accumulated during training

# During training, PyTorch automatically tracks:
# running_mean = momentum * running_mean + (1 - momentum) * batch_mean
# running_var = momentum * running_var + (1 - momentum) * batch_var

⚠️ Critical: Always set model.eval() before inference to use running statistics!

Where to Place BatchNorm?

# Option 1: After linear, before activation (original paper)
nn.Linear(256, 128),
nn.BatchNorm1d(128),
nn.ReLU(),

# Option 2: After activation (sometimes used)
nn.Linear(256, 128),
nn.ReLU(),
nn.BatchNorm1d(128),

Both work in practice. Option 1 is more common.

Benefits of Batch Normalization

1. Enables higher learning rates: Normalized activations are more stable
2. Regularization effect: Batch statistics add noise (like dropout)
3. Reduces sensitivity to initialization: Normalization handles bad init
4. Smoother loss landscape: Easier optimization

PyTorch Implementation

import torch.nn as nn

# For 1D data (fully connected layers): BatchNorm1d
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.BatchNorm1d(256),  # Normalize 256 features
    nn.ReLU(),
    nn.Linear(256, 10)
)

# For 2D data (CNNs): BatchNorm2d
cnn = nn.Sequential(
    nn.Conv2d(3, 64, kernel_size=3),
    nn.BatchNorm2d(64),  # Normalize 64 channels
    nn.ReLU(),
)

# CRITICAL: Set mode correctly!
model.train()  # Use batch statistics
model.eval()   # Use running statistics

Layer Normalization (Alternative)

For transformers and RNNs, Layer Normalization is preferred:

Figure: Batch Normalization vs Layer Normalization. BatchNorm (left) computes mean and variance across the batch dimension for each feature (highlighted column). LayerNorm (right) computes statistics across the feature dimension for each sample independently (highlighted row). This makes LayerNorm batch-independent and suitable for variable-length sequences.

# LayerNorm normalizes across features, not batch
# For a tensor of shape (batch, seq_len, hidden), normalizes across hidden
nn.LayerNorm(hidden_size)

RMSNorm: The Modern LLM Standard

RMSNorm (Root Mean Square Layer Normalization) has become the standard for modern LLMs like LLaMA, Gemma, Qwen, and Mistral. It simplifies LayerNorm by removing the mean-centering step.

LayerNorm formula: \[\text{LayerNorm}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta\]

RMSNorm formula: \[\text{RMSNorm}(x) = \gamma \cdot \frac{x}{\text{RMS}(x) + \epsilon} = \gamma \cdot \frac{x}{\sqrt{\frac{1}{n}\sum_i x_i^2 + \epsilon}}\]

Key differences:

Why does RMSNorm work without mean centering?

The key insight is that for well-behaved activations (especially after pre-norm architecture), the mean is often close to zero anyway. The primary purpose of normalization is to control the scale of activations, and RMS captures scale just as well as standard deviation:

# For zero-mean data: RMS ≈ std
x = torch.randn(1000)  # zero-mean by construction
print(f"std: {x.std():.4f}, RMS: {(x**2).mean().sqrt():.4f}")
# std: 1.0012, RMS: 0.9987  ← Nearly identical!

Implementation:

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))  # γ (scale only, no β)
    
    def forward(self, x):
        # RMS = sqrt(mean(x^2))
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        return self.weight * (x / rms)

# Usage in a transformer block (Pre-LN architecture)
class TransformerBlock(nn.Module):
    def __init__(self, dim, n_heads):
        super().__init__()
        self.norm1 = RMSNorm(dim)
        self.attn = MultiHeadAttention(dim, n_heads)
        self.norm2 = RMSNorm(dim)
        self.ffn = FeedForward(dim)
    
    def forward(self, x):
        x = x + self.attn(self.norm1(x))  # Pre-norm: normalize BEFORE attention
        x = x + self.ffn(self.norm2(x))   # Pre-norm: normalize BEFORE FFN
        return x

Performance advantage:

The speedup comes from fewer operations: - LayerNorm: compute mean, subtract mean, compute variance, divide by std - RMSNorm: compute mean of squares, divide by RMS

For large hidden dimensions (4096+) common in LLMs, this ~15-20% speedup in normalization adds up across billions of tokens.

Normalization Comparison Summary

1.19 Vanishing and Exploding Gradients

Vanishing/exploding gradients happen because backprop through depth/time multiplies many Jacobians. If their norms are <1, gradients shrink exponentially; if >1, they blow up. We mitigate with good init (Xavier/He), normalization (LayerNorm), residual connections (identity gradient path), gating (LSTM/GRU), and gradient clipping for explosions

The Problem

In deep networks, gradients must flow through many layers during backpropagation. At each layer, the gradient is multiplied by the layer’s weights:

\[\frac{\partial L}{\partial W_1} = \frac{\partial L}{\partial z_n} \cdot W_n \cdot W_{n-1} \cdots W_2\]

Vanishing Gradients

If weights are small (< 1), repeated multiplication makes gradients exponentially smaller:

10 layers with gradient factor 0.5 each:
0.5^10 = 0.001

Gradient at layer 1 is 1000x smaller than at layer 10!
→ Early layers learn extremely slowly

Causes:

• Sigmoid/Tanh activations (derivatives < 1 in most regions)
• Small weight initialization
• Very deep networks

Symptoms:

• Early layers don’t learn
• Loss decreases very slowly (Training is slow, especially for early layers.)
• Weights near input barely change

Exploding Gradients

If weights are large (> 1), gradients grow exponentially:

10 layers with gradient factor 2 each:
2^10 = 1024

Gradient at layer 1 is 1024x larger than at layer 10!
→ Huge, unstable weight updates

Causes:

• Large weight initialization
• Unstable architecture
• No normalization

Symptoms:

• Loss becomes NaN or Inf, or spikes wildly
• Weights explode to very large values
• Optimizer steps become unstable
• Training completely fails

Why activations matter (classic culprit)

• Sigmoid/tanh saturate: derivative max is < 1 and often near 0 when inputs are large ⇒ repeated multiplication shrinks gradients.
• ReLU helps vanishing (derivative is 1 on the positive side), but can create dead ReLUs (derivative 0 if stuck negative).

Solutions

How modern deep learning fights it

1. Better initialization

   Goal: keep variance of activations and gradients roughly constant across depth.

• Xavier/Glorot (tanh-ish nets)
• He/Kaiming (ReLU-ish nets)

2. Normalization

BatchNorm / LayerNorm / RMSNorm keep activations in a nice range, reducing Jacobian chaos.

3. Residual connections (ResNets, Transformers)

Residual block: \(h_{l+1} = h_l + F(h_l)\)

Backprop derivative includes an identity path: \(\frac{\partial h_{l+1}}{\partial h_l} = I + \frac{\partial F}{\partial h_l}\) That \(I\) is huge: gradients can flow even if \(F\) is messy.

4. Gradient clipping (mostly for exploding)

Clip global norm: \(g \leftarrow g \cdot \min\left(1, \frac{\tau}{\|g\|}\right)\) Prevents rare giant steps from blowing training up.

5. Gating mechanisms (LSTM/GRU)

They add near-identity “memory highways” so gradients don’t have to be multiplied by unstable transforms every step.

6. Optimizer/schedule choices

AdamW + warmup often stabilizes early training; but clipping + residuals + normalization do the heavy lifting.

Gradient Clipping

Cap the gradient norm to prevent explosion:

# Clip gradient norm to max value of 1.0
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# In training loop:
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # Add this!
optimizer.step()

Skip Connections (ResNet)

Add the input directly to the output, creating a “gradient highway”:

\[y = F(x) + x\]

class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(dim, dim),
            nn.ReLU(),
            nn.Linear(dim, dim)
        )
    
    def forward(self, x):
        return x + self.layers(x)  # Skip connection!

Why it helps: Gradients can flow directly through the skip connection, bypassing problematic layers.

1.20 Numerical Stability

Numerical stability is about designing computations so rounding error, overflow, and underflow don’t blow up the result. In floating-point, operations like exp, log, subtracting nearly equal numbers, summing long lists, and dividing by small values can be unstable.

A stable approach typically means reparameterizing (work in log space), shifting/scaling (e.g., softmax with max-subtraction), using special functions (log1p, expm1), stable accumulators (Welford/Kahan), and adding eps/clamps where appropriate—especially under mixed precision.

Numerical stability is controlling floating-point error/overflow/underflow. I watch for exp/log, subtracting close numbers, and large reductions. I use stable reformulations like max-shifted softmax and log-sum-exp, BCE-with-logits, Welford variance, log1p/expm1, and Kahan/pairwise summation; in mixed precision I keep critical ops in fp32 and use loss scaling/clipping.

What is Numerical Stability?

Numerical stability in neural networks refers to the ability of computations to produce consistent, finite, and meaningful results despite the inherent limitations of floating-point arithmetic. A numerically stable implementation:

1. Keeps values in representable ranges: Prevents overflow (values too large → inf) and underflow (values too small → 0) as activations pass through many layers
2. Minimizes accumulation of rounding errors: Ensures small errors don’t compound into large errors over many operations
3. Produces consistent results: Given the same inputs, returns the same outputs (determinism)

Why it matters for training: When values become very large or very small, gradients can explode or vanish, loss becomes NaN, and training fails. Normalization techniques (BatchNorm, LayerNorm, RMSNorm) help keep activations in a “numerically healthy” range where floating-point precision is good and gradients flow properly.

The Original Sin: Floating-Point Non-Associativity

The fundamental source of numerical variation in neural networks is floating-point non-associativity:

\[ (a + b) + c \neq a + (b + c) \quad \text{in floating-point!} \]

# This is NOT a bug - it's how floating-point works!
>>> (0.1 + 1e20) - 1e20
0.0
>>> 0.1 + (1e20 - 1e20)
0.1

Why does this happen?

Floating-point numbers use a format like \(\text{mantissa} \times 10^{\text{exponent}}\) (in base 2, but the principle is the same). This allows representing both very small and very large values with limited precision.

When adding two numbers with different scales (different exponents), information is lost:

  1.23 × 10³   (1230)
+ 2.34 × 10¹   (23.4)
= 1.2534 × 10³ (exact: 1253.4)

But with 3 digits of precision, we can only store: 1.25 × 10³ (1250)
The "34" is lost!

Every time we add floating-point numbers in a different order, we can get different results:

import random
vals = [1e-10, 1e-5, 1e-2, 1, -1e-10, -1e-5, -1e-2, -1]

results = set()
for _ in range(10000):
    random.shuffle(vals)
    results.add(sum(vals))

print(f"Summing 8 values in different orders gives {len(results)} unique results!")
# Output: 102 unique results!

Key insight: The mathematically “correct” answer depends on which order you perform the additions. There is no single “right” answer—just different answers for different orderings.

Determinism vs Batch Invariance (Advanced)

A common misconception is that GPU nondeterminism comes from “concurrent threads finishing in random order using atomic adds.” While this can cause nondeterminism, the real culprit in modern LLM inference is lack of batch invariance.

Run-to-run determinism vs user-observable determinism:

The batch invariance problem:

import torch
torch.set_default_device('cuda')

# Matrix-vector multiply (batch size = 1)
a = torch.linspace(-1000, 1000, 2048*4096).reshape(2048, 4096)
b = torch.linspace(-1000, 1000, 4096*4096).reshape(4096, 4096)

out1 = torch.mm(a[:1], b)         # Batch size 1
out2 = torch.mm(a, b)[:1]         # Same element, but batch size 2048

print((out1 - out2).abs().max())  # tensor(1669.25) - NOT zero!

The same element gives different results depending on batch size! When server load varies, batch size varies, and individual requests become nondeterministic—even though the kernel itself is “deterministic.”

Why this happens: Different batch sizes may trigger different:

• Parallelization strategies (data-parallel vs split-k)
• Tile sizes in matrix multiplication
  
• Reduction orderings in attention

To achieve true determinism, kernels must be batch-invariant: using fixed reduction strategies regardless of batch size. This is an active area of research for reproducible LLM inference.

The Log-Sum-Exp Trick

Computing softmax naively can overflow or underflow:

# BAD: Can overflow!
z = np.array([1000, 1001, 1002])
exp_z = np.exp(z)  # [inf, inf, inf] - OVERFLOW!

# BAD: Can underflow!
z = np.array([-1000, -1001, -1002])
exp_z = np.exp(z)  # [0, 0, 0] - UNDERFLOW!

Solution: Subtract the maximum value before exponentiating:

\[\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_j e^{z_j}} = \frac{e^{z_i - \max(z)}}{\sum_j e^{z_j - \max(z)}}\]

# GOOD: Numerically stable softmax
def stable_softmax(z):
    z_shifted = z - np.max(z)  # Shift to prevent overflow
    exp_z = np.exp(z_shifted)
    return exp_z / np.sum(exp_z)

Cross-Entropy with Logits

Never compute cross-entropy from probabilities — always from logits:

# BAD: Numerically unstable
probs = softmax(logits)
loss = -np.sum(y * np.log(probs))  # log(0) = -inf!

# GOOD: Use built-in function that handles stability
loss = nn.CrossEntropyLoss()(logits, labels)  # PyTorch handles it!

PyTorch’s CrossEntropyLoss combines log-softmax with NLL loss in a numerically stable way.

Epsilon for Logarithms

Always add a small epsilon when taking logarithms:

# BAD: log(0) = -inf
loss = -np.mean(y * np.log(y_pred))

# GOOD: Add epsilon
eps = 1e-15
loss = -np.mean(y * np.log(y_pred + eps))

Common Numerical Issues and Fixes

Further Reading: Defeating Nondeterminism in LLM Inference — Deep dive into floating-point non-associativity and batch invariance in production LLM systems.

1.21 Optimizers: SGD, Momentum, and Adam

Why Different Optimizers?

In Section 1.3, we introduced vanilla gradient descent: \[w \leftarrow w - \alpha \nabla L(w)\]

This works but has problems:

• Gets stuck in ravines (oscillates back and forth)
• Same learning rate for all parameters (suboptimal)
• Sensitive to learning rate choice

Different optimizers address these issues.

SGD with Momentum

The problem: Vanilla SGD oscillates in ravines—directions with high curvature see the gradient flip sign repeatedly, slowing convergence.

The solution: Add “momentum” like a ball rolling downhill. Accumulate velocity from past gradients:

\[v_t = \beta v_{t-1} + \nabla L(w_t)\] \[w_{t+1} = w_t - \alpha v_t\]

where \(\beta\) (typically 0.9) controls how much past gradients matter.

Why it works:

• Gradients that consistently point the same direction accumulate → faster progress
• Gradients that oscillate cancel out → dampens oscillations

# Momentum SGD (simplified)
v = 0
for step in range(num_steps):
    grad = compute_gradient(w)
    v = beta * v + grad           # Accumulate velocity
    w = w - lr * v                # Update weights

# PyTorch
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

Adam: Adaptive Moment Estimation

Key insight: Different parameters need different learning rates!

Adam combines two ideas:

1. Momentum (first moment): Accumulate gradient direction
2. RMSprop (second moment): Adapt learning rate per parameter based on gradient magnitude

\[m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t \quad \text{(momentum/first moment)}\] \[v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \quad \text{(squared gradients/second moment)}\] \[\hat{m}_t = m_t / (1-\beta_1^t) \quad \text{(bias correction)}\] \[\hat{v}_t = v_t / (1-\beta_2^t) \quad \text{(bias correction)}\] \[w_{t+1} = w_t - \alpha \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}\]

Typical values: \(\beta_1 = 0.9\), \(\beta_2 = 0.999\), \(\epsilon = 10^{-8}\)

Why it works:

• Parameters with consistently large gradients get smaller effective learning rate (denominator is big)
• Parameters with small gradients get larger effective learning rate (denominator is small)
• This “adapts” to each parameter’s scale automatically

Figure: Adam combines momentum (accumulating gradient direction) with adaptive learning rates (scaling by inverse gradient magnitude). Parameters with large, consistent gradients move faster in the right direction. Parameters with noisy gradients get dampened. The bias correction ensures early steps aren’t dominated by initialization.

# PyTorch
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

AdamW: Decoupled Weight Decay

The problem with Adam + L2: In Adam, L2 regularization gets scaled by the adaptive learning rate, weakening its effect.

AdamW decouples weight decay from the gradient-based update:

# Adam with L2 (problematic):
grad_with_l2 = grad + lambda * w
# ... Adam update using grad_with_l2
# Problem: weight decay is scaled by adaptive LR

# AdamW (correct):
# ... Adam update using grad only
w = w - lr * lambda * w  # Weight decay applied separately
# Weight decay applied consistently

# PyTorch
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)

AdamW is the default for modern transformers/LLMs.

Quick Comparison

When to Use Which?

The SGD vs Adam Generalization Debate

Empirical observation: SGD with momentum often finds solutions that generalize better than Adam, particularly for vision models.

Hypotheses:

• Adam converges to sharper minima (overfit)
• SGD’s noise helps escape bad minima
• Adam’s adaptivity can be too aggressive

In practice: For LLMs, AdamW wins. For vision, try both.

Deep dive: See Part 2 (Core Theory) and Part 5 (Optimization) for detailed mathematical treatment of convergence, learning rate schedules, and second-order methods.

1.22 Common Interview Questions Summary

Architecture & Training

Optimization

Regularization

Softmax & Classification

1.23 Debugging Checklist

When training doesn’t work, check these in order:

1. Loss is NaN or Inf

□ Learning rate too high? → Try 10x smaller
□ Numerical instability? → Check for log(0), exp overflow
□ Bad initialization? → Use proper init (He/Xavier)
□ Data issue? → Check for NaN/Inf in inputs
□ Exploding gradients? → Add gradient clipping

2. Loss Doesn’t Decrease

□ Learning rate too low? → Try 10x larger
□ Learning rate too high? → Loss oscillates, try smaller
□ Bug in loss computation? → Verify loss formula manually
□ Data not shuffled? → Enable shuffle in DataLoader
□ Model in eval mode? → Ensure model.train() during training
□ Gradients are zero? → Check gradient flow, dead ReLUs
□ Optimizer not stepping? → Verify optimizer.step() is called

3. Training Loss Good, Validation Loss Bad (Overfitting)

□ Add/increase dropout
□ Add/increase weight decay (L2)
□ Get more data
□ Data augmentation
□ Early stopping
□ Reduce model size
□ Use regularization (BatchNorm helps too)

4. Both Losses High (Underfitting)

□ Model too simple → More layers/neurons
□ Training long enough? → More epochs
□ Learning rate too low? → Increase LR
□ Too much regularization? → Reduce dropout/weight decay
□ Feature engineering needed? → Better input features

5. Quick Sanity Checks

# 1. Can model overfit single batch?
# If not, there's a bug in your model/training loop
small_batch = next(iter(train_loader))
for _ in range(100):
    loss = train_step(small_batch)
print(f"Should be near 0: {loss}")

# 2. Are gradients flowing?
loss.backward()
for name, param in model.named_parameters():
    if param.grad is not None:
        print(f"{name}: grad norm = {param.grad.norm():.4f}")
    else:
        print(f"{name}: NO GRADIENT!")

# 3. Are weights updating?
old_weights = model.fc1.weight.clone()
optimizer.step()
new_weights = model.fc1.weight
print(f"Weight change: {(new_weights - old_weights).abs().mean():.6f}")

Part 2: Core Theory

2.1 The Core Update Rule

The Question

In gradient descent, the update rule is:

\[ w_{t+1} = w_t - \alpha \nabla \ell(w_t) \]

Why do we pass \(w_t\) to the loss function while also subtracting from \(w_t\)?

Short Answer

Both uses of \(w_t\) refer to the same thing: your current position in parameter space. The update says:

“From where I currently am (\(w_t\)), take a step in the direction that reduces loss, where that direction is determined by the gradient computed at my current location (\(\nabla \ell(w_t)\)).”

Breaking Down Each Component

• \(w_t\) — Current parameters: where we are in weight space at step \(t\)
• \(\ell(w_t)\) — Loss evaluated at current parameters: how bad our model is right now
• \(\nabla \ell(w_t)\) — Gradient of loss at current parameters: direction of steepest increase
• \(-\alpha \nabla \ell(w_t)\) — The step we take: move opposite to gradient (toward lower loss)
• \(w_{t+1}\) — New parameters: where we’ll be after this update

2.2 Why Gradients Point the Way

Intuition: Hiking in the Fog

Imagine you’re hiking on a mountain in dense fog. You can’t see the valley below, but you want to get there. What do you do?

You feel the slope under your feet. If the ground tilts left, going left takes you uphill. So you go right — the opposite direction. You take a small step, feel the slope again, and repeat.

This is exactly what gradient descent does:

1. Feel the slope → compute the gradient at current position
2. Identify uphill → gradient points toward steepest ascent
   
3. Walk downhill → move in the negative gradient direction
4. Repeat → keep stepping until you reach flat ground (the minimum)

The gradient is like your feet sensing the tilt in every direction simultaneously. In high dimensions, there are millions of “directions” — but the gradient tells you the single steepest way up, so you go the opposite way.

The Gradient as a Direction

For a scalar function \(f: \mathbb{R}^n \to \mathbb{R}\), the gradient \(\nabla f(x)\) is a vector that:

1. Points in the direction of steepest increase of \(f\)
2. Has magnitude equal to the rate of increase in that direction

\[ \nabla f(x) = \begin{bmatrix} \frac{\partial f}{\partial x_1} \\ \frac{\partial f}{\partial x_2} \\ \vdots \\ \frac{\partial f}{\partial x_n} \end{bmatrix} \]

Directional Derivatives

The rate of change of \(f\) in direction \(u\) (unit vector) is:

\[ D_u f(x) = \nabla f(x) \cdot u = \|\nabla f(x)\| \cos(\theta) \]

where \(\theta\) is the angle between \(\nabla f(x)\) and \(u\).

Conclusion: To decrease \(f\) fastest, move opposite to \(\nabla f\).

Taylor Expansion Perspective

Expanding loss around current point \(w_t\):

\[ \ell(w_t + \Delta w) \approx \ell(w_t) + \nabla \ell(w_t)^\top \Delta w + \frac{1}{2} \Delta w^\top H \Delta w \]

where \(H = \nabla^2 \ell(w_t)\) is the Hessian matrix.

First-order methods (GD, SGD, Adam) use only the gradient term. Second-order methods (Newton, BFGS) also use curvature from \(H\).

2.3 Stochastic Gradient Descent (SGD)

Stochastic Gradient Descent (SGD) is a fundamental optimization algorithm used in machine learning and deep learning to minimize a model’s loss. Unlike standard Gradient Descent, which calculates the gradient over the entire dataset before updating parameters, SGD approximates this by using only a single random data point (or a small “mini-batch”) at each step.

How SGD Works

The algorithm iteratively adjusts model parameters \(\theta\) (weights and biases) to reach the “valley” or minimum of the loss function:

1. Initialize: Start with random parameter values
2. Random Selection: Pick one random sample (or a small batch) from the training data
3. Compute Gradient: Calculate the slope (gradient) of the loss function based only on that sample
4. Update: Adjust parameters in the opposite direction of the gradient using a learning rate \(\eta\): \[\theta := \theta - \eta \cdot \nabla_\theta \ell(x_i, y_i; \theta)\]
5. Repeat: Continue until the loss stops decreasing significantly (convergence)

Concrete Example: Linear Regression with SGD

To make this concrete, let’s apply SGD to linear regression. Suppose we want to fit a straight line:

\[\hat{y} = w_1 + w_2 x\]

to a training set with observations \(\{(x_1, y_1), (x_2, y_2), \ldots, (x_n, y_n)\}\).

The objective function (using least squares) is:

\[Q(w) = \sum_{i=1}^n Q_i(w) = \sum_{i=1}^n \left(\hat{y}_i - y_i\right)^2 = \sum_{i=1}^n \left(w_1 + w_2 x_i - y_i\right)^2\]

The SGD update for a single randomly selected point \((x_i, y_i)\):

\[ \begin{bmatrix} w_1 \\ w_2 \end{bmatrix} \leftarrow \begin{bmatrix} w_1 \\ w_2 \end{bmatrix} - \eta \begin{bmatrix} \frac{\partial}{\partial w_1} (w_1 + w_2 x_i - y_i)^2 \\ \frac{\partial}{\partial w_2} (w_1 + w_2 x_i - y_i)^2 \end{bmatrix} \]

Computing the partial derivatives:

\[ = \begin{bmatrix} w_1 \\ w_2 \end{bmatrix} - \eta \begin{bmatrix} 2 (w_1 + w_2 x_i - y_i) \\ 2 x_i(w_1 + w_2 x_i - y_i) \end{bmatrix} \]

Key insight: In each iteration, the gradient is evaluated at only a single point \((x_i, y_i)\) — not the entire dataset. This is what makes it “stochastic.”

Python implementation:

import numpy as np

def sgd_linear_regression(X, y, lr=0.01, epochs=100):
    """SGD for linear regression: y = w1 + w2 * x"""
    w1, w2 = 0.0, 0.0  # Initialize weights
    n = len(X)
    
    for epoch in range(epochs):
        # Shuffle data for randomness
        indices = np.random.permutation(n)
        
        for i in indices:
            # Prediction error for single point
            error = (w1 + w2 * X[i]) - y[i]
            
            # Update weights based on this ONE point
            w1 = w1 - lr * 2 * error           # ∂L/∂w1 = 2 * error
            w2 = w2 - lr * 2 * error * X[i]    # ∂L/∂w2 = 2 * error * x
    
    return w1, w2

# Example usage
X = np.array([1, 2, 3, 4, 5])
y = np.array([2.1, 4.0, 5.8, 8.1, 9.9])  # Approximately y = 2x

w1, w2 = sgd_linear_regression(X, y, lr=0.01, epochs=100)
print(f"Fitted line: y = {w1:.2f} + {w2:.2f}x")
# Output: Fitted line: y = 0.12 + 1.98x (close to y = 0 + 2x)

Figure: SGD training dynamics. Left: Loss decreases rapidly then stabilizes. Middle: Parameter trajectory from (0,0) to near the true values. Right: Fitted line closely matches true relationship.

Full-Batch vs Mini-Batch Gradient Descent

For a dataset with \(N\) examples, the true gradient is:

\[ \nabla \ell(\theta) = \frac{1}{N} \sum_{i=1}^{N} \nabla \ell_i(\theta) \]

Computing this exactly is expensive for large datasets. Instead, we estimate it using a mini-batch \(B\):

\[ g_t = \frac{1}{|B|} \sum_{i \in B} \nabla \ell_i(\theta_t) \]

Key property: \(\mathbb{E}[g_t] = \nabla \ell(\theta_t)\) — the mini-batch gradient is an unbiased estimator of the true gradient!

Why Use SGD?

1. Efficiency: SGD is significantly faster for massive datasets because it updates parameters immediately after seeing a few samples, rather than waiting to process millions of data points.

2. Memory Friendly: Only one batch needs to be in memory at a time, making it essential for training large neural networks that wouldn’t fit otherwise.

3. Escaping Local Minima: The “noise” introduced by random sampling causes the optimization path to zig-zag. This randomness can help the algorithm “jump out” of shallow local minima or saddle points to find better solutions.

Why Noise Can Help (Implicit Regularization)

Here’s a counterintuitive fact: the “noisiness” of SGD is actually a feature, not a bug.

When you use a small mini-batch, your gradient estimate bounces around the true gradient. You might think this would lead to worse solutions. But empirically, the opposite often happens: noisy SGD often finds better solutions than exact gradient descent.

Why? Think about it geometrically. Imagine the loss landscape has two valleys:

• A narrow, sharp valley — training loss is very low here, but the slightest change in weights causes loss to spike
• A wide, flat valley — training loss is similar, but nearby weights also have low loss

With exact gradients (large batch), you might descend into that sharp valley and stay there. But with noisy SGD, random fluctuations can bounce you out of the sharp valley. The wide, flat valley is more stable — noise doesn’t push you out because there’s more “room.”

This matters for generalization: test data is slightly different from training data. A sharp minimum that’s perfect for training data might be terrible for test data (small perturbations → big loss increases). A flat minimum is robust — exactly what you want.

Bottom line: SGD noise provides implicit regularization that biases training toward solutions that generalize better.

Batch Size Tradeoffs

Larger batches give more accurate gradient estimates but fewer updates per epoch. Smaller batches are noisier but provide more frequent updates and often better generalization.

SGD vs Full-Batch Gradient Descent

Key Variations

Mini-batch Gradient Descent: The modern standard that uses a small group of samples (e.g., 32, 64, 128) instead of just one. It balances the speed of SGD with more stable gradient estimates.

Momentum: Adds a “velocity” term that keeps updates moving in a consistent direction, reducing erratic zig-zagging and speeding up convergence. (See Section 2.4)

Adaptive Optimizers: Methods like Adam, RMSprop, and AdaGrad automatically adjust the learning rate for each parameter during training. (See Section 2.5)

PyTorch Implementation

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

# Basic SGD
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# SGD with momentum (recommended)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Training loop
for batch in dataloader:
    optimizer.zero_grad()          # Reset gradients
    loss = compute_loss(batch)     # Forward pass
    loss.backward()                # Compute gradients
    optimizer.step()               # Update parameters

2.4 Momentum Methods

The Problem: Oscillations in Ill-Conditioned Landscapes

Consider optimizing a loss surface shaped like an elongated valley:

Figure: Left: Vanilla SGD oscillates in ill-conditioned landscapes (high condition number). Right: Momentum dampens oscillations and converges faster.

The problem:

• Gradient points toward the steepest descent (across the valley)
• But the minimum is along the gentle direction (down the valley)
• SGD overshoots across the valley, oscillates back and forth

Classical Momentum

Add a velocity term that accumulates gradient direction:

\[ \begin{aligned} v_{t+1} &= \beta v_t + \nabla \ell(w_t) \\ w_{t+1} &= w_t - \alpha v_{t+1} \end{aligned} \]

Intuition: Ball Rolling Downhill

Figure: Without momentum (left), SGD oscillates back and forth. With momentum (right), the path is smooth and direct.

• Perpendicular oscillations cancel: \(+g\) then \(-g\) averages to 0
• Consistent direction accumulates: gradients pointing the same way add up
• Result: faster progress along the consistent direction

Effective Step Size

With momentum, the effective learning rate in a consistent direction is:

\[ \alpha_{\text{eff}} = \frac{\alpha}{1 - \beta} \]

For \(\beta = 0.9\): effective LR is 10× larger in consistent directions!

Why β = 0.9?

Nesterov Accelerated Gradient (NAG)

Key idea: “Look ahead” before computing the gradient.

Standard momentum computes gradient at current position, then moves. Nesterov computes gradient at the predicted next position:

\[ \begin{aligned} v_{t+1} &= \beta v_t + \nabla \ell(\underbrace{w_t - \alpha \beta v_t}_{\text{look-ahead position}}) \\ w_{t+1} &= w_t - \alpha v_{t+1} \end{aligned} \]

Standard Momentum:          Nesterov:
                           
w_t ──gradient──→ w_{t+1}   w_t ──lookahead──→ w̃ ──gradient──→ w_{t+1}
                           
                            "Where am I going? Let me check
                             the gradient there first."

Why It Helps

• Can “correct” before overshooting
• Better theoretical convergence rate
• In practice: similar to momentum, sometimes slightly better

PyTorch Implementation

import torch

# Standard momentum
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Nesterov momentum
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)

2.5 Adaptive Learning Rate Methods

The Problem with Fixed Learning Rates

Imagine you’re training a language model, and your vocabulary includes both “the” (appears in every sentence) and “quasar” (appears once in 10,000 sentences).

With a fixed learning rate, every weight gets updated with the same step size. But the embedding weights for “the” see gradients constantly — thousands of updates pushing them around. Meanwhile, “quasar” barely gets any signal.

What happens? Either: - Learning rate is tuned for frequent words → rare words learn too slowly - Learning rate is tuned for rare words → frequent words oscillate wildly

The insight: Different parameters need different learning rates. Frequent features need smaller updates (they’ve already seen lots of data). Rare features need larger updates (every bit of signal is precious).

This is what adaptive methods solve: they automatically adjust learning rates for each parameter based on the history of gradients that parameter has seen.

The Evolution: AdaGrad → RMSprop → Adam → AdamW

The story of adaptive optimizers is one of progressive refinement, each method solving a problem introduced by the previous one:

AdaGrad (2011)  →  RMSprop (2012)  →  Adam (2015)  →  AdamW (2019)
    ↓                   ↓                 ↓               ↓
 Per-param LR     + Decay factor    + Momentum      + Decoupled
                                                      weight decay

AdaGrad (2011): The First Adaptive Method

The idea: Keep a running sum of squared gradients for each parameter. Divide the learning rate by the square root of this sum.

\[ \begin{aligned} v_{t+1} &= v_t + g_t^2 \\ w_{t+1} &= w_t - \frac{\alpha}{\sqrt{v_{t+1}} + \epsilon} g_t \end{aligned} \]

What this achieves: Parameters that have seen large gradients (frequent features) accumulate a large \(v\), so their effective learning rate \(\alpha / \sqrt{v}\) shrinks. Rare features keep their original learning rate because they’ve accumulated little.

The fatal flaw: \(v\) only grows. It never shrinks. Over time, all learning rates decay toward zero and training stalls. This is fine for convex problems (where you want to converge and stop), but disastrous for deep learning where you need to keep adapting.

RMSprop (2012): Forgetting the Past

The fix: Instead of accumulating all past squared gradients, use an exponential moving average. This “forgets” old gradients, allowing learning rates to recover.

\[ \begin{aligned} v_{t+1} &= \beta v_t + (1-\beta) g_t^2 \\ w_{t+1} &= w_t - \frac{\alpha}{\sqrt{v_{t+1}} + \epsilon} g_t \end{aligned} \]

With \(\beta = 0.9\), we’re essentially looking at the average squared gradient over roughly the last 10 steps. If gradients suddenly become small (e.g., after the loss landscape changes), \(v\) will decay and learning rate can increase again.

The remaining issue: RMSprop doesn’t have momentum. We’re adapting the scale of updates, but not the direction. Can we get both?

Adam (2015): The Best of Both Worlds

The breakthrough: Combine RMSprop (adaptive learning rates) with momentum (smooth, accelerated updates). Add bias correction to handle initialization issues.

Think of Adam as having two “memories”: - First moment (\(m\)): Exponential average of gradients — the momentum term, pointing toward the consistent direction - Second moment (\(v\)): Exponential average of squared gradients — the RMSprop term, scaling by gradient magnitude

\[ \begin{aligned} m_t &= \beta_1 m_{t-1} + (1-\beta_1) g_t & \text{(momentum: which way?)} \\ v_t &= \beta_2 v_{t-1} + (1-\beta_2) g_t^2 & \text{(scale: how big?)} \\ \hat{m}_t &= m_t / (1 - \beta_1^t) & \text{(bias correction)} \\ \hat{v}_t &= v_t / (1 - \beta_2^t) & \text{(bias correction)} \\ w_t &= w_{t-1} - \alpha \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \end{aligned} \]

Note on indexing: Here \(t\) starts at 1 (step 1, step 2, …). At step \(t\), we compute \(m_t, v_t\) from the gradient \(g_t\), apply bias correction using \(\beta^t\), and update weights.

The update \(\frac{\hat{m}}{\sqrt{\hat{v}}}\) has a beautiful interpretation: we take the smoothed gradient direction (\(\hat{m}\)) and normalize it by the smoothed gradient magnitude (\(\sqrt{\hat{v}}\)). This gives us a kind of “standardized” step — regardless of whether gradients are large or small, the update magnitude stays roughly bounded.

Why these defaults? \(\beta_2 = 0.999\) is much closer to 1 than \(\beta_1 = 0.9\). This means Adam has a “long memory” for gradient magnitudes (scale changes slowly) but a “short memory” for gradient direction (adapts quickly to new directions). This makes sense: the scale of gradients in a layer tends to be stable, but the direction can change rapidly.

Why Bias Correction Matters

Here’s a subtle but important issue. Both \(m\) and \(v\) are initialized to zero. At step 1:

\[m_1 = 0.9 \times 0 + 0.1 \times g_1 = 0.1 \cdot g_1\]

We wanted an estimate of the gradient, but we got only 10% of it! The same happens with \(v\). Without correction, early updates would be way too small.

The fix divides by \((1 - \beta^t)\), which starts large and approaches 1 as \(t\) grows:

• Step 1: divide by \((1 - 0.9^1) = 0.1\) → multiply by 10 → \(\hat{m}_1 = g_1\) ✓
• Step 10: divide by \((1 - 0.9^{10}) \approx 0.65\)
• Step 100: divide by \((1 - 0.9^{100}) \approx 1.0\) → no correction needed

AdamW (2019): The Interview Question!

AdamW seems like a minor tweak, but understanding why it matters reveals deep insight into how adaptive optimizers interact with regularization.

The Problem: Adam Breaks Weight Decay

You want to regularize your model with L2 (weight decay). The standard approach: add \(\lambda w\) to the gradient.

\[g_{\text{regularized}} = g + \lambda w\]

With SGD, this works perfectly. The update becomes:

\[w \leftarrow w - \alpha(g + \lambda w) = (1 - \alpha\lambda)w - \alpha g\]

Every weight shrinks by the same factor \((1 - \alpha\lambda)\) at every step. Fair and consistent.

But Adam does something different. It divides the gradient by \(\sqrt{v}\):

\[\text{Update} = \frac{g + \lambda w}{\sqrt{v}}\]

Now the regularization term \(\lambda w\) also gets divided by \(\sqrt{v}\)! Parameters with large historical gradients (large \(v\)) get less regularization. Parameters with small gradients get more regularization.

This is backwards! We want consistent regularization across all parameters, but Adam’s adaptive scaling breaks it.

The Fix: Decouple Weight Decay

AdamW applies weight decay outside the adaptive mechanism:

\[w_{t+1} = w_t - \alpha \cdot \underbrace{\frac{\hat{m}_{t}}{\sqrt{\hat{v}_{t}} + \epsilon}}_{\text{Adam update}} - \underbrace{\alpha \lambda w_t}_{\text{separate weight decay}}\]

The key difference:

Adam + L2:    gradient ← gradient + λw,  then divide by √v
AdamW:        divide by √v,              then subtract αλw separately

Now every weight decays by exactly \(\alpha\lambda\) of its current value, regardless of gradient history. Simple, consistent, correct.

Why It Matters in Practice

This isn’t just theoretical. Empirically, AdamW consistently outperforms Adam + L2 on Transformers and other architectures. The BERT and GPT papers all use AdamW.

Bottom line: If you’re using Adam and want regularization, use AdamW, not Adam with weight_decay parameter (which does L2 the wrong way).

PyTorch Implementation

import torch

# Adam (don't use weight_decay with Adam!)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# AdamW (correct way to regularize)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)

# SGD with momentum
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4)

2.6 Learning Rate Schedules

Why Schedules Matter

Think of training a neural network like hiking down into a valley, but the terrain changes as you descend.

Early in training: The loss landscape is rough and you’re far from any minimum. Big steps are fine — even if they’re imprecise, they get you moving in the right direction quickly. A large learning rate lets you explore broadly and make rapid progress.

Late in training: You’re near a minimum now. The same big steps that helped earlier now cause you to overshoot and oscillate around the minimum. You need smaller, more careful steps to settle into the valley floor.

This is why a fixed learning rate is rarely optimal: - Too high throughout → Can’t converge precisely at the end (oscillates) - Too low throughout → Wastes time early on when you could move faster

The solution: Start with a high learning rate (fast progress through the easy terrain), then decay to a low rate (careful fine-tuning near the minimum).

Warmup: Why Transformers Need It

The Problem with Cold-Starting Adam

Adam tracks running estimates of gradient statistics (\(m\) for direction, \(v\) for scale). But at step 1, these estimates are based on a single gradient — they’re extremely noisy and unreliable.

Even with bias correction, the first few hundred steps can be problematic:

Step 1:  v₁ = 0.001 × g₁²      ← Based on ONE gradient
         update = g₁/√v₁       ← Scaling by noisy estimate → unstable!

For Transformers specifically, there’s another issue: the attention mechanism creates sharp, spiky loss landscapes early in training. Large learning rates on spiky terrain → disaster.

The Solution: Linear Warmup

Start with tiny LR, increase linearly:

\[ \alpha_t = \alpha_{\max} \cdot \frac{t}{T_{\text{warmup}}} \quad \text{for } t < T_{\text{warmup}} \]

Figure: Linear warmup gradually increases learning rate from 0 to max, stabilizing early training.

Typical warmup: 1-5% of total training steps (e.g., 2000 steps for 100k total)

Common Decay Schedules

Step Decay

Drop LR by factor at fixed milestones:

\[ \alpha_t = \alpha_0 \cdot \gamma^{\lfloor t / S \rfloor} \]

Figure: Step decay drops LR by 10× at epochs 30, 60, 90 (typical for CNN training).

Typical: Divide by 10 at epochs 30, 60, 90 (for 100 epoch training)

Cosine Annealing

Smooth decay following cosine curve:

\[ \alpha_t = \alpha_{\min} + \frac{1}{2}(\alpha_{\max} - \alpha_{\min})\left(1 + \cos\left(\frac{\pi t}{T}\right)\right) \]

Figure: Cosine annealing provides smooth LR decay, preferred for Transformers and LLMs.

Used by: GPT-3, LLaMA, most modern LLMs

Linear Decay

Simple linear decrease:

\[ \alpha_t = \alpha_{\max} \cdot \left(1 - \frac{t}{T}\right) \]

Comparison

Warmup + Cosine (Standard LLM Recipe)

Figure: Standard LLM training recipe — linear warmup followed by cosine decay.

PyTorch Implementation:

from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

# Warmup for first 2000 steps
warmup = LinearLR(optimizer, start_factor=0.01, total_iters=2000)

# Cosine decay for rest of training
cosine = CosineAnnealingLR(optimizer, T_max=100000 - 2000, eta_min=1e-5)

# Combine them
scheduler = SequentialLR(optimizer, [warmup, cosine], milestones=[2000])

# In training loop
for step in range(100000):
    loss.backward()
    optimizer.step()
    scheduler.step()  # Update LR

Optimizer + Schedule Decision Guide

Quick Recommendations

If unsure, start with:

# For Transformers / NLP
optimizer = AdamW(lr=1e-4, weight_decay=0.01)
scheduler = warmup(2000 steps) + cosine(to 1e-5)

# For CNNs / Vision  
optimizer = SGD(lr=0.1, momentum=0.9, weight_decay=1e-4)
scheduler = step(divide by 10 at 30%, 60%, 90% of training)

Decision Flowchart

┌─────────────────────────────────────────────────────────────────┐
│                OPTIMIZER & SCHEDULE DECISION GUIDE              │
└─────────────────────────────────────────────────────────────────┘
                                │
                 What are you training?
                                │
        ┌───────────┬───────────┼───────────┬───────────┐
        ↓           ↓           ↓           ↓           ↓
   Transformer     CNN      Fine-tune    RL (PPO)     GAN
        │           │           │           │           │
        ↓           ↓           ↓           ↓           ↓
     AdamW      SGD+Mom      AdamW       Adam        Adam
     warmup     step decay   linear     constant    constant
     cosine     0.1→0.001    decay       lr=3e-4    lr=1e-4
    lr=1e-4                  lr=2e-5

Complete Training Loop Example

Here’s a complete, minimal training loop using AdamW with warmup + cosine schedule — the standard recipe for Transformers:

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR

# =============================================================================
# 1. Model Definition
# =============================================================================
class SimpleTransformerBlock(nn.Module):
    """Minimal transformer-style block for demonstration."""
    def __init__(self, d_model=256, nhead=4, dim_ff=512, dropout=0.1):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
        self.ff = nn.Sequential(
            nn.Linear(d_model, dim_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim_ff, d_model),
            nn.Dropout(dropout),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
    
    def forward(self, x):
        # Self-attention with residual
        attn_out, _ = self.attn(x, x, x)
        x = self.norm1(x + attn_out)
        # Feed-forward with residual
        x = self.norm2(x + self.ff(x))
        return x

class SimpleModel(nn.Module):
    def __init__(self, vocab_size=1000, d_model=256, num_layers=2, num_classes=10):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.blocks = nn.ModuleList([SimpleTransformerBlock(d_model) for _ in range(num_layers)])
        self.classifier = nn.Linear(d_model, num_classes)
    
    def forward(self, x):
        x = self.embedding(x)
        for block in self.blocks:
            x = block(x)
        # Global average pooling + classification
        x = x.mean(dim=1)
        return self.classifier(x)

# =============================================================================
# 2. Training Configuration
# =============================================================================
# Hyperparameters
BATCH_SIZE = 32
LEARNING_RATE = 1e-4
WEIGHT_DECAY = 0.01
WARMUP_STEPS = 100
TOTAL_STEPS = 1000
GRAD_CLIP = 1.0

# Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# =============================================================================
# 3. Create Model, Optimizer, Scheduler
# =============================================================================
model = SimpleModel().to(device)

# AdamW with decoupled weight decay (correct way to regularize)
optimizer = AdamW(
    model.parameters(),
    lr=LEARNING_RATE,
    betas=(0.9, 0.999),
    weight_decay=WEIGHT_DECAY
)

# Warmup + Cosine schedule
warmup_scheduler = LinearLR(
    optimizer, 
    start_factor=0.01,  # Start at 1% of max LR
    total_iters=WARMUP_STEPS
)
cosine_scheduler = CosineAnnealingLR(
    optimizer,
    T_max=TOTAL_STEPS - WARMUP_STEPS,
    eta_min=1e-6  # Minimum LR
)
scheduler = SequentialLR(
    optimizer,
    schedulers=[warmup_scheduler, cosine_scheduler],
    milestones=[WARMUP_STEPS]
)

# Loss function
criterion = nn.CrossEntropyLoss()

# =============================================================================
# 4. Training Loop
# =============================================================================
def train_step(model, batch, optimizer, scheduler, criterion, grad_clip):
    """Single training step with all best practices."""
    inputs, targets = batch
    inputs, targets = inputs.to(device), targets.to(device)
    
    # Forward pass
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    
    # Backward pass
    loss.backward()
    
    # Gradient clipping (prevent exploding gradients)
    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    
    # Update weights and learning rate
    optimizer.step()
    scheduler.step()
    
    return loss.item()

def train(model, dataloader, optimizer, scheduler, criterion, total_steps, grad_clip=1.0):
    """Full training loop."""
    model.train()
    step = 0
    running_loss = 0.0
    
    while step < total_steps:
        for batch in dataloader:
            if step >= total_steps:
                break
            
            loss = train_step(model, batch, optimizer, scheduler, criterion, grad_clip)
            running_loss += loss
            step += 1
            
            # Logging
            if step % 100 == 0:
                avg_loss = running_loss / 100
                current_lr = scheduler.get_last_lr()[0]
                print(f"Step {step}/{total_steps} | Loss: {avg_loss:.4f} | LR: {current_lr:.2e}")
                running_loss = 0.0
    
    return model

# =============================================================================
# 5. Example Usage (with dummy data)
# =============================================================================
if __name__ == "__main__":
    # Create dummy dataset
    dummy_inputs = torch.randint(0, 1000, (320, 16))  # 320 samples, seq_len=16
    dummy_targets = torch.randint(0, 10, (320,))      # 10 classes
    dataset = TensorDataset(dummy_inputs, dummy_targets)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
    
    # Train
    model = train(model, dataloader, optimizer, scheduler, criterion, TOTAL_STEPS, GRAD_CLIP)
    print("Training complete!")

Key best practices demonstrated: 1. AdamW with decoupled weight decay (not Adam + L2) 2. Linear warmup → cosine decay schedule 3. Gradient clipping to prevent instability 4. Proper device handling (CPU/GPU) 5. Learning rate logging for debugging

2.7 Convergence and Loss Landscapes

The loss function defines a landscape over your parameter space — a surface with peaks, valleys, saddle points, and plateaus. Understanding this landscape is key to understanding why optimization works (or doesn’t).

What Does the Loss Landscape Look Like?

For neural networks, the loss landscape is incredibly high-dimensional — millions or billions of dimensions. We can’t visualize it directly, but research has revealed several important properties:

1. Many Local Minima, But They’re Usually Good

Early fears about neural networks centered on “getting stuck in bad local minima.” In practice, this rarely happens. Why?

• In high dimensions, most critical points are saddle points, not local minima
• True local minima tend to have similar loss values to the global minimum
• Different minima often generalize equally well

2. The Landscape is Non-Convex But Surprisingly Navigable

The loss surface is highly non-convex (unlike linear regression, where there’s one obvious valley). But SGD can usually find good solutions because:

• Gradient descent tends to avoid saddle points (noise helps escape)
• The “connected valley” hypothesis: good solutions are connected by paths of low loss

Sharp vs Flat Minima: Why It Matters for Generalization

Here’s a key insight that connects optimization to generalization:

Figure: Sharp minima (left) have large loss changes from small weight perturbations. Flat minima (right) are robust and generalize better.

Sharp minimum: The loss increases steeply when you move away from the optimal weights. Think of a narrow spike in the landscape.

Flat minimum: The loss stays low even when weights are perturbed slightly. Think of a broad, shallow valley.

Why does this matter? At test time, your data is different from training data — it’s like adding a small perturbation to everything. A model at a sharp minimum has weights that are “brittle” — small changes cause large loss increases. A model at a flat minimum has weights that are “robust” — nearby weight values also give low loss.

This is why flat minima generalize better: they’re stable under the perturbations that test data introduces.

How SGD Finds Flat Minima

Here’s the beautiful connection: SGD noise naturally biases toward flat minima.

In a sharp valley, the gradients are steep and variable — SGD’s noise bounces you out. In a flat valley, gradients are small and stable — you stay put.

It’s like rolling a ball with random kicks: - In a narrow groove, kicks bounce it out - In a wide basin, kicks just slosh it around without escaping

This is another form of implicit regularization from SGD (beyond what we discussed in section 2.3).

When Optimization “Converges”

What does it mean for training to converge? Several things:

1. Loss stabilizes: Stops decreasing meaningfully
2. Gradients shrink: \(\|\nabla L\| \to 0\) (at a critical point)
3. Weights stabilize: Updates become tiny

Signs of problems:

The Effect of Architecture on the Landscape

Different architectures create different loss landscapes:

ResNets (skip connections): Create smoother landscapes with more direct gradient paths. Easier to optimize.

Transformers: Attention creates sharp, spiky landscapes early in training. This is why warmup is critical.

Very Deep Networks (without skip connections): Pathological landscapes with vanishing gradients in most directions.

2.8 Practical Considerations

Gradient Clipping

Prevent exploding gradients:

\[ g \leftarrow \min(1, \frac{\tau}{\|g\|}) \cdot g \]

Weight Initialization

Debugging Checklist

2.9 Regularization: L1 and L2

The Problem: Overfitting

A model that fits training data too well may fail on new data:

Figure: Left: Training data. Middle: Overfitting (wiggly, fits noise). Right: Good fit (smooth, captures pattern).

Overfitting symptoms:

• Low training loss, high validation loss
• Model memorizes noise instead of learning patterns
• Weights become very large

The Solution: Penalize Large Weights

Add a regularization term to the loss:

\[ \text{Loss}_{\text{reg}} = \text{Loss}_{\text{data}} + \lambda \cdot R(w) \]

where \(\lambda\) controls regularization strength and \(R(w)\) penalizes model complexity.

L2 Regularization (Ridge / Weight Decay)

Formula

\[ R_{L2}(w) = \|w\|_2^2 = \sum_i w_i^2 \]

Full loss: \[ \text{Loss} = \frac{1}{N}\sum_{i=1}^{N} L(y_i, \hat{y}_i) + \frac{\lambda}{2} \sum_j w_j^2 \]

Effect on Gradient

\[ \frac{\partial \text{Loss}}{\partial w_j} = \frac{\partial L}{\partial w_j} + \lambda w_j \]

Update rule: \[ w_j \leftarrow w_j - \alpha\left(\frac{\partial L}{\partial w_j} + \lambda w_j\right) = (1 - \alpha\lambda)w_j - \alpha\frac{\partial L}{\partial w_j} \]

Interpretation: Each update shrinks weights toward zero by factor \((1 - \alpha\lambda)\).