Easy Primitives

Implement Gradient Descent Step

Gradient descent updates parameters in the direction opposite the loss gradient. The gradient tells you how the loss changes locally as each parameter changes. The learning rate controls how far to move.

For weights [5, -3], gradients [2, -1], and learning rate 0.1, the updated weights are [4.8, -2.9]. A positive gradient decreases that parameter; a negative gradient increases it.

The rule is new_weights = weights - lr * gradients. Here gradients means the derivative of the loss with respect to the weights, not the derivative of the weights themselves. A finite step need not reduce the loss if the learning rate is too large.

Your task

Implement gradient_descent_step(weights, gradients, lr) for matching floating-point tensors of any shape and a nonnegative Python float lr. Return the updated weights with the same shape, applying the rule to each element. Do not reduce or normalize the gradient.

The gradients are supplied. You do not compute a loss, call backward, zero a gradient buffer, or run an optimizer object. Return the result without changing either input tensor. With lr=0, the weight values stay the same; a zero gradient leaves its corresponding weight unchanged.

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…

Implement Gradient Descent Step

Your task

Update parameters in a direction that reduces the loss

Compute all gradients before changing the parameters

Keep batch weighting consistent

Your code, saved here

Stuck?

PyTorch reference solution