Source-linked AI summary
Optimization Benchmark for Diffusion Models on Dynamical Systems
Fabian Schaipp
TL;DR
Optimization benchmarks have largely left diffusion-model training untested, so transfer of newer methods beyond language and image tasks remains uncertain. This paper benchmarks optimizers on diffusion-based denoising of fluid-dynamics trajectories, finding that Muon and SOAP reach 18% lower final loss than AdamW. It also reports a loss–quality mismatch for some schedules and an AdamW–SGD gap that class imbalance does not explain in this setting.
Problem
Optimization benchmarking has largely omitted diffusion models, leaving uncertain whether methods developed for language or image applications transfer to diffusion training.
Source: §1 · p. 1Method
Results
18% lower final loss than AdamW is achieved by Muon and SOAP over 1024 epochs, while the study also finds an AdamW–SGD validation-loss gap.
Source: §3.1 · p. 4, §3.3 · p. 8Takeaways & Limitations
Muon and SOAP are convincing alternatives to AdamW for this diffusion-training task, while generative quality may depend on the training trajectory as well as final loss.
Source: §4 · p. 8The benchmark is limited to relatively small-scale problems of about 23M parameters, a constraint imposed to permit multiple seeds and hyperparameter configurations.
Source: §1 · p. 1
Abstract
from arXiv · showhide
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.
1 Introduction
This work tests whether optimization findings from image and language tasks transfer to diffusion-model training, an area largely absent from optimization benchmarks. On a fluid-dynamics trajectory task, it compares recent optimizers and revisits the Adam–SGD performance gap.
1 Introduction
- Optimization benchmarks largely omit diffusion models, leaving unclear whether newer methods such as SOAP and Muon transfer beyond LLM pretraining.The cited extensive benchmarking efforts do not report diffusion-model results.Source: §1 · p. 1
- The benchmark trains a diffusion model to denoise dynamical-system trajectories from fluid simulations and compares SOAP, Muon, and ScheduleFree with AdamW.The task differs from LLM pretraining in architecture, loss, data domain, and training regime.Source: §1 · pp. 1–2
- Computational constraints limited experiments to relatively small-scale problems of about 23M parameters so the authors could run multiple seeds and hyperparameter configurations.Source: §1 · p. 1
- The benchmark findings may be relevant to diffusion-based data assimilation for regional and global weather and climate simulation.The benchmark problem has previously been applied to those scientific applications.Source: §1 · p. 2
Performance gap between Adam and SGD
- The study tests whether an AdamW–SGD performance gap persists for diffusion training without class labels, while prior work reports SGD closing the gap under some language-task settings.The cited conditions include very small batch sizes or applying Adam only to embedding-layer weights.Source: §1 · p. 2
Summary and main findings
- Muon and SOAP achieve lower final loss than AdamW despite higher per-step runtime; ScheduleFree and WSD can attain good loss while generative quality degrades.The authors conjecture that the full training trajectory, not only final loss, may matter for model quality.Source: §1 · p. 2
2 Experimental Setup
The setup trains a U-Net score model on dynamical-system trajectories using a standard DDPM denoising approach. The authors tune optimizer settings and average metrics across seeds under a common training schedule.
2 Experimental Setup
Hyperparameter tuning
- For each optimizer, learning rate and weight decay are tuned separately; runs typically use three seeds, 1024 epochs, and a linear-decay schedule.Warmup and gradient clipping are added by default and produce a minute loss reduction.Source: §2 · p. 3
3 Results
The benchmark finds that Muon and SOAP outperform AdamW on the denoising task, while optimizer and schedule choices also affect training dynamics and generated-trajectory quality. The schedule comparisons and AdamW–SGD results show that lower loss and better generative quality do not always align.
3.1 Main Benchmark
- Over 1024 epochs (26.6K steps), Muon and SOAP achieve 18% lower final loss than AdamW.SOAP performs best by steps, closely followed by Muon; runtime comparisons favor Muon, while SOAP converges as fast as AdamW and reaches lower final loss.Source: Figure 1 · p. 4, §3.1 · p. 4
- AdamW trained for 2048 epochs still does not match Muon or SOAP’s 1024-epoch final loss, even with the same or lower runtime budget.The runtime advantage of Muon and SOAP over AdamW is reduced when runtime, rather than steps, is compared.Source: §3.1 · p. 4, Figure 2 · p. 5
- ScheduleFree’s similar loss values do not ensure similar generated-trajectory quality; adding a linear cooldown improves quality for at least some configurations.The authors conjecture that missing learning-rate annealing partially explains the mismatch.Source: §3.1 · p. 5
- Prodigy roughly matches AdamW’s second-best learning-rate result without learning-rate tuning, reducing tuning effort for preliminary runs with similar or slightly inferior model quality.Its adaptive learning rate reaches 6 · 10^-4 within a few epochs, reasonably close to AdamW’s best tuned learning rate.Source: §3.1 · pp. 5–6, Figure 3 · p. 6
3.2 Impact of Learning-Rate Schedule
- WSD matches or surpasses cosine in final loss, but its best-loss learning rate yields lower-quality trajectories; sqrt has worse final loss than both but more stable trajectory quality.The sqrt schedule is presented as an alternative when training length cannot be specified in advance.Source: §3.2 · pp. 6–7, Figure 4 · p. 7
3.3 Gap Between AdamW and SGD
- SGD has significantly higher validation loss than AdamW, and its generated trajectories are clearly inferior despite extensive hyperparameter tuning.The authors conclude that factors such as architecture properties must explain the gap in this problem instance.Source: Figure 5 · p. 8, §3.3 · p. 8
4 Conclusion
Muon and SOAP are presented as alternatives to AdamW for diffusion-model training despite higher per-step runtime. The conclusion also flags unexplained optimizer-performance gaps and a possible link between the training trajectory and generative quality.
4 Conclusion
- Muon and SOAP are convincing alternatives to AdamW for diffusion-model training despite their larger runtime per step.Source: §4 · p. 8
- ScheduleFree and the wsd learning-rate schedule can yield degraded generative quality despite reasonably good loss, and the authors leave the possible role of the full training trajectory open for future work.Source: §4 · p. 8
- The performance differences between Adam and SGD and between Muon/SOAP and Adam remain unexplained for this benchmark, which lies outside the domain of previously offered explanations.Source: §4 · p. 8
A Supplementary Material on Experiments
The authors make the benchmark code and training logs publicly available and built their codebase from the official Rozet & Louppe (2023) implementation.
A Supplementary Material on Experiments
- The benchmark code and training logs are public, and the codebase starts from the official Rozet & Louppe (2023) implementation.Source: A Supplementary Material on Experiments · p. 13
A.1 Overview of Model and Dataset
The study focuses on training diffusion models to learn a score function for dynamical-system trajectories, within a data-assimilation task that estimates true trajectories from noisy measurements. The inputs are fluid-velocity snapshots, sampled in five-snapshot windows for a U-Net model.
A.1 Overview of Model and Dataset
- Data assimilation estimates the distribution of true dynamical-system trajectories from noisy measurements; this study focuses on training the score function used for that task.Training and estimation can be performed entirely decoupled.Source: §A.1 · p. 13
- The diffusion model uses snapshots of a two-dimensional Navier–Stokes velocity field; 1,024 trajectories are split into 80% training, 10% validation, and 10% test data.Trajectories are generated at 256 × 256 resolution, down-sampled to 64 × 64, and filtered on their second half.Source: §A.1 · pp. 13–14
- Each training example samples a random five-snapshot window from a trajectory, producing input of shape (b, 10, 64, 64).Source: §A.1 · p. 14
- The score model is a U-Net with three hidden convolutional layers and 22.9 million trainable parameters.The hidden channel sizes are (96, 192, 384), with a time-embedding dimension of 64.Source: §A.1 · p. 14
A.2 Hyperparameters
The section specifies optimizer settings, learning-rate schedules, and sampling choices for the benchmark. Muon combines orthogonalized matrix-gradient updates with AdamW for selected parameters, while schedule factors and other hyperparameters are tuned by method.
A.2 Hyperparameters
- SGD uses heavy-ball momentum 0.9 with dampening 0.9; AdamW, SOAP, and ScheduleFree use (β1, β2) = (0.9, 0.999).Muon uses Nesterov momentum 0.9 for its parameters and (β1, β2) = (0.9, 0.999) for AdamW-trained parameters.Source: §A.2 · p. 14
- Muon approximates a matrix gradient’s closest orthogonal matrix; biases and time embeddings use AdamW, while higher-dimensional parameters are reshaped for Newton–Schulz updates.A heuristic aligns update magnitudes so Muon- and AdamW-trained parameters can share one learning rate and weight decay.Source: §A.2 · p. 14
- Sampling follows Rozet & Louppe (2023), including a cosine diffusion-process schedule; after training, two trajectories are sampled for 64 steps with the same seed.Source: §A.2 · p. 14
- Learning-rate schedules are epoch-wise, with each schedule’s multiplicative factor γ tuned independently; warmup shifts schedules by five epochs.For wsd and sqrt, T0 = ⌊0.8T⌋, giving a cooldown of 20% of training.Source: §A.2 · pp. 14–15
- For wsd and cosine, γ is tuned while weight decay stays at 10^-3; sqrt also tests weight decay values 10^-2 and 10^-4.Source: §A.2 · p. 15
- Table 1 lists default hyperparameters, while Tables 2 and 3 list method-specific hyperparameters for Muon and SOAP.Source: Table 1 · p. 15, Table 2 · p. 15, Table 3 · p. 15
A.3 Additional Plots
The additional plots compare loss-based model selection with generated-trajectory quality across learning-rate schedules and methods. They also show the benchmark’s loss curves, hyperparameter grids, and generated-field visualizations.
A.3 Additional Plots
- The figures show loss curves for each method’s best setup, one-standard-deviation bands across three runs, and learning-rate/weight-decay grids of final validation loss.The grid color scale differs by method; generated-field vorticity is also shown for selected trajectories, including the best SGD setting.Source: Figure 7 · p. 16, Figure 6 · p. 16, Figure 9 · p. 17
- For wsd, the learning rate with minimal validation loss produces lower-quality trajectories; this finding holds across all three seeds, unlike cosine and sqrt.Source: Figure 10 · p. 18
- ScheduleFree’s generated trajectories are significantly worse than AdamW’s despite comparable, slightly higher AdamW loss values.Source: Figure 11 · p. 18
- Adding wsd with a 20% cooldown to ScheduleFree improves generative quality for some hyperparameter configurations.The plotted example uses learning rate 0.001 and weight decay 0.00032.Source: Figure 12 · p. 19