DiM: Mamba de Difusión para la Síntesis Eficiente de Imágenes de Alta Resolución

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

Resumen

Los modelos de difusión han logrado un gran éxito en la generación de imágenes, con la arquitectura principal evolucionando desde U-Net hasta Transformers de Visión. Sin embargo, el costo computacional de los Transformers es cuadrático respecto al número de tokens, lo que plantea desafíos significativos al trabajar con imágenes de alta resolución. En este trabajo, proponemos Diffusion Mamba (DiM), que combina la eficiencia de Mamba, un modelo de secuencias basado en Modelos de Espacio de Estados (SSM), con el poder expresivo de los modelos de difusión para una síntesis eficiente de imágenes de alta resolución. Para abordar el desafío de que Mamba no puede generalizar a señales 2D, realizamos varios diseños arquitectónicos, incluyendo escaneos multidireccionales, tokens de relleno aprendibles al final de cada fila y columna, y mejora ligera de características locales. Nuestra arquitectura DiM logra eficiencia en tiempo de inferencia para imágenes de alta resolución. Además, para mejorar aún más la eficiencia en el entrenamiento para la generación de imágenes de alta resolución con DiM, investigamos una estrategia de entrenamiento "de débil a fuerte" que preentrena DiM en imágenes de baja resolución (256x256) y luego lo ajusta en imágenes de alta resolución (512x512). También exploramos estrategias de aumento de resolución sin entrenamiento para permitir que el modelo genere imágenes de mayor resolución (por ejemplo, 1024x1024 y 1536x1536) sin necesidad de ajustes adicionales. Los experimentos demuestran la efectividad y eficiencia de nuestro 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.