Adamax

Author: Chengcong Xu (cx253), Jessica Liu (hl2482), Xiaolin Bu (xb58), Qiaoyue Ye (qy252), Haoru Feng (hf352) (ChemE 6800 Fall 2024)

Stewards: Nathan Preuss, Wei-Han Chen, Tianqi Xiao, Guoqing Hu

Introduction

Adamax is an optimization algorithm introduced by Kingma and Ba in their Adam optimizer paper (2014). It improves upon the Adam algorithm by replacing the second moment's root mean square (RMS) norm with the infinity norm ($\ell _{\infty }$). This change makes Adamax more robust and numerically stable, especially when handling sparse gradients, noisy updates, or optimization problems with significant gradient variations.

Adamax dynamically adjusts learning rates for individual parameters, making it well-suited for training deep neural networks, large-scale machine learning models, and tasks involving high-dimensional parameter spaces.

Algorithm Discussion

The Adamax optimization algorithm follows these steps:

• Step 1: Initialize Parameters

Set the learning rate $\alpha$, exponential decay rates $\beta _{1}$ and $\beta _{2}$, and numerical stability constant $\epsilon$. Initialize the first moment estimate $m=0$ and infinity norm estimate $u=0$.

• Step 2: Gradient Computation

Compute the gradient of the loss function with respect to the model parameters, $g_{t}$.

• Step 3: Update First Moment Estimate

$m_{t}=\beta _{1}m_{t-1}+(1-\beta _{1})g_{t}$

• Step 4: Update Infinity Norm Estimate

$u_{t}=\max(\beta _{2}u_{t-1},|g_{t}|)$

• Step 5: Bias-Corrected Learning Rate

${\hat {\alpha }}={\frac {\alpha }{1-\beta _{1}^{t}}}$

• Step 6: Parameter Update

$\theta _{t}=\theta _{t-1}-{\frac {{\hat {\alpha }}\cdot m_{t}}{u_{t}+\epsilon }}$

Numerical Examples

To illustrate the Adamax optimization algorithm, we will minimize the quadratic function $f(x)=x^{2}$ with step-by-step calculations.

Problem Setup

• Optimization Objective: Minimize $f(x)=x^{2}$, which reaches its minimum at $x=0$ with $f(x)=0$.

• Initial Parameter: Start with $x_{0}=2.0$.

• Gradient Formula: $g_{t}={\frac {\partial f}{\partial x}}=2x_{t}$, which determines the direction and rate of parameter change.

• Hyperparameters:

- Learning Rate: $\alpha =0.1$ controls the step size. - First Moment Decay Rate: $\beta _{1}=0.9$, determines how past gradients influence the current gradient estimate. - Infinity Norm Decay Rate: $\beta _{2}=0.999$, governs the decay of the infinity norm used for scaling updates. - Numerical Stability Constant: $\epsilon =10^{-8}$, prevents division by zero.

• Initialization:

$m_{0}=0,u_{0}=0,t=0$

Step-by-Step Calculations

Iteration 1 $t=1$

• Gradient Calculation

$g_{1}=2x_{0}=2\cdot 2.0=4.0$

The gradient indicates the steepest direction and magnitude for reducing $f(x)$. A positive gradient shows $x_{0}$ must decrease to minimize the function.

• First Moment Update

$m_{1}=\beta _{1}m_{0}+(1-\beta _{1})g_{1}=0.9\cdot 0+0.1\cdot 4.0=0.4$

The first moment $m_{1}$ is a running average of past gradients, smoothing out fluctuations.

• Infinity Norm Update

$u_{1}=\max(\beta _{2}u_{0},|g_{1}|)=\max(0.999\cdot 0,4.0)=4.0$

The infinity norm $u_{1}$ ensures updates are scaled by the largest observed gradient, stabilizing step sizes.

• Bias-Corrected Learning Rate

${\hat {\alpha }}={\frac {\alpha }{1-\beta _{1}^{t}}}={\frac {0.1}{1-0.9^{1}}}=1.0$

The learning rate is corrected for bias introduced by initialization, ensuring effective parameter updates.

• Parameter Update

$x_{1}=x_{0}-{\frac {{\hat {\alpha }}\cdot m_{1}}{u_{1}+\epsilon }}=2.0-{\frac {1.0\cdot 0.4}{4.0+10^{-8}}}=1.9$

The parameter moves closer to the function's minimum at $x=0$.

Iteration 2 $t=2$

• Time Step Update

$t=2$

• Gradient Calculation

$g_{2}=2x_{1}=2\cdot 1.9=3.8$

• First Moment Update

$m_{2}=\beta _{1}m_{1}+(1-\beta _{1})g_{2}=0.9\cdot 0.4+0.1\cdot 3.8=0.758$

• Infinity Norm Update

$u_{2}=\max(\beta _{2}u_{1},|g_{2}|)=\max(0.999\cdot 4.0,3.8)=4.0$

• Bias-Corrected Learning Rate

${\hat {\alpha }}={\frac {\alpha }{1-\beta _{1}^{t}}}={\frac {0.1}{1-0.9^{2}}}=0.526$

• Parameter Update

$x_{2}=x_{1}-{\frac {{\hat {\alpha }}\cdot m_{2}}{u_{2}+\epsilon }}=1.9-{\frac {0.526\cdot 0.758}{4.0+10^{-8}}}=1.802$

The parameter continues to approach the minimum at $x=0$.

Summary

Through these two iterations, Adamax effectively adjusts the parameter $x$ based on the computed gradients, moving it closer to the minimum. The use of the infinity norm stabilizes the updates, ensuring smooth convergence.

Applications

Natural Language Processing

Adamax is particularly effective in training transformer-based models like BERT and GPT. Its stability with sparse gradients makes it ideal for tasks such as text classification, machine translation, and named entity recognition.

Computer Vision

In computer vision, Adamax optimizes deep CNNs for tasks like image classification and object detection. Its smooth convergence behavior has been observed to enhance performance in models like ResNet and DenseNet.

Reinforcement Learning

Adamax has been applied in training reinforcement learning agents, particularly in environments where gradient updates are inconsistent or noisy, such as robotic control and policy optimization.

Generative Models

For training generative models, including GANs and VAEs, Adamax provides robust optimization, improving stability and output quality during adversarial training.

Time-Series Forecasting

Adamax is used in financial and economic forecasting, where it handles noisy gradients effectively, resulting in stable and accurate time-series predictions.

Advantages over Other Approaches

• Stability: The use of the infinity norm ensures Adamax handles gradient variations smoothly.

• Sparse Gradient Handling: Adamax is robust in scenarios with zero or near-zero gradients, common in NLP tasks.

• Efficiency: Adamax is computationally efficient for high-dimensional optimization problems.

Conclusion

Adamax is a robust and efficient variant of the Adam optimizer that replaces the RMS norm with the infinity norm. Its ability to handle sparse gradients, noisy updates, and large parameter spaces makes it a widely used optimization method in natural language processing, computer vision, reinforcement learning, and generative modeling.

Future advancements may involve integrating Adamax with learning rate schedules and regularization techniques to further enhance its performance.

References

• Kingma, D. P., & Ba, J. (2014). Adam: A Method for Stochastic Optimization. arXiv preprint arXiv:1412.6980.
• TensorFlow Documentation. Adamax Optimizer.
• PyTorch Documentation. Adamax Optimizer.