Escalando Transformadores de Difusión a 16 Mil Millones de Parámetros

Scaling Diffusion Transformers to 16 Billion Parameters

July 16, 2024

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

cs.AI

Resumen

En este documento, presentamos DiT-MoE, una versión escasa del Transformer de difusión, que es escalable y competitiva con redes densas mientras exhibe una inferencia altamente optimizada. El DiT-MoE incluye dos diseños simples: enrutamiento de expertos compartidos y pérdida de equilibrio a nivel de experto, capturando así conocimientos comunes y reduciendo la redundancia entre los diferentes expertos enrutados. Cuando se aplica a la generación condicional de imágenes, un profundo análisis de las ganancias de especialización de los expertos revela algunas observaciones interesantes: (i) La selección de expertos muestra preferencia por la posición espacial y el paso de tiempo de eliminación de ruido, mientras es insensible a la información condicional de clase diferente; (ii) A medida que las capas MoE se vuelven más profundas, la selección de expertos se desplaza gradualmente desde una posición espacial específica hacia la dispersión y el equilibrio; (iii) La especialización de los expertos tiende a ser más concentrada en el paso de tiempo temprano y luego gradualmente uniforme después de la mitad. Lo atribuimos al proceso de difusión que primero modela la información espacial de baja frecuencia y luego la información compleja de alta frecuencia. Basándonos en las pautas anteriores, una serie de experimentos DiT-MoE logra experimentalmente un rendimiento similar al de las redes densas, pero requiere mucha menos carga computacional durante la inferencia. Más alentadoramente, demostramos el potencial de DiT-MoE con datos de imágenes sintetizadas, escalando el modelo de difusión a 16.5B parámetros que alcanza un nuevo puntaje FID-50K de SoTA de 1.80 en configuraciones de resolución de 512x512. La página del proyecto: 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.