SaMam横空出世！全球首个Mamba风格迁移框架，线性复杂度+局部增强，双指标碾压SOTA！

🌐 社群导航

🔗 点击加入➡️【AIGC/LLM/MLLM/3D/自动驾驶】 技术交流群

🔗 点击加入➡️【2-3月CCF投稿】交流群

最新论文解读系列

论文名：SaMam: Style-aware State Space Model for Arbitrary Image Style Transfer     

论文链接：https://arxiv.org/pdf/2503.15934

项目链接：暂无

导读

全局有效感受野对于图像风格迁移（ST）获得高质量的风格化结果起着至关重要的作用。然而，现有的ST主干网络（如卷积神经网络和Transformer）在实现全局感受野时会面临巨大的计算复杂度。最近，状态空间模型（SSM），特别是改进的变体Mamba，在线性复杂度下对长距离依赖建模显示出了巨大的潜力，这为解决上述困境提供了一种途径。在本文中，我们开发了一个基于Mamba的风格迁移框架，称为SaMam。具体来说，设计了一个Mamba编码器来高效地提取内容和风格信息。此外，开发了一个风格感知的Mamba解码器，以灵活适应各种风格。此外，为了解决现有SSM存在的局部像素遗忘、通道冗余和空间不连续的问题，我们引入了局部增强和之字形扫描机制。定性和定量结果表明，我们的SaMam在准确性和效率方面都优于现有最先进的方法。

简介

风格迁移（ST）旨在捕捉图像风格以生成艺术图像，自开创性工作以来，它引起了越来越多的关注。随着现代深度学习技术的发展，如卷积神经网络、Transformer和扩散模型，风格迁移性能在过去几年中不断提高。我们认为这种改进部分归因于感受野的增加。首先，相对较大的感受野允许模型从更广泛的区域提取足够的图像模式，使其能够更好地捕捉风格模式。其次，有了更大的感受野，模型能够利用内容图像中的更多像素来促进锚点像素的风格迁移。

图1. 不同方法在推理时间 （毫秒）和ArtFID之间的权衡。圆圈的大小表示乘累加运算次数（G）。

方法

1. 预备知识

结构化状态空间序列模型（S4）和Mamba受连续系统的启发，该系统通过一个隐式潜在状态  映射一个一维函数或序列 。具体而言，连续时间状态空间模型（SSMs）可以表述为如下的线性常微分方程（ODEs）：

其中  和  是加权参数。

图2. 我们的SaMam框架概述（a）以及视觉Mamba和VMamba选择性扫描方法的示意图（b）。

之后，通常采用离散化过程将公式1集成到实际的深度学习算法中。该过程引入一个时间尺度参数 ，将连续参数  转换为离散参数 。常用的转换方法是零阶保持（ZOH），定义如下：

离散化后，步长为  的公式1的离散版本可以重写为如下的循环神经网络（RNN）形式：

最后，模型通过全局卷积计算输出。

其中  是输入序列的长度， 是一个结构化卷积核， 表示卷积运算。作为近期先进的状态空间模型，Mamba提出了S6来改进  并使  依赖于输入，从而实现动态特征表示。

2. 总体架构

我们的SaMam由一个风格Mamba编码器、一个内容Mamba编码器和一个风格感知Mamba解码器组成，如图2（a）所示。首先，内容图像   和风格图像  分别输入到编码器中，以获得内容特征  和风格嵌入 。接下来， 被用作风格条件信息，用于调整解码器参数。最后， 输入到解码器中以获得风格化图像 。

3. 风格/内容Mamba编码器

图像首先被嵌入为降采样的图像特征。然后，这些图像特征被输入到视觉状态空间模块（VSSMs）中以提取深度特征。此外，在编码器的末尾引入了额外的局部增强（LoE）来增强从VSSM中提取的特征。

由于VMamba中SS2D模块的计算效率和长距离建模能力，我们也采用线性  深度可分离卷积  SS2D  线性的流程。

之字形扫描：先前的研究已证明使用多种扫描顺序来提高性能的有效性（例如，如图 2(b) 所示，在多个方向上进行逐行和逐列扫描）。先前的扫描顺序只能覆盖一种二维方向（例如，从左到右），这在移动到新的行或列时会导致空间不连续性。此外，由于公式 4 中的参数  作为衰减项，空间不连续性会导致相邻标记的衰减程度发生突变，加剧语义不连续性，并导致生成不自然的风格化纹理。受提出用于语义连续性的连续二维扫描的文献的启发，我们实现了之字形扫描（如图 2(a) 所示）。所提出的方法从 4 个顶点开始，以第一个顺时针方向的列（或行）作为起始扫描线，旨在保持空间和语义的连续性，并生成和谐的风格化结果。

局部增强：正如文献中所提到的，由于风格化状态空间模型（SSMs）将扁平化的特征图作为一维标记序列进行处理，序列中相邻像素的数量会受到扁平化策略的极大影响。在一维标记序列中，空间上相邻的像素之间的距离过远会导致局部像素遗忘（例如，在行  和列  处的块在按行优先扫描时不再与行  和列  处的块相邻）。此外，由于需要大量隐藏状态来记忆非常长距离的依赖关系，风格化状态空间模型（SSMs）会导致明显的通道冗余。为避免这些问题，我们在视觉风格化状态空间模型（VSSM）的末尾添加了局部增强（LoE）模块。具体来说，局部增强（LoE）模块由一个卷积层和一个通道注意力层组成，卷积层用于补偿局部特征，通道注意力层用于增强不同通道的表达能力。

图3.  风格感知视觉状态空间模块（SAVSSM）的详细架构。

4. 风格感知的曼巴解码器

在解码器中，内容特征  和风格嵌入  首先被输入到风格感知的视觉状态空间组（SAVSSGs）中，以获得风格化特征 ，每个风格感知的视觉状态空间组（SAVSSG）包含多个风格感知的视觉状态空间模块（SAVSSMs）。此外，在每个风格感知的视觉状态空间组（SAVSSG）的末尾实现了一个局部增强（LoE）模块，以细化从风格感知的视觉状态空间模块（SAVSSMs）中提取的特征。最后，引入了一个类似于的轻量级解码器，以生成风格化图像 。

4.1 风格感知的视觉状态空间模块

原始的曼巴模块（Mamba Module）是为一维序列设计的，不适用于需要空间感知理解的时空（ST）任务。为此，我们引入了风格感知的多方向序列建模模块（SAVSSM），它将多方向序列建模应用于视觉任务。此外，为了实现灵活的风格感知自适应，我们提出了一种风格可插拔机制（如图3所示）。具体而言，SAVSSM的操作如算法1所示。风格嵌入  作为条件信息，扩展为风格感知结构的参数。然后，内容特征  首先通过风格感知实例归一化（SAIN）进行归一化处理，接着线性投影到  ，再通过风格感知卷积（SConv）进行处理。此外，我们从4个方向处理  。对于每个风格感知S6块（S7块），我们分别将令牌序列  线性投影到  。然后使用  对  和  进行离散化处理，以获得  和  。然后我们通过序列状态模型（SSM）计算  。之后，将输出相加并归一化，得到输出令牌序列  。我们对  进行线性投影，并将其与残差相加，得到风格化特征  。此外，在残差分支中实现了风格感知通道调制（SCM）。

风格感知S6块（S7块）：与标准S6块中来自某个具体嵌入空间的A和D不同，我们引入了一种动态权重生成方案。具体而言，我们从风格嵌入  中预测  和  ：

其中  和  分别表示扩展维度大小和序列状态模型（SSM）维度。我们从两个方面设计S7块。（1）风格选择性：标准S6块仅基于内容更新隐藏状态。然而，隐藏状态应该同时受到内容和风格的影响。此外，S6块中的具体嵌入  也可以通过公式2获得选择性。为了在隐藏状态更新中引入风格信息，我们通过从风格嵌入空间预测  的选择性，而不是使用具体嵌入。（2）效率：如公式4所示，加权参数  和D分别扩展为卷积核和通道级缩放因子，这与类似。动态全局卷积核通过并行操作保持高效计算，同时适应各种风格。额外的风格感知模块：为了实现更好的视觉质量，我们还实现了几个额外的风格感知结构，以融合内容和风格信息。

算法1 风格感知的多方向序列建模模块（SAVSSM）处理流程

(1) 风格卷积（SConv）：受自适应卷积（AdaConv）的启发，该方法提出了一种风格感知的深度可分离卷积，以更好地保留风格图像的局部几何结构，我们用风格卷积（SConv）取代了深度可分离卷积（DWConv）。具体来说，风格嵌入  被输入到一个嵌入器中，以生成风格卷积（SConv）中的卷积核 。注意， 。然后，预测的卷积核  对内容图像特征  执行深度可分离卷积操作。

(2) 风格通道调制（SCM）：受控制残差多尺度模块（CResMD）的启发，该方法使用控制变量对不同通道进行重新缩放，以处理多种图像退化问题，我们的风格通道调制（SCM）学习基于风格嵌入  生成调制系数，以进行通道级特征自适应。具体来说， 被输入到嵌入器和Sigmoid激活层中，以生成通道级调制系数 。然后， 用于对  中的不同通道分量进行重新缩放。

(3) 风格感知实例归一化（SAIN）：除了局部几何结构外，全局属性对最终结果也至关重要。遵循自适应归一化在风格迁移（ST）、视觉推理和图像生成中的广泛应用，我们探索用自适应归一化取代标准归一化，以从风格图像中传递全局属性。与VSS块中层归一化的通道级调制相比，实例归一化的特征级调制在风格迁移（ST）领域更有前景。因此，我们提出了一种风格感知实例归一化（SAIN）。具体来说，风格嵌入  被输入到一个嵌入器中，以预测风格的均值  和方差 。此外，关于残差网络（ResNets）的先前工作发现，将每个残差块初始化为恒等函数是有益的。例如，发现，将每个块中的最终批量归一化缩放因子初始化为零可以加速监督学习环境下的大规模训练。扩散变压器（DiT）使用了类似的初始化策略，将每个块中的层归一化初始化为零。受先前探索的启发，我们将风格感知实例归一化（SAIN）和风格通道调制（SCM）的嵌入器初始化为输出零向量。这将风格感知可变尺度风格迁移模块（SAVSSM）初始化为恒等函数。

图4. 不同归一化策略的比较。

5. 损失函数

总体损失函数由内容项、风格项和恒等项组成，定义如下：

其中  和  分别设置为10、1和50。内容损失和风格损失：与先前的工作类似，我们将内容损失和风格损失定义如下：

其中  指的是从预训练的VGG - 19的第  层提取的特征。 和  分别表示提取特征的均值和方差。

身份损失：为了学习更准确的内容和风格信息，我们采用身份损失：

(9)

其中 （或 ）指的是由两张具有相同内容（或风格）的图像合成的输出图像。

图 5. 与先前最先进方法的定性比较。

表 1. 风格迁移（ST）方法的定量比较。分别突出显示了最佳和次佳结果。运行时间和每秒乘累加运算次数（MACs）

实验

1. 实验设置

实现细节：我们使用MS - COCO作为内容数据集，并从WikiArt中选择风格图像。在算法1中，图像特征通道数、扩展维度大小和SSM维度分别设置为256、512和16。VSSM中的和设置为与算法1中相同。在训练期间，内容图像和风格图像被重新缩放为像素。随机选择8对内容 - 风格图像块作为一个小批量。我们采用Adam优化器对整个模型进行次迭代训练。初始学习率设置为，并每次迭代将其降低一半。

评估指标：遵循StyleID的协议，我们使用ArtFID和内容特征结构距离（CFSD，Content Feature Structural Distance）作为指标。具体来说，ArtFID等于。由于这两个指标与人类判断高度一致，LPIPS衡量内容保真度，而FID评估风格相似度。此外，CFSD是一个额外的内容保真度指标，用于衡量图像块之间的空间相关性。

2. 与现有技术的比较

我们将我们的SaMam与最近的最先进的风格迁移（ST，Style Transfer）方法进行比较，包括基于卷积神经网络（CNN，Convolutional Neural Network）的方法（AesPA、EFDM、ATK、UCAST）、基于Transformer的方法（StyTr2、S2WAT、StyleFormer、STTR）、基于可逆神经网络（Reversible - NN）的方法（ArtFlow、CAPVST）和基于扩散模型的方法（DiffuseIT、ZStar、StyleID、VCT）。我们通过遵循这些方法的官方代码并使用默认配置来获取结果。

2.1 定性比较

我们在图5中展示了视觉比较结果。可以观察到，我们的SaMam从风格图像中捕捉到全局属性（例如，纹理和颜色，如行所示），同时也关注风格图案的局部几何特征（例如，行和行中的斑点）。除了获取足够的风格信息外，我们的方法还能准确地保留内容结构（例如，行中的建筑物），并产生更清晰的细节，实现更高的感知质量（例如，行中的文本和行中的车牌）。相比之下，StyleID和Zstar严重破坏了内容细节（例如，行中的漫画肖像和行中女孩的脸）。虽然CAPVST擅长捕捉风格图像的颜色，但它在风格化结果中破坏了局部几何结构和内容细节（例如，行）。S2WAT和AesPA也难以取得令人满意的结果。

2.2 定量比较

我们借助一些定量指标来更好地评估所提出的方法。需要注意的是，由于DiffuseIT和VCT在推理时进行训练，并且需要相当大的计算成本，因此未报告它们的乘加运算次数（MACs，Multiply - Accumulate Operations）。

（1）风格化质量：我们收集了20张内容图像和40张风格图像，为每种方法合成800张风格化图像，并在表1中展示它们的平均指标得分。可以观察到，基于扩散模型的方法在平衡内容和风格方面面临巨大挑战。为了更好地融合内容和风格特征，先前基于CNN（例如，AesPA和ATK）和基于Transformer的方法利用注意力机制来建立长距离依赖关系，以提取结构信息。然而，这种机制给这些方法提取完整的图像属性（例如，局部几何特征和细节）并生成满意的结果带来了巨大挑战。相比之下，除了通过Mamba建立长距离依赖关系外，我们还设计了更多风格感知架构（例如，SConv），这些架构能够更灵活地适应各种风格。因此，我们的SaMam在4个质量指标上取得了最佳结果，这表明它在传递足够的风格图案的同时，能更好地保留内容细节。

(2) 效率：如图1和表1所示，我们的SaMam在计算量和推理时间具有竞争力的情况下，实现了显著的性能提升。基于扩散的方法需要大量时间进行DDIM反演和采样，甚至在单一风格上进行训练需要更多时间。基于Transformer和可逆神经网络（Reversible-NN）的方法也很耗时。相比之下，在乘加运算次数（MACs）和推理时间方面，我们的方法仅次于基于卷积神经网络（CNN）的方法（即EFDM）。这是因为我们基于Mamba的方法执行全局卷积，以并行方式处理每个图像令牌。这种先进的方案继承了基于CNN方法的高推理效率，同时保持了长距离依赖关系。这进一步证明了我们方法的优越性。

表2. 对所提出组件的消融研究。rp.b. 代表“替换为”，r.m. 代表“移除”。

图6. 我们的SaMam模型的有效感受野（ERF）可视化结果。

3. 模型分析

有效感受野（ERF）：我们在图6中展示了有效感受野（ERF）。较广泛分布的深色区域表示更大的有效感受野。可以观察到，经过训练后，我们的SaMam模型展现出全局有效感受野，能够捕捉风格和内容方面的长距离依赖关系。

之字形扫描：为了在扫描过程中保持空间连续性，我们采用了一种四方向之字形扫描方法。为了证明其有效性，我们用另一种四方向扫描方法（即十字扫描）替代之字形扫描，得到配置B。如表2所示，十字扫描方法同时降低了艺术领域的弗里德距离（ArtFID）和内容风格特征距离（CSFD）。我们在图7中进一步提供了可视化结果。可以观察到，之字形扫描方法产生的背景更清晰，且与内容图像更接近的伪影更少。这是因为空间连续性不会给内容信息带来突然的变化，这使得自适应风格转移模块（SSM）的参数调整更加困难。

局部增强：我们引入了一个局部增强（LoE）模块来缓解局部像素遗忘问题。在图8中，可以观察到配置C存在不自然的噪声伪影，破坏了图像内容的平滑性。表2中的定量结果也证明了局部增强模块的有效性。风格感知的可变风格转移模块（SAVSSM）：我们设计了一个风格感知的可变风格转移模块（SAVSSM），以灵活适应不同的风格。我们进一步证明了所提出组件的有效性。

图7. 不同模型配置的消融研究。

图8. 局部增强的消融实验。请放大以获得最佳观看效果。

图9. S7模块的消融实验。

(1) S7模块：我们提出了一种新颖的S7模块，以更好地捕捉风格特征。我们用S6模块替换S7模块，得到配置D。可以观察到，我们的S7模块比配置D取得了显著更高的分数。从图9可以看出，配置A再现的图像颜色和对比度更接近风格图像。此外，它保留了内容细节，并产生了感知质量更高的更清晰边缘（例如，图9中的场景）。S7模块继承了全局有效感受野，以适应各种风格。

(2) 风格卷积（SConv）：我们提出了风格卷积（SConv），用于将风格图像中的局部几何结构复制到内容图像中。为了验证其有效性，我们用普通的深度可分离卷积（DWConv）层替换风格卷积（SConv），得到配置E。风格卷积（SConv）有助于取得显著更好的指标。例如，在图10的第一个场景中，我们的风格卷积（SConv）产生的电路图案更接近风格图像。

图10. 风格卷积（SConv）的消融实验。

(3) 风格上下文模块（SCM）：由于将风格上下文模块（SCM）添加到残差分支中，我们的方法取得了更好的视觉质量。此外，配置F的内容得分显著降低，这表明风格化图像存在更多的图像失真问题（例如，图7中的第一个场景）。

(4) 风格感知实例归一化（SAIN）：我们在训练过程中测量了各种归一化策略的艺术领域的弗里德距离（ArtFID）和内容风格特征距离（CSFD）。图4显示了结果，这表明初始化为零的风格感知实例归一化（SAIN）优于其他策略。然后，我们在方法中采用初始化为零的风格感知实例归一化（SAIN-zero）来捕捉风格图像的全局特征。为了证明其有效性，我们在SaMam模型中用普通的实例归一化（IN）替换风格感知实例归一化（SAIN），得到配置G，该配置在高色度边缘存在明显的伪影（例如，图7中的第二个场景）。

总结

在本文中，我们探索了最近先进的状态空间模型（即Mamba）在任意图像风格迁移中的潜力。为此，我们提出了一种风格感知的Mamba（SaMam）模型，以在计算效率和全局有效感受野之间取得平衡。具体来说，我们引入了一个Mamba编码器和一个风格感知的Mamba解码器。此外，我们基于风格嵌入设计了一个风格感知的可变风格转移模块（SAVSSM），以灵活适应各种风格。实验结果表明，我们的模型在风格迁移（ST）任务中取得了最先进的性能。