Optimization Benchmark for Diffusion Models on Dynamical Systems

TL;DR

This paper benchmarks various optimization algorithms for training diffusion models on dynamical systems, revealing that Muon and SOAP outperform AdamW and analyzing training dynamics and optimizer performance.

Contribution

It introduces a comprehensive benchmark for diffusion model optimization algorithms and compares their efficiency, highlighting Muon and SOAP as superior alternatives to AdamW.

Findings

Muon and SOAP are 18% more efficient than AdamW.

Learning-rate schedules significantly affect training dynamics.

There is a notable performance gap between Adam and SGD in this context.

Abstract

The training of diffusion models is often absent in the evaluation of new optimization techniques. In this work, we benchmark recent optimization algorithms for training a diffusion model for denoising flow trajectories. We observe that Muon and SOAP are highly efficient alternatives to AdamW (18% lower final loss). We also revisit several recent phenomena related to the training of models for text or image applications in the context of diffusion model training. This includes the impact of the learning-rate schedule on the training dynamics, and the performance gap between Adam and SGD.