Easy End-to-End

Training Loop

Implement training_loop for a linear model with one input feature: prediction = x @ w + b. Starting from the supplied weight and bias, take exactly n_steps full-batch gradient-descent updates and report the final mean squared error.

Inputs and return value

x and y: shape (N, 1), with N > 0.
w_init: shape (1, 1), the initial weight.
b_init: shape (1,), the initial bias, broadcast across all examples.
lr: a finite, nonnegative learning rate.
n_steps: a nonnegative integer.

Data and parameters are finite floating-point tensors with the same dtype and device. Leave all inputs unchanged. Start from w_init and b_init, not zeros.

Return a dictionary with exactly these keys:

w: final weight tensor, shape (1, 1).
b: final bias tensor, shape (1,).
final_loss: an ordinary Python float, calculated with the final w and b.

Preserve the parameter tensors' dtype and device. Do not return the loss from before the last update. With n_steps=0 or lr=0, the returned parameters have their initial values, and final_loss still needs to be evaluated.

One training step

Every step uses all N examples in their supplied order:

prediction = x @ w + b
error = prediction - y
grad_w = (2/N) * x.T @ error
grad_b = (2/N) * sum(error, over the example axis)
w = w - lr * grad_w
b = b - lr * grad_b

The loss is mean((prediction-y)**2), with no factor of one half. Compute both gradients from the same pre-update prediction, then update both parameters. The next step uses the new w and b. Calculate gradients directly; do not use autograd or an optimizer helper. There is no sampling, regularization, momentum, clipping or early stopping.

After all updates, recompute prediction and return its mean squared error. A large learning rate can make loss increase or oscillate; still perform the requested number of updates. This exercise does not guarantee convergence.

Worked example

For x=[[1],[2]], y=[[2],[4]], w_init=[[0]], b_init=[0], lr=0.1 and n_steps=1:

Initial predictions are [0,0], so errors are [-2,-4].
grad_w = (2/2)(1(-2)+2*(-4)) = -10 and grad_b = (2/2)*(-2-4) = -6.
The updated weight is [[1]] and bias is [0.6].
Final predictions are [1.6,2.6]. The returned loss is (0.42+1.42)/2 = 1.06.

The initial loss is 10, but final_loss must be 1.06 for this one-step call.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Update parameters in a direction that reduces the loss

Gradient descent updates model parameters theta using the loss gradient:

$$\theta_{t+1}=\theta_t-\eta\nabla L(\theta_t).$$

The gradient describes how the loss changes near the current parameters. Its negative points in the direction of steepest local decrease when distance is measured in the usual Euclidean way. The learning rate eta sets the step size. A step that is too large can increase the loss, even when its direction is correct.

For L(w) = (w - 3)^2 at w = 0, the gradient is -6:

With learning rate 0.1, the new weight is 0 - 0.1 * (-6) = 0.6. The loss falls from 9 to 5.76.
With learning rate 1, the new weight is 6. The step passes the minimum at 3, and the loss is still 9.

How quickly the gradient changes, and how the input features are scaled, affect which learning rates work well.

Compute all gradients before changing the parameters

For an ordinary gradient-descent step, compute every gradient from the same parameter values. Updating one weight before computing another weight's gradient changes the algorithm.

When recording a loss history, state whether each value comes before or after an update. Including the initial loss gives one more entry than the number of updates.

Keep batch weighting consistent

A mini-batch gradient estimates the gradient over the full training set. Different batches can give different gradients, so a noisy loss curve is not by itself a bug.

Gradient accumulation can reproduce a larger batch's gradient when the losses and model behavior match. If smaller batches have different sizes, weight their mean gradients by their example counts. Simply averaging those means gives the smaller batches too much influence. Layers such as batch normalization can also behave differently when a batch is split.

Detecting runtime…