Adam: Combining Momentum and RMSprop

Optimization algorithms for machine learning often adapt learning rates based on gradient moments. Some approaches, such as RMSprop, adjust learning rates using only the second moments of gradients. Other approaches, such as Momentum, use first moments to build velocity. This leads to a natural question: can the benefits of both be combined? Can an algorithm use momentum to accelerate convergence along consistent gradient directions while also adapting the step size for each parameter based on the historical magnitude of its gradients? The Adam optimizer provides an affirmative answer and has become one of the most widely used optimization algorithms in deep learning.

Adam, which stands for Adaptive Moment Estimation, computes adaptive learning rates for each parameter by keeping track of exponentially decaying averages of both past gradients (the first moment) and past squared gradients (the second moment).

The Adam Algorithm

Let's break down the mechanics of Adam. At each timestep $t$, with parameters $\theta_t$ and loss function $J(\theta)$, Adam performs the following steps:

1. Compute Gradient: Calculate the gradient of the loss function with respect to the current parameters: $g_t = \nabla_{\theta} J(\theta_t)$

2. Update Biased First Moment Estimate: Update the moving average of the gradient (the first moment estimate). This is similar to the momentum update but uses an exponential moving average controlled by the hyperparameter $\beta_1$. $m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t$ Here, $m_t$ is the estimate of the mean of the gradients. Common values for $\beta_1$ are typically close to 1, such as 0.9.

3. Update Biased Second Moment Estimate: Update the moving average of the squared gradients (the second raw moment estimate). This resembles the update in RMSprop, using an exponential moving average controlled by the hyperparameter $\beta_2$. $v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2$ Here, $v_t$ estimates the uncentered variance of the gradients (element-wise square is used for $g_t^2$). $\beta_2$ is also usually close to 1, for example, 0.999.

4. Compute Bias-Corrected First Moment Estimate: The moment estimates $m_t$ and $v_t$ are initialized as vectors of zeros. This initialization introduces a bias towards zero, especially during the initial timesteps. Adam corrects for this bias by computing bias-corrected estimates: $\hat{m}_t = \frac{m_t}{1 - \beta_1^t}$ As $t$ increases, the bias correction term $1 - \beta_1^t$ approaches 1, making the correction less significant.

5. Compute Bias-Corrected Second Moment Estimate: Similarly, correct the bias in the second moment estimate: $\hat{v}_t = \frac{v_t}{1 - \beta_2^t}$ The term $1 - \beta_2^t$ corrects the bias for the second moment estimate.

6. Update Parameters: Finally, update the parameters. The update rule uses the bias-corrected first moment estimate $\hat{m}_t$ (similar to momentum) and divides it by the square root of the bias-corrected second moment estimate $\hat{v}_t$ (similar to RMSprop), effectively providing a per-parameter learning rate scaling. $\theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$ Here, $\eta$ is the learning rate (step size), and $\epsilon$ is a small constant (e.g., $10^{-8}$) added for numerical stability to prevent division by zero in cases where $\hat{v}_t$ is very close to zero.

Hyperparameters

Adam introduces a few hyperparameters:

• $\eta$: The learning rate or step size. Controls the overall magnitude of the parameter updates. Values like 0.001 are often a good starting point.
• $\beta_1$: The exponential decay rate for the first moment estimates. Recommended default is 0.9.
• $\beta_2$: The exponential decay rate for the second moment estimates. Recommended default is 0.999.
• $\epsilon$: A small value to prevent division by zero. Recommended default is $10^{-8}$.

One of the appealing aspects of Adam is that its default hyperparameter values often work well across a range of problems, reducing the need for extensive tuning compared to some other optimizers.

Intuition and Advantages

Adam effectively combines the benefits observed in RMSprop and Momentum:

• Adaptive Learning Rates: Like RMSprop, the division by $\sqrt{\hat{v}_t}$ provides per-parameter learning rates. Parameters receiving large or frequent gradients will have smaller effective learning rates, while parameters with small or infrequent gradients will have larger effective learning rates. This helps particularly in sparse gradient scenarios.
• Momentum: The use of the first moment estimate $\hat{m}_t$ incorporates momentum, helping the optimizer accelerate in directions of low curvature and dampen oscillations in directions of high curvature, leading to faster convergence.
• Bias Correction: The bias correction steps are particularly important early in training when the moment estimates are still heavily influenced by their zero initialization. Without correction, the initial steps might be much smaller than intended.

Adam is computationally efficient, requires little memory overhead (it only needs to store the first and second moment vectors, which are the same size as the parameters), and is generally well-suited for problems with large datasets and high-dimensional parameter spaces, which are common in deep learning.

Relationship to Other Methods

It's helpful to see how Adam relates to algorithms we've already discussed:

• If we set $\beta_1 = 0$, the first moment estimate $m_t$ only depends on the current gradient $g_t$ (scaled by $1-\beta_1$), removing the momentum aspect.
• The second moment update and scaling are very similar to RMSprop, differing mainly in the specific form of the exponential moving average and the inclusion of bias correction.

Adam can be viewed as a sophisticated blend of RMSprop (handling adaptive scaling based on second moments) and momentum (using first moments to guide direction), with an added bias correction mechanism to improve stability during the initial phases of optimization. Its widespread adoption stems from its strong empirical performance across various machine learning tasks, particularly in training deep neural networks. While later sections will discuss variations like AMSGrad that address specific potential convergence issues with Adam, the original Adam algorithm remains a foundational and highly effective optimization technique.