渐进分割Mamba（PS-Mamba）图像恢复方法

引言

图像恢复是计算机视觉中的核心低层任务，旨在从退化输入（如噪声、模糊、JPEG压缩伪影等）中重建高质量图像。这类任务需要同时满足两个看似矛盾的需求：保留细粒度的局部结构（如纹理、边缘）和维持长程空间一致性。

传统的卷积神经网络（CNN）虽然在局部特征提取方面表现优异，但其固定且有限的感受野难以有效建模长程依赖关系。而 Vision Transformer（ViT）通过全局自注意力机制虽然能够捕获长程依赖，但计算复杂度随图像尺寸呈二次方增长，给高分辨率图像处理带来沉重的计算负担。¹

状态空间模型（State Space Models, SSM），尤其是 Mamba，提供了一种兼顾全局建模能力与线性时间复杂度的解决方案。然而，将一维的 Mamba 直接应用于二维图像时，会面临两个根本性挑战：局部性失真（Locality Distortion）和长程衰减（Long-Range Decay）。²

渐进分割 Mamba（Progressive Split Mamba, PS-Mamba）正是为解决这些问题而提出的新型框架，其核心思想是通过几何一致性划分保留空间拓扑结构，并通过跨尺度快捷通路对抗信息衰减。

问题背景

局部性失真

Mamba 等 SSM 最初设计用于处理一维序列数据。当将二维图像特征 $\mathbf{X}\in {\mathbb{R}}^{H\times W\times C}$ 展平为一维序列时，自然的四连通邻域像素可能被 $O(W)$ 的距离分隔开。这种raster扫描（逐行展平）方式会破坏图像的内在空间邻接关系，迫使模型通过长程过渡来建模本应短程可达的局部纹理特征。

数学上，给定展平后的序列长度 $L=H\times W$，原本空间上相邻的两个像素 $(i,j)$ 和 $(i,j+1)$ 在序列中的距离可能达到 $W$，这与它们的空间邻近性严重不符。

长程衰减

Mamba 的选择性状态空间机制虽然在理论上能够建模任意长距离的依赖关系，但其稳定性约束要求状态转移矩阵的谱半径小于 1。这意味着信息沿序列传播时会呈指数衰减：

$${\mathbf{h}}_i=A{\mathbf{h}}_{i-1}+B{\mathbf{x}}_i,{\mathbf{y}}_i=C{\mathbf{h}}_i+D{\mathbf{x}}_i$$

其中 $|{\lambda }_{\max }(A)|<1$ 以保证训练稳定性。然而，对于一幅 $1024\times 1024$ 的高分辨率图像，序列长度达到百万量级，长程像素的信息在到达目标位置时几乎完全衰减。这种长程衰减问题严重削弱了全局一致性的保持能力。

核心方法

几何一致性划分

PS-Mamba 的核心创新在于几何一致性划分（Geometry-Consistent Partitioning）。不同于简单地展平整个特征图，PS-Mamba 在进行状态空间处理之前，先将图像划分为保持邻域完整性的几何区块。

定义分割算子 ${\mathcal{S}}_k(\cdot )$ 将特征图划分为 $k$ 个连续的几何区域：

$${\mathcal{S}}_k:{\mathbb{R}}^{H\times W\times C}\to \left\{{\mathbf{X}}^{(1)},{\mathbf{X}}^{(2)},\dots ,{\mathbf{X}}^{(k)}\right\}$$

其中每个区块 ${\mathbf{X}}^{(j)}\in {\mathbb{R}}^{H_j\times W_j\times C}$ 保持其原始的二维空间结构。划分后的序列长度变为 $L_j=H_j\times W_j\ll L=H\times W$，从而确保：

1. 邻域完整性：每个区块内的像素保持原有的空间邻接关系
2. 递归深度缩短：每个区块的 Mamba 递归深度大幅减少，信息衰减减轻
3. 线性复杂度保持：总计算复杂度仍为 $O(HW)$

渐进分割层次结构

PS-Mamba 采用**渐进分割层次（Progressive Split Hierarchy）**策略，分三个阶段对特征图进行处理：

┌─────────────────────────────────────────────────────────────┐
│                    PS-Mamba 渐进分割流程                     │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   输入特征图 (H × W × C)                                    │
│          ↓                                                  │
│   ┌─────────────────────────────────────────────────┐       │
│   │  阶段1: 对半划分 (Halves)                        │       │
│   │  分割为 2 个区域，每个区域独立 Mamba 处理          │       │
│   └─────────────────────────────────────────────────┘       │
│          ↓                                                  │
│   ┌─────────────────────────────────────────────────┐       │
│   │  阶段2: 四象限划分 (Quadrants)                    │       │
│   │  分割为 4 个区域，进一步缩短递归深度                │       │
│   └─────────────────────────────────────────────────┘       │
│          ↓                                                  │
│   ┌─────────────────────────────────────────────────┐       │
│   │  阶段3: 八象限划分 (Octants)                      │       │
│   │  分割为 8 个区域，最大程度保留局部性               │       │
│   └─────────────────────────────────────────────────┘       │
│          ↓                                                  │
│   通道维度扩展 (+48) → Mamba处理 → 通道维度还原 (-48)        │
│          ↓                                                  │
│   跨尺度快捷通路（对称连接）                                  │
│          ↓                                                  │
│   特征融合与输出                                            │
│                                                             │
└─────────────────────────────────────────────────────────────┘

每个阶段的通道维度会扩展 48 以增强表示能力，在合并阶段对称地还原通道数，以保持平衡的快捷连接。

跨尺度快捷通路

为对抗长程衰减，PS-Mamba 引入了对称跨尺度快捷通路（Symmetric Cross-Scale Shortcut Pathways）。这些快捷通路绕过冗长的 Mamba 链，直接在不同层级之间传递全局结构信息。

快捷通路的设计遵循以下原则：

1. 直接传递：低频全局上下文信息可以通过快捷通路直接传递，无需经过完整的递归计算
2. 对称性：各层级的快捷连接呈对称分布，确保信息流的双向平衡
3. 低频增强：快捷通路主要传递低频分量，这些分量对全局一致性至关重要但在递归过程中最易衰减

数学上，跨尺度融合可以表示为：

$${\mathbf{F}}_{\text{fused}}={\mathcal{M}}_k\left(\text{Shortcut}\left({\mathbf{F}}_{\text{low}}\right),\text{SSM}\left({\mathbf{F}}_{\text{high}}\right)\right)$$

其中 ${\mathcal{M}}_k(\cdot )$ 为合并算子，$\text{Shortcut}(\cdot )$ 表示跨尺度快捷连接操作。

PS-Mamba 模块架构

整体结构

PS-Mamba 模块由三个主要组件构成：

┌─────────────────────────────────────────────────────────────┐
│                    PS-Mamba 模块                            │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   输入 x                                                    │
│      ↓                                                     │
│   ┌─────────────────────────────────────────────────┐       │
│   │  卷积预处理分支 (Conv Preprocessing)              │       │
│   │  f_conv = Conv_{3×3}(ReLU(Conv_{3×3}(x)))       │       │
│   │  提取短程纹理，增强局部连续性                      │       │
│   └─────────────────────────────────────────────────┘       │
│      ↓                                                     │
│   ┌─────────────────────────────────────────────────┐       │
│   │  渐进分割 Mamba 分支 (Progressive Split Mamba)   │       │
│   │  • 对半/四象限/八象限分割                        │       │
│   │  • 各区域独立 Mamba 处理                         │       │
│   │  • LayerNorm 归一化                             │       │
│   └─────────────────────────────────────────────────┘       │
│      ↓                                                     │
│   ┌─────────────────────────────────────────────────┐       │
│   │  注意力融合模块 (Attention-based Fusion)         │       │
│   │  • 门控机制: g = σ(W_2 δ(W_1 GAP(f_conv+f_m)))  │       │
│   │  • 混合特征: f_mix = g ⊙ f_conv + (1-g) ⊙ f_m   │       │
│   │  • 通道注意力 (CA) + 空间注意力 (SA)             │       │
│   └─────────────────────────────────────────────────┘       │
│      ↓                                                     │
│   残差连接: y = x + α · SA(CA(f_mix))                      │
│                                                             │
└─────────────────────────────────────────────────────────────┘

融合机制

PS-Mamba 采用内容自适应门控融合策略，将卷积分支（局部先验）与 Mamba 分支（全局建模）的特征动态融合：

门控向量 $\mathbf{g}$ 通过两层 MLP 从融合特征中学习：

$$\mathbf{g}=\sigma \left(W_2\cdot \delta \left(W_1\cdot \text{GAP}({\mathbf{f}}_{\text{conv}}+{\mathbf{f}}_m)\right)\right)$$

混合特征：

$${\mathbf{f}}_{\text{mix}}=\mathbf{g}\odot {\mathbf{f}}_{\text{conv}}+(1-\mathbf{g})\odot {\mathbf{f}}_m$$

最终的注意力精炼：

$$\mathbf{y}=x+\alpha \cdot \text{SA}\left(\text{CA}({\mathbf{f}}_{\text{mix}})\right)$$

其中 $\alpha$ 控制残差连接的强度，CA 为通道注意力，SA 为空间注意力。

复杂度分析

设输入特征图为 $\mathbf{X}\in {\mathbb{R}}^{H\times W\times C}$：

方法	序列长度	递归深度	计算复杂度
标准 Mamba	$L=HW$	$O(L)$	$O(L)$
PS-Mamba	${\sum }_j{L}_j=HW$	$O(\max L_j)$	$O(HW)$

PS-Mamba 保持了 SSM 的线性复杂度特性，但有效递归深度从 $O(HW)$ 降低到 $O(\max L_j)$，显著改善了信息传播的稳定性。

训练目标

PS-Mamba 采用统一的 L1 损失函数监督最终输出：

$$\mathcal{L}=\Vert y-\hat{y}{\Vert }_1$$

对于去噪任务，使用 Charbonnier 损失（可微 L1 变体）：

$$\mathcal{L}=\sqrt{\Vert y-\hat{y}{\Vert }_2^2+{ϵ}^2}$$

其中 $ϵ=10^{-3}$ 用于数值稳定性。

实验结果

PS-Mamba 在多个图像恢复基准任务上取得了优异表现：

轻量级超分辨率

方法	放大倍数	参数量	Set5 PSNR	Urban100 PSNR
SwinIR-light	×2	910K	38.14	32.76
MambaIR-light	×2	905K	38.13	32.85
MambaIRv2-light	×2	774K	38.26	33.26
PS-Mamba light	×2	970K	38.31	33.37

JPEG 伪影去除

方法	Classic-5 PSNR	LIVE1 PSNR
SwinIR	29.04	29.46
MambaIR	29.11	29.52
PS-Mamba	29.24	29.68

图像去噪

方法	Set12 PSNR (σ=15)	Set68 PSNR (σ=25)
SwinIR	33.85	30.94
MambaIR	33.79	30.89
PS-Mamba	33.98	31.08

实验结果表明，PS-Mamba 在保持线性复杂度的同时，在各项图像恢复任务上均优于现有的 Mamba-based 和 Transformer-based 方法。

与其他 Vision Mamba 方法的对比

PS-Mamba 与同期提出的其他 Vision Mamba 改进方法有着不同的设计理念：

方法	核心创新	局部性处理	全局建模
Vim	双向 SSM	无专门处理	双向扫描
VMamba	十字扫描	四方向独立处理	方向融合
LocalMamba	窗口限制	固定窗口划分	有限
MaIR	S形嵌套扫描	条纹划分	序列混洗注意力
PS-Mamba	渐进分割	几何一致性划分	跨尺度快捷通路

PS-Mamba 的独特之处在于：它不是通过改变扫描顺序来部分缓解局部性失真，而是通过先划分再处理的策略，从根本上保证了邻域完整性。

总结与展望

PS-Mamba 提出了一种针对图像恢复任务的拓扑感知层级状态空间框架，通过以下三大核心贡献解决了将 Mamba 应用于 2D 图像的根本性挑战：

1. 渐进分割状态空间模块：通过对半、四象限、八象限的渐进分割层次结构，在保持线性复杂度的同时，大幅缩短了每个区域的递归深度
2. 几何一致性划分：在状态空间处理之前保留空间拓扑结构，确保邻域像素保持空间邻接性
3. 对称跨尺度快捷通路：直接传递低频全局上下文，有效对抗长程衰减，稳定梯度流

这些设计使 PS-Mamba 能够在细粒度局部结构保持与长程空间一致性建模之间取得良好平衡，为基于状态空间模型的图像恢复开辟了新的方向。

参考文献

Footnotes

1. Hassanin M, Moustafa N, Deng W, Radwan I. Progressive Split Mamba: Effective State Space Modelling for Image Restoration. arXiv:2603.09171, 2026. ↩

2. 相关研究可参见 Vision Mamba 综述 中关于 Mamba 在视觉任务中局限性的讨论。 ↩