Robustly overfitting latents for flexible neural image compression
Authors: Yura Perugachi-Diaz
Arwin Gansekoele
Sandjai Bhulai
Venue: NeurIPS 2024 poster (2024)
Keywords: neural image compression
latent optimization
Primary area: deep_learning_architectures
TL;DR: This paper introduces SGA+ a method that works better and is less sensitive to hyperparameter settings

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.

Robustly overfitting latents for flexible neural image compression
Yura Perugachi-Diaz Vrije Universiteit Amsterdam y m. perugachi diaz@vu. ml
Arwin Gansekoele Centrum Wiskunde & Informatica awgBcwi...
Sandjai Bhulai Vrije Universiteit Amsterdam s. bhulai@vu mn1
Abstract
Neural image compression has made a great deal of progress. State-of-the-art models are based on variational autoencoders and outperforming classical models. Neural compression models learn to encode an image into quantized latent representation that can be efficiently sent the decoder, which decodes the quantized latent into a reconstructed image. While these models have proven successful in practice, they lead sub- sub-optimal results due imperfect optimization and limitations in the encoder and decoder capacity. Recent work shows how to use stochastic Gumbel annealing (SGA) 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 and CLIC dataset. Our method is deployed for a pre-trained hyperprior and for a more flexible model. Further, 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.
both the Tecnick
give detailed analysis of proposed methods
Introduction
Image compression allows efficient sending of 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 thus be restored in their original format, such the PNG format. Lossy compression, such as BPG Bellard (2014), JPEG Wallace (1992) or JPEG2000 Skodras al. (00T), loses some quality of the compressed image. Lossy compression aims preserve as much of the quality of the reconstructed image as possible, compared to its original format, while allowing significantly larger reduction in required storage.
Traditional methods Wallace (1992); Skodras et al. (200T), especially lossless methods, can lead to limited compression ratios degradation quality. With the rise of deep learning, neural image compression is becoming 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. (022): Habibian et al. (2019): [uu al. (2019).
In practice, neural lossy compression methods have proven to be successful and achieve state-of-the- art performance [Blle al. (2018): Cheng al. (2020)) Minnen et al. (2018); (0008);Les Lee et (211)) These models are frequently based on variational autoencoders (VAEs) with an encoder-decoder structure
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Kingma and Welling 2013). The models are trained to minimize the expected rate- dist-riitotion (R-D) cost: R + AD. Inturtively, one learns a mapping that encodes an image into compressible latent representation. The latent representation is sent to decoder and is decoded into a reconstructed image. The aim is to train the compression model in such way that 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 al. (2019); Guo 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 encoder/Aytekin AAtekin al. (2018); ; Lu et al. (2020) or
finetuning a full compression model after training to improve the video compression performance van Rozendaal al. 202T
The advantage of refining latents Campos et al. (2019); Yang et al. (2020) is that improved com- pression 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 perfor- mance 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 compres- sion 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. (020). Further, we show how SSL outperforms baselines in an R-D plot on the Kodak dataset, terms of peak hignal-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 eurve when refining the latents with another À than pre-trained model is trained on Gao et al. (2022); Xu et al. (2023). The code can be retrieved neuralleeeera_l.mmmeee
from: http:g/www..../.ppprruahhaaaa//a
2 Preliminaries and related work
In lossy compression, the aim is to find mapping of image where the distortion of the reconstructed image à is as little as possible compared to the original one while using as little storage as possible. Therefore, training lossy neural image compression model presents 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 al. (2018) Theis al. 2017).
Neural image compression models from Ballé et al. 2017 ; Cheng et al. (200)); Minnen et al. (2018);
Theis et al. (2017). also known as hyperpriors, accomplish this kind of mapping with latent variables. An image œ' is encoded onto a latent representation y = ga (x). where (-) is the encoder. Next, y is quantized Q(v) ú into a discrete variable that is sent losslessly the decoder. The reconstructed image is given by: & = go (ū). where (-) represents the decoder. The rate-distortion objective that needs to be minimized for this specific problem is given by:
L = R+ AD
Egyps [- loga Po(0) + Eppp, [d(x,
(I)
distortion
where X a Lagrange multiplier determining the rate-distortion trade-off, is the expected bitstream length to encode ÿ and D is the metric to measure the distortion of the reconstructed image o

---

compared to the original one æ. Specifically for the rate, Pa is the (unknown) image distribution and P3 represents the entropy model that is learned over the data distribution Px. frequently used distortion measure for d(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 2, which is modeled with relatively simple distribution p(z), and a larger variable, which is modeled by a distribution for which the parameters are predicted with neural network using 2, the distribution p(y|z)) We typically combine these two variables into a single symbol y for brevity. Furthermore, frequent method of quantizing Q(-) used to train hyperpriors consists of adding uniform noise to the latent variable.
2.1 Latent optimization
Neural image compression models have been trained over huge set of images 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 test time. The aim to find more optimal discrete latent representation g. Therefore, the following minimization problem needs to be solved for an image .r:
arg ming [- log2 Pe (0) Ad(zz
This is a powerful method that can fit to a test image x directly without the need to further train an entire compression model.
2.2 Stochastic Gumbel Annealing
Campos 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 difference between the true rate- distortion loss (É) 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 smaller discretization gap. SGA is soft-to-hard quantization method that quantizes a continuous variable v into the discrete representation for which gradients can be computed. A variable is quantized as follows. First, vector v, ([v),, [w]]) 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: 2t - [+] and ceiling: UR = [w] v. With temperature rate E (0, 1), that is decreasing over time, this variable determines the soft-to-hardness where 1 indicates training with fully continuous variable v and indicates training while fully rounding variable v. To obtain unnormalized log probabilities (logits), the inverse hyperbolic tangent (atanh) function is used as follows:
logits = (-
- (tan://))
To 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 = [u]). ). or ceiled: p(y = [v]). This is approximated by the Gumbel-softtax distribution. Then, samples are drawn: y ~ Gumbel-SSftmax (logitss 7) 7) Jang et al. (0016) and are multiplied and summed with the vector v, to obtain the quantized representation: 0 Σ, (Ur * y:). 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 T.
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 deterministic approach by computing probabilities with the Softmax from the logits. In practice, this method has been shown to suffer from unstable optimization behavior.
2.3 Other methods
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. (202T) proposed soft-then-har strategy alongside learned scaling factor for the uniform noise to achieve better
3

---

Probability space
Probability space
1.0
1.0
0.8
0.8
probs
probs
0.4
0.4
0.2
0.2
0.0
0.2
0.4
0.6
0.8
1.0
1.5
-1.0
00.5
0.0
0.5
1.0
(a) Probability space for several SGA+ methods, along with atanh. Solid lines denote the probability of floor- ing [n] and dotted lines the probability of ceiling [v].
(b) Three-class rounding for the extended version of linear. Solid lines denote two-class rounding, dashed lines denote three-class rounding and dotted lines in- dicate the smoothness.
Figure 1: Probability space for (a) Two-class rounding (b) Three-class rounding
compression and smoother latent. These methods are used to fine-tune network parameters but the latents directly. [Zhu al.] (2022) proposed using Swin- tew---ranfffrr--eee based coding instead of ConvNet-based coding. They showed that these transforms can achieve better compression with fewer parameters and shorter decoding times. /va Rozendaal et al. (202T) proposed to also fine-tune the decoder alongside the latent for video compression. While accommodating the additional cost of saving the model update, they demonstrated a gain of 1d.. Zhang et al. (202I) and Dupont et al. (2021) proposed using implicit neural representations for video and image compression, respectively, He et al. (2022) proposed an improved context model (SCCTX) and a change to the main transform (ELIC) that achieve strong compression results together. El-Nouby et al. (2023) revisited vector quantization for neural image compression and demonstrated performs on par with hyperprior- based methods. Li al. (2020) proposed 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, method that can further optimize latents using the correlation between the gradient of the reconstruction error and the gradient of the entropy.
3 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, extend SGA by introducing SGA+ containing three other methods for the computation of the unnormalized log probabilities (logits) 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.
3.1 Two-class rounding
Recall from SGA that a variable v is quantized to indicate the rounding direction to two classes and centered between (0,1). Computation of the unnormalized log probabilities is obtained with atanh from Equation (3). Recall, that general the probabilities are given by a softmax over the logits with function of choice. As example, for SGA the logits are computed with atanh. The corresponding
stanbiug)
probabilities for rounding down is then equal to: gatanhtng ? 4 (tanb(o)) Then looking at the probability space from this function, see Figure Ia the atanh function can lead to sub-optimal performance when used determine rounding probabilities. The problem is that gradients tend to infinity when the function approaches the limits of ( and 1, see Appendix for the proof that gradients at 0 tend to Co. 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 is rounded down by:
p(y [0]),
(4)
where y represents the random variable whose outcome can be either rounded down or up. The probability that u is rounded up is conversely: p(y [v]] = 1 - ply [w])-
Linear probabilities To prevent gradient saturation 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 = [x]) = - [+])-
(5)
easy to see that: p(y = [v]) =
[e]. In Figure Ta the linear probability is shown.
Cosine probabilities As can be seen in Figure la the atanh tends to have gradients that go to infinity for close to the corners. Subsequently, method that has low gradients in that area is by modeling the cosine probability as follows:
f = [vJ) = cos² (v (U-Lu))T
(6)
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 round- ing. We introduced two probabilities that overcome sub-optimality issues from atanh: the linear probability from Equation ( 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 different function from the ones already mentioned. Therefore, we introduce the sigmoid scaled logit (SSL), which can interpolate between different probabilities with its hyperparameter and is defined as follows:
ply = [w]) = aa-aa ( (x))),
(7)
where a is the factor determining the shape of the function. SSI. is exactly the linear for a = 1. For 1.6 and 0.65 SSL roughly resembles the cosine and atanh. For oc 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(z). 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.
3.2 Three-class rounding
As described in the previous section, the values for can either be floored 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 is rounded is denoted by: p(u [c]] x ful( n). where v [v] is centered around zero. Further, we define the probability that v is rounded 1 and rounded - is respectively given by: ply v 1) [w 1/r, n) and p(y [r] + 1) fae(w 1 n). Recall, that CL, Un E [0, 1]. whereas w E [-0.5, 0.5]. Defining w like this is more helpful for the 3-class since has a center class. The general notation for the three-class functions given by:
fao n) = fclipp(w- ".
(8)
where 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 7 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.
5

---

0.72
0.25
0.24
10.25
0.88
2002
250
-000
1355
1000
Oplimaation Linddoo
@Oytemication Foaebaaro
+Dttmicalkn Raruiera
(a) True R-D loss
(b) Difference loss
(c) PSNR
(d) BPP
Figure 2: Performance plots of (a) True R-D Loss (b) Difference in loss (c) PSNR (d) BPP.
Extended linear follows:
Recall that the linear probability can now be extended to three-class rounding as
(w) = |w|-
(9)
special case is J3earn (wlr = 1, 1), where the function equivalent the linear of the J3.. two-class rounding from Equation (5). For < 1 this function rounds to three classes, and for n o 1
this function not linear anymore.
In Figure 1b three-class rounding for the extension of Equation 5 can found. As can be seen, solid lines denote the special case of two-class rounding with 1 and 1, dashed lines denote three-class rounding with r = 0.9 and n = and dotted lines denote the two-class rounding with 1 and 3. which shows a less peaked function. For example of two- versus three-class rounding, consider the case where we have variable v = - 00.95. For two-class rounding there is only the chance of rounding to -1 with ply [v]) (red solid line), chance round to 0 with p(y= [w] + +1 (green solid line) and zero chance to round to -2 with ply = [v] 1) (yellow solid line). For three-class rounding, with r 0.9 and
-0.95 we find a high chance to round to -1 with p(y = [v]) (red dashed line) and small chance to round to 0 with p(y [ + 1) (green dashed line) and a tiny chance round to -2 with p(y [o] 1) (yellow dashed line).
1, when
Extended cosine Similarly, can transform the cosine probability from Equation 6 to three-class rounding:
[w
I Scosine (w) cos
(10)
When (escosine(Wl 1, 2), this function exactly resembles the cosine for two-class rounding, and for r < 1 this function rounds to three classes
Extended SSL follows:
Additionally, SSL from Equation can be transformed to three-class rounding as
fss (-ao- ([w])).
(11)
where a is the factor determining the shape of the function. When fac,sst 1, = 1), this function exactly resembles the two-class rounding case, and for 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.
Experiments
this section, we show the performance our best-performing method an R-D plot and compare it to the baselines. Further, we evaluate and compare the methods with the true R-D loss performance (C), the difference between the method loss and true loss (C L), and corresponding PSNR and BPP plot that respectively expresses the image quality, and cost over training steps. Finally, we show how our berf-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(t) min(exp(-oo Tmaz), where e is the temperature rate determining how fast temperature de- creasing over time, L. is the number of train steps for the refinement of the latents and Tmaz E (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- -re-trained hyperprior models to test SGA+, for both models we use the package from Ce 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 [ implementation details and results of this model can be found. All experiments in this model which is based on the architecture of of(hheng ct 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 Compress Béégiit ct al. (202)) 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 A's channel size is set to N 192 and the refining procedure takes approximately 35 minutes. We perform our experiments on a single NVIDIA A A100 GPU.
Implementation details
section are run with more recent hyperprior-basee
Baseline methods The Straight- Through Estimator (STE) is a method to round up or down to the nearest integer with rounding bound set to half. This rounding noted as 1.1. The derivative of STE for the backward pass is equal to [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 5, 5) to latent variable
We compare our methods against the methods that already exist in the literature.
y. Thus: 1 y u. In this manner becomes differentiable [Ballé et al. (0177. As discussed in Section 2.2 we compare against Stochastic Gumbel Annealing, which is a soft-to-hard quantization method that quantizes continuous variable into discrete representation for which gradients can be computed.
4.1 Overall performance
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. the results after / = 500 iterations are shown, where the performance is more pronounced.
To show how each of the methods behave, we take closer look the performance plots shown in Figure The refinement results are obtained from the pre-trained model, trained with À = 0.0075. The true R-D loss in Figure [a 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 O. 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, use the Tecnick Asuni and Giachetti (2014) and CLIC dataset. We run baselines atanh and the base model and compare against SSL with - 2.3. Figure 3b shows the corresponding R-D curves on the Tecnick dataset. Note that Appendix 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 = 500 iterations for both datasets can be found in Appendix B.I 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- Bjontegaard for Kodak, Tecnick and CLIC in Table [T] 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 atanh. The difference between the two becomes smaller at 2000 steps. This underlines the faster convergence behavior of the SSL method.
BD-Rate and BD-PSNR
4.2 Qualitative results
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
7

---

Table Pairwise Comparison between atanh and SSL BD-PSNR and BD-Rate.
BD-PSNR (dB)
BD Rate (%i
500 steps Tecnick
2000 steps Tecnick
500 steps Tecnick
2000 steps Tecnick
Kodak
CLIC
Kodak
CLIC
Kodak
CLIC
Kodak
CLIC
Base vs
0.50
0.57
0.56
0.82
0.95
0.89
-10.30
-11.60
-13.18
-16.23
-18.77
-20.11
Base Atanh
0.26
0.28
0.31
0.69
0.79
0.78
-5.52
-5.11
-7.37
-3.82
-1.99
-17.70
Atanh vs SSL
0.24
0.28
0.26
0.14
0.16
0.11
-3.04
-5.97
-6.20
-2.86
-3.34
-2.76
the image compressed by SSL, we observe that there less artifacts visible in the overall image. For instance, looking at the window we see more texture compared to the base model and alanh method.
Hyperparamettr settings Refinement of the latents with pre-trained models, similar to the one trained in Yang al. 22220) use the same optimal learning rate of 0.005 for each method. Refinement of the latents with the models Chengg et al. (2200) use a 10 times lower learning rate of 0.0005. the settings for atanh with temperature rate Tmaz 0.5, and for STE we use the smaller learning rate of 0.0001, yet STE still has trouble converging. Note that, we tuned STE just as the other baselines. However, the STE method is the only method that has lot of trouble converging. Even with smaller learning rates, the method performed poorly. The instability of training not only observed by us, but also observed in Yang et al.] (2020); Yin et al. (2019). For SGA+, we use optimal convergence settings, which are a fixed learning rate of 0.0005, and Tman 1. Experimentally, find approximately best performance for SLL with a = 2.3.
Following Yang et al. 2020),
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 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-multi- rate 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.
Kodak
38 Sami-M..G-Raa Performance
34
ESHR 1000
12
Base
STE
Unirorm nolce
Base model
Mare
Atart
20.2
0.0
0.2
0.4
0.6
0.8
1.0
36.
0.2
0.5
0.4
0.5
0.6
26
0.2
0.4
0.5
0..5
1.0
Bits per pizei (BP9)
Bits per piaei (BPP)
Biba per (bop)
(a) Kodak
(b) Tecnick
(c) Semi-multi-ratc
Figure 3: R-D performance for SSL on (a) Kodak with the baselines, (b) Tecnick with the base model and atanh and (c) Kodak for cami-multi-rate behavior with atanh. Best viewed electronically.
(a) Original Image
(b) Base Model
(c) atanh
(d) SSL
(0.19 BPP vs 25.91 PSNR) (0.19 BPP 25.98 PSNR) Figure 4: Qualitative comparison of Kodak image from pre-trained model trained with À 0.0016. Best viewed electronically.
10.19 BPP 25.75 PSNR)

---

Temperature sensitivity Table [ represents the stability of atanh and the SGA+ methods, expressed in true R-D loss, for different Tmaz settings for the temperature schedule. As can be seen, the most optimal setting is around Tman = 1 for each of the meth- ods. In the table we find that the linear function is least sensitive to changes in Tmaz. To further examine the sta- bility of the linear function compared to atanh, we subtract the best Timax, 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 Tmaz settings of the other methods and has reasonable performance. While SSL has the largest drop in performance when reducing Tmaz, it achieves the highest per- formance overall for higher values of Tmaz- 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 Tmaz = 1, we found loss of 0.6234. This is a difference of 0.0012, which implies that optimizing the hyper-latent aids the final loss.
Table 2: True R-D loss for different Tmaa settings of: atanh(e), linear, cosine and SSL with a = 2.3. The lowest R-D loss per column is marked with: 4. Note that the func- tion containing atanh is unnormalized.
Function Vinaz
0.2
0.4
0.6
0.8
1.0
exp expalaanl())
0.6301
0.6273
0.6267
0.6260
0.6259
(linear)
0.6291 4
0.6229
0.6225
0.6222
0.6220
cos" (
0.6307
0.6233
0.6194 4
0.6186
0.6187
of (v))
0.6311
0.6233
0.6196
0.6181
0.6175 |
cxp atanh(e)
0,0010
0.0044
0.0073
0.0079
0.0084
(lincur)
0
0
0.0031
0.0011
0.0015
Interpolation Table represents the interpolation between different functions, expressed in true R-D loss. In Appendix B.2 the corresponding loss plots can be found. Values for < 1 indicate methods that tend to have larger gradients for v close to the corners, while high values of a represent method that tends to a (re- versed) step function. The smallest a = 0.01 diverges and results in large loss value compared to the rest. For a = 2.3 we find a lowest loss of 0.6175 and for = we find fastest convergence compared to the rest. Comparing these model results with the model results with the one of Yang et al. (2020), see Appendix C.2 we find that the Cheng et al. (2020) model obtains more stable curves.
Table 3: True R-D loss of two- versus three-class rounding for
SGA+ with the extended version of the linear, cosine, and SSL method at iteration 2000 and in brackets after 500 iterations.
Function /Roundingg
Two
Three
Sao,linerr (wii = 0.98, n 2.5)
0.6220 g0esp4
0. 0.6175
faccomiue (w/r 0.98, n 3)
0.6187 (0.0516)
0.6175 (0.0440)
facstgmadiingeteee 0.93, n 2.5)
0.6175 (0644
0.6203 (16500))
Three-class rounding Table 3 the true R-D loss for two versus three-class rounding can be found at iteration 1 = 2000 and in brackets = 500 iterations. For each method, we performed a grid search over the hyperparameters and 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 = 2000 with loss difference 0.0045 and t 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 == 500. In Appendix C.2 the three-class experiments for the pre-trained model similar Yang et al. (2020) can be found. We find similar behavior as for the model of Cheng et al. (2020).
Table 4: True R-D loss results for the interpolation between different functions by chang- ing of the SSL.
R-D Loss
0.01
0.7323
0.3
0.6352
0.65 (approx atanh)
0.6260
0.8
0.6241
1 (lincar)
0.6220
1.33
0.6199
1.6 (approx cosine)
0.6186
2.3
0.6175
5
0.6209
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 = 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 linea (w] r = 0.98, n 2.5). At the first iteration, the probability is distributed
as follows: p(v- = [0]) 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 lot when taking into account that many samples are taken for large dimensional latent. In

---

conclusion, three-class rounding may be attractive under a constraint optimization budget, possibly because 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 use the same À during refinement of the latents, as used during training. This is also mentioned in Gao et al. (2022) for image Xu et al. (202)) for video compression. As 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 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., For future work it would be interesting how SGA+ compares to the methods mentioned Gao et al. (2022)
Xu al. 2023), since SSL outperforms atanh.
Societal impact 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 materials.
6 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 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 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-ran behavior over SGA. In conclusion, especially when limited computational budget is available, SGA+ offers the option to improve the compression performance without the need to re- trairain entire network and can be used as a drop-in replacement for SGA.
Besides being effeettiv, 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 limited computational budget is available. Future work may focus on reducing the number 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 room for tuning. Additionally, although three-class rounding improves the compression performance in general, it comes with the cost fine-tuning extra hyperparameters. Finally, it applies for each method that the temperature rate has reached stable setting, the performance will be less pronounced, the longer you train, but in return requires extra computation time at inference.
10

---

Acknowledgments
This work was carried out on the Dutch national e-infrastructure with the support of SURF Coopera- tive.
References
Eirikur Agustsson, Fabian Mentzer, Michael Tschannen, Lukas Cavigelli, Radu Timofte, Luca Benini, and Luc V Gool. Soft-to-hard vector quantization for end-to-end learning compressible representations. Advances neural information processing systems, 30, 2017.
Eirikur Agustsson, David Minnen, Nick Johnston, Johannes Balle, Sung Jin Hwang, and George Toderici. Scale-space flow for end-to-end optimized video compression. In IEEE conference on Computer Vision and Pattern Recognition, 2020.
N. Asuni and A. Giachetti. TESTIMAGES: large-scale archive for testing visual devices and basic image processing algorithms (SAMPLING 1200 RGB set). In STAG: Smart Tools and Apps for Graphics, 2014.
OLD
Caglar Aytekin, Xingyang Ni, Francesco Cricri, Jani Lainema, Emre Aksu, and Miska Hannuk- sela. Block-optimized variable bit rate neural image compression. In Proceedings of the IEEE Conference on Computer Vision and Putlern Recognifion Workshops, pages 2551-2554, 2018.
Muhammet Balcilar, Bharath Bhushan Damodaran, Karam Naser, Franck Galpin, and Pierre Hellier. Latent-shift: Gradient of entropy helps neural codecs. In 2023 IEEE International Conference on Image Processing (ICIP), pages
IEEE, 2023.
Johannes Ballé, Valero Laparra, and Eero P Simoncelli. End-to-end optimized image compression. International Conference on Learning Representations, 2017.
Johannes Ballé, David Minnen, Saurabh Singh, Sung Jin Hwang, and Nick Johnston. Variational image compression with scale hyperprior. arXiv preprint arXiv:1802.01436. 2018.
Jean Bégaint, Fabien Racapé, Simon Feltman, and Akshay Pushparaja. Compressai: a pytorch library and evaluation platform for end-to-end compression research. arXiv preprint 1.X0v22011.03029, 2020.
Fabrice Bellard. BPG specification, 2014. URL https::/bbblardd.rrgg/////////// httg//bgbbpgssee.... txt (accessed June 3, 2020).
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv arXiv preprint / 2013.
G. Bjontegaard. Calculation of average psnr differences between rd-curves. VCEG-M33.
Joaquim Campos, Simon Meierhans, Abdelaziz Djelouah, and Christopher Schroers. Content adaptive optimization for neural image compression. arXiv preprint arxiv: 2019.
Zhengxue Cheng, Heming Sun, Masaru Takeuchi, and Jiro Katto. Learned image compression with discretized gaussian mixture likelihoods and attention modules. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
CLIC. CLIC: Challenge on learned image compression. URL htttp:://compressioon
Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoencoders.
In International Conference on Machine Learning, pages 1078-1086. PMLR, 2018.
Emilien Dupont, Adam Golinski, Milad Alizadeh, Yee Whye Teh, and Arnaud Doucet. COIN: compression with implicit neural representations. CoRR, abs/2103.03123, 2021.
Alaaeldin El-Nouby, Matthew J. Muckley, Karen Ulirich, Ivan Laptev, Jakob Verbeek, and Hervé Jégou. Image compression with product quantized masked image modeling. Trans. Mach. Learn. Res., 2023, 2023.

---

Chenjian Gao, Tongda Xu, Dailan He, Yan Wang, and Hongwei Qin. Flexible neural image com- pression via code editing. Advances in Neural Information Processing Systems, 15:12184- 12196, 2022.
Tiansheng Guo, Jing Wang, Ze Cui, Yihui Feng, Yunying Ge, and Bo Bai. Variable rate image compression with content adaptive optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pages 22-123, 2020.
Zongyu Guo, Zhizheng Zhang, Runsen Feng, and Zhibo Chen. Soft then hard: Rethinking the quantization in neural image compression. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 3920-392 3929. PMLR, 18-24 Jul 2021.
Amirhossein Habibian, Ties van Rozendaal, Jakub M Tomezak, and Taco S Cohen. Video compression with rate-distortion autoencoders. In IEEE International Conference on Computer Vision, 2019.
Dailan He, Ziming Yang, Weikun Peng, Rui Ma, Hongwei Qin, and Yan Wang. ELIC: efficient learned image compression with unevenly grouped space-channel contextual adaptive coding. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, IEEE, 2022. doi: 10.1 10.109//CVRR22682222.0000
LA, USA, June 18-24, 2022, pages
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv: 1.01 144, 2016.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. 14, 2013.
arXiv preprint
Eastman Kodak. Kodak lossless true color image suite (PhotoCD PCD0992). URL http://rOk us/ assgraphiis/kooall
Jooyoung Lee, Seunghyun Cho, and Seong-Kwon Beack. adaptxt- entropy model for end-to-end optimized image compression. In the 7th Int. Conf. on Learning Representations, May 2019.
Binglin Li, Mohammad Akbari, Jie Liang, and Yang Wang. Deep learning-based image compression with trellis coded quantization. In 2020 Data Compression Conference (DCC), pages 13-22. IEBE, 2020.
Guo Lu, Wanli Ouyang. Dong Xu, Xiaoyun Zhang, Chunlei Cai, and Zhiyong Gao. Dvc: An end-to-end deep video compression framework. In IEEE conference on Computer Vision and Pattern Recognition, 2019.
Guo Lu, Chunlei Cai, Xiaoyun Zhang, Li Chen, Wanli Ouyang, Dong Xu, and Zhiyong Gao. Content adaptive and error propagation aware deep video compression. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part II 16, pages 456-472. Springer, 2020.
David Minnen, Johannes Ballé, and George Toderici. Joint autoregressive and hierarchical priors for learned image compression. Advances in neural information processing systems, 31, 2018.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115 (3):21 1-252, 2015. doi: 1007// 1263-015-0816) 081-y.
A. Skodras, C. Christopoulos, and T. Ebrahimi. The jpeg 2000 still image compression standard. IEEE Signal Processing Magazine, 2001.
Lucas Theis, Wenzhe Shi, Andrew Cunningham, and Ferenc Huszár. Lossy image compression with compressive autoencoders. arXiv preprint arXiv: / 703.00395, 2017.
George Toderici, Damien Vincent, Nick Johnston, Sung Jin Hwang, David Minnen, Joel Shor, and Michele Covell. Full resolution image compression with recurrent neural networks. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition, pages 5306-5314, 2017.
12

---

George Toderici, Wenzhe Shi, Radu Timofte, Lucas Theis, Johannes Balle, Eirikur Agustsson, Nick Johnston, and Fabian Mentzer. Workshop and challenge on learned image compression (clic2020), 2020. htt.tttp::////tt
wttp://sww.oomm..e......ooo.ooo. cc.
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Ties van Rozendaal, Iris A. M. Huijben, and Taco Cohen. Overfitting for fun and profit: Instance- adaptive data compression. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Gregory K Wallace. The jpeg still picture compression standard. IEEE transactions on consumer electronics, xxvxxxxvii-xxx 1992.
Tongda Xu, Han Gao, Chenjian Gao, Yuanyuan Wang, Dailan He, Jinyong Pi, Jixiang Luo, Ziyu Zhu, Mao Ye, Hongwei Qin, et al. Bit allocation using optimization. In International Conference on Machine Learning, pages 38377-38399. PMLR, 2023.
Yibo Yang, Robert Bamler, and Stephan Mandt. Improving inference for neural image compression. Advances in Neural Information Processing Systems, 33:573-584, 2020.
Penghang Yin, Jiancheng Lyu, Shuai Zhang, Stanley Osher, Yingyong Qi, and Jack Xin. Under- standing straight-through estimator in training activation quantized neural nets. arXiv preprint 2019.
Yunfan Zhang, Ties van Rozendaal, Johann Brehmer, Markus Nagel, and Taco Cohen. Implicit neural video compression. CoRR, abs/2112.1 1312, 2021.
Yinhao Zhu, Yang Yang, and Taco Cohen. Transformer- based transform coding. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
13

---

A Proof
In this appendix we will proof that the normalization from the (Gumbel) softmax causes infinite gradients at O.
Recall that the probability is given by 2-class softmax is defined by:
ef(w)
( ef(0) + eg(D))
where S(v) atanh(u) and g(v) = 1 atanh(1 v). We will study the softmax for the first class O, since the softmax symmetric this also holds for the second class. The problem that the gradients of the function ((v) will tend to co for both gradients also tend oo for v + O, via the normalization with the term g(v). First, take the derivative u:
1 but also for - 0. Here show that the
dK(u) du
dK(v) df(v)
dK(v) dg(v)
(f(r) du dg(r) du
df()) where dk(v) K(v)(1 (()) and dK(u)
4/(v) 1-12 and dg(e) dg(1) ((v)
cafe)
Recall that datanb( dK(v) gives us: dg(")
+eat
T-
therefore
11-- Plugging this in and computing dK (c) K(b) K(o))
du
(- 1 02
K(v) ef(w) + calu) 1 (1 - 1)2
dv
Taking the limit to O, (recall that lim K(v) 1, lim ef(r) 1 and lim e0(v) 0) allows the following simplifications:
(- (-1-0²) 0² - 11. eg(0)
lim 0 -00
1 1 (1 v)
For simplicity we substitute = 11 (when g + 1, then v - 0) which will result in the following:
lim atanh((-e)) 0 -00
1
atanh(())
(1 - (1=
q
Recall, - ag)nh(g) = -- 1 In ltg I-g' so e ataah((g)
= =1/v
thus:
(1 4)
lim 9-
lim 9+
1
1 9²
2 2 - q2
Since a constant and limg V = oo the final step is to simplify and solve:
(1- 9) lim (1- q2)2
-1
lim g+1
--
V
(g 1)( 1)2
This concludes the proof that the gradients tend to -c0 for + 0.
14

---

B Additional results
In this appendix, additional experimental results for refinement of the latents with the pre- pre-trained models of Cheng al. (2020) can 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 et al. (2020) with À 0.0032, running five refinement procedures for 2000 iterations, results in mean of 0.3969 and a standard deviation of 2.41 10-5.
B.1 Additional overall performance
Figure 1 show the R-D curve for the Kodak dataset. We observe that atanh is similar to uniform noise at 500 iterations, while SSL manages to achieve better R-D gains. After 2000 iterations SLL achieves the best R-D trade-off, but the gain is a bit smaller compared to atanh at t 500 iterations.
Tecnick and CLIC Figure and and Figure B.3 | show the R-D results for respectively, Tecnick and CLIC at (500, 2000) iterations. We observe that SSL achieves best performance at 2000, compared to / 500 iterations. However, running the method for 2000 iterations is in practice very 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.
long. Looking at the results
Kodak (500 Steps)
Kodak (2000 Steps)
38
38
36
36
34
34
PSNR
PSNR
Base model
Base model
32
32
STE
STE
Uniform noise
Uniform noise
30
Atanh
30
Atanh
SSL
SSL
28.0
28.0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Bits per pixel (BPP)
Bits per pixel (BPP)
(a) Kodak 500)
(b) Kodak 2000)
Figure B.1: Comparison of atanh and SSL on the Kodak dataset for t (500, 2000) iterations.
B.2 Interpolation
In Figure 18.4athh the true loss, difference in loss, BPP and PSNR curves can be found. As can be seen for 0.01, the function diverges and 0.3 does not seem to reach stable behavior, both resulting in large loss values. For a 2 1, the difference in losses start close to zero (see Figure C.7b). SSL with results in the fastest convergence and quickly finds stable point but ends a higher loss than most methods.
B.3
Two- versus three- class rounding
Besides an improved RD performance the three-class rounding, we also found that three-class rounding leads to faster convergence. In Figure B.5 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 constraint budget.
15

---

Tecnick Steps)
Tecnick (2000 Steps)
40
40
38
38
36
36
PSNR
PSNR
Base model
Base model
Atanh
32
Atanh
SSL
SSL
30.
38.1
0.2
0.3
0.4
0.5
0.6
0.2
0.3
0.4
0.5
0.6
Bits pixel (BPP)
Bits per pixel (BPP)
(a) Teenick (t 500)
(b) Tecnick (f 2000)
Figure B.2: Comparison of atanh and SSL the Teenick dataset for t {500, 2000} iterations.
CLIC (500 Steps)
CLIC (2000 Steps)
40
38
36
36
PSNR
PSNR
34
34
Base model
Base model
32
Atanh
32
Atanh
SSL
38.1
38.1
0.2
0.3
0.4
0.5
0.6
0.2
0.3
0.4
0.5
0.6
Bits per pixel (BPP)
Bits per pixel (BPP)
(a) CLIC (f 500)
(b) CLIC 2000)
Figure B.3: Comparison of atanh and SSL on the CLIC dataset for = {500, 2000} iterations.
B.4 Semi-multi-rate behavior
In Figure B.6 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 X A's. For each model trained using A {0.0016, 0.0032, 0.0075, 0.015, 0.03, 0.045), we run atanh and SSL with a = 2.3 for t = 2000 iterations for all À E {0.0004, 0.0008, 0.0016, 0.0032, 0.0075, 0.015. 0.03, 0.045, 0.06, 0.09]. We depicted the base curve alongside the curves for each base model. We observe that the improved performance by SSL is especially noticeable 500 iterations in Figure B.6a
16

---

0.2
0.72
Difference R-D Loss (z-
True R D Loss (2)
0.68
0.1
0.64
0.0
0.60
-0.1
500
1000
1500
2000
0
500
1000
1500
2000
Optimization iterations
Optimization iterations
(a) True R-D Loss
(b) Loss difference: L
0.305
32.4
0.300
32.0
PSNA
B9P
0.295
31.6
31.2
0.285
500
1500
2000
500
1000
# Optimization iterations
# Optimiization Iterations
(c) PSNR
(d) BPP
Figure B.4: Interpolation performance plots of different settings for SLL (a) True R-D Loss (b) Difference in loss (c) PSNR (d) BPP.
0.72
True R-D Loss (2)
Lincar (2-class) Linear (3-class)
0.68
0.64
0.60
500
1000
1500
2000
# Optimization iterations
Figure B.5: True R-D loss curves for two- versus three-class rounding of the linear method.
Semi-Multi-RRte Performance (500 Steps)
Semi- Multi-Rate Performance (2000 Steps)
38
38
36
36
34
34
PSNR (dB)
made!
SUse madel
Atanh 0.0036
Atnn A 9.6036
SSL 0.0916
SSL 9.0916
Atat 0.0032
PSNR
32
Aran x 0.0032
504 X 0.033
54 x: 5.001
Abath 0.0375
Aran x: 3.0035
30
SSL J: 0.0315
30
551. 251.1: 3.0075
Atanh x: 0.015
Alan x: 0.005
134. h: 0.015
5S1. i: 3:015
28
SSL X: 0.03
Auanh x 0.03
28
Alans N: 2.03
354 3.03
Ataii 0.045
Atain A. 2.045
X.L 0.045
SSL 9.045
26
26
8.0
0.2
0.4 Bits per Pixel (bpp)
0.8
1.0
0.2
0.6
1.0
Bits per Pixel (bpp)
(a) f = 500 iterations
(b) = 200 iterations
Figure B.6: R-D performance on Kodak of the Cheng et (2020) model when varying the target Best viewed electronically.
17

---

C Mean-Scale Hyperprior
To make a clear comparison we trained a similar mean-scale hyperprior as in Yang et al. (0020)
Therefore, we use the architecture of Minnen et al. (2018), except for the autoregressive part context model. Instead, we use the regular convolutional architecture of Ballé et al. (018)). The
model package for the mean-scale hyperprior is from CompressAl Bégaint et al. (2220) The details
and results this model can be found in this section.
Similar as for Cheng et al. (2020) we run all experiments with temperature schedule =()) = min(exp(- (f Tmaz). Additionally, we refine the latents for t 2000 train iterations, unless specified otherwise.
Implementations details The pre-trained mean-scale hyperpriors are trained from scratch the full-size CLIC 2020 Mobile dataset Toderici et al. (2020). mixed with the ImageNet 2012 dataset Russakovsky et al. (2015) with randomly cropped image patches taken of size 256 x 256. For ImageNet, only images with a size larger than 256 for height and width are used to prevent bilin- ear up up-sampling that negatively affects the model per- formance. During training, each model is evaluated on the Kodak [ataset[Kodak The models were trained with X = {0.001, 0.0025, 0.005, 0.01, 0.02, 0.04, 0.08). with a batch size of 32 and Adam optimizer with a learning rate set to 1e-4. The models are trained for 2M steps, except for model À= = 0.001, which is trained for 1M steps and model X- = 0.08, which is trained for 3M steps. Training runs took half week for the IM 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
Table C.I: True R-D loss results for the interpolation between different functions by changing of the SSL..
R-D Loss
0.01
1.15
0.3
0.7528
0.65 (approx atanh)
0.7410
0.8
0.7396
(lincar)
0.7386
1.33
0.7380 4
1.6 (approx cosinc)
0.7382
2.25
0.7388
0.7415
single NVIDIA A100 GPU. Further, the models for À {0.04, 0.08) ) are trained with 256 hidden channels and the model for X = 0.001 is trained with 128 hidden channels. The remaining models are trained with hidden channels set 192.
C.1 R-D Performance
We evaluate our best-performing method SSL 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 scan-scale hyperprior, trained with À = {0.001, 0.0025, 0.005, 0.01, 0.02, 0.04, 0,083. For SSL we choose a = 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 reported for Cheng al. (2020) but with two main differences. We found that could increase the learning rate by a factor 10 to 0.005 for atanh and SGA+. We also found that lower E [1.3, 1.4] was optimal.
Kodak Figure C.2 shows the R-D curve for refining the latents, evaluated on Kodak. We compare SLL against baselines: STE, uniform noise, atanh and the base model at iteration
500 (see
Figure C.2a) and after full convergence at 2000 (see Figure C.2bj. As can be seen in Figure C.2a STE performs slightly better than the base model, while after 1 = 2000 iterations the method performs worse, this is also reflected in the corresponding true loss curve for X = 0.01 (see Figure C.1a). which diverges. Remarkably, for the smallest = 0.001, STE performs better than at t = 500. Adding uniform noise results in better performance when running the method longer. Comparing the R-D curves, Figure|..r to Figure[C.) [C22 we find that most impact is made at t = 500 iterations. However, for the model similar to Yang al. (2020) the performance lay closer to the performance of atanh, than for the model Cheng et al. 2020)).
Teenick Figure C.3 shows the R-D curve when refining latents on the Tecnick dataset, after / 500 and / = 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.
18

---

Difference Loss (z- -2
Uniform notse
Aaanh
0.05
SSL
0.000
0.00
0.775
0.05
1000
2000
# Optimization iterations
Optimization iterations
(a) True R-D loss
(b) Difference
33.5
0.42
PSNR
BPP
33.0
0.40
32.5
0.38
500
1000
1500
2000
500
1000
1500
2000
# Optimization iterations
Optimization iterationn
(c) PSNR
(d) BPP
Figure C.I: Performance plots of (a) True R-D Loss (b) Difference in loss (c) PSNR (d) BPP.
Kodak (500 Steps)
Kodak (2000 Steps)
PSNR
Base model
Base model
STE
STE
30
Uniform noise
Uniform noise
28
Atanh
28
Atanh
SSL
SSL
00.00
0.75
28.00
Bits per pixel (BPP)
Bits per pixel (BPP)
(a) R-D curve iteration 500
(b) curve iteration 2000
Figure C.2: R-D performance Kodak of the base mean-scale hyperprior model, STE, Uniform
noise, SGA atanh and SSL with a
CLIC Figure shows the R-D curve when refining latents on the CLIC dataset, after / = 500 and 2000 iterations. Similar to the previous results, 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.
BD-Rate Gain In Table C.2 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
C.2 Analysis
In this appendix, we analyze additional experiments for the model, similar to those Yang al. 2202)). The results for the analysis are obtained from pre-trained model trained with À 0.01.
19

---

Tecnick (500 Steps)
Tecnick (2000 Steps)
38
38
PSNR
Base model
Base model
30
30
Sigmoid logit
28.0
28.0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Bits per (BPP)
Bits per (BPP)
(a) R-D curve iteration / = 500
(b) R-D curve at iteration 2000
Figure C.3: R-D performance on Tecnick the base mean-scale hyperprior model, SGA atanh and SSL with
3- Best viewed electronically.
CLIC (500 Steps)
CLIC (2000 Steps)
36
36
34
34
Base model
Base model
Atanh
Atanh
SSL
SSL
0.0
28.0
Bits per pixel (BPP)
Bits per pixel (BPP)
(a) R-D curve iteration 500
(b) R-D curve iteration 2000
Figure C.4: R-D performance CLIC of the base mean-scale hyperprior model, SGA atanh and SSL with a 3. Best viewed electronically.
Semi Somi-Multi-ae Rate Performance (500 Steps)
Multi-Rate
Semi Multi Rate Performance (2000 Steps)
38
38
SOK rodd 0.833
SEA reOs
PSTRR (dB)
2514R (d9)
7aam N 0.8255
MAW 7. 3023
3.025
32
SSL.3 6075
1.05
Alam N.O.t:
Karhh 2.000
351
SoL a 602
52. 3.00
28
Aarm
28
K0ath 3.04
Alam N
3.00
551 000
2:.0
26
0.2
0.4
0.6
0.6
1.0
1.2
26
0.2
0.4
0.6
0.8
1.0
1.2
Bits per Pixel (bop)
Bits per Pixeel (bpo)
(a) R-D semi-multi-eeee curve at iteration
(b) R-D semi-multi-ran curve at iteration / 2000
500
Figure C.5: R-D performance on Kodak of the base mean-scale hyperprior model, SGA
atanh and SSL with 3. Each point is optimized with different target À (0.001, 0.0025, 0.005, 0.01, 0.02, 0.04 0.08).
20

---

Table C.2: Pairwise Comparison BD-PSNR and BD-Rate for the Kodak, Tecnick, and CLIC dataset on the Mean-Scar Hyperprior Model.
BD-PSNR (dB)
BD Rate (%)
500 steps Tecnick
2000 steps Tecnick
500 steps Tecnick
2000 steps Tecnick
Kodak
CLIC
Kodak
CLIC
Kodak
CLIC
Kodak
CLIC
Base SSL
0.68
1.21
1,03
0.91
1.50
1.33
-13.52
-22.77
-21.87
-17.69
-28.17
-2786
Base vs Atanh
0.59
1.06
0.89
0.87
1.46
1.30
-11.90
-20200
-99114
-17.03
22.49
27.28
Atanh SSL
0.09
0.15
0.14
0.04
0.04
-1.80
-3.22
-3.14
-0.82
-1.03
-0.74
Loss
0.850
ssl (Ir=0.02)
True R-D Loss (2)
0.825
ssi (Ir=0.01)
ssi (Ir=0.005)
0.800
atanh (Ir=0.02)
atanh (Ir=0.01)
0.775
atanh (Ir=0.005)
0.750
0.725
500
1000
1500
2000
# Optimization iterations
Figure C.6: True R-D loss curves for different learning rates settings for method SSL and atanh.
Table C.3: True R-D loss for different learning settings of: atanh and SSL with a = $. At 2000 iterations and in brackets f = 500 iterations
Learning rates We run SSL and atanh with higher learning rate set- tings of 0.02 and 0.01 and compare it to the results obtained with learning rate 0.005. Figure [.6]shows the cor- responding loss curves and Table [.. shows the corresponding loss values at t= ={500, 2000} iterations. We find that for learning rate of (1.11 the gap between atanh versus SSL at 500 iterations is only around 14.3% smaller and remains pronounced, while with a learning rate of 0.01 the gap at 2000 iterations around 28.6% better. This concludes that SSL benefits more than atanh at 2000 iterations, with higher learning rate. More interestingly, for a learning rate 0.02, atanh diverges whereas SSL still reaches comparable performance. This highlights the sensitivity of atanh.
Method
atauh
0.02
0.7386 1006
0.7491 co n6.m
0. 0.7375 @ 7478
0.7411 (0.734m)
0.005 (base)
0.7380 co 10.101
0.7408 co 70.757)
Temperature sensitivity Table C.44 represents the stability of atanh and the SGA+ methods, expressed in true R-D loss, for different Tmaz settings for the temperature schedule. As can be seen, the most optimal setting is with Tyaz = for each of the SGA+ methods. atanh obtains equal loss for Tyax E [0.4, 0.5). In general, we find that the linear method of SGA+ is least sensitive to changes in Tmaz and has equal loss between Fmaz [0.7, 1]. To further examine the stability of the linear function compared to atanh, we subtract the best Twaas column-wise, from the linear and alanh of that column. We now see that the linear function is not only least sensitive to changes in Tmar, but overall varies little compared to the best Tmaz settings of the other methods. While the SSL has the largest drop in performance when reducing Tmaz, it achieves the highest performance overall for higher values of Tmaz-
Table C.4: True R-D loss for different Tima settings of: atanh(o). linear, cosine and SSL with a $ Lowest R-D loss per column is marked with: J. Note that the function containing atanh is unnormalized.
Function Vimas
0.4
0.8
eXT atanh(r)
0.7445 J
0.7408
0.7412
0.7416
0.7418
(linear)
0.7458
0.7406
0.7390
0.7386
0.7386
co0"
0.7496
0.7414
0.7393
0.7387
0.7384
of 27-000 00 (o))
0.7578
0.7409
0.7391
0.7383
0.7380 L
exp atanlff)
0.0002
0.0022
0.0033
0.0035
(lineur)
0.0013
0.0003
0.0006
21

---

presents the true R-D loss results for the interpolation with different 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 (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 Tmea = 1 for fair comparison. Additionally, note that SSL with 0.65 obtains true R-D loss of 0.7410 compared to 0.7418 for atanh with the same settings. This 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.
Interpolation Table C.I
Remarkably, for 2 1, the difference in losses start close to zero (see Figure C.7b). b.. SSL with 5 results in the fastest convergence and quickly finds a stable point but ends at 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 = 500 and in brackets t 2000 iterations. For each method, we performed a grid search over the hyperparameters and Additionally, for the extended SSL, we also performed 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 three-class rounding at iteration = 500 with loss difference 0.0035 and = 2000 with 0.0006 difference. There is small difference t 500 for the extended cosine version. In general, find that running models longer results in convergence to similar values. SSL converges to equal values for two- and three-class rounding.
Table C.5: True R-D loss two versus three-class rounding for SGA+ with the extended version of the linear, cosine, and SSL method at iteration 500 and in brackets after 2000 iterations.
Function RRondding
Two
Three
Sactinoas (wir = 0.98, n 1.5)
0.7552 (0 7566)
0.751 7500
Jsc,ounina(1 (1rr 0.98, = 2)
0.7512 (0 (3.758))
0.7513 (0:75w
Janaiggooiieggt (wir 0.93, n = 1.5)
0.7524 ( 73800
0.7504 (i 7a
Semi-multi-rate behavior Similar as in Appendix B.4 we experimented with different values for À to obtain muni-multi--ate curve. For every pre-trained model, we SSL and atanh using we have plotted the R-D curve of
E {0.001, 0.0025, 0.005, 0.01, 0.02, 0.04, 0.08). In Figure C.5 the base model (lime green line) and its corresponding R-D curves, obtained when refining the latents with the proposed X'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.
22

---

0.850
0.2
Dffference R-D Loss -2
0.825
0.1
0.725
500
1000
15:0
2000
-0.1
0
500
1000
1500
2000
Optimization iterations
Optimization iterations
(a) True R-D Loss
(b) Loss difference: L -
33.6
0.44
0.43
PS/RR
041
1500
1500
Optimization Reeations
Optimization iterations
(c) PSNR
(d) BPP
Figure C.7: Interpolation performance plots different a settings for SLL (a) True R-D Loss (b) Difference in loss (c) PSNR (d) BPP under the mean-scale hyperprior model.
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
See the Methods and Experiments sections.
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
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 not discussed in the paper.
The authors are encouraged to create 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]

---

Justification: See Appendix A
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: See Methods and Experiments sections, where we provide the methods imple- mented, data required and packages used.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
25

---

Answer:
[Yes]
Justification:
The data and packages used are all open source. The code is publicly accessible at GitHub.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
http:/uuiddes/oodeSuuiisssss
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https
ssicnpolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See the Experiments section.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: The changes per run are very small with negligible changes in standard deviation, see Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
26

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
See Experiments section.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics /
Answer: [Yes]
Justification:
Yes we adhere to the code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. 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 socictal impact 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.
27

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses for existing assets
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?
Question:
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 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
28

---

If this information is not available online, the authors are encouraged to reach out to the asset's '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/moodl 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.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer:[NA] [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS 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.
29


Reviews:
[{"review_id": "U1QP4CKqyL", "reviewer": "Reviewer_N6DE", "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.", "questions": "1. Whether the gradients problem of the SGA at the corners cannot be solved with tuning the temperature parameter?\n\n2. 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.\n\n3.  The analysis of BD rate gain [b] is important to quantify the average gain of the proposed method. \n\n4. 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]\n\n[b] G. Bjontegaard, \"Calculation of average PSNR differences between RD-curves\", VCEG-M33, Austin, TX, USA, April 2001.\n[c] Deep Learning-based Image Compression with Trellis Coded Quantization   https://arxiv.org/pdf/2001.09417", "limitations": "Yes, limitations are addressed", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. modeling of the probabilities using three methods to overcome the limitations of the gradients in the corners of the $atanh$ function\n2. 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?  \n\n2. 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.\n3. 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.  \n4. Results are missing with recent end to end neural image compression methods.\n\n5. The authors missed few works in the related section\n [a] M. Balcilar et. al, \"Latent-Shift: Gradient of Entropy Helps Neural Codecs\", 2023 IEEE International Conference on Image Processing (ICIP).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you reviewer N6DE for your comments and questions.\n\n_To answer the weaknesses:_\n\nRegarding 1, whether the method might work better for higher BPP points in the R-D curve: For certain experiments this seems to be the case (e.g. on Kodak), for other experiments not. Due to the complexity of optimization, it is very difficult to predict which type of model suffers more from these spikes. E.g. it could also be the case that a lower BPP model would be more sensitive to gradient issues.\n\nRegarding 2: Thank you for pointing out that we did not provide the B-D rate gain of our method with respect to the existing methods. We will include these results in the attached PDF file. \n\nRegarding 3: The reviewer suggests an analysis of what percentage of the latents were assigned to the 3-classes. Therefore, we ran an extra experiment for the best settings of the 3-class extended version of the linear with $r=0.98$ and $n=2.5$. At the first iteration, the probability is distributed as follows: $p(y=\\lfloor v \\rceil) = 0.9329$, for $p(y=\\lfloor v \\rceil -1)= 0.0312$, and  $p(y=\\lfloor v \\rceil +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.\n\nRegarding 4: Missing results with recent end-to-end neural image compression methods. We choose two different models to show the effect of latent optimization. The first model is trained from scratch (see Appendix), which is similar to the one trained in [1] to make a fair comparison. The other model is a pre-trained model for which we also showed similar improvements. We acknowledge that it would be interesting to see the method on other models as well. \n\nRegarding 5: Missing a few works in the related section. Thank you for pointing this out, we will add [2] in the paper.\n\n_To answer the questions:_\n\nRegarding 1: the gradients problem at the corners, by tuning the temperature parameter. This is difficult. The problem is that when the temperature increases (which traditionally stabilizes) the diverging gradient of the atanh logits becomes worse. Lowering the temperature mitigates the gradient issue but makes the Gumbel sampling itself more unstable. Experimentally, we demonstrate this in Table 1, where atanh logits over a wide range of temperatures lead to worse performance. \n\nRegarding 2: the reason for the behavior of the STE method. We tuned the STE method, just as the other baselines. However, the STE method is the only method that has a lot of trouble converging. Even with smaller learning rates, the method performed poorly. \nThe instability of training is not only observed by us, but is something that is also mentioned in [1] and [4]. In [1] they have tried to overcome this following [4], by changing the gradient of the backward pass to (clipped)ReLU instead of using the identity. However, this did not work. We will highlight this clarification in the paper. \n\nRegarding 3: The missing B-D rate gain, we will add this in the paper and included the results in the attached PDF file, as mentioned above.\n\nRegarding 4: Trellis-coded quantization (TCQ) is related to our approach in that regard. They also consider a softmax over different quantization levels, although they consider all possible quantization levels. However, while we maintain a scalar quantization approach, TCQ is a method to more efficiently perform vector quantization (VQ). Their rounding approach is also not stochastic, but deterministic based on the optimal Trellis path. We will add a discussion in the paper concerning [3].\n\nWe hope this clarifies your questions. Feel free to reach out for further comments.\n\n[1] Yang, Y., et al., \"Improving inference for neural image compression\", 2020 Advances in Neural Information Processing Systems (NeurIPS)\n\n[2] M. Balcilar et. al, \"Latent-Shift: Gradient of Entropy Helps Neural Codecs\", 2023 IEEE International Conference on Image Processing (ICIP)\n\n[3] Deep Learning-based Image Compression with Trellis Coded Quantization [https://arxiv.org/pdf/2001.09417](https://arxiv.org/pdf/2001.09417)\n\n[4] Yin, P., et al., \"Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets\", 2019 International Conference on Learning Representations (ICLR)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Fo1Z2MkI9d", "reviewer": "Reviewer_5M8z", "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.", "questions": "- What made you choose the 3 probability functions you chose? Technically there could be many functions $p$ which would fit the requirements.\n- 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": "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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper presents some relatively straight-forward approximations which should be easy to implement in practice\n- 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.\n- 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)\n- 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.\n- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you reviewer 5M8z for your comments and questions. The reviewer thinks the method is straightforward and easy to implement. \n\n_To answer the weaknesses:_\n\nConcerning point 1: the motivation behind the 3-class rounding. Besides our scientific curiosity, we did find that 3-class rounding improved performance. When there is a limited budget, it may be worth it to use the 3-class, since at earlier iterations, this method converges a bit faster than its 2-class counterparts. We will add a loss plot in the paper that shows the difference between the 2- versus 3-class rounding. The added complexity may not be worth it, which is why we see the 3-class as an extended version of the main functions. Therefore, we always recommend the 2-class linear as a starting point that does not require much tuning, we will highlight this in the paper. We will include a loss plot over the iterations in the paper that shows the difference between the 2- versus 3-class rounding.\n\nRegarding point 2, whether it is necessary to apply this algorithm to the hyper latents: We performed a run using the linear 2-class approach where we optimized only the latents (and not the hyper latents). We found a loss of 0.6233 when optimizing only the latents, and a loss of 0.6220 when optimizing both latents and hyper latents. We will add a discussion of this in the paper.\n\nRegarding point 3: for the six distinct methods, we wanted to be precise and show the possible options and their possible equivalences (such as SSL being able to interpolate between the other functions, and 3-class having settings that match the 2-class). Nevertheless, we understand this can be overwhelming and we hope that the recommendation from the point above (for the 2-class linear as a starting point) helps readers to make a quick and robust choice.\n\n_To answer the questions:_\n\nRegarding point 1, the main reason behind the three probability functions: We tried various methods while looking into the probability space. 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 $-\\text{atanh}(x)$ which shows that it is more stable compared to the $-\\text{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. We will add this in the paper.\n\nRegarding point 2: additional advantages of 3-class rounding, see our answer for weaknesses 1. We will add a loss plot over iterations, where the 3-class approach indeed converges faster.\n\nWe hope this clarifies your points. Feel free to reach out for any further comments.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lEqkPe149F", "reviewer": "Reviewer_Vgvs", "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.\n\nSpecifically, 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$.\n\nIn the present paper, the authors propose three alternatives for $f_\\tau$ and conduct experiments analogous to the ones in Yang et al. (2020).\n\n## References\n\nYang, Y., Bamler, R., & Mandt, S. (2020). Improving inference for neural image compression. Advances in Neural Information Processing Systems, 33, 573-584.", "questions": "n/a", "limitations": "n/a", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n\nIn 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.\n\nThe experiments are also not insightful and do not demonstrate that any of the proposed modifications lead to a significant improvement.\n\nFurthermore, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you reviewer Vgvs for your comments. \n\nWe think there is a misunderstanding, the reviewer believes that the probability function used in [1] only contains infinite gradients at 1. However, the reviewer may have overlooked the **normalization** from the Gumbel softmax that causes infinite gradients at 0. See the proof below, this will be included in the appendix.\n\nRecall that the probability given by a 2-class softmax is defined as:\n$$\nK(v) = \\frac{e^{f(v)}}{e^{f(v)} + e^{g(v)}},\n$$\nwhere $f(v) = -\\text{atanh}(v)$ and $g(v)= -\\text{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 $-\\infty$ for both $v \\to 1$ (as agreed by the reviewer) but also for $v \\to 0$. Here we show that the gradients also diverge for $v \\to 0$, via the normalization with the term $g(v)$. First, take the derivative to $v$:\n$$\n\\frac{d K(v)}{dv} = \\frac{dK(v)}{df(v)} \\cdot\\frac{df(v)}{dv} + \\frac{dK(v)}{dg(v)} \\cdot\\frac{dg(v)}{dv},\n$$\nwhere $\\frac{dK(v)}{df(v)} = K(v)(1-K(v))$ and $\\frac{dK(v)}{dg(v)} = -K(v)\\frac{e^{g(v)}}{e^{f(v)} + e^{g(v)}}$. Recall that $\\frac{d\\text{atanh}(v)}{dv} = \\frac{1}{1 - v^2}$ therefore $\\frac{df(v)}{dv} = - \\frac{1}{1-v^2}$ and $\\frac{dg(v)}{dv} = \\frac{1}{1-(1-v)^2 }$. \nPlugging this in and computing $\\frac{dK(v)}{dg(v)}$ gives us:\n$$\n\\frac{dK(v)}{dv}= K(v)(1-K(v))\\left(- \\frac{1}{1-v^2}\\right) -K(v) \\cdot \\frac{e^{g(v)}}{e^{f(v)} + e^{g(v)}} \\cdot \\frac{1}{1-(1-v)^2 }.\n$$\nTaking the limit to $0$, (recall that $\\lim_{v \\to 0} K(v) = 1, \\lim_{v \\to 0} e^{f(v)} = 1$ and $\\lim_{v \\to 0} e^{g(v)} = 0$) allows the following simplifications:\n$$\n\\lim_{v \\to 0} 0 \\cdot \\left(-\\frac{1}{1-0^2} \\right) - 1 \\cdot \\frac{e^{g(v)}}{1 + 0} \\cdot \\frac{1}{1-(1-v)^2 }\n$$\nFor simplicity we substitute $q=1-v$ (when $q \\to 1$, then $v \\to 0$) which will result in the following:\n$$\n\\lim_{v \\to 0} -e^{-\\text{atanh}(1 - v)} \\cdot \\frac{1}{1-(1 - v)^2} = \\lim_{q \\to 1} -e^{-\\text{atanh}(q)} \\cdot \\frac{1}{1-q^2 },\n$$\nRecall, $-\\text{atanh}(q) = -\\frac{1}{2}\\ln{\\frac{1+q}{1-q}}$, so $e^{-\\text{atanh}(q)} = 1 / \\sqrt{\\frac{1 + q}{1 - q}}$ thus:\n$$\n-\\lim_{q \\to 1} \\sqrt{\\frac{1 - q}{1 + 1}} \\cdot \\frac{1}{1-q^2 } = -\\lim_{q \\to 1} \\sqrt{\\frac{1}{2}\\frac{(1 - q)}{(1-q^2)^2} }\n$$\nSince $\\frac{1}{2}$ is a constant and $\\lim_{x \\to \\infty} \\sqrt{x} = \\infty$ the final step is to simplify and solve:\n$$\n-\\lim_{q \\to 1} \\sqrt{\\frac{(1 - q)}{(1-q^2)^2}} = -\\lim_{q \\to 1}\\sqrt{\\frac{-1}{(q-1)(q+1)^2}} = -\\infty.\n$$\nThis concludes the proof that the gradients tend to $-\\infty$ for $v \\to 0$.\n\nFurther, the reviewer mentions that \"I am not sure what function the authors have plotted in Figure 1a under the name 'atanh' \". Again, this is after normalization from the softmax. \n\nRegarding the temperature of 0 in the Gumbel, this is a separate issue and the temperature is therefore annealed to values above 0, both in [1] and our paper.\n\nRegarding the experiments, if the reviewer has more suggestions, we are happy to include them. The experiments do show issues with the sensitivity of [1].\n\nRegarding notations: as we mention in the paper, directly above Eq. (3):  $\\log \\pi$ represents unnormalized log probabilities. In contrast $p(y)$ is normalized. We will define Eq. (3) with normalization and with $p(y)$ for clarity. Regarding $v_L, v_R, w$, these are different $v_L, v_R \\in [0, 1]$ whereas $w \\in [-0.5, 0.5]$. $w$ is more helpful to define the 3-class because it has a center class. It is fair to point out that $n$ and $\n\\tau$ can be fused into a new temperature in this function and we will point this out. However, following [1] ([https://github.com/mandt-lab/improving-inference-for-neural-image-compression/blob/6a97aba5b17c70847465f5865bd9e2bf58ccbe73/sga.py](https://github.com/mandt-lab/improving-inference-for-neural-image-compression/blob/6a97aba5b17c70847465f5865bd9e2bf58ccbe73/sga.py) #L118) the temperature is used twice, once on the logits and once again in the Gumbel softmax. However, $n$ is only applied to logits to scale the function and not to modify the Gumbel temperature.\n\n[1] Yang, Y., Bamler, R., \\& Mandt, S. (2020). Improving inference for neural image compression. Advances in Neural Information Processing Systems, 33, 573-584.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Contains an additional experiment requested by reviewer N6DE.

Meta-review: 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.
Decision: Accept (poster)