• 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

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

• We use a very standard residual UNet, and we adopt a residual score matching loss: $$ \mathcal{L}_{DSM} = \underset{\boldsymbol{x} \sim P}{\mathbb{E}} \underset{\begin{subarray}{c} \boldsymbol{u} \sim \mathcal{N}(0, I) \\ \sigma_s \sim \mathcal{N}(0, s^2) \end{subarray}}{\mathbb{E}} \parallel \boldsymbol{u} + \sigma_s \boldsymbol{r}_{\theta}(\boldsymbol{x} + \sigma_s \boldsymbol{u}, \sigma_s) \parallel_2^2$$ $\Longrightarrow$ direct estimator of the score $\nabla \log p_\sigma(x)$


• Lipschitz regularization to improve robustness:
  
  
  Without regularization

  With regularization

• Mutual information between $X$ and $Y$:

  “"amount of information" obtained about one random variable through observing the other random variable”



• Given a parametric summarizing function $y = f_\phi(\kappa(\theta))$ optimizing $f_\phi$ can be done by maximizing $I(y, \theta)$.


• In practice, $f_\phi$ is a CNN, trained to maximize a variational lower bound on the mutual information: $$ I(y ; \theta) \ \ge \ \mathbb{E}_{y, \theta} [ \log q_\phi(\theta | y) ] + H(\Theta) $$

• Deep Residual Network $y = f_\phi(x)$ followed by neural density estimator $q_\phi(\theta | y)$


• Training on weak lensing maps simulated for different cosmologies






• Training by Variational Mutual Information Maximization: $$\mathbb{E}_{(x, \theta) \in \mathcal{D}} [ \log q_\phi(\theta | f_\phi(y) ) ]$$

      										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

• Conventional strategy: Use a photoz code to estimate redshifts, then bin galaxies based on their photoz.


• Strategy with Differentiable Physics:
  • Introduce a parametric bin assignement function $f_\theta(x_{phot})$
  • Optimize $\theta$ by back-propagating through the challenge metrics.

      													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
• Mesh TensorFlow implementation for distribution on supercomputers

• 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
  • Runtime: 3 mins



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