Katkat
Back

Implement MSE Loss and Linear Regression from Scratch

IndeedMachine Learning Engineer · Technical Screen
MediumMachine Learning EngineerTechnical ScreenMachine Learning

Last updated

Implement the mean squared error loss and a linear regression model with fit and predict methods, written from scratch rather than with a library model. It tests deriving and coding the gradients, choosing a training method, feature scaling, input validation, and verifying the implementation.

Implement the mean squared error (MSE) loss function, and then implement linear regression trained to minimize it. Write the code yourself: do not call a library's ready-made regression model or loss function.

Clarifying Questions

  • May you use NumPy for vectorized array math, or must the code be plain Python?
  • Should the model be trained with gradient descent, or is the closed-form least-squares solution acceptable?
  • Does the model need an intercept (bias) term, and must it support multiple input features?
  • Should any regularization be included?
  • How should the code handle invalid input, such as mismatched lengths or an empty dataset?

Part 1 — MSE loss

Write a function that takes the true targets and the predicted values and returns their mean squared error, $\frac{1}{n}\sum_{i=1}^{n}(y_i - \hat{y}_i)^2$. Then give its gradient with respect to the predictions, which Part 2 will need.

What This Part Should Cover

  • A correct, vectorized loss computation with input validation
  • The gradient with respect to the predictions, including its constant factor
  • Behavior on empty input and on very large values

Part 2 — Linear regression

Implement a linear regression model with a fit(X, y) method that learns weights and a bias by minimizing the MSE from Part 1, and a predict(X) method. X has one row per example and one column per feature.

What This Part Should Cover

  • Correct gradients with respect to the weights and the bias, and a working update loop
  • Step size, stopping criterion, and feature scaling
  • A way to verify the implementation, such as known data, a closed-form comparison, or a numerical gradient check
  • Time and memory cost per training iteration

Hints

Hint 1

Shapes first: Decide which input shapes you accept and what happens when the two inputs do not match, before you write the arithmetic.

Hint 2

Chain rule: Your Part 1 gradient is with respect to the predictions. Work out how a small change in each weight, and in the bias, changes the predictions.

Hint 3

Know when to stop: Decide how you will choose the step size and when training should end, and think about what happens when features have very different scales.

Follow-up questions

  • How would you add L2 regularization, and how does it change both the gradient and the closed-form solution?
  • The dataset no longer fits in memory. How do you train, and what changes about convergence?
  • When is MSE a poor loss for regression, and what would you use instead?
  • The closed-form solution fails on your data because two features are perfectly correlated. Why, and how do you fix it?