Medium End-to-End

Train with Adam End-to-End

Implement train_with_adam: fit a linear regressor using full-batch mean squared error and the Adam update rule. Compute the gradient and optimizer updates yourself, without a packaged optimizer or automatic differentiation.

Return the final weights and both moment estimates so the caller can inspect the complete state produced by this call.

Inputs and output

x: (N, d), with one example per row.
y: (N,), regression targets.
w0: (d,), initial weights.
m0: (d,), initial first-moment estimate, which averages past gradients.
v0: (d,), initial second-moment estimate, which averages squared gradients.
lr: a finite, nonnegative learning rate.
beta1 and beta2: decay factors in [0, 1).
eps: a finite, positive number added to the update denominator.
n_steps: a nonnegative integer, the number of full-batch updates.

N and d are positive. All shapes are consistent. Tensor values are finite, v0 is elementwise nonnegative, and all tensors share a floating-point dtype and device. There is no separate bias parameter; a supplied column of ones can act as one. Leave all inputs unchanged.

Return a vector of shape (3*d,), preserving dtype and device, in this order:

[final_w, final_m, final_v]

Each group contributes d entries. Return the stored moments m and v, not the bias-corrected m_hat and v_hat used in the weight update.

Loss and gradient

At the current weights w, predictions are x @ w. The loss is the mean of the N squared prediction errors, with no factor of one half:

loss = sum((x @ w - y)**2) / N
grad = (2/N) * x.T @ (x @ w - y)

Recompute the gradient from the current weights at every step. Use all examples each time, without shuffling or mini-batches.

Adam state and update order

Initialize w, m and v from w0, m0 and v0. For each step t from 1 through n_steps, compute grad and then apply:

m = beta1 * m + (1 - beta1) * grad
v = beta2 * v + (1 - beta2) * grad**2
m_hat = m / (1 - beta1**t)
v_hat = v / (1 - beta2**t)
w = w - lr * m_hat / (sqrt(v_hat) + eps)

Products, squares, square roots and division between moment vectors are elementwise. Use the newly updated m and v for bias correction, and put eps outside the square root. Carry w, m and v into the next step. Do not add weight decay, AMSGrad or other optimizer variants.

Starting moments at zero initially pulls their averages toward zero. The bias-correction denominators compensate for that initialization and approach one as t grows. This exercise also accepts nonzero m0 and v0: preserve them, but still start t at one for this call and apply exactly the formulas above. There is no previous-step-count argument, so this is not a general API for resuming an Adam run with its original step count.

Boundaries and worked example

With n_steps=0, return the initial weights and moments concatenated. With lr=0, weights stay fixed, but moments still update. A zero current gradient can also leave a nonzero weight update if m still contains earlier gradients.

For x=[[1]], y=[1], w0=m0=v0=[0], lr=0.3, beta1=0.5, beta2=0.75, eps=1 and one step:

grad = -2
m = -1; v = 1
m_hat = -2; v_hat = 4
w = 0 - 0.3 * (-2) / (2 + 1) = 0.2

Return [0.2, -1, 1]. The large eps here makes its position in the formula easy to check.

See Adam: A Method for Stochastic Optimization for the optimizer's derivation.

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

Use previous gradients to guide the next update

Plain gradient descent uses the current gradient. Momentum also keeps a running update direction:

$$v_t=\beta v_{t-1}+g_t.$$

Here, g_t is the current gradient and beta controls how much of the previous direction remains. The parameter update subtracts the learning rate times v_t. Some definitions multiply the new gradient by (1 - beta) as well. That changes the scale of the updates, so use the convention in the Problem tab.

Adam keeps two running averages

Adam tracks both the gradient and its square, element by element:

$$m_t=\beta_1m_{t-1}+(1-\beta_1)g_t,$$ $$v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2.$$

The first average helps estimate the update direction. The second tracks the size of recent gradients. It is a mean square, not a variance: Adam does not subtract the mean before squaring.

Starting both averages at zero makes their early values too small. Adam corrects for that initialization using

$$\hat m_t=\frac{m_t}{1-\beta_1^t},\qquad \hat v_t=\frac{v_t}{1-\beta_2^t}.$$

It then subtracts the learning rate times m_hat / (sqrt(v_hat) + epsilon) from the parameters. The first update uses t = 1.

For a first scalar gradient of 2, the corrected averages are 2 and 4. Ignoring epsilon, the update direction is 2 / sqrt(4) = 1. A gradient of 20 would also give direction 1 on that first step. This is why Adam's learning rate does not have the same effect as the learning rate in plain SGD.

Preserve state between steps

Use the newly updated averages when computing the parameter update. Keep those averages and the step counter for the next step; resetting them changes the optimizer.

Follow the specified placement of epsilon. Adding it inside the square root is a different formula. Also remember that a zero current gradient does not always mean no parameter change: stored momentum or weight decay can still cause an update.

Detecting runtime…