Implementing Adam Optimizer from Scratch

Recommended resources

The freely available resources this lesson will direct you to.

Introduction

In our last lesson, we developed two powerful, yet distinct, solutions to the challenges faced by Stochastic Gradient Descent. We derived Momentum, which acts like a heavy ball accumulating velocity to smooth out oscillations and accelerate through flat regions. We also derived RMSProp, which adaptively tunes the learning rate for each parameter, taking small steps in steep directions and large steps in flat ones. Momentum modifies the update direction, while RMSProp modifies the step size.

This naturally leads to a compelling question: can we get the best of both worlds?

In this lesson, we will do precisely that. We will combine the velocity accumulation of Momentum with the adaptive learning rates of RMSProp to derive the Adam (Adaptive Moment Estimation) optimizer. You will learn how it synthesizes these two ideas into a single, robust algorithm. We will also uncover a critical subtlety in its implementation—bias correction—and derive from first principles why it is necessary for stable and rapid convergence at the start of training. Finally, you will be prepared to implement the full Adam algorithm from scratch.

The Anatomy of Adam: Fusing Two Optimizers

Let's begin by recalling the core components of our previous optimizers. Momentum maintains an exponentially weighted moving average of the gradient, which is the first moment. RMSProp maintains an exponentially weighted moving average of the squared gradient, which is related to the uncentered second moment.

Adam's central idea is to maintain and utilize both of these moving averages.

First Moment (Momentum): It tracks the "velocity" of the gradients. Let's denote this vector as $m$ , with a decay rate $β_{1}$ .
$m_{k} = β_{1} m_{k - 1} + (1 - β_{1}) g_{k}$
where $g_{k}$ is the gradient at step $k$ . This is identical to the velocity term we saw in the Momentum optimizer.
Second Moment (RMSProp): It tracks the "variance" of the gradients. Let's denote this vector as $v$ , with a decay rate $β_{2}$ .
$v_{k} = β_{2} v_{k - 1} + (1 - β_{2}) g_{k}^{2}$
This is identical to the squared gradient accumulator in RMSProp. The squaring is done element-wise.

The following video provides a great high-level overview of how Adam brings these two concepts together.

Adam Optimizer from scratch | Gradient descent made better | Foundations for ML [Lecture 26]

The video "Adam Optimizer from scratch" by Vizuara introduces Adam as the synthesis of Momentum and RMSProp.

Please watch from this introduction. The presenter recaps the update rules for Momentum and RMSProp and then sets the stage for how Adam will combine an adaptive learning rate with a momentum term.

A naive combination would involve simply using the momentum term $m_{k}$ as the update direction, scaled by the adaptive learning rate from the RMSProp term $v_{k}$ . This would look like:

w_{k + 1} = w_{k} - η \frac{m _{k}}{v _{k} + ϵ} (Naive Update)

This is very close to the final algorithm, but it suffers from a subtle problem, particularly in the initial stages of training.

My question to you: Both moment vectors, $m$ and $v$ , are initialized to zero. Given the update rules above and a typical $β_{1}$ of 0.9, what will the value of $m_{1}$ be in relation to the first gradient $g_{1}$ ? What does this imply about the magnitude of the first few parameter updates?

The Problem of Initialization Bias

Your analysis is correct. Since $m_{0} = 0$ , the first moment update is $m_{1} = (1 - β_{1}) g_{1}$ . If $β_{1} = 0.9$ , then $m_{1} = 0.1 g_{1}$ . Similarly, if $β_{2} = 0.999$ , the second moment update is $v_{1} = 0.001 g_{1}^{2}$ .

Both moment estimates are heavily skewed, or biased, towards their zero initialization. This causes the initial update steps to be undesirably small, slowing down the start of the learning process. The creators of Adam proposed a simple and elegant fix: bias correction.

The idea is to compute bias-corrected estimates, $m^_{k}$ and $v^_{k}$ , that account for this initialization bias. Let's derive the correction for the first moment.

By unrolling the recursion for $m_{k}$ starting from $m_{0} = 0$ , we get:

m_{k} = (1 - β_{1}) i = 1 \sum k β_{1}^{k - i} g_{i}

To analyze the bias, we can take the expectation of $m_{k}$ , assuming the gradients $g_{i}$ are drawn from a stationary distribution with mean $E [g]$ .

E [m_{k}] = E [(1 - β_{1}) i = 1 \sum k β_{1}^{k - i} g_{i}] = (1 - β_{1}) i = 1 \sum k β_{1}^{k - i} E [g_{i}]

Assuming $E [g_{i}] = E [g]$ for all $i$ , and using the formula for a finite geometric series, we get:

E [m_{k}] = E [g] (1 - β_{1}) \frac{1 - β _{1}^{k}}{1 - β _{1}} = E [g] (1 - β_{1}^{k})

This shows that the expected value of our moment estimate $m_{k}$ is not the true expected gradient $E [g]$ , but is scaled by a factor of $(1 - β_{1}^{k})$ . To counteract this, we can define a bias-corrected estimate $m^_{k}$ by dividing $m_{k}$ by this factor:

m^_{k} = \frac{m _{k}}{1 - β _{1}^{k}}

The exact same logic applies to the second moment $v_{k}$ :

v^_{k} = \frac{v _{k}}{1 - β _{2}^{k}}

At the beginning of training (e.g., $k = 1$ ), $1 - β_{1}^{k}$ is small, so this correction significantly boosts the moment estimates. As training proceeds and $k$ gets large, $β_{1}^{k}$ approaches zero, $1 - β_{1}^{k}$ approaches 1, and the correction factor vanishes.

The following video and text provide excellent intuition for why this correction is so important.

Adam Optimizer from scratch | Gradient descent made better | Foundations for ML [Lecture 26]

This segment from the Vizuara video explains the mathematical steps for bias correction.

Please watch from this section on bias correction. Pay close attention to how the division by (1-\beta^k) counteracts the initial bias towards zero, especially for the second moment \mathbf{v} where \beta_2 is very close to 1.

For a deeper textual analysis, the following resource provides an exceptionally clear explanation.

Why is it important to include a bias correction term ...

This StackExchange Q&A is one of the best explanations available for Adam's bias correction.

First, read the problem statement to solidify your understanding of why the bias is an issue, especially with sparse gradients or a small \beta_2. Next, read how the bias is corrected, which recaps the mathematical derivation we just performed. Finally, read the excellent intuitive summary under the heading "What should you understand from all this, intuitively?". This part crystallizes the core idea: it's easier to correct for a known bias towards zero than for an unknown bias towards a noisy initial gradient.

The Complete Adam Algorithm and Implementation

With bias correction in place, we can now state the full Adam algorithm.

Algorithm: Adam

Initialize parameters $w_{0}$ , first moment vector $m_{0} \leftarrow 0$ , second moment vector $v_{0} \leftarrow 0$ , and timestep $k \leftarrow 0$ .
Choose hyperparameters: learning rate $η$ , decay rates $β_{1}, β_{2}$ , and numerical stability constant $ϵ$ . Default values are often $η = 0.001, β_{1} = 0.9, β_{2} = 0.999, ϵ = 1 0^{- 8}$ .
while not converged do:
1. $k \leftarrow k + 1$
2. Compute gradient: $g_{k} \leftarrow \nabla L (w_{k - 1})$
3. Update biased first moment: $m_{k} \leftarrow β_{1} m_{k - 1} + (1 - β_{1}) g_{k}$
4. Update biased second moment: $v_{k} \leftarrow β_{2} v_{k - 1} + (1 - β_{2}) g_{k}^{2}$
5. Compute bias-corrected first moment: $m^_{k} \leftarrow \frac{m _{k}}{1 - β _{1}^{k}}$
6. Compute bias-corrected second moment: $v^_{k} \leftarrow \frac{v _{k}}{1 - β _{2}^{k}}$
7. Update parameters: $w_{k} \leftarrow w_{k - 1} - η \frac{m ^ _{k}}{v ^ _{k} + ϵ}$

This algorithm elegantly combines the adaptive step size of RMSProp with the accelerated direction of Momentum. Let's see how this translates into code.

Adam Optimizer from Scratch in Python

The video "Adam Optimizer from Scratch in Python" by Deep Learning with Yacine provides a clean, step-by-step implementation.

Watch the segment from the Python implementation. The presenter maps each step of the algorithm directly to NumPy operations. Note how the function must maintain and return the state of the moments (m and v) for the next iteration. The bias correction is applied just before the final parameter update and is not part of the state that gets passed along.

Your interest in systems programming and performance might lead you to consider a C++ implementation. The logic remains identical, but you would manage memory for the moment vectors directly. The resource "Gradient Descent With Adam in Plain C++" by DemoFox (ID LINK) provides a great example of this, encapsulating the Adam state within a struct for each parameter. This is an excellent exercise to undertake on your own to solidify your understanding.

Comparing Optimizer Performance

The theoretical advantages of Adam become strikingly clear when we visualize its performance compared to the other optimizers we've studied.

This graph plots the training loss versus the number of update steps for SGD, Momentum, RMSProp, and Adam on the same task. The y-axis represents the loss, and the x-axis represents the training step.

Let's analyze the trajectories:

SGD (blue): The loss decreases, but the path is extremely noisy and exhibits high variance. Its convergence is slow and erratic.
Momentum (orange): The path is significantly smoother than SGD's, thanks to the averaging effect of the velocity term. It converges faster and to a lower loss than plain SGD.
RMSProp (green): Shows very rapid initial convergence, as it quickly adapts the learning rates. Its performance is comparable to Adam's in the early stages.
Adam (red): This is the clear winner. It converges as fast as RMSProp initially and continues to a very stable, low loss value, effectively combining the rapid adaptation of RMSProp with the smooth, steady progress of Momentum.

Conclusion

In this lesson, we have reached the pinnacle of our study of foundational optimization algorithms. By synthesizing the two major advancements over SGD, we have derived Adam, one of the most successful and widely used optimizers in deep learning.

Key Takeaways:

Adaptive Moment Estimation (Adam): Adam combines the first moment estimation (like Momentum) with second moment estimation (like RMSProp) to get the best of both worlds.
Dual Moving Averages: It maintains two separate exponentially weighted moving averages: $m$ for the gradients and $v$ for the squared gradients.
Bias Correction: Adam's key innovation is correcting the inherent bias of the moment estimates towards their zero initialization. This is achieved by dividing the moments $m_{k}$ and $v_{k}$ by $(1 - β_{1}^{k})$ and $(1 - β_{2}^{k})$ respectively, a correction that is crucial early in training and fades over time.
The Update Rule: The final parameter update uses the corrected first moment $m^_{k}$ to determine the update's direction and magnitude, while using the corrected second moment $v^_{k}$ to adaptively scale the learning rate for each parameter.

Preview of the Next Lesson:

We have now assembled a complete toolkit for numerical optimization. We understand how to define a loss function and, more importantly, how to iteratively minimize it in a computationally efficient and stable manner. With these powerful tools in hand, we are ready to build and train our first true learning machine.

In the next module, we will shift our focus from optimization to architecture. We'll begin with the historical building block of neural networks, the perceptron, and see why its limitations led to the development of multi-layer networks. This will set the stage for deriving the backpropagation algorithm, the engine that allows us to compute the gradients that our optimizers need to function.