Mamba原理及在low-level vision的工作[持续更新]

Mamba原理

Mamba详细介绍和RNN、Transformer的架构可视化对比

选择性扫描（Selective Retain Information）：选择有关/无关信息

状态空间模型（SSM）

用以描述状态，并根据输入预测下一状态的模型。
其核心公式如下图

其流程可以表示为

输入的信号$X(t)$与矩阵B相乘，矩阵B描述输入信号对系统的影响

状态与矩阵A相乘，矩阵A描述内部状态是如何连接的，创建状态前应用矩阵A，状态更新之后，更新矩阵A

矩阵C将状态转换为输出

矩阵D计算输入到输出的直接信号，SSM通常说的是不包括跳转连接的部分

Mamba的选择性保留信息

SSM将A、B和C矩阵固定，由于线性时不变

1. 不能进行内容感知推理，在选择性复制任务中表现不佳。
2. 不能选择获取历史中先前token，提问Q回答历史信息A表现不佳。

如何选择性的保留信息？
Mamba通过结合输入的序列长度和batchsize大小，使得矩阵B和C甚至步长$\Delta$依赖于输入

较小的步长$\Delta$会使当前单词忽略长距离单词的影响，更多的依赖前一个的上下文。较大的步长$\Delta$使得模型更多地考虑当前输入对上下文的影响，而非依赖前一个上下文决定当前单词特征。

Mamba的扫描操作（The Scan Operation）

由于矩阵B、C和$\Delta$是动态的，无法通过卷积并行计算，只能递归表示，并且要在知道前一个状态的情况下才能预测下一个状态，因此Mamba设计了并行扫描算法

Mamba.py: 状态空间模型的并行扫描

硬件感知（Hardware-aware Algorithm）：合理存储中间结果

最近gpu的一个缺点是它们在小但高效的SRAM和大但效率稍低的DRAM之间的传输(IO)速度有限。在SRAM和DRAM之间频繁地复制信息成为瓶颈。
Mamba的DRAM和SRAM分配的具体实例如下:

中间状态不被保存，但对于反向传播计算梯度是必要的。作者重新计算了反向传递过程中的中间状态。尽管这看起来效率很低，但它比从相对较慢的DRAM读取所有这些中间状态的成本要低得多。

Mamba架构

Mamba在底层视觉的工作

参考Awesome-Mamba-in-Low-Level-Vision

图像修复

MambaIR: A Simple Baseline for Image Restoration with State-Space Model，ECCV2024

MambaIR：使用状态空间模型进行图像恢复的简单baseline
paper&code

现有的修复主干网络往往面临全局感受野和高效计算之间的困境，阻碍了其在实践中的应用。最近，选择性结构化状态空间模型，尤其是改进版 Mamba，在线性复杂度的长程依赖建模方面显示出了巨大的潜力，为解决该困境提供了一种方法。标准 Mamba 在低级视觉方面仍面临局部像素遗忘、通道冗余等挑战。在这项工作中，我们引入了一个名为 MambaIR 的简单但有效的基线，它引入了局部增强和通道关注，以改善原版 Mamba。MambaIR 利用了局部像素相似性，减少了通道冗余。

空间和通道

VmambaIR: Visual State Space Model for Image Restoration

VmambaIR：用于图像恢复的视觉状态空间模型
paper&code

VmambaIR将具有线性复杂性的状态空间模型（SSM）引入全面的图像恢复任务。利用 Unet 架构来堆叠提出的全面选择性扫描 （OSS）块，它由OSS模块和**高效前馈网络(EFFN)**组成。提出的全向选择性扫描机制通过对所有六个方向的图像信息流进行有效建模，克服了 SSM 的单向建模限制。

扫描局限

Serpent: Scalable and Efficient Image Restoration via Multi-scale Structured State Space Models

Serpent：通过多尺度结构化状态空间模型进行可扩展且高效的图像恢复
paper&code

Serpent是一种用于高分辨率图像恢复的高效架构，它结合了状态空间模型（SSM）的最新进展及其核心计算块中的多尺度信号处理。SSM 最初是为序列建模而引入的，可以保持一个全局感受野，并在输入大小上具有良好的线性缩放。我们提出了一种新颖的分层架构，其灵感来自传统的信号处理原理，它将输入图像转换为序列集合，并以多尺度方式处理它们。

多尺度

CU-Mamba: Selective State Space Models with Channel Learning for Image Restoration

CU-Mamba： 用于图像恢复的具有通道学习的选择性状态空间模型
paper&code

提出了通道感知 U 形曼巴（CU-Mamba）模型，该模型将双状态空间模型（SSM）框架整合到 U-Net 架构中。CU-Mamba采用空间SSM模块进行全局上下文编码，并使用通道 SSM 组件来保留通道相关特征，两者的线性计算复杂性都与特征图大小相关。

空间和通道

Hi-Mamba: Hierarchical Mamba for Efficient Image Super-Resolution

Hi-Mamba：分层 Mamba，实现高效的图像超分辨率
2024 年 10 月 14 日
paper&code
SSM 的连续性需要在不同方向进行多次扫描，以补偿在将图像展开为 1D 序列时空间依赖性的损失。这种多方向扫描策略会显著增加计算开销，并且无法承受高分辨率图像处理。为了解决这个问题，我们提出了一种新的分层 Mamba 网络，即 Hi-Mamba，用于图像超分辨率（SR）。Hi-Mamba 由两个关键设计组成：（1） 由本地 SSM （L-SSM） 和区域 SSM （R-SSM） 组装的分层 Mamba 块 （HMB） 均采用单向扫描，聚合多尺度表示以增强上下文建模能力。（2） 方向交替分层曼巴组 （DA-HMG） 将同分异构体单向扫描分配到级联HMB中，以丰富空间关系建模。

内容和通道
扫描限制

MambaIRv2: Attentive State Space Restoration

MambaIRv2：专注状态空间恢复
2024 年 11 月 22 日
paper&code

Manba 固有的因果建模限制，即每个token仅取决于扫描序列中的前身，限制了整个图像中像素的充分利用，从而在图像恢复方面提出了新的挑战。这项工作提出了 MambaIRv2，它为 Mamba 配备了类似于 ViTs 的非因果建模能力，以达到专注的状态空间恢复模型。具体来说，所提出的注意力状态空间方程允许在扫描序列之外进行关注，并仅通过一次扫描即可促进图像展开。此外，我们进一步引入了一种语义导向的相邻机制，以鼓励距离遥远但相似的像素之间的交互。

扫描限制
空间

MaIR: A Locality- and Continuity-Preserving Mamba for Image Restoration

MaIR：用于图像恢复的保持局部性和连续性的 Mamba
2024 年 12 月 28 日
paper&code

i） 自然图像中固有的局部关系和空间连续性，ii） 序列之间的差异通过完全不同的方式展开。为了克服这些缺点，探讨了基于Mamba的恢复方法中的两个问题：i） 如何设计一种扫描策略，在促进恢复的同时保持局部性和连续性，以及 ii） 如何聚合以完全不同的方式展开的不同序列。

提出了一种新的基于 Mamba 的图像恢复模型 （MaIR），它由嵌套 S 形扫描策略 （NSS） 和序列随机注意力块（SSA）组成。具体来说，NSS 分别通过基于条带的扫描区域和 S 形扫描路径保留输入图像的位置性和连续性。SSA 通过计算不同序列对应通道内的注意力权重来聚合序列。


扫描限制

Contrast: A Hybrid Architecture of Transformers and State Space Models for Low-Level Vision

用于底层视觉的 transformer 和状态空间模型的混合架构
2025 年 1 月 23 日
paper&code
Transformer 具有强大的全局上下文建模功能，然而它们的二次计算复杂性需要使用基于窗口的注意力机制，这限制了感受野并限制了有效的上下文扩展。Mamba 架构能够避免窗口机制并保持较大的感受野，然而，当需要高像素级精度时，Mamba 在处理长上下文依赖关系时面临挑战。我们提出了Contrast，这是一种混合 SR 模型，它结合了 Convolutional、Transformer 和 State Space 组件，有效地融合了Transformer和 Mamba的优势，以解决它们各自的局限性。通过集成 transformer 和状态空间机制，Contrast 弥补了每种方法的缺点，增强了全局上下文建模和像素级精度。

混合Trans+Mamba

Directing Mamba to Complex Textures: An Efficient Texture-Aware State Space Model for Image Restoration

将 Mamba 引导至复杂纹理：一种用于图像恢复的高效纹理感知状态空间模型
2025 年 1 月 27 日
paper&code
现有方法主要基于 CNN、Transformer 或其混合方法，这些方法通常难以有效地对长期依赖关系进行建模，并且在很大程度上忽视了图像退化的空间特性（纹理更丰富的区域往往会遭受更严重的损坏），因此很难在恢复质量和效率之间实现最佳权衡。

上图中，（a）-（b）计算DIV2K和MANGATE109数据集中的高分辨率和低分辨率图像之间的残差误差图，发现降解程度在图像的不同区域之间变化。（c）我们将图像贴片分为10组，这些组通过纹理复杂性（通过统计变量衡量）排序，并计算所有人类109个数据集中每组的平均PSNR。可以观察到，质地较丰富的区域遭受更严重的降解，导致PSNR值较低。

为了解决这些问题，提出了一种新颖的纹理感知图像恢复方法 TAMambaIR，它可以同时感知图像纹理并在性能和效率之间实现权衡。具体来说，我们引入了一种新的纹理感知状态空间模型，它通过改进状态空间方程的过渡矩阵并专注于具有复杂纹理的区域来增强纹理感知并提高效率。此外，我们设计了一个Multi-Directional Perception Block来改善多向感受野，同时保持较低的计算开销。

MatIR: A Hybrid Mamba-Transformer Image Restoration Model

MatIR：一种混合 Mamba-Transformer 图像修复模型
2025 年 1 月 30 日
paper&code
Mamba 目前在上下文学习能力方面落后于 Transformers。为了克服这两种模型的局限性，我们提出了一种称为 MatIR的Mamba-Transformer 混合图像修复模型。

（a）Transformer-based的有效感受野（ERF）的可视化[15，35]。上图显示了上下文近邻感受野中比Mamba大的优点。（b）Mamba基本模型的ERF的可视化[33]。下图显示了长序列线性接受场中比Transformer大的优点。 （c）我们提出的Matir：一种结合Mamba- Transformer图像恢复模型实现了更重要的感受野。

MatIR 对 Transformer 层和 Mamba 层的块进行交叉循环以提取特征，从而充分利用这两种架构的优势。在 Mamba 模块中，我们引入了图像修复状态空间 （IRSS） 模块，该模块沿 4 条扫描路径遍历，以实现对长序列数据的高效处理。在 Transformer 模块中，将基于三角窗口的局部注意力与基于通道的全局注意力相结合，在更广泛的图像像素范围内有效地激活注意力机制。

RestorMamba: An Enhanced Synergistic State Space Model for Image Restoration, ICASSP

RestorMamba：用于图像恢复的增强协同状态空间模型
2025 年 3 月 7 日
paper&code
这种方法在网络内结合了高效的远程依赖性建模，特别适用于高纹理和高分辨率图像恢复场景的复杂性。为了在保持全局感受野的同时从更广泛的环境中获益，我们设计了两个关键模块：Skip Scan 和 Enhanced Synergistic Mamba （ESM）Block。