Why does Stable Diffusion use VAE instead of AE?

Question

I am currently studying the Latent Diffusion Models (LDMs) and am interested in training my own model using a unique dataset. In my research, I came across Stable Diffusion (SD). Some sources suggest that SD employs VAEs for the encoding and decoding of images, as opposed to AEs.

In my understanding, Diffusion Models learn the inverse of the noising process which perturbs the distribution of the images, transforming it into a Gaussian distribution. On the other hand, in VAEs, the distribution of the latent variables is also Gaussian. As a result, when using the LDM in conjunction with the VAE, it appears that we might merely be handling a redundant diffusion process, transitioning from one Gaussian distribution to another.

My main question is: Does SD really employ VAEs, rather than AEs, throughout its learning process?

The LDM does not seem to be a good match for the VAE. If my understanding is wrong, please let me know.

Answer 1

To my knowledge, when it comes to stable diffusion, the VQ-VAE is the commonly used method. This differs slightly from vanilla VAE which assumes the encoded features to be a normal distribution and the sampled values from the distribution are considered as the VAE's embedding. However, VQ-VAE doesn't make this assumption and instead finds the most similar embedding from the codebook with the encoder's output. Its focus is to find more meaningful representation from the VAE, resulting in more discrete values than those obtained from AE, as it quantizes the range of values into a single value.

Answer 2

Stabilizing diffusion uses variational autoencoders (VAEs) instead of autoencoders (AEs) because VAEs allow for the generation of continuously distributed representations in the latent space, which better captures the complexity and diversity of the data, whereas AEs typically generate discrete representations. This makes VAE more suitable for generating highly accurate and diverse data.

Answer 3

Yes, SD uses VAEs instead of AEs in learning process.

There are reasons behind it like

• AEs create deterministic latent space representation. That means for any input given AEs produce only specific points in latent space meanwhile VAEs generate distribution over latent space. deterministic approach caps the diversity of the generated image.

• AEs and VAEs don't share the same kind of regularization either. VAE has Gaussian Distribution in latent space which helps in generalization better and also prevents overfitting. AE can Overfit and lead to a less generalized model.

• VAEs generated results are smooth in continuous latent space. while AE results are are disjoint in latent space which affects the overall idea of SD.

Answer 4

The VAE is already trained to compress the input image into a lower dimensional Gaussian distribution. This distribution is close to the unit Gaussian distribution but is not exact. The diffusion model corrects this mismatch, leading to higher quality samples.

Answer 5

Vanilla autoencoders are perfectly fine. For example, here's one trained on MNIST. VAEs are often used because they smooth out the latent space (see this StackExchange answer). But you can smooth them out however you want, e.g. with R1 regularization.