Latent Diffusion: Diffusion in VAE latent space.

Implements the Latent Diffusion Model (LDM) concept from "High-Resolution Image Synthesis with Latent Diffusion Models" (Rombach et al., CVPR 2022). Instead of diffusing in the full input space, LDM first compresses data with a VAE encoder, runs diffusion in the compact latent space, then decodes back.

Key Innovation: Perceptual Compression + Diffusion

Training:
  1. Train VAE: input -> encoder -> z -> decoder -> reconstruction
  2. Freeze VAE
  3. Train diffusion in latent space:
     z_0 = encoder(input)
     z_t = add_noise(z_0, t)
     eps_hat = denoiser(z_t, t)
     loss = MSE(eps, eps_hat)

Inference:
  1. Sample z_T ~ N(0, I) in latent space
  2. Denoise: z_0 = diffusion_sample(z_T)
  3. Decode: output = decoder(z_0)

Advantages

Architecture

Returns a tuple of three models: {encoder, decoder, denoiser}

Input [batch, input_size]
      |
      v
+-------------------+
| Encoder (frozen)  |  -> z_0 [batch, latent_size]
+-------------------+
      |
      v (add noise)
+-------------------+
| Denoiser          |  -> eps_hat [batch, latent_size]
| (z_t, t) -> eps   |
+-------------------+
      |
      v (denoise)
+-------------------+
| Decoder (frozen)  |  -> output [batch, input_size]
+-------------------+

Usage

{encoder, decoder, denoiser} = LatentDiffusion.build(
  input_size: 287,
  latent_size: 32,
  hidden_size: 256,
  num_layers: 4
)

# Train VAE first, then freeze and train denoiser

Reference

• Paper: "High-Resolution Image Synthesis with Latent Diffusion Models"
• arXiv: https://arxiv.org/abs/2112.10752