支持码控的学习型图像压缩

题目: Practical Learned Image Compression with Rate Control 内容整理：胡锦堋 这篇文章主要介绍了图像压缩中的新技术。首先，我们设计了一种邻居补丁注意机制，该机制明确地使用空间注意分支中的附近补丁信息，以便变换网络更好地利用局部冗余。其次，我们设计了一个空间分配上下文模型，以有效利用特征之间的相关性，同时实现高度并行解码。第三，我们提出了一种新颖的速率控制模块来实现可变速率编码和特定速率编码。

• 讨论

引言

图像压缩是一种广泛使用的技术。在过去的几十年里，传统的图像压缩方法发挥了重要作用。JPEG 是一种基本的图像压缩方法，自20世纪90年代以来一直被使用，并且仍然是主流的压缩格式。JPEG 中应用了 DCT 变换和霍夫曼熵编码。后来，WEBP 、BPG 和 VVC 涉及更复杂的工具来增强速率失真性能。尽管这些强大的工具极大地提高了图像压缩性能，但手动设计的搜索空间和变换仍然可能限制性能。与传统的图像压缩方法相比，LIC是一种数据驱动的方法，不需要手动设计特定的规则。

基于 VAE 的 LIC 方法最吸引研究人员。超先验模型通过条件高斯分布极大地提高了率失真性能。然后提出了的顺序上下文模型来减少熵冗余并抑制 BPG，然而由于串行处理，效率不高。随后提出了一种称为棋盘上下文模型的并行空间上下文模型来加速解码速度。通道自回归熵模型还通过沿通道维度但在有限循环中解码特征来加速速度。不均匀分组的通道空间上下文模型进一步沿通道维度提取更多增益。除了上下文熵模型之外，非线性变换也是 LIC 方法中值得注意的部分。广义归一化变换 (GDN) 是一种优雅的归一化函数，非常适合对自然图像中的数据进行高斯化。简化的注意力模块增强了网络的非线性能力，基于窗口的注意力进一步将转换器引入到基于 CNN 的注意力模块中。所有这些建议使得非线性表现力更加强大，并按预期提高了压缩性能。

对于实际应用而言，LIC 中的速率自适应是一项重要的技术，使用单一模型来满足各种压缩比是 LIC 广泛使用的决定性因素之一，特别是在资源受限的场景中。提出了几种在LIC中实现速率自适应的解决方案。在一些带宽有限的场景下，码率控制策略对于图像和视频传输起着重要作用。速率控制的目的是在给定速率约束下实现最佳重建质量。R-

模型是 LIC 通过给定目标λ来控制速率的实用方法。模型是通过反向传播算法进行端到端训练的，因此它们可能更有潜力找到最佳表示和压缩规则。目前，主要有两种架构：一种是基于 RNN 的图像压缩，另一种是是基于变分自动编码器（VAE）的图像压缩。在 VAE 的变换方法中，是基于小波的，是基于变换器的，这与其他基于 CNN 的变换略有不同。

受上述研究的启发，我们提出了三个主要贡献，总结如下：

• 我们提出了邻居补丁注意模块（NPA），以更好地分析局部信息进行变换。NPA 包括一种新的局部感知机制，该机制通过移动卷积核的中心将附近补丁的信息视为输入源。这是一种设备友好的方法，可以轻松捕获显着结构。
• 我们进一步研究发现，当前的并行空间上下文模型由于掩码卷积而浪费了计算资源。因此，我们设计了一个空间分配上下文模型（SSC），其中包括一种新颖的张量排列来处理这个问题。结果表明，我们提出的 SSC 不仅优于并行上下文模型，而且优于顺序上下文模型。
• 我们设计了一种称为速率控制模块（RCM）的新颖机制来实现速率自适应和速率控制。RCM 使 LIC 能够通过一组标量参数和一个微型超网络来控制比特率，只需一点额外的编码时间成本和零 BD 速率增加。

方法

图1：紫色块（NPA）表示邻居补丁注意力，蓝色块（RCM）表示速率控制模块。橙色部分（CSSC）表示通道空间混洗上下文模型，它将通道自回归模型和SSC集成在一起。绿色块代表剩余块

如图 1 所示，我们的模型使用堆叠残差块、GDN 和 NPA 来增强非线性能力。SSC 用于增强压缩性能并实现并行解码，通道方式自回归策略还通过将通道均匀划分为四组来集成。我们的速率控制模块实现了速率自适应和比特率控制以获得所需的比特率。我们将变换对设置为

和

，超先验对设置为

和

，通道空间混洗上下文模型设置为 cssc，hyper-feed 是

的输出。

相邻块注意力机制

对于图像压缩任务，附近的像素通常具有很强的相关性。在图 2 中，NPA涉及附近 patch 的信息作为额外的输入源，并通过简洁的方式捕获周围的信息。与使用 Transformer 作为感知局部信息的基于窗口的注意力相比，NPA 仅移动卷积核的中心，使模型能够进一步增加感知局部相关性。实现上述机制的关键单元是 NP-Unit。

图2：NPA的结构

如图 3 所示，NP 单元将四个卷积核从原始位置移开。N1、N2、N3 和 N4 是四个新的卷积中心，NP 单元将连接每个移位卷积的输出作为 NPA 感知的附近信息的来源。因此，NPA 可以明确地知道相邻信息。与采用Transformer 实现局部感知注意力的基于窗口的注意力相比，NPA 更适用于当前的神经网络推理设备。

图3：NP-UNIT 中卷积核移位方式的图示。在左图中，灰色部分表示内核大小为 3×3 的卷积，绿色块是该内核的中心。在右图中，该内核的中心随后分别向四个方向移动

我们将图 4 中的第一个注意力模块的注意力尺度图可视化，该模块由 Sigmoid 操作激活。很明显，带有 NP-Unit 的注意力模块(b)比没有 NP-Unit 的注意力模块(c)在显着结构上的得分更高。捕获显着结构对于LIC的变换非常重要，NP-Unit 可以有效地感知结构，从而提高率失真性能。

图4 ：注意力图的可视化。(a) 是原始图像，(b) 和 (c) 是有/没有 NP-Unit 的模块的注意力图

为了揭示 NPA 的优越性，我们用简化的注意力模块和基于窗口的注意力模块 替换我们提出的模型中的 NPA。表 1 显示 NPA 在 Kodak 和 CLIC2020valid 图像集上表现最佳。

表1：NPA、SA、WBA 的 BD 速率比较

空间分配上下文模型

基于超先验架构，在LIC中设计的顺序上下文模型和并行棋盘上下文模型，但是卷积的屏蔽部分仍然是计算过程的一部分，这会导致计算能力的浪费。论文提出的 SSC 引入了一种新的张量排列方式，以充分利用卷积计算。该过程如图所示。SSC 由编码和解码过程的两个步骤组成。对于每个大小为 HxWxC 的 patch，SSC 将 patch 从空间维度分配到四个通道维度，切片并分组为两个子部分，每个具有 H/2×W/2×C×2 大小的子部分将被分配回 H× W×C/2 尺寸。作为第一个编码或解码部分收集的橙色和黄色块没有额外的上下文信息。图 1 中引用的 Hyper-feed 用于生成

和

以进行熵近似。作为第二编码或解码部分收集的蓝色和绿色块可以利用来自第一解码部分的上下文信息。超先验和上下文信息将被连接起来生成

和

以进行熵近似。

图 5：SSC 模型。SSC 对输入进行打乱，并收集橙色和黄色色块作为第一部分，分别收集蓝色和绿色色块作为第二部分。黑线是为第一部分生成相应 µ1 和 σ1 的第一步，橙色线是为第二部分生成相应的 µ1 和 σ1。

顺序上下文模型和棋盘上下文模型在熵参数模块中仅使用1x1内核大小的卷积，因为顺序上下文模型只能使用1x1内核大小的卷积，以保证要解码的特征可以使用适当的上下文信息。对于棋盘上下文模型，只有具有 1x1 内核大小的卷积才能确保上下文模型不会考虑未解码的特征。通常，与较大内核尺寸相比，具有 1x1 内核尺寸的卷积的数据拟合能力较弱。所提出的ssc通过预先将潜在特征分为两部分来避免这个缺点，SSC 可以在熵参数模块中应用具有 3x3 或 5x5 内核大小的卷积，以更好地提取上下文信息

为了揭示SSC的优越性，我们评估了联合上下文模型、棋盘上下文模型和SSC之间的率失真性能。我们通过用上面提到的上下文模型替换我们提出的模型中的 SSC 来进行比较。如表 2 所示，在具有与棋盘上下文模型相同的解码复杂度的情况下，SSC 在 Kodak 和 CLIC2020valid 图像集上取得了比顺序上下文模型更好的压缩性能。

表2：SSC、AR、CB 的 BD 速率比较

RCM

图6：RCM 的结构

如图 6 所示，RCM由两个子模块组成，一个是质量单元，另一个是微小网络，记为tiny

。质量单元是一个简单而有效的模块，其中包含一组称为质量因子的标量参数，而 tiny

充当原始合成先验

和上下文模型

的代理。为了实现速率自适应和速率控制，该模块有两个步骤。首先，在训练过程中，随机选取损失函数

来训练可变速率的模型。对于每个单独的

，潜在特征 y 在量化前乘以特定的品质因数，并在量化后除以该品质因数，所有这些品质因数的值都在 (0, 1) 范围内。质量因素是后期码率控制策略的基础。编码和解码中的速率自适应方程如下：

QF 表示品质因数，每个 y 对应的超潜在 z 在量化前都乘以之前相同的品质因数。速率自适应是通过使用不同的品质因数来实现的。

其次，tiny

通过拉普拉斯熵模型和分解熵模型进行熵估计，这两个熵模型估计潜在特征

̂和

的分布。我们发现，当在训练期间使用这个微型最小化熵时，最终的估计速率与原始大型网络的估计速率近似线性。在实践中，我们发现如果我们仅估计

并将

设置为零，则可以实现最低的速率误差。给定速率约束和输入，我们可以通过使用插值和二分搜索将品质因数调整到所需值来实现所需的速率。我们使用最小和最大的品质因数来找到输入的最大和最小可能的比特率，然后计算用于插值的比率

，bpp表示每像素的位数。

然后我们可以使用这个比率插入一个新的品质因数，

然后通过熵估计我们得到一个中间比特率。因此，我们再次计算比率并使用新比率进行插值，就像二分搜索过程一样，直到达到所需的比率。由于所需速率的品质因数应用于整个潜在特征，因此该过程不会导致 BD 速率增加。这个过程取决于熵估计，并且可以通过微型超级代理大大加速。通过将 z 与 y 分离，我们设法简化了重复熵估计的过程，因为我们不再需要每次都运行

从不同的潜在特征 y 中获取 z。现在 z 可以通过乘以相同的品质因数与 y 一起变化。

实验

设置

我们从 Open Images 数据集和 NIC 训练数据集中随机选择 200k 图像，处理图像数据的过程遵循 L3C 的方法。在训练过程中，我们随机将图像裁剪为256x256的大小，并将输入值归一化为[0,1]，如果标准差小于0.05，则简单的图像块将被丢弃。我们使用 Adam 优化器，批量大小为 16。失真损失针对均方误差（MSE）进行了优化。训练质量单元有两个阶段：在第一阶段，我们将

设置为 8192，并以学习率 0.0001 训练模型 1000k 次迭代。第二阶段，从(128,384,1024,3072,8196)中随机选取 lambda 来训练质量单元模型，其学习率和对应的迭代次数为(0.0001,0.00001)和（1500k，500k）。RCM 模块经过训练，可在所有其他过程完成后最小化熵。

评估时为了公平比较，我们通过平均 BPP、PSNR 和 MS-SSIM 来计算所有指标。我们在 Kodak 数据集上比较结果，更多比较结果可以在补充材料中找到。JPEG、JPEG2000、WEBP、BPG、VTM、自回归超先验模型 (Minnen2018)、通道自回归超先验模型 (Minnen2020)、具有注意力和高斯混合模型的自回归超先验（Cheng2020）、具有基于窗口注意力的通道自回归熵模型（Zou2022） 以及不均匀分组的空间通道上下文模型（He2022) 被包括在内以进行比较。

RD性能比较

在图 7 和表 3 中，我们提出的模型使用单个模型来实现可变速率，并在 Kodak 的 PSNR 和 MS-SSIM 指标中实现最佳速率失真性能。请注意，Cheng2020-anchor 是 CompressAI 提供的锚版本，而 Zou2022(CNN) 是基于 CNN 的模型，这比基于 Transformer 的版本稍好。VTM 的版本是 14.0。Zou2022（CNN）没有发布MSE优化模型的MSSSIM，因此我们不将其纳入比较。图 8 可视化了使用不同编解码器解码的重建图像，我们的图像明显保留了更多纹理和细节。

图7：我们的方法优于使用 PSNR 作为 Kodak 图像集上的图像质量指标的其他 LIC 方法和标准编解码器。

图8：与 VTM、BPG、WEBP 和 JPEG 的视觉质量比较。指标为 Bpp/MS-SSIM/PSNR。

推理速度对比

我们还评估了运行速度，如表 3 所示。提供的推理时间结果是 Kodak 图像集的平均编码/解码时间，这是在配备 NVIDIA GeForce RTX 2080 Ti 和 AMD Ryzen 7 3800X 的服务器上获得的 8 核 4.5 GHz。结果表明，我们的方法以实际的实际编码和解码速度实现了最佳的 RD 性能。

表3：RD 性能和推理时间。BD 速率是通过 Kodak 的 VTM-14.0 计算的。

评估RCM的性能

我们在三个图像集上测试码率控制的准确性，结果如表 4 所示。从表中我们可以看到，我们在 Kodak 上实现了约 1.89％ 的平均码率误差，1.87％和2.12 分别占 NIC 验证集 C 类和 B 类的百分比。在测试过程中，我们将 RCM 的目标 bpp 和估计 bpp 之间的差异阈值设置为 1%。RCM通过二分查找达到这个阈值要求，并确定具体的品质因数。通过我们的实验，我们发现通常 2 遍二分搜索就足以达到这个阈值。RCM估计的熵与实际熵模型估计的熵之间的差异引入了大于1%的速率误差。我们还测试了 RCM 的推理速度，发现它引入了大约 5% 的编码时间开销。

表4：所提出的 RCM 的几个目标 bpp 的平均速率控制误差（|Rfinal−Rtarget|/Rtarget× 100%）。测试时，我们挑选出目标速率超出给定模型的速率适应范围的测试样本

讨论

在本文中，我们提出了一种实用的学习图像压缩模型，具有三个主要贡献。首先，为了更好地捕获显着结构，我们提出了一种设备友好的本地感知注意机制，称为邻居补丁注意。然后我们继续设计一个没有掩码卷积的并行空间上下文模型，称为空间分配上下文模型。实验结果表明它优于顺序上下文模型。最后我们提出了一种新颖的速率控制模块，它可以同时实现速率自适应和精确的比特率控制。在最先进的率失真性能的前提下，我们提出的图像压缩方法具有快速的推理速度，并且可以准确地控制比特率。未来，我们将探索更高效、更有效的方法，以实现实时推理解码，并将我们的码率控制模块应用到学习视频压缩中。