• 1000 images each night, 15 TB/night for 10 years


• 18,000 square degrees, observed once every few days


• Tens of billions of objects, each one observed $\sim1000$ times

• Deep Belief Network
  (Hinton et al. 2006)


• Variational AutoEncoder
  (Kingma & Welling 2014)


• Generative Adversarial Network
  (Goodfellow et al. 2014)


• Wasserstein GAN
  (Arjovsky et al. 2017)


• Midjourney v5 Guided Latent Diffusion (2023)

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

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

• 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}}$$

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

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

• Most papers applying neural techniques for inference have used sub-optimal compression techniques, e.g. Mean Square Error $$ \mathcal{L} = || f_\varphi(x) - \theta ||_2^2 $$ $\Longrightarrow$ This is ok, contours will simply be larger than they could be.



• A lot of papers are still relying on assuming a proxy Gaussian likelihoods, i.e. estimating a mean and covariance from simulations.
  $\Longrightarrow$ This is dangerous, can lead to biased contours.

											import jax.numpy as np
											import jax_cosmo as jc
					
											# Defining a Cosmology
											cosmo = jc.Planck15()
					
											# Define a redshift distribution with smail_nz(a, b, z0)
											nz = jc.redshift.smail_nz(1., 2., 1.)
					
											# Build a lensing tracer with a single redshift bin
											probe = probes.WeakLensing([nz])
					
											# Compute angular Cls for some ell
											ell = np.logspace(0.1,3)
											cls = angular_cl(cosmo_jax, ell, [probe])
										

• Current cosmological MCMC chains take days, and typically require access to large computer clusters.


• Gradients of the log posterior are required for modern efficient and scalable inference techniques:
  • Variational Inference
  • Hamiltonian Monte-Carlo


• In jax-cosmo, we can trivially obtain exact gradients:

  
  															def log_posterior( theta ):
  																return gaussian_likelihood( theta ) + log_prior(theta)
  					
  															score = jax.grad(log_posterior)(theta)
  															



• On a DES Y1 analysis, we find convergence in 70,000 samples with vanilla HMC, 140,000 with Metropolis-Hastings

																							 import tensorflow as tf
																							 import flowpm
																							 # Defines integration steps
																							 stages = np.linspace(0.1, 1.0, 10, endpoint=True)
					
																							 initial_conds = flowpm.linear_field(32,       # size of the cube
																																100,       # Physical size
																																ipklin,    # Initial powerspectrum
																																batch_size=16)
					
																							 # Sample particles and displace them by LPT
																							 state = flowpm.lpt_init(initial_conds, a0=0.1)
					
																							 # Evolve particles down to z=0
																							 final_state = flowpm.nbody(state, stages, 32)
					
																							 # Retrieve final density field
																							 final_field = flowpm.cic_paint(tf.zeros_like(initial_conditions),
																															final_state[0])
																						 

• Seamless interfacing with deep learning components

• We developed a Mesh TensorFlow implementation that can scale on GPU clusters (horovod+NCCL).


• For a $2048^3$ simulation:
  • Distributed on 256 NVIDIA V100 GPUs


• Don't hesitate to reach out if you have a use case for model parallelism!
  
  
  https://github.com/DifferentiableUniverseInitiative/mesh


• Now developing the next generation of these tools in JAX
  • pmwd Differentiable PM library, (Li et al. 2022) arXiv:2211.09958
  • jaxdecomp: Domain Decomposition and Parallel FFTs

• Sampling from a discrete random distribution is classically not a differentiable operation
• Relaxed and reparameterized HOD sampling
  • Relaxed Bernoulli distributions for centrals
    $$ N_{\rm cen} = \frac{1}{1 + \exp( - (\log(\frac{p}{1 - p}) + \epsilon)/\tau) } \mbox{ with } \epsilon \sim \mathrm{Logistic}(0,1) \;. $$ where $\tau$ is a temperature parameter
  • Relaxed Binomial distribution for satelittes
    $$ N_{\rm sat} \sim \mathrm{Binomial}\left(N, p=\frac{ \left\langle N_{\rm sat}\right\rangle }{N} \right)$$