02 - Linear Regression from Scratch in Python

You have a table of examples, each with a few measurements and a number to predict, and you want a formula that returns that number from the measurements. Linear regression is the simplest formula you can write: a weighted sum of the measurements, plus a constant term. Choosing the weights is a least-squares problem and has an exact solution; the real work lies in computing it reliably. This article derives it, computes it in three ways that give the same coefficients, shows with a measurement when one of the three fails, and also obtains it by trial and error with gradient descent, checking the gradient against a numerical computation. You need chapter 01 for the words "train", "test" and "baseline".

Everything runs on numpy 2.4.6 and scikit-learn 1.9.0, with the Diabetes dataset bundled in the library: 442 patients, ten baseline measurements (age, sex, body mass index, mean blood pressure and six blood tests) and, as the value to predict, a quantitative index of disease progression one year after the measurement (description from the scikit-learn documentation). No downloads; every output shown is identical across two runs.

The model and the loss

Let X be the data matrix, one row per example and one column per measurement, with n rows and d columns, and y the vector of values to predict. The model predicts ŷ = Xβ, where β holds one weight per column. To also have a constant term (the intercept: the prediction when all the measurements are zero) a column of ones is added to X, and its weight becomes the intercept. An example with numbers: to predict the price of a flat from square metres and storey, β will have three weights, one for the square metres, one for the storey and the intercept.

You choose β by minimising the mean squared error (MSE), the mean of the squares of the residuals, that is, of the differences between true and predicted value: MSE(β) = (1/n)·||y − Xβ||². Squaring makes every error positive and punishes large errors far more, which has a useful side (the function is smooth and has a single minimum) and an inconvenient one (a single anomalous value pulls the line towards itself). To judge the model you also use R², equal to 1 minus the ratio between the sum of the squared residuals of the model and that of a model that always predicts the mean of y: it is 1 for a perfect fit, 0 if you do no better than the mean, and it can become negative on the test set.

The normal equations and the geometry of the residuals

The gradient of ||y − Xβ||² with respect to β is −2Xᵀ(y − Xβ). Setting it equal to zero gives the normal equations: XᵀX β = Xᵀy. If the columns of X are linearly independent, XᵀX is invertible and the solution is unique; the loss is a parabola in several dimensions and the point where the gradient vanishes is its global minimum. The textbook by James, Witten, Hastie and Tibshirani (An Introduction to Statistical Learning) derives the same solution in section 3.1.1, "Estimating the Coefficients".

The geometric meaning helps to remember the formula: Xβ̂ is the orthogonal projection of y onto the space spanned by the columns of X, so the residuals are perpendicular to every column, that is, Xᵀ(y − Xβ̂) = 0. With the column of ones, in particular, the residuals sum to zero. We check it on an exercise to do by hand.

Worked exercise: three points

Given the points (1, 1), (2, 2), (3, 2), find the line y = b₀ + b₁x that minimises the MSE. By hand: XᵀX is the matrix with rows (3, 6) and (6, 14), because there are three points, the sum of the x is 6 and the sum of the squares is 14; Xᵀy is (5, 11). The determinant is 3·14 − 6·6 = 6, so β̂ = (1/6)·(14·5 − 6·11, −6·5 + 3·11) = (4/6, 3/6), that is, intercept 0.667 and slope 0.5. The program confirms the computation and the perpendicularity of the residuals.

import numpy as np

x = np.array([1.0, 2.0, 3.0])
y = np.array([1.0, 2.0, 2.0])
X = np.c_[np.ones(3), x]
print("X'X =", (X.T @ X).tolist())
print("X'y =", (X.T @ y).tolist())
print("determinant:", round(float(np.linalg.det(X.T @ X)), 6))
b = np.linalg.solve(X.T @ X, X.T @ y)
print("intercept, slope:", b.round(4).tolist())
print("residuals:", (y - X @ b).round(4).tolist(), " sum:", round(float((y - X @ b).sum()), 12))
print("X' residuals (orthogonality):", (X.T @ (y - X @ b)).round(12).tolist())

# Output:
# X'X = [[3.0, 6.0], [6.0, 14.0]]
# X'y = [5.0, 11.0]
# determinant: 6.0
# intercept, slope: [0.6667, 0.5]
# residuals: [-0.1667, 0.3333, -0.1667]  sum: 0.0
# X' residuals (orthogonality): [0.0, 0.0]

The residuals are −0.1667, 0.3333 and −0.1667: they sum to zero and, weighted with x, to zero as well, which are the two conditions Xᵀr = 0 written out in full. If you change a point, for example the third one to (3, 5), the slope goes from 0.5 to 2 and the intercept from 0.667 to −1.333 (same computation, with Xᵀy equal to (8, 20)): this is the sensitivity to anomalous values mentioned above.

Three ways to compute the same solution

The textbook method solves the normal equations. A numerically safer method, numpy.linalg.lstsq, solves the least-squares problem without forming XᵀX, using a singular value decomposition (SVD); scikit-learn, says its documentation in the section on OLS complexity, computes the solution with the SVD of X, at a cost proportional to n·d² if n ≥ d. The listing compares them on the same dataset, after setting aside 25% of the examples.

import numpy as np
from sklearn.datasets import load_diabetes
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

X, y = load_diabetes(return_X_y=True, scaled=False)
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.25, random_state=0)
sc = StandardScaler().fit(Xtr)                       # mean and std from the training set only
A = np.c_[np.ones(len(Xtr)), sc.transform(Xtr)]      # column of 1s = intercept
A_te = np.c_[np.ones(len(Xte)), sc.transform(Xte)]

b_eq = np.linalg.solve(A.T @ A, A.T @ ytr)           # normal equations
b_ls = np.linalg.lstsq(A, ytr, rcond=None)[0]        # least squares via SVD
sk = LinearRegression().fit(sc.transform(Xtr), ytr)
b_sk = np.r_[sk.intercept_, sk.coef_]

print("train/test rows, features:", len(Xtr), len(Xte), X.shape[1])
print("max |normal equations - lstsq|:", np.abs(b_eq - b_ls).max() < 1e-9)
print("max |lstsq - sklearn|:", np.abs(b_ls - b_sk).max() < 1e-9)
print("intercept (= mean of y in train):", round(b_ls[0], 3), round(ytr.mean(), 3))
print("R2 train:", round(r2_score(ytr, A @ b_ls), 3), " R2 test:", round(r2_score(yte, A_te @ b_ls), 3))

# Output:
# train/test rows, features: 331 111 10
# max |normal equations - lstsq|: True
# max |lstsq - sklearn|: True
# intercept (= mean of y in train): 151.921 151.921
# R2 train: 0.555  R2 test: 0.359

The three solutions coincide to within 1e-9: they are mathematically the same number, computed with three algorithms. In np.linalg.solve(A.T @ A, A.T @ ytr) I use solve and not the explicit inverse, because solving the system is more precise and cheaper than inverting and multiplying. The StandardScaler is computed on the training set only, and the same mean and standard deviation are applied to the test set; with centred columns the intercept coincides with the mean of y on the training set (151.921 for both printed values). The value to read most carefully is R²: 0.555 on the training set and 0.359 on the test set. The model explains less on the test set, and that is normal; with only 111 test examples the estimate is also noisy, so do not compare these two figures with those of another seed. Chapter 05 shows how to estimate the variability.

When the normal equations betray you

The condition number κ(X), the ratio between the largest and the smallest singular value, says how much a small error on the data is amplified in the solution. Forming XᵀX squares that number: κ(XᵀX) = κ(X)². With double-precision numbers (about 16 decimal digits) a condition number of 1e18 leaves zero reliable digits. To see it we build a situation in which the right answer is known: the data are a polynomial with all coefficients equal to 1, computed without noise on 100 points in [0, 1], and we compare the maximum error on the coefficients of the two methods.

import numpy as np

x = np.linspace(0, 1, 100)
for degree in (5, 10, 14):
    X = np.vander(x, degree + 1, increasing=True)     # 1, x, x^2, ..., x^degree
    beta = np.ones(degree + 1)
    y = X @ beta                                      # noiseless data: the true solution is known
    b_eq = np.linalg.solve(X.T @ X, X.T @ y)
    b_ls = np.linalg.lstsq(X, y, rcond=None)[0]
    print(f"degree {degree:2d}  cond(X)={np.linalg.cond(X):.1e}  cond(X'X)={np.linalg.cond(X.T @ X):.1e}"
          f"  normal-eq. error={np.abs(b_eq - beta).max():.1e}  lstsq error={np.abs(b_ls - beta).max():.1e}")

# Output:
# degree  5  cond(X)=3.7e+03  cond(X'X)=1.4e+07  normal-eq. error=1.8e-10  lstsq error=1.2e-13
# degree 10  cond(X)=2.1e+07  cond(X'X)=4.5e+14  normal-eq. error=2.4e-03  lstsq error=2.8e-10
# degree 14  cond(X)=2.3e+10  cond(X'X)=2.0e+18  normal-eq. error=4.8e+01  lstsq error=2.1e-06

Up to degree 5 the normal equations work (error 1.8e-10). At degree 10 the error rises to 2.4e-3 while lstsq stays at 2.8e-10; at degree 14 the normal equations are off by 48 on coefficients that equal 1, that is, the answer is meaningless, and lstsq is still at 2e-6. The condition number of X is 2.3e10 and that of XᵀX is 2.0e18, beyond the precision of the machine. The last digits of these outputs depend on the linear algebra library NumPy is compiled with, but the orders of magnitude and the point where the naive method gives way are robust. This is the reason why scikit-learn never solves the normal equations, and the reason why it pays to standardise the measurements: columns on different scales worsen the conditioning.

Gradient descent

If n is huge, or the data arrive as a stream, you do not form XᵀX and you look for the minimum by trial and error: at each step β is moved by a stretch η (the learning rate) in the direction opposite to the gradient, β ← β − η·∇MSE(β), with ∇MSE = (2/n)·Xᵀ(Xβ − y). Before trusting a hand-written gradient you check it against centred finite differences, (f(β+h·eⱼ) − f(β−h·eⱼ)) / 2h, which only require computing the loss. The test catches almost all errors of sign and of a factor 2 or n.

The maximum step is fixed by the curvature: the Hessian of the MSE is (2/n)·XᵀX, and the method converges only if η < 2/λ_max, where λ_max is its largest eigenvalue. Each eigendirection of the error is multiplied at every step by (1 − η·λᵢ): if the modulus exceeds 1 the error grows, and the steepest direction is the first to leave the limit.

import numpy as np
from sklearn.datasets import load_diabetes
from sklearn.preprocessing import StandardScaler

X, y = load_diabetes(return_X_y=True, scaled=False)
A = np.c_[np.ones(len(X)), StandardScaler().fit_transform(X)]
n, d = A.shape

def mse(b):  return np.mean((A @ b - y) ** 2)
def grad(b): return 2 / n * A.T @ (A @ b - y)

# 1) the analytic gradient against centred finite differences
rng = np.random.default_rng(0)
b = rng.standard_normal(d)
h, num = 1e-5, np.zeros(d)
for j in range(d):
    e = np.zeros(d); e[j] = h
    num[j] = (mse(b + e) - mse(b - e)) / (2 * h)
rel = np.linalg.norm(num - grad(b)) / np.linalg.norm(grad(b))
print("relative error analytic gradient vs finite differences < 1e-6:", rel < 1e-6)

# 2) step threshold: eta < 2 / lambda_max of the Hessian (2/n) A^T A
lam = np.linalg.eigvalsh(2 / n * A.T @ A).max()
threshold = 2 / lam
b_closed = np.linalg.lstsq(A, y, rcond=None)[0]
print("lambda_max Hessian:", round(lam, 3), " eta threshold:", round(threshold, 4))
for frac in (0.9, 1.1):
    b = np.zeros(d)
    for _ in range(500):
        b = b - frac * threshold * grad(b)
    print(f"eta = {frac} x threshold, after 500 steps: MSE =", f"{mse(b):.4g}")
b = np.zeros(d)
for t in range(5000):
    b = b - 0.9 * threshold * grad(b)
print("eta = 0.9 x threshold, after 5000 steps: MSE", round(mse(b), 3), "versus closed form", round(mse(b_closed), 3))
print("coefficients equal to 1e-6:", np.abs(b - b_closed).max() < 1e-6)

# Output:
# relative error analytic gradient vs finite differences < 1e-6: True
# lambda_max Hessian: 8.048  eta threshold: 0.2485
# eta = 0.9 x threshold, after 500 steps: MSE = 2860
# eta = 1.1 x threshold, after 500 steps: MSE = 2.776e+82
# eta = 0.9 x threshold, after 5000 steps: MSE 2859.696 versus closed form 2859.696
# coefficients equal to 1e-6: True

The check passes: the relative error between analytic and numerical gradient is below 1e-6. For the standardised measurements of the example λ_max is 8.048 and the threshold 0.2485. With η equal to 90% of the threshold, 500 steps bring the MSE to 2860, and after 5000 steps the result coincides with the closed form (2859.696) with coefficients equal to within 1e-6; with 110% of the threshold the MSE after 500 steps is 2.8e82, so the method diverges. The slowness of the convergence depends instead on the smallest eigenvalue: the more unbalanced the eigenvalues are (that is, the higher the condition number), the more steps are needed. This is why gradient descent on non-standardised data can require many more iterations. The mini-batch variants and momentum are not run here; chapter 10 uses them for neural networks.

Regularisation: Ridge and Lasso

When the columns are almost dependent, XᵀX is almost singular and the coefficients become unstable. A penalty on the size of the weights is then added. Ridge regression minimises ||y − Xβ||² + α·||β||² (L2 norm) and still has a closed-form solution: (XᵀX + αI)β = Xᵀy. For α > 0 the matrix is always invertible. The Lasso uses the L1 norm, the sum of the absolute values, which is not differentiable at zero: there is no closed formula, but coordinate descent solves one weight at a time with the soft-thresholding function S(z, t) = sign(z)·max(|z| − t, 0), and brings some weights exactly to zero.

Beware of the convention. According to the scikit-learn documentation, Ridge minimises ||Xw − y||² + α||w||², while Lasso minimises (1/(2n))·||Xw − y||² + α||w||₁: the same α does not carry the same weight in the two cases, and a hand-written Lasso with the other convention gives results different from those of Lasso(alpha=...). In the listing the coordinate Lasso follows the scikit-learn convention, and the comparison checks it.

import numpy as np
from sklearn.datasets import load_diabetes
from sklearn.linear_model import Lasso, Ridge
from sklearn.preprocessing import StandardScaler

X, y = load_diabetes(return_X_y=True, scaled=False)
X = StandardScaler().fit_transform(X)
yc = y - y.mean()                                    # centred data: no intercept
n, d = X.shape

# Ridge: min ||y - Xb||^2 + a ||b||^2  ->  (X'X + a I) b = X'y
a = 10.0
b_ridge = np.linalg.solve(X.T @ X + a * np.eye(d), X.T @ yc)
sk_r = Ridge(alpha=a, fit_intercept=False).fit(X, yc)
print("Ridge equal to sklearn:", np.allclose(b_ridge, sk_r.coef_, atol=1e-8))

# Lasso: min (1/2n)||y - Xb||^2 + a |b|_1, coordinate descent with soft-thresholding
def soft(z, t): return np.sign(z) * np.maximum(np.abs(z) - t, 0)
def lasso_cd(X, y, a, iters=2000):
    b = np.zeros(X.shape[1]); n = len(y)
    for _ in range(iters):
        for j in range(X.shape[1]):
            r = y - X @ b + X[:, j] * b[j]            # residual without feature j
            b[j] = soft(X[:, j] @ r / n, a) / (X[:, j] @ X[:, j] / n)
    return b
a_l = 5.0
b_l = lasso_cd(X, yc, a_l)
sk_l = Lasso(alpha=a_l, fit_intercept=False, tol=1e-10, max_iter=100000).fit(X, yc)
print("Lasso equal to sklearn:", np.allclose(b_l, sk_l.coef_, atol=1e-6))
print("coefficients at zero (Lasso):", int((b_l == 0).sum()), "out of", d)
print("L1 norm OLS vs Ridge vs Lasso:",
      [round(float(np.abs(v).sum()), 1) for v in (np.linalg.lstsq(X, yc, rcond=None)[0], b_ridge, b_l)])

# Output:
# Ridge equal to sklearn: True
# Lasso equal to sklearn: True
# coefficients at zero (Lasso): 5 out of 10
# L1 norm OLS vs Ridge vs Lasso: [164.6, 105.0, 65.0]

The two comparisons with scikit-learn pass (tolerances 1e-8 and 1e-6). The Lasso with alpha=5.0 zeroes 5 coefficients out of 10, that is, it chooses by itself which measurements to keep; the weights of OLS, Ridge and Lasso have a sum of absolute values of 164.6, 105.0 and 65.0 respectively: the penalty shrinks the weights, and this is the phenomenon called shrinkage. y is centred and X is standardised because the penalty treats all the weights in the same way: without standardising, a measurement expressed in small units would have a large weight and would be penalised more than the others. How to choose α is the subject of chapter 04, which links the penalty to the trade-off between bias and variance.

When not to use linear regression

• The relationship between measurements and value is clearly curved or has thresholds: the residuals then show a systematic pattern if you plot them against the predictions, and you need polynomial terms or trees (chapter 06).
• The value to predict is a category: you use logistic regression, because a line can return "probabilities" below 0 or above 1.
• The measurements outnumber the examples (d > n): XᵀX is not invertible, and a penalty is needed.

The three items are not a list of defects but three different signals to look for in the data; the first shows in the residuals, the second in the type of target, the third in the shape of the table. In all three cases the remedy starts from the comparison with a baseline and from an evaluation on unseen data, as in chapter 01.

What you can do now

You can derive the normal equations and read them as a projection, compute the coefficients with three algorithms and know when the first is not reliable, check a gradient with finite differences and set the descent step from the largest eigenvalue, and use Ridge and Lasso with the right convention. In the next chapter the target is no longer a number but a probability: 03 - Logistic regression and binary classification.

The same series also includes: 12 - ML Deployment in Production: Serving and Monitoring.