DiMSUM: Diffusion Mamba -- A Scalable and Unified Spatial-Frequency Method for Image Generation

TL;DR

DiMSUM introduces a scalable diffusion architecture that merges spatial processing in Mamba with wavelet-based frequency analysis to better capture local structure and long-range frequency relations in images. By embedding frequency subbands via Wavelet Mamba and fusing them with spatial features through a cross-attention layer, plus augmenting with a globally-shared transformer, the method reduces inductive bias from fixed scanning orders while preserving global context. Empirical results on CelebA-HQ, LSUN Church, and ImageNet demonstrate state-of-the-art FID and recall with faster convergence and competitive parameter counts. This hybrid spatial-frequency diffusion approach advances image generation by leveraging multi-scale frequency information and efficient long-range modeling, with potential for broader diffusion and multimodal tasks.

Abstract

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.

Figures (24)

• Figure 1: Overview of DiMSUM architecture.
• Figure 2: Illustration of Wavelet Mamba (Best view in color). For illustration purpose, we plot wavelet representations of an input image but our real process is performed on encoded features of the input. Giving an image of size $(8, 8)$, for example, it is first decomposed to four wavelet subbands of size $(4, 4)$ where each is further transformed to 2nd-level subbands of size $(2, 2)$. Green dots indicate pixel points within each wavelet subband and a window of size $2 \times 2$ is used to perform scanning across multiple wavelet subbands like the CNN kernel.
• Figure 3: Comparison of window scanning on image and wavelet space. For illustration, one-level wavelet transformation is applied and each subband is half the resoluton of original image.
• Figure 4: Size $256 \times 256$. † is our reproduced result and ‡ is adopted results from LFM paper.
• Figure 5: Qualitative results.