Analyse und Verbesserung der Trainingsdynamik von Diffusionsmodellen
Analyzing and Improving the Training Dynamics of Diffusion Models
December 5, 2023
Autoren: Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine
cs.AI
Zusammenfassung
Diffusionsmodelle dominieren derzeit das Feld der datengesteuerten Bildsynthese
mit ihrer beispiellosen Skalierbarkeit auf große Datensätze. In diesem Artikel
identifizieren und beheben wir mehrere Ursachen für ungleichmäßiges und
ineffektives Training in der populären ADM-Diffusionsmodellarchitektur, ohne
dessen übergeordnete Struktur zu verändern. Durch die Beobachtung unkontrollierter
Größenänderungen und Ungleichgewichte in den Netzwerkaktivierungen und -gewichten
während des Trainingsprozesses gestalten wir die Netzwerkschichten neu, um die
Größenordnungen von Aktivierungen, Gewichten und Updates im Erwartungswert zu
erhalten. Wir stellen fest, dass die systematische Anwendung dieses Prinzips die
beobachteten Drifts und Ungleichgewichte beseitigt, was zu deutlich besseren
Netzwerken bei gleicher Rechenkomplexität führt. Unsere Modifikationen verbessern
den bisherigen Rekord-FID von 2,41 in der ImageNet-512-Synthese auf 1,81, erreicht
mit schnellem deterministischem Sampling.
Als eigenständigen Beitrag präsentieren wir eine Methode zur nachträglichen
Festlegung der Parameter des exponentiellen gleitenden Durchschnitts (EMA), d.h.
nach Abschluss des Trainingslaufs. Dies ermöglicht eine präzise Feinabstimmung der
EMA-Länge ohne die Kosten mehrerer Trainingsläufe und offenbart deren überraschende
Wechselwirkungen mit Netzwerkarchitektur, Trainingsdauer und Guidance.
English
Diffusion models currently dominate the field of data-driven image synthesis
with their unparalleled scaling to large datasets. In this paper, we identify
and rectify several causes for uneven and ineffective training in the popular
ADM diffusion model architecture, without altering its high-level structure.
Observing uncontrolled magnitude changes and imbalances in both the network
activations and weights over the course of training, we redesign the network
layers to preserve activation, weight, and update magnitudes on expectation. We
find that systematic application of this philosophy eliminates the observed
drifts and imbalances, resulting in considerably better networks at equal
computational complexity. Our modifications improve the previous record FID of
2.41 in ImageNet-512 synthesis to 1.81, achieved using fast deterministic
sampling.
As an independent contribution, we present a method for setting the
exponential moving average (EMA) parameters post-hoc, i.e., after completing
the training run. This allows precise tuning of EMA length without the cost of
performing several training runs, and reveals its surprising interactions with
network architecture, training time, and guidance.