All Lessons

Bias Correction in Adaptive Moment Optimization

Adaptive optimizers begin with zero-valued moment estimates, creating a measurable early-step bias. This lesson derives the correction factors and shows why the first and second moment biases do not simply cancel.

AI Narration Press play to listen
0  / 6 paragraphs
Click any paragraph to jump · Scroll freely without breaking narration

Adaptive moment optimization uses the gradient's recent direction and scale. Given mini-batch gradient $g_t$, define $m_t=\beta_1m_{t-1}+(1-\beta_1)g_t$ and $v_t=\beta_2v_{t-1}+(1-\beta_2)(g_t\odot g_t)$. The first moment smooths direction; the second raw moment tracks squared magnitude coordinate by coordinate.

Both states normally start at zero. If gradients are locally stationary with $\mathbb{E}[g_t]=\mu$, then $\mathbb{E}[m_t]=(1-\beta_1^t)\mu$. Likewise, if $\mathbb{E}[g_t\odot g_t]=\nu$, then $\mathbb{E}[v_t]=(1-\beta_2^t)\nu$. Early estimates are therefore pulled toward zero by known multiplicative factors.

The corrected estimates are $\hat m_t=m_t/(1-\beta_1^t)$ and $\hat v_t=v_t/(1-\beta_2^t)$. A standard adaptive update is $\theta_t=\theta_{t-1}-\alpha\hat m_t/(\sqrt{\hat v_t}+\epsilon)$. The division and square root are element-wise, and $\epsilon$ protects the denominator from numerical instability.

At the first scalar step, let $g_1=2$, $\beta_1=0.9$, and $\beta_2=0.999$. Then $m_1=0.2$ and $v_1=0.004$, while correction gives $\hat m_1=2$ and $\hat v_1=4$. Ignoring $\epsilon$, the corrected ratio is $2/\sqrt4=1$, but the uncorrected ratio is $0.2/\sqrt{0.004}\approx3.16$. Different decay rates prevent the two raw biases from canceling.

Bias correction has a precise scope. Under stationary first and second moments, it removes the bias caused by zero initialization from each estimator. It does not make $\hat m_t/\sqrt{\hat v_t}$ an unbiased estimator of a population-preconditioned gradient, because a ratio of correlated random variables is nonlinear. It also does not repair a poor learning rate, corrupted gradients, or an unsuitable objective.

The optimizer step counter, $m_t$, and $v_t$ must be checkpointed together. Restoring only parameters restarts the moment estimates and correction schedule, changing the training trajectory. When gradients are accumulated across microbatches, increment the correction step once per optimizer update, not once per backward pass. Log the step count during resume tests so an off-by-one error is visible.