Escalonando Transformadores de Difusão para 16 Bilhões 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

Resumo

Neste artigo, apresentamos o DiT-MoE, uma versão esparsa do Transformer de difusão, que é escalável e competitiva com redes densas, ao mesmo tempo que exibe inferência altamente otimizada. O DiT-MoE inclui dois projetos simples: roteamento de especialistas compartilhados e perda de equilíbrio de nível de especialista, capturando assim conhecimento comum e reduzindo a redundância entre os diferentes especialistas roteados. Quando aplicado à geração condicional de imagens, uma análise aprofundada dos ganhos de especialização dos especialistas revela algumas observações interessantes: (i) A seleção de especialistas mostra preferência pela posição espacial e pelo passo de tempo de remoção de ruído, sendo insensível a informações condicionais de classe diferentes; (ii) À medida que as camadas MoE se aprofundam, a seleção de especialistas gradualmente se desloca de posições espaciais específicas para dispersão e equilíbrio; (iii) A especialização dos especialistas tende a ser mais concentrada no início do passo de tempo e depois uniforme gradualmente após a metade. Atribuímos isso ao processo de difusão que primeiro modela as informações espaciais de baixa frequência e depois as informações complexas de alta frequência. Com base nas orientações acima, uma série de experimentos DiT-MoE alcança experimentalmente desempenho comparável a redes densas, mas requer uma carga computacional muito menor durante a inferência. Mais encorajadoramente, demonstramos o potencial do DiT-MoE com dados de imagem sintetizados, escalando o modelo de difusão para um parâmetro de 16,5B que alcança um novo escore FID-50K de SoTA de 1,80 em configurações de resolução de 512 vezes 512. A página do projeto: 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.