• 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

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

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


• Efficient training of parameters $\theta$ can be achieved by:
  • Adversarial Learning
  • Amortized Variational Inference

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

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



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

• Neural network parametrising a correction to the gravitational potential as a Fourier-based isotropic filter.

					def neural_nbody_ode(a, state):
				     pos = state[0]; vel = state[1]

					 	  #### Compute gravitational forces
				     delta = flowpm.cic_paint(tf.zeros([batch_size, nc, nc, nc]), pos)
				     delta_k = r2c3d(delta)
				     pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec)

						  # Neural potential correction
						  pot_k = pot_k * (1 + apply_model(kvec, a))
						  ######################################

				     forces = tf.stack([ flowpm.cic_readout(
				             c2r3d(pot_k *gradient_kernel(kvec, i)), pos)
				         for i in range(3) ], axis=-1)
				     forces = forces * 1.5 * cosmo.Omega_m
				  		#######################################

				     #Computes the update of position (drift)
				     dpos = 1. / (a**3 * flowpm.tfbackground.E(cosmo, a)) * vel

				     #Computes the update of velocity (kick)
				     dvel = 1. / (a**2 * flowpm.tfbackground.E(cosmo, a)) * forces

				   return tf.stack([dpos, dvel], axis=0)
							

• Introduce a Recurrent Neural Network (RNN) $h_\phi$, and state variable $s$, so that: $$ s_{i+1} = h^*_\phi( \nabla \log p(y|x_{i}), x_i, s_{i})$$ $$ x_{i+1} = x_i + h_\phi(\nabla \log p(y|x_{i}), x_i, s_{i+1})$$

• A few notable differences to a vanilla RIM:
  • We provide gradients of both prior and likelihood to the model.
  
  
  • Because our forward model couples scales, we use a multiscale U-Net architecture.
  
  
  • Input gradients are pre-scaled with the ADAM formula.