DiM: Mamba de Difusão para Síntese Eficiente de Imagens de Alta Resolução

DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis

May 23, 2024

Autores: Yao Teng, Yue Wu, Han Shi, Xuefei Ning, Guohao Dai, Yu Wang, Zhenguo Li, Xihui Liu

cs.AI

Resumo

Os modelos de difusão têm alcançado grande sucesso na geração de imagens, com a arquitetura evoluindo de U-Net para Vision Transformers. No entanto, o custo computacional dos Transformers é quadrático em relação ao número de tokens, o que representa desafios significativos ao lidar com imagens de alta resolução. Neste trabalho, propomos o Diffusion Mamba (DiM), que combina a eficiência do Mamba, um modelo de sequência baseado em State Space Models (SSM), com o poder expressivo dos modelos de difusão para a síntese eficiente de imagens de alta resolução. Para enfrentar o desafio de que o Mamba não pode generalizar para sinais 2D, realizamos vários projetos de arquitetura, incluindo varreduras multidirecionais, tokens de preenchimento aprendíveis no final de cada linha e coluna, e aprimoramento leve de características locais. Nossa arquitetura DiM alcança eficiência no tempo de inferência para imagens de alta resolução. Além disso, para melhorar ainda mais a eficiência do treinamento na geração de imagens de alta resolução com o DiM, investigamos uma estratégia de treinamento "fraco para forte" que pré-treina o DiM em imagens de baixa resolução (256x256) e, em seguida, o ajusta em imagens de alta resolução (512x512). Exploramos ainda estratégias de upsampling sem treinamento para permitir que o modelo gere imagens de resolução ainda maior (por exemplo, 1024x1024 e 1536x1536) sem a necessidade de ajuste adicional. Experimentos demonstram a eficácia e eficiência do nosso DiM.

English

Diffusion models have achieved great success in image generation, with the backbone evolving from U-Net to Vision Transformers. However, the computational cost of Transformers is quadratic to the number of tokens, leading to significant challenges when dealing with high-resolution images. In this work, we propose Diffusion Mamba (DiM), which combines the efficiency of Mamba, a sequence model based on State Space Models (SSM), with the expressive power of diffusion models for efficient high-resolution image synthesis. To address the challenge that Mamba cannot generalize to 2D signals, we make several architecture designs including multi-directional scans, learnable padding tokens at the end of each row and column, and lightweight local feature enhancement. Our DiM architecture achieves inference-time efficiency for high-resolution images. In addition, to further improve training efficiency for high-resolution image generation with DiM, we investigate ``weak-to-strong'' training strategy that pretrains DiM on low-resolution images (256times 256) and then finetune it on high-resolution images (512 times 512). We further explore training-free upsampling strategies to enable the model to generate higher-resolution images (e.g., 1024times 1024 and 1536times 1536) without further fine-tuning. Experiments demonstrate the effectiveness and efficiency of our DiM.