PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Statistical Fidelity for Neural Image Compression with Implicit Local Likelihood Models

Matthew J. Muckley, Alaaeldin El-Nouby, Karen Ullrich, Hervé Jégou, Jakob Verbeek

arXiv:2301.11189v3eess.IVcs.AIcs.CVcs.ITmath.IT

TL;DR

Distortion-focused image compression can produce reconstructions whose statistics diverge from the originals, particularly at low bitrates. The paper introduces a VQ-VAE-conditioned local discriminator for compression. It reports better statistical fidelity than HiFiC at matched PSNR, and higher PSNR at matched FID/KID.

  • Problem

    At low rates, optimizing distortion can pull reconstruction statistics away from the original image distribution, while existing compression approaches have primarily used discriminators designed for image generation.

    Source: §1 · p. 2

  • Method

    The proposed ILLM uses VQ-VAE autoencoders to condition a local adversarial discriminator, combined with a Mean-Scale Hyperprior compressor as MS-ILLM.

    Source: §1 · p. 2

  • Results

    At matched PSNR, MS-ILLM has uniformly better FID and KID than HiFiC across bitrates; at matched FID/KID, it achieves higher PSNR.

    Source: §5.3 · p. 7

  • Takeaways & Limitations

    MS-ILLM improves statistical fidelity while maintaining distortion performance relative to HiFiC, and the authors report improved FID and KID metrics.

    Source: §5.3 · p. 7, §7 · p. 9

    Improved FID/KID does not guarantee better human preference, and the authors advise further assessment of demographic bias before production use.

    Source: §6 · p. 9

Abstract

from arXiv · show

Lossy image compression aims to represent images in as few bits as possible while maintaining fidelity to the original. Theoretical results indicate that optimizing distortion metrics such as PSNR or MS-SSIM necessarily leads to a discrepancy in the statistics of original images from those of reconstructions, in particular at low bitrates, often manifested by the blurring of the compressed images. Previous work has leveraged adversarial discriminators to improve statistical fidelity. Yet these binary discriminators adopted from generative modeling tasks may not be ideal for image compression. In this paper, we introduce a non-binary discriminator that is conditioned on quantized local image representations obtained via VQ-VAE autoencoders. Our evaluations on the CLIC2020, DIV2K and Kodak datasets show that our discriminator is more effective for jointly optimizing distortion (e.g., PSNR) and statistical fidelity (e.g., FID) than the PatchGAN of the state-of-the-art HiFiC model. On CLIC2020, we obtain the same FID as HiFiC with 30-40\% fewer bits.

1. Introduction

The introduction frames rate-distortion optimization as a central compression objective, but distortion-focused optimization can diverge from the original image statistics, especially at low rates. It proposes a VQ-VAE-based local discriminator and reports improved distortion–statistical-fidelity trade-offs against PatchGAN.

1. Introduction

  • Compression has traditionally combined an autoencoding transform, lossy quantization, and lossless entropy coding; neural end-to-end systems have surpassed JPEG, JPEG2000, and BPG in rate-distortion performance.The introduction positions neural systems as a route to further improvements in the transform component.Source: §1 · p. 1
  • Rate-distortion optimization is common to traditional and learned compressors, often using handcrafted distortion metrics such as PSNR or MS-SSIM.The introduction presents this objective as the context for the statistical-fidelity trade-off.Source: §1 · pp. 1–2
  • At very low rates, distortion optimization necessarily pulls reconstruction statistics away from the true distribution, typically through blurring or smoothing.Prior adversarial approaches balance distortion loss with a discriminator; decode-time balancing has also been studied.Source: §1 · p. 2
  • The proposed implicit local likelihood model (ILLM) uses VQ-VAE autoencoders to optimize likelihoods in local image neighborhoods, and forms MS-ILLM with the Mean-Scale Hyperprior architecture.The approach adapts adversarial training to emphasize locality, aligning discriminator modeling with compression’s task.Source: §1 · p. 2
  • Across CLIC2020, DIV2K, and Kodak, MS-ILLM surpasses HiFiC’s FID scores without sacrificing PSNR.Figure 1 also reports better distortion–statistical-fidelity trade-offs than PatchGAN across compression regimes.Source: Figure 1 · p. 1, §1 · p. 2
  • The authors ablate latent dimensions for both the VQ-VAE labeler and U-Net discriminator to validate their architectural choices.The introduction identifies this as an experimental contribution.Source: §1 · p. 2

2. Related work

Related work covers distortion and distributional metrics, neural compression objectives, and the theoretical rate-distortion-perception trade-off. It also distinguishes this compression approach from semantic image generation methods such as OASIS.

Distortion and divergence metrics

  • MSE and MS-SSIM are handcrafted distortion metrics, while LPIPS and PIM use learned features; FID and KID assess distributional alignment without measuring distance to reference data.The section also discusses no-reference NIQE and contrastively learned distortion metrics.Source: §2 · p. 2

Neural image compression methods

  • Neural compression can optimize handcrafted distortion metrics, while perceptual optimization can produce artifacts; weighted GAN and distortion losses are used for perceptual benefits with stability.Other approaches use Wasserstein distance or KL-divergence of deep latent-variable models.Source: §2 · p. 2

Rate-distortion-perception tradeoff

  • Rate-distortion-perception studies find that improving distortion at a given rate can reduce perceptual quality, and realism generally comes at a rate-distortion cost.The section reports that perfect realism can be achieved with at most a two-fold increase in MSE.Source: §2 · p. 2

Minimizing distributional divergence

  • This work relates to OASIS but applies its model to compression and uses a latent-space VQ-VAE semantic map rather than pixel-space semantic maps.The section also reviews f-divergences and f-GANs as approaches to minimizing distributional divergence.Source: §2 · p. 3

3. Background

The background frames compression as minimizing bitrate under distortion and statistical-fidelity constraints, then describes neural codecs and adversarial divergence minimization as ways to optimize these goals.

3.1. Notation

  • The paper’s notation treats X and Y as random variables, uses lowercase letters for realizations, and distinguishes distributions from probability mass functions.Conditional distributions and expectations are also specified.Source: §3.1 · p. 3

3. Background

  • Rate-distortion theory formalizes compression as finding the lowest bitrate that meets a predefined reconstruction-quality threshold.The background situates this theory within a broader review of rate-distortion-perception theory and neural codec design.Source: §3.2 · p. 3, §3 · p. 3

3.2. Rate-distortion-perception theory

  • Rate-distortion-perception theory adds a statistical-fidelity constraint to bitrate and distortion, with improved statistical fidelity often costing distortion at constant rate.The paper uses “statistical fidelity” to distinguish distributional similarity from human perception and metrics such as LPIPS.Source: §3.2 · p. 3

3.3. Lossy codec optimization

  • A lossy codec encodes a source into a quantized latent, entropy-codes it, and decodes it into a reconstruction; its learned components include an encoder, decoder, and entropy model.The codec is optimized through a parameterized training objective that relaxes the rate-distortion constraints.Source: §3.3 · pp. 3–4, §3.2 · p. 3

3.4. Approximating the distributional divergence

  • GAN-based Jensen-Shannon divergence minimization can approximate distributional divergence, and the compression task motivates redesigning its discriminator for local projection.The generator loss can replace the divergence term in the relaxed objective, with alternating discriminator minimization.Source: §3.4 · p. 4, §3.3 · p. 4

4. Method

The method combines a mean-scale hyperprior codec with a discriminator that uses spatially distributed labels from a VQ-VAE. Its training and architecture are designed to enforce local likelihood comparisons for reconstructed images.

4. Method

  • MS-ILLM follows HiFiC while replacing its discriminator with an OASIS-type vector-valued design for modeling likelihoods near compressed images.The system encodes and quantizes an image, entropy-codes the latent, decodes a reconstruction, and trains a discriminator to match pretrained-labeler outputs.Source: §4 · p. 4, Figure 2 · p. 4

4.1. Autoencoder architecture

  • The codec uses Minnen et al.’s mean-scale hyperprior and follows Mentzer et al.’s larger, deeper encoder-decoder architecture.The hyperprior models the latent marginal through a hierarchical latent-variable model.Source: §4.1 · p. 4

4.2. Implicit local likelihood models

  • The discriminator predicts original-image labels as well as a fake class, while its generator loss makes reconstructions maximize the likelihood of their corresponding real-image labels.The vector-valued label loss conditions the discriminator on local label information and enforces locality in the loss rather than merely offering latent input.Source: §4.2 · p. 5

4.3. Choice of labeling function

  • VQ-VAE labels partition latent vectors into codebook clusters, making locality in label space imply locality in image space.The default setup uses a 32×32 latent space for 256×256 images and a codebook of 1024; the paper reports little variation across these parameters.Source: §4.3 · p. 5

4.4. Discriminator architecture

  • The discriminator is a U-Net that predicts labels from the chosen labeling function rather than manually annotated semantic classes.Its output path ends at the 32 × 32 latent resolution, and the authors found no convolutional normalization most effective among the configurations considered.Source: §4.4 · p. 5

4.5. Training

  • Training first pretrains the autoencoder without a discriminator, then fine-tunes its decoder with the full rate-distortion-perception loss and the discriminator.During fine-tuning, the discriminator learning rate is 4 × 10−4 and the generator rate is 1 × 10−4; the authors found the discriminator needed to train faster, particularly at higher bitrates.Source: §4.5 · p. 6

5. Experiments

Experiments compare distortion and statistical-fidelity metrics across datasets and baselines, finding that MS-ILLM improves statistical fidelity at matched distortion. Qualitative examples and ablations characterize image-level trade-offs and design sensitivity.

5.1. Datasets and metrics

  • Evaluation uses CLIC2020, DIV2K, and Kodak, with most main-body results on CLIC2020; metrics include reference-based MS-SSIM/PSNR and no-reference FID/KID.DIV2K and Kodak show trends similar to CLIC2020.Source: §5.1 · p. 6

5.2. Baseline models

  • Baselines include the pretrained No-GAN model, a PatchGAN-finetuned MS-PatchGAN (HiFiC*), handcrafted codecs, and published HiFiC and Multi-Realism.MS-PatchGAN uses the authors’ own HiFiC-like training recipe and is not an exact replica of the original HiFiC model.Source: §5.2 · pp. 6–7

5.3. Main results

  • At matched PSNR, MS-ILLM has better FID and KID than HiFiC across bitrates; at matched FID/KID, it achieves higher PSNR.The comparison with PatchGAN also finds better FID at all investigated MSE distortion points.Source: §5.3 · p. 7, Figure 3 · p. 7
  • On cityscape and pet examples, MS-ILLM sharpens or restores detail with fewer artifacts than HiFiC, while text quality remains degraded versus No-GAN.The text example reports slightly better legibility than HiFiC for “centrum.”Source: §5.3 · p. 8, Figure 4 · p. 8, §5.4 · p. 8

5.4. Ablations

  • No normalization was ideal without major observed training instability; spectral normalization took up to 25 times as many steps, while latent size mildly affected the FID–PSNR trade-off and codebook size had almost no effect.The latent and codebook-size comparison is for models at approximately 0.121 bpp.Source: §5.4 · p. 8, Table 1 · p. 8, Table 2 · p. 9

6. Limitations

The authors identify demographic-bias risk, substantial compute requirements, and uncertainty about whether improved FID/KID corresponds to human preference.

6. Limitations

  • Adversarial neural compression could produce demographic bias related to factors such as race or gender.The authors restrict their results to research use pending further assessment of these factors for production systems.Source: §6 · p. 9
  • The HiFiC autoencoder requires substantial compute, and real-world deployment would require model miniaturization and heavy weight quantization.Source: §6 · p. 9
  • Improved FID/KID does not guarantee better human preference than competing methods.Source: §6 · p. 9

7. Conclusion

The conclusion presents MS-ILLM as a neural image compressor that uses local adversarial discriminators to improve statistical fidelity. The authors report better FID and KID and say the method moves image compression closer to the theoretical optimum of perfect statistical fidelity.

7. Conclusion

  • MS-ILLM improves statistical fidelity using local adversarial discriminators, emphasizing locality for compression.Source: §7 · p. 9
  • The authors report that the local-discriminator approach yields better FID and KID metrics.Source: §7 · p. 9
  • The authors say their empirical findings concur with rate-distortion-perception theory and move image compression closer to perfect statistical fidelity.Source: §7 · p. 9

Improving Statistical Fidelity for Neural Image Compression with Implicit Local Likelihood Models — Supplementary Material

The supplementary material gives training details for the autoencoder, label function, and discriminator fine-tuning. It reports that perceptual loss was critical when training the VQ-VAE label function.

Improving Statistical Fidelity for Neural Image Compression with Implicit Local Likelihood Models — Supplementary Material

A.1. Autoencoder pretraining

  • Autoencoder pretraining used eight rate targets and adaptive rate targeting, with a linear learning-rate ramp for 10,000 steps followed by cosine decay.The authors observed a small PSNR boost on the CLIC 2020 test set with this schedule.Source: §A.1 · p. 12

A.2. Labeler pretraining

  • The VQ-VAE label-function loss combines distortion, embedding, and commitment losses; the authors found LPIPS critical to successful label-function training.The labeler uses a VGG backbone for the distortion term.Source: §A.2 · pp. 12–13

A.3. Fine-tuning with GAN loss

  • For discriminator fine-tuning, the encoder and autoencoder bottleneck were frozen and only the decoder was fine-tuned; ILLM used a higher learning rate than PatchGAN.The reported learning rates were 0.0001 for the autoencoder, 0.0001 for PatchGAN, and 0.0004 for ILLM.Source: §A.3 · p. 13

B. Calculation of baseline metrics

The evaluation pipeline compresses each dataset with each method and then evaluates all methods consistently. The section documents the codec-specific steps used for VTM, BPG, and HiFiC.

B. Calculation of baseline metrics

C. Note on metrics computation: bitrate, FID, and KID

The authors estimate bitrate from the saved bitstream and required image-size metadata, and standardize FID and KID calculations with torch-fidelity. They also recalculate baseline metrics because published FID/KID values could not be reproduced exactly.

C. Note on metrics computation: bitrate, FID, and KID

  • BPG and VTM file sizes slightly overestimate rate because their files include extra metadata, a small per-image overhead.Source: C. Note on metrics computation:… · p. 14
  • FID/KID calculations use torch-fidelity, whose Inception resizing and KID subset sampling differ from tensorflow-gan.torch-fidelity resizes inputs to 299 × 299 with TensorFlow 1.0-compatible resizing and bootstraps KID subsets with replacement.Source: C. Note on metrics computation:… · p. 14
  • HiFiC-Lo’s reported FID of 4.19736528 differed from the authors’ 4.18498420715332, so they recalculated baseline FID and KID with torch-fidelity for fair comparison.The recalculated plots use the more optimistic HiFiC FID of approximately 4.185.Source: C. Note on metrics computation:… · p. 14

D. Further analysis of the rate-distortion-perception tradeoff

Increasing the statistical-fidelity weight improves FID while worsening distortion, consistent with the rate-distortion-perception tradeoff. At matched PSNR, ILLM also yields better FID and fewer artifacts than PatchGAN in the qualitative comparison.

D. Further analysis of the rate-distortion-perception tradeoff

E. Additional qualitative examples

  • At matched PSNR, ILLM has better FID and fewer artifacts than HiFiC’s PatchGAN in the qualitative examples.Increasing discriminator weight lets ILLM add more sharpening and detail without the artifact reduction seen with PatchGAN.Source: E. Additional qualitative examples · p. 15

F. Experimental results on DIV2K and Kodak

On DIV2K, MS-ILLM matches HiFiC on reference-based metrics while outperforming it on FID and KID; Kodak supports comparisons of distortion but is too small for distributional metrics. SwAV-based FID results also retain MS-ILLM’s statistical-fidelity advantage, indicating gains beyond ImageNet class alignment.

F. Experimental results on DIV2K and Kodak

G. Investigation of ImageNet feature alignment

Loading 2301.11189v3…