Latent Diffusion Model Can't Learn the Latent Space of a VAE for the MNIST-Fashion Dataset

I'm currently playing around with LDMs on the MNIST-Fashion dataset. I thought the VQVAEs used in the original paper were a bit overkill for what I'm doing (and I don't fully understand how they construct the discretized codebook latent space), so I went with a simple convolutional autoencoder with a kl-regularizer to map to an approximately gaussian latent space. I've run this model a few times and verified that it does reconstruct the original image inputs fairly well. I run into issues when I try to use this model with my LDM implementation. I first made sure standard diffusion directly in image space works, which it does. I then tried latent diffusion with my trained autoencoder and I can't get the loss to drop below a certain threshold (~1.0), and the image outputs are pretty much gaussian noise still. As stated in the paper, I sample latent vectors from the encoder part of the autoencoder and scale them with stats collected from the first batch of data like so: batch = autoencoder.encode(batch).sample().detach() # rescale the embeddings to be unit variance if epoch == 0 and step == 0: print("Calculating scale factor...") std = batch.flatten().std() scale_factor = 1. / std cfg.scale_factor = scale_factor.item() batch *= scale_factor And then pretty much everything else (applying noise, calculating loss, etc) is the same as standard diffusion. Am I missing something, or is the latent space of my simple conv-autoencoder hard to learn for some reason? I would think that, s…