Ridge Regression

Implement ridge regression from scratch with NumPy using the closed-form normal equations. Ridge is ordinary least squares plus an L2 penalty on the coefficients, which shrinks them toward zero, keeps the fit stable when features are correlated, and makes X^T X invertible even when it otherwise would not be. Interviewers use it to check that you know the normal equations, that you handle the bias term correctly, and that you solve a linear system instead of inverting a matrix.

Definition

Augment X with a leading column of ones so the first weight is the intercept: X_aug has shape (n, d + 1). Ridge minimizes the penalized squared error

J(w) = ||X_aug @ w - y||^2 + lam * ||w[1:]||^2

where the sum in the penalty runs over the feature coefficients only, not the bias. Setting the gradient to zero gives the normal equations

(X_aug^T X_aug + lam * R) w = X_aug^T y

with R = diag(0, 1, 1, ..., 1) a (d + 1, d + 1) matrix whose first diagonal entry is 0 so the bias is left unpenalized. The lam * R term also lifts the diagonal, which is why the system stays solvable even when X_aug^T X_aug is singular.

Task

Complete ridge_regression(X, y, lam) so it returns the weight vector w of shape (d + 1,), where w[0] is the intercept and w[1:] are the feature coefficients. Build the augmented matrix, form A = X_aug^T X_aug + lam * R and b = X_aug^T y, and solve A w = b with np.linalg.solve. Do not call np.linalg.inv, and do not use sklearn, scipy, or any built-in ridge/regression helper.

Example

X = np.array([[1.0], [2.0], [3.0], [4.0]])
y = np.array([2.0, 4.0, 6.0, 8.0])   # exactly y = 2*x, intercept 0

w = ridge_regression(X, y, lam=0.0)
w   # -> array([0., 2.])   # lam=0 recovers ordinary least squares: bias 0, slope 2

w = ridge_regression(X, y, lam=1.0)
w   # -> array([0.83333333, 1.66666667])   # penalty shrinks the slope; bias absorbs the offset

With lam = 0 ridge is exactly OLS and recovers the true slope of 2 with a zero intercept. Turning up lam shrinks the slope while the unpenalized intercept shifts to compensate, so predictions stay reasonable rather than collapsing to zero.

Constraints

• 1 <= n <= 10^5, 1 <= d <= 500; use vectorized NumPy, not Python loops over rows or features.
• lam >= 0. At lam = 0 the result must equal the ordinary least squares solution (when it exists).
• The bias term (weight index 0) must never be regularized.
• Solve the linear system with np.linalg.solve; do not form an explicit inverse. Values fit in float64.

Approach

Prepend a column of ones to X so the first weight is the intercept, giving X_aug of shape (n, d + 1). Build the penalty matrix R = diag(0, 1, ..., 1) — an identity with its top-left entry zeroed so the bias is not shrunk. Form the normal-equation system A = X_aug^T X_aug + lam * R and b = X_aug^T y, then solve A w = b with np.linalg.solve, which is more stable and faster than computing A^-1 explicitly.

Solution

import numpy as np

def ridge_regression(X: np.ndarray, y: np.ndarray, lam: float) -> np.ndarray:
    X = np.asarray(X, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    n, d = X.shape

    # Augment with a leading column of ones for the bias term.
    X_aug = np.hstack([np.ones((n, 1)), X])          # shape (n, d + 1)

    # Penalty matrix: identity but with the bias (index 0) left unpenalized.
    R = np.eye(d + 1)
    R[0, 0] = 0.0

    A = X_aug.T @ X_aug + lam * R                     # (d + 1, d + 1)
    b = X_aug.T @ y                                   # (d + 1,)

    w = np.linalg.solve(A, b)                         # no explicit inverse
    return w

The only subtlety is R[0, 0] = 0.0: without it the intercept would be shrunk toward zero, biasing every prediction. np.linalg.solve(A, b) returns the exact solution of A w = b; because lam * R lifts the non-bias diagonal, A is positive definite for any lam > 0 and the solve succeeds even when X_aug^T X_aug alone is singular.

Walkthrough

On X = [[1],[2],[3],[4]], y = [2,4,6,8], lam = 1.0:

• X_aug = [[1,1],[1,2],[1,3],[1,4]], shape (4, 2).
• X_aug^T X_aug = [[4, 10], [10, 30]]; R = [[0,0],[0,1]], so lam*R = [[0,0],[0,1]].
• A = [[4, 10], [10, 31]] — only the feature diagonal (30 -> 31) is lifted; the bias entry stays 4.
• b = X_aug^T y = [20, 60].
• Solving A w = b: 4*w0 + 10*w1 = 20 and 10*w0 + 31*w1 = 60. From the first, w0 = 5 - 2.5*w1; substituting gives 6*w1 = 10, so w1 = 1.6667 and w0 = 0.8333, i.e. w = [0.8333, 1.6667].

With lam = 0, A = [[4,10],[10,30]], b = [20,60], and the solve returns w = [0, 2] — the exact OLS fit y = 2x with zero intercept.

Complexity & notes

• Time O(n d^2 + d^3), space O(n d + d^2) — building X_aug^T X_aug costs O(n d^2) and solving the (d+1)-dimensional system costs O(d^3). For d << n the O(n d^2) term dominates.
• Use np.linalg.solve(A, b), not np.linalg.inv(A) @ b. The solve is faster (it does not compute the full inverse) and numerically more stable; forming an explicit inverse is a common red flag in interviews.
• Zeroing R[0, 0] is the detail interviewers watch for. Penalizing the intercept shrinks it toward zero and biases predictions whenever the target has a nonzero mean, which is almost always.
• Convention matters: standard ridge does not penalize the bias, and features are usually standardized first so a single lam treats each coefficient comparably. This closed form assumes features are on sensible scales.
• The lam * R term guarantees invertibility. When features are collinear, X_aug^T X_aug is singular and OLS blows up, but any lam > 0 makes A positive definite and the solution well defined — the core reason ridge is used.