How does the bottleneck z dimension affect the reconstruction loss in VAEs

Question

I have come across few VAE papers that all report a similar metric bits/dim

Many (if not all) fail to mention the bottleneck size of the z space. I do know that this would directly affect the reconstruction loss, e.g. a too small bottleneck and you can't get any meaningful reconstructions (posterior collapse).

However it is not clear, is there an upper limit to this bottleneck.

For example: a z dimension of 512 and 2048 produces identical benchmarks for me. Is this number upper bounded by the architectural choices rather than the bottleneck size?

I would also imagine a bottleneck of size the same dimensionality as your input should lead to perfect reconstructions, but then wouldn't find the definition of a bottleneck.

Couldn't find any relevant papers. Any thoughts appreciated.

Answer 1

It seems to me that the link you are missing here is to the probabilistic / information theory interpretation of VAEs. When the capacity of your networks is large enough you will reach a point where the solution with a larger latent space does not keep more information than a smaller one. This is possible in VAEs because they produce a noisy representation inside.

To clarify things: First the bits/dim metric is per dimension of the input. You can read more about this metric in links collected here: What is bits per dimension (bits/dim) exactly (in pixel CNN papers)?

Maybe the limit for infinitely big networks and infinite data is instructive here: VAEs optimize a variational bound on the evidence for the model. This is bounded from above by the true entropy of your data. At or near this point your bits/dim will converge and adding more complexity anywhere will no longer improve performance. With limited data this point will come earlier.

As you appear to think in terms of bottlenecks & auto-encoders: For VAEs the bottleneck is not really the number of dimensions in the latent space, but the noise. Without noise even a single continuous number has infinite capacity. As VAEs are allowed to adjust the noise on their representation they may very well represent the same amount of information in fewer dimensions with less noise. Thus, the number of latent dimensions much less informative about the capacity of the encoding than for classic autoencoders. Indeed VAEs regularly have latent space units who converge towards always being equal to the prior, i.e. not carrying any information about images.

In practice more latent units become harder to train and costly, such that you would avoid using too many latent units which "die", but from a theoretical viewpoint many dimensions does not equal an open bottleneck for VAEs.

Thus, overall I would say the dimensionality of z for VAEs is one of many knobs changing the expressiveness/complexity of the encoder/decoder and does not affect the reconstruction loss directly beyond this.

Answer 2

As mentioned in Variational Autoencoder − Dimension of the latent space, there is a heuristic upper-bound for the latent variable dimension: the size of the training data.

If you encoder is sufficiently powered (we assume it is anyways for VAEs), if your latent variable is $N$ dimensional and you have $N$ training samples, then your encoder can simply encode each sample in one dimension of the latent, severely overfitting your model.

In practice, however, that bound is probably tighter, since the encoder is non-linear it can fit the training data into a lower dimensional latent and yet overfit.

One possible test for this is to decode random latents and see what they look like (assuming you are modeling images).