Multi-Sample Training for Neural Image Compression

TL;DR

This paper introduces a multiple-sample training approach for neural image compression that improves upon current methods by using a tighter importance weighted autoencoder objective, leading to better compression performance.

Contribution

It proposes MS-NIC, an enhanced training method using multiple samples for neural image compression, addressing variance issues and improving state-of-the-art results.

Findings

MS-NIC outperforms existing NIC methods in experiments.

The approach is plug-and-play and adaptable to other neural compression tasks.

Analysis of the uniform posterior's properties informs variance reduction strategies.

Abstract

This paper considers the problem of lossy neural image compression (NIC). Current state-of-the-art (sota) methods adopt uniform posterior to approximate quantization noise, and single-sample pathwise estimator to approximate the gradient of evidence lower bound (ELBO). In this paper, we propose to train NIC with multiple-sample importance weighted autoencoder (IWAE) target, which is tighter than ELBO and converges to log likelihood as sample size increases. First, we identify that the uniform posterior of NIC has special properties, which affect the variance and bias of pathwise and score function estimators of the IWAE target. Moreover, we provide insights on a commonly adopted trick in NIC from gradient variance perspective. Based on those analysis, we further propose multiple-sample NIC (MS-NIC), an enhanced IWAE target for NIC. Experimental results demonstrate that it improves sota NIC methods. Our MS-NIC is plug-and-play, and can be easily extended to other neural compression tasks.