Source-linked AI summary
Non-local Attention Optimized Deep Image Compression
Haojie Liu, Tong Chen, Peiyao Guo, Qiu Shen, Xun Cao, Yao Wang, Zhan Ma
TL;DR
Learned image compression methods can have limited receptive fields, treat latent features equally, or require signaling for importance maps. NLAIC combines non-local operations, attention, and joint context modeling in a VAE framework; it outperforms existing learned and traditional methods on both MS-SSIM and PSNR. The paper also notes that hyperprior bit allocation needs further study.
Problem
Existing methods may have limited receptive fields or treat latent features equally, and importance-map approaches can require extra signaling overhead.
Source: §1 · p. 1Method
NLAIC embeds non-local operations and attention in a VAE compression framework and uses hyperpriors with spatial-channel latent contexts for entropy coding.
Source: §1 · pp. 1–2Results
Takeaways & Limitations
The reported comparison supports NLAIC's rate-distortion advantage over conventional and learning-based methods at the same bit rate.
Source: §5 · p. 7The authors call for further study of hyperprior bit allocation after observing different distributions across joint and baseline models and distortion-loss schemes.
Source: §4.3 · p. 7
Abstract
from arXiv · showhide
This paper proposes a novel Non-Local Attention Optimized Deep Image Compression (NLAIC) framework, which is built on top of the popular variational auto-encoder (VAE) structure. Our NLAIC framework embeds non-local operations in the encoders and decoders for both image and latent feature probability information (known as hyperprior) to capture both local and global correlations, and apply attention mechanism to generate masks that are used to weigh the features for the image and hyperprior, which implicitly adapt bit allocation for different features based on their importance. Furthermore, both hyperpriors and spatial-channel neighbors of the latent features are used to improve entropy coding. The proposed model outperforms the existing methods on Kodak dataset, including learned (e.g., Balle2019, Balle2018) and conventional (e.g., BPG, JPEG2000, JPEG) image compression methods, for both PSNR and MS-SSIM distortion metrics.
1. Introduction
Existing compression methods have limited receptive fields or treat latent features equally, while explicit importance maps add signaling overhead. NLAIC combines non-local processing, attention masks, and context modeling, and its joint model outperforms existing learned and traditional methods on both distortion metrics.
1. Introduction
- Existing methods often have limited receptive fields or treat latent features equally, while explicit importance maps require signaling overhead.These limitations can overlook differing visual sensitivities to content such as texture and edges.Source: §1 · p. 1
- Attention masks generated at multiple layers guide adaptive latent-feature processing and allocate more bits to features important for reducing distortion.The masks are produced through end-to-end training rather than being limited to quantized features.Source: §1 · p. 2
- NLAIC joint outperforms existing learned and traditional compression methods in rate-distortion efficiency for both MS-SSIM and PSNR.The joint model uses hyperpriors and spatial-channel latent neighbors for context modeling.Source: §1 · p. 2
- Ablations examine variants that remove non-local operations or attention mechanisms layer by layer, alongside visual comparisons.The paper presents these experiments as further evidence of performance over a broad dataset.Source: §1 · p. 2
- NLAIC introduces non-local operations into the VAE compression framework to capture local and global correlations among pixels and feature maps.The framework embeds non-local attention optimization in the main and hyperprior encoder-decoder pairs.Source: §1 · pp. 1–2, Figure 1 · p. 1
- A masked 3D CNN models spatial and cross-channel latent correlations, then combines its output with hyperpriors to estimate conditional statistics for entropy coding.The contribution passage identifies the masked 3D CNN as one layer.Source: §1 · p. 2
2. Related Work
Related work draws on non-local methods, self-attention, and learned compression architectures. Prior attention-based bit allocation can require explicit signaling, while entropy modeling combines latent-feature contexts with hyperpriors.
2. Related Work
- Non-local operations have been used in video classification and image restoration, with other applications including HEVC intra block copy.The cited restoration tasks include denoising, artifact removal, and super-resolution.Source: §2 · p. 2
- Self-attention supports spatially adaptive activation and information allocation in low-level vision, while also being widely used in NLP.The cited vision examples emphasize challenging areas such as rich textures and salient regions.Source: §2 · p. 2
- Prior compression methods use quantized attention masks for bit allocation, but require extra signaling; NLAIC instead generates masks automatically from the input image.The prior methods cited are Li et al. [11] and Mentzer et al..Source: §2 · p. 2
- Learned compression uses autoencoders and approximations to non-differentiable quantization, while PixelCNNs and VAEs support entropy estimation.Conditional estimates jointly using autoregressive latent neighbors and hyperpriors have improved entropy coding.Source: §2 · p. 2
3. Non-Local Attention Implementation
NLAIC builds on a variational autoencoder, using non-local attention modules in its main and hyperprior encoder-decoder pairs. Its entropy model uses hyperpriors and autoregressive latent neighbors to estimate conditional probabilities for coding.
3.1. General Framework
- NLAIC uses a variational autoencoder with non-local attention modules in both the main and hyperprior encoder-decoder pairs.The main encoder and quantization generate latent features; the hyperprior encoder-decoder generates side information used with latent neighbors for conditional probability estimates.Source: §3.1 · pp. 2–3
3.4. Entropy Rate Modeling
- The entropy-rate model estimates latent-feature conditional probabilities from hyperpriors and autoregressive neighbors.The paper presents hyperprior and neighbor context modeling as components of the entropy engine for rate-distortion efficiency.Source: §3.1 · p. 3, §3.4 · p. 4
3.2. Non-local Module
- The non-local module computes each output location as a weighted average of transformed input features across accessible positions.Its embedded Gaussian affinity is softmax-normalized, and a residual connection can be added to the output.Source: §3.2 · pp. 3–4
3.3. Non-local Attention Module
- The non-local attention module combines a non-local mask branch with a conventional feature branch to produce spatial-channel attention masks.The mask is multiplied elementwise with the main-branch features, with a residual connection; experiments selected two NLAM pairs in the main encoder-decoder and one layer in the hyperprior encoder-decoder.Source: §3.3 · p. 4
3.4.1 Context Modeling Using Hyperpriors
- Hyperpriors predict the conditional Gaussian parameters for quantized latent features, and their bitstream share varies with loss and model variant.The hyperprior share is higher for MSE-optimized than MS-SSIM-optimized models; joint versus baseline shares reverse between the two loss settings.Source: §3.4.1 · p. 4, Figure 3 · p. 4
3.4.2 Context Modeling Using Neighbors
- A one-layer 5×5×5 masked 3D convolution uses spatial and cross-channel correlations to improve conditional modeling while reducing parameters versus a 2D masked CNN.The autoregressive context comprises causal, possibly reconstructed latent features preceding the current feature.Source: §3.4.2 · p. 5
4. Experiments
On Kodak, NLAIC is reported to outperform existing learned and traditional codecs on rate-distortion performance, with results assessed using MS-SSIM and PSNR. The experiments compare training variants, evaluate NLAM and context-modeling contributions, and examine visual quality and metric trade-offs.
4.1. Training
- Refining all modules in NLAIC joint yields more than 9% bit-rate reduction over NLAIC baseline at the same quality, versus 3% when updating only the context model.Training uses COCO and CLIC with 192×192×3 crops and end-to-end rate-distortion optimization.Source: §4.1 · pp. 5–6
4.2. Performance Efficiency
- In coding-efficiency comparison with JPEG as anchor, NLAIC joint achieves the best BD-Rate gains among popular algorithms, reducing BD-Rate by 64.39% against JPEG420 and 12.26% against BPG444.Source: §4.2 · p. 6, Figure 6 · p. 6
4.3. Ablation Studies
- Removing the first NLAM mask branches drops PSNR by about 0.1 dB at the same bit rate; removing all main and hyperprior mask branches gives the worst performance.The ablation progressively removes mask branches and retrains models using baseline context modeling, with MSE loss and PSNR measurement.Source: §4.3 · pp. 6–7
- With NLAM and joint hyperprior-neighbor contexts, NLAIC joint has more powerful latent features and more compact normalized prediction errors, supporting its leading coding efficiency.The reported error distributions are associated with lower bit rate.Source: §4.3 · p. 7
- Hyperprior bit share decreases as overall bit rate increases, differs between MSE- and MS-SSIM-optimized models, and has contrasting joint-versus-baseline distributions across those schemes.The paper identifies further study of hyperprior bit allocation as needed.Source: §4.3 · p. 7
4.4. Visual Comparison
- On BSD500, NLAIC is reported to provide the best subjective quality at a relatively smaller bit rate, while MS-SSIM-optimized results show worse detail than PSNR-optimized results at high rates.The paper suggests exploring a metric that combines PSNR's high-rate advantages with MS-SSIM's low-rate advantages.Source: §4.4 · p. 7
5. Conclusion
NLAIC combines non-local processing, attention-based feature allocation, and joint spatial-channel and hyperprior contexts for image compression. It reports state-of-the-art rate-distortion performance, outperforming existing methods on both MS-SSIM and PSNR.
5. Conclusion
- Non-local operations capture local and global correlations, while attention maps allocate more bits to important areas and joint contexts improve entropy coding efficiency.The joint contexts combine autoregressive spatial-channel neighbors with hyperpriors.Source: §5 · p. 7
- At the same bit rate, NLAIC outperforms existing methods, including BPG, JPEG2000, JPEG, and recent learned schemes [5] [20], on both MS-SSIM and PSNR.Some BPG and JPEG bit-rate points were unavailable, so the closest rates were selected.Source: §5 · p. 7
- The authors identify a deeper context model as future work to improve image compression performance.Source: §5 · p. 7
- The visual comparison reports that NLAIC reconstructions have more texture and fewer blocky or blurring artifacts than JPEG420 and BPG444.Figure 9 compares JPEG420, BPG444, NLAIC joint MSE optimization, MS-SSIM optimization, and the original image.Source: Figure 9 · p. 8
- Acceleration is identified as important for practical deployment, particularly on mobile platforms.Source: §5 · p. 8
- Extending the framework to end-to-end video compression with additional spatial and temporal priors is identified as a future direction.Source: §5 · p. 9