用Mamba结构加速扩散模型,效率提升三倍且画质更好。
U-Shape Mamba: State Space Model for faster diffusion
- 在U-Net结构中嵌入Mamba层,分阶段压缩和恢复序列长度
- 计算量降为原有模型1/3,内存占用更低,速度更快
- 适合追求高效生成的图像合成研究者
扩散模型虽是高质量图像生成主流方法,但计算开销大仍是难题。本文提出U-Shape Mamba(USM),将基于Mamba的层嵌入类U-Net的分层结构中,通过编码器逐步缩短序列长度、解码器恢复序列,显著降低计算开销,同时保持强大生成能力。与当前最高效的Mamba扩散模型Zigma相比,USM在计算量上减少至1/3(GFlops),内存需求更低,运行更快,且在图像质量上全面超越:在AFHQ、CelebAHQ和COCO数据集上的弗雷歇起始距离(FID)分别提升15.3、0.84和2.7点。结果表明,USM是一种高效且可扩展的生成模型方案,使高质量图像合成更易实现且成本更低。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。