OmniOpt: Taxonomy, Geometry, and Benchmarking of Modern Optimizers
A two-stage benchmark comparing more than twenty optimizers for large language model pre-training at scale, first screening them broadly on a shared architecture and then testing the strongest ones for generalization across several model architectures.
OmniOpt's first stage sweeps over 24 optimizers, including AdamW, Muon, Shampoo, SOAP, Lion and several MARS variants, on the C4 dataset under a LLaMA-3 architecture at four scales from 60 million to 1 billion parameters, using final validation perplexity to compare them. A strict controlled-variable protocol tunes only each optimizer's own hyperparameters, holding architecture, data and schedule fixed across all runs.
The second stage carries the stronger stage-one optimizers forward to FineWeb-Edu with longer 32k-token sequences, at 340 million and 1 billion parameters, across four architectures: Transformer++, GLA, DeltaNet and Gated DeltaNet, to test how well stage-one results generalize.
It runs on PyTorch with Transformers and Datasets, with the second stage additionally requiring Flash Linear Attention for its non-Transformer architectures, and supports a Hugging Face mirror for use in mainland China.