One Adam step with bias correction
Medium14 min10 hidden tests
A recommendation team writes its own optimizer so it can log every internal value. Write adam_step(w, grad, m, v, t, lr, beta1, beta2, eps) that applies one Adam update. Adam keeps m, a running average of the gradients, and v, a running average of their squares. Both start at 0, so early on they are too small, and Adam divides them by 1 - beta ** t to correct that. Return the new weights and the new m and v to store for the next step. On the very first step from zeros, each weight moves by almost exactly lr against the sign of its gradient: [1.0, -2.0] becomes [0.99, -1.99] with lr = 0.01.
Drag the step t
w ← w − lr · m̂ / (√v̂ + eps)Gradient 0.5 at every step. At t = 10: m = 0.3257, v = 0.002489, still small because both started at 0. Corrected, m̂ = 0.5 and v̂ = 0.25, so the weight moves 0.01, which is lr. Without the correction it would move 0.0653, 6.5 times lr.
Step 10. Corrected step 0.01, uncorrected 0.0653.
Examples
First step: each weight moves by lr
adam_step([1.0, -2.0], [0.5, -1.0], [0.0, 0.0], [0.0, 0.0], 1, 0.01)
returns ([0.99, -1.99], [0.05, -0.1], [0.00025, 0.001])
A later step that carries momentum from earlier gradients
adam_step([0.5], [0.2], [0.1], [0.01], 3, 0.1)
returns ([0.477812], [0.11], [0.01003])
Requirements
w,grad,mandvare 1-D lists or arrays of the same length.tis the number of this step, counting from 1.m = beta1 * m + (1 - beta1) * gradandv = beta2 * v + (1 - beta2) * grad ** 2.m_hat = m / (1 - beta1 ** t),v_hat = v / (1 - beta2 ** t), thenw = w - lr * m_hat / (sqrt(v_hat) + eps), withepsoutside the square root.- Return a tuple
(w, m, v)of three float arrays. Store the uncorrectedmandv.
The idea behind it
Adam (adaptive moment estimation) is the default optimizer for most deep learning. It keeps two exponential moving averages per weight: m, the mean of recent gradients, which acts as momentum, and v, the mean of recent squared gradients, which measures how large the gradient usually is. The step m / sqrt(v) gives each weight its own scale: a weight with large, noisy gradients takes smaller steps than one with small, steady ones. Since both averages start at 0, they are too small for the first few steps. After t steps, a constant gradient g gives m = (1 - beta1 ** t) * g, so dividing by 1 - beta1 ** t restores g. This is bias correction. Typical values are beta1 = 0.9, beta2 = 0.999 and eps = 1e-8, the defaults in PyTorch.