[Paper Review] Neural Image Compression via Non-Local Attention Optimization and Improved Context Modeling

Tong Chen, Haojie Liu+4 more|arXiv (Cornell University)|Oct 11, 2019

Advanced Data Compression TechniquesComputer Science37 references36 citations

TL;DR

This paper introduces NLAIC, a CNN-based variational autoencoder for image compression that uses non-local attention and improved 3D context modeling to capture local and global correlations, achieving state-of-the-art rate-distortion performance on Kodak and CLIC datasets.

ABSTRACT

This paper proposes a novel Non-Local Attention optmization and Improved Context modeling-based image compression (NLAIC) algorithm, which is built on top of the deep nerual network (DNN)-based variational auto-encoder (VAE) structure. Our NLAIC 1) embeds non-local network operations as non-linear transforms in the encoders and decoders for both the image and the latent representation probability information (known as hyperprior) to capture both local and global correlations, 2) applies attention mechanism to generate masks that are used to weigh the features, which implicitly adapt bit allocation for feature elements based on their importance, and 3) implements the improved conditional entropy modeling of latent features using joint 3D convolutional neural network (CNN)-based autoregressive contexts and hyperpriors. Towards the practical application, additional enhancements are also introduced to speed up processing (e.g., parallel 3D CNN-based context prediction), reduce memory consumption (e.g., sparse non-local processing) and alleviate the implementation complexity (e.g., unified model for variable rates without re-training). The proposed model outperforms existing methods on Kodak and CLIC datasets with the state-of-the-art compression efficiency reported, including learned and conventional (e.g., BPG, JPEG2000, JPEG) image compression methods, for both PSNR and MS-SSIM distortion metrics.

Motivation & Objective

• Motivate improved compression by capturing both local and global image correlations.
• Develop a VAE-based coder with embedded non-local attention to adapt bit allocation without explicit signaling.
• Enhance entropy modeling with joint 3D autoregressive contexts and hyperpriors for better latent feature prediction.
• Propose practical extensions to reduce memory and computation for real-world deployment.
• Demonstrate superior rate-distortion performance against learned and traditional codecs on standard benchmarks.

Proposed method

• Embed non-local network operations as non-linear transforms in main and hyperprior encoders/decoders to capture local and global correlations.
• Apply an attention mechanism to generate masks that implicitly adapt bit allocation across features without signaling overhead.
• Use a 3D masked convolution based autoregressive context model jointly with hyperpriors for conditional entropy estimation of latent features.
• Incorporate a non-local attention module (NLAM) with multi-layer attention masks to guide end-to-end learning.
• Introduce complexity-reduction options: sparse non-local processing, parallel 3D context modeling, and a unified variable-rate model via quality scaling factors.
• Train end-to-end with rate-distortion objective combining latent and hyperprior rates and a distortion term (MSE or MS-SSIM).

Experimental results

Research questions

• RQ1 How can non-local correlations be exploited in learned image compression to improve rate-distortion performance?
• RQ2 Can attention-generated implicit masks enable adaptive bit allocation without explicit signaling overhead?
• RQ3 Does joint 3D autoregressive context together with hyperpriors improve entropy modeling for latent features?
• RQ4 What practical extensions (sparse NLAM, parallel 3D contexts, unified variable-rate model) improve speed and memory with negligible performance loss?
• RQ5 How does NLAIC perform against traditional codecs and other learned compressors on standard benchmarks (Kodak, CLIC)?

Key findings

• NLAIC outperforms existing learned and traditional codecs on Kodak and CLIC datasets for both PSNR and MS-SSIM.
• Introduces non-local attention masks across multiple layers to enable adaptive processing without explicit bit allocation signaling.
• A 3D masked convolution context model jointly using autoregressive spatial-channel neighbors and hyperpriors yields more accurate conditional statistics for entropy coding.
• Proposed practical extensions (sparse NLAM, parallel 3D context modeling, and unified variable-rate model) reduce memory and computation with negligible coding efficiency loss.
• BD-Rate gains compared to several anchors: 64.39% over JPEG anchors, 59.84% over Minnen2018, 59.46% over BPG (YCbCr 4:4:4) HM, 56.19% over Ballé2018, and 38.02% over JPEG2000.