Robustly overfitting latents for flexible neural image compression

Authors: yura  perugachi-diaz
arwin  gansekoele
sandjai  bhulai
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 20:13:06

Abstract: 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.

Main paper: 
 introduction: 
 Image compression allows efficient sending of an image between systems by reducing their size. There are two types of compression: lossless and lossy. Lossless image compression sends images perfectly without losing any quality and can thus be restored in their original format, such as the PNG format. Lossy compression, such as BPG Bellard (2014) , JPEG Wallace (1992) or JPEG2000 Skodras et al. (2001) , loses some quality of the compressed image. Lossy compression aims to preserve as much of the quality of the reconstructed image as possible, compared to its original format, while allowing a significantly larger reduction in required storage.Traditional methods Wallace (1992); Skodras et al. (2001) , especially lossless methods, can lead to limited compression ratios or degradation in quality. With the rise of deep learning, neural image compression is becoming a popular method Theis et al. (2017) ; Toderici et al. (2017) . In contrast with traditional methods, neural image compression methods have been shown to achieve higher compression ratios and less degradation in image quality Ballé et al. (2018) ; Minnen et al. (2018) ; Lee et al. (2019) . Additionally, neural compression techniques have shown improvements compared to traditional codecs for other data domains, such as video. Agustsson et al. (2020) ; Habibian et al. (2019) ; Lu et al. (2019) .In practice, neural lossy compression methods have proven to be successful and achieve state-of-theart performance Ballé et al. (2018) ; Cheng et al. (2020) ; Minnen et al. (2018) ; Lee et al. (2019) . These models are frequently based on variational autoencoders (VAEs) with an encoder-decoder structure Kingma and Welling (2013) . The models are trained to minimize the expected rate-distortion (R-D) cost: R + λD. Intuitively, one learns a mapping that encodes an image into a compressible latent representation. The latent representation is sent to a decoder and is decoded into a reconstructed image. The aim is to train the compression model in such way that it finds a latent representation that represents the best trade-off between the length of the bitstream for an image and the quality of the reconstructed image. Even though these models have proven to be successful in practice, they do have limited capacity when it comes to optimization and generalization. For example, the encoder's capacity is limited which makes the latent representation sub-optimal Cremer et al. (2018) . Recent work Campos et al. (2019) ; Guo et al. (2020) ; Yang et al. (2020) proposes procedures to refine the encoder or latents, which lead to better compression performance. Furthermore, in neural video compression, other work focuses on adapting the encoder Aytekin et al. (2018) ; Lu et al. (2020) or finetuning a full compression model after training to improve the video compression performance van Rozendaal et al. (2021) .The advantage of refining latents Campos et al. (2019) ; Yang et al. (2020) is that improved compression results per image are achieved while the model does not need to be modified. Instead, the latent representations for each individual image undergo a refining procedure. This results in a latent representation that obtains an improved bitstream and image quality over its original state from the pre-trained model. As mentioned in Yang et al. (2020) , the refining procedure for stochastic rounding with Stochastic Gradient Gumbel Annealing (SGA) considerably improves performance.In this paper, we introduce SGA+, an extension of SGA that further improves compression performance and is less sensitive to hyperparameter choices. The main contributions are: (i) showing how changing the probability space with more natural methods instead of SGA boosts the compression performance, (ii) proposing the sigmoid scaled logit (SSL), which can smoothly interpolate between the approximate atanh, linear, cosine and round, (iii) demonstrating a generalization to rounding to three classes, that contains the two classes as a special case, and (iv) showing that SGA+ not only outperforms SGA on a similar pre-trained mean-scale hyperprior model as in Yang et al. (2020) , but also achieves an even better performance for the pre-trained models of Cheng et al. (2020) .Further, we show how SSL outperforms baselines in an R-D plot on the Kodak dataset, in terms of peak signal-to-noise ratio (PSNR) versus the bits per pixel (BPP) and in terms of true loss curves. Additionally, we show how our method generalizes to the Tecnick and CLIC dataset, followed by qualitative results. We analyze the stability of all functions and show the effect of interpolation between different methods with SSL. Lastly, we analyze a refining procedure at compression time that allows moving along the R-D curve when refining the latents with another λ than a pre-trained model is trained on Gao et al. (2022) ; Xu et al. (2023) . The code can be retrieved from: https://github.com/yperugachidiaz/flexible_neural_image_compression. 
preliminaries and related work: 
 In lossy compression, the aim is to find a mapping of image x where the distortion of the reconstructed image x is as little as possible compared to the original one while using as little storage as possible. Therefore, training a lossy neural image compression model presents a trade-off between minimizing the length of the bitstream for an image and minimizing the distortion of the reconstructed image Ballé et al. (2017) ; Lee et al. (2019) ; Minnen et al. (2018) ; Theis et al. (2017) .Neural image compression models from Ballé et al. (2017) ; Cheng et al. (2020) ; Minnen et al. (2018) ; Theis et al. (2017) , also known as hyperpriors, accomplish this kind of mapping with latent variables. An image x is encoded onto a latent representation y = g a (x), where g a (•) is the encoder. Next, y is quantized Q(y) = ŷ into a discrete variable that is sent losslessly to the decoder. The reconstructed image is given by: x = g s (ŷ), where g s (•) represents the decoder. The rate-distortion objective that needs to be minimized for this specific problem is given by:EQUATIONwhere λ is a Lagrange multiplier determining the rate-distortion trade-off, R is the expected bitstream length to encode ŷ and D is the metric to measure the distortion of the reconstructed image x compared to the original one x. Specifically for the rate, p x is the (unknown) image distribution and p ŷ represents the entropy model that is learned over the data distribution p x . A frequently used distortion measure for d(x, x), is the mean squared error (MSE) or PSNR.In practice, the latent variable y often consists of multiple levels in neural compression. Namely, a smaller one named z, which is modeled with a relatively simple distribution p(z), and a larger variable, which is modeled by a distribution for which the parameters are predicted with a neural network using z, the distribution p(y|z). We typically combine these two variables into a single symbol y for brevity. Furthermore, a frequent method of quantizing Q(•) used to train hyperpriors consists of adding uniform noise to the latent variable.Neural image compression models have been trained over a huge set of images to find an optimal encoding. Yet, due to difficulties in optimization or due to constraints on the model capacity, model performance is sub-optimal. To overcome these issues, another type of optimizing compression performance is proposed in Campos et al. (2019) ; Yang et al. (2020) where they show how to find better compression results by utilizing pre-trained networks and keeping the encoder and decoder fixed but only adapting the latents. In these methods, a latent variable y is iteratively adapted using differentiable operations at test time. The aim is to find a more optimal discrete latent representation ŷ. Therefore, the following minimization problem needs to be solved for an image x:EQUATIONThis is a powerful method that can fit to a test image x directly without the need to further train an entire compression model. et al. (2019) proposes to optimize the latents by iteratively adding uniform noise and updating its latents. While this method proves to be effective, there is still a difference between the true ratedistortion loss ( L) for the method and its discrete representation ŷ. This difference is also known as the discretization gap. Therefore, Yang et al. (2020) propose the SGA method to optimize latents and show how it obtains a smaller discretization gap. SGA is a soft-to-hard quantization method that quantizes a continuous variable v into the discrete representation for which gradients can be computed. A variable v is quantized as follows. First, a vector v r = (⌊v⌋, ⌈v⌉) is created that stacks the floor and ceil of the variable, also indicating the rounding direction. Next, the variable v is centered between (0, 1) where for the flooring: v L = v -⌊v⌋ and ceiling: v R = ⌈v⌉ -v. With a temperature rate τ ∈ (0, 1), that is decreasing over time, this variable determines the soft-to-hardness where 1 indicates training with a fully continuous variable v and 0 indicates training while fully rounding variable v. To obtain unnormalized log probabilities (logits), the inverse hyperbolic tangent (atanh) function is used as follows:EQUATIONTo obtain probabilities a softmax is used over the logits, which gives the probability p(y) which is the chance of v being floored: p(y = ⌊v⌋), or ceiled: p(y = ⌈v⌉). This is approximated by the Gumbel-softmax distribution. Then, samples are drawn: y ∼ Gumbel-Softmax(logits, τ ) Jang et al.(2016) and are multiplied and summed with the vector v r to obtain the quantized representation:v = i (v r,i * y i ).As SGA aids the discretization gap, this method may not have optimal performance and may not be as robust to changes in its temperature rate τ .Besides SGA, Yang et al. (2020) propose deterministic annealing Agustsson et al. (2017) , which follows almost the same procedure as SGA, but instead of sampling stochastically from the Gumbel Softmax, this method uses a deterministic approach by computing probabilities with the Softmax from the logits. In practice, this method has been shown to suffer from unstable optimization behavior.While methods such as SGA aim to optimize the latent variables for neural image compression at inference time, other approaches have been explored in recent research. Guo et al. (2021) proposed a soft-then-hard strategy alongside a learned scaling factor for the uniform noise to achieve better 2020) proposed a method to incorporate trellis-coded quantization in neural codecs. While these approaches change the training process, our work differs in that we only consider the inference process. Balcilar et al. (2023) proposes latent shift, a method that can further optimize latents using the correlation between the gradient of the reconstruction error and the gradient of the entropy.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: See the Conclusion and Limitations section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methods: 
 As literature has shown, refining the latents of pre-trained compression models with SGA leads to improved compression performance Yang et al. (2020) . In this section, we extend SGA by introducing SGA+ containing three other methods for the computation of the unnormalized log probabilities (logits) to overcome issues from its predecessor. We show how these methods behave in probability space. Furthermore, we show how the methods can be extended to three-class rounding.Recall from SGA that a variable v is quantized to indicate the rounding direction to two classes and is centered between (0,1). Computation of the unnormalized log probabilities is obtained with atanh from Equation (3). Recall, that in general the probabilities are given by a softmax over the logits with a function of choice. As an example, for SGA the logits are computed with atanh. The corresponding probabilities for rounding down is then equal to:e atanh(v L )e atanh(v L ) +e atanh(v R ) . Then looking at the probability space from this function, see Figure 1a , the atanh function can lead to sub-optimal performance when used to determine rounding probabilities. The problem is that gradients tend to infinity when the function approaches the limits of 0 and 1, see Appendix A for the proof that gradients at 0 tend to ∞. This is not ideal, as these limits are usually achieved when the discretization gap is minimal. In addition, the gradients may become larger towards the end of optimization. Further analyzing the probability space, we find that there are a lot of possibilities in choosing probabilities for rounding to two classes. However, there are some constraints: the probabilities need to be monotonic functions, and the probabilities for rounding down (flooring) and up (ceiling) need to sum up to one. Therefore, we introduce SGA+ and propose three methods that satisfy the above constraints and can be used to overcome the sub-optimality that the atanh function suffers from. We opted for these three as they each have their own interesting characteristics. However, there are many other functions that are also valid and would behave similarly to these three.We will denote the probability that v is rounded down by: p(y = ⌊v⌋), (4) where y represents the random variable whose outcome can be either rounded down or up. The probability that v is rounded up is conversely: p(y = ⌈v⌉) = 1 -p(y = ⌊v⌋).Linear probabilities To prevent gradient saturation or vanishing gradients completely, the most natural case would be to model a probability that linearly increases or decreases and has a gradient of one everywhere. Therefore, we define the linear:p(y = ⌊v⌋) = 1 -(v -⌊v⌋).(5) It is easy to see that: p(y = ⌈v⌉) = v -⌊v⌋. In Figure 1a , the linear probability is shown.Cosine probabilities As can be seen in Figure 1a , the atanh tends to have gradients that go to infinity for v close to the corners. Subsequently, a method that has low gradients in that area is by modeling the cosine probability as follows:EQUATION)This method aids the compression performance compared to the atanh since there is less probability of overshooting the rounding value.Sigmoid scaled logit There are a lot of possibilities in choosing probabilities for two-class rounding. We introduced two probabilities that overcome sub-optimality issues from atanh: the linear probability from Equation ( 5), which has equal gradients everywhere, and cosine from Equation ( 6), that has little gradients at the corners. Besides these two functions, the optimal probability might follow a different function from the ones already mentioned. Therefore, we introduce the sigmoid scaled logit (SSL), which can interpolate between different probabilities with its hyperparameter a and is defined as follows: p(y = ⌊v⌋) = σ(-aσ -1 (v -⌊v⌋)), (7) where a is the factor determining the shape of the function. SSL is exactly the linear for a = 1. For a = 1.6 and a = 0.65 SSL roughly resembles the cosine and atanh. For a → ∞ the function tends to shape to (reversed) rounding.Note that the main reason behind the linear version is the fact that it is the only function with constant gradients which is also the most robust choice, the cosine version is approximately mirrored across the diagonal of the -atanh(x) which shows that it is more stable compared to the -atanh(x), and the reason behind the SSL is that it is a function that can interpolate between all possible functions and can be tuned to find the best possible performance when necessary.As described in the previous section, the values for v can either be floored or ceiled. However, there are cases where it may help to round to an integer further away. Therefore, we introduce three-class rounding and show three extensions that build on top of the linear probability Equation (5), cosine probability Equation ( 6), and SSL from Equation (7).The probability that v is rounded is denoted by: p(y = ⌊v⌉) ∝ f 3c (w|r, n), where w = v -⌊v⌉ is centered around zero. Further, we define the probability that v is rounded +1 and rounded -1 is respectively given by: p(y = ⌊v⌉ -1) ∝ f 3c (w -1|r, n) and p(y = ⌊v⌉ + 1) ∝ f 3c (w + 1|r, n). Recall, that v L , v R ∈ [0, 1], whereas w ∈ [-0.5, 0.5]. Defining w like this is more helpful for the 3-class since it has a center class. The general notation for the three-class functions is given by:EQUATIONwhere clip(•) clips the value at 0 and 1, r is the factor determining the height and steepness of the function and power n controls the peakedness of the function. Note that n can be fused with temperature τ together, to scale the function. This only accounts for the computation of the logits and not to modify the Gumbel temperature, therefore, τ still needs a separate definition. Extended linear Recall that the linear probability can now be extended to three-class rounding as follows:EQUATIONA special case is f 3c,linear (w|r = 1, n = 1), where the function is equivalent to the linear of the two-class rounding from Equation ( 5). For r < 1 this function rounds to three classes, and for n ̸ = 1 this function is not linear anymore.In Figure 1b , three-class rounding for the extension of Equation ( 5) can be found. As can be seen, solid lines denote the special case of two-class rounding with r = 1 and n = 1, dashed lines denote three-class rounding with r = 0.9 and n = 1 and dotted lines denote the two-class rounding with r = 1 and n = 3, which shows a less peaked function. For an example of two-versus three-class rounding, consider the case where we have variable v = -0.95. For two-class rounding there is only the chance of rounding to -1 with p(y = ⌊v⌉) (red solid line), a chance to round to 0 with p(y = ⌊v⌉ + 1) (green solid line) and zero chance to round to -2 with p(y = ⌊v⌉ -1) (yellow solid line). For three-class rounding, with r = 0.9 and n = 1, when v = -0.95 we find a high chance to round to -1 with p(y = ⌊v⌉) (red dashed line) and a small chance to round to 0 with p(y = ⌊v⌉ + 1) (green dashed line) and a tiny chance to round to -2 with p(y = ⌊v⌉ -1) (yellow dashed line).Extended cosine Similarly, we can transform the cosine probability from Equation ( 6) to three-class rounding:EQUATION)When f 3c,cosine (w|r = 1, n = 2), this function exactly resembles the cosine for two-class rounding, and for r < 1 this function rounds to three classes.Extended SSL Additionally, SSL from Equation ( 7) can be transformed to three-class rounding as follows:EQUATION)where a is the factor determining the shape of the function. When f 3c,SSL (w|r = 1, n = 1), this function exactly resembles the two-class rounding case, and for r < 1, the function rounds to three classes. Recall that this function is capable of exactly resembling the linear function and approximates the cosine from two-class rounding for a = 1 and a = 1.6, respectively.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer Justification: Yes we adhere to the code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
experiments: 
 In this section, we show the performance of our best-performing method in an R-D plot and compare it to the baselines. Further, we evaluate and compare the methods with the true R-D loss performance ( L), the difference between the method loss and true loss (L -L), and corresponding PSNR and BPP plot that respectively expresses the image quality, and cost over t training steps. Finally, we show how our best-performing method performs on the Tecnick and CLIC dataset and show qualitative results.Following Yang et al. (2020) , we run all experiments with temperature schedule τ (t) = min(exp{-ct}, τ max ), where c is the temperature rate determining how fast temperature τ is decreasing over time, t is the number of train steps for the refinement of the latents and τ max ∈ (0, 1) determines how soft the latents start the refining procedure. Additionally, we refine the latents for t = 2000 train iterations, unless specified otherwise. See Section 4.2 for the hyperparameter settings.We use two pre-trained hyperprior models to test SGA+, for both models we use the package from CompressAI Bégaint et al. (2020) . The first model is similar to the one trained in Yang et al. (2020) . Note, that the refining procedure needs the same storage and memory as theirs.In Appendix C, implementation details and results of this model can be found. All experiments in this section are run with a more recent hyperprior-based model which is based on the architecture of Cheng et al. (2020) . Additionally, this model reduces the R-D loss for atanh on average by 7.6% and for SSL by 8.6%, compared to the other model. The model weights can be retrieved from CompressAI Bégaint et al. (2020) . The models were trained with λ = {0.0016, 0.0032, 0.0075, 0.015, 0.03, 0.045}. The channel size is set to N = 128 for the models with λ = {0.0016, 0.0032, 0.0075}, refinement of the latents on Kodak with these models take approximately 21 minutes. For the remaining λ's channel size is set to N = 192 and the refining procedure takes approximately 35 minutes. We perform our experiments on a single NVIDIA A100 GPU.Baseline methods We compare our methods against the methods that already exist in the literature.The Straight-Through Estimator (STE) is a method to round up or down to the nearest integer with rounding bound set to a half. This rounding is noted as ⌊•⌉. The derivative of STE for the backward pass is equal to 1 Bengio et al. (2013) ; Van Den Oord et al. (2017) ; Yin et al. (2019) .The Uniform Noise quantization method adds uniform noise from u ∼ U (-1 2 , 1 2 ) to latent variable y. Thus: ŷ = y + u. In this manner ŷ becomes differentiable Ballé et al. (2017) . As discussed in Section 2.2, we compare against Stochastic Gumbel Annealing, which is a soft-to-hard quantization method that quantizes a continuous variable v into a discrete representation for which gradients can be computed.Figure 3a shows the rate-distortion curve, using image quality metric PSNR versus BPP, of the base model and for refinement of the latents on the Kodak dataset with method: STE, uniform noise, atanh and SSL. We clearly see how SSL outperforms all other methods. In Appendix B.1a, the results after t = 500 iterations are shown, where the performance is more pronounced.To show how each of the methods behave, we take a closer look at the performance plots shown in Figure 2 . The refinement results are obtained from the pre-trained model, trained with λ = 0.0075. The true R-D loss in Figure 2a shows that STE performs worse and has trouble converging, which is reflected in the R-D curve. Uniform noise quickly converges compared to atanh and SSL. We find that SSL outperforms all other methods, including atanh in terms of the lowest true R-D loss at all steps. Looking at the difference between the method loss and true loss Figure 2b , we find that both SSL and atanh converge to 0. Yet, the initial loss difference is smaller and smoother for SSL, compared to atanh. Additionally, uniform noise shows a big difference between the method and true loss, indicating that adding uniform noise overestimates its method loss compared to the true loss.Tecnick and CLIC To test how our method performs on other datasets, we use the Tecnick Asuni and Giachetti (2014) and CLIC dataset. We run baselines atanh and the base model and compare against SSL with a = 2.3. Figure 3b shows the corresponding R-D curves on the Tecnick dataset. Note that Appendix B.3 contains the results of the CLIC dataset which shows similar behavior as on Tecnick. We find that both refining methods improve the compression performance in terms of the R-D trade-off. Additionally, our proposed method outperforms atanh and shows how it improves performance significantly at all points. The R-D plot after t = 500 iterations for both datasets can be found in Appendix B.1. Note, that these results show even more pronounced performance difference.To evaluate the overall improvements of our method, we computed the BD-Rate and BD-PSNR Bjontegaard for Kodak, Tecnick and CLIC in Table 1 . We observe that across the board SSL achieves an improvement in both BD-PSNR and BD-Rate. After 500 steps, SSL achieves almost double the BD-Rate reduction as atanh. The difference between the two becomes smaller at 2000 steps. This underlines the faster convergence behavior of the SSL method.In Figure 4 , we demonstrate the visual effect of our approach. We compare the original image with the compressed image from base model with λ = 0.0016, refinement method atanh, and SSL. For  
analysis and societal impact: 
 In this section we perform additional experiments to get a better understanding of SGA+. An in-depth analysis shows the stability of each proposed method, followed by an experiment that expresses changes in the true R-D loss performance when one interpolates between functions. Further, we evaluate three-class rounding for each of our methods. Finally, we show how SGA+ for semi-multirate behavior improves the performance of its predecessor and we discuss the societal impact. Note, the results of the additional experiments are obtained from the model trained with λ = 0.0075. Temperature sensitivity Table 2 represents the stability of atanh and the SGA+ methods, expressed in true R-D loss, for different τ max settings for the temperature schedule. As can be seen, the most optimal setting is around τ max = 1 for each of the methods. In the table we find that the linear function is least sensitive to changes in τ max . To further examine the stability of the linear function compared to atanh, we subtract the best τ max , column-wise, from the linear and atanh of that column. Also taking into account the sensitivity results of the one of Yang et al. (2020) in Appendix C.2 we find in general, that overall performance varies little compared to the best τ max settings of the other methods and has reasonable performance.While SSL has the largest drop in performance when reducing τ max , it achieves the highest performance overall for higher values of τ max . If there is no budget to tune the hyperparameter of SGA+, the linear version is the most robust choice. Further, we evaluated the necessity of tuning both the latents and hyper-latents. When only optimizing the latents with the linear approach for τ max = 1, we found a loss of 0.6234. This is a difference of 0.0012, which implies that optimizing the hyper-latent aids the final loss. Three-class rounding In Table 3 , the true R-D loss for two versus three-class rounding can be found at iteration t = 2000 and in brackets t = 500 iterations. For each method, we performed a grid search over the hyperparameters r and n. As can be seen in the table, most impact is made with the extended version of the linear of SGA+, in terms of the difference between the two versus threeclass rounding at iteration t = 2000 with loss difference 0.0045 and t = 500 with 0.0159 difference. In Appendix B.3 a loss plot of the two-versus three-class rounding for the extended linear method can be found. Concluding, the three-class converges faster. For the extended cosine version, there is a smaller difference and for SSL we find that the three-class extension only boosts performance when run for t = 500. In Appendix C.2, the three-class experiments for the pre-trained model similar to Yang et al. (2020) can be found. We find similar behavior as for the model of Cheng et al. (2020) , whereas with the linear version most impact is made, followed by the cosine and lastly SSL. The phenomenon that three-class rounding only improves performance for t = 500 iterations for SSL may be due to the fact that SSL is already close to optimal. Additionally, we ran an extra experiment to asses what percentage of the latents are assigned to the 3-classes. This is run with the best settings for the linear version f 3c,linear (w|r = 0.98, n = 2.5). At the first iteration, the probability is distributed as follows: p(y = ⌊v⌉) = 0.9329, for p(y = ⌊v⌉ -1) = 0.0312, and p(y = ⌊v⌉ + 1) = 0.0359. This indicates that the class probabilities are approximately 3.12 for class -1 and 3.6 for class +1. This is a lot when taking into account that many samples are taken for a large dimensional latent. In conclusion, three-class rounding may be attractive under a constraint optimization budget, possibly because it is easier to jump between classes. Additionally, for the extended linear three-class rounding also results in faster convergence.Semi-multi-rate behavior An interesting observation is that one does not need to use the same λ during refinement of the latents, as used during training. This is also mentioned in Gao et al. (2022) for image and Xu et al. (2023) for video compression. As a consequence of this approach, we can optimize to a neighborhood of the R-D curve without the need to train a new model from scratch. We experimented and analyze the results with methods: atanh and SSL. Figure 3c shows the performance after t = 500 iterations for the model of Cheng et al. (2020) . We find that SSL is moving further along the R-D curve compared to atanh. Note the refinement does not span the entire curve and that the performance comes closer together for running the methods longer, see Appendix B.4.For future work it would be interesting how SGA+ compares to the methods mentioned in Gao et al. (2022) ; Xu et al. (2023) , since SSL outperforms atanh.The improvement of neural compression techniques is important in our data-driven society, as it allows quicker development of better codecs. Better codecs reduce storage and computing needs, thus lowering costs. However, training these codes requires significant computational resources, which harms the environment through power consumption and the need for raw materials. 
conclusion and limitations: 
 In this paper we proposed SGA+, a more effective extension for refinement of the latents, which aids the compression performance for pre-trained neural image compression models. We showed how SGA+ has improved properties over SGA and we introduced SSL that can approximately interpolate between all of the proposed methods. Further, we showed how our best-performing method SSL outperforms the baselines in terms of the R-D trade-off and how it also outperforms the baselines on the Tecnick and CLIC dataset. Exploration of SGA+ showed how it is more stable under varying conditions. Additionally, we gave a general notation and demonstrated how the extension to three-class rounding improves the convergence of the SGA+ methods. Lastly, we showed how SGA+ improves the semi-multi-rate behavior over SGA. In conclusion, especially when a limited computational budget is available, SGA+ offers the option to improve the compression performance without the need to re-train an entire network and can be used as a drop-in replacement for SGA.Besides being effective, SGA+ also comes with some limitations. Firstly, we run each method for 2000 iterations per image. In practice this is extremely long and time consuming. We find that running the methods for 500 iterations already has more impact on the performance and we would recommend doing this, especially when a limited computational budget is available. Future work may focus on reducing the number of iterations and maintaining improved performance. Note, higher values for a flatten out quickly, but they achieve much better gains with low-step budgets. Further, the best results are obtained while tuning the hyperparameter of SSL and for each of our tested models this lead to different settings. Note, that the experiments showed that the linear version of SGA+ is least sensitive to hyperparameter changes and we would recommend using this version when there is no room for tuning. Additionally, although three-class rounding improves the compression performance in general, it comes with the cost of fine-tuning extra hyperparameters. Finally, it applies for each method that as the temperature rate has reached a stable setting, the performance will be less pronounced, the longer you train, but in return requires extra computation time at inference. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: See the Analysis and Societal Impact section.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: Yes all creators and original owners of assets are cited and respected.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: [NA] Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In this appendix we will proof that the normalization from the (Gumbel) softmax causes infinite gradients at 0.Recall that the probability is given by a 2-class softmax is defined by: v) , where f (v) = -atanh(v) and g(v) = -atanh(1 -v). We will study the softmax for the first class 0, since the softmax is symmetric this also holds for the second class. The problem is that the gradients of the function K(v) will tend to ∞ for both v → 1 but also for v → 0. Here we show that the gradients also tend ∞ for v → 0, via the normalization with the term g(v). First, take the derivative to v:K(v) = e f (v) e f (v) + e g(dK(v) dv = dK(v) df (v) • df (v) dv + dK(v) dg(v) • dg(v) dv ,wheredK(v) df (v) = K(v)(1 -K(v)) and dK(v) dg(v) = -K(v) e g(v) e f (v) +e g(v) . Recall that datanh(v) dv = 1 1-v 2 therefore df (v) dv = -1 1-v 2 and dg(v) dv = 1 1-(1-v) 2 .Plugging this in and computing dK(v) dg(v) gives us:dK(v) dv = K(v)(1 -K(v)) - 1 1 -v 2 -K(v) • e g(v) e f (v) + e g(v) • 1 1 -(1 -v) 2 .Taking the limit to 0, (recall that lim v→0 K(v) = 1, lim v→0 e f (v) = 1 and lim v→0 e g(v) = 0) allows the following simplifications:lim v→0 0 • - 1 1 -0 2 -1 • e g(v) 1 + 0 • 1 1 -(1 -v) 2For simplicity we substitute q = 1 -v (when q → 1, then v → 0) which will result in the following:lim v→0 -e -atanh(1-v) • 1 1 -(1 -v) 2 = lim q→1 -e -atanh(q) • 1 1 -q 2 ,Recall, -atanh(q) = -1 2 ln 1+q 1-q , so e -atanh(q) = 1/ 1+q 1-q thus:-lim q→1 1 -q 1 + 1 • 1 1 -q 2 = -lim q→1 1 2 (1 -q) (1 -q 2 ) 2Since 1 2 is a constant and lim x→∞ √ x = ∞ the final step is to simplify and solve:-lim q→1 (1 -q) (1 -q 2 ) 2 = -lim q→1 -1 (q -1)(q + 1) 2 = -∞.This concludes the proof that the gradients tend to -∞ for v → 0.In this appendix, additional experimental results for refinement of the latents with the pre-trained models of Cheng et al. (2020) can be found.Difference across runs Although we do not report the standard deviation for every run, the difference across runs is very small. For the model of Cheng et al. (2020) with λ = 0.0032, running five refinement procedures for 2000 iterations, results in a mean of 0.3969 and a standard deviation of 2.41 • 10 -5 . We observe that SSL achieves best performance at t = 2000, compared to t = 500 iterations. However, running the method for 2000 iterations is in practice very long. Looking at the results at t = 500 iterations we see that there is already great improvement of performance for SSL, compared to the base model for both Tecnick and CLIC dataset. In Figure B.4a the true loss, difference in loss, BPP and PSNR curves can be found. As can be seen for a = 0.01, the function diverges and a = 0.3 does not seem to reach stable behavior, both resulting in large loss values. For a ≥ 1, the difference in losses start close to zero (see Figure C.7b). SSL with a = 5 results in the fastest convergence and quickly finds a stable point but ends at a higher loss than most methods.Besides an improved RD performance for the three-class rounding, we also found that three-class rounding leads to faster convergence. In Figure B .5 a true loss plot over the iterations for the linear method, can be found. As one can see, the three-class converges faster and especially on 500 iterations, boosts performance which makes it attractive under a constraint budget. Semi-Multi-Rate Performance (500 Steps)To make a clear comparison we trained a similar mean-scale hyperprior as in Yang et al. (2020) . Therefore, we use the architecture of Minnen et al. (2018) , except for the autoregressive part as a context model. Instead, we use the regular convolutional architecture of Ballé et al. (2018) . The model package for the mean-scale hyperprior is from CompressAI Bégaint et al. (2020) . The details and results of this model can be found in this section.Similar as for Cheng et al. (2020) we run all experiments with temperature schedule τ (t) = min(exp{-ct}, τ max ). Additionally, we refine the latents for t = 2000 train iterations, unless specified otherwise. Training runs took half a week for the 1M step model, around a week for the 2M step models, and around 1.5 weeks for the larger 3M step model. We ran all models and methods on a single NVIDIA A100 GPU. Further, the models for λ = {0.04, 0.08} are trained with 256 hidden channels and the model for λ = 0.001 is trained with 128 hidden channels.The remaining models are trained with hidden channels set to 192.We evaluate our best-performing method SSL on the Kodak and Tecnick datasets, by computing the R-D performance, average over each of the datasets. The R-D curves use image quality metric PSNR versus BPP on the Kodak and Tecnick dataset. Recall that as base model we use the pre-trained mean-scale hyperprior, trained with λ = {0.001, 0.0025, 0.005, 0.01, 0.02, 0.04, 0.08}. For SSL we choose a = 4 3 as we found that this setting achieves the best R-D loss overall at 2000 iterations. This is a lower setting compared to the model by Cheng et al. (2020) . The hyperparameters are similar to what we reported for Cheng et al. (2020) but with two main differences. We found that we could increase the learning rate by a factor 10 to 0.005 for atanh and SGA+. We also found that a lower a ∈ [1.3, 1.4] was optimal. .2, we find that most impact is made at t = 500 iterations. However, for the model similar to Yang et al. (2020) the performance lay closer to the performance of atanh, than for the model of Cheng et al. (2020) .Tecnick Figure C.3 shows the R-D curve when refining latents on the Tecnick dataset, after t = 500 and t = 2000 iterations. As can be seen in the plot, we find that the longer the methods run, the closer the performance lies to each other. The improvement by SSL compared to atanh is greater for Tecnick than for Kodak. .4 shows the R-D curve when refining latents on the CLIC dataset, after t = 500 and t = 2000 iterations. Similar to the previous results, we find that the longer the methods run, the closer the performance lies to each other and that running the method shorter already gives better performance, compared to the base model.In Table C .2, we computed the change in BD-PSNR and BD-rate for the mean-scale hyperprior model. We observe that SSL is slightly better than atanh, although the difference between them is smaller than reported on the Cheng et al. (2020) model. The gap between 500 steps and 2000 steps is also smaller compared to the results in 1.In this appendix, we analyze additional experiments for the model, similar to those in Yang et al. (2020) . The results for the analysis are obtained from a pre-trained model trained with λ = 0.01. Learning rates We run SSL and atanh with higher learning rate settings of 0.02 and 0.01 and compare it to the results obtained with learning rate 0.005. Figure C.6 shows the corresponding loss curves and Table C .3 shows the corresponding loss values at t = {500, 2000} iterations. We find that for a learning rate of 0.01 the gap between atanh versus SSL at 500 iterations is only around 14.3% smaller and it remains pronounced, while with a learning rate of 0.01 the gap at 2000 iterations is around 28.6% better. This concludes that SSL benefits more than atanh at 2000 iterations, with a higher learning rate. More interestingly, for a learning rate of 0.02, atanh diverges whereas SSL still reaches comparable performance. This highlights the sensitivity of atanh.Table C .4: True R-D loss for different τ max settings of: atanh(v), linear, cosine and SSL with a = 4 3 . Lowest R-D loss per column is marked with: ↓. Note that the function containing atanh is unnormalized. Temperature sensitivity Table C .4 represents the stability of atanh and the SGA+ methods, expressed in true R-D loss, for different τ max settings for the temperature schedule. As can be seen, the most optimal setting is with τ max = 1 for each of the SGA+ methods. atanh obtains equal loss for τ max ∈ [0.4, 0.5]. In general, we find that the linear method of SGA+ is least sensitive to changes in τ max and has equal loss between τ max ∈ [0.7, 1]. To further examine the stability of the linear function compared to atanh, we subtract the best τ max , column-wise, from the linear and atanh of that column. We now see that the linear function is not only least sensitive to changes in τ max , but overall varies little compared to the best τ max settings of the other methods. While the SSL has the largest drop in performance when reducing τ max , it achieves the highest performance overall for higher values of τ max .Interpolation Table C .1 presents the true R-D loss results for the interpolation with different a settings for SSL for the mean-scale hyperprior model. In Figure C.7, the corresponding overall performance of the methods can be found. As can be seen in Figure C .7a, for a = {0.01, 0.30}, the functions diverge, resulting in large loss values. For a = 0.65, we find that the loss curve is slightly unstable at the beginning of training, which can be seen in the bending of the curve, indicating non-optimal settings. This may be due to the fact that we run all methods with the same τ max = 1 for a fair comparison. Additionally, note that SSL with a = 0.65 obtains a true R-D loss of 0.7410 compared to 0.7418 for atanh with the same settings. This is due to the fact that SSL, especially in the tails of the probability, is slightly more straight-curved compared to the atanh when looking at its probability space.Remarkably, for a ≥ 1, the difference in losses start close to zero (see Figure C.7b) . SSL with a = 5 results in the fastest convergence and quickly finds a stable point but ends at a higher loss than most methods. Three-class rounding Table C .5 shows the true R-D loss for two-versus three-class rounding, at iteration t = 500 and in brackets t = 2000 iterations.For each method, we performed a grid search over the hyperparameters r and n. Additionally, for the extended SSL, we also performed a grid search over a and found the best setting to be a = 1.4. As can be seen in the table, most impact is made with the extended version of the linear of SGA+, in terms of the difference between the two versus three-class rounding at iteration t = 500 with loss difference 0.0035 and t = 2000 with 0.0006 difference. There is a small difference at t = 500 for the extended cosine version. In general, we find that running models longer results in convergence to similar values. SSL converges to equal values for two-and three-class rounding.Semi-multi-rate behavior Similar as in Appendix B.4, we experimented with different values for λ to obtain a semi-multi-rate curve. For every pre-trained model, we ran SSL and atanh using λ ∈ {0.001, 0.0025, 0.005, 0.01, 0.02, 0.04, 0.08}. In Figure C .5, we have plotted the R-D curve of the base model (lime green line) and its corresponding R-D curves, obtained when refining the latents with the proposed λ's for atanh and SSL. As can be seen, running the methods for t = 500 iterations, SSL obtains best performance. While the longer you train, the closer together the performance will be. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: This paper presents three improvements to stochastic Gumbel annealing, which is used to refine the discrete latents in neural compression methods. Reviewers appreciated the simplicity of the method, solid improvements and extensive ablations. Initially there was some misunderstanding on the part of some reviewers that was addressed in the discussion, and graciously acknowledged by the reviewers. Overall this paper represents a solid contribution to the neural compression literature.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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
limitations: limitations: Yes, limitations are addressed
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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)?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: n/a
limitations: limitations: n/a

Comments: comment: 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.
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.