[Paper Review] Learned Image Compression with Discretized Gaussian Mixture Likelihoods and Attention Modules

Zhengxue Cheng, Heming Sun+2 more|arXiv (Cornell University)|Jan 6, 2020

Advanced Data Compression TechniquesComputer Science26 references44 citations

TL;DR

The paper introduces discretized Gaussian mixture likelihoods for entropy modeling and simplified attention modules, achieving state-of-the-art rate-distortion performance among learned methods and comparable PSNR to VVC, with superior MS-SSIM visuals.

ABSTRACT

Image compression is a fundamental research field and many well-known compression standards have been developed for many decades. Recently, learned compression methods exhibit a fast development trend with promising results. However, there is still a performance gap between learned compression algorithms and reigning compression standards, especially in terms of widely used PSNR metric. In this paper, we explore the remaining redundancy of recent learned compression algorithms. We have found accurate entropy models for rate estimation largely affect the optimization of network parameters and thus affect the rate-distortion performance. Therefore, in this paper, we propose to use discretized Gaussian Mixture Likelihoods to parameterize the distributions of latent codes, which can achieve a more accurate and flexible entropy model. Besides, we take advantage of recent attention modules and incorporate them into network architecture to enhance the performance. Experimental results demonstrate our proposed method achieves a state-of-the-art performance compared to existing learned compression methods on both Kodak and high-resolution datasets. To our knowledge our approach is the first work to achieve comparable performance with latest compression standard Versatile Video Coding (VVC) regarding PSNR. More importantly, our approach generates more visually pleasant results when optimized by MS-SSIM. This project page is at this https URL https://github.com/ZhengxueCheng/Learned-Image-Compression-with-GMM-and-Attention

Motivation & Objective

• Motivate reducing redundancy in learned image compression by improving entropy modeling.
• Propose a flexible, accurate entropy model using discretized Gaussian mixture likelihoods.
• Integrate lightweight attention modules to enhance network capacity without excessive training cost.
• Demonstrate state-of-the-art performance on standard benchmarks (Kodak, CLIC) against traditional codecs and prior learned methods.

Proposed method

• Model latent code distributions with discretized Gaussian mixture likelihoods to better approximate true marginals (p(y|z)).
• Use a hyperprior framework augmented by a Gaussian mixture to capture spatial and contextual redundancy.
• Incorporate a simplified attention module within the encoder/decoder to enhance focus on complex regions.
• Train with rate-distortion objective combining bits for y and z plus distortion term, across multiple lambdas.
• Clip y range and use discretized convolution with cumulative probability for stable training.
• Compare against JPEG, JPEG2000, HEVC/VVC, and prior learned methods on standard datasets.

Experimental results

Research questions

• RQ1 Can discretized Gaussian mixture likelihoods provide a more accurate entropy model than existing Gaussian/hyperprior approaches?
• RQ2 Does integrating an attention module improve rate-distortion performance without prohibitive training costs?
• RQ3 How close can a learned compression method get to the PSNR performance of VVC while maintaining or improving MS-SSIM quality?
• RQ4 What is the impact of model capacity (N) on RD performance when using Gaussian mixture likelihoods?
• RQ5 Do learned methods with the proposed entropy model outperform traditional codecs on Kodak and high-resolution datasets?

Key findings

• The Gaussian mixture likelihoods yield smaller scales and better spatial redundancy reduction, improving the entropy model.
• The proposed method achieves state-of-the-art performance among learned methods on Kodak and high-resolution datasets.
• The approach reaches comparable PSNR to VVC and yields superior MS-SSIM results compared to existing methods.
• A simplified attention module provides performance gains with moderate training cost, outperforming non-attention variants.
• Ablation studies show benefits of Gaussian mixture modeling across different capacity settings (N).
• Table 1 indicates Joint vs. Ours: PSNR 33.435 vs 33.623 dB; MS-SSIM 0.980 vs 0.981; Rate 0.533 vs 0.519 bpp.