DiM: Диффузионная мамба для эффективного синтеза изображений высокого разрешения

DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis

May 23, 2024

Авторы: Yao Teng, Yue Wu, Han Shi, Xuefei Ning, Guohao Dai, Yu Wang, Zhenguo Li, Xihui Liu

cs.AI

Аннотация

Модели диффузии достигли большого успеха в генерации изображений, пройдя путь от U-Net к Vision Transformers. Однако вычислительная сложность трансформеров квадратично зависит от количества токенов, что создает значительные трудности при работе с изображениями высокого разрешения. В данной работе мы предлагаем модель Diffusion Mamba (DiM), которая объединяет эффективность Mamba, модели последовательности на основе моделей пространства состояний (SSM), с выразительной мощностью моделей диффузии для эффективного синтеза изображений высокого разрешения. Для решения проблемы того, что Mamba не способна обобщаться на 2D сигналы, мы предлагаем несколько архитектурных решений, включая многонаправленные сканирования, обучаемые токены-заполнители в конце каждой строки и столбца, а также легкое улучшение локальных признаков. Наша архитектура DiM обеспечивает эффективность времени вывода для изображений высокого разрешения. Кроме того, для дальнейшего улучшения эффективности обучения генерации изображений высокого разрешения с помощью DiM мы исследуем стратегию обучения "от слабого к сильному", которая предварительно обучает DiM на изображениях низкого разрешения (256 на 256), а затем донастраивает его на изображениях высокого разрешения (512 на 512). Мы также исследуем стратегии повышения разрешения без обучения, чтобы позволить модели генерировать изображения более высокого разрешения (например, 1024 на 1024 и 1536 на 1536) без дополнительной донастройки. Эксперименты демонстрируют эффективность и эффективность нашей модели 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.