Linear Regression (Gradient Descent)

Fit an ordinary least squares linear regression from scratch using batch gradient descent. Rather than solving the normal equations in closed form, you iteratively step the weights downhill along the gradient of the mean squared error. Interviewers use this to check that you can derive the MSE gradient, wire in a bias term correctly, and keep the matrix shapes straight in vectorized NumPy.

Definition

The model predicts y_hat = X_b @ w, where X_b is X with a leading column of ones so that w[0] acts as the bias. The mean squared error and its gradient are:

loss(w)  = (1/n) * sum((X_b @ w - y) ** 2)
grad(w)  = (2/n) * X_b.T @ (X_b @ w - y)

Each gradient-descent step moves the weights against the gradient:

w <- w - lr * grad(w)

Repeat this n_iters times starting from w = 0.

Task

Complete linear_regression_gd(X, y, lr, n_iters) so it returns the learned weight vector of length d + 1, where index 0 is the bias and indices 1..d are the feature coefficients. Prepend the bias column yourself, initialize w to zeros, and run exactly n_iters full-batch updates using the MSE gradient above. Do not call sklearn, scipy, or a closed-form np.linalg solver.

Example

rng = np.random.default_rng(0)
X = rng.normal(size=(200, 1))
y = 3.0 * X[:, 0] + 2.0 + rng.normal(scale=0.1, size=200)  # true bias 2, slope 3

w = linear_regression_gd(X, y, lr=0.1, n_iters=2000)
w            # -> array([~2.0, ~3.0])   (bias first, then slope)
round(w[0], 1), round(w[1], 1)  # -> (2.0, 3.0)

The recovered weights converge to the true bias 2.0 and slope 3.0 that generated the data.

Constraints

• 1 <= n <= 10^5, 1 <= d <= 100; use vectorized NumPy, not Python loops over rows.
• The only loop is over n_iters gradient steps.
• lr is small enough that the loss converges (does not diverge to inf/nan) for reasonably scaled features.
• Values fit in float64.

Approach

Augment X with a leading column of ones so the bias is folded into the weight vector as w[0]; then a single X_b @ w produces predictions and a single X_b.T @ residual produces the full gradient including the bias. Initialize w to zeros and, on each of n_iters steps, compute the residual X_b @ w - y, form the MSE gradient (2/n) * X_b.T @ residual, and step w against it scaled by the learning rate. Everything is vectorized: the only Python loop is over the gradient steps.

Solution

import numpy as np

def linear_regression_gd(
    X: np.ndarray,
    y: np.ndarray,
    lr: float = 0.01,
    n_iters: int = 1000,
) -> np.ndarray:
    X = np.asarray(X, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64).ravel()
    n = X.shape[0]

    # Prepend a bias column of ones: w[0] becomes the intercept.
    X_b = np.hstack([np.ones((n, 1)), X])   # shape (n, d + 1)
    w = np.zeros(X_b.shape[1])              # shape (d + 1,)

    for _ in range(n_iters):
        residual = X_b @ w - y              # shape (n,)
        grad = (2.0 / n) * (X_b.T @ residual)  # shape (d + 1,)
        w -= lr * grad

    return w

The bias column means w[0] multiplies a constant 1 for every row, so it learns the intercept with no special-casing. The gradient (2/n) * X_b.T @ residual is the derivative of the mean squared error with respect to w; subtracting lr * grad walks the weights downhill.

Walkthrough

On the example with true bias 2.0 and slope 3.0:

• X is (200, 1), so X_b is (200, 2) with column 0 all ones and column 1 the feature. w starts as [0.0, 0.0].
• Iteration 1: residual = X_b @ [0, 0] - y = -y. The gradient (2/n) * X_b.T @ (-y) has a large negative entry for both bias and slope (since y is mostly positive), so w steps up toward positive values.
• Over successive iterations the residual shrinks as predictions approach y; the gradient magnitude falls and w settles.
• After 2000 steps at lr = 0.1, w converges to roughly [2.0, 3.0] — w[0] recovers the intercept and w[1] the slope. round(w[0], 1), round(w[1], 1) gives (2.0, 3.0).

Complexity & notes

• Time O(n_iters * n * d), space O(n * d) — each step does two matrix-vector products, each O(n * d), repeated n_iters times; the augmented X_b dominates memory.
• The bias-column trick is the cleanest way to include an intercept: no separate scalar to track, and the same gradient formula covers it. The alternative is keeping b separate with gradient (2/n) * sum(residual).
• Learning rate is the classic pitfall: too large and the residual grows each step, sending the loss to inf/nan; too small and it will not converge within n_iters. In practice you standardize features so a single lr works across dimensions.
• The 2/n factor comes from differentiating the mean of squared errors. Dropping the constant 2 (or the 1/n) still converges, just with an effectively rescaled learning rate — but matching the exact MSE gradient is what an interviewer checks.
• For a well-conditioned problem you would use the closed form w = (X_b.T @ X_b)^{-1} @ X_b.T @ y; gradient descent is the answer they want here because it scales to large n and generalizes to models with no closed form.