Chapter 1: Optimization Fundamentals

This chapter builds the foundation for the entire series. You will learn what "optimization" means in machine learning, why the shape of the loss landscape matters, and how the family of gradient-based optimizers — from plain gradient descent to Adam — actually works. Every algorithm is implemented from scratch with NumPy so you can see exactly what happens at each update step.

Learning Objectives

By completing this chapter, you will master the following:

• ✅ Understand what optimization means in machine learning: minimizing a loss function over parameters
• ✅ Distinguish convex from non-convex loss landscapes, and explain local minima and saddle points
• ✅ Implement gradient descent from scratch and visualize its trajectory on 2D test functions
• ✅ Explain the differences between SGD, momentum, RMSProp, and Adam, including their update rules
• ✅ Choose and implement learning-rate schedules (step decay, cosine annealing, warmup)

1.1 What is Optimization in Machine Learning?

Learning is Minimization

Every time a machine learning model "learns", the same mathematical event takes place: a computer searches for parameter values that make a loss function (a numerical score of how wrong the model's predictions are) as small as possible. Training a neural network, fitting a regression line, and tuning a recommendation system are all instances of the same problem:

$$ \theta^* = \arg\min_\theta \mathcal{L}(\theta) $$

• $\theta$: the parameters (the adjustable numbers inside the model, such as weights and biases)
• $\mathcal{L}(\theta)$: the loss function, which maps a parameter setting to a single number measuring prediction error
• $\theta^*$: the optimal parameters — the setting that achieves the smallest loss
• $\arg\min$: "the argument that minimizes" — we want the location of the minimum, not just its value

This framing may feel abstract at first, so let us make it concrete with the simplest possible model: a straight line $\hat{y} = w x + b$ fitted to noisy data. Here $\theta = (w, b)$, and a natural loss is the mean squared error (MSE, the average of the squared differences between predictions and true values):

$$ \mathcal{L}(w, b) = \frac{1}{N} \sum_{i=1}^{N} \left( w x_i + b - y_i \right)^2 $$

The Parameter Space

Because $\theta = (w, b)$ has two components, we can picture $\mathcal{L}$ as a surface over a 2D plane: every point $(w, b)$ has a height equal to its loss. Optimization is the search for the lowest point of this surface. In deep learning the same picture holds, except the "plane" has millions of dimensions — we can no longer draw it, but the mathematics is identical.

Three Components of Every Optimization Problem

Code Example 1: Evaluating a Loss Surface

Let us generate synthetic data and evaluate the MSE loss over a grid of $(w, b)$ values. Seeing the loss as a lookup table over parameters is the single most useful mental model in this series.

# Requirements:
# - Python 3.9+
# - numpy>=1.24.0, <3.0.0

"""
Example 1: Evaluating an MSE loss surface over a parameter grid

Purpose: Show that a loss function assigns one number to every
         parameter setting, and that "training" = finding the lowest one
Target: Beginner
Execution time: under 5 seconds
Dependencies: NumPy only
"""

import numpy as np

# --- Synthetic data: y = 2x + 1 plus noise ---
rng = np.random.default_rng(42)
x_data = rng.uniform(-3, 3, 50)
y_data = 2.0 * x_data + 1.0 + rng.normal(0, 0.5, 50)

def mse_loss(w, b, x, y):
    """Mean squared error of the line y_hat = w*x + b."""
    y_pred = w * x + b
    return np.mean((y_pred - y) ** 2)

# --- Evaluate the loss on a grid of (w, b) values ---
w_grid = np.linspace(-1.0, 5.0, 121)
b_grid = np.linspace(-2.0, 4.0, 121)
loss_surface = np.zeros((len(b_grid), len(w_grid)))

for i, b in enumerate(b_grid):
    for j, w in enumerate(w_grid):
        loss_surface[i, j] = mse_loss(w, b, x_data, y_data)

# --- Locate the lowest point on the grid ---
i_min, j_min = np.unravel_index(np.argmin(loss_surface), loss_surface.shape)
print("=== Loss Surface Analysis ===")
print(f"Grid size: {loss_surface.shape[0]} x {loss_surface.shape[1]} "
      f"= {loss_surface.size} loss evaluations")
print(f"Best grid point: w = {w_grid[j_min]:.2f}, b = {b_grid[i_min]:.2f}")
print(f"Loss at best grid point: {loss_surface[i_min, j_min]:.4f}")
print(f"Loss at true parameters (w=2, b=1): {mse_loss(2.0, 1.0, x_data, y_data):.4f}")
print(f"Loss at a bad guess (w=-1, b=-2):   {mse_loss(-1.0, -2.0, x_data, y_data):.4f}")

Sample Output:

=== Loss Surface Analysis ===
Grid size: 121 x 121 = 14641 loss evaluations
Best grid point: w = 2.00, b = 0.90
Loss at best grid point: 0.1439
Loss at true parameters (w=2, b=1): 0.1506
Loss at a bad guess (w=-1, b=-2):   37.9308

Notice that the best grid point ($b = 0.90$) is not exactly the true value ($b = 1.0$): the noise in this particular sample of 50 points shifts the empirical minimum slightly. This is expected behavior, not a bug.

Two important observations follow from this experiment:

1. Grid search does not scale. With 2 parameters and 121 values each, we needed 14,641 evaluations. A model with just 10 parameters would need $121^{10} \approx 6.7 \times 10^{20}$ — completely infeasible. Modern networks have millions of parameters, so we need a smarter search strategy. That strategy is gradient descent, coming in Section 1.3.
2. The minimum loss is not zero. Because the data contains noise, even the true parameters cannot fit it perfectly. This is normal and healthy — a loss of exactly zero on noisy data usually signals overfitting (memorizing noise instead of learning the pattern).

1.2 Convexity and Loss Landscapes

Why the Shape of the Loss Matters

Whether optimization is easy or hard depends almost entirely on the loss landscape (the shape of the loss function viewed as a surface over parameter space). The key property is convexity.

A function $f$ is convex if the line segment between any two points on its graph never dips below the graph. Formally, for all $x, y$ and all $\lambda \in [0, 1]$:

$$ f(\lambda x + (1 - \lambda) y) \;\leq\; \lambda f(x) + (1 - \lambda) f(y) $$

Intuitively, a convex function is bowl-shaped. This gives an enormous practical guarantee:

For a convex function, every local minimum is a global minimum. If you keep walking downhill, you cannot get trapped in a wrong valley — there is only one valley.

Convex vs Non-Convex in Machine Learning

Local Minima and Saddle Points

Two kinds of "flat spots" (points where the gradient is zero) can stall an optimizer in non-convex landscapes:

• Local minimum: a point lower than all its neighbors, but possibly higher than the global minimum elsewhere. The optimizer stops here because every direction goes uphill.
• Saddle point: a point that is a minimum along some directions and a maximum along others — like the center of a horse saddle. The gradient is zero, yet the point is not a minimum at all. In high-dimensional spaces, saddle points vastly outnumber local minima, and escaping them is a central concern in deep learning. The noise in stochastic gradient descent (Section 1.4) turns out to help here.

Code Example 2: Visualizing Convex, Non-Convex, and Saddle Landscapes

# Requirements:
# - Python 3.9+
# - numpy>=1.24.0, <3.0.0
# - matplotlib>=3.7.0

"""
Example 2: Convex vs non-convex functions and a saddle point

Purpose: Visualize why non-convex landscapes are harder to optimize
Target: Beginner
Execution time: under 10 seconds
Dependencies: NumPy, Matplotlib
"""

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(-2.5, 2.5, 400)
convex = x ** 2                          # single bowl
nonconvex = x ** 4 - 3 * x ** 2 + x      # two valleys of different depth

# --- Find discrete local minima of the non-convex curve ---
is_local_min = np.logical_and(nonconvex[1:-1] <= nonconvex[:-2],
                              nonconvex[1:-1] <= nonconvex[2:])
min_idx = np.where(is_local_min)[0] + 1

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

axes[0].plot(x, convex, color='tab:blue')
axes[0].set_title('Convex: $f(x) = x^2$\n(one valley, easy)')
axes[0].set_xlabel('x')
axes[0].set_ylabel('f(x)')

axes[1].plot(x, nonconvex, color='tab:orange')
axes[1].scatter(x[min_idx], nonconvex[min_idx], color='red', zorder=3,
                label='local minima')
axes[1].set_title('Non-convex: $f(x) = x^4 - 3x^2 + x$\n(two valleys, one is a trap)')
axes[1].set_xlabel('x')
axes[1].legend()

# --- Saddle: z = x^2 - y^2 (minimum along x, maximum along y) ---
xx, yy = np.meshgrid(np.linspace(-2, 2, 200), np.linspace(-2, 2, 200))
zz = xx ** 2 - yy ** 2
cs = axes[2].contour(xx, yy, zz, levels=15, cmap='coolwarm')
axes[2].scatter([0], [0], color='black', zorder=3, label='saddle at (0, 0)')
axes[2].set_title('Saddle: $f(x, y) = x^2 - y^2$\n(gradient is zero, but not a minimum)')
axes[2].set_xlabel('x')
axes[2].set_ylabel('y')
axes[2].legend()

plt.tight_layout()
plt.savefig('landscapes.png', dpi=110)
plt.close()

print("=== Local Minima of the Non-Convex Function ===")
for i in min_idx:
    print(f"x = {x[i]:+.3f}, f(x) = {nonconvex[i]:+.3f}")
print("The deeper valley is the global minimum; the shallower one is a trap.")
print("Figure saved to landscapes.png")

Sample Output:

=== Local Minima of the Non-Convex Function ===
x = -1.297, f(x) = -3.514
x = +1.134, f(x) = -1.070
The deeper valley is the global minimum; the shallower one is a trap.
Figure saved to landscapes.png

A gradient-based optimizer started at $x = 2$ will slide into the shallow right valley at $x \approx 1.13$ and stay there, never discovering the deeper valley at $x \approx -1.30$. This is the fundamental limitation of purely local methods — and the motivation for the metaheuristic methods in Chapter 2 and Bayesian optimization in Chapter 3, which are designed to search globally.

You might reasonably worry that deep learning is doomed, since neural network losses are highly non-convex. In practice, research has shown that in very high-dimensional landscapes most local minima reachable by SGD have similarly good loss values, so "getting trapped" is less catastrophic than the 1D picture suggests. Still, understanding the geometry is essential for diagnosing training problems.

1.3 Gradient Descent from Scratch

The Core Idea

The gradient (the vector of partial derivatives $\nabla \mathcal{L}(\theta)$, which points in the direction of steepest increase of the loss) gives us exactly what grid search lacked: a local compass. To decrease the loss, step in the opposite direction of the gradient:

$$ \theta_{t+1} = \theta_t - \eta \, \nabla \mathcal{L}(\theta_t) $$

• $\theta_t$: parameters at step $t$
• $\eta$ (eta): the learning rate (the step size — how far we move per update)
• $\nabla \mathcal{L}(\theta_t)$: gradient of the loss at the current parameters

The learning rate is the single most important hyperparameter in this entire series:

Two Test Functions

We will exercise our implementation on two classic 2D test functions:

1. Ill-conditioned quadratic: $f(x, y) = x^2 + 10 y^2$. Convex, but the valley is 10 times steeper in $y$ than in $x$. This mismatch of curvatures — called ill-conditioning — forces a small learning rate and causes zigzagging.
2. Rosenbrock function: $f(x, y) = (1 - x)^2 + 100 (y - x^2)^2$. A famous benchmark with a narrow, curved, banana-shaped valley. The global minimum is at $(1, 1)$. Gradient descent finds the valley quickly but then crawls along it very slowly — an honest preview of why we will need momentum and adaptive methods.

Code Example 3: Gradient Descent with Trajectory Visualization

# Requirements:
# - Python 3.9+
# - numpy>=1.24.0, <3.0.0
# - matplotlib>=3.7.0

"""
Example 3: Gradient descent from scratch on a quadratic and Rosenbrock

Purpose: Implement the update rule theta = theta - lr * grad and
         visualize the optimization path on contour plots
Target: Beginner-Intermediate
Execution time: 10-20 seconds
Dependencies: NumPy, Matplotlib
"""

import numpy as np
import matplotlib.pyplot as plt

# --- Test function 1: ill-conditioned quadratic ---
def quad(p):
    x, y = p
    return x ** 2 + 10 * y ** 2

def quad_grad(p):
    x, y = p
    return np.array([2 * x, 20 * y])

# --- Test function 2: Rosenbrock ---
def rosenbrock(p):
    x, y = p
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

def rosenbrock_grad(p):
    x, y = p
    return np.array([-2 * (1 - x) - 400 * x * (y - x ** 2),
                     200 * (y - x ** 2)])

def gradient_descent(grad_fn, p0, lr, n_steps):
    """Plain gradient descent. Returns the full path of visited points."""
    p = np.array(p0, dtype=float)
    path = [p.copy()]
    for _ in range(n_steps):
        p = p - lr * grad_fn(p)
        path.append(p.copy())
    return np.array(path)

# --- Run on the quadratic (note: lr must satisfy lr < 2/20 = 0.1 to converge) ---
path_quad = gradient_descent(quad_grad, p0=(-9.0, 2.5), lr=0.09, n_steps=60)

# --- Run on Rosenbrock (steep walls force a tiny learning rate) ---
path_rosen = gradient_descent(rosenbrock_grad, p0=(-1.0, 1.0), lr=0.001,
                              n_steps=2000)

# --- Contour plots with trajectories ---
fig, axes = plt.subplots(1, 2, figsize=(13, 5))

gx, gy = np.meshgrid(np.linspace(-10, 10, 300), np.linspace(-4, 4, 300))
axes[0].contour(gx, gy, gx ** 2 + 10 * gy ** 2, levels=25, cmap='viridis')
axes[0].plot(path_quad[:, 0], path_quad[:, 1], 'r.-', linewidth=1,
             markersize=3, label='GD path')
axes[0].scatter([0], [0], marker='*', s=180, color='gold',
                edgecolor='black', zorder=3, label='minimum (0, 0)')
axes[0].set_title('Quadratic $x^2 + 10y^2$: zigzag from ill-conditioning')
axes[0].legend()

rx, ry = np.meshgrid(np.linspace(-1.6, 1.6, 300), np.linspace(-0.6, 1.6, 300))
rz = (1 - rx) ** 2 + 100 * (ry - rx ** 2) ** 2
axes[1].contour(rx, ry, rz, levels=np.logspace(-1, 3, 20), cmap='viridis')
axes[1].plot(path_rosen[::20, 0], path_rosen[::20, 1], 'r.-',
             linewidth=1, markersize=3, label='GD path (every 20th step)')
axes[1].scatter([1], [1], marker='*', s=180, color='gold',
                edgecolor='black', zorder=3, label='minimum (1, 1)')
axes[1].set_title('Rosenbrock: fast into the valley, slow along it')
axes[1].legend()

plt.tight_layout()
plt.savefig('gradient_descent_paths.png', dpi=110)
plt.close()

print("=== Gradient Descent Results ===")
print(f"Quadratic:  start loss = {quad(path_quad[0]):9.4f}, "
      f"final loss = {quad(path_quad[-1]):.2e} after {len(path_quad) - 1} steps")
print(f"Rosenbrock: start loss = {rosenbrock(path_rosen[0]):9.4f}, "
      f"final loss = {rosenbrock(path_rosen[-1]):.2e} after {len(path_rosen) - 1} steps")
print(f"Rosenbrock final point: ({path_rosen[-1][0]:.4f}, {path_rosen[-1][1]:.4f})"
      f"  (true minimum is (1, 1))")
print("Figure saved to gradient_descent_paths.png")

Sample Output:

=== Gradient Descent Results ===
Quadratic:  start loss =  143.5000, final loss = 3.83e-09 after 60 steps
Rosenbrock: start loss =    4.0000, final loss = 7.42e-02 after 2000 steps
Rosenbrock final point: (0.7279, 0.5286)  (true minimum is (1, 1))

Study the two panels carefully — they summarize the two chronic diseases of plain gradient descent:

1. Zigzagging on ill-conditioned problems. On the quadratic, the path bounces across the narrow axis while creeping along the shallow one. The learning rate is capped by the steepest direction ($\eta \lt 2/20 = 0.1$ here), so progress in the shallow direction is slow.
2. Crawling along curved valleys. On Rosenbrock, 2,000 steps leave us far from the minimum (loss still 0.074). The gradient along the valley floor is tiny, so each step barely moves — reaching the minimum takes roughly 20,000 steps at this learning rate, and the rate cannot be raised because the steep valley walls would cause divergence.

Both diseases are treated by the optimizers in Section 1.5. But first, we must confront a practical issue: computing the exact gradient over a huge dataset is expensive.

1.4 Stochastic Gradient Descent and Mini-Batching

The Cost Problem

In machine learning the loss is an average over $N$ training examples: $\mathcal{L}(\theta) = \frac{1}{N} \sum_{i=1}^{N} \ell_i(\theta)$, where $\ell_i$ is the loss on example $i$. Computing the exact ("full-batch") gradient touches all $N$ examples per step. With millions of examples, that is prohibitively slow.

Stochastic gradient descent (SGD, gradient descent using a small random subset of the data to estimate the gradient at each step) solves this by replacing the exact gradient with an estimate from a mini-batch (a small random subset of $B$ training examples):

$$ \theta_{t+1} = \theta_t - \eta \cdot \frac{1}{B} \sum_{i \in \text{batch}} \nabla \ell_i(\theta_t) $$

The Variance / Noise Trade-off

The mini-batch gradient is an unbiased estimator (on average over random batches it equals the true gradient), but any single batch gives a noisy estimate. The batch size $B$ controls a fundamental trade-off:

The variance of the gradient estimate scales roughly as $1/B$: quadrupling the batch size halves the noise standard deviation. Interestingly, some noise is beneficial in non-convex landscapes — it can kick the optimizer out of saddle points and shallow local minima. This is one reason SGD generalizes well in deep learning.

One more term you will meet constantly: an epoch (one complete pass through the entire training dataset). With $N = 200$ examples and $B = 32$, one epoch consists of $\lceil 200 / 32 \rceil = 7$ update steps.

Code Example 4: Batch vs Mini-Batch vs Stochastic on Linear Regression

# Requirements:
# - Python 3.9+
# - numpy>=1.24.0, <3.0.0
# - matplotlib>=3.7.0

"""
Example 4: Comparing batch, mini-batch, and stochastic gradient descent

Purpose: Show the noise/cost trade-off controlled by batch size on a
         synthetic linear regression problem
Target: Beginner-Intermediate
Execution time: 10-20 seconds
Dependencies: NumPy, Matplotlib
"""

import numpy as np
import matplotlib.pyplot as plt

# --- Synthetic regression data: y = 3x - 1.5 + noise ---
rng = np.random.default_rng(0)
n_samples = 200
X = rng.uniform(-2, 2, n_samples)
y = 3.0 * X - 1.5 + rng.normal(0, 0.5, n_samples)

def predict(theta, x):
    return theta[0] * x + theta[1]           # theta = [w, b]

def full_loss(theta):
    return 0.5 * np.mean((predict(theta, X) - y) ** 2)

def batch_grad(theta, idx):
    """Gradient of the loss computed on the examples in idx."""
    err = predict(theta, X[idx]) - y[idx]
    return np.array([np.mean(err * X[idx]), np.mean(err)])

def train(batch_size, lr=0.1, n_epochs=40, seed=1):
    """Train with the given batch size; record full-data loss each epoch."""
    local_rng = np.random.default_rng(seed)
    theta = np.zeros(2)
    history = [full_loss(theta)]
    n_updates = 0
    for _ in range(n_epochs):
        perm = local_rng.permutation(n_samples)
        for start in range(0, n_samples, batch_size):
            idx = perm[start:start + batch_size]
            theta = theta - lr * batch_grad(theta, idx)
            n_updates += 1
        history.append(full_loss(theta))
    return theta, np.array(history), n_updates

configs = [("Full batch (B=200)", n_samples),
           ("Mini-batch (B=32)", 32),
           ("Stochastic (B=1)", 1)]

print("=== Batch Size Comparison (40 epochs, lr=0.1) ===")
plt.figure(figsize=(8, 5))
for label, B in configs:
    theta, history, n_updates = train(batch_size=B)
    plt.semilogy(history, label=label)
    print(f"{label:22s} final loss = {history[-1]:.5f}, "
          f"w = {theta[0]:.3f}, b = {theta[1]:.3f}, updates = {n_updates}")

plt.xlabel('Epoch')
plt.ylabel('Full-data loss (log scale)')
plt.title('Effect of batch size on convergence')
plt.legend()
plt.tight_layout()
plt.savefig('batch_size_comparison.png', dpi=110)
plt.close()
print("True parameters: w = 3.000, b = -1.500")
print("Figure saved to batch_size_comparison.png")

Sample Output:

=== Batch Size Comparison (40 epochs, lr=0.1) ===
Full batch (B=200)     final loss = 0.13114, w = 2.953, b = -1.500, updates = 40
Mini-batch (B=32)      final loss = 0.13078, w = 2.982, b = -1.569, updates = 280
Stochastic (B=1)       final loss = 0.13778, w = 3.056, b = -1.616, updates = 8000

Read the three curves in the saved figure together with the update counts:

• Full batch descends smoothly but makes only 1 update per epoch — on large datasets each of those updates is very expensive.
• Mini-batch reaches essentially the same loss with 7 cheap updates per epoch, converging in far less wall-clock time on real problems. This is why $B$ in the tens-to-hundreds range is the default everywhere.
• Stochastic ($B=1$) makes rapid early progress but then rattles around a noise floor: with a constant learning rate, the parameter never settles because every single-example gradient yanks it in a different direction. Reducing $\eta$ over time fixes this — precisely the topic of Section 1.6.

1.5 Momentum, RMSProp, and Adam

We now treat the two diseases diagnosed in Section 1.3. Each modern optimizer adds a small amount of per-parameter memory to the plain SGD update.

Momentum: A Heavy Ball Rolling Downhill

Momentum (an exponentially decaying moving average of past gradients, used as the update direction) treats the parameter vector like a heavy ball: gradients act as forces, and velocity accumulates.

$$ \begin{aligned} v_{t+1} &= \gamma \, v_t + \eta \, \nabla \mathcal{L}(\theta_t) \\ \theta_{t+1} &= \theta_t - v_{t+1} \end{aligned} $$

• $v$: the velocity vector (initialized to zero)
• $\gamma$ (gamma): the momentum coefficient, typically $0.9$ — how much past velocity is retained

Why it helps: along a consistent downhill direction (the valley floor) the velocity keeps growing, up to $1/(1-\gamma) = 10\times$ the plain step. Across the valley (the zigzag direction) successive gradients point in opposite directions and cancel inside $v$. Momentum therefore accelerates the slow direction and damps the oscillation — attacking both diseases at once.

RMSProp: A Per-Parameter Learning Rate

RMSProp (Root Mean Square Propagation, an adaptive method that divides each parameter's step by a running estimate of its typical gradient magnitude) keeps an exponential moving average of squared gradients:

$$ \begin{aligned} s_{t+1} &= \rho \, s_t + (1 - \rho) \, \big(\nabla \mathcal{L}(\theta_t)\big)^2 \\ \theta_{t+1} &= \theta_t - \frac{\eta}{\sqrt{s_{t+1}} + \epsilon} \, \nabla \mathcal{L}(\theta_t) \end{aligned} $$

• $s$: running average of squared gradients, element-wise (initialized to zero)
• $\rho$ (rho): decay rate, typically $0.9$
• $\epsilon$: a tiny constant (e.g., $10^{-8}$) that prevents division by zero
• All operations are element-wise: each parameter gets its own effective learning rate

Parameters with persistently large gradients (steep directions) get their steps shrunk; parameters with small gradients (shallow directions) get relatively larger steps. This directly fixes ill-conditioning without hand-tuning $\eta$ per direction.

Adam: Momentum + RMSProp + Bias Correction

Adam (Adaptive Moment Estimation) combines both ideas. It tracks the first moment $m$ (mean of gradients, like momentum) and the second moment $v$ (mean of squared gradients, like RMSProp), and corrects the startup bias caused by initializing both at zero:

$$ \begin{aligned} m_{t+1} &= \beta_1 m_t + (1 - \beta_1) \nabla \mathcal{L}(\theta_t) \\ v_{t+1} &= \beta_2 v_t + (1 - \beta_2) \big(\nabla \mathcal{L}(\theta_t)\big)^2 \\ \hat{m} &= \frac{m_{t+1}}{1 - \beta_1^{t+1}}, \qquad \hat{v} = \frac{v_{t+1}}{1 - \beta_2^{t+1}} \\ \theta_{t+1} &= \theta_t - \frac{\eta \, \hat{m}}{\sqrt{\hat{v}} + \epsilon} \end{aligned} $$

The defaults $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$ work well across a remarkable range of problems, which is why Adam is the most common default optimizer in deep learning today.

Code Example 5: Four Optimizers from Scratch on Rosenbrock

# Requirements:
# - Python 3.9+
# - numpy>=1.24.0, <3.0.0
# - matplotlib>=3.7.0

"""
Example 5: SGD, Momentum, RMSProp, and Adam implemented from scratch

Purpose: Compare the convergence paths of four optimizers on the
         Rosenbrock function
Target: Intermediate
Execution time: 10-30 seconds
Dependencies: NumPy, Matplotlib
"""

import numpy as np
import matplotlib.pyplot as plt

def rosenbrock(p):
    x, y = p
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

def rosenbrock_grad(p):
    x, y = p
    return np.array([-2 * (1 - x) - 400 * x * (y - x ** 2),
                     200 * (y - x ** 2)])

def run_sgd(p0, lr=0.001, n_steps=5000):
    p = np.array(p0, dtype=float)
    path = [p.copy()]
    for _ in range(n_steps):
        p = p - lr * rosenbrock_grad(p)
        path.append(p.copy())
    return np.array(path)

def run_momentum(p0, lr=0.0002, gamma=0.9, n_steps=5000):
    p = np.array(p0, dtype=float)
    v = np.zeros_like(p)
    path = [p.copy()]
    for _ in range(n_steps):
        v = gamma * v + lr * rosenbrock_grad(p)
        p = p - v
        path.append(p.copy())
    return np.array(path)

def run_rmsprop(p0, lr=0.002, rho=0.9, eps=1e-8, n_steps=5000):
    p = np.array(p0, dtype=float)
    s = np.zeros_like(p)
    path = [p.copy()]
    for _ in range(n_steps):
        g = rosenbrock_grad(p)
        s = rho * s + (1 - rho) * g ** 2
        p = p - lr * g / (np.sqrt(s) + eps)
        path.append(p.copy())
    return np.array(path)

def run_adam(p0, lr=0.02, beta1=0.9, beta2=0.999, eps=1e-8, n_steps=5000):
    p = np.array(p0, dtype=float)
    m = np.zeros_like(p)
    v = np.zeros_like(p)
    path = [p.copy()]
    for t in range(1, n_steps + 1):
        g = rosenbrock_grad(p)
        m = beta1 * m + (1 - beta1) * g
        v = beta2 * v + (1 - beta2) * g ** 2
        m_hat = m / (1 - beta1 ** t)
        v_hat = v / (1 - beta2 ** t)
        p = p - lr * m_hat / (np.sqrt(v_hat) + eps)
        path.append(p.copy())
    return np.array(path)

p0 = (-1.0, 1.0)
runs = [("SGD (lr=0.001)", run_sgd(p0), 'tab:red'),
        ("Momentum (lr=0.0002, gamma=0.9)", run_momentum(p0), 'tab:blue'),
        ("RMSProp (lr=0.002)", run_rmsprop(p0), 'tab:green'),
        ("Adam (lr=0.02)", run_adam(p0), 'tab:purple')]

# --- Contour plot with all four paths ---
rx, ry = np.meshgrid(np.linspace(-1.6, 1.6, 300), np.linspace(-0.6, 1.6, 300))
rz = (1 - rx) ** 2 + 100 * (ry - rx ** 2) ** 2

plt.figure(figsize=(9, 6))
plt.contour(rx, ry, rz, levels=np.logspace(-1, 3, 20), cmap='gray')
for label, path, color in runs:
    plt.plot(path[::50, 0], path[::50, 1], '.-', color=color,
             linewidth=1, markersize=3, label=label)
plt.scatter([1], [1], marker='*', s=200, color='gold',
            edgecolor='black', zorder=3, label='minimum (1, 1)')
plt.title('Optimizer paths on the Rosenbrock function (5000 steps)')
plt.xlabel('x')
plt.ylabel('y')
plt.legend(fontsize=8)
plt.tight_layout()
plt.savefig('optimizer_comparison.png', dpi=110)
plt.close()

print("=== Final Results after 5000 Steps ===")
for label, path, _ in runs:
    p_final = path[-1]
    print(f"{label:34s} point = ({p_final[0]:+.4f}, {p_final[1]:+.4f}), "
          f"loss = {rosenbrock(p_final):.2e}")
print("Figure saved to optimizer_comparison.png")

Sample Output (your exact numbers will match, since there is no randomness):

=== Final Results after 5000 Steps ===
SGD (lr=0.001)                     point = (+0.9398, +0.8830), loss = 3.63e-03
Momentum (lr=0.0002, gamma=0.9)    point = (+0.9927, +0.9855), loss = 5.28e-05
RMSProp (lr=0.002)                 point = (+0.9960, +0.9949), loss = 9.08e-04
Adam (lr=0.02)                     point = (+0.9999, +0.9999), loss = 3.40e-09

Adam reaches the minimum to nine decimal places while plain SGD is still six orders of magnitude behind, and momentum alone already buys a 70-fold improvement over SGD. Two honest caveats are worth stating. First, these results depend on tuning: RMSProp with lr=0.01 actually ends up worse than SGD on this function (loss $2.2 \times 10^{-2}$), because without momentum its normalized steps keep bouncing across the valley — try it yourself. Every optimizer's learning rate here was chosen by a small manual search. Second, Rosenbrock is a single deterministic function, not a full ML training run; on real tasks with noisy mini-batch gradients, well-tuned SGD with momentum sometimes generalizes better than Adam. The practical advice is: start with Adam, and consider SGD + momentum + a good schedule when squeezing out final performance.

1.6 Learning-Rate Schedules

Why a Fixed Learning Rate is Rarely Optimal

Section 1.4 showed that SGD with a constant learning rate rattles around a noise floor: the step size that was ideal early in training is too large near the minimum. A learning-rate schedule (a rule that changes the learning rate over the course of training) resolves this tension: start large to make fast progress, end small to settle precisely.

Three Standard Schedules

1. Step decay — multiply the rate by a factor $\gamma \lt 1$ every $s$ epochs:

$$ \eta_t = \eta_0 \cdot \gamma^{\lfloor t / s \rfloor} $$

2. Cosine annealing — decay smoothly along a half-cosine from $\eta_0$ to $\eta_{\min}$ over $T$ epochs:

$$ \eta_t = \eta_{\min} + \frac{1}{2} (\eta_0 - \eta_{\min}) \left( 1 + \cos \frac{\pi t}{T} \right) $$

3. Warmup (+ cosine) — ramp the rate linearly from near zero to $\eta_0$ over the first few epochs, then decay. Warmup (a brief initial phase of gradually increasing learning rate) prevents violent early updates when the parameters are still random and gradients are large or poorly scaled. It is standard practice when training Transformers.

Code Example 6: Plotting the Schedules

# Requirements:
# - Python 3.9+
# - numpy>=1.24.0, <3.0.0
# - matplotlib>=3.7.0

"""
Example 6: Implementing and plotting learning-rate schedules

Purpose: Implement step decay, cosine annealing, and warmup + cosine
         as plain functions of the epoch index
Target: Beginner-Intermediate
Execution time: under 5 seconds
Dependencies: NumPy, Matplotlib
"""

import numpy as np
import matplotlib.pyplot as plt

def step_decay(epoch, lr0=0.5, drop=0.5, every=20):
    """Multiply lr by `drop` every `every` epochs."""
    return lr0 * (drop ** (epoch // every))

def cosine_schedule(epoch, lr0=0.5, total=80, lr_min=0.005):
    """Smooth half-cosine decay from lr0 to lr_min over `total` epochs."""
    progress = min(epoch / total, 1.0)
    return lr_min + 0.5 * (lr0 - lr_min) * (1 + np.cos(np.pi * progress))

def warmup_cosine(epoch, lr0=0.5, total=80, warmup=8, lr_min=0.005):
    """Linear warmup for `warmup` epochs, then cosine decay."""
    if epoch < warmup:
        return lr0 * (epoch + 1) / warmup
    progress = (epoch - warmup) / max(1, total - warmup)
    return lr_min + 0.5 * (lr0 - lr_min) * (1 + np.cos(np.pi * min(progress, 1.0)))

total_epochs = 80
epochs = np.arange(total_epochs)

schedules = [("Constant", [0.5 for _ in epochs]),
             ("Step decay (x0.5 every 20)", [step_decay(e) for e in epochs]),
             ("Cosine annealing", [cosine_schedule(e) for e in epochs]),
             ("Warmup (8) + cosine", [warmup_cosine(e) for e in epochs])]

plt.figure(figsize=(8, 5))
for label, lrs in schedules:
    plt.plot(epochs, lrs, label=label)
plt.xlabel('Epoch')
plt.ylabel('Learning rate')
plt.title('Learning-rate schedules')
plt.legend()
plt.tight_layout()
plt.savefig('lr_schedules.png', dpi=110)
plt.close()

print("=== Learning Rate at Selected Epochs ===")
print(f"{'Epoch':>6s} {'Step':>8s} {'Cosine':>8s} {'Warmup+Cos':>11s}")
for e in [0, 4, 10, 20, 40, 60, 79]:
    print(f"{e:6d} {step_decay(e):8.4f} {cosine_schedule(e):8.4f} "
          f"{warmup_cosine(e):11.4f}")
print("Figure saved to lr_schedules.png")

Sample Output:

=== Learning Rate at Selected Epochs ===
 Epoch     Step   Cosine  Warmup+Cos
     0   0.5000   0.5000      0.0625
     4   0.5000   0.4970      0.3125
    10   0.5000   0.4812      0.4991
    20   0.2500   0.4275      0.4668
    40   0.1250   0.2525      0.2955
    60   0.0625   0.0775      0.0934
    79   0.0625   0.0052      0.0052

Code Example 7: How Schedules Change Convergence

Now we attach these schedules to mini-batch SGD on the regression problem from Section 1.4, deliberately starting with a learning rate that is too large to settle ($\eta_0 = 0.8$).

# Requirements:
# - Python 3.9+
# - numpy>=1.24.0, <3.0.0
# - matplotlib>=3.7.0

"""
Example 7: Effect of learning-rate schedules on SGD convergence

Purpose: Show that decaying the learning rate lets noisy SGD settle
         far below the noise floor of a constant learning rate
Target: Intermediate
Execution time: 10-20 seconds
Dependencies: NumPy, Matplotlib
"""

import numpy as np
import matplotlib.pyplot as plt

# --- Same synthetic regression data as Example 4 ---
rng = np.random.default_rng(0)
n_samples = 200
X = rng.uniform(-2, 2, n_samples)
y = 3.0 * X - 1.5 + rng.normal(0, 0.5, n_samples)

def predict(theta, x):
    return theta[0] * x + theta[1]

def full_loss(theta):
    return 0.5 * np.mean((predict(theta, X) - y) ** 2)

def batch_grad(theta, idx):
    err = predict(theta, X[idx]) - y[idx]
    return np.array([np.mean(err * X[idx]), np.mean(err)])

def step_decay(epoch, lr0=0.8, drop=0.5, every=20):
    return lr0 * (drop ** (epoch // every))

def cosine_schedule(epoch, lr0=0.8, total=80, lr_min=0.005):
    progress = min(epoch / total, 1.0)
    return lr_min + 0.5 * (lr0 - lr_min) * (1 + np.cos(np.pi * progress))

def warmup_cosine(epoch, lr0=0.8, total=80, warmup=8, lr_min=0.005):
    if epoch < warmup:
        return lr0 * (epoch + 1) / warmup
    progress = (epoch - warmup) / max(1, total - warmup)
    return lr_min + 0.5 * (lr0 - lr_min) * (1 + np.cos(np.pi * min(progress, 1.0)))

def train_with_schedule(schedule_fn, batch_size=16, n_epochs=80, seed=1):
    local_rng = np.random.default_rng(seed)
    theta = np.zeros(2)
    history = [full_loss(theta)]
    for epoch in range(n_epochs):
        lr = schedule_fn(epoch)
        perm = local_rng.permutation(n_samples)
        for start in range(0, n_samples, batch_size):
            idx = perm[start:start + batch_size]
            theta = theta - lr * batch_grad(theta, idx)
        history.append(full_loss(theta))
    return theta, np.array(history)

# The irreducible loss: MSE of the best possible line (least squares)
A = np.column_stack([X, np.ones(n_samples)])
theta_star, *_ = np.linalg.lstsq(A, y, rcond=None)
loss_star = full_loss(theta_star)

schedules = [("Constant lr=0.8", lambda e: 0.8),
             ("Step decay", step_decay),
             ("Cosine annealing", cosine_schedule),
             ("Warmup + cosine", warmup_cosine)]

print("=== Schedules with Mini-Batch SGD (B=16, 80 epochs) ===")
print(f"Best achievable loss (least squares): {loss_star:.6f}\n")

plt.figure(figsize=(8, 5))
for label, fn in schedules:
    theta, history = train_with_schedule(fn)
    gap = history[-1] - loss_star
    plt.semilogy(history - loss_star + 1e-12, label=label)
    print(f"{label:18s} final loss = {history[-1]:.6f}  "
          f"(excess over optimum: {gap:.2e})")

plt.xlabel('Epoch')
plt.ylabel('Excess loss over optimum (log scale)')
plt.title('Learning-rate schedules: how close SGD gets to the optimum')
plt.legend()
plt.tight_layout()
plt.savefig('schedule_convergence.png', dpi=110)
plt.close()
print("Figure saved to schedule_convergence.png")

Sample Output:

=== Schedules with Mini-Batch SGD (B=16, 80 epochs) ===
Best achievable loss (least squares): 0.130268

Constant lr=0.8    final loss = 0.262402  (excess over optimum: 1.32e-01)
Step decay         final loss = 0.131259  (excess over optimum: 9.92e-04)
Cosine annealing   final loss = 0.130282  (excess over optimum: 1.39e-05)
Warmup + cosine    final loss = 0.130284  (excess over optimum: 1.59e-05)

The log-scale plot makes the story vivid: the constant learning rate plateaus about 10,000 times further from the optimum than the cosine schedules. All schedules make the same fast early progress; the decaying ones then keep improving as the shrinking steps average away the mini-batch noise. On this small convex problem warmup gives no extra benefit — its value appears on large non-convex problems with poorly scaled initial gradients — but note that it costs nothing either.

Choosing a Schedule in Practice

1.7 Chapter Summary

What We Learned

1. Optimization is the engine of learning

   • Training a model means solving $\theta^* = \arg\min_\theta \mathcal{L}(\theta)$
   • Grid search is infeasible beyond a handful of parameters; gradients give a scalable local compass

2. Landscape geometry decides difficulty

   • Convex problems have a single valley — every local minimum is global
   • Non-convex problems (all neural networks) contain local minima and, especially, saddle points

3. Gradient descent and its two diseases

   • Update rule: $\theta_{t+1} = \theta_t - \eta \nabla \mathcal{L}(\theta_t)$; the learning rate governs everything
   • Disease 1: zigzagging on ill-conditioned problems; Disease 2: crawling along curved valleys

4. SGD trades noise for speed

   • Mini-batch gradients are unbiased but noisy; variance scales as $1/B$
   • Moderate batch sizes (32–256) are the practical sweet spot; some noise even helps escape saddles

5. Modern optimizers add memory

   • Momentum accumulates velocity; RMSProp adapts per-parameter step sizes; Adam does both with bias correction
   • Default advice: start with Adam, consider tuned SGD + momentum for final performance

6. Schedules finish the job

   • Start large, end small: step decay, cosine annealing, and warmup let noisy SGD settle orders of magnitude closer to the optimum

Next Chapter

Gradient-based methods require a differentiable loss and only find local optima. In Chapter 2, Metaheuristic Optimization, we lift both restrictions:

• Simulated annealing: escaping local minima with controlled randomness
• Genetic algorithms: evolving populations of candidate solutions
• Particle swarm optimization: collective search inspired by flocking
• When to prefer derivative-free methods over gradient descent

Exercises

Problem 1 (Difficulty: easy)

Using the function $f(x) = x^2$ with gradient $f'(x) = 2x$ and starting point $x_0 = 5$, implement gradient descent for 30 steps with three learning rates: $\eta = 0.01$, $\eta = 0.5$, and $\eta = 1.05$. Report the final $x$ for each and explain the three behaviors you observe.

# Requirements:
# - Python 3.9+
# - numpy>=1.24.0, <3.0.0

import numpy as np

def gd_on_parabola(lr, x0=5.0, n_steps=30):
    x = x0
    for _ in range(n_steps):
        x = x - lr * (2 * x)      # gradient of x^2 is 2x
    return x

print("=== Learning Rate Experiment on f(x) = x^2 ===")
for lr in [0.01, 0.5, 1.05]:
    x_final = gd_on_parabola(lr)
    print(f"lr = {lr:5.2f}: final x = {x_final:.6e}")

=== Learning Rate Experiment on f(x) = x^2 ===
lr =  0.01: final x = 2.727422e+00
lr =  0.50: final x = 0.000000e+00
lr =  1.05: final x = 8.724701e+01

Problem 2 (Difficulty: medium)

On the strongly ill-conditioned quadratic $f(x, y) = x^2 + 100 y^2$ (largest curvature $L = 200$, smallest curvature $\mu = 2$, condition number $\kappa = L/\mu = 100$), implement plain gradient descent with $\eta = 0.009$ (just under its stability limit $2/L = 0.01$) and heavy-ball momentum with the theoretically optimal parameters

$$ \eta = \frac{4}{(\sqrt{L} + \sqrt{\mu})^2}, \qquad \gamma = \left( \frac{\sqrt{L} - \sqrt{\mu}}{\sqrt{L} + \sqrt{\mu}} \right)^2 $$

Both start from $(-9, 2.5)$. Count how many steps each method needs to bring the loss below $10^{-8}$, and explain the size of the gap.

# Requirements:
# - Python 3.9+
# - numpy>=1.24.0, <3.0.0

import numpy as np

L, mu = 200.0, 2.0      # largest / smallest curvature of f

def loss(p):
    return p[0] ** 2 + 100 * p[1] ** 2

def grad(p):
    return np.array([2 * p[0], 200 * p[1]])

def steps_gd(lr=0.009, tol=1e-8, max_steps=200000):
    p = np.array([-9.0, 2.5])
    for t in range(1, max_steps + 1):
        p = p - lr * grad(p)
        if loss(p) < tol:
            return t
    return None

def steps_momentum(lr, gamma, tol=1e-8, max_steps=200000):
    p = np.array([-9.0, 2.5])
    v = np.zeros(2)
    for t in range(1, max_steps + 1):
        v = gamma * v + lr * grad(p)
        p = p - v
        if loss(p) < tol:
            return t
    return None

# Heavy-ball optimal hyperparameters for a quadratic
lr_opt = 4 / (np.sqrt(L) + np.sqrt(mu)) ** 2
gamma_opt = ((np.sqrt(L) - np.sqrt(mu)) / (np.sqrt(L) + np.sqrt(mu))) ** 2

n_gd = steps_gd()
n_mom = steps_momentum(lr_opt, gamma_opt)
print("=== Steps to reach loss below 1e-8 (kappa = 100) ===")
print(f"Optimal momentum hyperparameters: lr = {lr_opt:.4f}, "
      f"gamma = {gamma_opt:.3f}")
print(f"Plain GD (lr=0.009): {n_gd} steps")
print(f"Heavy-ball momentum: {n_mom} steps")
print(f"Speedup: {n_gd / n_mom:.1f}x")

=== Steps to reach loss below 1e-8 (kappa = 100) ===
Optimal momentum hyperparameters: lr = 0.0165, gamma = 0.669
Plain GD (lr=0.009): 629 steps
Heavy-ball momentum: 88 steps
Speedup: 7.1x

Problem 3 (Difficulty: hard)

Combine what you built in Sections 1.5 and 1.6: implement Adam with (a) a constant learning rate $\eta = 0.02$ and (b) a warmup + cosine schedule peaking at $\eta_0 = 0.1$ (warmup 100 steps, total 3000 steps), and run both on the Rosenbrock function from $(-1, 1)$ for 3000 steps. Compare final losses and describe the effect of the schedule.

# Requirements:
# - Python 3.9+
# - numpy>=1.24.0, <3.0.0

import numpy as np

def rosenbrock(p):
    x, y = p
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

def rosenbrock_grad(p):
    x, y = p
    return np.array([-2 * (1 - x) - 400 * x * (y - x ** 2),
                     200 * (y - x ** 2)])

def warmup_cosine_step(t, lr0=0.1, total=3000, warmup=100, lr_min=1e-4):
    """Schedule indexed by step t (0-based)."""
    if t < warmup:
        return lr0 * (t + 1) / warmup
    progress = (t - warmup) / max(1, total - warmup)
    return lr_min + 0.5 * (lr0 - lr_min) * (1 + np.cos(np.pi * min(progress, 1.0)))

def run_adam(schedule_fn, p0=(-1.0, 1.0), n_steps=3000,
             beta1=0.9, beta2=0.999, eps=1e-8):
    p = np.array(p0, dtype=float)
    m = np.zeros_like(p)
    v = np.zeros_like(p)
    for t in range(1, n_steps + 1):
        lr = schedule_fn(t - 1)
        g = rosenbrock_grad(p)
        m = beta1 * m + (1 - beta1) * g
        v = beta2 * v + (1 - beta2) * g ** 2
        m_hat = m / (1 - beta1 ** t)
        v_hat = v / (1 - beta2 ** t)
        p = p - lr * m_hat / (np.sqrt(v_hat) + eps)
    return p

p_const = run_adam(lambda t: 0.02)
p_sched = run_adam(warmup_cosine_step)

print("=== Adam: Constant vs Warmup + Cosine (3000 steps) ===")
print(f"Constant lr=0.02:  point = ({p_const[0]:+.5f}, {p_const[1]:+.5f}), "
      f"loss = {rosenbrock(p_const):.3e}")
print(f"Warmup + cosine:   point = ({p_sched[0]:+.5f}, {p_sched[1]:+.5f}), "
      f"loss = {rosenbrock(p_sched):.3e}")

=== Adam: Constant vs Warmup + Cosine (3000 steps) ===
Constant lr=0.02:  point = (+1.00000, +0.99999), loss = 1.521e-11
Warmup + cosine:   point = (+1.00000, +1.00000), loss = 2.155e-27

References

1. Ruder, S. (2016). An overview of gradient descent optimization algorithms. arXiv:1609.04747.
2. Kingma, D. P., & Ba, J. (2015). Adam: A Method for Stochastic Optimization. ICLR 2015.
3. Tieleman, T., & Hinton, G. (2012). Lecture 6.5 - RMSProp. COURSERA: Neural Networks for Machine Learning.
4. Polyak, B. T. (1964). Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5), 1-17.
5. Loshchilov, I., & Hutter, F. (2017). SGDR: Stochastic Gradient Descent with Warm Restarts. ICLR 2017.
6. Boyd, S., & Vandenberghe, L. (2004). Convex Optimization. Cambridge University Press.
7. Goodfellow, I., Bengio, Y., & Courville, A. (2016). Deep Learning, Chapter 8: Optimization for Training Deep Models. MIT Press.