Scalabilità dei Transformers Diffusivi fino a 16 Miliardi di Parametri

Scaling Diffusion Transformers to 16 Billion Parameters

July 16, 2024

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

cs.AI

Abstract

In questo articolo presentiamo DiT-MoE, una versione sparsa del Transformer di diffusione, scalabile e competitiva rispetto alle reti dense, pur mostrando un'inferenza altamente ottimizzata. Il DiT-MoE include due semplici design: il routing condiviso degli esperti e una funzione di perdita bilanciata a livello di esperto, catturando così conoscenze comuni e riducendo la ridondanza tra i diversi esperti instradati. Quando applicato alla generazione condizionata di immagini, un'analisi approfondita della specializzazione degli esperti rivela alcune osservazioni interessanti: (i) La selezione degli esperti mostra una preferenza per la posizione spaziale e il passo temporale di denoising, mentre è insensibile alle diverse informazioni condizionate per classe; (ii) Man mano che i livelli MoE diventano più profondi, la selezione degli esperti passa gradualmente da una posizione spaziale specifica a una dispersione e bilanciamento. (iii) La specializzazione degli esperti tende a essere più concentrata nei primi passi temporali per poi uniformarsi gradualmente dopo la metà. Attribuiamo ciò al processo di diffusione che modella prima le informazioni spaziali a bassa frequenza e poi le informazioni complesse ad alta frequenza. Sulla base di queste indicazioni, una serie di DiT-MoE raggiunge sperimentalmente prestazioni pari a quelle delle reti dense, pur richiedendo un carico computazionale inferiore durante l'inferenza. Ancora più incoraggiante, dimostriamo il potenziale di DiT-MoE con dati di immagini sintetizzate, scalando il modello di diffusione a 16,5 miliardi di parametri, ottenendo un nuovo punteggio SoTA FID-50K di 1,80 in impostazioni di risoluzione 512x512. La pagina del progetto: 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.