Robustly overfitting latents for flexible neural image compression

Neural image compression has made a great deal of progress. State-of-the-art models are based on variational autoencoders and are outperforming classical models. Neural compression models learn to encode an image into a quantized latent representation that can be efficiently sent to the decoder, which decodes the quantized latent into a reconstructed image. While these models have proven successful in practice, they lead to sub-optimal results due to imperfect optimization and limitations in the encoder and decoder capacity. Recent work shows how to use stochastic Gumbel annealing (SGA) to refine the latents of pre-trained neural image compression models. We extend this idea by introducing SGA+, which contains three different methods that build upon SGA. We show how our method improves the overall compression performance in terms of the R-D trade-off, compared to its predecessors. Additionally, we show how refinement of the latents with our best-performing method improves the compression performance on both the Tecnick and CLIC dataset. Our method is deployed for a pre-trained hyperprior and for a more flexible model. Further, we give a detailed analysis of our proposed methods and show that they are less sensitive to hyperparameter choices. Finally, we show how each method can be extended to three- instead of two-class rounding.

Paper

Similar papers

Peer review

Reviewer Vgvs5/10 · confidence 4/52024-06-29

Summary

The authors build on previous autoencoder-based neural image compression methods. Earlier works have shown that the quantized latent representations output by an end-to-end trained encoder for any given image are suboptimal for a decoder. Concretely, for a fixed set of decoder weights, one can usually find latent representations that yield a better rate-distortion than the representations output by the encoder. Hence, previous works have proposed optimizing the rate-distortion loss of an image using gradient descent for a fixed set of decoder weights. Since the latent representations we wish to optimize are discrete, previous works have developed "soft-quantization" approaches based on the Gumbel-Softmax trick, which introduces a temperature parameter $t$, such that for $t > 0$ the optimization problem becomes continuous, while $t \to 0$ recovers hard quantization. One starts with a temperature $t \gg 0$, which is then annealed towards $0$ during the optimization. Specifically, the authors build on stochastic Gumbel annealing (SGA, Yang et al., 2020), which uses the function $f_\tau(x) = -\tanh^{-1}(x) / \tau$ to compute the soft quantization log-probabilities for a given temperature $\tau$. In the present paper, the authors propose three alternatives for $f_\tau$ and conduct experiments analogous to the ones in Yang et al. (2020). ## References Yang, Y., Bamler, R., & Mandt, S. (2020). Improving inference for neural image compression. Advances in Neural Information Processing Systems, 33, 573-584.

Strengths

Unfortunately, I couldn't find any particular strengths worth highlighting.

Weaknesses

Unless I have completely misunderstood the paper, the authors' motivation for the paper is incorrect. Concretely, for the function $\exp(-\tanh^{-1}(x))$ used by SGA, the authors claim on Line 155 that "The problem is that the gradients tend to infinity when the function approaches the limits of $0$ and $1$." While this is true for the limit point $1$, this does not hold for $0$. This is a crucial point, as the derivative tending to $\infty$ at $1$ is desirable, as it prevents rounding to the wrong value, while if this occurred $0$, it would cause optimization difficulties. However, it can be easily verified that the function I mentioned above doesn't have an infinite derivative at $0$ by plotting it. In fact, I am not sure what function the authors have plotted in Figure 1a under the name "atanh", as the graph of $\tanh^{-1}$ is not centrally symmetric. In fact, the training instability of SGA arises because the derivative at $0$ becomes infinite as the temperature $\tau$ is annealed to $0$. All of the authors' proposed alternatives have the same issue, meaning they do not tackle the very problem they set out to solve. The experiments are also not insightful and do not demonstrate that any of the proposed modifications lead to a significant improvement. Furthermore, the paper's writing should be significantly improved. The authors review many unnecessary, basic details (such as the first two paragraphs of the introduction or the societal impact section, which should be cut altogether). The notation should also be cleared up; for example, the authors introduce duplicate or equivalent notations for the same objects: $\pi$ in Eq (3), $p(y)$ in Eq (4) for the soft quantization probabilities, $v_L$ and $w$ for the fractional part of $v$, and $n$ for the inverse temperature $\tau$. Finally, the font sizes should be increased in each figure as they are currently unreadable without zooming in significantly.

Questions

n/a

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

n/a

Reviewer 5M8z7/10 · confidence 3/52024-07-12

Summary

This paper proposes a technique which takes a neural compression method based on a VAE and runs an optimization using the same loss function as the original method was trained with, but instead of back-propagating into the weights of the network, the gradients accumulate into the quantized latents. Due to the fact that the latents are quantized, the authors propose 3 alternatives on how to handle this non-trivial issue. The results presented show that the method is improving RD (as one would expect) over the original latents, but the computational cost per image is rather high.

Strengths

- The paper presents some relatively straight-forward approximations which should be easy to implement in practice - Strong results, and good ablations in the appendix (I actually had a bunch of comments I ended up deleting because the answers were in the appendix, so I think it's a very detailed paper from multiple respects)

Weaknesses

- I am rather confused by the motivation behind 3-class rounding. - I would have liked to see a discussion on whether it's necessary to apply this algorithm to the hyper latents or not (technically the hyper latents can be computed from the latents, but I am unclear whether they don't suffer from the same problem as the latents) - The fact that there are 6 distinct possible methods that actually get presented is a bit confusing (2-class + 3-class) * 3 variants. The appendix does go into ablating over all of these, which is great, but it's still a bit confusing. - The method is acknowledged to be very expensive, and the gains are modest (though the authors do mention that this is a generic method that can be applied to any base architecture).

Questions

- What made you choose the 3 probability functions you chose? Technically there could be many functions $p$ which would fit the requirements. - Given the ablations, I saw that the 3-class rounding offers a modest improvement over 2-class, while being quite a bit more confusing to understand. Are there any other advantages besides "better RD performance" of 3-class rounding (i.e., do you need fewer iterations before convergence, etc)?

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors have addressed the limitations of their method adequately. The only comment here would be that the energy impact of using such a method on a global scale would be quite large, while the reduction in bits for the amount of energy used to achieve this would be... questionable, unless one were to chase the absolute limits of compression.

Reviewer N6DE6/10 · confidence 4/52024-07-19

Summary

In this paper, the authors proposed the method to improve the compression performance of the pre-trained end to end neural image compression methods by fine-tuning the latents of each image at the test time with the rate-distortion loss. They proposed three class rounding method, named SGA+ which is an extension of stochastic Gumbel annealing (SGA). The authors demonstrated that their method on the two pre-trained models with two different datasets and showed that their method was able to achieve better R-D performance than the existing methods.

Strengths

1. The proposed SGA+ method has the potential to round the latent's to slightly far way quantization grid. This might be optimal with respect to the rate-distortion loss function. The modeling of the probability for three classes are well formulated. 2. modeling of the probabilities using three methods to overcome the limitations of the gradients in the corners of the $atanh$ function 2. Interpolation based function using sigmoid scaled logit to model the probabilities

Weaknesses

1. The authors described the existing SGA method uses $atanh$ function to model the probabilities and as a result the gradients to tend infinity at the borders, and they proposed different ways to counter this problem by different functions. The instability of the gradients occurs only when the discretization gap is very minimal which might be the case at the higher BPP points in the R-D curve, so in these regions proposed method should provide higher gain. It is not evident from the results whether the proposed are able to achieve this? 2. The authors did not provide the B-D rate gain of their method with respect to the existing methods to quantify the average compression gain. 3. Analysis of what is percentage of the latents that were assigned into each category of three classes is missing to see the advantage of the proposed method. 4. Results are missing with recent end to end neural image compression methods. 5. The authors missed few works in the related section [a] M. Balcilar et. al, "Latent-Shift: Gradient of Entropy Helps Neural Codecs", 2023 IEEE International Conference on Image Processing (ICIP).

Questions

1. Whether the gradients problem of the SGA at the corners cannot be solved with tuning the temperature parameter? 2. The results of STE method is not consistent with iterations at 500 and 2000. At 500 iterations, the STE is above the base model and at 2000 iterations is below the base model. what is the reason for this behavior? Is the baselines are tunned properly? It seems there is a divergence in the optimization. 3. The analysis of BD rate gain [b] is important to quantify the average gain of the proposed method. 4. what are the differences and similarities between the Trellis Coded Quantization (TCQ) and your proposed method? In TCQ, also you have the possibility to have rounding the more classes. There are few works which make TCQ differentiable in the deep learning framework [c] [b] G. Bjontegaard, "Calculation of average PSNR differences between RD-curves", VCEG-M33, Austin, TX, USA, April 2001. [c] Deep Learning-based Image Compression with Trellis Coded Quantization https://arxiv.org/pdf/2001.09417

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, limitations are addressed

Reviewer Vgvs2024-08-10

I thank the authors for their rebuttal. Their response helped me realise that I misunderstood not only their paper but also the original SGA paper (I was familiar with it but never implemented it myself). I never realised that the original SGA paper had such a glaring suboptimality. As I now understand, the `atanh` reparameterisation proposed by the authors of the SGA paper is superfluous, as the authors did not realise that $v - \lfloor v \rfloor$ is already in the correct range. Could the authors comment on their understanding of why the `atanh` reparameterisation could have been chosen? After the authors' rebuttal, I also re-read the paper and now believe that they make a valuable contribution. I have updated my score accordingly. My main concerns that I would like to see the authors address in the camera-ready version of the papers are: - To avoid colossal misunderstandings like mine above, I would like to ask the authors to update the manuscript and be crystal clear about what function produces the probabilities and what is being plotted in Fig 1a. - Improve the Figures by making them vector instead of raster graphics and improving their readability by increasing their tick and legend font sizes, and increasing the line widths.

Authorsrebuttal2024-08-12

Thank you reviewer Vgvs for your reply and for updating your score. To answer your question regarding why the $\text{atanh}$ reparametrisation could have been chosen. This may be due to the fact that in the unnormalized log space, the $\text{atanh}$ looks like an appropriate function that satisfies some useful, and non-trivial properties, as mentioned in their paper. For example, the function $-\text{atanh}(v)$ being strictly decreasing on $(0,1)$ guarantees that the closer some value $v$ gets to an integer, the higher the probability it gets rounded to that $v$. Thank you for pointing out the following points, we will adjust these to avoid misunderstanding in the future: - We will add a clear explanation on how the probabilities of the functions in Fig 1a) are created. - We will improve the readability of each figure accordingly. We hope this answers your question. Feel free to reach out for any further comments.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC