Mini-Batch Training

Train a one-feature linear regression model for one pass through the data. Update its weight and bias after each consecutive mini-batch, so later batches use the parameters learned from earlier ones.

Your task

Implement mini_batch_train(x, y, w_init, b_init, lr, batch_size).

x and y have shape (N, 1), with N > 0.
w_init has shape (1, 1) and b_init has shape (1,).
These floating-point tensors share a dtype and device.
lr is a nonnegative learning rate; batch_size is a positive integer.

Return a dictionary containing the final "w" tensor of shape (1, 1), the final "b" tensor of shape (1,), and "final_loss" as a Python float. Preserve the parameter dtype and device. Leave all input tensors unchanged.

Process batches in order

Start with copies of w_init and b_init. Use rows 0 through batch_size-1 for the first batch, the next consecutive rows for the second, and so on. Do not shuffle, skip or repeat rows. The last batch may be smaller, and it still gets one update. If batch_size is at least N, there is just one full-data update.

For a batch xb, yb containing B actual rows, compute:

prediction = xb @ w + b
residual = yb - prediction
grad_w = (-2 / B) * (transpose(xb) @ residual)
grad_b = (-2 / B) * sum(residual over rows)
w_next = w - lr * grad_w
b_next = b - lr * grad_b

Both gradients use the same current w and b. Compute them before changing either parameter. Use B, not the configured batch_size or total N, in the mean-squared-error gradient. This matters for a short final batch.

For x=[1, 2], y=[2, 4], initial w=0, b=0, lr=0.1 and batch_size=2, the residuals are [2, 4]. The gradients are -10 and -6, giving w=1 and b=0.6. The updated predictions are [1.6, 2.6], so the final loss is (0.4**2 + 1.4**2) / 2 = 1.06.

Report loss after all updates

Using the final parameters, evaluate every original row again:

final_loss = mean((y - (x @ w + b))**2)

This is not the last batch's loss or the average of losses measured while parameters were changing. Perform exactly one epoch; do not run until the loss converges. With lr=0, the parameters stay unchanged, but still compute their loss over the full dataset.

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…