SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We conducted a large-scale compression evaluation involving multiple image quality metrics and test datasets. Besides metrics measuring the traditional distortion scores, we also consider metrics that can reflect perceptual quality. While some of these metrics are fixed, others are learned from data. We will refer to our approach as "Conditional Diffusion Compression" (CDC).
Metrics We selected sixteen metrics for image quality evaluations, of which we present eight most widely-used ones in the main paper and the remaining eight in the appendix. Specifically, several more recently proposed learned metrics (Heusel et al., 2017; Zhang et al., 2018; Prashnani et al., 2018; Ding et al., 2020) capture perceptual properties/realism better than other non-learned methods; we denote these metrics as perceptual metrics and the others as distortion metrics. It is important to note that when calculating FID, we follow Mentzer et al. (2020) by segmenting images into non-overlapping patches of  $256 \times 256$  resolution. A brief introduction about the metrics is also available in Appendix D.
Test Data To support our compression quality assessment, we consider the following datasets with necessary preprocessing: 1. Kodak (Franzen, 2013): The dataset consists of 24 high-quality images at  $768 \times 512$  resolution. 2. Tecnick (Asuni & Giachetti, 2014): We use 100 natural images with  $600 \times 600$  resolutions and then downsample these images to  $512 \times 512$  resolution. 3. DIV2K (Agustsson & Timofte, 2017): The validation set of this dataset contains 100 high-quality images. We resize the images with the shorter dimension being equal to  $768\mathrm{px}$ . Then, each image is center-cropped to a  $768 \times 768$  squared shape. 4. COCO2017 (Lin et al., 2014): For this dataset, we extract all test images with resolutions higher than  $512 \times 512$  and resize them to  $384 \times 384$  resolution to remove compression artifacts. The resulting dataset consists of 2695 images.
Model Training Our model was trained using the well-established Video-90k dataset (Xue et al., 2019), which includes 90,000 video clips, each containing 7-frame sequences with a resolution of  $448 \times 256$  pixels, curated from vimeo.com. For each iteration, we randomly selected one frame from every video clip, which was then subject to random cropping to achieve a uniform resolution of  $256 \times 256$  pixels.
(a) Ground Truth
(d) Ground Truth
Figure 3: Reconstructed Kodak images (cropped images, see full images in Appendix F).  $1^{\mathrm{st}}$  row: compared to HiFiC under similar bitrate, our model retains more details around the eyes of the parrot.  $2^{\mathrm{nd}}$  row: our model still gets slightly better visual reconstruction than HiFiC while using less bitrate.
(b)  $\mathrm{CDC}(\rho = 0.9)$  (bpp=0.205)
(e)  $\mathrm{CDC}(\rho = 0.9)$  (bpp=0.398)
(c) HiFiC (bpp=0.207)
(f) HiFiC (bpp=0.456)
The training procedure initiated with a warm-up phase, setting  $\lambda$  to  $10^{-6}$  and running the model for approximately 700,000 steps. Subsequently, we increased  $\lambda$  to align with the desired bitrates and continued training for an additional 1,000,000 steps until the model reached convergence.
For the  $\epsilon$ -prediction model, our training utilized diffusion process comprising  $N_{\mathrm{train}} = 20,000$  steps. Conversely, the number of diffusion steps for the  $\mathcal{X}$ -prediction model is  $N_{\mathrm{train}} = 8,193$ . We implemented a linear variance schedule to optimize the  $\epsilon$ -prediction model, while a cosine schedule was selected for the  $\mathcal{X}$ -prediction model optimization. Throughout the training regime, we maintained a batch size of 4. The Adam optimizer (Kingma & Ba, 2014) was employed to facilitate efficient convergence. We commenced with an initial learning rate of  $lr = 5 \times 10^{-5}$ , which was reduced by  $20\%$  after every 100,000 steps, ultimately clipped to a learning rate of  $lr = 2 \times 10^{-5}$ .
4.1 Baseline Comparisons
Baselines and Model Variants We showed two variants of our  $\mathcal{X}$ -prediction CDC model. Our first proposed model is optimized in the presence of an additive perceptual reconstruction term at  $\rho = 0.9$ , which is the largest  $\rho$ -value we chose. For this variant, we used  $\mathbf{x}_N \sim \mathcal{N}(0, \gamma^2 I)$  with  $\gamma = 0.8$  to reconstruct the images. The other proposed version is the base model, trained without the additional perceptual term ( $\rho = 0$ ) and using a deterministic decoding with  $\mathbf{x}_N = 0$ . As discussed below, this base version performs better in terms of distortion metrics, while the stochastic and LPIPS-informed version performs better in perceptual metrics.
We compare our method with several neural compression methods. The best reported perceptual results were obtained by HiFiC (Mentzer et al., 2020). This model is optimized by an adversarial network and employs additional perceptual and traditional distortion losses (LPIPS and MSE). In terms of rate-distortion performance, two VAE models are selected: DGML (Cheng et al., 2020) and NSC (Wang et al., 2022). Both are the improvements over the MSE-trained Mean-Scale Hyperprior (MS-Hyper) architecture (Minnen et al., 2018). For a fair comparison, we did not use the content-adaptive encoding for NSC model. For comparisons with classical CODECs, we choose BPG as a reference.
Figure 2 illustrates the tradeoff between bitrate and image quality using the DIV2K dataset. We only employ 17 steps to decode the images with  $\mathcal{X}$ -prediction model, which is much more efficient than  $\epsilon$ -prediction model that requires hundreds of steps to achieve comparable performance (see Appendix H for results on other datasets and comparison with  $\epsilon$ -prediction). In the figure, dashed lines represent the baseline models, while solid lines depict our proposed CDC models. The eight shown plots present two different types of metrics, distinguished by their respective frame colors.
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.
- Perceptual Metrics (red). The red subplots depict perceptual metrics that assess the compression for realism performance. Our findings reveal that our proposed CDC model ( $\rho = 0.9$ ) achieves the best performance in three out of four metrics. The closest competitor is the HiFiC baseline. Notably, HiFiC demonstrates the highest score in LPIPS but exhibits suboptimal performance in all other metrics.
- Distortion Metrics (blue). The blue subfigures present distortion-based metrics. We note that the CDC model with  $\rho = 0$  produces relatively favorable results in distortion metrics, excluding PSNR. It shows on-par scores with the best baselines in FSIM, SSIM, and MS-SSIM scores, despite none of the shown models being specifically optimized for these three metrics. In contrast, "classical" neural compression models (Minnen et al., 2018; Wang et al., 2022; Cheng et al., 2020) directly target MSE distortion by minimizing an ELBO objective with Gaussian decoders, resulting in better PSNR scores.
Our proposed versions CDC  $(\rho = 0)$  and CDC  $(\rho = 0.9)$  show qualitative differences in perceptual and distortion metrics. Setting  $\rho = 0$  only optimizes a trade-off between bitrate and the diffusion loss; compared to  $\rho = 0.9$ , this results in better performance in model-based distortion metrics (i.e., except PSNR). Fig. 3 qualitatively shows that the resulting decoded images show fewer over-smoothing artifacts than VAE-based CODECs (Cheng et al., 2020). In contrast, CDC  $(\rho = 0.9)$  performs the best in perceptual metrics. These are often based on extracted neural network features, such as Inception or VGG (Szegedy et al., 2016; Simonyan & Zisserman, 2014), and are more susceptible to image realism. By varying  $\rho$ , we can hence control a three-way trade-off among distortion, perceptual quality, and bitrate (See Appendix E for results with other  $\rho$  values).
Distortion vs. Perception Our experiments revealed the aforementioned distortion-perception tradeoff in learned compression (Blau & Michaeli, 2019). In contrast to perceptual metrics, distortions such as PSNR are very sensitive to imperceptible image translations (Wang et al., 2005). The benefit of distortions is that they carry out a direct comparison between the reconstructed and original image, albeit using a debatable metric (Dosovitskiy & Brox, 2016). The question of whether distortion or perceptual quality is more relevant may ultimately not be easily solvable; yet it is plausible that most compression gains can be expected when targeting a combination of perception/realism and distortion, rather than distortion alone (Mentzer et al., 2020; Yang et al., 2023). Especially, our method's strong performance in terms of FID, one of the most widely-adopted perceptual evaluation schemes (Ho et al., 2020; Song & Ermon, 2019; Mentzer et al., 2020; Brock et al., 2019; Song et al., 2021a,b), seems promising in this regard.
(b)  $\gamma = 0$  MS-SSIM  $= 0.979$  LPIPS  $= 0.078$
(c)  $\gamma = 0.6$   
MS-SSIM=0.977  
LPIPS=0.064
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.
(d)  $\gamma = 0.8$  MS-SSIM  $= 0.973$  LPIPS  $= 0.054$
(e)  $\gamma = 1$  MS-SSIM  $= 0.964$  LPIPS  $= 0.058$
Influence of decoding steps In the previous section, we demonstrated that the CDC  $\mathcal{X}$ -prediction model can achieve decent performance with a small number of decoding steps. In Figure 4, we further investigate the compression performance of the  $\mathcal{X}$ -prediction  $(\rho = 0)$  model using different decoding steps. Our findings reveal that when employing stochastic decoding, the model consistently produces better perceptual results as the number of decoding steps increases. However, in the case of deterministic decoding, more decoding steps do not lead to a substantial improvement in distortion.
Our findings show that the  $\mathcal{X}$ -prediction model can behave similarly to a Gaussian VAE decoder. In this scenario, the latent code  $\mathbf{z}$  becomes the primary determinant of the decoding outcome, enabling the model to reconstruct the original image  $\mathbf{x}_0$  with a single decoding step. However, even a single decoding step has a tendency to reconstruct data closer to the data mode, only guaranteeing an acceptable distortion score. To effectively improve perceptual quality, it is crucial to incorporate more iterative decoding steps, particularly when utilizing stochastic decoding. Thus, we further explore the impact of stochastic decoding through the following ablation experiment.
Stochastic Decoding By adjusting the noise level parameter, denoted as  $\gamma$ , during the image decoding process, we can achieve different decoding outcomes. In order to investigate the impact of the noise on the decoding results, we present Figure 5, which provides both quantitative and qualitative evidence for 4 candidates  $\gamma$  values. Our findings indicate that larger values lead to improved perceptual quality and higher distortion, as evidenced by lower LPIPS and lower MS-SSIM. Values of  $\gamma$  greater than 0.8 not only increase distortion but also diminish perceptual quality. In terms of finding the optimal balance, a  $\gamma$  value of 0.8 offers the lowest LPIPS and the best qualitative outcomes as shown in Figure 5. From a qualitative standpoint, we notice that the noise introduces plausible high-frequency textures. Although these textures may not perfectly match the uncompressed ones (which is impossible), they are visually appealing when an appropriate  $\gamma$  is chosen. For additional insights into our decoding process, we provide visualizations of the decoding steps in Appendix G. These visualizations showcase how "texture" variables evolve during decoding. We observe that this phenomenon becomes particularly pronounced when employing the  $\epsilon$ -prediction, as  $\mathcal{X}$ -prediction models exhibit a stronger inclination towards reconstruction.
We proposed a transform-coding-based neural image compression approach using diffusion models. We use a denoising decoder to iteratively reconstruct a compressed image encoded by an ordinary neural encoder. Our loss term is derived from first principles and combines rate-distortion variational autoencoders with denoising diffusion models. We conduct quantitative and qualitative experiments to compare our method against several GAN and VAE based neural codecs. Our approach achieves
promising results in terms of the rate-perception tradeoff, outperforming GAN baselines in three out of four metrics, including FID. In terms of classical distortion, our approach still performs comparable to highly-competitive baselines.
Limitations & Societal Impacts In our quest to enhance compression performance, further improvement can be achieved by integrating advanced techniques such as autoregressive entropy models or iterative encoding. We leave such studies for future research.
One critical societal concern associated with neural compression, particularly when prioritizing perceptual quality, is the risk of misrepresenting data in the low bitrate regime. Within this context, there is a possibility that the model may generate hallucinated lower-level details. For example, compressing a human face may lead to a misrepresentation of the person's identity, which raises important considerations regarding fairness and trustworthiness in learned compression.
The authors acknowledge support by the IARPA WRIVA program, the National Science Foundation (NSF) under the NSF CAREER Award 2047418; NSF Grants 2003237 and 2007719, the Department of Energy, Office of Science under grant DE-SC0022331, as well as gifts from Intel, Disney, and Qualcomm.
Agustsson, E. and Timofte, R. Ntire 2017 challenge on single image super-resolution: Dataset and study. In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, July 2017.
Agustsson, E., Tschannen, M., Mentzer, F., Timofte, R., and Gool, L. V. Generative adversarial networks for extreme learned image compression. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 221-231, 2019.
Agustsson, E., Minnen, D., Toderici, G., and Mentzer, F. Multi-realism image compression with a conditional generator. arXiv preprint arXiv:2212.13824, 2022.
Asuni, N. and Giachetti, A. Testimages: a large-scale archive for testing visual devices and basic image processing algorithms, stag - smart tools & apps for graphics conference, 2014., 2014.
Balle, J., Minnen, D., Singh, S., Hwang, S. J., and Johnston, N. Variational image compression with a scale hyperprior. International Conference on Learning Representations, 2018.
Bégaint, J., Racapé, F., Feltman, S., and Pushparaja, A. Compressai: a pytorch library and evaluation platform for end-to-end compression research. arXiv preprint arXiv:2011.03029, 2020.
Bellard, F. Bpg image format, 2018. URL https://bellard.org/bpg/.
Blau, Y. and Michaeli, T. Rethinking lossy compression: The rate-distortion-perception tradeoff. In International Conference on Machine Learning, pp. 675-685. PMLR, 2019.
Brock, A., Donahue, J., and Simonyan, K. Large scale gan training for high fidelity natural image synthesis. In International Conference on Learning Representations, 2019.
Cheng, Z., Sun, H., Takeuchi, M., and Katto, J. Learned image compression with discretized gaussian mixture likelihoods and attention modules. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7939-7948, 2020.
Cisco. Visual network index cisco. forecast and methodology. White Paper, 2017.
Ding, K., Ma, K., Wang, S., and Simoncelli, E. P. Image quality assessment: Unifying structure and texture similarity. IEEE transactions on pattern analysis and machine intelligence, 2020.
Dosovitskiy, A. and Brox, T. Generating images with perceptual similarity metrics based on deep networks. Advances in neural information processing systems, 29, 2016.
Wang, Y., Schiff, Y., Gokaslan, A., Pan, W., Wang, F., De Sa, C., and Kuleshov, V. Infodiffusion: Representation learning using information maximizing diffusion models. arXiv preprint arXiv:2306.08757, 2023.
Wang, Z., Bovik, A. C., and Simoncelli, E. P. Structural approaches to image quality assessment. Handbook of image and video processing, 7(18), 2005.
Xue, T., Chen, B., Wu, J., Wei, D., and Freeman, W. T. Video enhancement with task-oriented flow. International Journal of Computer Vision (IJCV), 127(8):1106-1125, 2019.
Yang, R., Srivastava, P., and Mandt, S. Diffusion probabilistic modeling for video generation. arXiv preprint arXiv:2203.09481, 2022.
Yang, Y., Bamler, R., and Mandt, S. Improving inference for neural image compression. Advances in Neural Information Processing Systems, 33:573-584, 2020.
Yang, Y., Mandt, S., and Theis, L. An introduction to neural data compression. Foundations and Trends in Computer Graphics and Vision, 2023.
Zhang, G., Qian, J., Chen, J., and Khisti, A. Universal rate-distortion-perception representations for lossy compression. Advances in Neural Information Processing Systems, 34:11517-11529, 2021.
Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang, O. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586-595, 2018.
Zhao, S., Song, J., and Ermon, S. Towards deeper understanding of variational autoencoding models. arXiv preprint arXiv:1702.08658, 2017.
Zheng, H., Nie, W., Vahdat, A., Azizzadenesheli, K., and Anandkumar, A. Fast sampling of diffusion models via operator learning. arXiv preprint arXiv:2211.13449, 2022.
Zhu, Y., Yang, Y., and Cohen, T. Transformer-based transform coding. In International Conference on Learning Representations, 2021.
A Pretrained Baselines
We refer to Bégaint et al. (2020) for pretrained MS-Hyper and DGML models. For HiFiC model, we use the pretrained models implemented in the publicly available repositories<sup>1</sup>. Both models were sufficiently trained on natural image datasets (Xue et al., 2019; Kuznetsova et al., 2020). For NSC (Wang et al., 2022) baseline, we use the official codebase<sup>2</sup> and DIV2k training dataset to train the model.
B Architectures
Figure 6: Visualization of our model architecture
The denoising module's design adopts a U-Net architecture similar to that of the DDIM (Song et al., 2021a) and DDPM (Ho et al., 2020) projects. Each U-Net unit is composed of two ResNet blocks (He et al., 2016), an attention block, and a convolutional up/downsampling block. We employ six U-Net units for both the downsampling and upsampling processes. In the downsampling units, the channel dimension is  $64 \times j$ , where  $j$  represents the index of the layer ranging from 1 to 6. The upsampling units follow the reverse order. Each encoder module comprises a ResNet block and a convolutional downsampling block. For embedding conditioning, we employ ResNet blocks and transposed convolutions to upscale  $\mathbf{z}$  to match the spatial dimension of the inputs of the initial four U-Net downsampling units. This allows us to perform conditioning by concatenating the output of the embedder with the input of the corresponding U-Net unit.
Figure 6 also describes our design choice of the model. We list the additional detailed specifications that we did not clarify in the main paper as follow:
- The hyper prior structure shares the same design as Minnen et al. (2018). The channel number of the hyper latent  $\mathbf{y}$  is set as 256.
<sup>1</sup>https://github.com/Justin-Tan/high-fidelity-generative-compression
2https://github.com/Dezhao-Wang/Neural-Syntax-Code
- We use 3x3 convolution for most of the convolutional layers. The only exceptions are the 1st conv-layer of the first DU component and the 1st layer of the 1st ENC component, where we use 7x7 convolution for wider receptive field.
-  $i_{N}^{n}$  is embedded by a linear layer, which expands the 1-dimensional scalar to the same channel size as the corresponding DU/UU units. We then add the expanded tensor to the intermediate ResBlock of each DU/UU unit.
C Model Efficiency
We provide information on the model parameter size of the proposed model and baselines, and the corresponding time cost of running a full forward pass in Table 1. We run benchmarking on a server with a RTX A6000. We decode 24 images from Kodak dataset and calculate the average neural decoding time, which does not include entropy-coding process, as all the model share a similar entropy model.
CDC (1 step)	CDC (17 steps)	HiFiC	MS-hyper	DGML
Number of Parameters	53.8M	53.8M	181.47M	17.5M	26.5M
Decoding Time (Seconds)	0.015	1.04	0.0051	0.0011	0.0025
Table 1: Model and decoding time.
Our model exhibits superior memory efficiency compared to HiFiC. However, diffusion models suffer from slow decoding speed owing to their iterative denoising process. In the benchmark model utilized in our main paper, the decoding of an image takes approximately 1 second. Although this is slower than the baselines, it remains within an acceptable time range. Further optimization of the neural network module, such as the removal of the attention module, holds the potential to enhance efficiency even further.
D Additional explanation on experiment metrics
FID, as the most popular metric for evaluating the realism of images, measures the divergence (Frechet Distance) between the statistical distributions of compressed image latent features and ground truth ones. The model extracts features from Inception network and calculates the latent features' corresponding mean and covariance. LPIPS measures the 12 distance between two latent embeddings from VGG-Net/AlexNet. Likewise, PieAPP provides a different measurement of perceptual score based on a model that is trained with the pairwise probabilities. DISTS measures the structural and textural similarities based on multiple layers of network feature maps and an algorithm inspired by SSIM. CKDN leverages a distillation method that can extract a knowledge distribution from reference images, which can help calculate the likelihood of the restored image under such distribution. Both MUSIQ and DBCNN are non-reference metrics, as they both use deep network models (transformer and CNN, respectively) that are pre-trained on labeled image data with Mean Opinion Score. For non-trained metrics (model-based methods), FSIM uses the phase congruency and the color gradient magnitude of two images to calculate the similarity. The (CW/MS-)SSIM family uses insights about human perception of contrasts to construct a similarity metric to mimic human perception better. GMSD evaluates the distance between image color gradient magnitudes. NLPD means normalized Laplacian pyramid distance, which derives from a simple model of the human visual system and is also sensitive to the contrast of the images. VSI reflects a quantitative measure of visual saliency that is widely studied by psychologists and neurobiologists. MAD implements a multi-stage algorithm also inspired by the human visual system for distortion score calculation.
E Supplemental Ablation Study
By manipulating the trade-off parameter  $\rho$ , our model can be trained to prioritize either perceptual quality or traditional distortion performance. In Figure 7, we present the rate-distortion curves for the COCO dataset. Throughout our study, we examine four distinct values of  $\rho$  (0, 0.32, 0.64, 0.9). The outcomes illustrate that higher values of  $\rho$  result in improved perceptual quality but at the cost of worsened distortions in most scenarios. It is worth noting that values exceeding  $\rho > 0.9$  are not viable, as perceptual quality ceases to exhibit noticeable improvements.
Figure 7: rate-distortion curves with different  $\rho$  values