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

• Training set: GalSim COSMOS HST/ACS postage stamps
  
  • 80,000 deblended galaxies from I < 25.2 sample
  • Drawn on 128x128 stamps at 0.03 arcsec resolution
  • Each stamp comes with:
    • PSF
    • Noise power spectrum
    • Bulge+Disk parametric fit


• Auto-Encoder model:
  • Deep residual autoencoder:
    7 stages of 2 resnet blocs each
  • Dense bottleneck of size 16.
  • Outputs positive, noiseless, deconvolved, galaxy surface brightness.


• Generative models are conditioned on:
  mag_auto, flux_radius, zphot

• Hybrid Physical-Deep Learning Model for Astronomical Inverse Problems
  F. Lanusse, P. Melchior, F. Moolekamp
  

• 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, let us build a forward model of the observables.


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

A non-exhaustive list of methods:
• Hamiltonian Monte-Carlo
• Variational Inference
• MAP+Laplace
• Gold Mining
• Dimensionality reduction by Fisher-Information Maximization

• The numerical scheme:
  
  
  • Estimate the density of particles on a mesh
    => compute gravitational forces by FFT
  
  
  • Interpolate forces at particle positions
  
  
  • Update particle velocity and positions, and iterate


• Fast and simple, at the cost of approximating short range interactions.

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

										with tf.Session() as sess:
										sim = sess.run(final_field)
									

• Seamless interfacing with deep learning components
• Mesh TensorFlow implementation for distribution on supercomputers

• Simulations of scientifically interesting sizes do not fit on a single GPU RAM
  e.g. $128^3$ operational, need $1024^3$ for survey volumes
  $\Longrightarrow$ We need a distributed Machine Learning Framework
  


• Most common form of distribution is data-parallelism $\Longrightarrow$ Reached Exascale on scientific deep learning applications


• What we need is model-parallelism on HPC environments

• Fisher matrices are notoriously unstable
  $\Longrightarrow$ they rely on evaluating gradients by finite differences.



• They do not scale well to large number of parameters.



• The main reason for using Fisher matrices is that full MCMCs are extremely slow, and get slower with dimensionality.

• 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