Gradient Descent for Linear Regression

Definition

How does a model find the best weights when the formula is too expensive to compute?

Gradient descent is an iterative optimisation method that repeatedly moves the parameters a small step in the direction of the negative gradient of the loss. For mean squared error in linear regression, the update is w ← w − η·(2/n)·Xᵀ(Xw − y), where η is the learning rate. Mini-batch and stochastic variants use a subset of rows per step.

In simple words

Imagine standing on a foggy hillside, wanting the valley floor. Feel which way is downhill, take a step, and repeat. The gradient says which way is uphill; you go the opposite way. Too big a step and you overshoot and bounce around; too small and it takes forever.

Why it matters

The normal equation costs about p³ operations and needs all data in memory. Gradient descent scales to millions of rows and features, and the same idea trains logistic regression and every neural network.

How it works

Start with w = 0.
Compute predictions Xw and the error Xw − y.
Compute the gradient (2/n)Xᵀ(Xw − y).
Update w ← w − η × gradient.
Repeat until the loss stops decreasing.

The maths

Mean squared error

J (w) = \frac{1}{n} ∥ X w - y ∥^{2}

Its gradient

\nabla J (w) = \frac{2}{n} X^{⊤} (X w - y)

The update rule

w \leftarrow w - η \frac{2}{n} X^{⊤} (X w - y)

$η$: The learning rate: the step size

In code

Batch gradient descent from scratch, with three learning rates, compared with the exact least-squares answer. The largest rate overshoots and the loss explodes.

gradient_descent.py

import numpy as np

rng = np.random.default_rng(0)
n = 100
x = rng.uniform(-2, 2, n)
y = 3 * x + 2 + rng.normal(scale=0.3, size=n)
X = np.column_stack([np.ones(n), x])

def gradient_descent(lr, steps=300):
    w = np.zeros(2)
    for step in range(steps):
        gradient = 2 / n * X.T @ (X @ w - y)
        w -= lr * gradient
    return w, float(np.mean((X @ w - y) ** 2))

exact = np.linalg.lstsq(X, y, rcond=None)[0]
print("Exact least squares:", np.round(exact, 3))
for lr in [0.01, 0.1, 0.8]:
    w, loss = gradient_descent(lr)
    print(f"lr = {lr:<4} -> w = {np.round(w, 3)}, final MSE = {loss:.4g}")

w = np.zeros(2)
for step in range(1, 51):
    w -= 0.1 * 2 / n * X.T @ (X @ w - y)
    if step in (1, 5, 10, 50):
        print(f"step {step:>2}: loss {np.mean((X @ w - y) ** 2):.4f}")

Output

Exact least squares: [1.978 2.99 ]
lr = 0.01 -> w = [1.975 2.991], final MSE = 0.08461
lr = 0.1  -> w = [1.978 2.99 ], final MSE = 0.0846
lr = 0.8  -> w = [-1.74076932e+54 -5.15840624e+54], final MSE = 4.661e+109
step  1: loss 9.4811
step  5: loss 0.6169
step 10: loss 0.1071
step 50: loss 0.0846

Try it yourself

What does the learning rate do?

Match each learning-rate setting to its typical behaviour.

Very small (0.0001)
Well chosen (0.1 on scaled data)
Too large
Decaying over time

Common mistakes

Forgetting to scale features, so convergence is painfully slow.
Using one learning rate without checking the loss curve.
Updating weights one at a time with stale values instead of all together.

Check yourself

1. Which way does gradient descent move the weights?

2. If the loss grows every step, the learning rate is probably…

3. Why use gradient descent instead of the normal equation?

Summary

Gradient descent steps against the gradient: w ← w − η∇J.
For MSE, ∇J = (2/n)Xᵀ(Xw − y).
Scale features and tune the learning rate by watching the loss.

Read it. Now try it yourself.

Open the interactive activity

Finished reading?

Sources: Goodfellow et al., Deep Learning, chapter 4 · James et al., An Introduction to Statistical Learning