Learning Switchable Priors for Neural Image Compression

TL;DR

A lightweight NIC model, namely FastNIC, is presented, which outperforms BPG with encoding and decoding complexities below 12 and 10 KMACs/pixel, respectively and the switchable priors are implanted into state-of-the-art NIC models and observed improved compression efficiency.

Abstract

Neural image compression (NIC) usually adopts a predefined family of probabilistic distributions as the prior of the latent variables, and meanwhile relies on entropy models to estimate the parameters for the probabilistic family. More complex probabilistic distributions may fit the latent variables more accurately, but also incur higher complexity of the entropy models, limiting their practical value. To address this dilemma, we propose a solution to decouple the entropy model complexity from the prior distributions. We use a finite set of trainable priors that correspond to samples of the parametric probabilistic distributions. We train the entropy model to predict the index of the appropriate prior within the set, rather than the specific parameters. Switching between the trained priors further enables us to embrace a skip mode into the prior set, which simply omits a latent variable during the entropy coding. To demonstrate the practical value of our solution, we present a lightweight NIC model, namely FastNIC, together with the learning of switchable priors. FastNIC obtains a better trade-off between compression efficiency and computational complexity for neural image compression. We also implanted the switchable priors into state-of-the-art NIC models and observed improved compression efficiency with a significant reduction of entropy coding complexity.

Figures (15)

• Figure 1: Diagram of coding process of the proposed switchable priors method. The Cumulative Distribution Function is abbreviated as CDF. The prior set contains various prior distributions described by a skip mode and CDF tables shared by both the encoder and decoder. The entropy model predicts the appropriate prior index within the set for each latent, rather than estimating the specific parameters of parameterized probabilistic models. The encoder selects latents based on skip mode and entropy encodes the selected latents using the prior set and predicted indexes. The decoder recovers the coded latents from bitstreams and then reshapes them to the original position.
• Figure 2: Diagram of a typical neural image compression model. $g_a$ and $g_s$ denote the analysis and synthesis transforms. $h_a$ and $h_s$ denote the hyper-analysis and synthesis transform. $g_{cm}$ denotes the context model. $g_{ep}$ denotes the entropy parameters module, which is used to generate the entropy parameters of probabilistic models. Our switchable priors method focuses on the improvement in $g_{ep}$, the rate estimation process during training, and the entropy coding during testing.
• Figure 3: Diagram of training and testing stages for dynamically computing CDF, LUTs-based method, and the proposed switchable priors method.
• Figure 4: Frequency of the rate of latent variables collected from the Kodak set in the FM-intra model li2024fm. The x-axis is log-scaled.
• Figure 5: Framework of the proposed FastNIC model. (a) shows the diagram of FastNIC. In the entropy model, the module $h_m(\text{mean})$ is used to predict the mean parameter, while the module $h_e(\text{entropy})$ predicts other parameters of the probabilistic models. The detailed network structure is shown in (b).