DiM: Diffusion Mamba per la Sintesi Efficiente di Immagini ad Alta Risoluzione
DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis
May 23, 2024
Autori: Yao Teng, Yue Wu, Han Shi, Xuefei Ning, Guohao Dai, Yu Wang, Zhenguo Li, Xihui Liu
cs.AI
Abstract
I modelli di diffusione hanno ottenuto un grande successo nella generazione di immagini, con l'architettura di base che è evoluta da U-Net ai Vision Transformer. Tuttavia, il costo computazionale dei Transformer è quadratico rispetto al numero di token, portando a sfide significative quando si ha a che fare con immagini ad alta risoluzione. In questo lavoro, proponiamo Diffusion Mamba (DiM), che combina l'efficienza di Mamba, un modello di sequenza basato su State Space Models (SSM), con il potere espressivo dei modelli di diffusione per una sintesi efficiente di immagini ad alta risoluzione. Per affrontare la sfida che Mamba non può generalizzare a segnali 2D, abbiamo progettato diverse architetture, tra cui scansioni multi-direzionali, token di padding apprendibili alla fine di ogni riga e colonna, e un miglioramento leggero delle caratteristiche locali. La nostra architettura DiM raggiunge un'efficienza al momento dell'inferenza per immagini ad alta risoluzione. Inoltre, per migliorare ulteriormente l'efficienza dell'addestramento per la generazione di immagini ad alta risoluzione con DiM, abbiamo investigato una strategia di addestramento "da debole a forte" che prevede il pre-addestramento di DiM su immagini a bassa risoluzione (256x256) e poi il fine-tuning su immagini ad alta risoluzione (512x512). Abbiamo inoltre esplorato strategie di upsampling senza addestramento per consentire al modello di generare immagini a risoluzione ancora più alta (ad esempio, 1024x1024 e 1536x1536) senza ulteriore fine-tuning. Gli esperimenti dimostrano l'efficacia e l'efficienza del nostro 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.