What is the point of the means in the function quantize of EntropyModel ?

[gitlabspy]

https://github.com/InterDigitalInc/CompressAI/blob/master/compressai/entropy_models/entropy_models.py#L168
Why do we need outputs -= means?

[YodaEmbedding]

1. When $y$ is near the predicted mean, it reduces the error in the reconstructed $\hat{y}$.
2. The precomputed, discretized CDFs for EntropyBottleneck are centered around offsets (which we set means to).
3. The CDFs for GaussianConditional are evaluated for a zero-mean variable $Y \sim \mathcal{N}(0, \sigma^2)$ at values $y ∈ \{..., -2, -1, 0, 1, 2, ...\}$.

Example showing (1):

mean = 0.48
y = 0.47

y_hat1 = round(y) = 0.0
y_hat2 = round(y - mean) + mean = 0.48

Evidently, in this example, $|\hat{y}_2 - y| \leq |\hat{y}_1 - y|$. In general, for arbitrary $y$ distributed uniformly, the asymptotic average error of $\hat{y}_1$ is the same as $\hat{y}_2$. But near the mean, $\hat{y}_2$ is usually better.

Regarding (2), a mean-centered variable is much better for sharp PDFs such as pdf = (0.03, 0.94, 0.03). The PDFs of most "almost empty" channels look similar to this. Also, if e.g., offset = 0.48, then clearly by our example above, the variable $y =$

0.47, 0.56, 0.42, 0.51, 1.1, 0.38, 0.41, 0.44

is much better reconstructed as $\hat{y} =$

0.48, 0.48, 0.48, 0.48, 1.48, 0.48, 0.48, 0.48

than as $\hat{y} =$

0, 1, 0, 1, 1, 0, 0, 0

...which does not follow the expected distribution, causing the rate cost to skyrocket, and also has a large reconstruction error.

Regarding (3), for large σ, it doesn't really matter, but for small σ, it's the same as the example for (2) above.

[gitlabspy]

That's so clever! Thanks a lot!