Optimizers for Diffusion Models: A Controlled Benchmark
Abstract Overview
This paper presents a controlled benchmark of seven optimizers for diffusion training across four formulations: masked diffusion on text8, uniform diffusion on QM9, DUO on LM1B, and Gaussian image diffusion on CelebA-64. All optimizers are evaluated under a shared tuning protocol, and selected winners are retrained at full budget with three seeds, making the comparison more controlled than inheriting default settings from prior work. The study finds that optimizer performance depends strongly on the diffusion formulation rather than transferring consistently across tasks. It also shows that objective-based metrics and sample-quality metrics can disagree, especially on QM9 and CelebA-64.
Novelty
The work appears to be the first controlled optimizer benchmark focused on discrete diffusion, while also including a Gaussian diffusion image setting for comparison. Its distinctive contribution is the combination of a uniform search protocol, full-budget retraining of winners across three seeds, and a cross-formulation analysis of whether optimizer rankings transfer.
Results
A tuned AdamW is a strong baseline: no alternative clearly beats it on text8 or CelebA-64, while SOAP, Lion, and Muon achieve small but resolved gains over it on QM9. On LM1B, Muon and MARS-M outperform AdamW by a larger margin, MARS also finishes ahead, and Lion diverges on two of three seeds after the curriculum switch. Across all four tasks, Schedule-Free finishes last, and optimizer rankings vary substantially by formulation.
Key Points
- The benchmark compares seven optimizers across four diffusion formulations under the same tuning and retraining protocol, with released code for end-to-end reproducibility.
- AdamW remains a strong default, but Muon, MARS-M, and SOAP each outperform tuned AdamW on at least one formulation, showing that optimizer choice matters in diffusion training.
- Optimizer rankings do not transfer reliably across formulations, and lower training loss or likelihood-style metrics do not always correspond to better sample-quality metrics such as FID or FCD.