arXiv 2405.142242024-07-10cs.CV

DiM: 面向高效高分辨率图像合成的扩散Mamba

DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis

The University of Hong Kong（香港大学）；Huawei Noah’s Ark Lab（华为诺亚方舟实验室）；Tsinghua University（清华大学）；Shanghai Jiao Tong University（上海交通大学）

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

备注：: The code of our work is available here: {\url{https://github.com/tyshiwo1/DiM-DiffusionMamba/}}

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

针对Transformer类扩散模型高分辨率图像生成计算成本高的问题，提出结合Mamba高效性与扩散模型表达力的DiM架构，搭配多方向扫描等设计及弱到强训练、免训练上采样策略，实现高效高分辨率图像合成。

摘要：

扩散模型在图像生成领域已取得巨大成功，其骨干网络从U-Net逐步演进为Vision Transformers。但Transformers的计算成本与token数量呈二次方关系，在处理高分辨率图像时面临显著挑战。本研究提出Diffusion Mamba（DiM，扩散Mamba），将基于状态空间模型（SSM）的序列模型Mamba的高效性，与扩散模型的表达能力相结合，用于高效高分辨率图像合成。为解决Mamba无法泛化到二维信号的难题，我们开展了多项架构设计，包括多方向扫描、每行每列末端的可学习填充token，以及轻量级局部特征增强。所提出的DiM架构在高分辨率图像推理阶段具备高效性。此外，为进一步提升DiM生成高分辨率图像的训练效率，我们研究了“弱到强”训练策略：先在低分辨率（256×256）图像上预训练DiM，再在高分辨率（512×512）图像上微调。我们进一步探索了免训练上采样策略，使模型无需额外微调即可生成更高分辨率的图像（例如1024×1024和1536×1536）。实验验证了DiM的有效性与高效性。本研究代码可在以下地址获取：https://github.com/tyshiwo1/DiM-DiffusionMamba/

英文摘要：

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 ($256\times 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., $1024\times 1024$ and $1536\times 1536$) without further fine-tuning. Experiments demonstrate the effectiveness and efficiency of our DiM. The code of our work is available here: {\url{https://github.com/tyshiwo1/DiM-DiffusionMamba/}}.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。