Skip to main content

Generative Models

Generative models learn the distribution of data so they can produce new samples — new digits, new images — rather than just labeling existing ones. This track builds four families from scratch, each adding one new idea about how to model that distribution.

What you'll build​

  • Autoencoders — compress and reconstruct images through a bottleneck. Introduces nnx.ConvTranspose for learned upsampling (the decoder primitive every later guide reuses).
  • Variational Autoencoders (VAE) — put a probability distribution over the latent space and sample new digits, using the reparameterization trick and the ELBO objective.
  • Generative Adversarial Networks (GAN) — a generator and discriminator locked in a game; learn density implicitly with two optimizers and spectral normalization.
  • Diffusion Models (DDPM) — generate by iteratively denoising pure noise; a small U-Net predicts the noise at each step.
  • Normalizing Flows — an invertible network that gives you the exact likelihood via the change-of-variables formula (RealNVP coupling layers).

Suggested order​

Read top to bottom: each guide builds on the previous one's intuition (reconstruction → latent distribution → adversarial → iterative denoising).

Prerequisites​

You should be comfortable with CNNs and the training loop. VAE and beyond lean on Flax NNX state and RNGs.

Next steps​