• 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

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

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

• 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