PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Image Compression via Non-Local Attention Optimization and Improved Context Modeling

Tong Chen, Haojie Liu, Zhan Ma, Qiu Shen, Xun Cao, Yao Wang

arXiv:1910.06244v1eess.IV

TL;DR

Existing image compression methods are limited in capturing non-local correlations and accounting for unequal feature importance. NLAIC addresses these gaps with non-local attention and improved context modeling, achieving state-of-the-art rate-distortion performance on Kodak and CLIC for PSNR and MS-SSIM.

  • Problem

    Existing compression methods rely largely on local operations and treat latent features with equal importance despite diverse visual sensitivities.

  • Method

    NLAIC integrates non-local attention modules, importance masks, and masked 3D CNN context prediction within a VAE-based image compression framework.

  • Results

    NLAIC achieves state-of-the-art rate-distortion performance on Kodak and CLIC for both MS-SSIM and PSNR, outperforming learned and traditional methods.

  • Takeaways & Limitations

    Joint context modeling produces more compact latent features and lower normalized prediction error, supporting the model’s leading coding efficiency.

Abstract

from arXiv · show

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.

I. INTRODUCTION

NLAIC addresses the limited receptive fields and equal feature treatment of existing learned image compression methods by combining non-local processing, attention-based weighting, and improved context modeling. It also adds practical simplifications while reporting state-of-the-art rate-distortion performance.

  • Motivation: Existing learned coders mainly use local operations and often treat latent features equally, limiting correlation modeling and adaptive bit allocation.Existing importance-map methods generally operate only at the bottleneck and may require explicit signaling.
  • Results: 64.39% BD-Rate gains are reported for NLAIC against the same JPEG anchors, exceeding the listed learned and conventional baselines.The paper reports evaluation on Kodak and CLIC using PSNR and MS-SSIM rate-distortion performance.
  • Contributions: NLAIC introduces non-local operations into the compression framework to capture local and global correlations among pixels and latent features.
  • Contributions: Attention masks are applied at multiple layers to adapt processing and allocate more bits to important feature areas.
  • Contributions: A masked 3D CNN jointly models spatial-channel correlations and combines its output with hyperpriors for conditional entropy estimation.
  • Practical enhancements: Sparse non-local processing, parallel context modeling, and a unified variable-rate model reduce complexity and improve implementation friendliness.The reported reductions concern computational complexity and memory storage, with negligible coding-efficiency sacrifice.

II. RELATED WORK

Related work spans non-local processing, self-attention, and learned image-compression architectures. NLAIC builds on these directions while targeting explicit-signaling overhead, local context limits, and model complexity.

  • Non-local Operations: Non-local methods improve image restoration by exploiting correlations beyond spatially neighboring pixels.
  • Self Attention: Self-attention provides spatially adaptive activation and has been applied to low-level vision and compression.
  • Self Attention: Compression attention masks support adaptive bit allocation, but explicit signaling can add overhead or induce probability-estimation errors.
  • Image Compression Architectures: RNN-based compression methods can support variable rates with a single model, while CNN-based approaches have attracted broad research attention.
  • Image Compression Architectures: Learned image compression uses autoencoders, quantization, and conditional entropy modeling with autoregressive contexts and hyperpriors.

III. NLAIC: NON-LOCAL ATTENTION OPTIMIZED IMAGE COMPRESSION

NLAIC uses a VAE with non-local attention modules in the main and hyperprior transforms, then estimates latent-feature probabilities from hyperpriors and autoregressive spatial-channel contexts. Its non-local block aggregates globally related feature values and can augment them through residual processing.

  • Optimization: Training minimizes a rate-distortion objective combining latent and hyperprior entropy rates with weighted image distortion.The framework supports MSE or negative MS-SSIM distortion in the described study.
  • Architecture: NLAIC is built on a VAE whose main and hyperprior encoder-decoder pairs contain non-local attention modules.The main encoder and decoder generate and reconstruct quantized latent features, while the hyperprior pair supplies side information.
  • Entropy modeling: Hyperpriors and autoregressive spatial-channel neighbors are processed by a conditional context model to estimate probabilities for entropy coding.
  • Non-local Network Processing: The non-local network computes each output position as a normalized weighted average of transformed features across accessible input positions.The correlation function compares input positions, while the representation function transforms the input values.
  • Non-local Network Processing: A residual connection applies a channel-mixing 1×1 convolution to produce an output with augmented global and local correlations.

B. Non-local Attention Module (NLAM)

NLAM combines non-local processing and attention to generate multi-scale spatial-channel masks that adaptively weight latent features within the VAE. This design captures local and global correlations while avoiding explicit mask signaling.

  • Attention masks: NLAM generates implicit attention masks to allocate more bits to visually important latent features.The masks weight feature-map elements adaptively, supporting rate-distortion efficiency without requiring explicit mask signaling.
  • Mask generation: The mask branch combines NLN processing, residual blocks, 1×1 convolution, and sigmoid activation to produce a joint spatial-channel mask.Each mask element satisfies 0 < M_k < 1 and is multiplied element-wise with the corresponding feature-map element.
  • Multi-scale integration: Multiple NLAM pairs are embedded across main and hyperprior encoder-decoder layers to exploit local and non-local correlations at multiple scales.Unlike bottleneck-only masking, different layers provide masks with multiple levels of granularity.
  • Activation design: The model uses ReLU instead of GDN because NLAM captures non-local correlations efficiently while maintaining effective processing.The paper reports this choice as effective based on its experimental observations.

C. Conditional Entropy Rate Modeling

NLAIC models hyperpriors with a learned factorized density and models quantized latent features conditionally using hyperprior-predicted Gaussian parameters. The resulting probabilities provide the estimated rates used for entropy coding and rate-distortion optimization.

  • Hyperprior modeling: A non-parametric, fully factorized density model estimates the probability distribution of quantized hyperpriors.Each univariate distribution is parameterized by ψ(i).
  • Latent-feature modeling: Each quantized latent feature x̂_i is modeled with a conditional Gaussian distribution whose μ_i and σ_i are predicted from hyperpriors ẑ.These parameters support conditional probability estimation for the latent representation.
  • Rate estimation: The estimated rates R̂_x and R̂_z are computed from the negative base-2 logarithms of the modeled latent and hyperprior probabilities.The equations evaluate the bits associated with quantized latent features and hyperpriors.
  • Side information: Hyperpriors serve as side information for estimating latent-feature distribution parameters and occupy a very small fraction of the total bits.The paper states that this small bit fraction is shown in Section VI and Fig. 13.

2) Context Modeling Using Joint Autoregressive Spatial-Channel Neighbors and Hyperpriors:

NLAIC uses joint 3D autoregressive context modeling to capture spatial and cross-channel dependencies while combining causal latent neighbors with hyperpriors for probability estimation. Shared kernels reduce the context model’s parameter count relative to 2D masked CNNs.

  • Parameter sharing: Unlike the 2D masked CNN approach, 3D masked convolutions implicitly capture correlations across adjacent channels without requiring a carefully structured channel order.The approach enforces shared kernels across the spatial-channel space, significantly reducing network parameters.
  • Conditional probability: The conditional model estimates each current latent element from causal preceding elements and hyperpriors.The causal elements are those prior to the current element in the feature map, and they contribute to estimating μ_i and σ_i.

IV. EXTENSIONS OF NLAIC FOR COMPLEXITY REDUCTION

NLAIC adds complexity-reduction extensions targeting memory, computation, and deployment across bitrates. These include sparse non-local processing, parallel masked 3D convolution, and a unified variable-rate model.

  • B. Parallel context modeling: Parallel masked 3D convolution reduces the computational complexity associated with sequential context modeling.The extension targets the throughput penalty caused by raster-scan autoregressive processing.
  • C. Unified variable-rate model: A unified neural model supports variable bitrates through quality mapping factors without requiring separate retraining for each rate.This reduces model complexity for practical application.
  • Practical impact: The proposed extensions reduce space and time complexity with negligible sacrifice of coding efficiency.The paper frames these changes as improvements in practical implementation friendliness.
  • A. Sparse NLAM: Sparse NLAM reduces memory consumption by downsampling before constructing the non-local correlation matrix.The downsampling scales down the matrix involved in non-local processing.

B. Parallel 3D Masked CNN based Context Modeling

The context model removes spatial dependencies to expose parallelism, reducing theoretical processing complexity while preserving coding performance. The broader system also supports variable rates and complexity reductions for practical deployment.

  • Parallel context modeling: Strictly sequential masked convolutions impose O(H × W × C) processing complexity because each latent element depends on raster-scan neighbors.
  • Parallel context modeling: Removing left-neighbor dependencies enables line-wise parallel processing and reduces theoretical complexity to O(H × C).
  • Parallel context modeling: Removing vertical neighbors allows all pixels within each channel to process simultaneously, reducing theoretical complexity to O(C) with negligible performance impact.
  • Unified variable-rate model: A unified variable-rate model is motivated by the memory and switching overhead of maintaining separate models for different bitrates.
  • Unified variable-rate model: Variable-rate retraining shows similar kernel and feature-map patterns with scaled intensities, motivating scaling factors instead of retraining complete models.
  • Unified variable-rate model: Scaling factors modify encoder feature maps before quantization, while inverse scaling restores entropy-decoded elements before decoding; the context model remains fixed.
  • Complexity reduction: Hyperpriors contribute only 2%-8% overhead, supporting their use as a fixed component while scaling adapts the main codec across rates.

V. EXPERIMENTAL STUDIES

The experiments train NLAIC on COCO and CLIC using cropped patches and end-to-end rate-distortion optimization with both MSE and negative MS-SSIM losses.

  • Training: NLAIC is trained on COCO and CLIC using randomly cropped 192×192×3 patches.
  • Training: End-to-end training minimizes L = λ·d(Ŷ, Y) + R_x + R_z across various bitrates, using MSE and negative MS-SSIM as distortion losses.
  • Training: The rate terms estimate latent-feature and hyperprior bitrates, while all NLAIC components are trained jointly.

B. Rate-Distortion Efficiency

On Kodak and BSD500 evaluations, NLAIC achieves leading objective rate-distortion performance and the strongest reported subjective quality among the compared codecs.

  • Objective measurement: NLAIC ranks first for both MS-SSIM-based and PSNR-based objective coding efficiency.
  • Subjective evaluation: Subjective comparisons on BSD500 report the best visual quality for NLAIC at a relatively smaller bitrate, with closer-rate matching for JPEG and BPG.
  • Implementation measurements: The tested model is about 262MB and requires 6172MB running memory for a 512×768×3 input during encoding.

VI. ABLATION STUDIES

The ablations show that non-local attention, joint context modeling, and rate-scaling choices each affect NLAIC’s coding efficiency and practical complexity. The studies also identify lower-bitrate limitations and loss-dependent reconstruction trade-offs.

  • Impacts of Loss Functions: MS-SSIM-optimized reconstructions exhibit worse details than PSNR-optimized models at high bit rate.The paper attributes this to differing emphasis on pixel distortion and structural similarity across bitrate ranges.
  • Impacts of Contexts: With NLAM and joint contexts, latent features become more compact and normalized prediction errors remain lower, supporting the leading coding efficiency.The joint modeling combines hyperpriors with autoregressive neighbors.
  • Impacts of Contexts: More than 9% bit rate reduction at the same quality is achieved when all NLAIC modules are refined over the baseline.Transfer learning that updates only the conditional context model yields a 3% reduction, while refining all modules performs better.
  • Impacts of NLAM: About 0.1dB PSNR is lost after removing the first NLAM mask branch, with further degradation when additional mask branches are removed.Disabling mask branches in both main and hyperprior encoder-decoders produces the worst performance.
  • Impacts of Parallel Context Modeling: ≈0% BD-Rate loss is measured with contexts excluding the left neighbor, compared with 3.73% using channel neighbors only and 6.76% with hyperprior only.The parallel context design removes left-neighbor dependency and greatly reduces computational complexity.
  • Hyperpriors: The percentage of hyperprior bits decreases as overall bit rate increases and is higher for MSE-optimized models than for MS-SSIM-optimized models.Hyperprior bits remain much smaller than the bits consumed by latent features.
  • Impacts of Unified Model for Variable Rates: Three unified models are suggested across low, medium, and high bitrate ranges because scaling and inverse scaling can degrade coding efficiency at lower bitrates.Fixed hyperprior bitrate across the range also causes performance loss at low bitrates.

VII. CONCLUDING REMARKS AND FUTURE WORKS

The paper concludes that NLAIC achieves state-of-the-art image-compression performance through non-local attention and improved conditional context modeling. It also introduces practical complexity reductions, while identifying video compression and embedded implementation as future directions.

  • Concluding Remarks: NLAIC achieves state-of-the-art performance for both MS-SSIM and PSNR at the same bit rate against conventional and learning-based methods.The comparisons include BPG, JPEG2000, JPEG, and recent learned schemes.
  • Concluding Remarks: Non-local transforms capture local and global correlations, attention masks weight salient areas, and joint 3D CNN contexts improve entropy modeling.These are the paper’s principal architectural novelties.
  • Concluding Remarks: Sparse non-local processing, parallel context modeling, and unified variable-rate models target lower space and time complexity and easier implementation.The practical enhancements are designed for deployment-oriented use.
  • Future Work: Future work extends the framework toward end-to-end video compression and simpler embedded-system implementations.Suggested implementation directions include fixed-point operation and platform-friendly network structures.
Loading 1910.06244v1…