Answers · The mathematics underneath

Why does Adam need bias correction at the start of training?

In shortAdam needs bias correction because its two running averages start at zero, and without correction the early steps are several times larger than the learning rate asks for.

Adam needs bias correction because both running averages start at zero, so at early steps they are too small. Correcting each by its exact shortfall makes the first update equal to the learning rate; without correction the update is not cautious but several times too large.

The two averages are not symmetric. The gradient average with $\beta_1=0.9$ looks back about ten steps. The squared-gradient average with $\beta_2=0.999$ looks back about a thousand. At step one the first is at a tenth of its final size and the second is at a thousandth. The numerator therefore warms up quickly and the denominator warms up slowly, and that mismatch is why the uncorrected ratio overshoots rather than undershoots.

The exact shortfall

Expand the recursion for the gradient average. After $t$ steps the weights on the observed gradients are $(1-\beta_1)\beta_1^{t-i}$ , and they sum to $1-\beta_1^t$ rather than to one. The missing weight was spent on the zero with which the average started. The same calculation applies to the squared-gradient average with $\beta_2$ . That is why the correction divides by $1-\beta^t$ rather than adding a constant.

m_t=(1-\beta_1)\sum_{i=1}^{t}\beta_1^{\,t-i}g_i,\qquad \mathbb{E}[m_t]=\big(1-\beta_1^{\,t}\big)\,\mathbb{E}[g]\\[2pt] v_t=(1-\beta_2)\sum_{i=1}^{t}\beta_2^{\,t-i}g_i^2,\qquad \mathbb{E}[v_t]=\big(1-\beta_2^{\,t}\big)\,\mathbb{E}[g^2]

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

At $t=1$ with $\beta_1=0.9$ and $\beta_2=0.999$ , take a gradient $g$ that is the same at every step for the calculation. Then $m_1=0.1g$ and $v_1=0.001g^2$ . The square root is $\sqrt{v_1}=\sqrt{0.001}\,g\approx0.0316g$ , so the uncorrected ratio is $0.1/0.0316\approx3.16$ . The corrected quantities are $\hat m_1=0.1g/0.1=g$ and $\hat v_1=0.001g^2/0.001=g^2$ , so their ratio is one. The corrected first step is $\alpha\,g/(|g|+\epsilon)$ , approximately $\alpha$ .

FIG 1Two averages warming up at different speeds

the gradient average, beta 0.9the squared-gradient average, beta 0.999

The gradient average has essentially finished warming up by step ten. The squared-gradient average needs hundreds of steps. Correcting neither, or correcting only one, leaves their ratio wrong throughout that stretch.

Why both being small does not cancel

The natural guess is that if $m_t$ and $v_t$ are both short, the ratio $m_t/\sqrt{v_t}$ is roughly right. It is not, because the two averages are short by different factors. The uncorrected update is inflated relative to the corrected one by $(1-\beta_1^t)/\sqrt{1-\beta_2^t}$ . At $t=1$ that factor is $0.1/\sqrt{0.001}\approx3.16$ , not one. At $t=2$ the uncorrected ratio is $0.19/\sqrt{0.001999}\approx4.25$ .

The gap gets worse before it gets better. At $t=10$ the uncorrected ratio is $0.6513/\sqrt{0.009955}\approx6.53$ . By $t=1000$ the squared-gradient average has warmed up and the uncorrected ratio is $1.000/\sqrt{0.6323}\approx1.26$ . A learning rate tuned for corrected Adam can therefore be exceeded by a large factor in the first few dozen steps. A ramped learning rate does a similar job, but bias correction does it exactly for a zero initialisation.

Correcting only one average is worse than correcting neither in one direction and not enough in the other. If only $m_t$ is corrected at $t=1$ , the update is $g/\sqrt{0.001g^2}\approx31.6$ , still far too large. If only $v_t$ is corrected, the update is $0.1g/|g|=0.1$ , too small. Both factors are needed because the ratio contains both.

The correction is exact for the zero start. If the gradients were stationary, $\hat m_t$ would be an unbiased estimate of $\mathbb{E}[g]$ and $\hat v_t$ an unbiased estimate of $\mathbb{E}[g^2]$ . If the distribution changes, both corrected and uncorrected quantities estimate different things, but the initialisation bias is still removed. The $\epsilon$ of about $10^{-8}$ only stops a coordinate with no history from dividing by zero; it is not what fixes the bias.

What it does not fix

Bias correction repairs the start, not the diagonal approximation. If a valley runs at forty-five degrees to both coordinates, both coordinates have gradients of the same typical size, so the adaptive denominators are equal and the method behaves like plain momentum. Only a full matrix preconditioner would notice that curvature. Weight decay placed inside the loss has a related problem: its gradient is divided by $\sqrt{\hat v_t}$ , so the decay applied to a parameter is scaled by that parameter's gradient history. Applying the decay directly to the parameters is the difference between AdamW and Adam.

Worth keeping

At step one with the usual settings, uncorrected Adam takes a step about 3.16 times the learning rate; by step ten the factor is about 6.53.
Bias correction divides each average by one minus beta to the power t, which is the exact factor by which a zero-started geometric average is short.
It fixes the initialisation, not the diagonal limitation: a valley at forty-five degrees still looks like plain momentum.

Where this goes next

This question is one step of The Mathematics of Gradient Descent, a course that works through its subject from the beginning. The page that teaches this part of it is A Step Size Per Coordinate, and it reads the whole way through whether or not you have an account.

What an account changes is the practice: problems standing under every idea, marked with a reason rather than a tick, in a copy of the course that is yours.

Read that page The whole course