Katkat
Back

Implement Gradient Descent Regression

DatabricksMachine Learning Engineer · Onsite
MediumMachine Learning EngineerOnsiteMachine Learning

Last updated

This problem tests your understanding of linear regression by requiring the implementation of a gradient descent optimizer from scratch to minimize the Mean Squared Error loss.

Problem Statement

Implement a linear regression model that learns parameters w (weights) and b (bias) using gradient descent. Your goal is to minimize the Mean Squared Error (MSE) loss function defined as: J(w, b) = (1/n) * sum((y_pred_i - y_true_i)^2)

where y_pred = Xw + b.

Requirements

Implement a class LinearRegressionGD with the following methods:

  1. fit(X, y, learning_rate, iterations): Updates weights and bias iteratively.
  2. predict(X): Returns the predicted values.

Constraints

  • X is a 2D numpy array of shape (n_samples, n_features).
  • y is a 1D numpy array of shape (n_samples,).
  • Ensure your gradient update steps are vectorized for performance.

Example

import numpy as np
X = np.array([[1], [2], [3]])
y = np.array([2, 4, 6])
model = LinearRegressionGD()
model.fit(X, y, learning_rate=0.01, iterations=1000)
print(model.predict(np.array([[4]]))) # Should be close to 8

Edge Cases

  • What happens if the learning rate is too large?
  • How should the model handle feature scaling (e.g., if one feature is in range [0, 1] and another in [0, 1000])?

Hints

Hint 1

The partial derivative of the MSE loss with respect to weight `w` is `(2/n) * X.T.dot(y_pred - y_true)`.

Hint 2

The partial derivative of the MSE loss with respect to bias `b` is `(2/n) * np.sum(y_pred - y_true)`.

Hint 3

Remember to initialize your weights and bias, typically to zeros or small random values, before starting the loop.

Follow-up questions

  • How would you modify this implementation to incorporate L2 regularization (Ridge regression)?
  • How does Stochastic Gradient Descent (SGD) differ from the batch gradient descent implemented here in terms of convergence speed and memory usage?
  • What diagnostic steps would you take if the loss function starts increasing during training?