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 a variant of the Adam optimization algorithm, introduced by Kingma and Ba in 2014. It modifies the adaptive learning rate mechanism of Adam by replacing the second-moment estimate with the infinity norm of past gradients. This adjustment simplifies the optimization process and improves stability when working with sparse gradients or parameters with large variations.

The algorithm is designed to adaptively adjust the learning rates for each parameter based on the first-moment estimate and the infinity norm of the gradient updates. This is particularly effective in high-dimensional parameter spaces, where the algorithm avoids issues caused by over-reliance on second-moment estimates, as seen in the original Adam algorithm.

Adamax is well-suited for tasks involving sparse gradients and has been successfully applied in various fields, including natural language processing, computer vision, and reinforcement learning. Its robustness and computational efficiency make it a preferred choice for optimizing deep learning models.

Algorithm Discussion

The Adamax optimizer, a variant of the Adam optimizer, adapts the learning rate for each parameter based on the first moment estimate and the infinity norm of past gradients. This approach makes it particularly robust for handling sparse gradients and stable under certain training conditions. By replacing the second moment estimate with the infinity norm, Adamax simplifies the parameter update while retaining the core benefits of adaptive learning rates.

Given the parameters $\theta$, a learning rate $\alpha$, and decay rates $\beta _{1}$ and $\beta _{2}$, Adamax follows these steps:

Initialize

• Initialize parameters $\theta _{0}$, the first-moment estimate $m_{0}=0$, and the exponentially weighted infinity norm $u_{0}=0$.
• Set hyperparameters:

  $\alpha$: Learning rate
  $\beta _{1}$: Exponential decay rate for the first moment
  $\beta _{2}$: Exponential decay rate for the infinity norm
  $\epsilon$: Small constant to avoid division by zero

For each time step

• Compute Gradient: $g_{t}=\nabla _{\theta }J(\theta _{t-1})$

• Update First Moment Estimate: $m_{t}=\beta _{1}\cdot m_{t-1}+(1-\beta _{1})\cdot g_{t}$

• Update Infinity Norm: $u_{t}=\max(\beta _{2}\cdot u_{t-1},|g_{t}|)$

• Bias Correction for the First Moment: ${\hat {m}}_{t}={\frac {m_{t}}{1-\beta _{1}^{t}}}$

• Parameter Update: $\theta _{t}=\theta _{t-1}-\alpha \cdot {\frac {{\hat {m}}_{t}}{u_{t}+\epsilon }}$

Pseudocode for Adamax

For $t=1$ to $T$:

 $g_{t}=\nabla _{\theta }J(\theta _{t-1})$
 $m_{t}=\beta _{1}\cdot m_{t-1}+(1-\beta _{1})\cdot g_{t}$
 $u_{t}=\max(\beta _{2}\cdot u_{t-1},|g_{t}|)$
 ${\hat {m}}_{t}={\frac {m_{t}}{1-\beta _{1}^{t}}}$
 $\theta _{t}=\theta _{t-1}-\alpha \cdot {\frac {{\hat {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$

• 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

Adamax has been widely used in various machine learning and deep learning tasks due to its robustness in handling sparse gradients and its computational efficiency. Some key application areas include:

Natural Language Processing (NLP)

Adamax performs well in NLP tasks, such as training word embeddings, text classification, and language modeling. The ability to handle sparse gradients makes it particularly effective in models like BERT and GPT. Its adaptive learning rate mechanism is advantageous for tasks where vocabulary size leads to large parameter spaces.

Computer Vision

Adamax has been applied in image classification and object detection tasks using deep convolutional neural networks (CNNs). For instance, its stability and adaptive learning rate have been shown to improve the training of models like ResNet and EfficientNet.

Reinforcement Learning

Adamax is particularly useful in reinforcement learning tasks, where it optimizes policy and value networks. Its robustness ensures stable convergence even with noisy and sparse reward signals.

Generative Models

Adamax has been used in training generative adversarial networks (GANs) and variational autoencoders (VAEs). The optimizer helps stabilize the training process, which can be sensitive to gradient updates.

Time Series Prediction

In time series forecasting tasks, Adamax efficiently handles models with recurrent neural networks (RNNs) and transformers. It has been applied to tasks like financial prediction and sensor data analysis.

Adamax is preferred in scenarios requiring robust handling of large parameter spaces, sparse gradients, or noisy data. Its wide adoption across different domains highlights its versatility and effectiveness.

Conclusion

Adamax is a robust and computationally efficient optimization algorithm that builds upon the Adam framework by replacing the second-moment estimate with the infinity norm. This modification simplifies the optimization process and enhances stability, particularly in handling sparse gradients and high-dimensional parameter spaces.

The algorithm's versatility makes it suitable for various deep learning tasks, including natural language processing, computer vision, reinforcement learning, generative models, and time series forecasting. Its robustness in dealing with sparse gradients, coupled with its adaptive learning rate mechanism, has contributed to its adoption in many state-of-the-art machine learning frameworks, such as TensorFlow and PyTorch.

Adamax’s ability to balance simplicity and performance ensures its ongoing relevance in optimizing complex models across diverse applications.