Gradient Accumulation

Gradient accumulation combines gradients from several small batches before updating the weights. Each small batch is called a microbatch. Here, compute one full-batch SGD update for a linear regressor by processing equal-sized microbatches in turn.

Your task

Implement gradient_accumulation_step(weights, xs, ys, lr, accum_steps):

weights has shape (d,) and contains the current parameters.
xs has shape (accum_steps, B, d).
ys has shape (accum_steps, B).
lr is a nonnegative learning rate.
accum_steps, B and d are positive integers.

accum_steps matches the first dimension of both input arrays. The model predicts x @ weights, with no bias or activation. Its objective is the mean squared error over all accum_steps * B examples.

Accumulate first, update once

Start a zero gradient accumulator with the same shape as weights. For each microbatch, compute predictions and its mean-loss gradient using the original weights:

grad_total = zeros_like(weights)
for k in range(accum_steps):
    x = xs[k]
    y = ys[k]
    prediction = x @ weights
    grad = (2 / B) * x.T @ (prediction - y)
    grad_total = grad_total + grad

average_grad = grad_total / accum_steps
updated_weights = weights - lr * average_grad

Use this explicit gradient formula rather than automatic differentiation. Return updated_weights of shape (d,) and preserve all inputs. Do not update weights inside the loop: later microbatches must use the same weights as earlier ones.

Each microbatch gradient already averages over its B examples. Averaging those gradients gives the gradient of the combined mean loss because all microbatches have the same size. Dividing by accum_steps twice would make the update too small; omitting it would make the update too large.

Example

Let weights = [1], xs = [[[1]], [[1]]] and ys = [[0], [2]]. The two microbatch gradients are [2] and [-2]. Their average is zero, so return [1] for any learning rate. Updating after the first microbatch would change the second gradient and produce a different result.

For this deterministic linear model, accumulation matches one update on the concatenated batch, apart from floating-point rounding. In a larger model, that equivalence also depends on operations such as batch normalization, which can behave differently across microbatches.

Accumulation can reduce peak intermediate-activation memory when a training loop loads and processes one microbatch at a time. This exercise supplies all inputs as stacked arrays; it demonstrates the gradient calculation, not a data-loading or memory-management system.

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…