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):
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.
