Easy Primitives

Ridge Regression

Ridge regression fits a linear model while penalizing large weights. The penalty favors smaller coefficients, but its usefulness depends on the data and regularization strength; it does not guarantee better predictions.

Your task

Implement ridge_regression(x, y, lam) for floating-point arrays:

x has shape (N, d), one example per row and one feature per column.
y has shape (N,).
lam is a nonnegative regularization strength.

N and d are positive. Return the fitted weight vector with shape (d,), without changing inputs. There is no separate intercept. Do not add a column of ones or center the data. Every supplied feature's coefficient is penalized, including a constant column if one is already present.

Minimize this objective, where the error term is a sum, not a mean:

sum((y - x @ w)**2) + lam * sum(w**2)

The weights satisfy the normal equations:

A = x.T @ x + lam * I
rhs = x.T @ y
A @ w = rhs

I is the d-by-d identity matrix. Solve the system with torch.linalg.solve or jnp.linalg.solve; do not form an explicit matrix inverse. In PyTorch, create I with x's dtype and device.

When is the solution unique?

For lam > 0, the added diagonal penalty makes A positive definite in exact arithmetic, even if x has dependent columns or fewer rows than columns. Floating-point computations can still be affected by conditioning.

For lam = 0, assume x has full column rank. Without that assumption, the unregularized normal equations may be singular. Positive regularization fixes that issue; zero regularization does not.

Example

For x = [[1], [2], [3]] and y = [2, 4, 6], x.T @ x is 14 and x.T @ y is 28. The single fitted weight is 28 / (14 + lam):

lam = 0 gives weight 2, fitting these targets exactly.
lam = 1 gives weight 28/15, approximately 1.866667.

Do not divide the data term by N when building this system. Using a mean squared error objective with the same numeric lam would require an N-scaled penalty in the normal equations and would produce different weights.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Predict a numerical value

Regression predicts a number, such as a price or a temperature. A linear model combines the input features using weights and a bias:

$$\hat y=Xw+b.$$

Mean squared error (MSE) averages the squared differences between predictions and targets. For N examples with one target each,

$$L=\frac{1}{N}\sum_i(\hat y_i-y_i)^2.$$

Squaring makes larger errors count more. An error of 4 contributes 16 to the sum, while an error of 2 contributes 4.

The gradient with respect to each prediction is 2 * (prediction - target) / N. For several outputs per example, check whether the loss averages over examples only or over all output elements. The denominator, and therefore the gradient, will differ.

More features can represent curves

Polynomial regression uses features such as [1, x, x²]. The predicted curve can be nonlinear in x while remaining linear in its learned coefficients. Large powers can create very different feature scales, making the fitting problem more sensitive to numerical error and learning-rate choices.

Regularization discourages large weights

Adding the penalty lambda * sum(w²) adds 2 * lambda * w to the weight gradient. Some definitions include a factor of one half, so follow the stated loss rather than assuming a convention.

This penalty is used in ridge regression. It can make fitting more stable when features are strongly correlated, but feature scaling still matters: rescaling a feature changes the coefficient needed to represent the same prediction.

Why squared error favors the mean

For targets [1, 3], predicting 2 for both gives squared errors [1, 1] and MSE 1. Predicting 1 for both gives errors [0, 4] and MSE 2. The mean minimizes squared error for a constant prediction. More generally, when the relevant expectations exist, the best prediction under expected squared loss is the target's conditional mean given the input.

Large outliers can strongly affect this mean and the fitted model. Evaluate on held-out data to check whether a low training loss carries over to new examples.

Detecting runtime…