DiM: Diffusion Mamba für effiziente Hochauflösungs-Bildsynthese

DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis

May 23, 2024

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

cs.AI

Zusammenfassung

Diffusionsmodelle haben große Erfolge bei der Bildgenerierung erzielt, wobei sich das Grundgerüst von U-Net zu Vision-Transformern entwickelt hat. Die Rechenkosten von Transformatoren sind jedoch quadratisch zur Anzahl der Tokens, was zu erheblichen Herausforderungen bei der Bearbeitung von hochauflösenden Bildern führt. In dieser Arbeit schlagen wir Diffusion Mamba (DiM) vor, das die Effizienz von Mamba, einem Sequenzmodell basierend auf Zustandsraummodellen (SSM), mit der Ausdruckskraft von Diffusionsmodellen für eine effiziente Bildsynthese hoher Auflösung kombiniert. Um das Problem anzugehen, dass Mamba nicht auf 2D-Signale verallgemeinern kann, haben wir mehrere Architekturentwürfe erstellt, darunter multidirektionale Scans, erlernbare Padding-Tokens am Ende jeder Zeile und Spalte sowie eine leichte lokale Merkmalsverbesserung. Unsere DiM-Architektur erreicht Effizienz zur Inferenzzeit für hochauflösende Bilder. Darüber hinaus untersuchen wir zur weiteren Verbesserung der Trainingseffizienz für die Bildgenerierung hoher Auflösung mit DiM eine „schwach-zu-stark“-Trainingsstrategie, die DiM auf niedrigauflösenden Bildern (256mal 256) vorab trainiert und dann auf hochauflösenden Bildern (512 mal 512) feinabstimmt. Wir erforschen außerdem trainingsfreie Upsampling-Strategien, um dem Modell die Generierung von Bildern mit höherer Auflösung (z. B. 1024mal 1024 und 1536mal 1536) ohne weitere Feinabstimmung zu ermöglichen. Experimente zeigen die Wirksamkeit und Effizienz unseres 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.