Amount of quantization

[Danish Nazir]

Hello,

Is it possible to somehow defined the level of quantization during training? For instance one can use the following code for hard quantization during training,

from .compressai.ops.ops import quantize_ste
_, y_likelihoods = self.gaussian_conditional(y, scales_hat)
y_hat = quantize_ste(y)

then if we look inside quantize_ste function, we have
return (torch.round(x) - x).detach() + x
This is basically dependant on the input x, however, I am not really sure what level of quanitzation will be applied to latent variable y.
Is this possible to know during training?

[Mateen Ulhaq]

The latent y is quantized uniformly to bins of width 1.

This is done in one of the following ways:

1. Adding uniform noise $\mathcal{U}(-\frac{1}{2}, \frac{1}{2})$ during training, and rounding during evaluation.
2. Applying the Straight-Through Estimator (STE), which rounds, but passes through the gradients during the backwards pass:

   ste(y) = round(y)   during forward
   ste(y) = y          during backward

   i.e., $\frac{d}{dy} \mathrm{STE}(y) = 1$

[Danish Nazir]

Thanks for your response. is it possible/reccomended to somehow change the bin width to influence the amount of quantization?

[Mateen Ulhaq]

Sure. Just replace round with your own quantize:

@torch.no_grad()
def quantize(x):
    return (torch.sign(x) * x.abs().sqrt()).round()

And then:

def quantize_ste(x):
    return x + (quantize(x) - x).detach()

Or for noise-based quantization, just modify each symbol's noise gain accordingly:

noise_gain = function_depending_on_bin_widths_and_y(y)
y_hat = y + noise_gain * 2 * (torch.rand(y.shape) - 0.5)