DiM : Diffusion Mamba pour la synthèse efficace d'images haute résolution

DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis

May 23, 2024

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

cs.AI

Résumé

Les modèles de diffusion ont connu un grand succès dans la génération d'images, avec une évolution de l'architecture de base allant de l'U-Net aux Vision Transformers. Cependant, le coût computationnel des Transformers est quadratique par rapport au nombre de tokens, ce qui pose des défis importants lors du traitement d'images haute résolution. Dans ce travail, nous proposons Diffusion Mamba (DiM), qui combine l'efficacité de Mamba, un modèle de séquence basé sur les State Space Models (SSM), avec la puissance expressive des modèles de diffusion pour une synthèse efficace d'images haute résolution. Pour résoudre le défi que Mamba ne peut pas généraliser aux signaux 2D, nous proposons plusieurs conceptions architecturales, incluant des scans multidirectionnels, des tokens de remplissage apprenables à la fin de chaque ligne et colonne, et une amélioration légère des caractéristiques locales. Notre architecture DiM permet une efficacité en temps d'inférence pour les images haute résolution. De plus, pour améliorer davantage l'efficacité de l'entraînement pour la génération d'images haute résolution avec DiM, nous explorons une stratégie d'entraînement « faible à fort » qui pré-entraîne DiM sur des images basse résolution (256 × 256) puis l'affine sur des images haute résolution (512 × 512). Nous explorons également des stratégies de suréchantillonnage sans entraînement pour permettre au modèle de générer des images de résolution encore plus élevée (par exemple, 1024 × 1024 et 1536 × 1536) sans affinage supplémentaire. Les expériences démontrent l'efficacité et l'efficience de notre 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.