Optimizers for Diffusion Models: A Controlled Benchmark
Abstractの概要
本論文では、text8でのmasked diffusion、QM9でのuniform diffusion、LM1BでのDUO、CelebA-64でのGaussian画像拡散という4つの定式化にわたり、拡散モデル学習における7種類のオプティマイザの制御されたベンチマークを提示しています。すべてのオプティマイザは共通のチューニングプロトコルの下で評価され、選定された上位手法は3つのシードを用いてフル予算で再学習されており、先行研究のデフォルト設定を引き継ぐよりも厳密に制御された比較が行われています。本研究により、オプティマイザの性能はタスク間で一貫して転移するのではなく拡散の定式化に強く依存することが明らかになりました。また、特にQM9やCelebA-64において、目的関数ベースの指標とサンプル品質の指標が一致しない場合があることも示されています。
新規性
本研究は、比較のためにGaussian拡散による画像設定も含めつつ、離散拡散モデルに焦点を当てた初の制御されたオプティマイザベンチマークであると考えられます。統一的な探索プロトコル、上位手法に対する3つのシードを用いたフル予算での再学習、およびオプティマイザの順位が定式化間で転移するかどうかの横断的分析を組み合わせた点が独自の貢献です。
成果
適切にチューニングされたAdamWは強力なベースラインであり、text8やCelebA-64でそれを明確に上回る代替手法は存在しなかった一方、QM9ではSOAP、Lion、Muonがわずかながら確実な改善を達成しました。LM1BではMuonとMARS-Mがより大きな差でAdamWを上回り、MARSも優位に立ちましたが、Lionはカリキュラムの切り替え後に3シード中2シードで発散しました。4つのタスク全体を通じてSchedule-Freeが最下位となり、オプティマイザの順位は定式化によって大幅に変動しました。
論文の注目点
- 本ベンチマークは、同一のチューニングおよび再学習プロトコルの下で4つの拡散定式化にわたり7種類のオプティマイザを比較し、エンドツーエンドの再現性のためのコードを公開しています。
- AdamWは依然として強力なデフォルトですが、Muon、MARS-M、SOAPはそれぞれ少なくとも1つの定式化でチューニングされたAdamWを上回り、拡散モデルの学習においてオプティマイザの選択が重要であることを示しています。
- オプティマイザの順位は定式化間で確実に転移するわけではなく、学習損失や尤度形式の指標の低さがFIDやFCDなどのサンプル品質指標の向上に必ずしも一致しないことが示されました。