清华+华为+鹏程实验室等 开源MambaVC | 比CNN和Transformer提高9.3%、15.6%，同时计算量大减！

点击下方卡片，关注 「AI视界引擎」 公众号

学习型视觉压缩是多媒体领域的一个重要且活跃的任务。现有方法已经探索了基于CNN和Transformer的各种设计，以建模内容分布并消除冗余，其中在平衡有效性（即速率-失真权衡）和效率方面仍然存在挑战。近来，状态空间模型（SSMs）因其长距离建模能力和效率而显示出潜力。

受此启发，作者首次探索将SSMs应用于视觉压缩。作者引入了MambaVC，这是一个简单、强大且高效的基于SSM的压缩网络。MambaVC开发了一个视觉状态空间（VSS）块，在每个下采样后使用2D选择性扫描（2DSS）模块作为非线性激活函数，这有助于捕捉信息丰富的全局上下文并增强压缩。在压缩基准数据集上，MambaVC以较低的运算和内存开销实现了卓越的速率-失真性能。

具体来说，在Kodak数据集上，它分别比CNN和Transformer变体提高了9.3%和15.6%，同时减少了42%和24%的计算量，并节省了12%和71%的内存。

MambaVC在高分辨率图像上显示出更大的改进，突显了其在实际应用中的潜力和可扩展性。作者还提供了不同网络设计的全面比较，强调了MambaVC的优势。

代码可在https://github.com/QinSY123/2024-MambaVC获取。

1 Introduction

视觉压缩是多媒体处理中一个长期存在的问题。在过去的几十年里，经典标准[6; 7]长期占据主导地位。随着深度神经网络结构如 CNNs 和 Transformers 的出现，学习型压缩方法已经崭露头角，并显示出不断提高的性能，逐渐超过传统方法，引起了越来越多的兴趣。

视觉压缩的核心是神经网络设计，以消除冗余信息并捕捉内容分布，这自然呈现了在率失真优化与模型效率之间的两难选择。尽管基于CNN的方法在许多资源受限的情境中仍受欢迎，得益于硬件高效的卷积运算符，但它们局部的感受野限制了全局上下文建模能力，从而限制了压缩性能。相比之下，基于Transformer的方法[23; 35; 52; 53]在全局感知方面表现出色，借助注意力机制从而有利于减少冗余。然而，它们在计算和内存上的二次复杂度提出了效率问题。

尽管一些混合方法如 TCM 结合了 CNN 和 Transformer 以平衡压缩效果和效率，但这并不是进一步发展的可持续方向。与之前的工作不同，作者致力于探索超出工程权衡的解决方案，以应对这一问题。

最近，状态空间模型（SSMs），尤其是结构化变体（S4）[16]，已经得到广泛研究。Mamba[15]作为一项代表工作脱颖而出，其数据相关选择机制增强了关键信息的提取同时消除了输入中的不相关噪声。这提示基于Mamba的模型可以有效收集全局上下文，从而在压缩方面具有优势。此外，Mamba集成了结构化重参化技巧，并利用了硬件高效的并行扫描算法，确保在GPU上更快地进行训练和推理。这些引人注目的特点激励作者研究Mamba在视觉压缩方面的潜力。

在本文中，作者介绍了MambaVC，一个简单、强大且高效的视觉压缩网络，带有选择性状态空间。受到Liu等人[26]的启发，作者在神经压缩网络中的每次下采样后设计了一个“视觉状态空间”（VSS）块作为非线性激活函数，该块集成了一个专门的“2D选择性扫描”（2DSS）机制进行空间建模。2DSS沿着4个预定义的遍历路径并行执行选择性扫描，有助于捕捉全面的全局上下文，并促进有效且高效的压缩。

作者进行了大量的图像和视频基准数据集实验。无需任何花哨的技巧，MambaVC在计算和内存开销方面优于基于CNN和Transformer的对应方法，如图1(a)所示的一些方法。更令人鼓舞的是，作者展示了MambaVC在高分辨率图像压缩上表现出更强的性能，如图1(b)所示。这些有利的结果与SSM的高效长距离建模能力一致，揭示了其在许多重要但具有挑战性的应用中的潜力，例如压缩高清医学图像和传输高分辨率卫星图像。作者还从各个方面比较和分析不同的设计，包括空间冗余、有效的感受野以及压缩过程中的信息丢失，以便全面理解MambaVC的有效性。

总结来说，作者的贡献如下：

• 作者开发了MambaVC，这是第一个带有选择性状态空间的视觉压缩网络。设计的2DSS改进了全局上下文建模，并有助于有效和高效的压缩。
• 在基准数据集上的大量实验展示了MambaVC在图像和视频压缩方面的卓越性能和竞争力。这些强有力的结果突显了超越CNN和Transformer的压缩网络设计新方向。
• 作者展示了MambaVC在高分辨率压缩方面的特别有效性和可扩展性，提示了其在许多重要但具有挑战性的应用中的潜力。
• 作者彻底比较和分析不同的网络设计，从各个方面展示了MambaVC的优势，以验证并理解其有效性。

学习型视觉压缩 在过去十年中，学习型视觉压缩展示了显著的潜力并留下了深刻的印象。主流方法可以分为基于CNN和基于Transformer的方法。早期工作，如带有广义分割归一化（GDN）层的CNN [4; 5; 34]，在图像压缩方面取得了良好的性能。后来，注意力机制和残差块[9; 49; 50]被整合到VAE架构中。然而，有限的感受野限制了这些模型的进一步发展。随着Vision Transformers [10; 27]的爆炸式增长，基于Transformer的压缩模型[28; 35; 52; 53]表现出了强大的竞争力。然而，它们大量的计算和存储需求是令人畏惧的。最近的研究[25]试图结合这两种方法的优点，但如Figure 1(a)所示，导致了计算复杂性的增加。在模型性能与效率之间的权衡仍然是一个需要解决的紧迫问题。

状态空间模型 最近提出的结合深度学习来捕捉长序列数据动态和依赖关系的模型。LSSL [17] 首次利用线性状态空间方程来建模序列数据。随后，结构化状态空间序列模型（S4）[16] 使用线性状态空间进行上下文化，并在各种序列建模任务上表现出色，特别是在处理长序列时。基于此，许多模型[13; 31; 38] 被提出，而 Mamba [15] 凭借其数据依赖性和并行扫描脱颖而出。因此，许多工作将 Mamba 从自然语言处理（NLP）扩展到视觉领域，如图像分类[26; 51]，多模态学习[36]以及其他[8; 18; 30]。然而，将 Mamba 应用于视觉压缩尚未被探索。在这项工作中，作者探讨如何将 Mamba 的成功转移到构建有效且高效的压缩模型上。

3 Method

Preliminaries: State-State Models and Mamba

状态空间模型（SSM）通过隐藏状态 将刺激 映射到响应 ，这里作者定义矩阵 为隐藏状态的演化映射，矩阵 和 分别作为隐藏状态的输入和读出映射。通常，作者可以通过线性常微分方程（ODEs）来制定这一过程：

现代SSM通过离散化近似这一连续时间ODE。具体来说，它们通过时间尺度 对连续参数 和 进行离散化，使用零阶保持技巧：

然后，方程（1）的离散化版本被重新制定如下：

Mamba [15] 进一步将数据依赖性整合到 ， 和 中，使能了一种输入感知的选择机制，以改善状态空间建模。虽然循环性质限制了完全并行的能力，但Mamba 巧妙地实现了结构重参化技巧和硬件效率并行扫描算法，以补偿整体效率。

The proposed MambaVC

3.2.1 Overview

作者在图2(a)中展示了MambaVC的架构。给定一个图像 ，作者首先使用编码器 和超编码器 分别获得潜在表示 和超潜在表示 ：

然后，量化的超潜在表示 用于熵编码以获得速率 ，其中 ，有一个学习到的分解先验 。

在解码器端，作者首先使用一个超解码器 来获得初始的均值和方差：

然后作者将潜在表示 划分为 个片段 并通过以下方式计算片段信息：

其中 表示通道自适应回归熵模型（CAM）[25]中的第 个网络，。作者将逐片估计的分布参数进行拼接，得到整体的 和 。作者计算 ，其中 服从参数为 的正态分布。

接下来，作者使用解码器 从量化的潜在表示 重建图像：

最后，作者优化以下训练目标：

其中 是拉格朗日乘数，用于控制率失真权衡。

3.2.2 Visual State Space (VSS) Block

受到刘等人[26]的启发，对于每个非线性变换 、、 和 ，作者在每次上采样或下采样操作后设计了一个视觉状态空间（VSS）块。图2（b）展示了这个结构。具体来说，每个VSS块由多个VSS层组成。遵循Mamba[15]，VSS层采用具有两个分支的门控结构，并在层归一化（LN）[2]之后。给定一个输入特征图 ，主分支通过以下方式处理它：

其中 表示层归一化。 表示2D选择性扫描模块，将在3.2.3节中详细阐述。 表示SiLU激活函数[37]。 表示深度卷积。 表示可学习的线性投影。

类似地，门控分支通过以下方式计算权重向量：

最后，两个分支结合在一起产生输出特征图：

其中 表示元素逐点乘积。#### 3.2.3 2D选择性扫描（2DSS）

原始Mamba[15]只能处理1D序列，不能直接应用于2D图像数据。为了有效地建模空间上下文，作者扩展了4种展开方式用于选择性扫描。具体来说，对于特征图 ，其中 表示特征图的第 行（）和第 $w$="" 列（$0\leq="" w<w$）的标记，展开模式定义为：<="" p=""></w$）的标记，展开模式定义为：<="">

(16) (17)

其中 ，$0\leq i<n$。$\mathbf{s}\_{1},\mathbf{s}\_{2},\mathbf{s}\_{3},\mathbf{s}\_{4}\in\mathbb{r}^{n\times c}$="" 是扩展并展平的标记序列。对于每个展平的标记序列，作者应用一个s6[15]算子进行选择性扫描，生成上下文标记序列="" $\mathbf{s}^{\prime}\_{1},\mathbf{s}^{\prime}\_{2},\mathbf{s}^{\prime}\_{3},\mathbf{s}^{\prime}\_{4="" }\in\mathbb{r}^{n\times="" c}$。<="" p="">

然后，作者通过以下折叠模式对上下文标记序列应用逆向操作：

(20) (21)

其中 表示 的扩展和变换后的特征图。

最后，作者将变换后的特征图合并以获得输出特征图：

3.2.3 Extension to Video Compression

作者还扩展了MambaVC到视频压缩领域，以探索其潜力。在这里，作者选择了一种著名的基于学习的视频压缩模型——尺度空间流（SSF）[1]作为扩展的基础框架。

作者将SSF中的基于CNN的变换在三部分（即I帧压缩、尺度空间流和残差）升级为开发的VSS块。作者将这个扩展称为MambaVC-SSF。

作者将在第4.4节展示并讨论实验结果。

4 Experiments

Experimental Setup

4.1.1 Datasets and Training Details

对于图像压缩，作者选择了文献[47]中的Flickr30k数据集，该数据集包含31,783张图片。每个模型都进行了2M次的训练。在前1.2M步中，每个批次包含8个随机裁剪的256×256图像；在接下来的0.8M步中，每个批次包含2个随机选择的512×512上采样图像。学习率从10^4开始，在1.8M步时降至10^5，最终在1.95M步时降至10^6。作者在率失真损失中使用了。

对于视频压缩，模型都在Vimeo-90k [45]上进行训练，先以10^4的学习率进行1M步训练，再以10^5的学习率进行额外的0.6M步训练。在第一阶段，每个批次包含8个随机裁剪的256×256图像；在第二阶段，每个批次包含8个随机裁剪的384×256图像。作者针对MSE失真度量优化视频模型。特别是，作者使用。受到[20; 32]的启发，在每次优化步骤中，作者分别处理原始和反向顺序的视频序列。

4.1.2 Baselines

作者对MambaVC在Kodak [12]，CLIC2020 [39]，JPEG-AI [22]和UHD [48]数据集上进行了两个方向的全面深入评估，这些数据集具有不同的图像分辨率。首先，作者将它与最先进的方法进行比较，包括MLIC+ [21]，Mixed [25]，GLLMM [14]，QResVAE [11]，ELIC [19]，STF [53]，WACNN [53]，Entroformer [35]，Swin-ChARM [52]，Invcompress [44]以及传统的编码方法BPG444 [6]和VTM-15.0 [7]。其次，作者验证了MambaVC在其卷积和Transformer变体方面的性能和效率优势。具体来说，作者将MambaVC中的VSS块分别替换为swin transformer和GDN层，将它们命名为SwinVC和ConvVC。详细结构展示在附录A中。

同时，作者在MCL-JCV [42]和UVG [33]数据集上评估了变体SSF，与标准的编解码器AVC(x264)，HEVC(x265)以及HEVC的测试模型实现（称为HEVC (HM)）进行了比较。所有方法将GOP大小固定为12。

Standard Image Compression

4.2.1 Comparison with the State-of-the-art Methods

在Kodak数据集[12]上的率失真性能展示在图3中。为了公平起见，所有展示的学习方法都是针对最小化MSE进行优化的。测试了PSNR和MS-SSIM，以证明MambaVC的鲁棒性。与之前最佳方法MLIC+[21]相比，作者的方法平均PSNR提高了0.1 dB，同时计算复杂度仅为前者的一半，内存开销为60%，如图1(a)所示。

4.2.2 Comparison of Variants

表1：各种变体的BD率（越低越好），以VTM-15.0为基准。

所有变体在Kodak [12]上的RD曲线如图11(a)所示。为了更清晰地比较不同变体之间的性能，图11(b)展示了相对于VTM-15.0实现等效PSNR时，各变体的比特率节省百分比。图11表明，MambaVC在多种情况下一致优于SwinVC和ConvVC。如前人工作所强调，SwinVC胜过ConvVC。

MambaVC和SwinVC相比于VTM-15.0均展现出更高的压缩效率，而ConvVC则有所不足。随着比特率的增加，SwinVC的性能优势略有减弱，而MambaVC则不受影响。在表1中，作者展示了与VTM-15.0相比，不同变体在四个数据集上的BD率。MambaVC实现了平均比特率节省13.35%，而SwinVC平均节省1.94%。相比之下，ConvVC平均消耗了4.76%的更多比特。值得注意的是，MambaVC是唯一一个在UHD [48]上超越VTM-15.0的变体，突显了其处理高分辨率图像的潜力，这将在下一节中讨论。更多详细信息见附录C.1。

高分辨率图像压缩

近期的工作已经证明了Mamba在长距离建模方面的优势。为了探索这一潜力在视觉压缩中的应用，作者以两种方式将作者的MambaVC与SwinVC和ConvVC在不同分辨率图像上进行比较。

具体来说，作者从UHD [48]的高分辨率图像开始，通过不同的因子下采样，创建了分布相同但大小不同的多组图像。如图1(b)所示，与其它变体相比，随着分辨率的增加，MambaVC节省的比特更多。为了减轻特定数据集分布的影响，作者在具有不同分辨率的四个数据集上进行测试。

如表2所示，MambaVC在高分辨率UHD [48]上的性能优势远大于在低分辨率Kodak [12]上的优势。对于大小相似的数据集，如CLIC2020 [39]和JPEG-AI [22]，性能优势相对一致。作者还记录了不同分辨率下计算成本的变化。如表3所示，随着图像大小的增加，计算差距从最初的0.23 TMACs和0.1 TMACs扩大到最终的12.96 TMACs和5.46 TMACs。这些结果表明MambaVC在压缩高分辨率图像方面具有明显优势。这种潜力可能会影响医学成像和卫星图像等专门领域的未来发展。

Video Compression with SSF Backbone

遵循Agustsson等人[1]的配置，作者在MCL-JCV[42]和UVG[33]数据集上评估了作者的方法。为了确保更全面的比较，作者还构建了基于CNN和Swin-Transformer的MambaVC-SSF的对标模型，分别称为SwinVC-SSF和ConvVC-SSF。不同模型的详细配置可以在第4.1.1节和第5节中找到。图4展示了MambaVC-SSF及其不同变体与传统方法的率失真（RD）曲线。基于mamba的模型优于其卷积和 Transformer 对标模型。然而，在视频压缩中性能的提升不如图像压缩那么显著，这可能是因为仅仅改变非线性变换结构不足以捕捉更多的冗余。此外，所有变体在MCL-JCV数据集上的性能仍低于HM，这表明还有很大的改进空间。

Computational and Memory Efficiencies

4.6.1 Latent Correlation and Distribution

学习到的视觉压缩冗余消除包括两个关键步骤：非线性编码转换，以及使用条件因子化高斯先验分布来去相关潜在变量 。

具体来说，前者将输入信号从图像域转换到特征域，而后者使用超网络学习潜在变量 的均值和方差 ，假设高斯分布，进一步减少相关性。随着各种相关性和冗余被消除，需要熵编码的信息更少，从而提高压缩效率。为此，作者可视化了 中每个空间像素与其周围位置之间的相关性，作者称之为潜在相关性。图5表明，与SwinVC和ConvVC相比，MambaVC在所有距离上的相关性都更低。

从理论上讲，去相关的潜在变量应遵循标准正态分布（SND）。为了验证这一点，作者拟合了不同方法的分布曲线，并计算了与SND的KL散度[24]，如图6所示。MambaVC的曲线明显更接近SND，且KL散度[24]更小，这表明基于Mamba的超网络可以更准确地学习 。作者还研究了超潜在相关性以及 与相关性之间的关系，如图14所示。

4.6.2 有效感受野

有效感受野（ERF）[29]表示神经网络中一个神经元“感知”的输入区域。更大的感受野使网络能够从更广泛的区域捕获相关信息。这一特性与视觉压缩中的非线性编码器完美契合，因为它通过特征提取和维度降低减少图像中的冗余。因此，作者非常关注MambaVC及其变体的感受野大小。

如图7所示，MambaVC是唯一具有全局ERF的模型，而ConvVC具有最小的感受野。这证实了在高分辨率场景中，MambaVC可以利用全局更多的像素来消除冗余，而SwinVC和ConvVC由于有限感受野，只能利用局部信息，导致了性能差异。

4.6.3 Quantize Deviation

在有损压缩中，量化是信息损失的主要来源。根据[44]，作者通过检查潜在变量与其量化版本之间的偏差来评估这种损失。图8展示了缩放后的偏差图及其具体数值。偏差图中的每个像素沿着通道维度缩放后绝对偏差的平均值。与MambaVC相比，SwinVC和ConvVC分别显示出平均信息损失增加了3.3%和17%。可视化结果也表明，在大多数位置上，MambaVC的信息损失更小（更深蓝色和更浅红色）。

5 Conclusions

在本文中，作者介绍了MambaVC，这是第一个基于状态空间模型的视觉压缩网络。MambaVC构建了一个视觉状态空间（VSS）块，并采用2D选择性扫描（2DSS）机制来改进全局上下文建模和内容压缩。

实验结果表明，MambaVC在保持计算和内存效率的同时，与卷积神经网络和Transformer变体相比，取得了更优的率失真性能。这些优势在高分辨率图像中更为显著，突显了MambaVC在实际应用中的潜力和可扩展性。

与其他设计相比，MambaVC显示出更强的冗余消除能力、更大的感受野和更低的量化损失，揭示了其在压缩方面的全面优势。作者希望MambaVC能为在压缩领域探索状态空间模型提供基础，并启发未来的工作。

作者的方法

MambaVC 详细架构已在第3.2节中描述。对于通道数和层数，作者分别设置为和。由于超高清图像的分辨率较高，会减慢推理速度，作者从UHD数据集中随机选择了20张图像，并将它们的中心部分裁剪到3328像素作为测试集。

MambaVC-SSF 对于SSF [1]中的编码器/解码器和超编码器/解码器，在每个上采样或下采样操作之后，除了生成重建图像或潜在层时，都会有一个VSS块，层数为。

Convolutional Variant

ConvVC ConvVC的结构如图A.2所示。具体来说，作者将VSS块替换为流行的GDN层[3]，这种层已经在高斯化自然图像的局部联合统计方面被证明是有效的。为了补偿卷积有限的有效感受野，作者将所有卷积核的大小设置为5。在架构方面，作者的基础模型的参数如下：。

Transformer Variant

SwinVC 在众多视觉Transformer变体中，作者选择Swin Transformer [10]作为网络组件，因其复杂度较低且建模能力卓越。如图A.3所示，所有实验中采用的层数量 和窗口大小 是一致的。对于通道数，作者设置 。

SwinVC-SSF 原始的下采样模块保持不变。类似于图像模型的架构，作者使用Swin Transformer [10]，但去除了所有的LayerNorm，并在其后添加了一个ReLU层。潜变量和超潜变量的通道数均设置为192。对于I帧压缩、尺度空间流和残差，作者分别采用8、4和8的窗口大小。层的数量与MambaVC-SSF相同。

图像压缩

BPG444：作者从http://bellard.org/bpg/获取BPG软件，并使用以下命令：

bpgenc -e x265 -q [质量] -f 444 -o [编码比特流文件] [输入图像文件] bpgdec -o [输出图像文件] [编码比特流文件]

VTM-15.0： VTM来源于https://vcgit.hhi.fraunhofer.de/jvet/VVCSoftware\_VTM。命令如下：

VVCSoftware_VTM/bin/EncoderAppStatic -i [输入YUV文件] -c [配置文件] -q [质量] -o /dev/null -b [编码比特流文件] -wdt 1976 -hpt 1312 -fr 1 -f 1 -mputChromaFormat=444 --InputBitDepth=8 --ConformanceWindowMode=1 VVCSoftware_VTM/bin/DecoderAppStatic -b [编码比特流文件] -o [输出YUV文件] -d 8

Video Compression

AVC(x264)

ffmpeg -y -pix_fmt yuv420p -s [分辨率] -r [帧率] -crf [质量] -i [输入yuv420原始视频] -c:v libx264 -present medium -tune zerolatency -x264-params "keyint=12:min-keyint=12:verbose=1" [输出mkv文件路径]

HEVC(x265)

ffmpeg -pix_fmt yuv420p -s [分辨率] -r [帧率] -tune zerolatency -y -i [输入视频] -c:v libx265 -present medium -crf [质量] -x265-params "keyint=12:min-keyint=12:verbose=1" [输出文件路径]

HEVC(HM)

HM/bin/TAppEncoderStatic -c HM/cfg/encoder_lowdelay_P_main.cfg -i [输入视频] --InputBitDepth=8 -wdt [宽度] -hgt [高度] -fr [帧率] -f [帧数] -o [输出视频] -b [编码比特流文件] -ip 12 -q [质量] 更多结果

Variant Visual Compression Performance on Different Datasets

在Kodak [12]，CLIC2020 [39]和JPEG-AI [22]上的额外率失真结果分别显示在图11、图12和图13中。

超前潜在相关性

图14展示了标准化先验潜在的空间相关性。水平比较不同的方法，MambaVC在所有上均表现出最佳性能。垂直比较结果，随着的减小，失真损失的比例减少，使得模型更专注于压缩比，从而消除更多的冗余。

Effective Receptive Field

在本节中，作者介绍了通过优化学习率调整和权重衰减正则化来提高网络训练效率的方法。

在3.1小节中，作者引入了一种动态学习率调整策略，该策略能够适应训练进度。

在3.2小节中，作者提出了一种新颖的权重衰减正则化方案，该方案考虑了损失景观的曲率。

参考

[1].MambaVC: Learned Visual Compression with Selective State Spaces.

点击上方卡片，关注 「AI视界引擎」 公众号