Ridge Regression From Scratch: Closed Form in NumPy

Given a design matrix X (n × d), targets y (length n) and a penalty lam ≥ 0, return the weights w that minimise ‖Xw − y‖² + lam·‖w‖². There is no intercept. When lam = 0 the tests guarantee XᵀX is invertible.

Examples

Input X = [[1, 0], [0, 1], [1, 1]], y = [1, 2, 3], lam = 0

Output [1, 2]

An exact fit: w = (1, 2) reproduces every target.

Input X = [[1, 0], [0, 1], [1, 1]], y = [1, 2, 3], lam = 1

Output [0.875, 1.375]

The penalty shrinks both weights towards zero.

Input X = [[2], [4]], y = [2, 4], lam = 100

Output [0.1667]

A heavy penalty: the weight is pulled far below the unpenalised 1.

Constraints

1 ≤ n ≤ 500, 1 ≤ d ≤ 50
0 ≤ lam ≤ 10⁴
Checked to a relative tolerance of 10⁻⁷

Hints

Hint 1

Differentiate ‖Xw − y‖² + λ‖w‖² with respect to w and set the result to zero.

Hint 2

You should get (XᵀX + λI)w = Xᵀy. Which NumPy function solves a linear system directly?

Hint 3

np.linalg.solve(X.T @ X + lam * np.eye(d), X.T @ y). Never np.linalg.inv.

Where candidates go wrong

Inverting XᵀX + λI and multiplying. Slower, less accurate, and the first thing an interviewer will question.
Adding λ to every entry of XᵀX instead of to the diagonal only.
Dividing the loss by n in your head but not in the algebra: here the loss is a plain sum, so λ is not scaled.

The idea

The loss is a convex quadratic in w. Its gradient is 2Xᵀ(Xw − y) + 2λw, and setting it to zero gives (XᵀX + λI)w = Xᵀy. For λ > 0 the matrix is positive definite, so the solution is unique whatever the shape of X.

np.linalg.solve uses an LU factorisation, which costs about the same as one inversion but avoids the extra multiplication and its rounding. For a symmetric positive definite matrix a Cholesky factorisation is twice as fast again; scipy.linalg.cho_solve does it.

Time O(nd² + d³) · Space O(d²)

Reference solution

import numpy as np

def ridge(X, y, lam):
    X = np.asarray(X, dtype=float)
    y = np.asarray(y, dtype=float)
    d = X.shape[1]
    # (X^T X + lam I) w = X^T y: solve the system, never form the inverse
    return np.linalg.solve(X.T @ X + lam * np.eye(d), X.T @ y)

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

It checks that you can differentiate a matrix expression and that you know the difference between solving a system and inverting a matrix. The follow-up, almost always, is what happens when there are more features than rows.

The theory behind it is in the lesson Linear and logistic regression from scratch.

Keep going

Easy10 min

Ridge regression in closed form