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

©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）机