Learned image compression possesses a unique challenge when incorporating non-differentiable quantization into the gradient-based training of the networks. Several quantization surrogates have been proposed to fulfill the training, but they were not systematically justified from a theoretical perspective. We fill this gap by contrasting uniform scalar quantization, the most widely used category with rounding being its simplest case, and its training surrogates. In principle, we find two factors crucial: one is the discrepancy between the surrogate and rounding, leading to train-test mismatch; the other is gradient estimation risk due to the surrogate, which consists of bias and variance of the gradient estimation. Our analyses and simulations imply that there is a tradeoff between the train-test mismatch and the gradient estimation risk, and the tradeoff varies across different network structures. Motivated by these analyses, we present a method based on stochastic uniform annealing, which has an adjustable temperature coefficient to control the tradeoff. Moreover, our analyses enlighten us as to two subtle tricks: one is to set an appropriate lower bound for the variance parameter of the estimated quantized latent distribution, which effectively reduces the train-test mismatch; the other is to use zero-center quantization with partial stop-gradient, which reduces the gradient estimation variance and thus stabilize the training. Our method with the tricks is verified to outperform the existing practices of quantization surrogates on a variety of representative image compression networks.

翻译：学习型图像压缩在将不可微量化融入基于梯度的网络训练时面临独特挑战。现有研究提出了多种量化替代方法以完成训练，但均未从理论角度进行系统论证。本文通过对比均匀标量量化（最广泛使用的量化类别，其最简单形式为取整操作）及其训练替代方法，填补了这一理论空白。原则上，我们确定了两个关键因素：一是替代方法与取整之间的差异导致训练-测试不匹配；二是因替代方法产生的梯度估计风险（包含梯度估计的偏差与方差）。我们的分析与模拟表明，训练-测试不匹配与梯度估计风险之间存在权衡关系，且该权衡随网络结构变化。基于上述分析，我们提出一种基于随机均匀退火的方法，该方法通过可调温度系数控制权衡。此外，我们的分析揭示了两项精妙技巧：其一，为估计的量化潜变量分布方差参数设置适当下限，有效降低训练-测试不匹配；其二，采用具有部分梯度停止的零中心量化，可降低梯度估计方差从而稳定训练。实验证明，我们提出的方法及其配套技巧在多种代表性图像压缩网络上优于现有量化替代方法。