Learning Objectives & Program

$+ p(x) =$

• First Hour: Combining Physical Models and Generative Modeling
  • Understanding the limits of black-box deep learning
  • Taking a Bayesian point of view on inverse problems
  • Using generative models to capture implicit distributions
  • Applications in astrophysics: mapping the invisible dark matter


• Second Hour: Practical Work: Conditional Forecasting for Heliophysics
  • Compare deterministic vs stochastic forecasting methods
  • Implement guidance with explicit likelihood on a diffusion model

A Motivating Example: Image Deconvolution

Observed $y$

Ground-Based Telescope

$f_\theta$

some deep Convolutional Neural Network

Unknown $x$

Hubble Space Telescope

• A standard approach would be to train a neural network $f_\theta$ to estimate $x$ given $y$.

• Step I: Assemble from data or simulations a training set of images $$\mathcal{D} = \{(x_0, y_0), (x_1, y_1), \ldots, (x_N, y_N) \}$$ $\Longrightarrow$ the dataset contains hardcoded assumptions about PSF $P$ noise $n$, and galaxy morphology $x$.
• Step II: Train the neural network $f_\theta$ under a regression loss of the type: $$ \mathcal{L} = \sum_{i=0}^N \parallel x_i - f_\theta(y_i)\parallel^2 $$

$$ y $$

$$f_\theta$$

$$f_\theta(y)$$

$$ \mbox{True } x $$

Let's try to understand the neural network output by looking at the loss function

• Using an $\ell_2$ loss learns the mean of the $p(x | y)$


• Using an $\ell_1$ loss learns the median of $p(x|y)$


• In general, training a neural network for regression doesn't achieve de mode of the distribution.
  
  
  Check this blogpost to learn more about that.

Example: GalaxyGAN (Schawinski et al. 2017)

Linear inverse problems

Deconvolution

Inpainting

Denoising

A Bayesian view of the problem

• $p(y | x)$ is the data likelihood, which contains the physics
  


• $p(x)$ is our prior knowledge on the solution.

Classical examples of signal priors

Sparse
$$ \log p(x) = \parallel \mathbf{W} x \parallel_1 $$

Gaussian $$ \log p(x) = x^t \mathbf{\Sigma^{-1}} x $$

Total Variation $$ \log p(x) = \parallel \nabla x \parallel_1 $$

But what about this?

The Variational Auto-Encoder

The famous Variational Auto-Encoder

Normalizing Flows

Still a latent variable model, except the mapping $f_\theta$ is made to be bijective.

Dinh et al. 2016

Normalizing Flows

• Assumes a bijective mapping between data space $x$ and latent space $z$ with prior $p(z)$: $$ z = f_{\theta} ( x ) \qquad \mbox{and} \qquad x = f^{-1}_{\theta}(z)$$
• Admits an explicit marginal likelihood: $$ \log p_\theta(x) = \log p(z) + \log \left| \frac{\partial f_\theta}{\partial x} \right|(x) $$

$\Longrightarrow$ The challenge is in designing mappings $f_\theta$ that are both: easy to invert, easy to compute the jacobian of.

One example of NF: RealNVP

Jacobian of an affine coupling layer

In a standard affine RealNVP, one Flow layer is defined as: $$ \begin{matrix} y_{1\ldots d} &=& x_{1 \ldots d} \\ y_{d+1\ldots N} &=& x_{d+1 \ldots N} ⊙ \sigma_\theta(x_{1 \ldots d}) + \mu_\theta(x_{1 \ldots d}) \end{matrix} $$ where $\sigma_\theta$ and $\mu_\theta$ are unconstrained neural networks.

We will call this layer an affine coupling.

Back to our motivating example

Hubble Space Telescope

some deep neural network

Simulated Ground-Based Telescope

A Physicist's approach to the deconvolution problem: let's build a model

Probabilistic model

$$ x \sim ? $$

$$ x \sim \mathcal{N}(z, \Sigma) \quad z \sim ? $$
latent $z$ is a denoised galaxy image

$$ x \sim \mathcal{N}( \mathbf{P} z, \Sigma) \quad z \sim ?$$
latent $z$ is a super-resolved and denoised galaxy image

$$ x \sim \mathcal{N}( \mathbf{P} (\Pi \ast z), \Sigma) \quad z \sim ? $$
latent $z$ is a deconvolved, super-resolved, and denoised galaxy image

$$ x \sim \mathcal{N}( \mathbf{P} (\Pi \ast g_\theta(z)), \Sigma) \quad z \sim \mathcal{N}(0, \mathbf{I}) $$
latent $z$ is a Gaussian sample
$\theta$ are parameters of the model

How to train your dragon model

• Training the generative amounts to finding $\theta_\star$ that maximizes the marginal likelihood of the model: $$p_\theta(x | \Sigma, \Pi) = \int \mathcal{N}( \Pi \ast g_\theta(z), \Sigma) \ p(z) \ dz$$
  $\Longrightarrow$ This is generally intractable


• Efficient training of parameter $\theta$ is made possible by Amortized Variational Inference.

Bayesian Inference a.k.a. Uncertainty Quantification

The Bayesian view of the problem: $$ p(z | x ) \propto p_\theta(x | z, \Sigma, \mathbf{\Pi}) p(z)$$ where:

• $p( z | x )$ is the posterior
• $p( x | z )$ is the data likelihood, contains the physics
• $p( z )$ is the prior

Data
$x_n$

Truth
$x_0$

Posterior samples
$g_\theta(z)$

$\mathbf{P} (\Pi \ast g_\theta(z))$

Median

Data residuals
$x_n - \mathbf{P} (\Pi \ast g_\theta(z))$

Standard Deviation

Now you try it!

Raise your hand when you reach the Maximum Likelihood Solution i.e. end of Step III.

How to solve the posterior inference problem by Variational Inference?

• The core idea of Variational Inference is to assume a model distribution $q_\phi$ and we fit it to the unknown posterior.
• How do we fit a distribution to another distribution?
  We minize the KL divergence $$ \mathbb{D}_{KL}\left( q_\phi(z) \parallel p(z | x) \right) = \mathbb{E}_{z \sim q_{\phi}} \left[ \log \frac{q_\phi(z)}{p(z | x)} \right] $$

  $$ = \mathbb{E}_{z \sim q_{\phi}} \left[ \log q(z) \right] - \mathbb{E}_{z \sim q_{\phi}} \left[ \log p(z | x) \right] $$

  $$ = \mathbb{E}_{z \sim q_{\phi}} \left[ \log q(z) \right] - \mathbb{E}_{z \sim q_{\phi}} \left[ \log p(x | z) + \log p(z) - \log p(x)\right] $$

  $$ = \underbrace{\mathbb{D}_{KL}\left( q_\phi(z) \parallel p(z) \right) - \mathbb{E}_{z \sim q_{\phi}} \left[ \log p(x | z) \right]}_{\mbox{Evidence Lower Bound}} + \log p(x) $$
• Finally, we need a flexible enough model distribution, why not a Normalizing Flow :-) ?

First realization: The score is all you need!

• Whether you are looking for the MAP or sampling with HMC or MALA, you only need access to the score of the posterior: $$\frac{\color{orange} d \color{orange}\log \color{orange}p\color{orange}(\color{orange}x \color{orange}|\color{orange} y\color{orange})}{\color{orange} d \color{orange}x}$$
  • Gradient descent: $x_{t+1} = x_t + \tau \nabla_x \log p(x_t | y) $
  • Langevin algorithm: $x_{t+1} = x_t + \tau \nabla_x \log p(x_t | y) + \sqrt{2\tau} n_t$

• The score of the full posterior is simply: $$\nabla_x \log p(x |y) = \underbrace{\nabla_x \log p(y |x)}_{\mbox{known explicitly}} \quad + \quad \underbrace{\nabla_x \log p(x)}_{\mbox{known implicitly}}$$ $\Longrightarrow$ "all" we have to do is model/learn the score of the prior.

Neural Score Estimation by Denoising Score Matching (Vincent 2011)

• Denoising Score Matching: An optimal Gaussian denoiser learns the score of a given distribution.
  • If $x \sim \mathbb{P}$ is corrupted by additional Gaussian noise $u \in \mathcal{N}(0, \sigma^2)$ to yield $$x^\prime = x + u$$
  • Let's consider a denoiser $r_\theta$ trained under an $\ell_2$ loss: $$\mathcal{L}=\parallel x - r_\theta(x^\prime, \sigma) \parallel_2^2$$
  • The optimal denoiser $r_{\theta^\star}$ verifies: $$\boxed{\boldsymbol{r}_{\theta^\star}(\boldsymbol{x}', \sigma) = \boldsymbol{x}' + \sigma^2 \nabla_{\boldsymbol{x}} \log p_{\sigma^2}(\boldsymbol{x}')}$$

$\boldsymbol{x}'$

$\boldsymbol{x}$

$\boldsymbol{x}'- \boldsymbol{r}^\star(\boldsymbol{x}', \sigma)$

$\boldsymbol{r}^\star(\boldsymbol{x}', \sigma)$

Score-Based Generative Modeling Song et al. (2021)

• The SDE defines a marginal distribution $p_t(x)$ as the convolution of the target distribution $p(x)$ with a noise kernel $p_{t|s}(\cdot | x_s)$: $$p_t(x) = \int p(x_s) p_{t|s}(x | x_s) d x_s$$
• For a given forward SDE that evolves $p(x)$ to $p_T(x)$, there exists a reverse SDE that evolves $p_T(x)$ back into $p(x)$. It involves having access to the marginal score $\nabla_x \log_t p(x)$.

But... We don't actually know the marginal posterior score!

• We know the following quantities:
  • Annealed likelihood (analytically): $p_\sigma(y | x) = \mathcal{N}(y; \mathbf{A} x, \mathbf{\Sigma} + \sigma^2 \mathbf{I})$
  • Annealed prior score (by score matching): $\nabla_x \log p_\sigma(x)$
• But, unfortunately: $\boxed{p_\sigma(x|y) \neq p_\sigma(y|x) \ p_\sigma(x)}$ $\Longrightarrow$ We don't know the marginal posterior score!

• We cannot use the reverse SDE/ODE of diffusion models to sample from the posterior. $$\mathrm{d} x = [f(x, t) - g^2(t) \underbrace{\nabla_x \log p_t(x|y)}_{\mbox{unknown}} ] \mathrm{d}t + g(t) \mathrm{d} w$$

Proposed sampling strategy (Remy et al. 2020)

• Even if not equivalent to the marginal posterior score, $\nabla_x \log p_{\sigma^2}(y | x) + \nabla_x \log p_{\sigma^2}(x)$ still has good properties:
  • Tends to an isotropic Gaussian distribution for large $\sigma$
  • Corresponds to the target posterior for $\sigma=0$


• If we simulate this SDE sufficiently slowly (i.e. timescale of change of $\sigma$ is much larger than the timescale of the SDE) we can expect to sample from the target posterior.

Approximation methods for the marginal likelihood: DPS (Chung et al. 2023)

• Let's express the marginal likelihood in terms of x_0:
  \begin{align*} p_t(y | x_t) &= \int p(y | x_0) \, p(x_0 | x_t) \, d x_0 \\ &= \mathbb{E}_{x_0 \sim p(x_0 | x_t)}[p(y | x_0)] \end{align*}


• The DPS approximation:
  $$p(y | x_t) \simeq p(y | \hat{x}_0)$$ with $\hat{x}_0 = \mathbb{E}[x_0 | x_t] = x_t - \sigma_t^2 \nabla_{x_t} \log p_t(x_t)$


• In the case of a Gaussian likelihood and usual linear inverse problem: $$ \nabla_{x_t} \log p_t( y | x_t ) \simeq - \frac{1}{\sigma^2} \nabla_{x_t} \parallel y - \mathbf{A} \hat{x}_0 \parallel_2^2 $$

DPS in practice

Better methods: e.g. Moment Matching Posterior Sampling Rozet et al. 2024

Advertisement!

Francois Rozet @ University of Liege

					
from azula.denoise import PreconditionedDenoiser
from azula.noise import VPSchedule
from azula.sample import DDPMSampler

# Choose the variance preserving (VP) noise schedule
schedule = VPSchedule()

# Initialize a denoiser
denoiser = PreconditionedDenoiser(
	backbone=CustomNN(in_features=5, out_features=5),
	schedule=schedule,
)

# Generate 64 points in 1000 steps
sampler = DDPMSampler(denoiser.eval(), steps=1000)

x1 = sampler.init((64, 5))
x0 = sampler(x1)
					
					

https://github.com/probabilists/azula

Let's set the stage: Gravitational lensing

Galaxy shapes as estimators for gravitational shear

$$ e = \gamma + e_i \qquad \mbox{ with } \qquad e_i \sim \mathcal{N}(0, I)$$

• We are trying the measure the ellipticity $e$ of galaxies as an estimator for the gravitational shear $\gamma$

Gravitational Lensing as an Inverse Problem

Shear $\gamma$

Convergence $\kappa$

Example of one chain during Annealed Posterior Sampling

Validating Posterior Sampling under a Gaussian prior

Illustration on $\kappa$-TNG simulations

Remy, Lanusse, et al. (2022)

True convergence map

Traditional Kaiser-Squires

Wiener Filter

Posterior Mean (ours)

Posterior samples

Reconstruction of the HST/ACS COSMOS field

• COSMOS shear data from Schrabback et al. 2010
• Prior learned from $\kappa$-TNG simulation from Osato et al. 2021.

Massey et al. (2007)

Remy et al. (2022) Posterior mean

Remy et al. (2022) Posterior samples