引言

在计算机视觉领域，图像风格迁移 (ST) 是最具视觉吸引力的任务之一。它使我们能够将一张内容图像 (如街道照片) 与一张风格图像 (如《星月夜》) 融合，让照片看起来像是梵高亲手绘制的一样。

虽然其艺术效果令人惊叹，但其背后的工程实现面临着一个显著的瓶颈: 生成质量与计算效率之间的权衡。

为了生成高质量的艺术图像，模型需要具备“全局感受野”——它需要一次性理解整个图像，以捕捉大尺度的风格模式和语义结构。

  • CNN (卷积神经网络) 通过堆叠多层来实现这一点，但这会带来沉重的计算负担。
  • Transformer (使用自注意力机制) 天生具有全局感受野，但受限于二次方复杂度 (\(O(N^2)\)) ，导致速度慢且极其消耗内存。
  • 扩散模型 (Diffusion Models) 能产生令人难以置信的细节，但需要大量的迭代步骤，推理时间显著增加。

SaMam (Style-aware Mamba，风格感知 Mamba) 由此应运而生。在最近的一篇论文中，研究人员提出利用 Mamba 架构——一种以线性复杂度 (\(O(N)\)) 著称的状态空间模型 (SSM) ——来解决这一难题。

图 1. 不同方法实现的推理时间 t (ms) 与 ArtFID 之间的权衡。圆圈大小代表 MACs (G)。

如图 1 所示，SaMam (绿色圆圈) 占据了一个“最佳平衡点”。它的 ArtFID 得分 (一种质量衡量标准，越低越好) 与重量级的扩散模型 (如 ZStar) 相当，但推理时间和计算成本 (MACs) 却只有后者的一小部分。

在这篇文章中，我们将剖析 SaMam 如何调整 Mamba 架构以适应视觉任务，特别是针对任意风格迁移的定制化设计。

背景: 效率难题

在深入了解 SaMam 之前，我们需要理解为什么现有方法难以平衡速度和质量。

感受野困境

对于神经网络来说，要对特定像素进行风格化处理，了解图像其余部分发生的情况非常有帮助。

  • 基于 CNN 的方法 通常具有有限的局部窗口。为了“看到”整张图像，它们必须通过深层的层级结构处理图像，这增加了浮点运算次数 (FLOPs) 。
  • 基于 Transformer 的方法 计算每一个补丁与其他所有补丁之间的关系。随着图像分辨率的增加，计算时间呈二次方爆炸式增长。

状态空间模型 (SSM) 解决方案

状态空间模型，特别是 Mamba 变体，最近通过以线性复杂度对长序列建模，彻底改变了自然语言处理 (NLP) 领域。它们通过一个隐式潜在状态 \(h(t)\) 将一维函数或序列 \(x(t) \to y(t)\) 进行映射。

连续系统由线性常微分方程 (ODEs) 定义:

连续时间 SSM 方程

这里，\(\mathbf{A}\)、\(\mathbf{B}\)、\(\mathbf{C}\) 和 \(\mathbf{D}\) 是权重参数。为了在深度学习中使用它，我们对这些方程进行“离散化” (将连续时间转换为离散步骤) 。这将方程转换为一种非常类似于循环神经网络 (RNN) 的递归形式:

SSM 的离散递归形式

Mamba 的魔力在于它允许参数如 \(\mathbf{B}\)、\(\mathbf{C}\) 和 \(\Delta\) (时间尺度) 是依赖于输入的 (input-dependent) 。 这使得模型能够选择性地“记住”或“忽略”信息，从而在不产生 Transformer 那样巨大计算成本的情况下创建全局感受野。

SaMam 框架

研究人员开发了 SaMam , 以适应风格迁移的具体挑战。该框架不仅仅是一个标准的 Mamba 模型；它引入了几种新颖的机制来处理 2D 图像和风格注入。

1. 整体架构

如图 2 所示，该架构由三个主要部分组成:

  1. 风格 Mamba 编码器 (Style Mamba Encoder) : 从艺术作品中提取特征。
  2. 内容 Mamba 编码器 (Content Mamba Encoder) : 从照片中提取特征。
  3. 风格感知 Mamba 解码器 (Style-aware Mamba Decoder) : 融合两者以生成最终图像。

图 2. 我们的 SaMam 框架概览 (a) 以及选择性扫描方法的图解 (b)。

输入图像首先被处理成补丁 (类似于 Vision Transformers) 。这些补丁通过 视觉状态空间模块 (VSSMs) 进行处理。

之字形扫描 (Zigzag Scan)

将 Mamba 应用于图像的一个主要挑战是，Mamba 是为 1D 序列 (如文本) 设计的。当你使用标准的栅格扫描 (逐行) 将 2D 图像展平为 1D 序列时，图像中垂直相邻的像素在序列中可能会相距甚远。这被称为空间不连续性 。

为了解决这个问题，SaMam 引入了 之字形扫描 (如图 2a 所示) 。

  • 标准扫描 (如图 2b 中的扫描或交叉扫描) 通常会从一行的右边缘跳到下一行的左边缘。
  • 之字形扫描以连续的蛇形模式遍历图像。这保留了语义连续性，确保序列中相邻的 Token 在图像空间上也是接近的。

2. 解码器与 S7 块

本文的核心创新位于解码器中的 风格感知视觉状态空间模块 (SAVSSM) 。 标准的 Mamba 块设计用于处理内容，但它们没有内在机制来以风格为“条件”。

研究人员创建了一个名为 S7 块 (Style-aware Selective Scan Structured State Space Sequence Block，风格感知选择性扫描结构化状态空间序列块) 的新模块。

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

S7 块的工作原理

在标准的 Mamba 块 (S6) 中，转换矩阵 \(\mathbf{A}\) 通常是固定的或作为静态参数学习的。在 S7 块中，权重参数是根据 风格嵌入 (\(\mathbf{E}_s\)) 动态预测的。

从风格嵌入预测 A 和 D 的方程

通过从风格图像预测 \(\mathbf{A}\) 和 \(\mathbf{D}\):

  1. 风格选择性: 隐藏状态更新 (\(h_t\)) 现在受到艺术风格的影响。模型创建了一个动力系统，其转换规则本身取决于风格是莫奈还是毕加索。
  2. 效率: 尽管这种适应是动态的，但在利用并行扫描的情况下，操作的复杂度仍然保持线性。

3. 解决 Mamba 的弱点

虽然 Mamba 效率很高，但在应用于视觉时也有局限性。作者针对这些问题引入了具体的修复方案。

局部增强 (LoE)

由于 Mamba 将图像展平为序列，存在“局部像素遗忘”的风险。即使有之字形扫描，一些局部邻域关系仍会丢失。

为了解决这个问题，研究人员在编码器的末端添加了一个 局部增强 (LoE) 模块。该模块使用标准的卷积层 (擅长处理局部特征) 和通道注意力机制来“补偿”Mamba 块可能遗漏的任何局部细节。

风格感知模块

为了确保风格渗透到生成的每一部分，SaMam 用风格感知版本替换了标准组件:

  1. SConv (风格感知卷积) : 卷积核不再是标准的深度卷积，而是根据风格嵌入生成的。这捕捉了局部几何结构 (如笔触) 。
  2. SCM (风格感知通道调制) : 这根据风格重新缩放残差分支中的通道，帮助模型强调特定特征 (如调色板) 。
  3. SAIN (风格感知实例归一化) : 归一化在风格迁移中至关重要。SAIN 直接从风格嵌入中预测用于归一化的均值和方差。

下图绘制了使用风格感知实例归一化 (SAIN) 与其他归一化策略的效果对比。SAIN (特别是零初始化版本) 收敛到了更好的 ArtFID 分数。

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

实验与结果

研究人员将 SaMam 与最先进的方法进行了比较，包括 CNN (AesPA) 、Transformer (StyTr2) 和扩散模型 (ZStar、StyleID) 。

定性比较

在视觉上，SaMam 在平衡内容保留与风格应用方面表现出色。

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

在图 5 中，请看带有文字的那一行 (第 7 行) 。

  • ZStar 和 StyleID (基于扩散模型) 倾向于产生幻觉或破坏文字结构。
  • SaMam 在准确应用纹理的同时保持了文字的可读性。
  • 同样，在第 6 行 (建筑物) 中，SaMam 比竞争对手更好地保留了建筑的直线，同时也应用了“斑点”风格。

定量分析

效率的提升是 SaMam 真正的亮点。下表 1 列出了性能指标。

  • LPIPS: 衡量内容保真度 (越低越好) 。
  • FID: 衡量风格相似度 (越低越好) 。
  • ArtFID: 两者的结合 (越低越好) 。
  • MACs / Time: 计算成本和速度。

表 1. 风格迁移方法的定量比较。

数据主要结论:

  1. 速度: SaMam 处理一张图像仅需 0.034 秒 。 相比之下，StyTr2 (Transformer) 需要 0.385 秒，而 ZStar (扩散模型) 更是需要惊人的 42.439 秒 。 SaMam 的速度大约是扩散方法的 1000 倍 。
  2. 质量: 尽管速度极快，SaMam 仍实现了最低 (最好) 的 ArtFID 分数 (26.305) 和 LPIPS 分数 (0.3884) 。它击败了那些体量大几个数量级的模型。

验证有效感受野

Mamba 的主要主张之一是它能实现全局感受野。研究人员通过可视化有效感受野 (ERF) 证明了这一点。

图 6. 我们 SaMam 的有效感受野 (ERF) 可视化。

如图 6 所示，训练后深色区域 (代表感受野) 广泛分布在整个图像上。这证实了 SaMam 确实利用了全局上下文来做出风格化决策，而不仅仅是查看微小的局部补丁。

消融实验: 组件是否重要？

研究人员进行了消融实验，以验证 SConv 层和 S7 块等单个组件的贡献。

SConv 的影响: 图 10 展示了使用标准深度卷积 (DWConv) 与提出的风格感知卷积 (SConv) 之间的差异。

图 10. SConv 的消融研究。

注意“电路板”示例 (顶行) 。 SConv 结果 (中间列) 成功复制了复杂的电路如同线条。 DWConv 结果 (右列) 看起来模糊，未能捕捉到风格中锐利的几何结构。

结论

SaMam 论文为在视觉任务中使用状态空间模型提供了令人信服的论据。通过采用风格感知机制 (特别是 S7 块) 和空间改进 (之字形扫描) 来调整 Mamba，作者创建了一个打破质量与效率传统权衡的框架。

主要启示:

  • 线性复杂度适用于视觉: 我们并不总是需要 Transformer 的二次方成本来获得全局上下文。
  • 动态参数: 根据条件嵌入 (如风格) 预测 SSM 参数 (\(\mathbf{A}\), \(\mathbf{D}\)) 是使这些模型适应生成任务的有力方式。
  • 实时潜力: 凭借低至 0.034 秒的推理时间，SaMam 为消费级硬件上的高质量实时视频风格迁移打开了大门，这对于扩散模型来说仍然是困难的。

SaMam 证明了 Mamba 不仅仅是 NLP 的有力竞争者，它也是计算机视觉和创意 AI 领域的强大竞争对手。