Reproduce “LaMamba-Diff: Linear-Time High-Fidelity Diffusion Models Based on Local Attention and Mamba” on OpenResearch

About this paper

Yunxiang Fu, Chaoqi Chen, Yizhou Yu

Researchers at The University of Hong Kong developed LaMamba-Diff, a novel U-Net architecture for diffusion models that integrates linear-time State Space Models with local attention, enabling high-fidelity image generation. This approach significantly reduces computational costs while achieving competitive or superior performance, exemplified by an FID of 2.04 on ImageNet 256x256 class-conditional generation using 57.6% fewer GFLOPs compared to DiT-XL/2.

Key insights

• A hybrid architecture that effectively combines the linear-time global context modeling of VSSM with linear-time local attention can achieve both computational efficiency and high-fidelity detail preservation in diffusion models.
• The U-Net structure with additive skip connections, when integrated with LaMamba blocks, significantly enhances performance and efficiency compared to isotropic architectures.
• Preserving fine-grained local information through 1x1 patch processing and enabling cross-window information propagation via alternating shifts in local attention are crucial for achieving state-of-the-art image generation quality efficiently.

Summary by alphaXiv.