Diffusion model for full field inference: application to weak lensing

Benjamin Remy

Survey Science Group, KICP

Shear

Convergence

\gamma
\kappa

Weak lensing mass-mapping as an inverse problem

\gamma_1 = \frac{1}{2}(\partial_1^2 -\partial_2^2)\Psi
\gamma_2 = \partial_1^2\partial_2^2\Psi
\kappa = \frac{1}{2}(\partial_1^2 + \partial_2^2)\Psi

Shear

Convergence

\gamma
\kappa

Weak lensing mass-mapping as an inverse problem

\gamma = \mathbf{P}\kappa

Mass-mapping with the Dark Energy Survey (DES) Y3

Jeffrey, Gatti, et al. 2021

What about using the prior from cosmological simulations?

The evolution of generative models

  • Deep Belief Network
    (Hinton et al. 2006)

The evolution of generative models

  • Deep Belief Network
    (Hinton et al. 2006)

 

  • Variational AutoEncoder
    (Kingma & Welling 2014)

The evolution of generative models

  • Deep Belief Network
    (Hinton et al. 2006)

 

  • Variational AutoEncoder
    (Kingma & Welling 2014)

 

  • Generative Adversarial Network
    (Goodfellow et al. 2014), WGAN (2017)

The evolution of generative models

  • Deep Belief Network
    (Hinton et al. 2006)

 

  • Variational AutoEncoder
    (Kingma & Welling 2014)

 

  • Generative Adversarial Network
    (Goodfellow et al. 2014), WGAN (2017)

 

  • Diffusion models (Ho et al. 2020, Song et al. 2020)

Midjourney v3

Generative modeling

Generative modeling aims to learn an implicit distribution

from which the training set

\mathbb{P}
X = \left\{ x_0, x_1, \dots , x_n \right\}

This usually means building a parametric model         that tries to be close to

\mathbb{P}_\theta
\mathbb{P}
\mathbb{P}
x \sim \mathbb{P}
\mathbb{P}_\theta

Model

Samples

True

Once trained, we can typically samples from        and evaluate the density 

\mathbb{P}_\theta
p_\theta(x)

is drawn.

Diffusion models

Once                       is learned, draw                                and solve the reverse SDE to get target distribution samples

\nabla \log p(x)
x_T \sim \mathcal{N}(0, \sigma_T)
x_0

The score                      is all you need

\nabla \log p_t(x)

We learn is by denoising our data

If               is corrupted by additional Gaussian noise

x \sim \mathbb{P}
x^\prime = x + n

Then a denoiser                    trained under an        loss

d_\theta(x^\prime, \sigma_t)
\ell_2
\ell_2 = \| x - d_\theta(x^\prime, \sigma_t) \|^2_2

Verifies the Tweedie's formula

\nabla \log p_t(x^\prime) = \frac{d_\theta(x^\prime, \sigma_t) - x^\prime}{\sigma_t^2}
\nabla \log p_t(x^\prime)
x
x + n
n \sim \mathcal{N}(0, \sigma_t)
d_\theta(x^\prime, \sigma_t)

Diffusion models for inverse problems

y = \mathbf{A}x + n

     is known and encodes our physical understanding of the problem.
⟹  When non-invertible or ill-conditioned, the inverse problem is ill-posed with no unique solution  x

\mathbf{A}

Deconvolution

Inpainting

Denoising

Diffusion models for inverse problems

y = \mathbf{A}x + n
p(x\mid y) \propto p(y\mid x)p(x)
p(y\mid x)
p(x)
\hat x = \arg \underset{x}{\max}\,\, \log p(y\mid x) + \log p(x)
\hat x = \arg \underset{x}{\max}\,\, -\frac{1}{2}\| y - \mathbf{A} \|^2 + \log p(x)

Or, we can aim to sample from the full posterior distribution by MCMC techniques

For instance, if      is Gaussian

n

is the data likelihood, which constains the physics

is our prior knowledge on the solution

We can target the Maximum A Posteriori (MAP)

Bayesian view of the problem:

Posterior sampling with diffusion models

Sampling from the prior

dx = \left[ f(x, t) - g^2(t) \color{yellow}{\nabla_x \log p(x_t)} \right]dt + g(t)dw
p(x)
\kappa\sim p(\kappa)

Posterior sampling with diffusion models

Sampling from the posterior

dx = \left[ f(x, t) - g^2(t) \color{yellow}{\nabla_x \log p(x_t\mid y)} \right]dt + g(t)dw
p(x\mid y)
\color{yellow}{\nabla_x \log p(x_t\mid y)}\, \color{white}{= \nabla_x \log p(y\mid x_t) + \nabla_x \log p(x_t)}
\color{yellow}{\nabla_x \log p(x_t\mid y)}\, \color{white}{= \nabla_x\, \underset{-\frac{1}{2\sigma^2} \| y-\mathbf{A}x \|^2}{\underbrace{\log p(y\mid x_t)}} \,+}\, \color{orange}{\underset{\text{learned}}{\underbrace{\nabla_x \log p(x_t)}}}

Ground truth simulation

Posterior samples

Posterior sampling with diffusion models

Application to HST/ACS COSMOS field

Massey et al. 2007

Remy et al. 2022 (Posterior mean)

Remy et al. 2022 (Posterior samples)

Posterior sampling with diffusion models

Remy et al. 2022 (Posterior mean)

We built a generative model of mass-maps, that we can condition on observation to get the full posterior distribution

p(\kappa\mid \gamma)

But we implicitly assumed a cosmological model when training our prior            , i.e. 

p(\kappa)
p(\kappa\mid \mathcal{M})

How to infer jointly the cosmology and the mass map?

But we implicitly assumed a cosmological model when training our prior            , i.e. 

Benjamin Remy, Francois Lanusse, Niall Jeffrey, Jia Liu, Jean-Luc Starck,

Ken Osato, Tim Schrabback, Probabilistic mass-mapping with neural score estimation, A&A 2022

Joint inference of mass-maps and cosmological parameters

with Chihway Chang and Rebecca Willet

Learning the joint distribution

p(\theta, \kappa)

We now have a dataset of pairs

\theta, \kappa \sim p(\theta)p(\kappa\mid \theta)
\theta_0, \kappa_0
\theta_1, \kappa_1
\theta_2, \kappa_2
\theta_3, \kappa_3

How can we build a diffusion model to learn the joint distribution?

Learning the joint distribution

p(\theta, \kappa)

We need to model to design a denoiser architecture

to learn the score function

\nabla \log p(\theta, \kappa)
\left[\theta^\prime, \kappa^\prime\right] = \left[\theta, \kappa\right] + \sigma_tn
d^\star(\theta^\prime, \kappa^\prime, \sigma_t) = \left[\theta^\star, \kappa^\star\right]

We want to learn the denoiser

U-net architecture based on convolutional layers

Only inputs and ouputs images (or volumes)....

How do we add the cosmology?

\nabla \log p(\theta, \kappa)

Use Tweedie's formula to get

p(\theta, \kappa)

Learning the joint distribution

p(\theta, \kappa)

We need to model to design a denoiser architecture

to learn the score function.

\left[\theta_t, \kappa_t\right] = \left[\theta, \kappa\right] + \sigma_tn
d^\star(\theta_t, \kappa_t, \sigma_t) = \left[\theta^\star, \kappa^\star\right]

We want to learn the denoiser

We now have a joint score function!

\nabla_{\kappa, \theta} \log p(\kappa, \theta)

Learning the joint distribution

p(\theta, \kappa)

Marginal

p(\theta\mid \gamma)

Coverage test

We built a generative model of mass-maps and cosmology, that we can condition on observation to get the joint posterior distribution

p(\theta, \kappa\mid \gamma)

Paper out soon...

Thank you!