AMSGrad: Improving Adam's Convergence

While the Adam optimizer, discussed previously, has become a default choice for many deep learning tasks due to its efficiency and effectiveness, further analysis revealed potential issues with its convergence proof. Specifically, the use of the exponentially decaying average of past squared gradients ($v_t$) doesn't guarantee that this term remains non-decreasing. In some scenarios, particularly late in training or with very informative gradients appearing after less informative ones, $v_t$ could decrease significantly. This decrease, when used in the denominator of the parameter update step ($\frac{\eta}{\sqrt{\hat{v}_t} + \epsilon}$), could lead to undesirably large learning steps, potentially causing the optimizer to diverge or fail to converge to an optimal solution.

AMSGrad (Adaptive Moment Estimation with improved convergence guarantees) was proposed to directly address this theoretical flaw in Adam's convergence analysis. The core idea is simple yet effective: ensure that the adaptive learning rate term used for scaling gradients does not increase over time.

The AMSGrad Modification

Instead of directly using the current estimate of the second moment ($v_t$), AMSGrad maintains the maximum value of this estimate seen so far. Let's break down the update steps:

1. Compute Gradient: Calculate the gradient $g_t$ of the loss function with respect to the parameters $\theta_t$ at timestep $t$.

   $$g_t = \nabla_{\theta} L(\theta_t)$$

2. Update Biased First Moment Estimate: This is the same as in Adam, using the hyperparameter $\beta_1$.

   $$m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t$$

3. Update Biased Second Moment Estimate: This is also the same as in Adam, using the hyperparameter $\beta_2$.

   $$v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2$$

   (Note: $g_t^2$ denotes element-wise squaring).

4. Maintain Maximum of Second Moment Estimates: This is the main difference introduced by AMSGrad. A new variable, $\hat{v}_t$, keeps track of the maximum value of $v$ encountered up to timestep $t$.

   $$\hat{v}_t = \max(\hat{v}_{t-1}, v_t)$$

   Initialize $\hat{v}_0 = 0$. The $\max$ operation is performed element-wise.

5. Compute Bias-Corrected First Moment Estimate: Same as Adam.

   $$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}$$

6. Update Parameters: Use the maximum second moment estimate $\hat{v}_t$ in the denominator, instead of the bias-corrected $\hat{v}_t$ used in Adam's update. Notice that $\hat{v}_t$ itself is not bias-corrected in the standard AMSGrad formulation used for the update step.

   $$\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t$$

   Here, $\eta$ is the step size (learning rate) and $\epsilon$ is a small constant for numerical stability.

By ensuring that the denominator term $\sqrt{\hat{v}_t} + \epsilon$ is non-decreasing, AMSGrad prevents the large, potentially destabilizing steps that could occur in Adam if $v_t$ were to shrink unexpectedly. This modification provides stronger theoretical convergence guarantees, particularly in online and non-convex settings.

Practical Implications and Usage

Does the theoretical improvement translate to better practical performance? The results are mixed and often problem-dependent.

• Stability: AMSGrad can sometimes offer more stable training and convergence, especially on datasets or architectures where Adam might exhibit instabilities or convergence issues.
• Performance: In many standard deep learning applications (e.g., image classification), standard Adam often converges faster and achieves comparable or even slightly better final performance than AMSGrad, despite the theoretical concerns. The conditions under which Adam's convergence fails might not be frequently encountered in these common benchmarks.
• Memory: AMSGrad requires slightly more memory as it needs to store the additional $\hat{v}_t$ vector.

When to Consider AMSGrad:

1. You observe convergence problems or instability with Adam on your specific task.
2. You are working in a setting (perhaps reinforcement learning or online learning) where the long-term behavior and convergence guarantees are more significant.
3. You prioritize robustness and theoretical soundness over potentially faster convergence in simpler cases.

Most deep learning frameworks (like TensorFlow and PyTorch) offer AMSGrad as a simple boolean flag within their Adam optimizer implementation (e.g., tf.keras.optimizers.Adam(amsgrad=True) or torch.optim.Adam(amsgrad=True)). This makes experimentation straightforward.

While Adam remains a very strong baseline and often the first choice, AMSGrad provides a valuable alternative, grounded in addressing a specific theoretical shortcoming of Adam. Knowing its mechanism and motivation allows you to make informed decisions when tuning optimization strategies for complex models.