Modelos de Difusão Matryoshka

Resumo

Modelos de difusão são a abordagem de fato para gerar imagens e vídeos de alta qualidade, mas aprender modelos de alta dimensão continua sendo uma tarefa formidável devido a desafios computacionais e de otimização. Métodos existentes frequentemente recorrem ao treinamento de modelos em cascata no espaço de pixels ou ao uso de um espaço latente reduzido de um autoencoder treinado separadamente. Neste artigo, introduzimos os Modelos de Difusão Matryoshka (MDM), uma estrutura de ponta a ponta para síntese de imagens e vídeos de alta resolução. Propomos um processo de difusão que remove ruídos de entradas em múltiplas resoluções de forma conjunta e utiliza uma arquitetura NestedUNet, onde características e parâmetros para entradas em pequena escala estão aninhados dentro daqueles de escalas maiores. Além disso, o MDM permite um cronograma de treinamento progressivo de resoluções mais baixas para mais altas, o que leva a melhorias significativas na otimização para geração de alta resolução. Demonstramos a eficácia de nossa abordagem em vários benchmarks, incluindo geração de imagens condicionadas por classe, geração de imagens de alta resolução a partir de texto e aplicações de texto para vídeo. Notavelmente, conseguimos treinar um único modelo no espaço de pixels em resoluções de até 1024x1024 pixels, demonstrando uma forte generalização zero-shot usando o conjunto de dados CC12M, que contém apenas 12 milhões de imagens.

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.