Масштабирование трансформеров диффузии до 16 миллиардов параметров

Scaling Diffusion Transformers to 16 Billion Parameters

July 16, 2024

Авторы: Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, Junshi Huang

cs.AI

Аннотация

В данной статье мы представляем DiT-MoE, разреженную версию диффузионного трансформера, которая масштабируема и конкурентоспособна по сравнению с плотными сетями, обладая высокооптимизированным выводом. DiT-MoE включает два простых дизайна: общую маршрутизацию экспертов и потерю баланса на уровне экспертов, тем самым захватывая общие знания и уменьшая избыточность среди различных маршрутизированных экспертов. При применении к условному созданию изображений глубокий анализ специализации экспертов приводит к нескольким интересным наблюдениям: (i) Выбор эксперта проявляет предпочтение к пространственному положению и шагу времени шумоподавления, при этом нечувствителен к различной классовой условной информации; (ii) По мере углубления слоев MoE, выбор экспертов постепенно смещается от конкретного пространственного положения к дисперсии и балансу. (iii) Специализация экспертов имеет тенденцию быть более сосредоточенной на раннем временном шаге, а затем постепенно становится равномерной после середины. Мы связываем это с процессом диффузии, который сначала моделирует низкочастотную пространственную информацию, а затем высокочастотную сложную информацию. Основываясь на вышеуказанных рекомендациях, ряд экспериментов с DiT-MoE экспериментально достигает производительности на уровне плотных сетей, требуя при этом значительно меньше вычислительной нагрузки во время вывода. Более того, мы продемонстрировали потенциал DiT-MoE на синтезированных изображениях, масштабируя модель диффузии на 16,5 млрд параметров, что дает новый показатель FID-50K в размере 1,80 в настройках разрешения 512х512. Страница проекта: 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.