Learnable Non-Uniform Quantization With Sampling-Based Optimization for Variable-Rate Learned Image Compression

Variable-rate coding is challenging but indispensable for learned image compression (LIC) that is in nature characterized by nonlinear transform coding (NTC). Existing methods for variable-rate LIC are restricted by the non-smooth quantization process with zero gradients almost everywhere, and consequently, suffer from training-test gap and degraded rate-distortion (R-D) performance. To address this problem, in this paper, we propose sampling-based optimization for training NTC models along with non-uniform quantizers. Different from gradient-based optimization, the proposed sampling-based optimization first randomly samples the parameters from Gaussian distributions with progressively reduced variance and then selects the optimal parameters with a R-D indicator. On the basis of sampling-based optimization, we develop a learnable non-uniform dead-zone quantizer by adaptively refining the quantization steps for variable-rate coding with nonlinear transforms. Furthermore, we incorporate the learnable dead-zone quantizer to achieve a variable-rate LIC model with enhanced R-D performance and design rate and distortion control algorithms to adapt to dynamic network conditions. Experimental results show that the proposed method achieves state-of-the-art R-D performance in variable-rate image compression. It obtains an average 8.82% BD-rate reduction compared to latest versatile video compression (VVC) standard, and simultaneously achieves precise rate and distortion control with an average variation of 0.0087 bpp in bit-rates and 0.1265 dB in distortion on the Kodak dataset.