Pre-Training LLMs on a budget: A comparison of three optimizers

TL;DR

This paper compares three optimizers—AdamW, Lion, and Sophia—for pre-training large language models, analyzing their efficiency and performance across different architectures and training strategies.

Contribution

It provides a systematic comparison of three major optimizers for LLM pre-training, highlighting their relative strengths in training speed and downstream performance.

Findings

Sophia had the lowest training and validation loss.

Lion was the fastest in GPU hours.

AdamW achieved the best downstream evaluation results.

Abstract

Optimizers play a decisive role in reducing pre-training times for LLMs and achieving better-performing models. In this study, we compare three major variants: the de-facto standard AdamW, the simpler Lion, developed through an evolutionary search, and the second-order optimizer Sophia. For better generalization, we train with two different base architectures and use a single- and a multiple-epoch approach while keeping the number of tokens constant. Using the Maximal Update Parametrization and smaller proxy models, we tune relevant hyperparameters separately for each combination of base architecture and optimizer. We found that while the results from all three optimizers were in approximately the same range, Sophia exhibited the lowest training and validation loss, Lion was fastest in terms of training GPU hours but AdamW led to the best downstream evaluation results.