《Mamba-ST：面向高效风格迁移的状态空间模型》学习笔记

最新推荐文章于 2025-05-01 16:54:33 发布

种花生的图图

最新推荐文章于 2025-05-01 16:54:33 发布

阅读量815

点赞数 24

分类专栏：生成式模型文章标签：学习笔记深度学习人工智能机器学习 mamba

本文链接：https://blog.csdn.net/weixin_56848903/article/details/145677503

版权

生成式模型专栏收录该内容

12 篇文章

订阅专栏

paper：2409.10385

GitHub：FilippoBotti/MambaST

会议：WACV

摘要

风格迁移的目标是：给定内容图像与风格源图像，生成一张保留原内容但具有风格源艺术表现的新图像。当前主流架构多采用Transformer或扩散模型，但存在高计算负担问题。具体而言，Transformer的自注意力与跨注意力层内存占用大，扩散模型推理时间长。为克服这些问题，本文提出一种基于新兴状态空间模型（SSM）的创新架构——Mamba-ST。我们改造Mamba的线性方程以模拟跨注意力层的行为（将两个独立嵌入合并为单一输出），同时大幅降低内存占用与时间复杂度。通过修改Mamba内部方程，使其能接收并融合双数据流输入。据我们所知，这是首次在不依赖跨注意力层或自定义归一化模块的情况下，将SSM方程适配视觉风格迁移任务的尝试。大量实验证明，相较于Transformer与扩散模型，本方法在风格迁移效果与效率上均具优势，ArtFID与FID指标均有提升。

1、介绍

风格迁移是一种深度学习技术，旨在生成一张新图像，该图像既包含给定图像（内容图像）的内容（如物体、布局），又具有另一图像（风格图像）的风格（如色彩或纹理结构）。风格迁移已被广泛研究[6,11,24,25,27,29,33,42]，现有多种模型能在质量与内容-风格一致性上取得良好效果。近期，通过用文本描述替代风格图像，风格迁移被扩展至基于文本的模型[25,27]。本文不涉及此类模型的比较，因其与标准风格迁移架构差异显著。主流方法遵循编码器-解码器流程，在中间层进行风格迁移[24,29,33,42]，但这些模型难以建立内容与风格间的关联，导致图像细节质量欠佳。为此，近期架构（如[11]）利用Transformer的跨注意力机制提升结果质量，但其内存需求较高。另一方面，部分研究[6,43,49]尝试以扩散模型为骨干进行风格迁移，但未能达到与Transformer相当的质量，或如[6]所述，推理时间较长。

相比之下，状态空间模型（SSM）[18,19]近期在长序列建模与视觉任务中展现出与Transformer相当的性能[20,23,50]。其中，Mamba[17]被证明是Transformer的高效替代方案，且内存需求显著降低。此外，Mamba的复杂度与序列长度呈线性关系（而非平方关系），推理速度远超扩散模型。Mamba 最初设计用于一维序列（如语句中的单词），但通过VMamba[32]已适配至二维视觉数据。近期，Wang等人[41]在文本风格迁移架构中使用Mamba模块，其性能与Transformer架构相当，但需额外引入自适应层归一化（AdaLN）以融合风格与内容。

本文提出一种改进Mamba内部方程的方法，实现风格迁移的新架构Mamba-ST（见图1）。

其核心是Mamba-ST解码器（MSTD）模块，可融合来自不同图像的内容与风格信息。具体而言，内容与风格均被建模为patch嵌入序列，输入至MSTD中。为模拟跨注意力层的功能，我们修改Mamba的内部矩阵，同时保留SSM的核心特性。该方案无需AdaLN等辅助模块即可实现内容-风格交互。本文贡献如下：

SSM内的类跨注意力方法：通过调整内部矩阵的数学公式，在无需AdaLN等额外层的情况下实现跨模态交互。
新型视觉Mamba架构Mamba-ST：风格迁移效果与Transformer及扩散模型相当。
高效推理：内存占用显著低于Transformer，推理速度远超扩散模型。

2、相关工作

2.1 风格迁移

风格迁移问题在文献中已被广泛研究[2,4,12,22,24,28,31]。Gatys等人[15]提出的首次尝试表明，通过求解优化问题合并CNN提取的内容和风格特征是可行的。随后，自适应实例归一化（AdaIN）[24]的引入允许通过调整内容图像的均值和方差并与风格图像对齐来执行任意风格迁移。得益于其高效性，AdaIN成为一种非常流行的架构。后来，Transformer的出现[31,34]展示了自注意力机制如何通过发现内容与风格之间更强的关联性来提升结果质量。随后，Deng等人[11]提出了一种完全基于Transformer的架构StyTr2，该架构结合新的内容感知位置编码方案，在风格迁移任务上超越了当时的最先进方法。尽管这些方法在质量方面达到了较高水平，但它们严重依赖Transformer，导致计算复杂度与图像尺寸的平方成正比，因此仅能处理小尺寸图像。

最近，基于扩散模型的风格迁移方法[6,43,49]展示了如何利用扩散模型的生成能力进行风格迁移。例如：InST[49]通过基于文本的反演方法捕获风格信息并进行迁移。StyleDiffusion[43]引入了基于CLIP的风格解缠损失，以在CLIP图像空间中分离风格与内容。

尽管这些架构生成的图像质量较高，但内容与风格未能完美融合，其结果仍不及基于Transformer的模型。最后，StyleID[6]提出了一种新的风格迁移方法，该方法利用Stable Diffusion 1.4[35]的知识，无需监督或优化。StyleID仅需在自注意力层中将内容的键（Key）和值（Value）替换为风格的对应项，但尽管其结果优于先前方法，其推理时间仍高于基于Transformer的架构。为解决上述内存占用与速度问题，同时保持风格迁移的质量与一致性，我们提出了一种全新的完全基于Mamba的架构。

2.2 状态空间模型

尽管基于Transformer的架构在视觉任务中具有显著优势[3,14,37,46,47]，但其平方复杂度与高内存需求仍是关键问题。为此，近期许多研究尝试解决此问题[5,7,8,13,40]。特别是受控制系统理论启发的状态空间模型（SSM）[18,19]被引入深度学习领域，以利用其线性复杂度优势[23,32,39,50]。结构化状态空间模型（S4）[18]提出了一种新的SSM参数化方法，以实现训练时的并行化优势与推理时的高速度。Mamba[17]作为一种改进的SSM被提出，其主要贡献在于使SSM参数具有输入依赖性。由于其在内存占用、时间复杂度和结果质量上的优势，Mamba已被广泛应用于从自然语言处理（NLP）[39]到视觉领域[32,50]的任务中，包括超分辨率[20,36]甚至扩散模型[23]。近期，Wang等人[41]首次尝试将Mamba用于文本驱动风格迁移，但未能充分利用SSM内部矩阵的优势，仍需依赖自适应层归一化（AdaLN）等定制归一化层进行风格迁移。在此工作中，Mamba并非直接用于融合内容与风格，而是主要作为特征提取器以利用其快速推理速度。值得注意的是，AdaLN与Mamba结合的模型泛化能力有限，需针对每个新的（文本，图像）对单独训练。因此，Mamba仅用于加速图像块间的交互，而风格融合过程仍由AdaLN层主导。相比之下，我们通过改造Mamba的内部方程，提出了一种完全基于Mamba的架构，无需AdaLN或跨注意力层即可融合内容与风格。此外，我们的方法能够学习风格迁移过程，训练完成后可直接应用于任意图像-风格对而无需重新训练。

3、方法

3.1 Mamba理论基础

状态空间模型（如S4[18]）学习将一维输入序列 $x(t)\in R$ 映射到一维输出序列 $y(t) \in R$ ，同时维护内部状态空间 $h(t) \in R^N$ （N为状态尺寸）。与Transformer等序列学习方法不同，SSM的复杂度与序列长度呈线性关系。SSM由线性常微分方程（ODE）系统描述：

$h'(t) = Ah(t) +Bx(t)$

$y(t) = Ch(t) + Dx(t)$

其中矩阵 $A \in R^{N \times N}$ 、 $B \in R^{N \times 1}$ 、 $C \in R^{1 \times N}$ 和标量 $D \in R$ 为可学习参数。为适配深度学习架构，需对系统进行离散化。本文采用零阶保持（ZOH）规则，其中 Δ 为步长参数。通过将离散化矩阵 $\bar{A}, \bar{B}$ 表示为RNN形式，方程变为：

$h_{k} = \overline{\mathbf{A}}h_{k-1} + \overline{\mathbf{B}}x_{k}$

$y_{k} = \mathbf{C}h_{k} + Dx_{k}$

循环网络虽能通过有限状态压缩上下文信息，但其核心局限在于状态对上下文信息的压缩能力（例如理解句子中的关键词语而忽略次要信息）。为此，Gu等人提出Mamba[17]，通过使矩阵B,C和Δ具有输入依赖性：

$B =Lin_B(x),C = Lin_C(x),\Delta = Lin_\Delta (x)$

其中 $Lin_{*}$ 为全连接层。该改进结合高效的选择性扫描算法，在保持计算效率的同时提升状态对上下文信息的记忆与理解能力。

零阶保持（ZOH）规则：信号处理中将离散信号转换为连续信号的核心方法，其规则可概括为：在相邻两个采样点之间，保持前一个采样点的值恒定，形成阶梯状连续信号。

3.2 整体架构

基于Mamba的上下文理解能力与循环网络效率，本文提出用于风格迁移的Mamba-ST架构（图2a）。系统包含三个核心组件：

双Mamba编码器：分别编码内容与风格图像(风格迁移必备)
Mamba风格迁移解码器（MSTD）：融合内容与风格信息
CNN解码器：将解码器输出重构为图像

内容与风格图像被切分为patch，通过PatchEmbed层[44]投影为1D嵌入序列。该层以图像 $I\in R^{C\times H\times W }$ 为输入，生成嵌入序列 $t \in R^D \times (h\cdot w)$ （其中 $h=\frac{H}{p}$ , $w=\frac{W}{p}$ ，p 为patch尺寸，D 为隐藏维度）。两个域编码器分别处理内容与风格嵌入，学习图像的视觉表征。编码后的内容与风格表征输入至MSTD（图2c）进行信息融合，最后通过反patch化块生成 $D \times h \times w$ 特征图，由CNN解码器转换为 $C \times H \times W$ 输出图像。

3.3 Mamba编码器

每个编码器由3个Mamba编码器层构成（图2b），其结构基于 VMamba[32]，但添加了层间跳跃连接以避免梯度消失。

嵌入序列经层归一化后输入基础视觉SSM（Base VSSM）（图2d）。Base VSSM通过将Mamba的S6模块替换为2D-SSM（图2e）实现。

2D-SSM学习矩阵A,B,C和 Δ，建模层的状态与输出（即图像的编码视觉表征）。其核心是适配2D序列的选择性扫描机制[17]（图3）。

我们遵循[32]，使用4种扫描方向以保留空间信息。对每种扫描方向，按式(2)计算状态与输出，通过重排与求和进行融合（算法1）。

Lin(p)：将输入 patch 投影到隐藏空间。

DWConv(x)：逐通道卷积，提取局部空间特征，参数量为标准卷积的 1/C（C为通道数）。

SiLU(x)：激活函数。

$Lin_B/Lin_C/Lin_\Delta$ ：生成输入依赖的动态参数。

$parameter^A$ ：可学习对角矩阵
$SelectiveScan$ ：多方向融合扫描。

LayerNorm：归一化层。

3.4 Mamba-ST解码器

Mamba-ST解码器（MSTD）的任务是将两个Mamba编码器提取的内容与风格视觉表征融合为单一表征，实现风格迁移。其整体结构与编码器类似。

编码器与解码器架构上的区别：

但有两个关键区别：

双输入流：接收内容嵌入（来自内容编码器）和风格嵌入（来自风格编码器），输入改进的ST-VSSM模块（见图2(f))。
ST-2D-SSM模块：替换原2D-SSM，专门设计用于融合内容与风格信息（见图2(g)）。

ST-VSSM 和 Base-VSSM 架构区别：

近期研究[1,9]揭示了Mamba方程与Transformer自注意力的对偶性。具体地，自注意力中的Q,K,V矩阵可映射为：

$Q \approx C,K \approx B,V \approx X$

其中，X是输入序列。基于Mamba与自注意力之间的这种对称性，我们的直觉是通过让A、B和∆矩阵依赖于风格源s，同时让C依赖于内容x，从而模拟交叉注意力机制，如下所示：

$B=Lin_B(s),\Delta =Lin_\Delta (s),C=Lin_C(x)$

类似地，我们决定传递编码后的风格特征，而不是内容，作为ST-2D-SSM的输入序列。更具体地，我们将风格信息融合到状态中，因为状态的方程仅依赖于A和B矩阵：

$h_k=\bar{A}h_{k-1}+\bar{B}s_k$

然后，输出同时依赖于风格和内容，因为C是从内容图像中提取的：

$y_k=Ch_k$

这使得我们能够有效地融合内容和风格信息，从而完成风格迁移。

解码器层内的选择性扫描机制与编码器中的机制相同。此外，为了去除可能破坏风格迁移的内容细节，我们决定对风格嵌入应用随机打乱，如图2（c）所示。通过这种方式，模型的隐藏状态失去了有关风格图像内容的所有信息，只保留了风格信息。在算法2中，我们提供了解码器块的完整算法描述。

3.5 损失函数

我们使用两种感知损失函数来训练我们的模型。内容损失 $L_C$ 主要关注于保持原始图像的内容，而风格损失 $L_S$ 则旨在将源图像的风格转移到目标图像中。

我们通过预训练的VGG19模型来实现这两种损失，参考了[2, 11, 24]的方法。设 $x_c$ 为内容图像， $x_s$ 为风格图像， $x_g$ 为生成图像。给定 $N_l$ 为从VGG19中选取的层数，我们定义 $\varphi_i(x)$ 为从第 i 层提取的特征，输入图像为 x 。内容损失定义为：

$L_C = \frac{1}{N_l} \sum_{i=1}^{N_l} \| \varphi_i(x_g) - \varphi_i(x_c) \|_2$

风格损失定义为：

$L_S = \frac{1}{N_l} \sum_{i=0}^{N_l} \left( \| \mu(\varphi_i(x_g)) - \mu(\varphi_i(x_s)) \|_2 + \| \sigma(\varphi_i(x_g)) - \sigma(\varphi_i(x_s)) \|_2 \right)$

其中 $\mu(\cdot)$ 表示给定特征图的均值， $\sigma(\cdot)$ 表示给定特征图的标准差。

风格损失关注图像中不同区域的统计特征（均值和方差），而不是直接比较图像内容的特征。风格损失的核心思想是，通过调整生成图像的风格特征，使其与目标风格图像的特征相匹配。

此外，我们还使用了两个恒等损失[11, 34]，它们有助于学习更好的内容和风格特征表示。

恒等损失的目的是确保当内容和风格图像相同时，生成图像也能保持与原始图像一致。也就是说，当输入内容图像或风格图像时，模型应能生成一个尽可能接近原始图像的结果。

设 $x^c_g$ 为使用内容图像 $x_c$ 作为内容和风格信息生成的图像， $x^s_g$ 为使用风格图像 $x_s$ 作为内容和风格信息生成的图像，则我们定义：

$L_{id1} = \| x^c_g - x_c \|_2 + \| x^s_g - x_s \|_2$

确保当内容图像被当作风格和内容一起传递时，生成图像应该尽可能接近内容图像。

$L_{id2} = \frac{1}{N_l} \sum_{i=0}^{N_l} \left( \| \varphi(x^c_g) - \varphi(x_c) \|_2 + \| \varphi(x^s_g) - \varphi(x_s) \|_2 \right)$

要求生成图像在某些VGG19层上的特征与内容和风格图像的特征匹配。

我们用于训练模型的最终损失函数为：

$L = \lambda_C L_C + \lambda_S L_S + \lambda_{id1} L_{id1} + \lambda_{id2} L_{id2}$

其中 $\lambda_C = 7, \lambda_S = 10, \lambda_{id1} = 70,\lambda_{id2}=1$ 用于平衡各个损失的大小[11]。

4、实验

4.1 实验细节

数据集：使用 COCO 数据集作为内容数据集，WikiArt 数据集作为风格数据集来训练模型。COCO 数据集包含各种各样的自然场景，适合用于图像内容的训练；而 WikiArt 数据集包含了大量的艺术风格图像，适用于训练风格迁移模型。

超参数设置：模型采用与文献[11]相同的超参数设置，除了学习率（学习率设置为 0.00005，没有使用预热阶段）。

训练：模型总共训练了 160,000 次迭代，使用批量大小为 8 的训练设置，训练设备为单个 NVIDIA L40S GPU。

补充细节：图像的 patch 大小设置为 8。

4.2 对比实验

我们将我们的模型与几个用于图像到图像风格迁移的现有最先进模型进行定性和定量的比较：StyleID[6]、AesPA-Net[22]、StyTr2[11]、AdaAttn[31]、AdaIN[24]。我们故意避免与使用文本描述作为风格条件而不是图像的风格迁移方法进行比较（包括 StyleMamba1 [41]），因为那样的比较不公平。

为了定量评估我们的模型，我们采用了四个主要的指标：ArtFID [45]、FID [21]、LPIPS [48] 和 CFSD [6]。我们没有考虑将内容损失 LC 和风格损失 LS 作为补充评估指标，因为[6]指出，若将这两种损失同时用于训练和评估，会引入评估偏差。

详细来说：

FID [21]：衡量生成图像和风格图像分布之间的相似度。
LPIPS [48]：评估源图像与风格化图像之间的内容保留，从而衡量内容在风格化图像中的保留效果。
ArtFID [45]：一种与人类判断高度相关的指标，将风格迁移和内容保留融合在一起。ArtFID 的计算方式为：(1 + FID) · (1 + LPIPS)。
CFSD (Content Feature Structural Distance) [6]：一种新的度量方式，旨在解决 LPIPS 的不足。LPIPS 使用了一个预训练的 AlexNet 特征空间，但由于 ImageNet 是偏向纹理的，因此 LPIPS 在风格注入时容易受到影响，导致内容保留评估不准确。CFSD 通过计算图像块之间的空间相关性来进行度量。

对于每个模型，我们计算了 800 个生成样本的指标，这些样本是通过随机选择 40 张风格图像和 20 张内容图像生成的[6]。此外，我们还计算了生成这 800 张图像的推理时间（秒）和内存使用量（MebiByte，MiB），并使用批量大小为 1 进行评估。

在表 1 中，我们展示了我们模型与几种最先进模型的定量评估结果。值得注意的是，我们的方法在 ArtFID 上的表现超过了以前的架构，正如前面讨论的那样，ArtFID 与人类判断具有很强的相关性。此外，我们的方法在 FID 上获得了最低的值，表明我们的风格迁移效果优于内容图像。

另一方面，在内容保留的度量（即 LPIPS 和 CFSD）上，我们的模型表现相对较弱，但性能差异非常小。这种折衷表明我们的方法在风格迁移的效果上略有牺牲，但总体上能更好地保持风格。

与 AdaIN 相比，后者由于其轻量级的骨干网络，是时间和内存使用最有效的模型。具体来说，AdaIN 生成 800 张图像（由 20 张内容图像和 40 张风格图像生成）只需 12.26 秒，同时内存消耗最小。然而，即使 AdaIN 在 ArtFID 和 FID 上取得了第二好的成绩，但它在 LPIPS 和 CFSD 上表现较差，意味着它未能很好地保持图像内容。

我们提出的模型在时间和内存使用方面取得了平衡，既保持了较低的内存消耗，又保证了可接受的推理时间，尤其是在与 StyleID 等扩散模型比较时，后者内存消耗大且推理时间长。

定性比较结果见图 4。从图中可以看出，我们能够在当前最先进的模型中实现相当的结果。具体来看：

AdaIN：能够正确地应用风格，但整体内容发生了较大变化，如表 1 中 LPIPS 和 CFSD 的得分所示。
StyTr2、AesPA-Net 和 AdaAttn：这些方法有时难以保持色彩一致性，尤其是在应用风格时。例如，在第二行中，人物裙子的颜色变成了蓝色，而原图中的颜色是绿色。
StyleID：虽然能够生成一致的图像，既展示了原始内容也展示了风格特征，但颜色对比度和饱和度过高，通常是扩散模型的典型表现。这可能会改变柔和风格的应用，或者导致过高对比度风格的过度应用。

在第四行的例子中，所有方法都能够生成令人满意的结果。我们的方法表现为在风格迁移的同时，能保持内容图像的正确色彩一致性，但不至于过度改变生成样本的饱和度和对比度。

尽管取得了优秀的结果，但模型有时会无法正确应用风格，如图 5 所示。有时它会在生成的图像中产生不均匀的补丁，且它们之间有间隙。一个可能的原因是，Mamba-ST 继承了 RNN 的局限性；在某些情况下，由于上下文信息在状态中的记忆，确保补丁之间的连续性变得困难。

4.3 消融实验

最后，我们进行了几项消融研究，以确定系统的最佳配置，结果见表 2。

内容图像 vs 风格图像：我们研究了将内容图像传递给选择性扫描模块的影响，结果表明，使用内容图像输入的效果较差。因此，我们最终选择了使用风格图像作为输入。
扫描方向数：当将扫描方向数从 4 个减少到 1 或 2 个时，性能大幅下降。采用单个方向时，性能显著下降；使用 2 个方向时，性能略有提升，但仍低于 4 个方向的最终架构。
Mamba 内部状态维度：当维度从 8 增加到 16 时，性能有所提升，但当增加到 32 时，性能下降。因此，最终选择 16 作为内部状态的维度。
去除 shuffling 模块：去除 shuffling 模块后，FID 得分有所改善，表示模型更准确地捕捉了参考图像的风格特征，生成了较好的风格化图像。然而，LPIPS 得分显著下降，表明模型可能过度强调风格，牺牲了内容的保留。没有洗牌时，风格图像中的空间和结构信息得到了保留，最终生成了一个融合了内容和风格的图像。