将小波变换融入Mamba,提升图像生成的局部结构与长程频率感知能力。
DiMSUM: Diffusion Mamba -- A Scalable and Unified Spatial-Frequency Method for Image Generation
- 用小波分解图像频率成分,增强Mamba对局部特征的感知。
- 跨注意力融合空间与频域信息,提升生成图像细节质量。
- 引入全局共享变压器,加速训练并改善生成效果,适合图像生成研究者。
我们提出一种新型状态空间架构用于扩散模型,有效结合图像的空间与频率信息,增强对输入图像局部特征的归纳偏置。尽管状态空间网络(如Mamba)作为循环神经网络的革新,通常按从左到右顺序扫描输入序列,但在处理图像数据时难以设计有效的扫描策略。本文表明,将小波变换融入Mamba可增强视觉输入的局部结构感知,并通过将频率解耦为小波子带,更有效地捕捉高频与低频之间的长程关系。这些小波基输出经由交叉注意力融合层与原始Mamba输出无缝融合,整合空间与频率信息,优化状态空间模型的顺序感知能力,这对图像生成的细节和整体质量至关重要。此外,引入全局共享变压器以显著提升Mamba性能,利用其强大的全局关系建模能力。在标准基准上的大量实验表明,该方法优于DiT和DIFFUSSM,实现更快训练收敛并产出高质量图像。代码与预训练模型已公开于 https://github.com/VinAIResearch/DiMSUM.git。
原文摘要 · Abstract (English)
We introduce a novel state-space architecture for diffusion models, effectively harnessing spatial and frequency information to enhance the inductive bias towards local features in input images for image generation tasks. While state-space networks, including Mamba, a revolutionary advancement in recurrent neural networks, typically scan input sequences from left to right, they face difficulties in designing effective scanning strategies, especially in the processing of image data. Our method demonstrates that integrating wavelet transformation into Mamba enhances the local structure awareness of visual inputs and better captures long-range relations of frequencies by disentangling them into wavelet subbands, representing both low- and high-frequency components. These wavelet-based outputs are then processed and seamlessly fused with the original Mamba outputs through a cross-attention fusion layer, combining both spatial and frequency information to optimize the order awareness of state-space models which is essential for the details and overall quality of image generation. Besides, we introduce a globally-shared transformer to supercharge the performance of Mamba, harnessing its exceptional power to capture global relationships. Through extensive experiments on standard benchmarks, our method demonstrates superior results compared to DiT and DIFFUSSM, achieving faster training convergence and delivering high-quality outputs. The codes and pretrained models are released at https://github.com/VinAIResearch/DiMSUM.git.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。