Bridging Simulators with Conditional Optimal Transport

Justine Zeghal, Benjamin Remy,

Yashar Hezaveh, François Lanusse,

Laurence Perreault-Levasseur

Field Level Meeting SkAI, Chicago

February 2026

\underbrace{p(\theta|x=x_0)}_{\text{posterior}}
\underbrace{p(x = x_0|\theta)}_{\text{likelihood}}
\propto

 Full-field inference: extracting all cosmological information

\underbrace{p(\theta)}_{\text{prior}}

Bayes theorem:

\underbrace{p(\theta|x=x_0)}_{\text{posterior}}
\underbrace{p(x = x_0|\theta)}_{\text{likelihood}}
\propto

 Full-field inference: extracting all cosmological information

\underbrace{p(\theta)}_{\text{prior}}

Bayes theorem:

\underbrace{p(\theta|x=x_0)}_{\text{posterior}}
\underbrace{p(x = x_0|\theta)}_{\text{likelihood}}
\propto

 Full-field inference: extracting all cosmological information

\underbrace{p(\theta)}_{\text{prior}}

Bayes theorem:

\underbrace{p(\theta|x=x_0)}_{\text{posterior}}
\underbrace{p(x = x_0|\theta)}_{\text{likelihood}}
\propto

 Full-field inference: extracting all cosmological information

\underbrace{p(\theta)}_{\text{prior}}
\theta

 Simulator

x

Bayes theorem:

\underbrace{p(\theta|x=x_0)}_{\text{posterior}}
\underbrace{p(x = x_0|\theta)}_{\text{likelihood}}
\propto

Bayes theorem:

 Full-field inference: extracting all cosmological information

\underbrace{p(\theta)}_{\text{prior}}
\theta

 Simulator

x

Two ways to get the posterior:

  • Explicit inference:
\rightarrow \text{we need } \nabla_{\theta, z} \log p(x,\theta, z)
\rightarrow \text{we need } (x, \theta) \sim p(x, \theta, z)
  • Implicit inference
\underbrace{p(\theta|x=x_0)}_{\text{posterior}}
\underbrace{p(x = x_0|\theta)}_{\text{likelihood}}
\propto

 Full-field inference: extracting all cosmological information

\underbrace{p(\theta)}_{\text{prior}}
\theta

 Simulator

x

Two ways to get the posterior:

  • Explicit inference:
\rightarrow \text{we need } \nabla_{\theta, z} \log p(x,\theta, z)
\rightarrow \text{we need } (x, \theta) \sim p(x, \theta, z)
  • Implicit inference

Has to be realistic!

Bayes theorem:

 Wrong models generate bias

Fast simulations

Costly simulations

 Wrong models generate bias

→ e.g. full nbody, hydro

\log p(x)

Fast simulations

Costly simulations

O(ms) runtime ❌
differentiable   ❌
❌
realistic ✅

 Wrong models generate bias

→ e.g. full nbody, hydro

\log p(x)
\log p(x)

→ e.g. log-normal, LPT, PM

O(ms) runtime ✅
differentiable   ✅
✅
realistic ❌

Fast simulations

Costly simulations

O(ms) runtime ❌
differentiable   ❌
❌
realistic ✅

 Learning the correction

We can learn 

the correction!

Fast simulations

Costly simulations

 Learning the correction

We can learn 

the correction!

Fast simulations

x_1 = \phi(x_0)
  • it preserves the conditioning,

Costly simulations

\text{We seek a mapping } \phi:
  •                                ,

such that

  • it minimally correct the simulation.
x_1 \sim p_1(x \mid \theta)

 Learning the correction

We can learn 

the correction!

Fast simulations

x_1 = \phi(x_0)
  • it preserves the conditioning,

Costly simulations

\text{We seek a mapping } \phi:
  •                                ,

such that

  • it minimally correct the simulation.

Requirements:

  •      has to map to a distribution sample.
\phi
  •      has to work in high dimensions.
\phi
  •      has to bridge any two distributions.
\phi
  •      has to bridge conditional distributions.
\phi
  •      has to be the solution of the OT problem.
\phi
x_1 \sim p_1(x \mid \theta)

Conditional Optimal Transport Flow Matching

Conditional Optimal Transport Flow Matching

(Lipman et al. 2023)

Flow matching

f^{-1}_1
f^{-1}_2
f_1
f_2
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
x_0 \sim p_0

(Lipman et al. 2023)

Flow matching

f^{-1}_1
f^{-1}_2
f_1
f_2
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
x_0 \sim p_0
\rightarrow \text{ learn discrete}\\ \text{ transformations } f_t

(Lipman et al. 2023)

Flow matching

f^{-1}_1
f^{-1}_2
f_1
f_2
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
x_0 \sim p_0
\rightarrow \text{ learn discrete}\\ \text{ transformations } f_t
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
\rightarrow \text{ learn continuous }\\ \text{ transformations } f_t \\ \text{ solution of }\\ \frac{d x_t}{dt} = \color{#ce6eff}{v_\varphi}\color{black}{(x_t, t)}
x_0 \sim p_0
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
x_0 \sim p_0
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
\rightarrow \text{ learn continuous }\\ \text{ transformations } f_t \\ \text{ solution of }\\ \frac{d x_t}{dt} = \color{#ce6eff}{v_\varphi}\color{black}{(x_t, t)}
f^{-1}_1
f^{-1}_2
f_1
f_2
x_0 \sim p_0
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
x_0 \sim p_0
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
x_0 \sim p_0
\rightarrow \text{ learn discrete}\\ \text{ transformations } f_t
\mathcal{L}_{FM}(\theta) = \mathbb{E}_{p(t)q(x_0,x_1)p_t(x_t\mid x_0,x_1)} \Big[ \| \color{#ce6eff}{v_\varphi} \color{black}{(x_t, t)- v(x_t,t) \|^2 \Big]}

(Lipman et al. 2023)

Flow matching

x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
\rightarrow \text{ learn continuous }\\ \text{ transformations } f_t \\ \text{ solution of }\\ \frac{d x_t}{dt} = \color{#ce6eff}{v_\varphi}\color{black}{(x_t, t)}
f^{-1}_1
f^{-1}_2
f_1
f_2
x_0 \sim p_0
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
x_0 \sim p_0
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
x_0 \sim p_0
\rightarrow \text{ learn discrete}\\ \text{ transformations } f_t
\mathcal{L}_{FM}(\theta) = \mathbb{E}_{p(t)q(x_0,x_1)p_t(x_t\mid x_0,x_1)} \Big[ \| \color{#ce6eff}{v_\varphi} \color{black}{(x_t, t)- v(x_t,t) \|^2 \Big]}
p(x_t\mid x_0,x_1) = \mathcal{N}((1-t)x_1 + t x_0, \sigma)

(Lipman et al. 2023)

Flow matching

x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
\rightarrow \text{ learn continuous }\\ \text{ transformations } f_t \\ \text{ solution of }\\ \frac{d x_t}{dt} = \color{#ce6eff}{v_\varphi}\color{black}{(x_t, t)}
f^{-1}_1
f^{-1}_2
f_1
f_2
x_0 \sim p_0
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
x_0 \sim p_0
x_0 \sim p_0
x_1 \sim p_1
x_t \sim p_t
x_0 \sim p_0
\rightarrow \text{ learn discrete}\\ \text{ transformations } f_t
\mathcal{L}_{FM}(\theta) = \mathbb{E}_{p(t)q(x_0,x_1)p_t(x_t\mid x_0,x_1)} \Big[ \| \color{#ce6eff}{v_\varphi} \color{black}{(x_t, t)- v(x_t,t) \|^2 \Big]}
p(x_t\mid x_0,x_1) = \mathcal{N}((1-t)x_1 + t x_0, \sigma)

(Lipman et al. 2023)

Flow matching

v(x_t, t) = x_1 - x_0

Requirements:

  •      has to map to a distribution sample.
\phi
  •      has to work in high dimensions.
\phi
  •      has to bridge any two distributions.
\phi
  •      has to bridge conditional distributions.
\phi
  •      has to be the solution of the OT problem.
\phi

✅

✅

✅

Conditional Optimal Transport Flow Matching

Requirements:

  •      has to map to a distribution sample.
\phi
  •      has to work in high dimensions.
\phi
  •      has to bridge any two distributions.
\phi
  •      has to bridge conditional distributions.
\phi
  •      has to be the solution of the OT problem.
\phi

✅

✅

✅

Conditional Optimal Transport Flow Matching

 Optimal Transport Flow matching (Tong et al. 2024)

Flow Matching loss function:

\rightarrow (x_0, x_1) \sim p_0(x_0)p_1(x_1)

Indepent coupling:

Optimal Transport coupling:

\rightarrow (x_0, x_1) \sim \pi^*(x_0,x_1)
W(q_0, q_1)^2_2 = \inf_{p_t, v_t} \int_{\mathbb{R}^d}\int^1_0 p_t(x)\|v_t(x)\|^2 dt dx

i.e. minimizes the path for all trajectories between      and     .

p_0
p_1
\text{ with } \pi^*= \inf_{\pi} \int_{{\mathbb{R}^d}^2} ||x_0 - x_1||^2 d\pi(x_0,x_1)

This coupling, combined with the linear interpolant, solve the dynamic OT:

x_0 \sim p_0
\color{blue}{x_1 \sim p_1}
p(x_t\mid x_0,x_1) = \mathcal{N}((1-t)x_1 + t x_0, \sigma)
\mathcal{L}_{FM}(\theta) = \mathbb{E}_{p(t)q(x_0,x_1)p_t(x_t\mid x_0,x_1)} \Big[ \| \color{#ce6eff}{v_\varphi} \color{black}{(x_t, t)- (x_1-x_0) \|^2 \Big]}

Requirements:

  •      has to map to a distribution sample.
\phi
  •      has to work in high dimensions.
\phi
  •      has to bridge any two distributions.
\phi
  •      has to bridge conditional distributions.
\phi
  •      has to be the solution of the OT problem.
\phi

✅

✅

✅

Conditional Optimal Transport Flow Matching

✅

Requirements:

  •      has to map to a distribution sample.
\phi
  •      has to work in high dimensions.
\phi
  •      has to bridge any two distributions.
\phi
  •      has to bridge conditional distributions.
\phi
  •      has to be the solution of the OT problem.
\phi

✅

✅

✅

Conditional Optimal Transport Flow Matching

✅

 Conditional Optimal Transport Flow matching (Kerrigan et al. 2024)

OT Flow Matching loss function:

\text{ where } \pi^*= \inf_{\pi} \int_{{\mathbb{R}^d}^2} ||x_0 - x_1||^2 d\pi(x_0,x_1)
v_\varphi\left[\begin{matrix} \theta \\ x \end{matrix}\right] = \left[\begin{matrix} v_\varphi(\theta, t) \\ v_\varphi(\theta, x, t) \end{matrix}\right] = \left[\begin{matrix} 0 \\ v_\varphi(\theta, x, t) \end{matrix}\right]
c(x_0,x_1, \theta_0, \theta_1)= \|\theta_1 - \theta_0 \|^2 + \epsilon \|x_1 - x_0 \|^2
p(x_t\mid x_0,x_1) = \mathcal{N}((1-t)x_1 + t x_0, \sigma) \text{ with } (x_0, x_1) \sim \pi^*(x_0,x_1)
\mathcal{L}_{FM}(\theta) = \mathbb{E}_{p(t)q(x_0,x_1)p_t(x_t\mid x_0,x_1)} \Big[ \| \color{#ce6eff}{v_\varphi} \color{black}{(x_t, t)- (x_1-x_0) \|^2 \Big]}

Dataset 1

Optimal Transport Plan

\pi(x_0, x_1)

Dataset 2

Requirements:

  •      has to map to a distribution sample.
\phi
  •      has to work in high dimensions.
\phi
  •      has to bridge any two distributions.
\phi
  •      has to bridge conditional distributions.
\phi
  •      has to be the solution of the OT problem.
\phi

✅

✅

✅

✅

✅

→ e.g. full nbody, hydro

\log p(x)
O(ms) runtime ✅
differentiable   ✅
✅
realistic ✅
\log p(x)

→ e.g. log-normal, LPT, PM

O(ms) runtime ✅
differentiable   ✅
✅
realistic ❌

Fast simulations

Emulated simulations

 Results on weak lensing maps

LPT

PM

Learned

Residuals

 NeurIPS Challenge: Weak Lensing Uncertainty

LogNormal

Emulated

Challenge simulation

VS

 NeurIPS Challenge: Weak Lensing Uncertainty

Power spectrum

PDF

LogNormal

Emulated

Challenge simulation

VS

 NeurIPS Challenge: Weak Lensing Uncertainty

🥳

LogNormal

Emulated

Challenge simulation

VS

 NeurIPS Challenge: Weak Lensing Uncertainty

Thank you for your attention!