Optimization Benchmark for Diffusion Models on Dynamical Systems

TL;DR

This work addresses the lack of optimization benchmarks for diffusion-model training by evaluating optimizers on a diffusion task that denoises dynamical-system trajectories derived from Navier–Stokes simulations. It compares AdamW, Muon, SOAP, ScheduleFree, Prodigy, and SGD, revealing that Muon and SOAP achieve roughly 18% lower final validation loss than AdamW, while ScheduleFree can lag in generative quality despite competitive loss. The study also investigates learning-rate schedules (wsd, cosine, sqrt) and shows that training dynamics and trajectory quality are strongly affected by scheduling choices, with Prodigy offering tuning-free performance nearing tuned AdamW. Overall, the results provide practical guidance for diffusion-model optimization in scientific contexts (e.g., data assimilation) and raise open questions about the SGD–AdamW gap and the impact of the full training trajectory on generative quality.

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.

Figures (13)

• Figure 1: (Left) Final validation loss (averaged over the last five epochs) for each method and learning rate. Enlarged dot marks best learning rate. (Right) Validation loss curve for the best found setup for each method. Legend indicates learning rate (lr) and weight decay (wd) values. To obtain smoother curves we plot exponential moving averages with coefficient 0.95. See also \ref{['fig:benchmark-2']}.
• Figure 2: (Left) Final validation loss (averaged over the last five epochs, with band of one standard deviation over three seeds). Horizontal line marks best final loss for Muon and SOAP after 1024 epochs. (Right) Validation loss curve for the best AdamW run over 1024 and 2048 epochs (smoothened by exponential moving averages with coefficient 0.95).
• Figure 3: (Left) Final validation loss (averaged over the last five epochs) across weight decay for Prodigy. Horizontal lines marks final loss for the three best learning-rates we have found for AdamW. (Right) Validation loss curve for the best Prodigy run (smoothened by exponential moving averages with coefficient 0.95).
• Figure 4: (A) Final validation loss (averaged over the last five epochs) for cosine, wsd and sqrt schedules across peak learning rate. Enlarged dot marks best learning rate. (B) Shape of the schedules, normalized by peak learning rate. For wsd and sqrt, we use linear cooldown over the final 20% of training. (C) Validation loss curve for the best found setup for each schedule (smoothened by exponential moving averages with coefficient 0.95). (D) Same as (C) for $\ell_2$-norm of the batch gradients.
• Figure 5: (Left) Final validation loss (averaged over the last five epochs) for each method and learning rate. (Right) Validation loss curve for the best found AdamW setup, and the best five SGD setups (smoothened by exponential moving averages with coefficient 0.95).