Gradient Descent From Scratch

Question

Implement linear regression trained with batch gradient descent. You are not allowed to use scikit-learn or any closed-form solver (so no np.linalg).
def fit_linear_regression(X: np.ndarray, y: np.ndarray, lr: float, n_iters: int) -> tuple:
    pass
Input
X is an (n_samples, n_features) numpy array, y is an (n_samples,) numpy array, lr is the learning rate and n_iters the number of iterations.

Output
Return the tuple (w, b): the learned weight vector and intercept.

Also answer
Derive the gradient that your implementation uses, and explain what happens if the learning rate is chosen too small or too large.

Example Test

 	On data generated as y = 2x + 1 + noise, the function should recover w ≈ 2 and b ≈ 1.

Accepted Answer

Gradient descent is the workhorse of machine learning: instead of solving for the optimal parameters in one algebraic step, we repeatedly take small steps downhill on the loss surface. For linear regression the model and the mean squared error (MSE) loss are: \begin{equation} \hat{y}_i = x_i^T w + b \end{equation} \begin{equation} L(w, b) = \frac{1}{n} \sum_{i=1}^{n} (\hat{y}_i - y_i)^2 \end{equation} Deriving the gradientTake a single weight $w_j$ and apply the chain rule. The outer function is the square, the inner function is the prediction: \begin{equation} \frac{\partial L}{\partial w_j} = \frac{1}{n} \sum_{i=1}^{n} 2(\hat{y}_i - y_i) \cdot \frac{\partial \hat{y}_i}{\partial w_j} \end{equation} The prediction $\hat{y}_i = w_1 x_{i1} + ... + w_p x_{ip} + b$ depends on $w_j$ only through the term $w_j x_{ij}$, so \begin{equation} \frac{\partial \hat{y}_i}{\partial w_j} = x_{ij} \end{equation} which gives \begin{equation} \frac{\partial L}{\partial w_j} = \frac{2}{n} \sum_{i=1}^{n} (\hat{y}_i - y_i) \, x_{ij} \end{equation} Stacking this over all $j$ produces the compact matrix form (with $e = \hat{y} - y$ the error vector): \begin{equation} 
abla_w L = \frac{2}{n} X^T e \end{equation} For the intercept, $\frac{\partial \hat{y}_i}{\partial b} = 1$, so \begin{equation} \frac{\partial L}{\partial b} = \frac{2}{n} \sum_{i=1}^{n} (\hat{y}_i - y_i) \end{equation} The update ruleThe gradient points in the direction of steepest increase of the loss, so we step in the opposite direction, scaled by the learning rate $\eta$: \begin{equation} w \leftarrow w - \eta 
abla_w L, \qquad b \leftarrow b - \eta \frac{\partial L}{\partial b} \end{equation} The implementation is now a direct translation:
import numpy as np

def fit_linear_regression(X: np.ndarray, y: np.ndarray, lr: float = 0.05, n_iters: int = 5000) -> tuple:
    n_samples, n_features = X.shape
    w = np.zeros(n_features)
    b = 0.0
    for _ in range(n_iters):
        # Forward pass: current predictions and errors
        y_pred = X @ w + b
        error = y_pred - y

        # Gradients (derived above)
        grad_w = (2 / n_samples) * (X.T @ error)
        grad_b = (2 / n_samples) * np.sum(error)

        # Step against the gradient
        w -= lr * grad_w
        b -= lr * grad_b
    return w, b

# Test: recover y = 2x + 1
rng = np.random.default_rng(42)
X = rng.uniform(-2, 2, size=(500, 1))
y = 2.0 * X[:, 0] + 1.0 + rng.normal(0, 0.1, size=500)
w, b = fit_linear_regression(X, y)
print(w, b)  # approximately [2.0] and 1.0

Running this prints $w \approx 1.998$ and $b \approx 0.994$ — the true line up to the noise level.

The learning rateThe MSE loss of linear regression is a convex bowl, so gradient descent is guaranteed to converge — but only if $\eta$ is sensible:

 	Too smallevery step barely moves the parameters, and you need an enormous number of iterations. The loss creeps down painfully slowly.
 	Too largeeach step overshoots the bottom of the bowl and lands on the opposite, higher side. The loss oscillates and, beyond a critical value, diverges to infinity.

Two practical notes. First, feature scaling matters here too: if features have very different scales, the bowl becomes a long narrow valley, and a learning rate that is safe along the steep direction is far too slow along the shallow one — standardising the features makes the bowl round and gradient descent fast. Second, on large datasets one usually replaces the full-batch gradient with the gradient on a random mini-batch (stochastic gradient descent): each step is noisier but much cheaper, and the noise even helps escape saddle points in non-convex problems like neural networks.

Title	Category	Subcategory	Difficulty	Status
Batch, Stochastic and Mini-Batch Gradient Descent	Machine Learning	Optimization	Medium	Unlock
Newton's Method and Second-Order Optimization	Machine Learning	Optimization	Hard	Unlock

The Interview Guide for Quants and Traders

Gradient Descent From Scratch

What do you think of this page?