DiM: Diffusion Mamba voor Efficiënte Synthese van Hoogresolutiebeelden

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

Samenvatting

Diffusiemodellen hebben grote successen geboekt in beeldgeneratie, waarbij de backbone is geëvolueerd van U-Net naar Vision Transformers. De rekenkosten van Transformers zijn echter kwadratisch ten opzichte van het aantal tokens, wat aanzienlijke uitdagingen met zich meebrengt bij het verwerken van hoge-resolutiebeelden. In dit werk stellen we Diffusion Mamba (DiM) voor, dat de efficiëntie van Mamba, een sequentiemodel gebaseerd op State Space Models (SSM), combineert met de expressieve kracht van diffusiemodellen voor efficiënte synthese van hoge-resolutiebeelden. Om de uitdaging aan te pakken dat Mamba niet kan generaliseren naar 2D-signalen, hebben we verschillende architectuurontwerpen gemaakt, waaronder multidirectionele scans, leerbare padding-tokens aan het einde van elke rij en kolom, en lichtgewicht lokale feature-verbetering. Onze DiM-architectuur bereikt efficiëntie tijdens de inferentie voor hoge-resolutiebeelden. Daarnaast onderzoeken we, om de trainingsefficiëntie voor hoge-resolutiebeeldgeneratie met DiM verder te verbeteren, een "weak-to-strong" trainingsstrategie waarbij DiM wordt voorgetraind op lage-resolutiebeelden (256x256) en vervolgens wordt gefinetuned op hoge-resolutiebeelden (512x512). We verkennen verder training-vrije upsamplingstrategieën om het model in staat te stellen hogere-resolutiebeelden te genereren (bijvoorbeeld 1024x1024 en 1536x1536) zonder verdere finetuning. Experimenten demonstreren de effectiviteit en efficiëntie van onze 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.