[论文解读] Soft then Hard: Rethinking the Quantization in Neural Image Compression

Zongyu Guo, Zhizheng Zhang+2 more|arXiv (Cornell University)|Apr 12, 2021

Advanced Data Compression Techniques参考文献 41被引用 7

一句话总结

本文提出一种用于神经图像压缩的软量化后硬量化策略，首先通过添加均匀噪声进行软量化训练以学习富有表现力的潜在表征，随后通过硬量化进行微调以弥合训练-测试差异。此外，引入缩放的加性均匀噪声，通过一种新的变分上界实现自适应量化粒度，相较于最先进模型实现8.9%的BD-rate节省。

ABSTRACT

Quantization is one of the core components in lossy image compression. For neural image compression, end-to-end optimization requires differentiable approximations of quantization, which can generally be grouped into three categories: additive uniform noise, straight-through estimator and soft-to-hard annealing. Training with additive uniform noise approximates the quantization error variationally but suffers from the train-test mismatch. The other two methods do not encounter this mismatch but, as shown in this paper, hurt the rate-distortion performance since the latent representation ability is weakened. We thus propose a novel soft-then-hard quantization strategy for neural image compression that first learns an expressive latent space softly, then closes the train-test mismatch with hard quantization. In addition, beyond the fixed integer quantization, we apply scaled additive uniform noise to adaptively control the quantization granularity by deriving a new variational upper bound on actual rate. Experiments demonstrate that our proposed methods are easy to adopt, stable to train, and highly effective especially on complex compression models.

研究动机与目标

• 为解决神经图像压缩中因可微量化近似导致的训练-测试不匹配问题。
• 提升潜在表征能力，该能力因直通估计器和软到硬的退火策略而被削弱。
• 在消除其训练-测试差异的同时，保持加性均匀噪声带来的富有表现力的潜在空间学习能力。
• 通过可学习的噪声尺度实现自适应量化粒度，提升率失真效率。
• 开发一种即插即用的方法，可在不改变网络架构的前提下增强现有噪声松弛型压缩模型。

提出的方法

• 首先，使用加性均匀噪声作为量化软近似来训练编码器，以学习平滑且富有表现力的潜在空间。
• 然后，通过硬量化对解码器进行事后微调，使其与真实测试时的行为对齐，从而弥合训练-测试差异。
• 通过推导实际率的新型变分上界，引入缩放的加性均匀噪声，实现逐元素自适应的量化步长。
• 噪声尺度通过端到端方式学习，实现与图像纹理和通道特异性复杂度相关的上下文感知量化粒度。
• 该方法为即插即用：软-硬策略与缩放噪声可应用于任意现有噪声松弛型压缩模型。
• 该方法结合了基于噪声训练的表征表达能力与推理阶段硬量化的准确性，从而提升率失真性能。

实验结果

研究问题

• RQ1 为何尽管消除了训练-测试不匹配，直通估计器与软到硬退火策略仍会降低率失真性能？
• RQ2 我们能否在解决其训练-测试差异的同时，保留加性均匀噪声带来的富有表现力的潜在空间？
• RQ3 如何使量化粒度自适应于图像内容与潜在通道特性？
• RQ4 两阶段训练策略（先软训练，后硬微调）是否能在多种模型上提升压缩性能？
• RQ5 可学习的自适应噪声尺度是否能超越固定量化步长，进一步提升率失真效率？

主要发现

• 当软-硬量化策略与缩放噪声联合应用时，该方法在最先进模型（Cheng et al., 2020）基础上实现8.9%的BD-rate节省。
• 该方法在所有比特率下均提升率失真性能，即使在高比特率下，其性能稳定优于以往的退火方法。
• 应用缩放均匀噪声可实现逐元素自适应量化，学习到的噪声尺度与图像纹理及通道特异性复杂度高度相关。
• 该方法优于基于STE和SGA的微调策略，后者存在训练不稳定与性能次优的问题。
• 在MS-SSIM优化下，该方法在所有比特率下均提供感知质量更优的重建结果，且性能提升稳定。
• 该方法具有鲁棒性与泛化能力，当部署于其他高性能模型（如Guo et al., 2020）时，仍能实现最先进性能，超越VVC标准。