[论文解读] ZigMa: A DiT-style Zigzag Mamba Diffusion Model

Vincent Tao Hu, Stefan Andreas Baumann+5 more|arXiv (Cornell University)|Mar 20, 2024

Cellular Automata and Applications被引用 4

一句话总结

ZigMa 提出了一种基于 DiT 的扩散模型，用一种新颖的之字形 Mamba 模块替代注意力机制，以提升 2D 和 3D 视觉生成中的空间连续性。通过将 Mamba 的扫描路径重新排列为之字形，该方法在 1024×1024 FacesHQ 和 MS-COCO 数据集上实现了最先进（SOTA）的 FID 分数，且在参数量零增加的前提下，展现出更优的速度和内存效率。

ABSTRACT

The diffusion model has long been plagued by scalability and quadratic complexity issues, especially within transformer-based structures. In this study, we aim to leverage the long sequence modeling capability of a State-Space Model called Mamba to extend its applicability to visual data generation. Firstly, we identify a critical oversight in most current Mamba-based vision methods, namely the lack of consideration for spatial continuity in the scan scheme of Mamba. Secondly, building upon this insight, we introduce a simple, plug-and-play, zero-parameter method named Zigzag Mamba, which outperforms Mamba-based baselines and demonstrates improved speed and memory utilization compared to transformer-based baselines. Lastly, we integrate Zigzag Mamba with the Stochastic Interpolant framework to investigate the scalability of the model on large-resolution visual datasets, such as FacesHQ $1024 imes 1024$ and UCF101, MultiModal-CelebA-HQ, and MS COCO $256 imes 256$ . Code will be released at https://taohu.me/zigma/

研究动机与目标

• 解决现有基于 Mamba 的视觉方法中缺乏空间连续性的问题，这些方法通常采用线性行序或列序扫描顺序。
• 通过启发式扫描路径引入空间连续性，提升 Mamba 模块在 2D 和 3D 视觉数据上的归纳偏置。
• 在随机插值框架内，利用基于 Mamba 的架构实现可扩展、高效且高保真度的图像与视频生成。
• 在 FacesHQ 1024×1024 和 UCF101 等大规模数据集上，证明基于 Mamba 的扩散模型的可行性与优越性。
• 探索随机插值框架在 256×256 分辨率以上的可扩展性，包括 1024×1024 分辨率和视频生成任务。

提出的方法

• 提出一种之字形 Mamba 模块，将 Mamba 的扫描路径从线性改为之字形，以保留空间局部性并增强位置感知能力。
• 引入一种即插即用、零参数的技巧，在不修改 Mamba 架构的前提下，将之字形扫描模式应用于所有层。
• 通过将空间与时间序列分解，将 2D 之字形扫描扩展至 3D，采用空间优先扫描策略以利用时间冗余性。
• 将之字形 Mamba 模块集成到 DiT 风格的扩散模型主干网络中，用 Mamba 模块替代注意力机制。
• 使用随机插值框架进行模型训练，实现扩散、流匹配与归一化流目标的统一训练。
• 采用潜在扩散设置，使用 2×2 分块潜变量，实现 1024×1024 分辨率下 4,096 个 token 的高分辨率生成。

Figure 1 : Motivation. Our Zigzag Mamba method improves the network’s position-awareness by arranging and rearranging the scan path of Mamba in a heuristic manner.

实验结果

研究问题

• RQ1 标准 Mamba 扫描顺序中缺乏空间连续性，如何影响视觉生成性能？
• RQ2 一种简单且无参数的扫描重排策略（如之字形扫描）是否能显著提升基于 Mamba 的视觉模型性能？
• RQ3 之字形 Mamba 在高分辨率图像与视频生成任务（如 1024×1024 FacesHQ 和 UCF101）上的可扩展性如何？
• RQ4 随机插值框架是否能有效支持大规模基于 Mamba 的扩散模型训练？
• RQ5 将 3D 序列分解为 2D 与 1D 之字形扫描，对视频生成质量与效率有何影响？

主要发现

• 在使用 32 张 A100 GPU 的情况下，之字形 Mamba 在 1024×1024 FacesHQ 数据集上实现了 26.6 的 FID 分数，优于双向 Mamba（51.1 FID），并展现出更优的可扩展性。
• 在 MS-COCO 256×256 数据集上，Zigzag-8 变体实现了 41.8 的 FID 分数，显著优于双向 Mamba（60.2 FID）和 Zigzag-1（73.1 FID）。
• 在 UCF101 视频生成任务中，分层 3D 之字形 Mamba 实现了 121.2 的 Frame-FID 和 140.1 的 FVD，优于双向 Mamba（146.2 Frame-FID，201.1 FVD），使用 16 张 A100 GPU。
• 在 FacesHQ 1024×1024 和 MultiModal-CelebA 512×512 的可视化结果中，该方法在不同分辨率下均保持了高质量的样本保真度。
• 消融实验表明，同时采用多种扫描策略可有效增强归纳偏置的引入，并带来一致的性能提升。
• 与基于 Transformer 的基线相比，该模型在大分辨率下也展现出更优的内存与速度效率。

Figure 2 : ZigMa. Our backbone is structured in L layers, mirroring the style of DiT [ 65 ] . We use the single-scan Mamba block as the primary reasoning module across different patches. To ensure the network is positionally aware, we’ve designed an arrange-rearrange scheme based on the single-scan