Linear regression by gradient descent
Fit ŷ = Xw + b by batch gradient descent on the mean squared error
L = (1/n) Σ (xi·w + b − yi)². Start from w = 0 and b = 0. At each of
steps iterations, compute both gradients from the current parameters, then update
w ← w − lr·∂L/∂w and b ← b − lr·∂L/∂b.
Return the final parameters as a list [w₁, …, w_d, b].
Examples
Input X = [[1], [2], [3]], y = [2, 4, 6], lr = 0.1, steps = 1
Output [1.8667, 0.8]
One step from zero. Both gradients come from the same residuals, which are −y.
Input X = [[1], [2], [3]], y = [3, 5, 7], lr = 0.05, steps = 2000
Output [2, 1]
Converged: y = 2x + 1.
Input X = [[1, 2], [3, 4]], y = [1, 2], lr = 0.01, steps = 0
Output [0, 0, 0]
Zero steps returns the starting point.
Constraints
- 1 ≤ n ≤ 1000, 1 ≤ d ≤ 20
- 0 ≤ steps ≤ 5000
- lr is small enough that the iteration converges
- Checked to a relative tolerance of 10⁻⁷
Hints
Hint 1
Write the residual vector r = Xw + b − y. Both gradients are simple functions of r.
Hint 2
∂L/∂w = (2/n)Xᵀr and ∂L/∂b = (2/n)Σr. Check the shapes: Xᵀr has shape (d,).
Hint 3
Compute r once per step, then update both parameters from it. Recomputing r between the two updates changes the algorithm.
Where candidates go wrong
- Updating w, then computing the residuals again before updating b. The tests expect a simultaneous update.
- Dropping the factor 2. The fitted line is the same eventually, but every intermediate value differs.
- Looping over examples in Python. Correct, and a hundred times slower than X.T @ r.
- Passing y with shape (n, 1): X @ w − y then broadcasts to (n, n).
The idea
With residuals r = Xw + b − y, the loss is (1/n)rᵀr, so ∂L/∂w = (2/n)Xᵀr and ∂L/∂b = (2/n)1ᵀr. One iteration is one matrix-vector product to form r and one to form Xᵀr: O(nd) work and no loop over examples.
The loss is a convex quadratic, and gradient descent on it converges for any learning rate below 2/λmax, where λmax is the largest eigenvalue of the Hessian (2/n)[X 1]ᵀ[X 1]. The rate of convergence is set by the ratio of the largest to the smallest eigenvalue, which is why standardising features matters so much in practice.
Time O(steps · nd) · Space O(nd)
Reference solution
import numpy as np
def fit_linear_gd(X, y, lr, steps):
X = np.asarray(X, dtype=float)
y = np.asarray(y, dtype=float)
n, d = X.shape
w, b = np.zeros(d), 0.0
for _ in range(steps):
r = X @ w + b - y # residuals, shape (n,)
gw = (2.0 / n) * (X.T @ r) # dL/dw, shape (d,)
gb = (2.0 / n) * r.sum() # dL/db
w -= lr * gw # both gradients came from the same r
b -= lr * gb
return [*w.tolist(), float(b)]Follow-ups an interviewer asks
- Turn it into mini-batch stochastic gradient descent with a fixed seed. What changes in the update?
- Add momentum: v ← βv + g, θ ← θ − lr·v.
- Stop early when the relative change in loss falls below 10⁻¹⁰ instead of running a fixed number of steps.
- Find the largest stable learning rate from the eigenvalues of the Hessian, and show what happens just above it.
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
The question is less about regression than about whether you can turn a loss into an update rule and vectorise it. Interviewers watch for the factor of 2, the 1/n, the order of the updates, and whether you loop over examples when you do not need to.
The theory behind it is in the lesson Linear and logistic regression from scratch.
Keep going
Ridge regression in closed form
Solve (XᵀX + λI)w = Xᵀy without forming an inverse.
Open the problem → Medium15 minLinear regression by gradient descent
Fit weights and a bias by batch gradient descent on mean squared error, vectorised in NumPy.
Open the problem → Medium20 minLogistic regression with an L2 penalty
Train logistic regression by gradient descent with an L2 penalty and a stable sigmoid, from scratch in NumPy.
Open the problem →