当前位置：首页 > news >正文

《MambaIR：一种基于状态空间模型的简单图像修复基线方法》学习笔记

news 来源：原创 2025/6/14 12:02:17

paper：2402.15648

摘要

一、引言

1、模型性能的提升依赖于网络感受野的扩大：

2、全局感受野和高效计算之间存在固有矛盾：

3、改进版 Mamba的巨大潜力

4、Mamba 在图像修复任务中仍面临以下挑战：

5、方法

6、主要贡献

二、相关工作

1、图像恢复

2、空间状态模型（State Space Models, SSMs）

三、框架

1、预备知识

2、整体架构

1. 浅层特征提取

2. 深层特征提取

3. 高质量重建

3、残差状态空间组 Residual State-Space Block (RSSB)

4、视觉状态空间模块 Vision State-Space Module (VSSM)

5、二维选择性扫描模块 2D Selective Scan Module (2D-SSM)

6、损失函数

四、实验

1、数据集

2、实验细节

3、消融实验

4、超分辨率对比实验

5、图像去噪对比实验

摘要

近年来，图像修复取得了显著进展，这主要得益于现代深度神经网络（如 CNN 和 Transformers）的发展。然而，现有的修复骨干网络在实际应用中往往面临全局感受野与高效计算之间的矛盾。最近，选择性结构化状态空间模型（Selective Structured State Space Model），尤其是改进版本 Mamba，在长距离依赖建模方面表现出巨大的潜力，并且其复杂度为线性级别，这为解决上述矛盾提供了一种可能。然而，标准的 Mamba 在低级视觉任务中仍面临一些挑战，例如局部像素信息丢失和通道冗余问题。

在这项工作中，我们提出了一种简单但有效的基线方法，称为 MambaIR。该方法在原始 Mamba 的基础上引入了局部增强和通道注意力机制，从而利用了局部像素的相似性并减少了通道冗余。大量实验表明，我们的方法具有显著的优越性。例如，在图像超分辨率（Image Super-Resolution）任务中，MambaIR 以类似的计算成本超越了 SwinIR，性能提升高达 0.45dB，同时还保留了全局感受野的优势。

一、引言

图像修复（Image Restoration）的目标是从给定的低质量输入中重建高质量图像。这是计算机视觉中的一个长期研究问题，并包括许多子任务，例如超分辨率（Super-Resolution）、图像去噪（Image Denoising）等。近年来，随着现代深度学习模型的引入（如 CNN 和 Transformer），图像修复技术的性能不断刷新。

1、模型性能的提升依赖于网络感受野的扩大：

更大的感受野可以让网络从更广泛的区域中捕获信息，有助于参考更多像素来重建目标像素。
较大的感受野能够提取图像中的高级模式和结构，这对于某些需要保持结构的任务（如图像去噪）至关重要。
基于 Transformer 的方法通常具有更大的感受野，在实验中优于基于 CNN 的方法，且研究表明，激活更多像素通常会带来更好的修复效果。

2、全局感受野和高效计算之间存在固有矛盾：

CNN 修复网络尽管感受野有限，但由于卷积的并行计算效率，适合在资源受限的设备上部署。
基于 Transformer 的方法通常需要处理与图像分辨率一致数量的“token”，导致计算复杂度呈二次增长，即使引入高效注意力机制（如窗口注意力），以牺牲全局感受野为代价。

3、改进版 Mamba的巨大潜力

Mamba 使用离散化的状态空间方程，可以通过特殊设计的结构化重新参数化来建模长距离依赖。
Mamba 的并行扫描算法允许在 GPU 等现代硬件上高效训练。

4、Mamba 在图像修复任务中仍面临以下挑战：

局部像素遗忘问题：Mamba 将图像展平为 1D 序列处理，导致空间上相邻的像素可能在序列中变得相距遥远，从而丢失局部信息。
通道冗余问题：为了记忆长序列依赖，状态空间方程通常需要大量隐藏状态，导致关键通道特征学习受限。

5、方法

为了解决上述问题，作者提出了 MambaIR，一种简单而有效的基线模型，通过以下三阶段适配 Mamba 于图像修复任务：

浅层特征提取：使用简单的卷积层提取浅层特征。
深层特征提取：使用堆叠的残差状态空间块（Residual State Space Block, RSSB）。RSSB 的核心改进包括：
- 局部卷积：缓解局部像素遗忘问题。
- 通道注意力：减少隐藏状态引起的通道冗余。
- 可学习跳跃连接因子：优化特征传递。
高质量图像重建：整合浅层和深层特征，输出高质量图像。

MambaIR 结合了全局有效感受野与线性计算复杂度，为图像修复提供了一种全新备选骨干网络。

6、主要贡献

首次将状态空间模型适配于低级图像修复：通过大量实验，提出了一个有效的替代方法 MambaIR。
设计残差状态空间块（RSSB）：通过局部增强和通道冗余降低，提升了标准 Mamba 的能力。
实验验证：在多项任务中，MambaIR 优于其他强基线，为图像修复任务提供了强大的骨干网络解决方案。

二、相关工作

1、图像恢复

图像恢复旨在从低质量图像重建出高质量图像，其研究已经因为深度学习的引入而取得显著进展。一些开创性工作为图像超分辨率（如 SRCNN ）、图像去噪（如 DnCNN）以及 JPEG 压缩伪影还原（如 ARCNN）提供了基础。早期的深度学习方法通常通过卷积神经网络（CNN），结合残差连接和密集连接等技术，提升模型的表达能力。然而，尽管这些方法取得了成功，CNN 在建模全局依赖性上仍面临挑战。

近年来，Transformer 在多任务场景中的表现（如时间序列、三维点云和多模态学习）证明了其强大的建模能力。这促使研究者将 Transformer 引入到图像恢复任务中。然而，Transformer 的自注意力机制因其二次方计算复杂度限制了其在高分辨率图像上的直接应用。为此，IPT将图像分割为多个小块，分别应用自注意力机制进行处理；SwinIR则通过引入滑动窗口注意力改进了性能。此外，许多工作专注于设计高效注意力机制来进一步提升恢复性能。尽管如此，这些高效注意力机制设计往往以牺牲全局感受野为代价，未能本质上解决计算效率和全局建模之间的平衡问题。

2、空间状态模型（State Space Models, SSMs）

状态空间模型（SSMs）源自经典的控制理论，近年来被引入深度学习领域，成为一种在状态转换中具有竞争力的骨干网络。其在线性扩展序列长度方面表现出的出色建模长距离依赖能力，吸引了广泛关注。例如，结构化状态空间序列模型（S4）是深度状态空间模型的开创性工作之一，通过结构化重参数化实现了对长距离依赖的建模；随后，S5 层在 S4 基础上引入了多输入多输出（MIMO）SSM 和高效的并行扫描。此外，H3实现了性能上的显著提升，几乎弥补了 SSM 与 Transformer 在自然语言任务中的表现差距。研究者还通过门控单元增强 S4，提出了门控状态空间层（Gated State Space Layer），进一步提高了其能力。

最近，Mamba]作为一种基于数据的 SSM，引入了选择性机制和高效硬件设计，不仅在自然语言任务上超越了 Transformer，还具有线性扩展输入长度的优异特性。此外，Mamba 已被初步应用于视觉任务，如图像分类、视频理解和生物医学图像分割等领域。然而，其在图像恢复任务中的潜力尚未充分挖掘。本研究通过针对图像恢复任务的特定设计，将 Mamba 应用于图像恢复，提出了一种简单但有效的基线方法，为未来研究提供借鉴。

三、框架

1、预备知识

结构化状态空间序列模型（S4）的最新进展受到连续线性时不变（Linear Time-Invariant, LTI）系统的启发。LTI 系统通过隐式的潜在状态 $h(t) \in \mathbb{R}^N$ 将一维输入函数或序列 $x(t)\in \mathbb{R}$ 映射为输出序列 $y(t) \in \mathbb{R}$ 。其数学形式可以通过以下线性常微分方程（ODE）表示：

$h'(t) = A h(t) + B x(t), \quad y(t) = C h(t) + D x(t)$

其中 N 为状态维度，矩阵 $A \in \mathbb{R}^{N \times N}$ $B \in \mathbb{R}^{N \times 1}$ $C \in \mathbb{R}^{1 \times N}$ ，以及标量 $D \in \mathbb{R}$ ]是模型参数。

为了将上述连续系统转化为适用于深度学习的形式，通常需要进行离散化操作。设 $\Delta$ 为时间步长参数，离散化过程通常采用零阶保持（Zero-Order Hold, ZOH）规则，定义如下：

$\tilde{A} = \exp(\Delta A), \quad \tilde{B} = (\Delta A)^{-1} (\exp(\Delta A) - I) \cdot B$

其中 $\tilde{A}$ 和 $\tilde{B}$ 分别为离散化后的参数。

离散化后，上述系统的表达式可以改写为递归神经网络（RNN）的形式：

$h_k = \tilde{A} h_{k-1} + \tilde{B} x_k, \quad y_k = C h_k + D x_k$

其中 $h_k$ 为离散时间步 $k$ 的隐状态， $x_k$ 和 $y_k$ 分别为输入和输出。

进一步地，等价的数学推导可以将上述 RNN 表达式转化为卷积神经网络（CNN）的形式：

$K \triangleq (CB, CAB, \cdots, CA^{L-1}B), \quad y = x \circledast K$

其中 $L$ 为输入序列的长度， $\circledast$ 表示卷积操作， $K \in \mathbb{R}^L$ 是一个结构化卷积核。

Mamba 模型的改进

在 S4 的基础上，Mamba模型通过引入输入依赖性优化了参数 $B$ 、 $C$ 和 $\Delta$ ，从而实现了动态特征表达能力。具体来说，Mamba 继承了 S4 的递归形式，使得模型能够记忆超长序列，并激活更多像素以辅助恢复任务。同时，其并行扫描算法允许模型在享有卷积形式所带来高效训练的同时，进一步提升计算效率。

在图像恢复任务中，Mamba 利用 S4 模型对长距离依赖建模的优势，不仅能够处理高分辨率图像，还能通过动态调整特征表示适应不同图像场景。这种特性为高效图像恢复提供了一种具有竞争力的解决方案。

2、整体架构

如图所示，MambaIR 模型的整体架构包含三个主要阶段：浅层特征提取（Shallow Feature Extraction）、深层特征提取（Deep Feature Extraction）和高质量重建（High-Quality Reconstruction）。

1. 浅层特征提取

给定一个低质量输入图像 $I_{\text{LQ}} \in \mathbb{R}^{H \times W \times 3}$ ，首先通过一个 $3 \times 3$ 的卷积层提取浅层特征：

$F_{\text{S}} = \text{Conv}_{3 \times 3}(I_{\text{LQ}})$

其中 $F_{\text{S}} \in \mathbb{R}^{H \times W \times C}$ 表示浅层特征， $H$ 和 $W$ 分别为输入图像的高度与宽度， $C$ 为通道数。该阶段旨在对输入图像进行初步特征编码，为后续深层特征提取阶段提供基础表示。

2. 深层特征提取

浅层特征 $F_{\text{S}}$ 进一步进入深层特征提取阶段，经过多层堆叠后生成深层特征 $F_{\text{D}}^{l} \in \mathbb{R}^{H \times W \times C}$ ，其中 $l \in \{1, 2, \cdots, L\}$ 表示第 $I$ 层。该阶段通过多个残差状态空间组（Residual State-Space Groups, RSSGs）实现特征提炼。

每个 RSSG 包含若干个残差状态空间块（Residual State-Space Blocks, RSSBs），每个 RSSB 基于状态空间模型设计，用以捕获超长依赖特征。为了进一步优化特征提取效果，每个 RSSG 末尾附加了一个卷积层，用于对 RSSB 输出特征进行细化。经过多个 RSSG 的堆叠，模型逐层深化输入图像的全局上下文表征。

RSSG 的结构优势在于通过状态空间模型的动态特性，对多尺度图像特征进行高效建模，兼顾了全局感受野和计算效率。

3. 高质量重建

完成深层特征提取后，将最终的深层特征 $F_{\text{D}}^{L}$ 与浅层特征 $F_{\text{S}}$ 进行逐元素加和（element-wise sum）：

$F_{\text{R}} = F_{\text{D}}^{L} + F_{\text{S}}$

其中 $F_{\text{R}} \in \mathbb{R}^{H \times W \times C}$ 是重建阶段的输入特征。通过重建模块， $F_{\text{R}}$ 被进一步处理以生成高质量的输出图像 $I_{\text{HQ}}$ ：

$I_{\text{HQ}} = \text{Reconstruction}(F_{\text{R}})$

重建阶段的设计旨在结合浅层和深层特征，以有效提升输出图像的细节质量和整体视觉效果。

3、残差状态空间组 Residual State-Space Block (RSSB)

RSSB 是专为 MambaIR 设计的基础模块，旨在结合状态空间模块（SSM）的长程依赖建模能力和传统卷积的局部特征提取优势，从而实现高效的图像恢复。

对比 Transformer 的传统模块设计
现有基于 Transformer 的图像恢复网络大多遵循 Norm → Attention → Norm → MLP 的设计模式。然而，尽管 Attention 和 SSM 都可以建模全局依赖，但二者在行为特性上存在差异。简单地用 SSM 替代 Attention 往往会导致次优结果，因此需要重新设计适配 SSM 的模块结构。
局部特征的补充问题
SSM 模块通过将特征图展平成一维序列进行处理，其局部像素感知能力受限于展平策略。例如，四方向展开策略会导致空间上相邻像素在 1D 序列中变得距离较远，导致局部像素特征丢失。为此，引入局部卷积层来补偿邻域特征。
通道冗余问题
SSM 往往引入大量隐藏状态以建模长程依赖，但这会导致显著的通道冗余现象。为了避免这一问题，引入通道注意力机制（CA）选择关键通道，从而提升通道表达能力。

RSSB 的设计如图所示，其输入为第 $I$ 层的深层特征 $F^l_D \in \mathbb{R}^{H \times W \times C}$ ，输出为下一层的特征 $F^{l+1}_D$ 。

以下是 RSSB 的具体处理流程：

长程依赖建模（Vision State-Space Module, VSSM）
- 首先通过 LayerNorm 对输入特征 $F^l_D$ 进行归一化。
- 将归一化特征传入 VSSM，提取空间维度上的长程依赖。
- 使用一个可学习的缩放因子 $s \in \mathbb{R}^C$ 对跳跃连接进行加权，形成第一阶段输出：

$Z^l = \text{VSSM}(\text{LN}(F^l_D)) + s \cdot F^l_D$

局部特征补偿
- 为解决 SSM 中的局部像素遗忘问题，对 $Z^l$ 进行 LayerNorm 归一化，并通过局部卷积层补偿邻域特征。
- 卷积层采用瓶颈结构：通道首先压缩为原来的 $1/\gamma$ ，然后再扩展回原始大小，公式如下：

$Z^l_{\text{conv}} = \text{Conv}(\text{LN}(Z^l))$

其中卷积操作包含通道压缩和扩展的过程。

通道注意力（Channel Attention, CA）
- 在局部卷积后的特征上，应用通道注意力机制选择关键通道，从而提升模块的表达能力并减少通道冗余：

$Z^l_{\text{CA}} = \text{CA}(Z^l_{\text{conv}})$

最终输出
- 最后，通过残差连接将 $Z^l_{\text{CA}}$ 与 $Z^l$ 融合，使用另一个可学习缩放因子 $s' \in \mathbb{R}^C$ 调节残差路径，形成模块的最终输出：

$F^{l+1}_D = Z^l_{\text{CA}} + s' \cdot Z^l$

4、视觉状态空间模块 Vision State-Space Module (VSSM)

VSSM（Vision State-Space Module）在图像恢复任务中引入了状态空间方程（State-Space Equation）来建模长程依赖性。与传统的Transformer网络通常通过分割图像为小块或采用平移窗口注意力（shifted window attention）来限制全图层级的交互不同，VSSM通过线性复杂度高效地捕捉长程依赖性，避免了上述限制。

VSSM的架构如图所示

输入特征 $X \in \mathbb{R}^{H \times W \times C}$ 会通过两个并行分支进行处理：

第一分支：
- 输入特征的通道数通过线性层扩展到 $\lambda C$ ，其中 $\lambda$ 是预定义的通道扩展因子。
- 接下来，特征通过深度卷积（depth-wise convolution）、SiLU 激活函数、2D SSM 层以及 LayerNorm 进行处理。
第二分支：
- 输入特征同样通过线性层将通道数扩展到 $\lambda C$ ，并经过 SiLU 激活函数。

两条分支的输出通过 Hadamard 乘积（元素级乘法）进行聚合，最后将通道数投影回原始的 $C$ ，以生成输出特征 $X_{\text{out}}$ ，其形状与输入特征相同。

第一分支：

对输入特征 $X$ 进行线性扩展、深度卷积、激活函数、2D SSM 和 LayerNorm 处理：

$X_1 = \text{LN}\left( \text{2D-SSM} \left( \text{SiLU} \left( \text{DWConv} \left( \text{Linear}(X) \right) \right) \right) \right)$

第二分支：

对输入特征进行线性扩展和 SiLU 激活：

$X_2 = \text{SiLU}\left( \text{Linear}(X) \right)$

聚合与输出：

通过 Hadamard 乘积将两个分支的特征进行聚合，并将通道数恢复到 $C$ ：

$X_{\text{out}} = \text{Linear}(X_1 \odot X_2)$

其中 $\odot$ 表示 Hadamard 乘积，即元素级的乘法。

5、二维选择性扫描模块 2D Selective Scan Module (2D-SSM)

在标准的 Mamba 网络中，由于其因果处理的特性，输入数据只能局部处理，这种方式适合处理具有序列性质的自然语言处理（NLP）任务。然而，当这种结构应用于图像等非因果数据时，便会遇到显著的挑战。为了更好地利用图像中的二维空间信息，我们采用了2D Selective Scan Module (2D-SSM)。

如图所示，2D-SSM 对输入的图像特征进行处理。首先，将二维图像特征展平成一维序列，并沿四个不同方向进行扫描，具体如下：

从左上角到右下角（top-left to bottom-right）
从右下角到左上角（bottom-right to top-left）
从右上角到左下角（top-right to bottom-left）
从左下角到右上角（bottom-left to top-right）

然后，使用离散的状态空间方程来捕捉每个序列的长程依赖性。最后，将所有扫描序列进行求和，并通过 reshape 操作恢复其原始的二维结构。

6、损失函数

采用 L1 损失函数来优化 MambaIR 以进行图像超分辨率（SR）。具体损失函数公式如下：

$L = ||I_{\text{HQ}} - I_{\text{LQ}}||_1$

其中 $|| \cdot ||_1$ 表示 L1 范数。

对于图像去噪任务，我们使用 Charbonnier 损失，其公式为：

$L = ||I_{\text{HQ}} - I_{\text{LQ}}||_2 + \epsilon^2$

其中 $\epsilon = 10^{-3}$ 为一个小的常数，用于稳定计算并避免数值问题。

四、实验

1、数据集

为了进行图像恢复任务的实验，我们遵循先前工作的设置，涵盖了图像超分辨率（包括经典超分辨率、轻量级超分辨率、真实超分辨率）、图像去噪（包括高斯彩色图像去噪和真实世界图像去噪）以及 JPEG 压缩伪影去除（JPEG CAR）。我们使用了以下数据集来训练和评估模型：

图像超分辨率（SR）：
- 训练数据集：DIV2K和 Flickr2K用于经典超分辨率模型的训练；仅使用 DIV2K 来训练轻量级超分辨率模型。
- 测试数据集：Set5 、Set14、B100、Urban100和 Manga109 用于评估不同超分辨率方法的效果。
高斯彩色图像去噪：
- 训练数据集：DIV2K 、Flickr2K、BSD500 和 WED。
- 测试数据集：BSD68、Kodak24 、McMaster和 Urban100。
真实图像去噪：
- 训练数据集：使用来自 SIDD 数据集的320张高分辨率图像进行训练。
- 测试数据集：使用 SIDD 测试集和 DND数据集进行测试。

模型评估：所有任务的性能均通过在 YCbCr 颜色空间的 Y 通道上计算 PSNR 和 SSIM 来进行评估。当测试时使用自集成策略时，模型被称为 MambaIR+。

2、实验细节

根据先前的工作，我们对数据进行了增强，方法包括水平翻转和随机旋转（90°、180°、270°）。此外，在训练过程中，我们将原始图像裁剪为 64×64 的图像块用于超分辨率任务，裁剪为 128×128 的图像块用于去噪任务。

图像超分辨率（SR）训练：
我们使用 ×2 模型的预训练权重来初始化 ×3 和 ×4 模型的权重，并通过减半学习率和训练总迭代次数来减少训练时间。
批量大小调整：
为了确保公平比较，我们将图像超分辨率的训练批量大小调整为 32，图像去噪的训练批量大小调整为 16。
优化器与学习率：
我们使用 Adam 优化器，其超参数为 $\beta_1 = 0.9$ , $\beta_2 = 0.999$ 。初始学习率设置为 $2 \times 10^{-4}$ ，并在训练达到特定的里程碑时将学习率减半。
硬件配置：
我们的 MambaIR 模型使用 8 台 NVIDIA V100 GPU 进行训练。