DiM: 高解像度画像合成のための効率的なDiffusion Mamba

要旨

拡散モデルは画像生成において大きな成功を収めており、そのバックボーンはU-NetからVision Transformersへと進化してきました。しかし、Transformersの計算コストはトークン数の二乗に比例するため、高解像度画像を扱う際に大きな課題となっています。本研究では、State Space Models (SSM) に基づくシーケンスモデルであるMambaの効率性と、拡散モデルの表現力を組み合わせたDiffusion Mamba (DiM)を提案し、効率的な高解像度画像合成を実現します。Mambaが2D信号に一般化できないという課題に対処するため、多方向スキャン、各行および各列の終端に学習可能なパディングトークンを追加、軽量な局所特徴強化など、いくつかのアーキテクチャ設計を行いました。我々のDiMアーキテクチャは、高解像度画像に対する推論時の効率性を実現しています。さらに、DiMを用いた高解像度画像生成の訓練効率をさらに向上させるため、低解像度画像（256×256）でDiMを事前訓練し、その後高解像度画像（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.

DiM: 高解像度画像合成のための効率的なDiffusion Mamba

DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis

要旨

Support