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

TL;DR

LaMamba-Diff introduces a linear-time diffusion model that combines local attention and Mamba to efficiently capture both global and local features, outperforming existing models in scalability and performance on ImageNet.

Contribution

The paper proposes LaMamba blocks that integrate self-attention and Mamba, enabling linear complexity diffusion models with high fidelity for visual generation.

Findings

Outperforms DiT models across various scales on ImageNet.

Reduces GFLOPs by up to 62% compared to DiT-XL/2.

Achieves superior performance with fewer parameters.

Abstract

Recent Transformer-based diffusion models have shown remarkable performance, largely attributed to the ability of the self-attention mechanism to accurately capture both global and local contexts by computing all-pair interactions among input tokens. However, their quadratic complexity poses significant computational challenges for long-sequence inputs. Conversely, a recent state space model called Mamba offers linear complexity by compressing a filtered global context into a hidden state. Despite its efficiency, compression inevitably leads to information loss of fine-grained local dependencies among tokens, which are crucial for effective visual generative modeling. Motivated by these observations, we introduce Local Attentional Mamba (LaMamba) blocks that combine the strengths of self-attention and Mamba, capturing both global contexts and local details with linear complexity. Leveraging the efficient U-Net architecture, our model exhibits exceptional scalability and surpasses the performance of DiT across various model scales on ImageNet at 256x256 resolution, all while utilizing substantially fewer GFLOPs and a comparable number of parameters. Compared to state-of-the-art diffusion models on ImageNet 256x256 and 512x512, our largest model presents notable advantages, such as a reduction of up to 62% GFLOPs compared to DiT-XL/2, while achieving superior performance with comparable or fewer parameters. Our code is available at https://github.com/yunxiangfu2001/LaMamba-Diff.