Research map / Image generation / Image restoration
Robustly overfitting latents for flexible neural image compression
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SGA+, three alternative logit/probability functions (linear, cosine, sigmoid scaled logit) replacing atanh in Stochastic Gumbel Annealing for test-time latent refinement of pre-trained neural image compression models, plus a three-class rounding extension.
Problem
Latent refinement of pre-trained neural image compression models with stochastic Gumbel annealing (SGA) suffers from a discretization gap, sensitivity to hyperparameters, and unstable gradients from the atanh-based rounding probabilities.
Approach
Replace the atanh-based two-class rounding probabilities in SGA with monotone probability functions (linear, cosine, and a sigmoid scaled logit that interpolates between them), and extend each to three-class rounding, improving R-D performance and robustness to temperature settings.
Key results
- SSL with a=2.3 gives the lowest true R-D loss of 0.6175 versus 0.6260 for atanh on the main model after 2000 steps (Kodak-based analysis).
- On Kodak, SSL vs atanh BD-Rate at 500 steps is -5.04% and at 2000 steps is -2.86% for the main model.
- For the mean-scale hyperprior, SSL achieves BD-Rate of -22.77% on Tecnick at 500 steps versus -20.20% for atanh vs base.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Refine latents with SSL (a=2.3) instead of leaving the pre-trained hyperprior unchanged | Pre-trained hyperprior · Kodak, Tecnick, and CLIC | BD-PSNR (dB): Kodak 0.82; Tecnick 0.95; CLIC 0.89 | Helped |
| Refine latents with SSL (a=2.3) instead of leaving the pre-trained hyperprior unchanged | Pre-trained hyperprior · Kodak, Tecnick, and CLIC | BD-Rate (%): Kodak -16.23; Tecnick -18.77; CLIC -20.11 | Helped |
| Replace atanh refinement with SSL (a=2.3) | Pre-trained hyperprior · Kodak, Tecnick, and CLIC | BD-PSNR (dB): Kodak 0.14; Tecnick 0.16; CLIC 0.11 | Helped |
| Refine latents with SSL (a=4/3) instead of leaving the mean-scale hyperprior unchanged | Mean-scale hyperprior · Kodak, Tecnick, and CLIC | BD-PSNR (dB): Kodak 0.91; Tecnick 1.50; CLIC 1.33 | Helped |
| Replace atanh refinement with SSL (a=4/3) | Mean-scale hyperprior · Kodak, Tecnick, and CLIC | BD-PSNR (dB): Kodak 0.04; Tecnick 0.04; CLIC 0.03 | Helped |
| Replace SSL (a=2.3) with linear probabilities (ablation) | SGA+ latent refinement · Kodak, tau_max=1.0 | True R-D loss: 0.6175 -> 0.6220 | Hurt |
| Replace SSL (a=2.3) with cosine probabilities (ablation) | SGA+ latent refinement · Kodak, tau_max=1.0 | True R-D loss: 0.6175 -> 0.6187 | Hurt |
| Use three-class rather than two-class rounding with extended linear (ablation) | SGA+ latent refinement · Kodak, 2000 steps | True R-D loss: 0.6220 -> 0.6175 | Helped |
| Use three-class rather than two-class rounding with extended SSL (ablation) | SGA+ latent refinement · Kodak, 2000 steps | True R-D loss: 0.6175 -> 0.6203 | Hurt |
| Set SSL interpolation parameter to a=0.01 instead of a=2.3 (ablation) | SGA+ latent refinement · Kodak | True R-D loss: 0.6175 -> 0.7323 | Hurt |
Surprising findings
- The atanh-based rounding probability diverges (gradients tend to infinity) at the corners, and atanh with a learning rate of 0.02 diverges while SSL still converges.
- STE has trouble converging and performs poorly even with smaller learning rates.
Limitations
- Each image requires 2000 refinement iterations, which is time-consuming at inference.
- Best SSL hyperparameter a varies per model and must be tuned.
- Three-class rounding adds extra hyperparameters (r, n) and gains are modest once trained long.
Experimental setup
Datasets: Kodak CLIC 2020 Mobile ImageNet 2012
Benchmarks: Kodak Tecnick CLIC
Metrics: PSNR bits per pixel true R-D loss BD-PSNR BD-Rate
Baselines: Straight-Through Estimator Uniform Noise Stochastic Gumbel Annealing (atanh) base model without refinement
Builds on: Stochastic Gumbel Annealing mean-scale hyperprior Gumbel-softmax CompressAI
Closest papers by meaning
- Improving Inference for Neural Image Compression (NeurIPS 2020 · similarity 0.92)
- Soft then Hard: Rethinking the Quantization in Neural Image Compression (ICML 2021 · similarity 0.91)
- Variational image compression with a scale hyperprior (ICLR 2018 · similarity 0.89)
- Fitted Neural Lossless Image Compression (CVPR 2025 · similarity 0.87)
- Progressive Compression with Universally Quantized Diffusion Models (ICLR 2025 · similarity 0.87)
- Conditional Latent Coding with Learnable Synthesized Reference for Deep Image Compression (AAAI 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.