Kürzlich haben Forscher der University of California, Irvine und anderer Institutionen die Trainingskosten für Diffusionsmodelle auf 1.890 US-Dollar reduziert, indem sie Strategien wie verzögerte Maskierung, MoE und hierarchische Skalierung einsetzten. Dies ist eine Größenordnung niedriger als die bisher günstigste Methode (Wuerstchen) mit 28.400 US-Dollar, während Modelle wie Stable Diffusion noch teurer sind.
Forscher haben verschiedene Ansätze ausprobiert, um diese enormen Kosten zu senken. Zum Beispiel benötigten ursprüngliche Diffusionsmodelle etwa 1000 Schritte von Rauschen zu Bild, was jetzt auf etwa 20 Schritte oder weniger reduziert wurde. Als die grundlegenden Module in Diffusionsmodellen schrittweise von Unet (CNN) zu DiT (Transformer) wechselten, folgten einige Optimierungen basierend auf Transformer-Eigenschaften, wie Quantisierung, Überspringen redundanter Berechnungen in der Attention und Pipeline-Verarbeitung.
Diese Studie trainierte ein Diffusionsmodell mit 11,6 Milliarden Parametern von Grund auf für nur 1.890 US-Dollar. Dies stellt eine Verbesserung um eine Größenordnung gegenüber dem Stand der Technik dar und gibt normalen Menschen Hoffnung auf Vortraining. Wichtiger noch, die kostensenkenden Techniken beeinträchtigten nicht die Modellleistung, und die 11,6 Milliarden Parameter lieferten sehr gute Ergebnisse. Neben dem visuellen Eindruck waren auch die Datenkennzahlen des Modells hervorragend, mit FID-Scores sehr nahe an Stable Diffusion 1.5 und DALL·E 2.
Die Geheimnisse der Kosteneinsparung umfassen:
-
Verzögerte Maskierungsstrategie: Verwendung eines Patch-Mixers zur Vorverarbeitung vor der Maskierung, um Informationen verworfener Patches in überlebende Patches einzubetten und so die Leistungseinbußen durch hohe Maskierung erheblich zu reduzieren.
-
Feinabstimmung: Leichte Feinabstimmung (unmaskiert) nach dem Vortraining (maskiert), um durch Maskierung verursachte unerwünschte Generierungsartefakte zu mildern.
-
MoE und hierarchische Skalierung: Verwendung vereinfachter MoE-Schichten basierend auf Expertenauswahl-Routing, um Modellparameter und Ausdrucksfähigkeit zu erhöhen, ohne die Trainingskosten signifikant zu steigern. Gleichzeitig wurde eine hierarchische Skalierungsmethode berücksichtigt, die die Breite der Transformer-Blöcke linear erhöht.
In der experimentellen Konfiguration verwendeten die Autoren zwei DiT-Varianten: DiT-Tiny/2 und DiT-Xl/2 mit einer Patch-Größe von 2. Alle Modelle wurden mit einem AdamW-Optimierer mit Cosinus-Lernratenabnahme und hohem Gewichtsverfall trainiert. Das Frontend des Modells verwendet den vierkanaligen variationellen Autoencoder (VAE) aus dem Stable-Diffusion-XL-Modell zur Extraktion von Bildmerkmalen. Zusätzlich wurde die Leistung des neuesten 16-Kanal-VAE im großen Maßstab (kostengünstige Version) getestet.
Die Autoren verwendeten das EDM-Framework als einheitliche Trainingseinstellung für alle Diffusionsmodelle und nutzten FID sowie CLIP-Scores zur Leistungsmessung der Bildgenerierungsmodelle. Als Textencoder wurde das am häufigsten verwendete CLIP-Modell gewählt.
Als Trainingsdatensatz wurden drei reale Bilddatensätze (Conceptual Captions, Segment Anything, TextCaps) mit 22 Millionen Bild-Text-Paaren verwendet. Da SA1B keine echten Bildunterschriften bereitstellt, wurden hier synthetische Bildunterschriften verwendet, die vom LLaVA-Modell generiert wurden.