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

Мы исследуем новый класс диффузионных моделей, основанных на архитектуре трансформера. Мы обучаем латентные диффузионные модели изображений, заменяя обычно используемый U-Net бэкенд на трансформер, работающий с латентными патчами. Мы анализируем масштабируемость наших Diffusion Transformers (DiT) с точки зрения сложности прямого прохода, измеряемой в GFLOPS. Мы обнаруживаем, что у DiT с более высоким значением GFLOPS — достигнутым за счёт увеличения глубины/ширины трансформера или увеличения числа входных токенов — последовательно ниже FID. Помимо хороших свойств масштабируемости, наши самые крупные модели DiT-XL/2 превосходят все предыдущие диффузионные модели на условных по классам бенчмарках ImageNet 512×512 и 256×256, достигая рекордного FID 2.27 для последнего.