Why does the latent space in Stable Diffusion have a shape of 64x64x3?

Since the encoding is performed by a Variational Autoencoder, the VAE encoder must output some mean and log variance that we can use to sample a latent vector of shape (latent_dim,). But then how does Stable Diffusion use a vector of shape 64x64x3?

Reference (https://arxiv.org/pdf/2112.10752.pdf, page 24, table 12):

• Can you provide a reference that tells that the latent space has the dimensionality you say?
– nbro
Commented Feb 10 at 16:19
• check updated question Commented Feb 10 at 18:20
• Can you provide a citation/link to the paper where you got the table? I assume it's from Rombach et al., 2022? Commented Feb 11 at 1:17

There's not really a restriction on the shape for variational autoencoders. If you really wanted a 1D vector, you could just flatten the matrix and get a vector of size 64 * 64 * 3.

The 64x64x3 size makes sense if you consider why they encode images into a latent dimension in the first place.

Consider this passage from the latent diffusion paper:

(i) By leaving the high-dimensional image space, we obtain DMs which are computationally much more efficient because sampling is performed on a low-dimensional space. (ii) We exploit the inductive bias of DMs inherited from their UNet architecture [71], which makes them particularly effective for data with spatial structure and therefore alleviates the need for aggressive, quality-reducing compression levels as required by previous approaches [23, 66].

and

This is in contrast to previous works [23, 66], which relied on an arbitrary 1D ordering of the learned space z to model its distribution autoregressively and thereby ignored much of the inherent structure of z.

The purpose of the latent encoding is to compress images to a lower-dimensional space that's easier to work with---but you also want to make sure that you aren't losing information during the compression process.

A 2D representation would better models spatial structures that subsequent UNet layers are designed to work with.

• Thank you for the feedback! But if we consider 3D (binary voxel) data instead of images, then what structure of latent representation would be useful: 2D or 3D? The issue is that I can't find any 3D latent diffusion implementations. Commented Feb 11 at 9:30
• Yeah, I'm not too familiar with this as well. Although this paper (snap-research.github.io/3DVADER/paper.pdf) seems to do things in a different way. Their latent has three dimensions (whereas you'd expect four if you were to extend the 2D diffusion model). Commented Feb 11 at 18:18