Adam Optimizer from Scratch

1. Combining Momentum with Adaptive Learning Rates

The central idea behind Adam is to synthesize two powerful concepts we've touched upon:

1. Momentum: Using a moving average of the gradient to determine the update direction, providing inertia and smoothing the optimization path. This is the first moment estimate.
2. Adaptive Learning Rates: Scaling the learning rate for each parameter individually based on its past gradients. Parameters that receive large or noisy gradients get smaller updates, while those with small or consistent gradients get larger updates. This is achieved using the second moment estimate. Predecessors like Adagrad and RMSProp pioneered this idea.

Adam combines these into a single, cohesive algorithm. The following video provides a concise overview of how these pieces fit together.

2. The Moment Estimates and the Bias Problem

As introduced in the video, Adam maintains two running averages:

• First Moment Estimate ($m_t$): This is the EWMA of the gradients, which functions as the momentum component. Let $g_t=\nabla f(w_t)$ be the gradient at timestep $t$.

  $$m_t={\beta }_1{m}_{t-1}+(1-{\beta }_1)g_t$$

  Here, ${\beta }_1$ is the exponential decay rate, typically close to 1 (e.g., 0.9). This is exactly analogous to the velocity vector in the classical momentum method we studied previously.

• Second Moment Estimate ($v_t$): This is the EWMA of the squared gradients. This term will be used to scale the learning rate.

  $$v_t={\beta }_2{v}_{t-1}+(1-{\beta }_2)g_t^2$$

  Here, $g_t^2$ denotes the element-wise square of the gradient vector. ${\beta }_2$ is the decay rate for this average, usually set higher than ${\beta }_1$ (e.g., 0.999) to ensure the adaptive learning rate is stable.

A subtle but critical issue arises from the standard practice of initializing the moment vectors $m_0$ and $v_0$ to zero. Because they start at zero, and ${\beta }_1$ and ${\beta }_2$ are close to one, the estimates $m_t$ and $v_t$ will remain biased towards zero during the initial stages of training. This can lead to undesirably small update steps right when the model could be making rapid progress.

The same video explains this "startup problem" and the high-level idea behind its solution.

3. Deriving the Bias Correction

Now, let's derive the correction factors that transform the biased estimates $m_t$ and $v_t$ into unbiased ones, ${\hat{m}}_t$ and ${\hat{v}}_t$. This derivation is a classic example of analyzing the expected value of a recursive estimator, a technique you are likely familiar with from your work with time-series or signal processing.

Let's focus on the second moment, $v_t$. The update rule is $v_t={\beta }_2{v}_{t-1}+(1-{\beta }_2)g_t^2$. If we unroll this recursion back to $t=0$, assuming $v_0=0$, we get:

To analyze the bias, we take the expectation of $v_t$. We make a simplifying assumption that the second moment of the gradient is stationary, i.e., $E[g_i^2]=E[g^2]$ for all $i$.

The summation is a finite geometric series ${\sum }_{j=0}^{t-1}{\beta }_2^j=\frac{1-{\beta }_2^t}{1-{\beta }_2}$. Substituting this in, we find:

This equation clearly shows the bias. The expected value of our estimator $v_t$ is not the true second moment $E[g^2]$, but is instead scaled by a factor of $(1-{\beta }_2^t)$.

To correct this, we can simply divide our estimate $v_t$ by this factor. This gives us the bias-corrected second moment estimate:

By construction, $E[{\hat{v}}_t]=E[g^2]$, so our new estimator is unbiased. An identical derivation for the first moment yields the bias-corrected first moment estimate:

The following document provides a concise summary of this derivation.

4. The Complete Adam Update Rule

With the bias-corrected moment estimates in hand, we can now formulate the final parameter update rule for Adam. At each step, the parameters $w$ are updated as follows:

Let's analyze this expression:

• $\alpha$ is the master learning rate.
• The update direction is given by ${\hat{m}}_t$, our momentum-based estimate of the mean gradient.
• The denominator, $\sqrt{{\hat{v}}_t}+ϵ$, provides the per-parameter adaptive learning rate. For parameters where the historical gradients have been large (large ${\hat{v}}_t$), the effective learning rate is reduced. For parameters with small historical gradients, the effective learning rate is increased.
• $ϵ$ is a small constant (e.g., $10^{-8}$) added for numerical stability to prevent division by zero, in case ${\hat{v}}_t$ is very close to zero.

This single, elegant equation combines the smooth, directed path of momentum with the fine-grained, adaptive step sizes of RMSProp-like methods, all while correcting for initialization bias.

The two resources below give a full picture of the algorithm, from pseudocode to a concrete numerical example.

Conclusion

In this lesson, we have derived the Adam optimizer from its constituent parts. We saw how it elegantly merges the concepts of momentum and adaptive learning rates, and we performed the crucial derivation of the bias correction terms needed to ensure its stability and effectiveness from the very first step of training.

Key Takeaways:

• Adaptive Moment Estimation (Adam) combines a first moment estimate (momentum) for update direction with a second moment estimate (uncentered variance) for per-parameter adaptive scaling of the learning rate.
• The moment estimates $m_t$ and $v_t$ are calculated using exponentially weighted moving averages of the gradient and its square, respectively.
• Initializing the moments to zero introduces a bias that causes initial update steps to be too small.
• This bias can be mathematically derived and corrected by dividing the raw moment estimates by factors of $(1-{\beta }_1^t)$ and $(1-{\beta }_2^t)$.
• The final Adam update rule uses these bias-corrected estimates to provide robust and efficient optimization across a wide range of deep learning problems.

While Adam automates the tuning of per-parameter learning rates, the global learning rate $\alpha$ remains a critical hyperparameter. In our next lesson, we will discuss strategies for managing this global learning rate over the course of training, such as learning rate schedules, to further improve convergence and final model performance.