Katkat
Back

Implement SGD for linear regression and derive gradients

AmazonApplied Scientist · Technical Screen
MediumApplied ScientistTechnical ScreenMachine Learning

Last updated

This question evaluates understanding of gradient-based optimization and linear model training, including competency in deriving gradients for mean squared error and implementing stochastic (or mini-batch) gradient descent for parameter estimation.

Prompt

You are given a dataset of (n) 1D samples ({(x_i, y_i)}_{i=1}^n), where (x_i) and (y_i) are real numbers.

We want to fit a linear model: [ \hat{y} = a x + b ] by minimizing the mean squared error (MSE).

Tasks

  1. Define the loss function for this problem (e.g., MSE over the dataset).
  2. Using the chain rule / backprop-style reasoning, derive the gradients (\frac{\partial L}{\partial a}) and (\frac{\partial L}{\partial b}).
  3. Describe (and optionally write pseudocode for) how to train (a) and (b) using SGD (or mini-batch SGD):
  • parameter initialization
  • per-step gradient computation
  • update rule
  • learning rate choice / scheduling
  • stopping criteria
  1. Discuss common pitfalls and edge cases (e.g., scaling, divergence, choosing batch size).

Output / Expected Result

After training, return the learned parameters (a) and (b) that approximately minimize the chosen loss on the provided data.