41 - Gradient descent: derivative, learning rate and three variants

You are on a mountain in thick fog and must get down to the valley. You cannot see the valley floor, but you can feel under your feet which way the ground drops most steeply. So you take a step in that direction, stop, feel the slope again, take another step. Step after step you reach the bottom, without ever having seen the map. Gradient descent is exactly this, applied to the cost function you met in the lesson on hypothesis, mean squared error and cost function: the "mountain" is J(θ), your "position" is the parameters θ, and "feeling the slope" means computing a derivative.

In the previous lesson the minimum of J was found with a closed formula. That road closes when the model has too many parameters or is not linear, which is the case of every neural network. Gradient descent is the algorithm that remains: it works for any differentiable function, and it is the engine of training for everything you will see from here on. Every listing was run twice with a fixed seed and the output is identical; the PyTorch block was run with PyTorch 2.12.0 on CPU (no GPU, 2 threads). This lesson builds it from scratch: what a derivative is, why "against the gradient" is the right direction, how to choose the step, and the three variants (batch, stochastic, mini-batch) used in practice.

The derivative: the slope at a point

Take a function f(θ) with a single parameter. If you move by a small amount h, f changes by f(θ + h) − f(θ). The ratio of the two is the average slope over that stretch; if you let h tend to zero, you get the derivative:

f'(θ) = lim_{h→0} [ f(θ + h) − f(θ) ] / h

The derivative says "by how much f changes for a small shift of θ, per unit of shift". If f'(θ) > 0, f grows moving right; if f'(θ) < 0, f falls moving right; if f'(θ) = 0, you are at a flat point (a minimum, a maximum or a saddle). An example you will use often: for f(θ) = θ², the derivative is f'(θ) = 2θ. At θ = 3 the slope is 6; at θ = −3 it is −6; at θ = 0 it is zero, and it is the minimum.

You do not have to trust the formula: you can check it numerically with a central finite difference, which approximates the derivative by choosing a tiny h and looking at both sides: [f(θ + h) − f(θ − h)] / 2h.

import numpy as np

f = lambda t: t ** 2
h = 1e-5
t = 3.0
print("numerical derivative at 3: %.6f   (formula 2*theta = %.1f)" % ((f(t + h) - f(t - h)) / (2 * h), 2 * t))

numerical derivative at 3: 6.000000   (formula 2*theta = 6.0)

The code defines f as an anonymous function (lambda), picks a small but not too small h (with h too close to the machine precision, the difference of two nearly equal numbers loses significant digits) and computes the symmetric ratio. The result matches 2θ = 6 up to the sixth digit. This technique will come in handy later as a tool to verify hand-computed gradients.

With more parameters: partial derivatives and the gradient

The cost function of linear regression depends on two parameters, J(θ₀, θ₁). The partial derivative ∂J/∂θ₁ is the derivative with respect to θ₁ holding θ₀ fixed, as if the other variable were a constant. The gradient is the vector that collects them all:

∇J(θ) = ( ∂J/∂θ₀ , ∂J/∂θ₁ )

The gradient has a fundamental geometric property: it points in the direction in which J grows fastest (the "steepest ascent"), and its length is the steepness. So going in the opposite direction, −∇J, is the steepest descent. Hence the update rule:

θ ← θ − α · ∇J(θ)

where α (alpha) is the learning rate: the length of the step, in units of the gradient. It is repeated until J stops dropping significantly. The step is automatically small near the minimum, where the gradient tends to zero, and large far from it.

The gradient for linear regression, derived with the chain rule

With h(x) = θ₀ + θ₁x and J = (1/2m) Σ (h(x_i) − y_i)², the chain rule says that the derivative of a composite function f(g(θ)) is f'(g(θ)) · g'(θ): you differentiate "from the outside in" and multiply. Call e_i = h(x_i) − y_i the error. The square is the outer function, e_i the inner one:

∂J/∂θ₀ = (1/2m) Σ 2·e_i · ∂e_i/∂θ₀ = (1/m) Σ e_i · 1
∂J/∂θ₁ = (1/2m) Σ 2·e_i · ∂e_i/∂θ₁ = (1/m) Σ e_i · x_i

Why ∂e_i/∂θ₀ = 1 and ∂e_i/∂θ₁ = x_i? Because e_i = θ₀ + θ₁x_i − y_i is linear in the parameters: θ₀ appears with coefficient 1 and θ₁ with coefficient x_i. In matrix form, with the matrix X of shape (m, 2) that has a column of ones and one with the x values, the gradient is a single expression ∇J = (1/m) Xᵀ(Xθ − y), the same one you set to zero to get the normal equation. Now, instead of solving it, you use it to take a step.

The learning rate: the parameter that decides everything

The picture shows three descents on the same bowl. With α too small you reach the bottom but it takes a huge number of steps. With a suitable α you converge in a few steps. With α too large each step overshoots the bottom and lands on the opposite side, higher than before: the descent diverges. We can understand it precisely on the simplest case, J(θ) = θ², with gradient 2θ. The update becomes:

θ_new = θ − α · 2θ = (1 − 2α) · θ        →        θ_k = (1 − 2α)^k · θ₀

At each step θ is multiplied by (1 − 2α). If |1 − 2α| < 1, i.e. 0 < α < 1, θ tends to zero (converges). If 0 < α < 0.5 the factor is positive and θ descends "on one side only"; if 0.5 < α < 1 the factor is negative and θ oscillates from one side to the other but with decreasing amplitude; α = 1 oscillates without ever shrinking; α > 1 oscillates with growing amplitude: it diverges. Let us verify:

for alpha in (0.1, 0.5, 0.9, 1.0, 1.1):
    theta, trajectory = 1.0, [1.0]
    for _ in range(10):
        theta = theta - alpha * 2 * theta        # one descent step on J = theta^2
        trajectory.append(theta)
    print("alpha = %.1f  first steps: %s   after 10 steps: %.4f" % (alpha, np.round(trajectory[:5], 3), theta))

alpha = 0.1  first steps: [1.    0.8   0.64  0.512 0.41 ]   after 10 steps: 0.1074
alpha = 0.5  first steps: [1. 0. 0. 0. 0.]   after 10 steps: 0.0000
alpha = 0.9  first steps: [ 1.    -0.8    0.64  -0.512  0.41 ]   after 10 steps: 0.1074
alpha = 1.0  first steps: [ 1. -1.  1. -1.  1.]   after 10 steps: 1.0000
alpha = 1.1  first steps: [ 1.    -1.2    1.44  -1.728  2.074]   after 10 steps: 6.1917

With α = 0.5 the factor (1 − 2α) is zero and you reach the minimum in one step (a lucky case due to the perfect symmetry of the parabola). With 0.9 it oscillates but converges; with 1.0 it bounces between +1 and −1 forever; with 1.1 the amplitude grows and the descent explodes. In linear regression the stability threshold depends on the curvature of J, i.e. on the largest eigenvalue λ_max of the Hessian matrix H = XᵀX/m: the descent converges if α < 2/λ_max. The larger the feature values, the larger λ_max and the smaller the admissible step. This leads us to the most common trap.

Hand-written gradient descent for linear regression

We implement the rule θ ← θ − α∇J on the same synthetic data as the previous lesson (y = 3 + 2x + noise) and compare it with the exact solution:

rng = np.random.default_rng(0)
m = 100
x = rng.uniform(0, 10, m)
y = 3 + 2 * x + rng.normal(0, 2, m)
X = np.column_stack([np.ones(m), x])                       # (m, 2): column of ones + x
theta_exact = np.linalg.solve(X.T @ X, X.T @ y)            # normal equation, for comparison

def descent(X, y, alpha, iterations):
    theta = np.zeros(X.shape[1])                           # start: all parameters at zero
    costs = []
    for _ in range(iterations):
        gradient = X.T @ (X @ theta - y) / len(y)          # (1/m) X^T (X theta - y)
        theta = theta - alpha * gradient                   # step against the gradient
        costs.append(((X @ theta - y) ** 2).sum() / (2 * len(y)))
    return theta, costs

H = X.T @ X / m
ev = np.linalg.eigvalsh(H)                                 # eigenvalues (H is symmetric), in increasing order
print("Hessian eigenvalues: %s  ->  stability threshold 2/lambda_max = %.4f, condition number %.0f" % (np.round(ev, 4), 2 / ev.max(), ev.max() / ev.min()))
print("exact solution:", np.round(theta_exact, 4))
for alpha, it in ((0.01, 1000), (0.02, 1000)):
    th, costs = descent(X, y, alpha, it)
    print("alpha=%.2f, %4d iterations -> theta = %s   J: %.2f (start) %.4f (end)" % (alpha, it, np.round(th, 4), costs[0], costs[-1]))
with np.errstate(all="ignore"):
    th, costs = descent(X, y, 0.06, 200)
print("alpha=0.06 (above the threshold): theta = %s   final J = %.3g" % (np.round(th, 3), costs[-1]))

Hessian eigenvalues: [ 0.2296 40.0196]  ->  stability threshold 2/lambda_max = 0.0500, condition number 174
exact solution: [2.9851 1.9739]
alpha=0.01, 1000 iterations -> theta = [2.7184 2.0113]   J: 43.14 (start) 1.8884 (end)
alpha=0.02, 1000 iterations -> theta = [2.9584 1.9776]   J: 7.18 (start) 1.8802 (end)
alpha=0.06 (above the threshold): theta = [-6.55996369e+28 -4.66845879e+29]   final J = 4.45e+60

The code does a single thing repeatedly: it computes the gradient with the matrix formula (a multiplication of Xᵀ by the vector of errors, divided by m), shifts θ by α times the gradient, and records the cost so the learning curve can be seen. First it prints the eigenvalues of the Hessian (NumPy's eigvalsh, designed for symmetric matrices, documentation) and the threshold 2/λ_max: with features x between 0 and 10, λ_max is large, and an α above the threshold makes the cost explode. np.errstate suppresses the overflow warnings of the last experiment, which diverges on purpose.

Important observation: with α below the threshold the descent approaches the exact solution, but even after 1000 iterations it does not coincide entirely (with α = 0.01 the intercept is 2.72 against 2.99). The reason lies in the two eigenvalues printed: the curvature is very different in the two directions, the step is limited by the steepest one (λ_max) and the flattest one (λ_min, about 0.23) advances slowly. The ratio between the two, the condition number, is the symptom that the section on feature scale takes up and that further below gets a name and a measurement.

thickness x (tenths of mm)    1     2     3     4
hours      y                  1     3     7     8

The feature scale trap

Imagine changing the unit of measure of x, for example from metres to centimetres: the values become a hundred times larger. The shape of J changes drastically: it becomes a very long and narrow bowl, steep in one direction and almost flat in the other. A single α cannot suit both: small enough not to diverge on the steep direction means very slow on the flat one. Let us check by multiplying x by 100:

xr = x * 100                                               # same information, scale x100
Xr = np.column_stack([np.ones(m), xr])
for alpha in (1e-7, 1e-6, 1e-5):
    with np.errstate(all="ignore"):
        th, costs = descent(Xr, y, alpha, 1000)
    print("original scale x100, alpha=%.0e -> theta = %s   final J = %s" % (alpha, np.round(th, 4), "nan (divergence)" if np.isnan(costs[-1]) else "%.3f" % costs[-1]))

mu, sd = xr.mean(), xr.std()
Xs = np.column_stack([np.ones(m), (xr - mu) / sd])         # standardized feature: mean 0, deviation 1
th_s, costs_s = descent(Xs, y, 0.1, 100)
print("standardized, alpha=0.1, 100 iterations -> final J = %.4f   (theoretical minimum %.4f)" % (costs_s[-1], ((Xs @ np.linalg.solve(Xs.T @ Xs, Xs.T @ y) - y) ** 2).sum() / (2 * m)))
# going back to original units: y = t0 + t1 * (x - mu)/sd  =>  slope t1/sd, intercept t0 - t1*mu/sd
print("in original units: slope = %.5f (true 0.02),  intercept = %.4f (exact 2.985)" % (th_s[1] / sd, th_s[0] - th_s[1] * mu / sd))

original scale x100, alpha=1e-07 -> theta = [0.0001 0.0239]   final J = 2.923
original scale x100, alpha=1e-06 -> theta = [0.0007 0.0239]   final J = 2.922
original scale x100, alpha=1e-05 -> theta = [nan nan]   final J = nan (divergence)
standardized, alpha=0.1, 100 iterations -> final J = 1.8801   (theoretical minimum 1.8801)
in original units: slope = 0.01974 (true 0.02),  intercept = 2.9850 (exact 2.985)

What the code does. For small values of α (1e-7, 1e-6) the descent does not diverge, but after 1000 iterations the cost is stuck at 2.92, against the minimum of 1.88: the direction of the intercept has hardly moved. With α = 1e-5 the step is already above the stability threshold and the values explode to nan. There is no α that works. After standardization (subtract the mean and divide by the standard deviation, for each feature) the bowl becomes round: α = 0.1 converges to the theoretical minimum in fewer than 100 iterations. The last lines show how to go back to the coefficients in the original units: since y = θ₀ + θ₁(x − μ)/σ, the original slope is θ₁/σ and the intercept θ₀ − θ₁μ/σ.

Operating rule: always standardize the features before a gradient descent (with mean and deviation computed on the training set only), and remember to apply the same transformation to the test data.

Three variants: batch, stochastic, mini-batch

The gradient (1/m) Σ ... sums over all m examples: it costs a pass over the whole dataset to take a single step. With millions of examples it is prohibitive. The gradient computed on a subset is a noisy but unbiased estimate of the full one, so you can take a step even with few examples. The three variants differ in how many:

• Batch: all m examples for each step. Smooth, deterministic trajectory, but a single step per pass over the data and a high cost per step.
• Stochastic (SGD): a single example per step. Very many cheap steps, but noisy: the cost oscillates around the minimum without stopping.
• Mini-batch: a small group (typically 16-512 examples). A compromise: a less noisy gradient estimate and efficient computation on GPUs. It is the one used in practice in deep learning.

An epoch is a complete pass over the whole training set. With batch size b you take m/b steps per epoch. Let us compare the three variants at equal epochs (10):

def train(X, y, batch_size, alpha, epochs, seed=0):
    rng = np.random.default_rng(seed)
    n = len(y); theta = np.zeros(X.shape[1]); steps = 0; costs = []
    for _ in range(epochs):
        order = rng.permutation(n)                         # reshuffle every epoch
        for start in range(0, n, batch_size):
            b = order[start:start + batch_size]
            g = X[b].T @ (X[b] @ theta - y[b]) / len(b)    # gradient estimated on the mini-batch
            theta = theta - alpha * g; steps += 1
        costs.append(((X @ theta - y) ** 2).sum() / (2 * n))
    return theta, steps, costs

Xz = np.column_stack([np.ones(m), (x - x.mean()) / x.std()])
minimum = ((Xz @ np.linalg.solve(Xz.T @ Xz, Xz.T @ y) - y) ** 2).sum() / (2 * m)
print("theoretical minimum cost: %.4f" % minimum)
for name, b, a in (("batch (100)", 100, 0.1), ("mini-batch (10)", 10, 0.05), ("stochastic (1)", 1, 0.01)):
    th, steps, costs = train(Xz, y, b, a, 10)
    print("%-16s steps: %4d   J after 1, 2, 5, 10 epochs: %s" % (name, steps, np.round([costs[0], costs[1], costs[4], costs[9]], 3)))

theoretical minimum cost: 1.8801
batch (100)      steps:   10   J after 1, 2, 5, 10 epochs: [93.589 76.164 41.358 15.645]
mini-batch (10)  steps:  100   J after 1, 2, 5, 10 epochs: [42.349 16.37   2.55   1.886]
stochastic (1)   steps: 1000   J after 1, 2, 5, 10 epochs: [16.991  3.891  1.885  1.881]

What the code does. The function reshuffles the indices every epoch (rng.permutation) and scans them in blocks of batch_size; for each block it computes the gradient only on those examples and takes a step. The variant changes only batch_size: 100 = batch, 10 = mini-batch, 1 = stochastic. At equal epochs, batch takes only 10 steps in total and has not yet reached the minimum; mini-batch (100 steps) and stochastic (1000 steps) get there. The learning rates differ because noisier steps want smaller α. The stochastic trajectory never fully stops: the cost keeps oscillating slightly around the minimum, which is why in practice α is reduced over time (learning rate decay).

Checking a gradient: the gradient check

When you write by hand the gradient of a more complicated function, a sign error or an omission is easy. The standard defence is comparing the analytic gradient with the numerical one (central finite differences), parameter by parameter:

def cost(theta):
    return ((X @ theta - y) ** 2).sum() / (2 * m)

def analytic_gradient(theta):
    return X.T @ (X @ theta - y) / m

def numeric_gradient(theta, eps=1e-6):
    g = np.zeros_like(theta)
    for k in range(len(theta)):
        plus, minus = theta.copy(), theta.copy()
        plus[k] += eps; minus[k] -= eps
        g[k] = (cost(plus) - cost(minus)) / (2 * eps)
    return g

t_test = np.array([1.0, 0.5])
ga, gn = analytic_gradient(t_test), numeric_gradient(t_test)
print("analytic:", np.round(ga, 6), " numeric:", np.round(gn, 6))
print("relative error: %.2e" % (np.linalg.norm(ga - gn) / np.linalg.norm(ga + gn)))

analytic: [-10.066094 -68.731636]  numeric: [-10.066094 -68.731636]
relative error: 5.23e-11

What the code does. numeric_gradient perturbs one parameter at a time forwards and backwards by a tiny amount and measures the change in cost. The relative error is of the order of 10⁻¹⁰ or less: if between the two versions there is a difference of the order of 10⁻³ or more, there is a bug in the analytic gradient. You will reuse this technique in the lesson on backpropagation.

The library: PyTorch computes the gradient for you

Writing the gradient by hand is instructive, but for complex models automatic differentiation is used: the library records the operations performed (a "graph" of operations) and derives the gradient by applying the chain rule backwards. The same regression in PyTorch. The listing reuses Xz, y and minimum from the previous blocks; it was run twice with PyTorch 2.12.0 on CPU (CUDA_VISIBLE_DEVICES="", 2 threads) with identical output.

import torch

Xt = torch.tensor(Xz, dtype=torch.float32)
yt = torch.tensor(y, dtype=torch.float32)
theta_t = torch.zeros(2, requires_grad=True)               # requires_grad: "keep track of the operations on this tensor"
optimizer = torch.optim.SGD([theta_t], lr=0.1)

for epoch in range(100):
    optimizer.zero_grad()                                  # reset the accumulated gradients
    cost_t = ((Xt @ theta_t - yt) ** 2).sum() / (2 * len(yt))
    cost_t.backward()                                      # computes d(cost)/d(theta) for every parameter
    optimizer.step()                                       # theta <- theta - lr * gradient

print("PyTorch: theta = %s   J = %.4f   (by hand: minimum J %.4f)" % (np.round(theta_t.detach().numpy(), 4), cost_t.item(), minimum))

PyTorch: theta = [13.8072  5.9826]   J = 1.8801   (by hand: minimum J 1.8801)

The result matches the exact solution computed with NumPy on the same standardized data (normal equation, np.linalg.solve), θ = [13.8075, 5.9827] and minimum J = 1.8801, up to the precision of float32: the difference in the fourth decimal (13.8072 against 13.8075) is exactly that. The printed cost is the one computed in the last epoch, before the last update.

What the code does. requires_grad=True tells PyTorch to build the graph of operations. backward() walks the graph backwards and deposits in theta_t.grad the gradient of the cost. optimizer.step() performs the rule θ ← θ − α·grad (with the default values of torch.optim.SGD, that is momentum 0, it is exactly the rule of this lesson; the default learning rate is 0.001, which is why the listing sets it to 0.1). zero_grad() is necessary because PyTorch accumulates gradients: without resetting them, at each epoch they would add to the previous ones (one of the most frequent traps). The θ values are in the units of the standardized feature, as in the previous blocks. The behaviour of zero_grad() and of backward() is described in the notes on autograd and on the page for torch.optim.SGD in the official documentation; what exactly backward() does you derive in full in backpropagation derived step by step.

Trade-offs and traps

• Learning rate: too large diverges, too small never arrives. Explore on a logarithmic scale (0.001, 0.01, 0.1...) and look at the cost curve: it must fall smoothly. If it rises or becomes nan, halve or divide α by ten.
• Not standardizing features: see above. It is the most frequent cause of descents that "do not work".
• Non-convex functions: linear regression has a single minimum; neural networks have a cost surface full of dips, saddles and flat areas. Gradient descent finds a local minimum (or a stationary point), it does not guarantee the best one. In practice it works surprisingly well, but the result depends on initialization.
• SGD without shuffling: if the data are sorted (for example by class), consecutive mini-batches are all of the same class and the gradient is misleading. Reshuffle every epoch.
• Stopping criterion: a fixed number of iterations is simple; better to stop when the cost on a validation set stops falling (early stopping).
• Gradients not reset (PyTorch) and hand-computed gradients without verification: two silent bugs; the second is prevented by the gradient check.

The six items are not independent. Learning rate, feature scale and stopping criterion answer the same question, namely whether the step suits the shape of the surface and when to stop; non-convexity decides where you end up; shuffling decides whether the gradient estimated on the mini-batch is reliable. The last two items are silent bugs: no error, only wrong numbers, which is why the gradient check and a look at the cost curve must be done before trusting the result.

The condition number: why gradient descent becomes slow

In the section on feature scale we saw that features with very different scales slow the descent. Now we give that phenomenon a number. For a quadratic cost function, like the MSE of linear regression, the curvature in each direction is an eigenvalue of the Hessian matrix H (the matrix of second derivatives). The largest learning rate that does not diverge is 2/λ_max, where λ_max is the largest eigenvalue; but the direction with the smallest eigenvalue λ_min progresses at a step λ_min · η, so very slowly. The ratio κ = λ_max / λ_min is the condition number, and the number of iterations needed grows in proportion to κ.

Intuition: a long, narrow valley, very steep on the sides and almost flat along the bottom; the learning rate must be small not to bounce from one side to the other, and for this reason you advance along the bottom in tiny steps.

Momentum (or the "heavy ball method") remedies this by accumulating a velocity: v ← β·v + g, w ← w − η·v (it is the form of torch.optim.SGD with momentum=β and dampening=0, as described in the documentation). The components of the gradient that oscillate from one step to the next (the walls of the valley) cancel in the sum, while those that always point in the same direction (the bottom) add up. With the optimal choice of parameters for a quadratic, η = 4/(√κ + 1)² and β = ((√κ − 1)/(√κ + 1))², the number of iterations grows like √κ instead of κ. We measure it on a function with two eigenvalues, 1 and κ, counting the steps needed to bring the norm of the point below 10⁻⁶.

def steps_to_converge(kappa, lr, momentum=0.0, tol=1e-6, max_steps=200000):
    curvature = np.array([1.0, kappa])                                    # the function is 0.5·(x₁² + κ·x₂²)
    point, velocity = np.array([1.0, 1.0]), np.zeros(2)
    for step in range(1, max_steps + 1):
        velocity = momentum * velocity + curvature * point                # the gradient is curvature · point
        point = point - lr * velocity
        if np.linalg.norm(point) < tol:
            return step
    return None

print("kappa    plain descent (lr=1/kappa)   with optimal momentum   ratio")
for kappa in (10, 100, 1000):
    plain = steps_to_converge(kappa, 1.0 / kappa)
    root = np.sqrt(kappa)
    heavy = steps_to_converge(kappa, 4 / (root + 1) ** 2, ((root - 1) / (root + 1)) ** 2)
    print("%-8d  %-30d  %-21d  %.1fx" % (kappa, plain, heavy, plain / heavy))

kappa    plain descent (lr=1/kappa)   with optimal momentum   ratio
10        132                             27                     4.9x
100       1375                            95                     14.5x
1000      13809                           321                    43.0x

What the code does. The function counts the steps needed to bring the point, started at (1, 1), below a norm of 10⁻⁶, on a quadratic with curvatures 1 and κ. Plain descent, with learning rate 1/κ (half the maximum stable 2/κ), needs 132 steps with κ = 10, 1,375 with κ = 100 and 13,809 with κ = 1,000: it grows almost exactly tenfold each time κ is multiplied by ten, as the theory predicts (the number of steps is about κ · ln(10⁶) ≈ 13.8 · κ). With optimal momentum the steps are 27, 95 and 321: multiplying κ by ten multiplies the steps by about three (√10 ≈ 3.2).

The saving grows with the difficulty of the problem: 4.9 times fewer steps at κ = 10, 14.5 at κ = 100, 43.0 at κ = 1,000.

Two honest caveats. First, the "optimal" parameters presuppose knowing the eigenvalues, which in a neural network is not true: in practice one fixes β = 0.9 and tunes the learning rate, and the gain is more modest. Second, this is a quadratic with curvatures aligned to the axes, the simplest case: neural networks have non-convex cost functions, and for them momentum is a well-tested heuristic, not a theorem.

Newton's method: a single step for linear regression

Gradient descent uses only the slope. Newton's method also uses the curvature: w ← w − H⁻¹ g, that is, it corrects the gradient with the inverse of the Hessian, and for this reason it "sees" the shape of the valley and goes straight to the bottom. For a quadratic function, like the MSE of linear regression, a single Newton step reaches the exact minimum (it is the normal equation in another guise). We verify it on data with two features of very different scale (deviation 1 and 30), comparing with gradient descent at the maximum stable learning rate.

rng_g = np.random.default_rng(0)
n_g = 200
X_g = np.column_stack([np.ones(n_g), rng_g.normal(0, 1, n_g), rng_g.normal(0, 30, n_g)])
y_g = X_g @ np.array([1.0, 2.0, -0.05]) + rng_g.normal(0, 0.5, n_g)

hessian = X_g.T @ X_g / n_g
w_star = np.linalg.solve(X_g.T @ X_g, X_g.T @ y_g)                        # exact solution
eigenvalues = np.linalg.eigvalsh(hessian)
print("Hessian eigenvalues: %s   condition number %.0f" % (np.round(eigenvalues, 3), eigenvalues.max() / eigenvalues.min()))

w = np.zeros(3)
for iteration in range(1, 100001):
    w = w - (1.0 / eigenvalues.max()) * (X_g.T @ (X_g @ w - y_g) / n_g)
    if np.abs(w - w_star).max() < 1e-6:
        break
print("gradient descent: %d iterations to get within 1e-6 of the solution" % iteration)

w_newton = np.zeros(3) - np.linalg.solve(hessian, X_g.T @ (X_g @ np.zeros(3) - y_g) / n_g)
print("Newton: 1 iteration, maximum error %.1e" % np.abs(w_newton - w_star).max())

Hessian eigenvalues: [9.1900e-01 9.9400e-01 9.5315e+02]   condition number 1037
gradient descent: 15002 iterations to get within 1e-6 of the solution
Newton: 1 iteration, maximum error 5.8e-15

What the code does. It builds 200 points with three columns (a constant term, a feature of deviation 1 and one of deviation 30) and a linear target with noise 0.5. The Hessian of the error (with the half-MSE convention, gradient XᵀXw/n) has eigenvalues 0.919, 0.994 and 953.15: the condition number is 1,037, almost entirely due to the scale-30 feature, whose square (900) appears in the curvature. Gradient descent, even with the maximum safe learning rate 1/λ_max, takes 15,002 iterations to get within 10⁻⁶ of the exact solution (the same count as the experiments above: about 14 · κ). Newton, in a single iteration, finds the solution with a maximum error of 5.8·10⁻¹⁵, i.e. at rounding level.

Why then isn't Newton used everywhere? Because it requires forming and inverting a p × p matrix: the cost is O(p³) and the memory O(p²). For a regression with three coefficients it is nothing; for a network with a hundred million parameters it is impossible. This is the underlying trade-off of optimization: second-order methods are powerful but expensive, first-order ones (with momentum, or with adaptive scales like Adam) are scalable. And feature standardization, which costs almost nothing, is a way of bringing the problem closer to the well-conditioned case where even first order works.

SGD noise and learning rate decay

x_n = rng_g.normal(0, 1, 1000)
y_n = 3 * x_n + rng_g.normal(0, 1, 1000)
w_best = (x_n @ y_n) / (x_n @ x_n)                                        # exact minimum for the model y = w·x

def sgd_final_error(schedule, seed, epochs=30):
    order_rng = np.random.default_rng(seed)
    w, last = 0.0, []
    for epoch in range(epochs):
        for i in order_rng.permutation(1000):
            w -= schedule(epoch) * (w * x_n[i] - y_n[i]) * x_n[i]         # one step per example
        last.append(w)
    return np.abs(np.array(last[-5:]) - w_best).mean()                    # mean error over the last 5 epochs

print("exact solution w = %.4f   (the true value is 3)" % w_best)
for name, schedule in (("fixed learning rate 0.05", lambda e: 0.05), ("decay 0.05/(1+epoch)", lambda e: 0.05 / (1 + e))):
    errors = [sgd_final_error(schedule, seed) for seed in range(10)]
    print("%-34s  mean distance from the minimum, 10 seeds: %.4f" % (name, np.mean(errors)))

exact solution w = 3.0497   (the true value is 3)
fixed learning rate 0.05            mean distance from the minimum, 10 seeds: 0.1259
decay 0.05/(1+epoch)                mean distance from the minimum, 10 seeds: 0.0123

What the code does. It trains a single-parameter model, y = w·x, with stochastic descent (one example per step, 30 epochs on 1,000 points), and measures the mean distance from the exact solution (3.0497, close to the true value 3) over the last five epochs, over ten different random orderings.

With a fixed learning rate of 0.05 the mean distance is 0.1259; with the learning rate decaying as 0.05/(1 + epoch) it is 0.0123, about ten times smaller. The reason is the nature of SGD: the gradient of a single example is a noisy estimate of the true gradient, and with a fixed step the parameter keeps oscillating around the minimum with an amplitude proportional to the learning rate, never stopping. If the step decreases, the oscillation shrinks. It is the same compromise seen for batch descent: at the start a large step to get close quickly, at the end a small step to refine. In real networks these ideas go by the name of learning rate schedule (step decay, cosine, initial warm-up). Note the limitation of the experiment: a single one-parameter problem, and 10 seeds; the tenfold ratio should not be taken as a constant, but as the order of magnitude of the effect in this case.

When the function is not convex: the end point depends on the start

f = lambda x: x ** 4 - 3 * x ** 2 + x                                     # two local minima
df = lambda x: 4 * x ** 3 - 6 * x + 1
print("start        end       value of the function")
for start in (-2.0, -0.5, 0.5, 2.0):
    x = start
    for _ in range(2000):
        x -= 0.01 * df(x)
    print("%+.1f       %+.4f   %+.4f" % (start, x, f(x)))

start        end       value of the function
-2.0       -1.3008   -3.5139
-0.5       -1.3008   -3.5139
+0.5       +1.1309   -1.0702
+2.0       +1.1309   -1.0702

What the code does. It applies gradient descent (2,000 steps of 0.01) to the function f(x) = x⁴ − 3x² + x, which has two valleys. The four starts give only two arrivals: starting from −2 or −0.5 you end in the left minimum, x = −1.3008, where the function is worth −3.5139; starting from +0.5 or +2 you end in the right one, x = +1.1309, where it is worth −1.0702. Both are points where the gradient is zero and the descent stops satisfied, but only the first is the global minimum: the other is a local minimum, with a cost 2.44 units higher. The algorithm has no way of noticing, because it sees only the slope here and now.

This is exactly what changes when passing from linear regression (convex function, a single minimum, the start does not matter) to neural networks (non-convex function, result dependent on initialization and on SGD noise). Very large networks are less afflicted by the problem than one-dimensional intuition suggests, because in many dimensions truly bad local minima seem rare compared with saddle points; but it is an empirical observation, not a guarantee, and for this reason one often trains with several seeds and reports the variability, as we do in these lessons.

To continue on your own

Three tests to do without a solution, each on the experiments already run above. Test A. Standardize the two features of the Newton experiment (mean 0, deviation 1) and recount the gradient descent iterations: does the condition number fall? by how much? What does it say about the relationship between preprocessing and optimization? Test B. In the momentum experiment try non-optimal values of β (0.5; 0.99): for which does convergence get worse than plain descent? Test C. In the non-convex experiment look, by hand or with a grid of starting points, for the separation point between the two basins of attraction: where is it with respect to the two minima?

Other lessons in the series that link to this one: Calculus 1: Limits and Derivatives, Calculus 2: Gradient, Integrals, Gradient Descent, Support Vector Machine: the maximum margin and the kernel trick, Boosting: AdaBoost, Gradient Boosting and a note on XGBoost, Logistic regression: from a line to a probability, The perceptron and activation functions: what non-linearity is for, Recurrent Neural Networks (RNN): Giving the Network a Memory, Backpropagation Through Time: Vanishing and Exploding Gradients.

Sources

• PyTorch, Autograd mechanics - the fact that gradients accumulate in the .grad field of leaf tensors with requires_grad=True, consulted on 2026-11-12
• PyTorch, torch.optim.SGD - default parameters (lr 0.001, momentum 0) and the update rule with momentum, consulted on 2026-11-12
• NumPy, numpy.linalg.eigvalsh - the eigenvalues of a symmetric matrix, used for λ_max and λ_min, consulted on 2026-11-12

Historical references quoted from memory and not reopened during the verification (the publisher's page for Robbins and Monro rejected an automated client): A.-L. Cauchy, Méthode générale pour la résolution des systèmes d'équations simultanées, 1847, for gradient descent; H. Robbins and S. Monro, A Stochastic Approximation Method, Annals of Mathematical Statistics, 1951, for the stochastic version.

In the next lesson you extend regression to multiple features and to curved relations in multivariate and polynomial regression and feature normalization, and see how the same gradient keeps working.