Connections between Schedule-Free Optimizers, AdEMAMix, and Accelerated SGD Variants

TL;DR

This paper links recent schedule-free optimizers with theoretical accelerated SGD, demonstrating that AdEMAMix performs best and introducing a simplified version that maintains performance with less complexity.

Contribution

It establishes explicit theoretical connections between schedule-free optimizers and accelerated SGD, and proposes a simplified AdEMAMix variant with comparable performance.

Findings

AdEMAMix outperforms other optimizers in preliminary experiments.

Simplified-AdEMAMix matches AdEMAMix performance across batch sizes.

Theoretical links between optimizer classes are explicitly demonstrated.

Abstract

Recent advancements in deep learning optimization have introduced new algorithms, such as Schedule-Free optimizers, AdEMAMix, MARS and Lion which modify traditional momentum mechanisms. In a separate line of work, theoretical acceleration of stochastic gradient descent (SGD) in noise-dominated regime has been achieved by decoupling the momentum coefficient from the current gradient's weight. In this paper, we establish explicit connections between these two lines of work. We substantiate our theoretical findings with preliminary experiments on a 150m language modeling task. We find that AdEMAMix, which most closely resembles accelerated versions of stochastic gradient descent, exhibits superior performance. Building on these insights, we introduce a modification to AdEMAMix, termed Simplified-AdEMAMix, which maintains the same performance as AdEMAMix across both large and small batch-size settings while eliminating the need for two different momentum terms. The code for Simplified-AdEMAMix is available on the repository: https://github.com/DepenM/Simplified-AdEMAMix/.