当Mamba遇上视觉压缩：首个具有选择性状态空间的视觉压缩网络

公众号新闻

2024-07-12 05:07

©PaperWeekly 原创 · 作者 | 秦诗雨

单位 | 清华大学硕士生

研究方向 | 图像压缩

摘要

在多媒体领域，学习型视觉压缩是一项重要的任务。现有的方法已经探索了各种基于 CNN 和 Transformer 的网络设计，以建模内容分布并消除冗余，但在平衡效率（即率失真权衡）和效能方面仍然是一个挑战。

最近，状态空间模型（SSM）因其长距离建模能力和效率而表现出潜力。受此启发，我们首次探索了 SSM 在视觉压缩中的应用。我们设计了 MambaVC，一个基于 SSM 的简单、强大且高效的压缩网络。MambaVC 开发了一个视觉状态空间（VSS）模块，并在每次下采样后使用 2D 选择扫描（2DSS）模块作为非线性激活函数，这有助于捕捉有信息的全局上下文并增强压缩效果代码已开源。

论文标题：

MambaVC: Learned Visual Compression with Selective State Spaces

论文地址：

https://arxiv.org/pdf/2405.15413

代码地址：

https://github.com/QinSY123/2024-MambaVC

Introduction

学习型视觉压缩的发展主要分为两个方向，一个方向是构建更有效的熵模型用于率估计；另一个方向是设计更好的神经网络以消除冗余信息并准确捕捉内容分布。

对于后者，我们可以根据网络结构的不同将现有模型分为以为 CNNs 基础的方法和以 Transformers 基础的方法。CNN 方法计算量更小但是感受野受限，无法捕捉更多的冗余；Transformer 的方法通过注意机制在全局感知方面表现出色，有效减少了冗余，然而计算效率十分低下。像 TCM 这样的混合方法试图融合二者平衡压缩效果和效率，但并不是一个可持续的长期解决方案。

最近，状态空间模型（SSMs）得到了广泛关注。它采用数据依赖的选择机制来增强关键信息提取，同时消除不相关的噪声。这种有效收集全局上下文的能力表明，基于 Mamba 的模型在压缩方面可能具有优势。Mamba 还结合了结构化重新参数化技术和硬件高效的并行扫描算法，确保了在 GPU 上的快速训练和推理。

因此，我们提出了 MambaVC，一种具有选择性状态空间的视觉压缩网络。我们将视觉状态空间模块作为神经压缩网络中每个下采样后的非线性激活函数。VSS 模块集成了专门的二维选择性扫描（2DSS）机制用于空间建模。该机制沿四条预定义的遍历路径并行执行选择性扫描，捕捉全面的全局上下文，促进有效和高效的压缩。

总而来说，我们的贡献如下：

1. 我们提出了 MambaVC，这是首个具有选择性状态空间的视觉压缩网络。其中二维选择性扫描机制（2DSS）改进了全局上下文建模，有助于实现有效和高效的压缩。

2. 在基准数据集上进行的广泛实验表明，MambaVC 在图像和视频压缩方面具有卓越的性能和计算效率。这些强劲的结果突显了在 CNN 和 Transformers 之外压缩网络设计的新方向。

3. 我们展示了 MambaVC 在高分辨率压缩中的特殊有效性和可扩展性，表明其在许多重要但具有挑战性的应用中具有潜力。

4. 我们对不同网络设计进行了全面比较和分析，展示了 MambaVC 相比于 CNNs 和 Transfromers 的方法在各方面的优势，以验证和理解其有效性。

方法

▲ 图1 网络结构

模型的编码器和解码器主要包括两部分，上/下采样和视觉状态空间模块（VSS Block）。对于下采样和上采样我们分别使用了卷积和反卷积操作处理。如图 1（b）所示，每个视觉状态空间模块由多个视觉状态空间层组成，在 Mamba 的基础上，VSS 层在层归一化（LN）后采用了一个带有两个分支的门控结构。给定输入特征图，主分支做如下处理：

类似地，门控分支计算权重向量 w：

最后，这两个分支结合起来生成输出特征图：

基础的 Mamba 只能处理一维序列，无法直接应用于二维图像数据。为了有效地建模空间上下文，我们扩展了四种展开方式用于选择性扫描。具体来说，对于特征图，表示特征图中第行和第列的标记，展开定义如下：

其中。是展开后的标记序列。我们应用 S6 操作进行选择性扫描生成上下文标记序列。然后，我们通过以下折叠模式对上下文标记序列应用反向操作：

最后，我们合并变换后的特征图以获得输出特征图：

实验

图像压缩

▲ 图 2 Kodak数据集上的率失真性能

与之前最好的方法 MLIC+ 相比，我们的方法在平均 PSNR 上提高了 0.1 dB，而只需一半的计算复杂度和 60% 的内存开销。

在不同分辨率的四个数据集上 MambaVC 的性能表现均比 CNN 变体和 Transformer 变体表现要更好。

视频压缩

▲ 图3 视频压缩性能

我们还将 MambaVC 扩展到视频压缩以探索其潜力。我们选择了学习型视频压缩模型 scale-space flow（SSF）作为扩展的基础框架，将其中的非线性变换替换为视觉状态空间模块。同时我们也比较了 Transformer 变体，可以看出 MambaVC-SSF 的性能超过了 CNN 变体和 Transformer 变体。

高分辨率场景压缩

▲ 图4 MambaVC 在 UHD 数据集上相对不同变体的 BD-rate

我们探究了 Mamba 在视觉压缩任务上的特有优势。我们将 UHD 中的高分辨率图像按不同的比例下采样，以创建具有相同分布但不同大小的多组图像。如图 4 所示，随着分辨率的增加，MambaVC 相比其他变体节省了更多的比特。

为了减轻特定数据集分布的影响，我们在四个不同分辨率的数据集上进行了测试。MambaVC 在高分辨率 UHD 上的性能优势明显大于在低分辨率 Kodak 上的性能优势。

计算存储开销

尽管 MLIC+ 由于采用了更先进的熵模型而带来了更高的计算成本，但它并未实现更优的性能。同时相比 CNN 和 Transformer 变体，计算量和存储量都达到了最优。

增益来源

为了更加清楚了解 Mamba 在视觉压缩领域相较于 CNN 和 Transformer 的优势，我们从潜在分布和相关性、有效感受野以及量化偏差方面进行了探究。

潜在分布和相关性

▲ 图5 潜在相关性

我们可视化了中每个空间像素与其周围位置之间的相关性，MambaVC 在所有距离上的相关性均低于 SwinVC 和 ConvVC。

▲ 图6 潜在分布

理论上，解相关的潜变量应遵循标准正态分布（SND），我们拟合了不同方法的分布曲线并计算了与 SND 的 KL 散度，MambaVC 的曲线明显更接近 SND，KL 散度更小，这表明基于 Mamba 的超先验网络可以更准确地学习 (µ,σ)。

有效感受野

▲ 图7 有效感受野

更大的感受野使得网络能够从更广的区域捕捉相关信息。这一特性与视觉压缩中的非线性编码器完美契合，因为它通过特征提取和降维来减少图像中的冗余。MambaVC 是唯一具有全局 ERF 的模型，而 ConvVC 的感受野最小。

量化偏差

在有损压缩中，量化是信息损失的主要来源。我们通过对比隐特征和量化解码后的之间偏差来评估这种损失。图 8 展示了缩放后的偏差图和具体数值。与 MambaVC 相比，SwinVC 和 ConvVC 的信息损失平均分别增加了 3.3% 和 17%。可视化结果还表明，在大多数位置上，MambaVC 的信息损失更小（更深的蓝色和更浅的红色）。

总结

在本文中，我们介绍了 MambaVC，这是第一个基于状态空间模型的视觉压缩网络。MambaVC 通过构建具有二维选择性扫描（2DSS）机制的视觉状态空间（VSS）块来改进全局上下文建模和内容压缩。

实验结果表明，与 CNN 和 Transformer 变体相比，MambaVC 在率失真性能上表现出色，同时保持了计算和内存效率。这些优势在高分辨率图像上更为显著，突显了 MambaVC 在实际应用中的潜力和可扩展性。与其他设计相比，MambaVC 展示了更强的冗余消除能力、更大的感受野和更低的量化损失，体现了其在压缩方面的综合优势。我们希望 MambaVC 能为在压缩中探索状态空间模型（SSMs）提供基础，并激发未来的研究工作。

更多阅读