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.

翻译：近年来，基于Transformer的扩散模型展现出卓越的性能，这主要归功于自注意力机制通过计算输入标记之间的全对交互，能够准确捕捉全局与局部上下文。然而，其二次复杂度对长序列输入构成了显著的计算挑战。相反，一种称为Mamba的最新状态空间模型通过将经过滤波的全局上下文压缩至隐藏状态，实现了线性复杂度。尽管效率很高，但压缩不可避免地导致标记间细粒度局部依赖关系的信息丢失，而这对于有效的视觉生成建模至关重要。基于这些观察，我们引入了局部注意力Mamba（LaMamba）模块，它结合了自注意力与Mamba的优势，以线性复杂度同时捕捉全局上下文与局部细节。利用高效的U-Net架构，我们的模型展现出卓越的可扩展性，并在ImageNet 256×256分辨率下，以显著更少的GFLOPs和相当的参数量，超越了不同模型规模的DiT性能。与ImageNet 256×256和512×512上的最先进扩散模型相比，我们最大的模型展现出显著优势，例如与DiT-XL/2相比，GFLOPs减少高达62%，同时以相当或更少的参数实现了更优的性能。