Матрёшечные модели диффузии

Аннотация

Диффузионные модели являются фактическим стандартом для генерации высококачественных изображений и видео, однако обучение моделей в высокоразмерных пространствах остается сложной задачей из-за вычислительных и оптимизационных трудностей. Существующие методы часто прибегают к обучению каскадных моделей в пространстве пикселей или используют понижающее латентное пространство отдельно обученного автоэнкодера. В данной работе мы представляем Matryoshka Diffusion Models (MDM) — сквозную архитектуру для синтеза изображений и видео высокого разрешения. Мы предлагаем диффузионный процесс, который совместно удаляет шум на нескольких разрешениях, и используем архитектуру NestedUNet, где признаки и параметры для входных данных малого масштаба вложены в параметры для больших масштабов. Кроме того, MDM позволяет реализовать прогрессивный график обучения от низких к высоким разрешениям, что приводит к значительному улучшению оптимизации для генерации высокого разрешения. Мы демонстрируем эффективность нашего подхода на различных бенчмарках, включая условную генерацию изображений по классам, генерацию изображений высокого разрешения по тексту и генерацию видео по тексту. Примечательно, что мы можем обучить единую модель в пространстве пикселей с разрешением до 1024x1024 пикселей, демонстрируя сильную обобщающую способность в условиях zero-shot на наборе данных CC12M, содержащем всего 12 миллионов изображений.

English

Diffusion models are the de facto approach for generating high-quality images and videos, but learning high-dimensional models remains a formidable task due to computational and optimization challenges. Existing methods often resort to training cascaded models in pixel space or using a downsampled latent space of a separately trained auto-encoder. In this paper, we introduce Matryoshka Diffusion Models(MDM), an end-to-end framework for high-resolution image and video synthesis. We propose a diffusion process that denoises inputs at multiple resolutions jointly and uses a NestedUNet architecture where features and parameters for small-scale inputs are nested within those of large scales. In addition, MDM enables a progressive training schedule from lower to higher resolutions, which leads to significant improvements in optimization for high-resolution generation. We demonstrate the effectiveness of our approach on various benchmarks, including class-conditioned image generation, high-resolution text-to-image, and text-to-video applications. Remarkably, we can train a single pixel-space model at resolutions of up to 1024x1024 pixels, demonstrating strong zero-shot generalization using the CC12M dataset, which contains only 12 million images.