Source-linked AI summary
MLIC: Multi-Reference Entropy Model for Learned Image Compression
Wei Jiang, Jiayu Yang, Yongqi Zhai, Peirong Ning, Feng Gao, Ronggang Wang
TL;DR
Most entropy models capture only one-dimensional correlations, although latent representations contain channel-wise, local spatial, and global spatial correlations. The paper proposes multi-reference entropy models that model these correlations jointly, and MLIC and MLIC+ achieve state-of-the-art performance, reducing Kodak PSNR BD-rate against VTM-17.0 by 8.05% and 11.39%.
Problem
Effectively capturing local, global, and channel-wise contexts within a single entropy model remains challenging, despite their potential to improve learned image compression.
Method
MEM and MEM+ divide latent representations into slices, use previously decoded slices as context, and combine channel-wise, local spatial, and global spatial context modeling.
Results
MLIC and MLIC+ achieve state-of-the-art performance, reducing Kodak PSNR BD-rate over VTM-17.0 by 8.05% and 11.39%, respectively.
Takeaways & Limitations
The work demonstrates the potential of entropy models that capture multiple correlations in latent representations for learned image compression.
Takeaways & Limitations
Global context models incorporated with serial autoregressive context models further increase decoding latency.
Abstract
from arXiv · showhide
Recently, learned image compression has achieved remarkable performance. The entropy model, which estimates the distribution of the latent representation, plays a crucial role in boosting rate-distortion performance. However, most entropy models only capture correlations in one dimension, while the latent representation contain channel-wise, local spatial, and global spatial correlations. To tackle this issue, we propose the Multi-Reference Entropy Model (MEM) and the advanced version, MEM$^+$. These models capture the different types of correlations present in latent representation. Specifically, We first divide the latent representation into slices. When decoding the current slice, we use previously decoded slices as context and employ the attention map of the previously decoded slice to predict global correlations in the current slice. To capture local contexts, we introduce two enhanced checkerboard context capturing techniques that avoids performance degradation. Based on MEM and MEM$^+$, we propose image compression models MLIC and MLIC$^+$. Extensive experimental evaluations demonstrate that our MLIC and MLIC$^+$ models achieve state-of-the-art performance, reducing BD-rate by $8.05\%$ and $11.39\%$ on the Kodak dataset compared to VTM-17.0 when measured in PSNR. Our code is available at https://github.com/JiangWeibeta/MLIC.
1 INTRODUCTION
The paper argues that learned compression needs entropy models that capture multiple latent correlations, and introduces MEM/MEM+ with MLIC/MLIC+ to address them.
- Motivation: Entropy coding improves learned compression by estimating the latent representation’s distribution for conditional-probability coding.The paper places entropy modeling within the transform, quantization, entropy-coding, and inverse-transform pipeline.
- Motivation: Most existing context models capture conditional probabilities in only one dimension, despite local spatial, global spatial, and channel-wise correlations.The paper identifies these three correlation types as distinct contexts in latent representations.
- Proposed approach: MEM and MEM+ combine local spatial, global spatial, channel-wise, and hyper-prior contexts within multi-reference entropy models.The proposed models divide latent representations into slices and use previously compressed slices as context.
- Proposed approach: MEM-based MLIC and MLIC+ use stacked checkerboard context modeling and checkerboard attention to retain two-pass decoding while addressing checkerboard degradation.The paper contrasts these techniques with checkerboard context modeling, which can cause performance degradation.
2 RELATED WORKS
This section reviews rate-distortion learning, entropy-model components, and context-modeling approaches, emphasizing the challenge of jointly capturing local, global, and channel-wise correlations.
- Learned compression framework: Rate-distortion optimization balances bit rate R and distortion D, with λ controlling the distortion weight and target bit rate.The stated objective is L = R + λD.
- Learned compression framework: The basic learned compression framework transforms an input image into a latent representation, quantizes it, entropy-codes it, and reconstructs the image.The analysis transform g_a produces y, Q produces ŷ, and synthesis transform g_s produces x̂.
- Entropy models: Hyper-prior models extract side information from the latent representation to estimate its conditional distribution for rate coding.Later work uses more flexible mean-scale Gaussian, asymmetric Gaussian, and Gaussian-mixture distributions.
- Context models: Local context models capture adjacent-symbol correlations, using either serial masked convolutions or two-pass parallel checkerboard decoding.Checkerboard decoding separates anchor and non-anchor parts and predicts the latter from the former.
- Context models: Global context models capture distant-symbol relationships, while channel-wise models condition slices on previously decoded channels or slices.Global and serial autoregressive context modeling can introduce extra bits or decoding latency.
- Open challenge: Effectively capturing local, global, and channel-wise contexts in one entropy model remains a challenge.The paper presents this multi-dimensional correlation problem as an opportunity for further performance improvement.
3 METHOD
MLIC and MLIC+ use multi-reference entropy models to capture channel-wise, local spatial, and global spatial correlations in sliced latent representations. Their context modules combine previously decoded slices, checkerboard-based local contexts, and attention-based global contexts.
- MLIC and MLIC+: MLIC uses MEM for lower complexity, whereas MLIC+ uses MEM+ for improved rate-distortion performance.Both models share the same analysis, synthesis, hyper-analysis, and hyper-synthesis transforms; they differ in their entropy models.
- MLIC and MLIC+: MEM and MEM+ combine channel-wise, local spatial, and global spatial contexts for latent probability estimation.The latent representation is divided into slices, with each current slice conditioned on previously decoded slices.
- Local spatial context: Checkerboard context modeling enables parallel decoding by conditioning non-anchor symbols on decoded anchor symbols, but can slightly degrade performance.The method addresses this limitation with stacked checkerboard context and shifted-window checkerboard attention modules.
- Stacked Checkerboard Context Module: The stacked checkerboard module increases context-model depth and nonlinearity by stacking convolutional layers.The design uses an odd number of convolutions so information alternately transfers between anchor and non-anchor parts; J=3 and K=5 are selected.
- Shifted-Window-based Checkerboard Attention: Checkerboard attention uses overlapped K × K windows and dynamically generated attention weights to capture content-adaptive local contexts.The attention formulation embeds the anchor part, applies an attention mask, and refines the result with convolution and a feed-forward network.
- Global spatial context: The intra-slice global context module uses the previous slice’s attention relationships to estimate global correlations in the current slice.A K × K convolution refines the attention map by aggregating global similarities of adjacent symbols, while masking prevents local interactions from dominating.
4 EXPERIMENTS
Experiments show that MLIC and MLIC+ achieve state-of-the-art rate-distortion performance while retaining competitive codec efficiency. Ablations attribute gains to multi-dimensional context modeling, including channel-wise, local spatial, and global spatial contexts.
- Rate-Distortion Performance: MLIC and MLIC+ achieve state-of-the-art performance across the evaluated datasets in both PSNR and MS-SSIM.The evaluation includes Kodak, Tecnick, CLIC Pro Val, CLIC’21 Test, CLIC’22 Test, and JPEGAI Test.
- Rate-Distortion Performance: 8.05% and 11.39% BD-rate reductions are achieved on Kodak over VVC for MLIC and MLIC+, respectively, measured in PSNR.Compared with Cheng’20, MLIC improves PSNR by up to 0.5–0.8 dB and MS-SSIM by up to 0.6 dB, while MLIC+ improves PSNR by up to 0.8–1.0 dB.
- Qualitative Results: MLIC reconstructed images retain more details at lower bpp and achieve PSNR values 1 dB higher than VTM-17.0 on the illustrated example.The qualitative comparison includes MLIC, MLIC+, Entroformer, Xie’21, Cheng’20, and VTM-17.0.
- Codec Efficiency: MLIC can encode and decode quite fast relative to other recent models, while MLIC+ takes slightly longer because its slice entropy model is more complex.The timing comparison includes arithmetic coding time and is conducted on Kodak.
- Ablation Studies: The channel-wise context module significantly improves performance, supporting redundancy among channels and access to symbols in the same or nearby positions of previous slices.This ablation supports using previously decoded slices as channel-wise context.
- Ablation Studies: 4.87% more bit-rates are saved by checkerboard attention than by the vanilla checkerboard context module in the ablation study.The stacked checkerboard module saves 1.81% more bit-rates than the vanilla checkerboard module, while vanilla checkerboard modeling causes slight degradation.
- Context Analysis: Global context modules capture distant correlations, but their quadratic computational complexity creates a practical efficiency constraint.Combining global and local context modules further improves performance; the inter-slice global-context gain is described as limited.
5 CONCLUSION
The conclusion presents MEM and MEM+ as entropy models that capture channel, local, and global correlations, yielding state-of-the-art MLIC and MLIC+ models. It also identifies computational overhead as a deployment barrier for mobile devices.
- Conclusion: MEM and MEM+ capture correlations across channel, local spatial, and global spatial dimensions in latent representations.The paper describes this as its first successful attempt to capture these three correlation types together.
- Conclusion: MLIC and MLIC+ achieve state-of-the-art performance using the proposed multi-reference entropy models.The models are presented as evidence of the potential of entropy modeling for learned image compression.
- Conclusion: Computational overhead prevents direct application of MLIC and MLIC+ to mobile devices.The paper suggests knowledge distillation, network pruning, structural re-parameterization, and other lightweight designs as possible remedies.
A MORE NETWROK ARCHITECTURE DETAILS
The supplementary architecture details split each slice into anchor and non-anchor parts and specify the modules used for channel, checkerboard, feed-forward, and inter-slice global context processing.
- Slice Structure: Each latent slice is divided into anchor and non-anchor parts during encoding.The notation distinguishes the anchor part from the non-anchor part within slice i.
- Module Architectures: The stacked checkerboard context module contains three 5 × 5 convolutional layers and uses GELU nonlinearities.The architecture details also reference the channel context, feed-forward, and inter-slice global context modules.
- Global Context: The inter-slice global context model is documented as the module g_gc,inter, with S denoting the channel number of a slice.Its architecture is illustrated in Figure 19.
B DETAILED EXPERIMENT SETTINGS
The implementation uses PyTorch and CompressAI on specified GPU and CPU hardware, with deterministic inference enabled for speed measurements.
- Implementation: MLIC and MLIC+ are implemented with PyTorch 1.10.0, CompressAI 1.2.0b3, and Python 3.9.7.Encoding and decoding latency is tested on a GeForce RTX 3090 GPU and Xeon Silver 4210R under Ubuntu 20.04.
- Speed Evaluation: Deterministic inference mode is enabled when testing model speeds.This setting applies to the reported encoding and decoding latency measurements.
C EXPERIMENTS ON OTHER BACKBONE AND SIMPLER GLOBAL CONTEXT MODULE
MLIC-EX replaces the Cheng’20 backbone with ELIC, simplifies global context by sharing the first slice’s attention map, and removes two modules to reduce complexity. Despite a slight performance drop from shared attention maps, it outperforms ELIC and ranked third in VCIP’22’s performance track.
- MLIC-EX replaces the Cheng’20 backbone with ELIC and evenly divides the latent representation into slices.
- MLIC-EX shares the first slice’s attention map to predict global correlations in other slices and removes inter-slice global context and latent residual prediction modules.These changes simplify the intra-slice global context module and reduce complexity.
- A slight performance drop from sharing attention maps still leaves MLIC-EX outperforming ELIC.The model was optimized for MSE with λ set to 0.08 for high-bit-rate evaluation.
- MLIC-EX ranked 3-rd in the performance track of the VCIP’22 Practical End-to-End Image Compression Challenge.
D MORE RATE-DISTORTION PERFORMANCE RESULTS
The paper reports rate-distortion evaluations across Kodak, Tecnick, CLIC, and JPEGAI datasets, comparing MLIC and MLIC+ with learned and traditional image compression methods. Dataset-specific preprocessing includes padding images to multiples of 64 for several benchmarks.
- The evaluations compare MLIC and MLIC+ with learned methods including ELIC, Contextformer, Entroformer, and STF, plus JPEG, JPEG2000, BPG, and VTM.
- Images are padded to multiples of 64 when evaluating CLIC Professional Valid, Tecnick, and CLIC2021 Test.
- VTM-17.0 rate-distortion data are provided for other researchers to compare performance because encoding takes a long time.
- Rate-distortion data are reported for MLIC and MLIC+ on Kodak, Tecnick, CLIC Professional Valid, CLIC 2021 Test, and JPEGAI Test datasets.The datasets contain 24, 100, 41, 60 or 30, and 16 raw images, respectively, as specified by the figure captions.
E.1 Visualization of Reconstructed Images
The paper visualizes reconstructed Kodak images and attention maps, and tests resolution effects by compressing non-overlapping 448 × 448 patches. Patch-based compression shows almost no performance degradation on the CLIC Professional Valid dataset.
- Visualization of Reconstructed Images: MLIC and MLIC+ reconstructed images are compared with Xie’21, Cheng’20, and VTM-17.0 on Kodak examples.
- Visualization of Reconstructed Images: Additional attention maps are provided to demonstrate the effectiveness of the global context modules.
- Influence of Resolution: Almost no performance degradation occurs when MLIC and MLIC+ compress non-overlapping 448 × 448 patches.The test uses CLIC Professional Valid images cropped to multiples of 448, avoiding padding zeros.
G PROGRESSIVE DECODING ANALYSIS
MLIC supports progressive decoding, but its performance is substantially worse and unstable compared with separately optimized models. The authors attribute these issues to global spatial context modules, especially inter-slice interactions.
- MLIC’s progressive decoding performance is much worse than that of separately optimized models.The progressive decoding results are illustrated on Kodak and JPEGAI Test datasets.
- Progressive decoding can produce unpleasant artifacts, noise, or crashes, and its performance is not stable.The paper includes examples of crashed progressive decoding.
- MLIC+ has even worse progressive decoding results than MLIC.
- The authors attribute the problems to global spatial context modules that change information distribution between slices and impose extra constraints.
- Channel-wise context models are described as suitable for progressive decoding, although changing information between slices remains necessary for optimization.