Source-linked AI summary
Learned Image Compression with Discretized Gaussian Mixture Likelihoods and Attention Modules
Zhengxue Cheng, Heming Sun, Masaru Takeuchi, Jiro Katto
TL;DR
Learned image compression remains behind established standards, especially in PSNR, despite rapid progress. This paper combines discretized Gaussian mixture likelihoods for flexible entropy modeling with attention modules, reporting state-of-the-art results and performance comparable to VVC in PSNR, alongside more visually pleasant MS-SSIM reconstructions.
Problem
Learned compression methods still exhibit a performance gap versus established standards, especially on the widely used PSNR metric.
Method
The method parameterizes latent-code distributions with discretized Gaussian mixture likelihoods and incorporates attention modules into the network architecture.
Results
The approach achieves state-of-the-art performance against learned methods and classical standards, with PSNR comparable to VVC and visually pleasant MS-SSIM reconstructions.
Takeaways & Limitations
More flexible entropy modeling and attention produce a learned codec competitive with established standards across Kodak and high-resolution evaluations.
Takeaways & Limitations
The joint entropy model remains imperfect because parameterized distributions do not fully exploit some spatial contexts and neighboring information.
Abstract
from arXiv · showhide
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
1. Introduction
The paper addresses remaining redundancy in learned image compression by combining a more accurate entropy model based on discretized Gaussian mixture likelihoods with attention modules. Experiments report state-of-the-art performance against classical standards and learned methods, approaching VVC in PSNR.
- Learned compression methods still lag established standards, particularly on PSNR, motivating further investigation of remaining redundancy.
- The proposed entropy model uses discretized Gaussian mixture likelihoods to represent latent-code distributions more accurately and flexibly.The authors motivate this design as a way to remove remaining redundancy and reduce required encoding bits.
- The network incorporates a simplified attention module that emphasizes complex image regions while maintaining moderate training complexity.
- The method achieves state-of-the-art performance on PSNR and MS-SSIM compared with classical standards and existing learned compression approaches.
- The method reaches performance very close to VVC in PSNR and produces visually pleasant reconstructions when optimized with MS-SSIM.The visual comparison is illustrated for Kodak images at approximately 0.1 bpp.
2. Related Work
Prior work progressed from hand-crafted compression standards and differentiable training techniques toward learned architectures with compact latent representations. Recurrent and context-adaptive methods form important strands of this development.
- Classical standards use hand-crafted modules including prediction, transforms, quantization, entropy coding, and rate-distortion mode selection.Examples include JPEG, JPEG2000, HEVC, and VVC.
- Early learned-compression research addressed non-differentiable quantization and rate estimation to enable end-to-end training.
- Subsequent methods designed network structures to extract compact latent representations and reconstruct high-quality images.
- Recurrent approaches recursively compress residual information but mainly rely on binary representations at each iteration for scalable coding.
3. Proposed Method
The proposed learned compression model improves entropy modeling by replacing a single Gaussian with discretized Gaussian mixtures and adds attention-based architecture enhancements. These choices target residual spatial redundancy and reduce the bits required to encode latent codes.
- Learned Compression Formulation: The baseline learned codec uses transform coding, quantization, and entropy coding, with a hyperprior providing side information to estimate latent-code distributions.The hyperprior models spatial dependencies in y through an auxiliary autoencoder and conditional distributions.
- Entropy-Model Motivation: Existing hyperprior and autoregressive entropy models can leave spatial redundancy because their parameterized distributions cannot fully represent neighboring-element information.The Joint model adds a 5 × 5 masked convolution, but residual redundancy remains in simple regions.
- Discretized Gaussian Mixture Likelihoods: The method uses discretized Gaussian mixture likelihoods to parameterize quantized latent-code distributions with more flexible and accurate likelihoods.The mixture model assigns learnable weights, means, and scales to K Gaussian components; experiments use K = 3.
- Discretized Gaussian Mixture Likelihoods: Gaussian mixtures adapt component weights across regions, allowing small scales for likely values and reducing the bits required for encoding.For Kodim21, the proposed model has smaller scales than Joint and HyperPrior, and its required bits are fewer than Joint under λ = 0.015.
4. Implementation Details
The models are trained on cropped ImageNet samples with Adam, using MSE or MS-SSIM objectives and specified rate-distortion evaluation procedures.
- Training uses 13,830 ImageNet crops of size 256 × 256, Adam optimization, batch size 8, and a learning rate reduced for the final 80k iterations.
- MSE-trained models use six λ values, with channel capacity N=128 for lower-rate models and N=192 for higher-rate models.
- MS-SSIM training defines distortion as D(x, ˆx) = 1−MS-SSIM(x, ˆx), uses four λ values, and trains each model to 10^6 iterations.
- Evaluation uses Kodak’s 24 uncompressed 768 × 512 images, measuring rate in bits per pixel and quality with PSNR or MS-SSIM.
5. Experiments
Experiments assess ablations and rate-distortion performance against learned and conventional codecs, including Kodak and high-resolution CLIC validation images. The proposed method achieves competitive PSNR with VVC and strong MS-SSIM and visual quality.
- Ablation Study: The ablation study evaluates Gaussian mixture modeling and simplified attention at N=128 and N=192 using MSE with λ = 0.015.Loss curves and rate-distortion points are used to assess their coding gains.
- Rate-distortion Performance: On Kodak, the method achieves competitive PSNR with VVC, exceeds previous learned methods, and obtains state-of-the-art MS-SSIM results.The comparison includes JPEG, JPEG2000, HEVC, VVC, and several learned compression methods.
- Rate-distortion Performance: On the CLIC validation dataset, the approach surpasses all compared codecs except VVC in PSNR and significantly outperforms them in MS-SSIM.The results indicate that the method also works for high-resolution images.
- Qualitative Results: At approximately 0.10–0.12 bpp, MS-SSIM-optimized reconstructions preserve more natural details, while MSE-optimized results reach comparable quality with VVC.Examples include hair, brick walls, flowers, clouds, and reduced blocking artifacts relative to older codecs.
6. Conclusion
The paper proposes learned image compression combining discretized Gaussian mixture likelihoods with attention modules. It reports state-of-the-art performance against learned methods and standards, comparable PSNR to VVC, and improved visual quality under MS-SSIM.
- The method uses discretized Gaussian mixture likelihoods and attention modules for learned image compression.The mixture likelihoods provide a more flexible and accurate entropy model, while the simplified attention module targets high coding efficiency.
- The approach achieves state-of-the-art performance against existing learned compression methods and coding standards including HEVC, JPEG2000, and JPEG.
- The method achieves comparable PSNR performance with the next-generation compression standard VVC and produces visually superior reconstructions when optimized by MS-SSIM.
7. Appendix
The appendix reports ablations of the network architecture and mixture count, latent-distribution visualizations, VVC configuration comparisons, and boundary-handling procedures. The reported findings favor the proposed backbone, Gaussian mixtures with K ≥ 3, and YUV444 VVC evaluation.
- 7.1.1 Backbone: The proposed backbone reduces the rate by about 6% from 0.49bpp to 0.46bpp while providing slightly better quality than prior network architectures.The comparison uses the proposed four 3 × 3 convolution design with residual connections as the anchor architecture.
- 7.1.2 Number of mixtures: Mixture models achieve lower loss than a single Gaussian model, while rate-distortion performance nearly saturates when K equals 3, 4, or 5.The paper uses K = 3 empirically.
- 7.1.3 Estimated distributions: Single Gaussian distributions are symmetric and fixed-shaped, whereas Gaussian mixtures represent more flexible and arbitrary discrete likelihood shapes for latent codes.The visualization examines representative latent-code locations in Kodim21 from the Kodak dataset.
- 7.2.1 Versatile Video Coding (VVC): VVC evaluation uses official VTM 5.2 commands with RGB images converted to YUV444 or YUV420 and quantization parameters from 22, 27, 32, 37, 42, and 47.The appendix specifies separate command-line configurations for the two chroma formats.
- 7.2.1 Versatile Video Coding (VVC): YUV420 performs worse than YUV444 because chroma-component sampling causes quality loss, especially at high rates; therefore, the paper uses YUV444 for comparison.
- 7.2.2 Boundary handling: For arbitrary image sizes, VVC inputs are reflect-padded to multiples of 8, while the learned codec requires padding to multiples of 64 before decoding and cropping.Image dimensions are encoded in the bitstream so the original size can be restored after decoding.