How to find bits/dim of a gaussian output distribution?

Question

I have images that are 64x64x3 and 64x64x1 8-bit. I transform those images down to [-1,1]. I now want to find the bits/dim for my VAE log probability. How do I find the bits/dim of the log likelihood? Doing $-\frac{ELBO}{log(2) D}$ where D is the dimensions of the image does not give a correct output, it seems that the bits/dim is too low. The output distribution I use is a Gaussian distribution as such the decoder outputs $\mu$ and the variance is fixed, i.e. p(x|z) I assume to follow a Gaussian distribution.

Thanks a lot.

Answer 1

A few hints:

1. Your images are discrete, but your model's distribution is continuous. bits/dim only makes sense in a discrete setting (it takes "infinite bits" to encode real numbers with infinite precision).
2. This suggests you should discretize your $p(x|z)$. A quick and dirty way to do this is to divide your density function by 128 and scale it, but this is not quite correct, as there's no reason this will sum to 1 and be a valid probability distribution. The proper way to do this would be binning your continuous values and integrating the density function across each bin.
3. While discretizing there are a few things left to your discretion, like what do you do outside the bounds (-1,1). Anything reasonable (truncating, clipping) seems fair here.
4. ELBO is a lower bound on the log-likelihood, and therefore you'll get an upper bound on your bits/dim. Appendix D of VAEs has some suggestions for how to best compute $p(x)$.

---

Edit: more details on the precise discretization procedure.

If you have some CDF $F: \mathbb{R} \rightarrow \mathbb{R}$ and some discrete set of values $Z = \{z_1, z_2, \ldots, z_n\}$ (aka the 256 possible pixel values), and for each $z_i$ you have a corresponding interval $(u_i, v_i)$, and all the intervals partition $\mathbb{R}$, then the obvious way to discretize your distribution is to make $P(Z = z_i) = F(v_i) - F(u_i)$.

If you arrived at your continuous distribution by modeling your ordinal discrete values as real, then the "obvious" intervals would be $(i-0.5,i+0.5)$ for each $i$, except for $i=0$ where the interval would be $(-\infty, i+0.5)$ and for the largest $i$, where the interval would be $(i-0.5, \infty)$

If you then applied some other mapping function $g$ (in your case, it sounds like $g(x) = \frac{x}{128}-1$, then you'd probably want to use $(g(i-0.5), g(i+0.5))$ as your intervals instead.

Note that $F(v_i) - F(u_i)$ in this case is pretty close to $f(g(i))\cdot |\frac{\partial g}{\partial x}|$ ($f$ being the density), so a rough approximation is just to call the density function and divide by 128.