ECCV 2024 | ZigMa：Mamba遇见扩散模型！强强联合！

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

点击进入—>【Mamba/多模态/扩散】交流群

添加微信：CVer5555，小助手会拉你进群！

扫描下方二维码，加入CVer学术星球！可以获得最新顶会/顶刊上的论文idea和CV从入门到精通资料，及最前沿应用！发论文/搞科研/涨薪，强烈推荐！

作者：没有星星的夜（已授权转载，不得二次转发）

链接：https://zhuanlan.zhihu.com/p/688473044

ZigMa: Zigzag Mamba Diffusion Model

主页：https://taohu.me/zigma/

代码：https://github.com/CompVis/zigma

论文：https://arxiv.org/abs/2403.13802

摘要

扩散模型长期以来一直受到可伸缩性和二次复杂性问题的困扰，特别是在基于变压器结构的情况下。在这项研究中，我们旨在利用称为Mamba的状态空间模型的长序列建模能力，将其适用性扩展到视觉数据生成领域。首先，我们发现了当前大多数基于Mamba的视觉方法中存在的一个关键疏忽，即Mamba扫描方案中缺乏对空间连续性的考虑。其次，基于这一认识，我们引入了一种简单的、即插即用、零参数的方法，名为Zigzag Mamba，它优于基于Mamba的基线，并且相对于基于变压器的基线表现出更好的速度和内存利用率。最后，我们将Zigzag Mamba与随机插值框架集成，以探究模型在大分辨率视觉数据集上的可伸缩性，例如FacesHQ 1024×1024、UCF101、MultiModal-CelebA-HQ和MS COCO 256×256。

绪论

扩散模型在各种应用中展示了显著的进展，包括图像处理[68]、视频分析[39]、点云处理[88]和人体姿势估计[28]。其中许多模型都建立在潜在扩散模型（LDM）[68]之上，通常基于UNet骨干。

然而，可伸缩性仍然是LDM的一个重大挑战[44]。最近，基于变压器的结构因其可伸缩性[8,66]和多模态训练的有效性而受到欢迎[9]。值得注意的是，基于变压器的结构DiT 甚至通过提高由OpenAI开发的高保真度视频生成模型SORA [65]。尽管通过诸如窗口化[60]、滑动[11]、稀疏化[16,49]和哈希化[17,75]等技术来缓解注意机制的二次复杂性的努力，但这仍然是扩散模型的瓶颈。

另一方面，状态空间模型[30,31,34]已经展现出与基于变压器的方法竞争的长序列建模的巨大潜力。已经提出了几种方法[26,29,31,70]来增强状态空间模型的鲁棒性[93]、可伸缩性[29]和效率[31,32]。其中，一种名为Mamba的方法[29]旨在通过工作效率高的并行扫描和其他数据相关的创新来缓解这些问题。

然而，Mamba的优势在于1D序列建模，将其扩展到2D图像是一个具有挑战性的问题。先前的工作[59,98]直接提出通过计算层次结构（例如按行和按列的主要顺序）来展开2D标记，但这种方法忽略了空间连续性，如图1所示。

图1：动机。我们的Zigzag Mamba方法通过以启发式的方式安排和重新安排Mamba的扫描路径来提高网络的位置感知能力。

其他工作[56,62]考虑单个Mamba块中的各种方向，但这引入了额外的参数和GPU内存负担。在本文中，我们旨在强调Mamba中空间连续性的重要性，并提出几种直观简单的方法，通过将基于连续性的归纳偏差纳入图像中，使Mamba块能够应用于2D图像。我们还通过在3D序列上进行时空因子分解将这些方法推广到3D。最后，随机插值提供了一个更加广义的框架，可以统一各种生成模型，包括归一化流[14]、扩散模型[38,71,73]、流匹配[4,54,58]和薛定谔桥[55]。先前的一些工作[63]探索了随机插值在相对较小的分辨率上的应用，例如256×256、512×512。在这项工作中，我们旨在在更复杂的场景中探索它，例如1024×1024的分辨率甚至是视频。总之，我们的贡献如下：

首先，我们确定了从1D序列建模到2D图像和3D视频建模的普遍化Mamba块中的空间连续性的关键问题。基于这一认识，我们提出了一种简单的、即插即用、零参数的范例，名为Zigzag Mamba（ZigMa），它利用空间连续性最大程度地纳入了从视觉数据中的归纳偏差。

其次，我们通过将空间和时间序列进行因式分解来将方法从2D扩展到3D，以优化性能。

其次，我们在扩散模型的范围内提供了对Mamba块的全面分析。

最后，我们证明了我们设计的Zigzag Mamba优于相关的基于Mamba的基线，代表了对大规模图像数据（1024×1024）和视频进行随机插值探索的首次尝试。

方法

在本节中，我们首先提供有关状态空间模型的背景信息，特别关注一种称为Mamba的特殊情况[29]。然后，我们突出显示了Mamba框架中空间连续性的关键问题，并基于这一认识提出了Zigzag Mamba。这种增强旨在通过纳入2D数据固有的连续性归纳偏差来提高2D数据建模的效率。此外，我们设计了一个基本的交叉注意力块，建立在Mamba块之上，以实现文本调节。随后，我们建议通过将模型因式分解为空间和时间维度，从而促进建模过程，将这种方法扩展到3D视频数据。最后，我们介绍了用于训练和采样的随机插值的理论方面，这些理论支撑了我们的网络架构。

背景：状态空间模型

状态空间模型（SSMs）[30, 31, 34]在理论上和实践中已被证明可以处理长程依赖关系[32]，并且在序列长度方面呈线性缩放。在它们的一般形式中，线性状态空间模型可以写成如下形式：

将一个一维输入序列 u(t)∈R 映射到一个一维输出序列 y(t)∈R ，通过一个隐式的 N 维潜在状态序列 x(t)∈Rn 。具体来说，深度状态空间模型（Deep SSMs）旨在在神经序列建模架构中使用这种简单模型的堆叠，其中每一层的参数 A,B,C 和 D 可以通过梯度下降学习得到。

最近，Mamba [29]通过放宽对SSM参数的时不变性约束，同时保持计算效率，大大提高了SSMs的灵活性。通过采用高效的并行扫描，Mamba减轻了循环的顺序性质所带来的影响，而融合GPU操作则消除了需要形成扩展状态的要求。在本文中，我们专注于探索Mamba在扩散模型中的扫描方案，以最大程度地利用来自多维视觉数据的归纳偏差。

图2：ZigMa。我们的骨干网络结构分为 L 层，反映了 DiT [66] 的风格。我们使用单扫描的 Mamba 块作为不同补丁之间的主要推理模块。为了确保网络具有位置感知能力，我们设计了基于单扫描 Mamba 的排列-重新排列方案。不同的层遵循一对唯一的重新排列操作 Ω 和反向重新排列 Ω¯，优化了方法的位置感知性。

扩散主干：Zigzag Mamba

DiT风格网络

我们选择使用AdaLN提出的ViT框架[66]，而不是侧重于跳层的U-ViT结构[8]，因为ViT在文献中已经被证明是一种可伸缩的结构[9,15,65]。考虑到上述观点，它指导了我们在图4中描述的Mamba网络设计。这种设计的核心组件是Zigzag扫描，将在下文中进行解释。

图4：我们的Zigzag Mamba块的详细信息。Mamba扫描的详细信息如图2所示。条件可以包括一个时间步和一个文本提示。这些被馈送到一个MLP中，该MLP分别调制了Mamba扫描用于长序列建模和交叉注意力用于多模态推理。

Mamba中的Zigzag扫描

先前的研究[82,91]在SSM框架内使用了双向扫描。这种方法已经扩展到包括额外的扫描方向[56,59,92]，以考虑2D图像数据的特性。这些方法将图像补丁沿着四个方向展开，导致四个不同的序列。随后，每个序列都通过每个SSM进行处理。然而，由于每个方向可能具有不同的SSM参数（A、B、C 和 D），增加方向的数量可能会导致内存问题。在这项工作中，我们研究了将Mamba的复杂性摊销到网络的每一层的潜力。

我们的方法围绕着在将它们输入前向扫描块之前对令牌进行重新排列的概念。对于来自第i层的给定输入特征zi，经过重新排列后前向扫描块的输出特征zi+1可以表示为：

Ωi 表示第i层的一维排列，它通过 Ωi 重新排列补丁令牌的顺序，并且 Ωi 和 Ω¯i 表示反向操作。这确保了 zi 和 zi+1 保持了原始图像令牌的样本顺序。现在我们探讨 Ωi 操作的设计，考虑到来自2D图像的额外归纳偏差。我们提出一个关键性质：空间连续性。

关于空间连续性，当前在图像中的Mamba的创新通常直接按照计算机层次结构，如行和列的主要顺序，挤压2D补丁标记。然而，这种方法可能不是最佳的，以将归纳偏差与相邻的令牌相结合，如图3所示。

图3：二维图像扫描。我们的Mamba扫描设计基于子图(a)中显示的扫描方案。从中，我们开发了一个在子图(b)中显示的Zigzag扫描方案，以增强补丁的连续性，从而最大限度地发挥Mamba块的潜力。由于这些连续扫描有几种可能的排列方式，我们在子图(c)中列出了八种最常见的Zigzag扫描方案。

为了解决这个问题，我们引入了一种新的扫描方案，设计在扫描过程中保持空间连续性。此外，我们考虑填充空间，这意味着对于一个大小为 N × N 的补丁，一维连续扫描方案的长度应该是 N2。这有助于有效地将令牌纳入Mamba块内，以最大限度地发挥长序列建模的潜力。为了实现上述属性，我们启发式地设计了八种可能的填充空间连续方案，表示为 Sj（其中 j ∈ [0, 7]），如图3所示。虽然可能存在其他可想象的方案，但为了简单起见，我们将使用限制在这八种方案中。因此，每个层的方案可以表示为 Ωi = S{i%8}，其中 % 表示取模运算符。

在Zigzag Mamba上部署文本条件

虽然Mamba提供了高效的长序列建模优势，但这是以注意机制为代价的。因此，对于基于Mamba的扩散模型，对于整合文本条件的探索是有限的。为了填补这一空白，我们提出了一个简单的跳层跨注意力块，建立在Mamba块之上，如图4所示。

图4：我们的Zigzag Mamba块的详细信息。Mamba扫描的详细信息如图2所示。条件可以包括一个时间步和一个文本提示。这些被馈送到一个MLP中，该MLP分别调制了Mamba扫描用于长序列建模和交叉注意力用于多模态推理。

这种设计不仅可以实现长序列建模，还可以促进多令牌调节，例如文本条件。此外，它还具有提供可解释性的潜力，因为跨注意力已经被应用在扩散模型中。

将设计推广到三维视频，通过因子分解空间和时间信息。在前面的章节中，我们的重点是在二维空间的Mamba上，我们设计了几种空间连续、填充空间的二维扫描方案。在本节中，我们旨在利用这一经验来帮助设计相应的三维视频处理机制。我们开始我们的设计过程，从图5中所示的传统方向性Mamba中推导出来。给定一个视频特征输入z∈RB×T×C×W×H ，我们提出了三种3D视频生成的Video Mamba Block变体。

图5：三维视频扫描。(a) 我们用扫描的双向Mamba进行说明，其中空间和时间信息被视为一组具有计算机层次结构顺序的令牌。(b) 对于三维Zigzag扫描，我们旨在通过采用空间连续扫描方案并采用图3所示的最佳Zigzag扫描解决方案来最大化Mamba的潜力。(c) 我们进一步将空间和时间信息之间的推理分开

(a) 扫描扫描：在这种方法中，我们直接将3D特征z展平，而不考虑空间或时间的连续性。值得注意的是，展平过程遵循计算机层次结构顺序，意味着在展平表示中没有保留连续性。

(b) 3D Zigzag：与前述子节中的2D Zigzag相比，我们采用类似的设计将其推广到3D Zigzag，以同时保持2D和3D中的连续性。潜在地，这个方案的复杂性要高得多。我们启发式地列出了8种方案。然而，我们凭经验发现，这种方案会导致次优的优化。

(c) 分解的3D Zigzag = 2D Zigzag + 1D扫描：为了解决次优的优化问题，我们建议将空间和时间的相关性因子分解为单独的Mamba块。它们的应用顺序可以根据需要进行调整，例如，“sstt”或“ststst”，其中“s”表示空间Zigzag Mamba，“t”表示时间Zigzag Mamba。

计算分析

对于一个视觉序列 T∈R1×M×D ，全局自注意力和k方向mamba以及我们的zigzag mamba的计算复杂度如下：

其中自注意力与序列长度M呈二次复杂度，而Mamba呈线性复杂度（N是一个固定参数，默认设置为16）。在这里，k表示单个Mamba块中的扫描方向数量。因此，k-mamba和zigzag相对于自注意力都具有线性复杂度。此外，我们的zigzag方法可以消除k系列，进一步降低总体复杂度。

完成改进的视觉归纳偏差集成Zigzag Mamba网络的设计后，我们继续将其与一个新的扩散框架结合起来，如下所示。

图5

扩散框架：随机插值器

基于向量v和分数s进行采样。xt 的时间相关概率分布 pt(x) 也与逆时SDE的分布重合：

其中 W¯t 是逆时间维纳过程，wt > 0是任意时间相关扩散系数，s(x, t) = ∇ log pt(x)是得分，而v(x, t)由条件期望给出。

其中αt是t的递减函数，而σt是t的递增函数。这里，α˙ t和σ˙ t分别表示αt和σt的时间导数。

只要我们能够估计速度v(x, t)和/或得分s(x, t)场，我们就可以利用它们进行采样过程，无论是通过概率流ODE[73]还是逆时间SDE(7)。从 XT=ε∼N(0,I) 向时间反向解决逆SDE(7)使我们能够从近似数据分布 p0(x)∼p(x) 中生成样本。在采样过程中，我们可以通过直接从ODE或SDE中进行采样来在采样速度和保真度之间平衡。如果我们选择进行ODE采样，我们可以通过将噪声项s设为零来简单实现。

在[3]中，它表明实际上需要估计sθ(x, t)和vθ(x, t)中的一个。这直接遵循了约束

这可以用来重新表示得分 s(x, t)，如下所示：

因此，v(x, t)和s(x, t)可以相互转换。我们将说明如何计算它们如下所示。

估计得分s和速度v。已经在基于得分的扩散模型中表明，得分可以通过使用损失参数化估计为sθ(x, t)。

公式（11）

类似地，速度v(x, t)可以通过损失参数化地估计为vθ(x, t)。

公式（12）

其中θ代表我们在前面部分描述的Zigzag Mamba网络，我们采用线性路径进行训练，因为它简单且轨迹相对直接：

公式（13）

我们注意到任何时间相关的权重都可以包含在(11)和(12)中的积分中。当T变大时，这些权重因子在基于得分的模型中起着至关重要的作用。因此，它们提供了一个考虑到时间相关权重和随机性的一般形式。

结论

本文提出了在随机插值器框架内开发的Zigzag Mamba扩散模型。我们最初的重点是解决空间连续性的关键问题。然后，我们设计了一个Zigzag Mamba块，以更好地利用2D图像中的归纳偏差。此外，我们将3D Mamba分解为2D和1D Zigzag Mamba以促进优化。我们经验性地设计了各种消融研究来检验不同的因素。这种方法允许对随机插值器理论进行更深入的探索。我们希望我们的努力能够激发对Mamba网络设计的进一步探索。在未来，我们计划研究这种网络设计用于判别任务。

限制和未来工作

我们的方法完全依赖于具有DiT风格布局和调节方式的Mamba块。然而，我们工作的一个潜在局限性是，我们无法穷尽地列出在给定特定全局补丁大小的情况下所有可能的空间连续Zigzag扫描方案。目前，我们凭经验设置这些扫描方案，这可能导致性能次优。此外，由于GPU资源限制，我们无法探索更长的训练时间，尽管我们预计结果可能类似。

未来的工作中，我们的目标是深入探讨Zigzag Mamba的各种应用，利用其在长序列建模方面的可扩展性。这种探索可能会导致更好地利用Mamba框架在不同领域和应用中的应用。

何恺明在MIT授课的课件PPT下载

在CVer公众号后台回复：何恺明，即可下载本课程的所有566页课件PPT！赶紧学起来！

CVPR 2024 论文和代码下载

在CVer公众号后台回复：CVPR2024，即可下载CVPR 2024论文和代码开源的论文合集

Mamba、多模态和扩散模型交流群成立

扫描下方二维码，或者添加微信：CVer5555，即可添加CVer小助手微信，便可申请加入CVer-Mamba、多模态学习或者扩散模型微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要备注：研究方向+地点+学校/公司+昵称（如Mamba、多模态学习或者扩散模型+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲扫码或加微信号: CVer5555，进交流群
CVer计算机视觉（知识星球）来了！想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料，欢迎扫描下方二维码，加入CVer计算机视觉（知识星球），已汇集近万人！

▲扫码加入星球学习

▲点击上方卡片，关注CVer公众号
整理不易，请赞和在看