AdamS: Momentum Itself Can Be A Normalizer for LLM Pretraining and Post-training

TL;DR

AdamS is a novel optimizer for large language models that uses a new normalization technique based on momentum, offering efficiency, simplicity, and improved performance over AdamW without requiring architectural changes.

Contribution

Introducing AdamS, an optimizer that replaces second-moment estimates with a momentum-based normalization, providing theoretical guarantees and practical benefits for LLM training.

Findings

AdamS matches SGD with momentum in efficiency.

It outperforms AdamW in LLM pretraining tasks.

It is easy to integrate into existing training pipelines.

Abstract

We introduce AdamS, a simple yet effective alternative to Adam for large language model (LLM) pretraining and post-training. By leveraging a novel denominator, i.e., the root of weighted sum of squares of the momentum and the current gradient, AdamS eliminates the need for second-moment estimates. Hence, AdamS is efficient, matching the memory and compute footprint of SGD with momentum while delivering superior optimization performance. Moreover, AdamS is easy to adopt: it can directly inherit hyperparameters of AdamW, and is entirely model-agnostic, integrating seamlessly into existing pipelines without modifications to optimizer APIs or architectures. The motivation behind AdamS stems from the observed $(L_0, L_1)$ smoothness properties in transformer objectives, where local smoothness is governed by gradient magnitudes that can be further approximated by momentum magnitudes. We establish rigorous theoretical convergence guarantees and provide practical guidelines for hyperparameter selection. Empirically, AdamS demonstrates strong performance in various tasks, including pre-training runs on GPT-2 and Llama2 (up to 13B parameters) and reinforcement learning in post-training regimes. With its efficiency, simplicity, and theoretical grounding, AdamS stands as a compelling alternative to existing optimizers.