a new era of wide-field galaxy surveys

One of the core probes: 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$

Illustration on the Dark Energy Survey (DES) Y3

The DES Y3 pipeline

the Rubin Observatory Legacy Survey of Space and Time

Previous generation survey: SDSS

Current generation survey: DES

LSST precursor survey: HSC

The challenges of modern surveys

A Blessing

• Unprecedented statistical power


• Great potential for new discoveries

A Curse

• Existing methods are reaching their limits at every step of the science analysis


• Control of systematics is paramount

LSST forecast on dark energy parameters

A few illustrations

Bosch et al. 2017

Jeffrey, Lanusse, et al. 2020

Cheng et al. 2020

• Galaxies are no longer blobs.
• Signals are no longer Gaussian.
• Cosmological likelihoods are no longer tractable.

The Deep Learning Explosion in Astrophysics

Review of the impact of Deep Learning in Galaxy Survey Science

Selected examples of Deep Learning in Astrophysics

• Classifying objects
  
  Galaxy Zoo: Probabilistic Morphology through Bayesian CNNs and Active Learning,Walmsley et al. 2019
• Inferring object properties
  
  Photometric redshifts from SDSS images using a Convolutional Neural Network, Pasquet et al. 2018
• Self-supervised Representation Learning
  
  Self-supervised similarity search for large scientific datasets, Stein et al. 2021
• Solve inverse problems
  
  Probabilistic Mass Mapping with Neural Score Estimation, Remy et al. 2022
• Emulate costly physical models
  
  SPECULATOR: Emulating stellar population synthesis for fast and accurate galaxy spectra and photometry, Alsing et al. 2020
• Accelerating numerical simulations
  
  AI-assisted super-resolution cosmological simulations, Li et al. 2021
• Enabling Simulation-Based Inference
  
  Likelihood-free inference with neural compression of DES SV weak lensing map statistics, Jeffrey et. al 2021

What is a neural network?

Simplest architecture: Multilayer Perceptron (MLP)

• Series of Dense a.k.a Fully connected layers: $$ h = \sigma(W x + b)$$ where:
  • $\sigma$ is the activation function (e.g. ReLU, Sigmoid, etc.)
  • $W$ is a multiplicative weight matrix
  • $b$ is an additive bias parameter


• This defines a parametric non-linear function $f_\theta(x)$


• MLPs are universal function approximators
  Nota bene: only asymptotically true!

How do you use it to approximate functions?

• Assume a loss function that should be small for good approximations on a training set of data points $(x_i, y_i)$ $$ \mathcal{L} = \sum_{i} ( y_i - f_\theta(x_i))^2 $$


• Optimize the parameters $\theta$ to minimize the loss function by gradient descent $$ \theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L} $$

Different neural architectures for different types of data

Performance can be improved for particular types of data by making use of inductive biases

$$ h = \sigma(W \ast x + b)$$

• Image -> Convolutional Networks
  • Convolutional layers are translation invariant
  • Pixel topology is incorporated by construction


• Time Series -> Recurrent Networks
  • Temporal information and variable length is incorporated by construction


• Graph-stuctured data -> Graph Neural Networks
  • Graph topology is incorporated by construction
  • Can be seen as generalisation of convolutional networks

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 and this notebook to learn how to do that:

Let's take a step back and think as Bayesians

Data $y$

Posterior samples $p(x | y)$

True $x$

• This deconvolution problem is an ill-posed inverse problem $y = A x + n$.
  $\Longrightarrow$There is no unique solution, the true answer can never be recovered.

The limits of traditional cosmological inference

HSC cosmic shear power spectrum

HSC Y1 constraints on $(S_8, \Omega_m)$

(Hikage et al. 2018)

• Measure the ellipticity $\epsilon = \epsilon_i + \gamma$ of all galaxies
  $\Longrightarrow$ Noisy tracer of the weak lensing shear $\gamma$


• Compute summary statistics based on 2pt functions,
  e.g. the power spectrum


• Run an MCMC to recover a posterior on model parameters, using an analytic likelihood $$ p(\theta | x ) \propto \underbrace{p(x | \theta)}_{\mathrm{likelihood}} \ \underbrace{p(\theta)}_{\mathrm{prior}}$$

Full-Field Simulation-Based Inference

• Instead of trying to analytically evaluate the likelihood of sub-optimal summary statistics, let us build a forward model of the full observables.
  $\Longrightarrow$ The simulator becomes the physical model.


• Each component of the model is now tractable, but at the cost of a large number of latent variables.

Benefits of a forward modeling approach

• Fully exploits the information content of the data (aka "full field inference").


• Easy to incorporate systematic effects.


• Easy to combine multiple cosmological probes by joint simulations.

...so why is this not mainstream?

The Challenge of Simulation-Based Inference

$$ p(x|\theta) = \int p(x, z | \theta) dz = \int p(x | z, \theta) p(z | \theta) dz $$ Where $z$ are stochastic latent variables of the simulator.

$\Longrightarrow$ This marginal likelihood is intractable!

Black-box Simulators Define Implicit Distributions

• A black-box simulator defines $p(x | \theta)$ as an implicit distribution, you can sample from it but you cannot evaluate it.
• Key Idea: Use a parametric distribution model $\mathbb{P}_\varphi$ to approximate the implicit distribution $\mathbb{P}$.

True $\mathbb{P}$

Samples $x_i \sim \mathbb{P}$

Model $\mathbb{P}_\varphi$

Deep Learning Approaches to Implicit Inference

Automated Summary Statistics Extraction

• Introduce a parametric function $f_\varphi$ to reduce the dimensionality of the data while preserving information.

Information-based loss functions

• Summary statistics $y$ is sufficient for $\theta$ if $$ I(Y; \Theta) = I(X; \Theta) \Leftrightarrow p(\theta | x ) = p(\theta | y) $$
• Variational Mutual Information Maximization $$ \mathcal{L} \ = \ \mathbb{E}_{x, \theta} [ \log q_\phi(\theta | f_\varphi(x)) ] \leq I(Y; \Theta) $$
  Jeffrey, Alsing, Lanusse (2021)

Unrolling the Probabilistic Learning Recipe

• I assume a forward model of the observations: \begin{equation} p( x ) = p(x | \theta) \ p(\theta) \nonumber \end{equation} All I ask is the ability to sample from the model, to obtain $\mathcal{D} = \{x_i, \theta_i \}_{i\in \mathbb{N}}$


• I am going to assume $q_\phi(\theta | x)$ a parametric conditional density


• Optimize the parameters $\phi$ of $q_{\phi}$ according to \begin{equation} \min\limits_{\phi} \sum\limits_{i} - \log q_{\phi}(\theta_i | x_i) \nonumber \end{equation} In the limit of large number of samples and sufficient flexibility \begin{equation} \boxed{q_{\phi^\ast}(\theta | x) \approx p(\theta | x)} \nonumber \end{equation}

Illustration on log-normal lensing simulations

DifferentiableUniverseInitiative/sbi_lens
JAX-based log-normal lensing simulation package

• 10x10 deg$^2$ maps at LSST Y10 quality, conditioning the log-normal shift parameter on $(\Omega_m, \sigma_8, w_0)$


• Infer full-field posterior on cosmology:
  • explicitly using an Hamiltonian-Monte Carlo (NUTS) sampler
  • implicitly using a learned summary statistics and conditional density estimation.

Example of application: Constraining Dark Matter Substructures

Brehmer, Mishra-Sharma, Hermans, Louppe, Cranmer (2019)

Example of application: Infering Microlensing Event Parameters

Zhang, Bloom, Gaudi, Lanusse, Lam, Lu (2021)

Example of application: Likelihood-Free parameter inference with DES SV

Jeffrey, Alsing, Lanusse (2021)

Suite of N-body + raytracing simulations: $\mathcal{D}$

What is generative modeling?

• The goal of generative modeling is to learn an implicit distribution $\mathbb{P}$ from which the training set $X = \{x_0, x_1, \ldots, x_n \}$ is drawn.


• Usually, this means building a parametric model $\mathbb{P}_\theta$ that tries to be close to $\mathbb{P}$.

True $\mathbb{P}$

Samples $x_i \sim \mathbb{P}$

Model $\mathbb{P}_\theta$

• Once trained, you can typically sample from $\mathbb{P}_\theta$ and/or evaluate the likelihood $p_\theta(x)$.

Why isn't it easy?

• The curse of dimensionality put all points far apart in high dimension

Distance between pairs of points drawn from a Gaussian distribution.

• Classical methods for estimating probability densities, i.e. Kernel Density Estimation (KDE) start to fail in high dimension because of all the gaps

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

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 Scarlet algorithm: deblending as an optimization problem

Where for a $K$ component blend:

• $P$ is the convolution with the instrumental response


• $A_i$ are channel-wise galaxy SEDs, $S_i$ are the morphology models


• $\mathbf{\Sigma}$ is the noise covariance


• $\log p_\theta$ is a PixelCNN prior


• $f_i$ and $g_i$ are arbitrary additional non-smooth consraints, e.g. positivity, monotonicity...

Training the morphology prior

Postage stamps of isolated COSMOS galaxies used for training, at WFIRST resolution and fixed fiducial PSF

isolated galaxy $\log p_\theta(x) = 3293.7$

artificial blend $\log p_\theta(x) = 3100.5 $

Scarlet in action

Input blend

Solution

Residuals

• Classic priors (monotonicity, symmetry).


• Deep Morphology prior.

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$

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)$

Example of one chain during annealing

Validating Posterior Sampling under a Gaussian prior

Some details: 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

• 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.

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

A Physicist's approach: 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.

The famous Variational Auto-Encoder

Sampling from the model

Woups... what's going on?

Tradeoff between code regularization and image quality

Latent space modeling with Normalizing Flows

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) $$

Let's again think as physicists

Probabilistic model

$$ x \sim \mathcal{N}(\Pi \ast (g_\theta(z) \otimes \gamma), \Sigma) \quad z \sim \mathcal{N}(0, \mathbf{I}) $$
latent $z$ are morphological parameters
$\theta$ are global parameters of the model
$\gamma$ are shear parameters

Joint inference using a parametric model for the morphology

Let's assume that $g(z)$ is a sersic model, i.e. $z = \{n, r_\text{hlr}, F, e_1, e_2, s_x, s_y\}$ and $$g(z) = F \times I_0 \exp \left( -b_n \left[\left( \frac{r}{r_\text{hlr}}\right)^{\frac{1}{n}} -1\right] \right)$$

The joint inference of $p(z, \gamma | \mathcal{D})$ leads to a biased posterior...

Marginal shear posterior $p(\gamma|\mathcal{D})$

Maximum a posteriori fit and residuals

We need a more realistic model of galaxy morphology

Joint inference using a generative model for the morpholgy

Remy, Lanusse, Starck (2022)

Let's use a learned $g_\theta(z)$

The joint inference of $p(z, \gamma | \mathcal{D})$ leads to an unbiased posterior!

Marginal shear posterior $p(\gamma|\mathcal{D})$

Maximum a posteriori fit and residuals