Optimizers impact training success of diffusion models on tasks
Optimizers for Diffusion Models: A Controlled Benchmark
Machine Learning
Summary
Diffusion models are a kind of AI that can generate data, like images or text, by reversing a noisy process. The way these models are trained depends heavily on the choice of optimizer—a tool that helps the model learn. The authors tested seven different optimizers on four types of diffusion tasks to see which worked best. They found that the popular default, AdamW, is good but not always the best, and that some newer optimizers can outperform it depending on the task. This work helps guide which optimizer to choose to get better results from diffusion models.
What this means in practice
- •For machine learning engineers: Select optimizers tailored to specific diffusion model tasks to improve training efficiency and model accuracy.
- •For ai product developers: Use the benchmarked optimizers to enhance the quality of AI-generated images or text in commercial applications.$Commercial implications: This enables better-performing AI generation tools for products requiring high-quality synthetic media.
Authors
Arman Bolatov, Egor Shulgin, David Li, Abduragim Shtanchaev, Sebastian U. Stich, Maxim Panov, Eric Moulines, Peter Richtárik, Martin Takáč
Abstract
Discrete diffusion models now match autoregressive language models on several benchmarks, while the question of how best to train them has received far less attention: the optimizer is inherited from one paper to the next and never compared. New optimizers, meanwhile, are validated almost exclusively on autoregressive pretraining, a different objective on a different loss surface. We present a controlled optimizer benchmark across four diffusion formulations, to our knowledge the first for discrete diffusion: seven optimizers (AdamW, Lion, Muon, SOAP, MARS, MARS-M, Schedule-Free) on masked diffusion (text8), uniform diffusion (QM9, and LM1B through the Gaussian duality) and Gaussian diffusion on images (CelebA-64), each on a task with published reference values. Every optimizer receives the same search protocol, and every winner is retrained at the full budget with three seeds. AdamW is a strong default but not always the right choice: it is beaten by a resolved margin on two of the four tasks, and the winner changes with the formulation, so the optimizer deserves the same care as the rest of the training recipe. Notably, methods validated on autoregressive language model pretraining transfer well: Muon, MARS-M and SOAP each beat the tuned AdamW on at least one diffusion formulation. The benchmark, all runs and every figure are reproducible end to end from the released code at https://github.com/armanbolatov/diffusion-baselines.