The Training Loop - PyTorch Basics

The Training Loop

What you'll build: train_linear_model(X, y, lr, epochs), which fits a linear model with full-batch gradient descent and records the loss at every epoch.

Background

Every PyTorch model, from linear regression to an LLM, is trained by the same five-line loop:

for epoch in range(epochs):
    optimizer.zero_grad()          # 1. clear the gradients from the last step
    pred = model(X)                # 2. forward pass
    loss = loss_fn(pred, y)        # 3. compute the loss
    loss.backward()                # 4. backprop: fill p.grad for every parameter
    optimizer.step()               # 5. update the parameters

Forgetting zero_grad() is the classic bug: .grad accumulates, so every step would use the sum of all past gradients.

Formula

Fit y_hat = X @ w + b with full-batch gradient descent on the mean squared error, with N samples:

L     = mean((X @ w + b - y) ** 2)
dL/dw = (2 / N) * X.T @ (X @ w + b - y)
dL/db = (2 / N) * sum(X @ w + b - y)
w     = w - lr * dL/dw
b     = b - lr * dL/db

Inputs

• X: float array (N, d); y: float array (N,).
• lr: learning rate (default 0.1); epochs: number of gradient steps (default 100).

Edge Cases and Conventions

• Start from zeros: w = 0 (shape (d,)), b = 0.
• Each epoch, compute the loss with the current parameters, append it to the history, then take one step. So losses[0] is the loss of the all-zero model, mean(y ** 2).

Constraints

• NumPy track: write the gradients by hand (what loss.backward() does for you in PyTorch).
• PyTorch track: use nn.Linear, nn.MSELoss and torch.optim.SGD, and let loss.backward() compute the gradients.

Expected Return Format

Return a dictionary with these keys:

• 'losses': list of epochs Python floats, the loss before each update
• 'w': learned weights, shape (d,)
• 'b': learned bias, a Python float

Follow-up

• Switch to minibatches (next problem). Why does the loss curve become noisy, and why is that often fine?
• Where would model.train() / model.eval() and torch.no_grad() go in a loop with a validation set?

Input: train_linear_model(np.array([[1.0], [2.0], [3.0]]), np.array([2.0, 4.0, 6.5]), lr=0.05, epochs=4)
Output: {'losses': [20.75, 4.2638, 0.9907, 0.3383], 'w': array([1.716]), 'b': 0.6924}

Input: X = np.random.randn(64, 2); y = X @ [2, -3] + 0.5; train_linear_model(X, y, lr=0.1, epochs=200)
Output: {'losses': [..., ~0.0], 'w': array([ 2., -3.]), 'b': 0.5}

import numpy as np

def train_linear_model(X, y, lr=0.1, epochs=100) -> dict:
    """
    Fit y ~ X @ w + b with full-batch gradient descent on the MSE.

    Args:
        X: Inputs of shape (N, d)
        y: Targets of shape (N,)
        lr: Learning rate
        epochs: Number of gradient steps

    Returns:
        dict with 'losses' (list of floats, loss before each step),
        'w' (array of shape (d,)) and 'b' (float)
    """
    w = np.zeros(X.shape[1])
    b = 0.0
    losses = []
    # Your code here: forward, loss, gradients, update
    pass