ICML 2023 Workshop | 使用量化整流器的神经图像压缩

用户1324186

发布于 2023-09-27 14:56:54

3150

题目：Neural Image Compression with Quantization Rectifier 作者：Wei Luo, Bo Chen 来源：ICML 2023 Workshop 文章链接：https://openreview.net/forum?id=IE9LsQ8SDx 整理：李江川神经图像压缩已被证明在率失真性能方面优于传统图像编解码器。然而，量化在压缩过程中引入误差，这会降低压缩图像的质量。本文提出了一种新的图像压缩量化整流器（QR）方法，该方法利用图像特征相关性来减轻量化的影响。作者设计了一种神经网络架构，从量化的特征中恢复量化前的特征，保留特征的表现力以获得更好的图像重建质量。作者同时提出了一种 soft-to-predictive（STP）训练方法，将 QR 集成到现有的神经图像编解码器中，并在 Kodak 数据集上进行测试。结果表明 QR 的编码效率得到了一致的提高，运行时间的增加可以忽略不计。

引言

量化通过将连续值映射到用于熵编码的有限离散值集来离散图像特征以压缩图像。虽然目前基于神经网络的图像压缩采用的量化方法解决了训练测试不匹配的问题，但量化对潜在特征的随机影响仍未解决。量化将连续值统一映射到单个离散值，根据特征可变性引入不同程度的噪声。例如，从

[-0.5,0.5)

到零的特征量化引入了

(-0.5,0.5]

范围内的噪声。更重要的是，量化以一种不可预测的方式改变了潜在特征的表现力。基于此，我们提出了一种新的量化整流器（QR），它利用图像中的空间相关性来减轻量化的影响。具体来说，QR 从量化的特征中重建未量化的特征。为了将 QR 无缝集成到神经图像编解码器中，我们引入了一种 STP 训练方法。在这里，我们首先对原始图像压缩模型进行端到端的训练，直到收敛。然后，我们冻结编码器网络，并优化解码器网络和 QR 网络。QR 弥合了原始特征和量化特征之间的差距，保留了特征的表现力，从而提高了图像重建质量。我们将提出的方法纳入了之前 SOTA 的基于神经网络的压缩方法中。在不影响比特率的情况下，我们将所有基线模型的表现提高了0.21 dB（PSNR）和 0.25 dB（MS-SSIM）。QR 是轻量级的，大多数基线模型的运行时间增加很小（0.7-5.4%）。本研究的贡献总结如下： • 我们提出了QR，这是一种通过预测校正量化图像特征的方法，保留了特征的表现力并提高了编码效率。 • 我们开发了 STP 训练程序和超参数探索算法，实现了 QR 与现有神经图像编解码器的无缝集成。 • 我们在 SOTA 的神经图像编解码器上广泛评估了 QR，这些编解码器始终如一地证明了 QR 的优越性。

提出的方法

量化整流器QR

在大多数现有的基于神经网络的图像压缩模型中，图像压缩的一般过程可以用以下公式表示：

y=g_a(x;\phi) \quad (1)

\hat{y}=Q(y) \quad (2)

\hat{x}=g_s(\hat{y};\theta) \quad (3)

其中，

和

\hat{x}

代表原始图像和重建的图像，

和

\hat{y}

分别代表潜在特征和量化后的潜在特征，

代表量化过程，

\phi

和

\theta

代表了编码器和解码器的参数。在推理过程中，

一般采用四舍五入操作，但由于这个过程是不可微分的，因此在训练时，通过加入

[-0.5,0.5]

的均匀噪声来实现。整个网络在训练时采用 RD 损失作为

Loss

：

\mathcal{L}_{\phi,\theta}=\mathcal{R}(\hat{y})+\lambda\mathcal{D}(x,\hat{x}) \quad (4)

现有的方法普遍采用超先验网络的设计，作者在这里将其省略了，当然，这种简化并不影响方法的普遍性。为了减轻量化的随机影响，我们引入了量化整流器（QR）。QR 网络位于量化操作和解码器之间，我们的见解是利用图像特征中的空间相关性，即使在有噪声的情况下也能恢复自身。受扩散模型在图像去噪方面的成功的启发，我们设计了如图 1 所示的 QR 网络，它由卷积层、残差块和注意力层组成。这些层在空间上关联量化特征

\hat{y}

，最后通过一个卷积生成特征的残差并将其添加到量化特征

\hat{y}

上，生成校正后的特征

\tilde{y}

。与原始的扩散模型相比，我们的网络在有效的同时配置了更少的层，以提高效率。QR 是一个多功能模块，可以无缝集成到任何神经图像压缩方法中，不需要对原始图像压缩模型的编码器和解码器组件进行显著修改。

图1 QR网络结构

加入 QR 后，图像压缩的过程可以用以下公式来表述：

\tilde{y}=QR(\hat{y}) \quad (5)

\hat{x}=g_s(\tilde{y};\theta) \quad (6)

为了能够端到端的训练 QR，需要对

Loss

函数进行修改：

\mathcal{L}_{\phi,\theta} = \mathcal{R}(\hat{y}) + \lambda\mathcal{D}(x,\hat{x}) + \alpha\mathcal{D}^f(y,\tilde{y}) \quad (7)

其中

\alpha

是控制特征损失的参数，关于

\alpha

参数的选取将在下一节具体介绍。

\mathcal{D}^f(y,\tilde{y})

是原始图像特征和量化图像特征之间的距离的度量，即特征距离。特征距离的设计对 QR 的学习至关重要，我们考虑了四个常用的距离，如

距离、

距离、平滑

距离和余弦相似性。我们的经验研究表明，最小化

距离可以使 QR 获得最好的生成质量。

STP 训练策略

训练是将 QR 集成到神经图像编解码器中的重要步骤。我们发现，由于编解码器和 QR 的训练是相互依赖的，直接进行端到端的训练是次优的。首先，QR 的学习依赖于其输入潜在特征的稳定性。其次，潜在特征的稳定性取决于 QR 网络的收敛性。即使潜在特征中的轻微扰动也会扰乱 QR 网络的训练过程。因此，受干扰的 QR 网络将进一步影响潜在特征的稳定性。在这样一个恶性循环中，潜在特征和 QR 网络不断收敛，整个训练过程是次优的。

为了解决次优训练问题，我们开发了一种由软训练和预测训练组成的 STP 训练策略。在软训练阶段，基于公式 (1)、(2) 和 (3) 来重建图像。同时，使用公式 (7) 对编解码器和 QR 网络进行优化，这里量化操作是通过添加均匀噪声来实现的。尽管这里没有采用公式 (5) 和 (6) 来生成

\hat{x}

，但仍将学习 QR 网络来预测特征，这将为下一阶段预热。

在预测训练阶段，将使用公式 (1)、(2)、(5) 和 (6) 来生成

\hat{x}

。同时编码器将被固定不再优化，采用四舍五入做量化而不是添加均匀噪声。在这一阶段将采用新的损失函数：

\mathcal{L}_{\theta} = \mathcal{D}(x,\hat{x}) + \alpha\mathcal{D}^f(y,\tilde{y}) \quad (8)

在这一阶段中，只有解码器和 QR 网络被优化，同时比特率

\mathcal{R}

被省略。因为编码器被固定后，潜在特征和比特率保持不变，这稳定了 QR 网络的训练。

学习参数

\alpha

的选取是十分重要的，它会直接影响到 QR 网络的效果，而且

\alpha

的最优值在不同的模型和压缩质量之间也会不断变化。图 2 比较了在不同学习参数下，图像压缩模型在 Kodak 数据集上测试的性能表现。在

\alpha = 10^{-3}

时，图像质量一开始就处于较高水平，并在几个 epoch 内快速收敛。当

\alpha = 10^{-6}

时，模型需要很长时间才能收敛。而当

\alpha = 10^{2}

时，QR 网络和解码器无法收敛，反而降低了图像的质量。为了解决这个问题，我们引入了一种学习参数探索方法，该方法可以自动找到不同模型和压缩质量下的最佳学习参数。

图2 不同

\alpha

下，模型性能表现

我们的一个关键发现是，存在一个最优学习参数，增加或减少它只会单调降低编码效率。基于这一发现，我们对特定模型和压缩质量的探索策略进行了如下描述： i）以初始学习参数

\alpha = \alpha_{max}

开始探索。 ii）按照 STP 训练策略使用

\alpha

训练编解码器，直到损失（公式(7)）在

个连续时期内停止改善。 iii）将学习参数

\alpha

乘以

0.1

。 iv）如果学习参数不小于预定义的最低学习参数

\alpha_{min}

，则继续步骤 ii），否则停止探索。具体的，设定

M=3

，且在一个较小的数据集（Flickr）上进行探索以增加探索的效率，探索得到的

\alpha

在大数据集上具有一致性。

具体的探索结果如表 1 所示，展示了在不同基线模型和压缩质量下的结果。具体的，使用的图像压缩模型为：Factorized Prior，Scale Hyperprior ，Joint Hyperprior 和 Attention-based Joint Hyperprior，简写为 “Factorized”，“Scale”，“Joint” 和 “Attn”。表 1 中最佳的 PSNR 结果将被选为训练时的

\alpha

参数，注意到，当参数增加到选定值以上或减少到选定值以下时，PSNR 性能单调下降。由于数据精度的问题，PSNR 只舍入显示到小数点后两位，实际结果仍有小数点两位后的可区分度。根据经验发现，PSNR 和 MS-SSIM 在不同模型上的最优参数分布相似，只是 MS-SSIM 的最优参数比 PSNR 的最优参数小大约 10 倍。因此，在预测训练阶段，我们将表 1 中 PSNR 的最佳参数乘以

0.1

来训练 MS-SSIM 模型。

表1 学习参数

\alpha

的探索结果

实验结果

我们将提出的 QR 网络添加到了前面提到的四个图像压缩模型中，并将增强后的模型分别用 “Factorized+QR”，“Scale+QR”，“Joint+QR”和“Attn+QR” 表示。除此之外，还设定了不同的压缩质量

q \in \left \{1,2,3,4 \right\}

，

越大代表在基线模型训练时使用的损失函数(公式(4))中的

\lambda

越大。基线模型使用从 Vimeo90K 数据集中随机提取和裁剪的

256 \times 256

大小的图像块进行训练，增强后的模型在 ImageNet 数据集上进行预测阶段的训练。

图 3 和图 4 比较了不使用和使用所提出的 QR 网络的基线模型的 RD 性能表现。曲线上的每个点表示在 Kodak 图像数据集上不同压缩质量

q \in \left \{1,2,3,4 \right\}

上平均的 bpp 和失真。对于任何给定的基线模型，应用 QR 网络后，平均 bpp 值保持不变。QR 网络在 PSNR 和 MS-SSIM 两方面在不同压缩质量下一致地改进了所有基线模型。此外，对于相对更复杂的模型，例如 Attn，QR 网络显示出比 Factorized 等简单模型更大幅度的改进。我们推测，具有更多参数的更复杂的模型可以更好地利用 QR 网络的重建效果来获得更好的图像质量。此外，与利用 PSNR 相比，利用 MS-SSIM 的 QR 网络的改进更加明显。表 2 给出了性能提升的数值结果。

图3 在PSNR下模型性能表现

图4 在MS-SSIM下模型性能表现

表2 模型的平均和最大性能提升

除了比较模型的性能表现，我们还提出了一种新的度量，量化误差

\epsilon_Q

，来体现我们的方法在保持图像特征方面的能力。具体的，量化误差定义为：

\epsilon_Q = \begin{cases} \left \| \tilde{y}-y \right \|_2 &\text{if using QR} \\ \left \| \hat{y}-y \right \|_2 &\text{if not using QR.} \end{cases} \quad (9)

图 5 展示了在各种压缩质量下，与所有基线模型相比，QR 网络减少的量化误差的百分比。对于像 Attn 和 Joint 这样更复杂的模型，这种减少通常更为显著。同时，较低的压缩质量往往会放大这种减少效果。

图5 QR网络引起的量化误差减少

表 3 中，我们比较了 Kodak 数据集上基线模型及其增强版本的平均每帧处理时间。对于大多数基线模型（Attn、Joint和Scale），我们的方法略微增加了

0.7\%-5.4\%

的处理时间，而 Factorized 受到的影响更大，因为其网络模型本身比较简单，处理时间较短。

表3 模型处理时间对比

结论

我们介绍了一种量化整流器（QR）方法来增强神经图像压缩。QR 在量化之后利用图像中的空间相关性来预测特征，从而保持其特征一致性。我们的方法包括一种 STP 的训练方法，该方法允许将 QR 无缝地集成到现有的神经图像编解码器中。实验结果一致地证明了 QR 在各种 SOTA 的神经图像编解码器中的有效性。

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2023-09-27 10:16，如有侵权请联系 cloudcommunity@tencent.com 删除

量化