Lossy Image Compression with Conditional Diffusion Models

TL;DR

<3-5 sentence high-level summary>The paper tackles efficient lossy image compression by replacing Gaussian decoders with a conditional diffusion-based decoder within a transform-coding framework. It introduces a diffusion variational autoencoder (DVAE) that uses a discrete content latent $\hat{\mathbf z}$ to condition a conditional diffusion process, while texture latents are synthesized at decode time, enabling a controllable rate-distortion-perception tradeoff. A variational upper bound on the RD objective leads to a diffusion-based distortion term and an entropy term for the content latent, with optional perceptual losses to bias reconstructions toward realism; the $\mathcal{X}$-prediction variant enables fast, high-quality reconstructions with few steps. Empirically, the method achieves state-of-the-art FID among learned codecs and competitive distortion metrics, demonstrates a clear tradeoff controlled by $\rho$, and highlights the practicality of fast decoding via $\mathcal{X}$-prediction, albeit with slower decoding than some baselines.

Abstract

This paper outlines an end-to-end optimized lossy image compression framework using diffusion generative models. The approach relies on the transform coding paradigm, where an image is mapped into a latent space for entropy coding and, from there, mapped back to the data space for reconstruction. In contrast to VAE-based neural compression, where the (mean) decoder is a deterministic neural network, our decoder is a conditional diffusion model. Our approach thus introduces an additional ``content'' latent variable on which the reverse diffusion process is conditioned and uses this variable to store information about the image. The remaining ``texture'' variables characterizing the diffusion process are synthesized at decoding time. We show that the model's performance can be tuned toward perceptual metrics of interest. Our extensive experiments involving multiple datasets and image quality assessment metrics show that our approach yields stronger reported FID scores than the GAN-based model, while also yielding competitive performance with VAE-based models in several distortion metrics. Furthermore, training the diffusion with $\mathcal{X}$-parameterization enables high-quality reconstructions in only a handful of decoding steps, greatly affecting the model's practicality. Our code is available at: \url{https://github.com/buggyyang/CDC_compression}

Figures (19)

• Figure 1: Overview of our proposed compression architecture. A discrete "content" latent variable $\hat{{\mathbf z}}$ contains information about the image. Upon decoding, this variable is used for conditioning a denoising diffusion process. The involved "texture" variables $\bar{{\mathbf x}}_{1:N}$ are synthesized on the fly.
• Figure 2: Tradeoffs between bitrate (x-axes, in bpp) and different metrics (y-axes) for various models tested on DIV2K. We consider both perceptual (red frames) and distortion metrics (blue frames). Arrows in the plot titles indicate whether high ($\uparrow$) or low ($\downarrow$) values indicate a better score. CDC (proposed) in its basic version (deterministic, without finetuning to LPIPS) compares favorably in distortion metrics, while CDC with stochastic decoding and added LPIPS losses performs favorably on perceptual metrics.
• Figure 3: Reconstructed Kodak images (cropped images, see full images in Appendix \ref{['fig:additional_vis']}). $1^\text{st}$ row: compared to HiFiC under similar bitrate, our model retains more details around the eyes of the parrot. $2^\text{nd}$ row: our model still gets slightly better visual reconstruction than HiFiC while using less bitrate.
• Figure 4: Compression performance with different numbers of decoding step. We use $\gamma=0$ (deterministic decoding) to plot distortion curves and $\gamma=1$ for perceptual quality curves.
• Figure 5: Qualitative comparison of deterministic and stochastic decoding methods. Deterministic decoding typically results in a smoother image reconstruction. By increasing the noise $\gamma$ used upon decoding the images, we observe more and more detail and rugged texture on the face of the sculpture. $(\gamma=0.8)$ show the best agreement with the ground truth image. All the images share the same bpp.