Skalierung von Diffusions-Transformatoren auf 16 Milliarden Parameter

Scaling Diffusion Transformers to 16 Billion Parameters

July 16, 2024

Autoren: Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, Junshi Huang

cs.AI

Zusammenfassung

In diesem Paper präsentieren wir DiT-MoE, eine spärliche Version des Diffusions-Transformers, die skalierbar ist und im Wettbewerb mit dichten Netzwerken steht, während sie eine hochgradig optimierte Inferenz aufweist. Das DiT-MoE umfasst zwei einfache Designs: gemeinsame Experten-Routing und Experten-Level-Balance-Verlust, wodurch gemeinsames Wissen erfasst und Redundanz unter den verschiedenen gerouteten Experten reduziert wird. Bei der Anwendung auf die bedingte Bildgenerierung ergibt eine eingehende Analyse der Expertenspezialisierung einige interessante Beobachtungen: (i) Die Expertenauswahl zeigt eine Präferenz für den räumlichen Standort und den Rauschunterdrückungs-Zeitschritt, während sie unempfindlich gegenüber unterschiedlichen klassenbedingten Informationen ist; (ii) Mit zunehmender Tiefe der MoE-Schichten verschiebt sich die Auswahl der Experten allmählich von einer spezifischen räumlichen Position zu Dispersion und Balance. (iii) Die Expertenspezialisierung neigt dazu, am Anfang des Zeitschritts konzentrierter zu sein und sich dann nach der Hälfte allmählich zu vereinheitlichen. Wir führen dies auf den Diffusionsprozess zurück, der zuerst die niederfrequenten räumlichen Informationen modelliert und dann hochfrequente komplexe Informationen. Basierend auf obiger Anleitung erreicht eine Reihe von DiT-MoE experimentell eine Leistung auf Augenhöhe mit dichten Netzwerken, erfordert jedoch während der Inferenz wesentlich weniger Rechenlast. Noch ermutigender ist, dass wir das Potenzial von DiT-MoE mit synthetisierten Bilddaten demonstrieren, indem wir das Diffusionsmodell auf 16,5 Milliarden Parametern skalieren, was einen neuen SoTA FID-50K-Score von 1,80 in Einstellungen mit einer Auflösung von 512mal512 erreicht. Die Projektseite: https://github.com/feizc/DiT-MoE.

English

In this paper, we present DiT-MoE, a sparse version of the diffusion Transformer, that is scalable and competitive with dense networks while exhibiting highly optimized inference. The DiT-MoE includes two simple designs: shared expert routing and expert-level balance loss, thereby capturing common knowledge and reducing redundancy among the different routed experts. When applied to conditional image generation, a deep analysis of experts specialization gains some interesting observations: (i) Expert selection shows preference with spatial position and denoising time step, while insensitive with different class-conditional information; (ii) As the MoE layers go deeper, the selection of experts gradually shifts from specific spacial position to dispersion and balance. (iii) Expert specialization tends to be more concentrated at the early time step and then gradually uniform after half. We attribute it to the diffusion process that first models the low-frequency spatial information and then high-frequency complex information. Based on the above guidance, a series of DiT-MoE experimentally achieves performance on par with dense networks yet requires much less computational load during inference. More encouragingly, we demonstrate the potential of DiT-MoE with synthesized image data, scaling diffusion model at a 16.5B parameter that attains a new SoTA FID-50K score of 1.80 in 512times512 resolution settings. The project page: https://github.com/feizc/DiT-MoE.