SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
R D _ {l o s s} = E _ {x \sim p _ {(x)}} \left[ - \log_ {2} p _ {\hat {y} | \hat {z}} (\hat {y} \mid \hat {z}) - \log_ {2} p _ {\hat {z}} (\hat {z}) \right] + \lambda \cdot E _ {x \sim p _ {(x)}} [ d (x, \hat {x}) ] + \alpha \cdot [ L _ {c o r r} ], \tag {13}
$$
where  $\alpha$  is the scaling factor for the correlation loss.
Given that the entropy estimate presumes independent latent elements, our proposed loss function increases compression efficiency by reducing the correlation between latent features, hence reducing the discrepancy between the assumed probability distribution of the entropy model and the actual latent probability distribution.
(a) Minnen's mean scale hyperprior
(b) Cheng's mean scale hyperprior
(c) SwinT mean scale hyperprior
Figure 7: RD rate comparison of (a) Minnen's (b) Cheng's and (c) SwinT mean scale hyperprior models with the inclusion of correlation loss.
We train all our models on the Vimeo-90k dataset Xue et al. [2019], with the images cropped to the resolution of  $256 \times 256$  for training. The models were optimized using the Adam optimizer Kingma and Ba [2015] with a batch size of 16 and trained for 1.5 million iterations with a learning rate of  $1 \times 10^{-4}$  for the first million iterations and then halved every 50,000 iterations till 1.25 million iterations.
RD loss given in eq. 9 is used to train the baseline models, whereas the modified loss in eq. 13 is used to train models with correlation loss. The rate-distortion tradeoff is guided by  $\lambda$ , whose value is contained in the set [0.0009, 0.0018, 0.0035, 0.0067, 0.0130, 0.0250].
We tested our model on a commonly used Kodak lossless images dataset Kodak [1993], with 24 uncompressed images of  $768 \times 512$  or  $512 \times 768$  resolution. Bits per pixel (bpp) provide the rate, while PSNR indicates the quality of the reconstructed image when evaluating rate-distortion performance. In order to demonstrate the coding efficiency, RD curves are drawn.
4.2 Rate-Distortion Performance
We employ Minnen's Minnen et al. [2018], Cheng's Cheng et al. [2020], and SwinT Liu et al. [2021] based mean scale hyperprior models as well as Cheng's with Checkerboard He et al. [2021] for our experiments. We perform all our experiments on the Pytorch framework Paszke et al. [2017] and use the CompressAI library Bégaint et al. [2020]. Minnen's and Cheng's models were trained using an NVIDIA 2080Ti, whereas the SwinT model was trained on an NVIDIA 3070Ti due to the transformers' high memory requirement. Note that including correlation loss has a negligible effect on the hardware requirement of these base models.
Minnens' Hyperior: Figure 7a depicts the RD curve for Minnen's model Minnen et al. [2018] with and without correlation loss; for Minnen's model, correlation loss achieves BD rate gains of  $3.20\%$ .
Cheng's Hyperprior: Figure 7b illustrates the RD curve for Cheng's model Cheng et al. [2020] with and without correlation loss. Cheng's model yields a BD rate gain of  $9.5\%$  with correlation loss. Cheng's model consists of attention modules along with CNN layers that significantly increase the overall efficiency of the correlation loss. Figure 9 depicts the spatial correlation of latent  $y$  averaged across all latent channels and compares the baseline Cheng's model with the one trained on correlation loss. Figure 9 demonstrates that the model trained with our proposed loss function exhibits significantly less correlation than the baseline model. This observation holds for the models trained with other  $\lambda$  values. The qualitative comparison in Figure 8 between Cheng's model with and without correlation loss demonstrates the superior performance achieved by incorporating correlation loss. It implies that models trained on correlation loss incur less redundancy across different spatial latent locations, resulting in a better rate-distortion trade-off overall.
Transformer Hyperprior: We also compared the effectiveness of correlation loss on the Transformer-based LIC model. We implemented our SwinT with a hyperprior model inspired by Zhu's work and employed a similar architectural setting Zhu et al. [2022]. Transformers tend to have a lesser correlation in their latent space than the convolutional-based LIC models Zhu et al. [2022]. The
Cheng's with Hyperprior
Cheng's with Correlation loss
BPP:0.06,PSNR:28.36dB
Figure 8: Comparison of reconstruction quality of the original image from Kodak dataset using Cheng's hyperprior and Cheng's hyperprior with correlation loss.
Cheng's Hyperprior
(λ=0.0067)
Kodak:PSNR=30.58dB，BPP=0.27
Cheng's with Correlation Loss
$(\lambda = 0.0067)$
Kodak: PSNR=30.84dB, BPP=0.26
Figure 10: RD rate comparison of Cheng's + Checkerboard with the inclusion of correlation loss.
Figure 9: Spatial correlation map comparison of Cheng's hyperprior (left) with Cheng's hyperprior with correlation loss (right)
Figure 11: RD rate comparison of Cheng's + ChARM with the inclusion of correlation loss.
Figure 7c illustrates the performance of the SwinT model with and without correlation loss. We found that with the correlation loss, the SwinT model achieves a BD-rate gain of about  $4.8\%$ . The effectiveness of correlation loss for the SwinT model further indicates the efficacy of the correlation loss, which in turn improves the mean scale based entropy model.
Cheng's with Checkerboard Context Model: Figure 10 illustrates the RD curve for Cheng's model with Checkerboard He et al. [2021] with and without correlation loss. Cheng's with Checkerboard yields a BD rate gain of  $6.13\%$  with correlation loss. Checkerboard uses a parallel decoding approach by encoding and decoding only half of the latent variables using a checkerboard-shaped context. The remaining latents, referred to as anchors, are coded solely based on the hyperprior. Since half of the latents are encoded by the hyperprior module. Incorporating correlation loss into the Checkerboard context model significantly improves the performance of the hyperprior module, resulting in an overall enhancement of the model's performance.
Cheng's with ChARM: Figure 11 illustrates the RD curve for Cheng's model with ChARM Minnen and Singh [2020], both with and without the integration of correlation loss. Notably, when correlation loss is incorporated into Cheng's with ChARM, there is a substantial BD rate gain of  $3.3\%$ . ChARM effectively employs channel context and latent residual prediction for the encoding and decoding of the latents. The combination of correlation loss with ChARM leads to a significant enhancement in its overall performance, marking a performance difference of approximately  $0.5\%$  when compared to a full AR model.
4.3 Complexity Tradeoff
The primary advantage of the correlation loss lies in its capacity to minimize the discrepancy between the actual and presumed probability distribution within the entropy model. Figure 1 illustrates that for
Table 1: Average encoding and decoding time on the Kodak dataset for different entropy models applied to Cheng's Hyperprior. BD-Rate gains are evaluated using Cheng's Hyperprior as the baseline.
Architecture	BD Rate Gains (%)	Inference Time (sec)
Cheng's Hyperprior (CH)	—	4.66
CH + Correlation Loss (Proposed)	9.5	4.66
CH + Checkerboard	10.37	5.33
CH + Checkerboard + Correlation Loss (Proposed)	16.50	5.33
CH + ChARM	14.69	8.14
CH + ChARM + Correlation Loss (Proposed)	17.99	8.14
CH + AR Context	18.47	251.65
Figure 12: Comparison between Cheng's AR context, Cheng's hyperprior, Cheng's hyperprior with correlation loss, Cheng's Checkerboard with correlation loss, and Cheng's ChARM with correlation loss
Figure 13: Visual representation depicting the relationship between bits per pixel (bpp) and correlation, and its influence on peak signal-to-noise ratio (PSNR) gains. For all the four models, both the average correlation and the gains decreases as the bpp value increases.
a given encoder and decoder architecture, which is Cheng's model in this case, the choice of entropy model determines the performance of the resulting model. Note that Cheng's hyperprior signifies the lower bound with the given encoder-decoder architecture, whereas Cheng's AR defines the upper limit with the same encoder-decoder architecture. Using Cheng's hyperprior model as the baseline, we compared the effectiveness of including an AR module vs. the correlation loss. The Figure 12 shows the RD curve comparison, while the Table 1 compares the inference time for processing the entire Kodak test dataset. Incorporating correlation loss into Cheng's hyperprior model results in a performance improvement roughly half that of the AR-context-based model, all while reducing inference costs to just 1/55th of the original. Similarly, when utilizing Cheng's Checkerboard with correlation loss, a substantial performance enhancement of approximately  $90\%$  is achieved compared to the AR-context-based model, with an inference cost as low as 1/50th of the latter. Additionally, Cheng's ChARM with correlation loss attains an impressive performance gain of about  $98\%$  when compared to the AR-context-based model, while incurring an inference cost of merely 1/30th of the AR-context-based model as shown in Figure 1.
It should be noted that when correlation loss is applied to Cheng's hyperprior, significant improvements can be achieved due to the large performance gap between Cheng's hyperprior baseline and the upper limit. On the other hand, since the gaps are diminished in models like Cheng's with Checkerboard and Cheng's with ChARM, the potential for improvement by applying correlation loss is also limited when compared to Cheng's hyperprior baseline. In summary, the introduction of correlation loss to these models would result in a closer approximation to the full AR performance, albeit with gains that are comparatively smaller.
5.1 Visual Illustration of Effectiveness of the Correlation Loss
Figure 4 visualizes some of the internal workings of our models. We compare Cheng's hyperprior model with and without correlation loss. By observing the first column of Figure 4, it is clear that our correlation loss makes the latent feature preserve more detailed texture information yielding higher PSNR compared with the base model. The mean  $\mu$  can also extract more distinctive and meaningful information from the latent space. Since the hyperpriors (mean  $\mu$  and scale  $\sigma$ ) contain more information, they require slightly more bits than the base model. Nevertheless, the rise in hyperprior rate is adequately offset by the relatively greater decrease in latent bitrate, as demonstrated in the final column of Figure 4, where the normalized latents exhibit improved decorrelation compared to the base model. Overall, our proposed method surpasses the base model by enhancing PSNR through more decorrelated latents while either maintaining or reducing bpp with better decorrelated latents.
5.2 Effect of Correlation score on the Model's performance
Our proposed loss function consistently demonstrates significant BD-rate gains at lower bit rates while maintaining comparable performance at high bit rates. This behavior is primarily attributed to learned image compression (LIC) models exhibiting higher correlation among latent variables at lower bpp values compared to higher bpp values as reported by Zhu et al. [2022] and also evident from the correlation maps in Figure 6 of the supplementary material. We conducted a comprehensive analysis and presented the findings in Figure 13. Figure 13 showcases the relationship between PSNR gains, bpp, and correlation for different models, including Cheng's hyperprior, Cheng's with Checkerboard, SwinT hyperprior, and Minnen's hyperprior. From the graphs in the Figure 13, a clear pattern can be observed: as the bpp decreases, the correlation of the latents increases, resulting in a higher gain in PSNR. However, as we move towards higher bpp, the correlation becomes notably reduced, resulting in decreased PSNR gain. This trend explains the reason why the efficacy of the correlation loss is larger in low bpp range and diminish at high bpp range.
Comprehensive ablation studies regarding various mask types, mask sizes, and  $\alpha$  values are presented in the supplementary material.
The hyperprior-based entropy models assume probabilistic independence, which leads to a discrepancy between the actual and the assumed probability distribution. We proposed correlation loss, which reduces the correlation among spatially neighbored elements in the latent space resulting in a better fit with the spatially independent probability model. Our proposed loss function does not require any model structure or capacity changes and acts as a plug-in method for the existing neural compression methods. The proposed approach, without modifying the entropy models and increasing the inference time, significantly improves the RD performance with BD rate gains up to  $17.99\%$ .
Authors are thankful to Muhammad Awais for his constructive and helpful feedback. Authors are also thankful to Taegoo Kang for managing and providing the required compute resources.
This work was supported in part by the Institute of Information and Communications Technology Planning and Evaluation (IITP) grant funded by the Korea Government [Ministry of Science and Information Communication Technology (MSIT)] (Development of Audio/Video Coding and Light Field Media Fundamental Technologies for Ultra Realistic Tera-media) under Grant 2017-0-00072; in part by the IITP grant funded by the Korea Government (MSIT) (Artificial Intelligence Convergence Innovation Human Resources Development, Kyung Hee University) under Grant RS-2022-00155911; in part by the IITP grant funded by the Korea Government (MSIT) (Artificial Intelligence Innovation Hub) under Grant 2021-0-02068.
Lucas Theis, Wenzhe Shi, Andrew Cunningham, and Ferenc Huszár. Lossy image compression with compressive autoencoders. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https://openreview.net/forum?id=rJiNwv9gg.
George Toderici, Sean M O'Malley, Sung Jin Hwang, Damien Vincent, David Minnen, Shumeet Baluja, Michele Covell, and Rahul Sukthankar. Variable rate image compression with recurrent neural networks. International Conference on Learning Representations (ICLR), 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Gregory K Wallace. The JPEG still picture compression standard. Communications of the ACM, 34(4): 30-44, 1991.
Tianfan Xue, Baian Chen, Jiajun Wu, Donglai Wei, and William T. Freeman. Video enhancement with task-oriented flow. Int. J. Comput. Vis., 127(8):1106-1125, 2019. doi: 10.1007/S11263-018-01144-2. URL https://doi.org/10.1007/s11263-018-01144-2.
Yibo Yang, Stephan Mandt, Lucas Theis, et al. An introduction to neural data compression. Foundations and Trends® in Computer Graphics and Vision, 15(2):113-200, 2023.
Zhisheng Zhong, Hiroaki Akutsu, and Kiyoharu Aizawa. Channel-level variable quantization network for deep image compression. In Christian Bessiere, editor, Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI-20, pages 467-473. International Joint Conferences on Artificial Intelligence Organization, 7 2020. doi: 10.24963/ijcai.2020/65. URL https://doi.org/10.24963/ijcai.2020/65. Main track.
Yinhao Zhu, Yang Yang, and Taco Cohen. Transformer-based transform coding. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https://openreview.net/forum?id=IDwN6xjHnK8.