Deep Neural Networks (DNNs) are regarded as one of the most effective tools for finding patterns in large datasets through training. At the core of the training problems, we have complex loss landscapes and the training of a DNN boils down to optimizing the loss as the number of iterations increases. A few of the most commonly used optimizers are Stochastic Gradient Descent, RMSProp (Root Mean Square Propagation), Adam (Adaptive Moment Estimation) etc.
Recently (September 2024), researchers from Apple (and EPFL) proposed a new optimizer, AdEMAMix¹, which they show to work better and faster than AdamW optimizer for language modeling and image classification tasks.
In this post, I will go into detail about the mathematical concepts behind this optimizer and discuss some very interesting results presented in this paper. Topics that will be covered in this post are:
- Review of Adam Optimizer
- Exponential Moving Average (EMA) in Adam.
- The Main Idea Behind AdEMAMix: Mixture of two EMAs.
- The Exponential Decay Rate Scheduler in AdEMAMix.