Diffusie Transformers opschalen naar 16 miljard parameters

Scaling Diffusion Transformers to 16 Billion Parameters

July 16, 2024

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

cs.AI

Samenvatting

In dit artikel presenteren we DiT-MoE, een sparse versie van de diffusion Transformer, die schaalbaar is en kan concurreren met dichte netwerken terwijl deze sterk geoptimaliseerde inferentie vertoont. De DiT-MoE omvat twee eenvoudige ontwerpen: gedeelde expert routing en expert-level balance loss, waardoor gemeenschappelijke kennis wordt vastgelegd en redundantie tussen de verschillende gerouteerde experts wordt verminderd. Wanneer toegepast op conditionele beeldgeneratie, levert een diepgaande analyse van expertspecialisatie enkele interessante observaties op: (i) Expertselectie vertoont voorkeur voor ruimtelijke positie en denoising-tijdstap, terwijl deze ongevoelig is voor verschillende klasse-conditionele informatie; (ii) Naarmate de MoE-lagen dieper gaan, verschuift de selectie van experts geleidelijk van specifieke ruimtelijke positie naar spreiding en balans. (iii) Expertspecialisatie neigt meer geconcentreerd te zijn in de vroege tijdstap en wordt daarna geleidelijk uniform na de helft. Wij schrijven dit toe aan het diffusieproces dat eerst de laagfrequente ruimtelijke informatie modelleert en vervolgens de hoogfrequente complexe informatie. Op basis van de bovenstaande richtlijnen bereikt een reeks DiT-MoE experimenteel prestaties die gelijk zijn aan die van dichte netwerken, maar vereist veel minder rekenkracht tijdens inferentie. Nog bemoedigender is dat we het potentieel van DiT-MoE aantonen met gesynthetiseerde beeldgegevens, waarbij het diffusiemodel wordt geschaald naar 16,5B parameters dat een nieuwe SoTA FID-50K-score van 1,80 behaalt in 512x512 resolutie-instellingen. De projectpagina: 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.