A Learning Framework for Perceptual Lossy Compression With Stochastic Coding

Recent studies in perceptual lossy compression have highlighted the advantage of stochastic coding with shared randomness between encoder and decoder, over deterministic encoding in the regime of “high perceptual quality”. While the theoretical benefits of stochastic coding have been well-established and demonstrated by analytic examples, its practical realization remains challenging and largely unexplored. In this work, we propose a practical learning framework for stochastic coding that effectively realizes its theoretical advantages. Starting with a theoretically optimal scheme, we develop an implementation closely approximates it through a two-stage training process: learning a stochastic encoder, followed by a stochastic decoder which is modeled as an optimal transport problem conditioned on minimum mean square error (MMSE) decoding. Additionally, for training stochastic coding models, we prove the equivalence between quantized representation and “noisy” representation. Based on this insight, we introduce a quantization-free training method that effectively addresses the non-differentiability challenge posed by quantization. Experiments on a circular distribution example and the MNIST dataset validate our findings and demonstrate the effectiveness of the proposed method.