Adamax and Nadam Variants

The Adam optimizer provides a combination of momentum and adaptive scaling based on the second moment of gradients. Researchers have proposed several variants aiming for further improvements or addressing specific behaviors. Two notable variants are Adamax and Nadam. These methods modify Adam's core components, specifically how it scales the learning rate or incorporates momentum, offering alternative approaches for certain optimization scenarios.

Adamax: Adaptive Learning Rates with Infinity Norm

Adam scales the learning rate for each parameter based on a decaying average of the squared past gradients ($v_t$), which effectively uses the $L_2$ norm. Adamax, introduced in the same paper as Adam by Kingma and Ba, uses the $L_\infty$ norm (infinity norm, or maximum norm) for scaling instead.

The motivation stems from generalizing Adam's update rule. Adam's per-parameter update involves dividing the first moment estimate $\hat{m}_t$ by $\sqrt{\hat{v}_t}$. The term $\sqrt{v_t}$ is roughly proportional to the $L_2$ norm of the current and past gradients. What happens if we generalize this to an $L_p$ norm? The update component $v_t$ would involve the $p$-th power of gradients:

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

The parameter update would then involve division by $v_t^{1/p}$. Adam corresponds to $p=2$. Adamax considers the case where $p \to \infty$. As $p \to \infty$, the $L_p$ norm converges to the $L_\infty$ norm (the maximum absolute value). This leads to a remarkably simple update for the scaling term, often denoted as $u_t$ to distinguish it from Adam's $v_t$:

$$u_t = \max(\beta_2 u_{t-1}, |g_t|)$$

Notice that this update doesn't require element-wise squaring or square roots. It relies on the max operation. Because $u_t$ relies on a max operation, it can be more stable than Adam's $v_t$ when dealing with very large but sparse gradients, as the max is less affected by infrequent large values compared to the sum of squares.

The full Adamax parameter update rule becomes:

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

Here, $\hat{m}_t$ is the bias-corrected first moment estimate, calculated exactly as in Adam:

$$m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t$$ $$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}$$

Important Points about Adamax:

• Scaling: Uses the $L_\infty$ norm (maximum) for scaling, derived as the limit of the $L_p$ norm.
• Update Rule ($u_t$): Simpler calculation ($u_t = \max(\beta_2 u_{t-1}, |g_t|)$) compared to Adam's $v_t$.
• Stability: Potentially more stable than Adam, especially with sparse gradients or sudden large gradients, due to the max operation being less sensitive than summing squares.
• Hyperparameters: Uses the same hyperparameters as Adam ($\eta$, $\beta_1$, $\beta_2$), with similar typical default values ($\beta_1 \approx 0.9$, $\beta_2 \approx 0.999$).

Adamax provides a computationally slightly simpler variant of Adam that might offer better stability in certain situations.

Nadam: Incorporating Nesterov Momentum into Adam

Nesterov Accelerated Gradient (NAG) is known to often improve the convergence speed of standard momentum methods by calculating the gradient after applying a preliminary momentum step (looking ahead). The Nadam (Nesterov-accelerated Adaptive Moment Estimation) optimizer, proposed by Dozat, aims to integrate this Nesterov momentum principle into the Adam framework.

Recall that Adam updates parameters using bias-corrected estimates of the first moment ($m_t$) and second moment ($v_t$) of the gradients. Nadam modifies how the first moment estimate influences the update step to incorporate the "lookahead" aspect of NAG.

The core idea in NAG is to compute the gradient at $\theta_t + \mu m_{t-1}$ (where $\mu$ is the momentum coefficient) instead of at $\theta_t$. Nadam applies a similar correction within Adam's update. Instead of using only the previous momentum term $\beta_1 \hat{m}_{t-1}$ to anticipate the next position, Nadam effectively incorporates the current gradient information more directly into the momentum part of the update.

Let's look at the update steps for Nadam:

1. Compute gradient: $g_t = \nabla f(\theta_{t-1})$
2. Update biased first moment estimate: $m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t$
3. Update biased second moment estimate: $v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2$
4. Compute bias-corrected first moment: $\hat{m}_t = \frac{m_t}{1 - \beta_1^t}$
5. Compute bias-corrected second moment: $\hat{v}_t = \frac{v_t}{1 - \beta_2^t}$
6. Compute the Nadam parameter update: $$\theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \left( \beta_1 \hat{m}_t + \frac{(1 - \beta_1) g_t}{1 - \beta_1^t} \right)$$

Compare the Nadam update (Step 6) to the standard Adam update:

$$\text{Adam Update: } \theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t = \theta_{t-1} - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \left( \frac{\beta_1 m_{t-1}}{1 - \beta_1^t} + \frac{(1 - \beta_1) g_t}{1 - \beta_1^t} \right)$$

The main difference lies in how the momentum term is applied. Nadam uses $\beta_1 \hat{m}_t$ which includes the influence of the current gradient $g_t$ via $m_t$, combined with another term involving the current gradient $\frac{(1 - \beta_1) g_t}{1 - \beta_1^t}$. This structure effectively applies the momentum step using the updated momentum $m_t$ and incorporates the gradient correction in a way that mimics NAG's lookahead behavior.

Important Points about Nadam:

• Momentum: Integrates Nesterov momentum with Adam's adaptive learning rates.
• Update Rule: Modifies the Adam update to use a Nesterov-style momentum calculation.
• Convergence: Often converges faster than Adam, especially on problems where NAG outperforms standard momentum. It can be particularly effective for training complex models like recurrent neural networks or deep CNNs.
• Hyperparameters: Uses the same hyperparameters as Adam ($\eta$, $\beta_1$, $\beta_2$), often with similar default values.

Nadam is frequently a strong choice when faster convergence is desired and the overhead of the slightly more complex update calculation is acceptable.

Practical Notes

Both Adamax and Nadam are readily available in popular machine learning libraries like TensorFlow (tf.keras.optimizers.Adamax, tf.keras.optimizers.Nadam) and PyTorch (torch.optim.Adamax, torch.optim.Nadam).

• When to Choose: While Adam remains a very popular and often effective default, consider trying Nadam if faster convergence is a priority. Adamax might be worth exploring if you observe instability possibly related to large or sparse gradients when using Adam.
• Hyperparameter Tuning: Although defaults often work well, tuning the learning rate $\eta$ and momentum decay parameters $\beta_1, \beta_2$ can still yield performance improvements, just as with Adam. $\beta_1$ typically influences the momentum aspect, while $\beta_2$ controls the adaptivity based on past gradient magnitudes.
• Empirical Evaluation: As with all optimization algorithms, the best choice often depends on the specific problem, dataset, and model architecture. It's recommended to empirically compare the performance of Adam, Adamax, and Nadam (along with potentially other optimizers like RMSprop or SGD with momentum) on your specific task.

By understanding the subtle but significant modifications introduced by Adamax and Nadam, you gain more tools for navigating optimization, potentially achieving more stable or faster training for your machine learning models. They represent valuable refinements built upon the foundational Adam algorithm.