Heidelberg Summer School

From Zero to Generative

Carolina Cuesta-Lazaro

Art: "The art of painting" by Johannes Vermeer

Learning Generative Modelling from scratch

Flatiron Institute

Institute for Advanced Studies

p(\mathrm{World}|\mathrm{Prompt})

["Genie 2: A large-scale foundation model" Parker-Holder et al (2024)]

p(\mathrm{Drug}|\mathrm{Properties})

["Generative AI for designing and validating easily synthesizable and structurally novel antibiotics" Swanson et al]

Probabilistic ML has made high dimensional inference tractable

1024x1024xTime

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

["Genie 3: A new frontier for world models" Parker-Holder et al (2025)]

https://parti.research.google

A portrait photo of a kangaroo wearing an orange hoodie and blue sunglasses standing on the grass in front of the Sydney Opera House holding a sign on the chest that says Welcome Friends!

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

BEFORE

Artificial General Intelligence?

AFTER

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

p(x)

p(y|x)

p(x|y) = \frac{p(y|x)p(x)}{p(y)}

p(x|y)

Generation vs Discrimination

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

p_\phi(x)

Data

A PDF that we can optimize

Maximize the likelihood of the data

Generative Models

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Generative Models 101

Maximize the likelihood of the training samples

\hat \phi = \argmax \left[ \log p_\phi (x_\mathrm{train}) \right]

x_1

x_2

Parametric Model

p_\phi(x)

Training Samples

x_\mathrm{train}

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

x_1

x_2

Trained Model

p_\phi(x)

Evaluate probabilities

Low Probability

High Probability

Generate Novel Samples

Simulator

Generative Model

Fast emulators

Testing Theories

Generative Model

Simulator

Generative Models: Simulate and Analyze

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

The Generative Zoo

GANS

VAEs

Normalizing

Flows

Diffusion Models

[Image Credit: https://lilianweng.github.io/posts/2018-10-13-flow-models/]

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Bridging two distributions

x_1

x_0

Base

Data

How is the bridge constrained?

Normalizing flows: Reverse = Forward inverse

Diffusion: Forward = Gaussian noising

Flow Matching: Forward = Interpolant

is p(x0) restricted?

Diffusion: p(x0) is Gaussian

Normalising flows: p(x0) can be evaluated

Is bridge stochastic (SDE) or deterministic (ODE)?

Diffusion: Stochastic (SDE)

Normalising flows: Deterministic (ODE)

(Exact likelihood evaluation)

"Creating noise from data is easy;

creating data from noise is generative modeling."

Yang Song

Flow: Change of Variables

X \sim \mathcal{N}(\mu,\sigma)

Y = g(X) = a X + b

How is

distributed?

p_Y(y) = p_X(g^{-1}(y)) \left| \frac{dg^{-1}(y)}{dy}\right|

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Y \sim N(aμ + b, a²σ²)

Transformation (flow):

\mathrm{Uniform(0,1)} \rightarrow U_1, U_2

Z_0 = \sqrt{-2 \ln U_1} \cos(2 \pi U_2)

Z_1 = \sqrt{-2 \ln U_1} \sin(2 \pi U_2)

Z_0, Z_2 \leftarrow N(0,1)

Box-Muller transform

Normalizing flows in 1934

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Tutorial 1

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

\mathrm{Uniform(0,1)} \rightarrow U_1, U_2

Z_0, Z_2 \leftarrow N(0,1)

On tutorials

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Different people learn differently

Take the data generator and start from scratch

Fill in the gaps

Start from a working solution and play around

On tutorials

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Base distribution

Target distribution

p_X(x) = p_Z(z) \left| \frac{dz}{dx}\right|

Z \sim \mathcal{N} (0,1) \rightarrow g(z) \rightarrow X

Invertible transformation

z \sim p_Z(z)

p_Z(z)

Normalizing flows

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Normalizing flows

[Image Credit: "Understanding Deep Learning" Simon J.D. Prince]

z \sim p(z)

x \sim p(x)

x = f(z)

p(x) = p(z = f^{-1}(x)) \left| \det J_{f^{-1}}(x) \right|

Bijective

Sample

Evaluate probabilities

Probability mass conserved locally

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

z_0 \sim p(z)

z_k = f_k(z_{k-1})

\log p(x) = \log p(z_0) - \sum_{k=1}^{K} \log \left| \det J_{f_k} (z_{k-1}) \right|

Image Credit: "Understanding Deep Learning" Simon J.D. Prince

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Invertible functions aren't that common!

Splines

Issues NFs: Lack of flexibility

Invertible functions
Tractable Jacobians

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Masked Autoregressive Flows

p(x) = \prod_i{p(x_i \,|\, x_{1:i-1})}

p(x_i \,|\, x_{1:i-1}) = \mathcal{N}(x_i \,|\,\mu_i, (\exp\alpha_i)^2)

\mu_i, \alpha_i = f_{\phi_i}(x_{1:i-1})

Neural Network

x_i = z_i \exp(\alpha_i) + \mu_i

z_i = (x_i - \mu_i) \exp(-\alpha_i)

Sample

Evaluate probabilities

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

y_i = \mu_i, \alpha_i = f_{\phi_i}(x_{1:i-1})

Masked MLP

x_1

x_2

x_3

y_1

y_2

y_3

h_1

h_2

h_3

h_4

h_5

h_6

\theta

Forward Model

Observable

\color{darkgray}{\Omega_m}, \color{darkgreen}{w_0, w_a},\color{purple}{f_\mathrm{NL}}\, ...

Dark matter

Dark energy

Inflation

Predict

Infer

Parameters

Inverse mapping

p(\theta|x)

Simulation-based Inference

Normalizing flow

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

\frac{dx_t}{dt} = v^\phi_t(x_t)

x_1 = x_0 + \int_0^1 v^\phi_t(x_t) dt

\frac{d p(x_t)}{dt} = - \nabla \left( v^\phi_t(x_t) p(x_t) \right)

In continuous time

Continuity Equation

[Image Credit: "Understanding Deep Learning" Simon J.D. Prince]

x_0 = x_1 + \int_1^0 v^\phi_t(x_t) dt

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Chen et al. (2018), Grathwohl et al. (2018)

x_1 = x_0 + \int_0^1 v_\theta (x(t),t) dt

Generate

x_0

x_1

Evaluate Probability

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Tutorial 2: Continuous time normalizing flows

x_0

Gaussian

Two

Moons

x_1

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

\log p_X(x) = \log p_Z(z) + \int_0^1 \mathrm{Tr} J_v (x(t)) dt

Loss requires solving an ODE!

Diffusion, Flow matching, Interpolants... All ways to avoid this at training time

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Conditional Flow matching

x_t = (1-t) x_0 + t x_1

\mathcal{L}_\mathrm{CFM} = \mathbb{E}_{t\sim U, x_1 \sim p_1, x_t \sim p_t}\left[\| u_t^\phi(x_t) - u_t^\mathrm{target}(x|x_1) \|^2 \right]

Assume a conditional vector field (known at training time)

The loss that we can compute

The gradients of the losses are the same!

\nabla_\phi \mathcal{L}_\mathrm{CFM} = \nabla_\phi \mathcal{L}_\mathrm{FM}

x_0

x_1

["Flow Matching for Generative Modeling" Lipman et al]

["Stochastic Interpolants: A Unifying framework for Flows and Diffusions" Albergo et al]

u_t^\mathrm{target}(x) = \int u_t(x|x_1) \frac{p_t(x|x_1) p_1(x_1)}{p_t(x)} \, dx_1

p_t(x) = \int p_t(x|x_1) q(x_1) \, dx_1

Intractable

\mathcal{L}_\mathrm{FM} = \mathbb{E}_{t \sim U,x \sim p_t}\left[\| u_t^\phi(x) - u_t^\mathrm{target}(x) \|^2 \right]

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Flow Matching

\frac{dz_t}{dt} = u^\phi_t(z_t)

x = z_0 + \int_0^1 u^\phi_t(z_t) dt

Continuity equation

\frac{d p(z_t)}{dt} = - \nabla \left( u^\phi_t(z_t) p(z_t) \right)

[Image Credit: "Understanding Deep Learning" Simon J.D. Prince]

Sample

Evaluate probabilities

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Tutorial 3: Flow Matching from scratch

x_0

Gaussian

Two

Moons

x_1

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Diffusion Models

Reverse diffusion: Denoise previous step

Forward diffusion: Add Gaussian noise (fixed)

Prompt

A person half Yoda half Gandalf

Denoising = Regression

Fixed base distribution:

Gaussian

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

["A point cloud approach to generative modeling for galaxy surveys at the field level"

Cuesta-Lazaro and Mishra-Sharma
International Conference on Machine Learning ICML AI4Astro 2023, Spotlight talk, arXiv:2311.17141]

Base Distribution

Target Distribution

Simulated Galaxy 3d Map

Prompt:

\Omega_m, \sigma_8

Prompt: A person half Yoda half Gandalf

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

\mathrm{R} \sim p(\theta|x)

\mathrm{F} \sim \hat{p}(\theta|x)

Real or Fake?

How good is my generative model?

["A Practical Guide to Sample-based Statistical Distances for Evaluating Generative Models in Science" Bischoff et al 2024
arXiv:2403.12636]

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Mean relative velocity

k Nearest neighbours

Pair separation

Varying cosmological parameters

Physics as a testing ground: Well-understood summary statistics enable rigorous validation of generative models

Reproducing summary statistics

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Has my model learned the underlying density?

Carolina Cuesta-Lazaro IAIFI/MIT - From Zero to Generative

["Generalization in diffusion models arises from geometry-adaptive harmonic representations" Kadkhodaie et al (2024)]

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Split training set into non-overlapping

S_1, S_2

Generate = Understand?

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

1) Compress into low dimensional latent

p(x|z)

z = f(x)

2) Use generative model as decoder

Encoder

Autoregressive in Frequency

["CosmoFlow: Scale-Aware Representation Learning for Cosmology with Flow Matching" Kannan, Qiu, Cuesta-Lazaro, Jeong (in prep)]

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Generate = Understand?

["CosmoFlow: Scale-Aware Representation Learning for Cosmology with Flow Matching" Kannan, Qiu, Cuesta-Lazaro, Jeong (in prep)]

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Generate = Understand?

["CosmoFlow: Scale-Aware Representation Learning for Cosmology with Flow Matching" Kannan, Qiu, Cuesta-Lazaro, Jeong (in prep)]

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Generative Priors: Plug and Play

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

EHT posterior samples with different priors

["Event-horizon-scale Imaging of M87* under Different Assumptions via Deep Generative Image Priors" Feng et al]

CIFAR-10

GRMHD

RIAF

CelebA

(Sims)

(LR Natural Images)

(Human Faces)

\log p(x|y) = \log p(y|x) + \log p(x) + \mathrm{constant}

Prior

Generative priors: learn directly?

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

["Learning Diffusion Priors from Observations by Expectation Maximization" Rozet et al]

Bridging two distributions

x_1

x_0

Base

Data

How is the bridge constrained?

Normalizing flows: Reverse = Forward inverse

Diffusion: Forward = Gaussian noising

Flow Matching: Forward = Interpolant

is p(x0) restricted?

Diffusion: p(x0) is Gaussian

Normalising flows: p(x0) can be evaluated

Is bridge stochastic (SDE) or deterministic (ODE)?

Diffusion: Stochastic (SDE)

Normalising flows: Deterministic (ODE)

(Exact likelihood evaluation)

Tutorial 4: Generating Galaxy Images

x_0

Gaussian

Two

Moons

x_1

Carolina Cuesta-Lazaro Flatiron/IAS - From Zero to Generative

Books by Kevin P. Murphy
- Machine learning, a probabilistic perspective
- Probabilistic Machine Learning: advanced topics
ML4Astro workshop https://ml4astro.github.io/icml2023/
ProbAI summer school https://github.com/probabilisticai/probai-2023
IAIFI Summer school
- https://github.com/iaifi/summer-school-2023
- Tutorials https://github.com/florpi/summer_school_generative
Blogposts

Carolina Cuesta-Lazaro IAIFI/MIT - From Zero to Generative

References

carolina.clzr@gmail.com