U-Shape Mamba: State Space Model for faster diffusion
This work addresses the computational bottleneck in diffusion models for image generation, making high-quality synthesis more accessible, though it is incremental as it builds on existing Mamba-based methods.
The paper tackled the high computational cost of diffusion models for image generation by proposing U-Shape Mamba (USM), which uses Mamba-based layers in a U-Net structure to reduce computational overhead while maintaining quality, achieving one-third the GFlops and improving FID scores by up to 15.3 points on datasets like AFHQ.
Diffusion models have become the most popular approach for high-quality image generation, but their high computational cost still remains a significant challenge. To address this problem, we propose U-Shape Mamba (USM), a novel diffusion model that leverages Mamba-based layers within a U-Net-like hierarchical structure. By progressively reducing sequence length in the encoder and restoring it in the decoder through Mamba blocks, USM significantly lowers computational overhead while maintaining strong generative capabilities. Experimental results against Zigma, which is currently the most efficient Mamba-based diffusion model, demonstrate that USM achieves one-third the GFlops, requires less memory and is faster, while outperforming Zigma in image quality. Frechet Inception Distance (FID) is improved by 15.3, 0.84 and 2.7 points on AFHQ, CelebAHQ and COCO datasets, respectively. These findings highlight USM as a highly efficient and scalable solution for diffusion-based generative models, making high-quality image synthesis more accessible to the research community while reducing computational costs.