+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ How do we make sense of this vast amount
+
+
+
+
+
+
+
+
+
+
+ +
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ ![]()
+ From left to right: raw reconstruction, posterior sample, ground truth +
+
+
+
+
+
+
+
+
+
+
+
+
+
+ ![]()
+
+
+
+
+
+
+
+
+
+ +
+
+
+
+
+ +
+
+ +
+
+
+ ![]()
+
+
+
+
+![]()
+
+
+
+
+
+
+
+
+
+
+ +
+
+
+ + +
+
+
+
+
+
+
+
+ Bridging Physics & AI for Cosmology in the Exascale Era
+BASE Workshop - September 28th, 2020
++
+
+
+
+
+
+
+
+
+
+
+
+ François Lanusse
++
+ +
+
+
+
+
+
+
+
+ +
+
+
+
+
+
+ +
+
+
+ + +
+
+ the Rubin Observatory Legacy Survey of Space and Time
+
+
+
+
+
+
+
+ -
+
- 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 +
+
+
+
$\Longrightarrow$ Incredible potential for discovery, along with unprecedented challenges.
+
+
+
+
+ How do we make sense of this vast amount
of complex data?
+ A Motivating Example: Deconvolution
++
+
+
+
+
+ ![]()
+
Hubble Space Telescope +
+
+ 
+ Hubble Space Telescope +
+ ![]()
+
some deep neural network +
+
+ 
+ some deep neural network +
+ ![]()
+
Simulated Ground-Based Telescope +
+ 
+ Simulated Ground-Based Telescope +
+
+ +
+
+
+
+ The issue with generic black box deep learning inference
+
+
+
+
+ -
+
- No explicit control of noise, PSF, depth.
+
-
+
- Unless covered by the training data, the result becomes unpredictable. + +
+ - No guarantees some physical properties are preserved
+
$\Longrightarrow$ In this case, the flux of the deconvolved object +
+ - Robust quantification of uncertainties is extremely difficult. + +
+
+
+
A Physicist's approach: let's build a model
++
+
+
+
+
+ ![]()
+
+ 
+
+ ![]()
+
+
+
+ ![]()
+
+
+ 
+
+ ![]()
+
+
+ 
+
+ ![]()
+
+
+
+
+
+ $\longrightarrow$
$g_\theta$
+ $g_\theta$
$\longrightarrow$
PSF
+ PSF
$\longrightarrow$
Pixelation
+ Pixelation
$\longrightarrow$
Noise
+ Noise
+
+
+
+
+ ![]()
+ ![]()
+ ![]()
+ ![]()
+ ![]()
+
+ 
+ 
+
+ 
+ 
+
+
+
+
+
+ Probabilistic model
+
+
+
+
+
+
+
+
+ $$ x \sim ? $$
+ $$ x \sim \mathcal{N}(z, \Sigma) \quad z \sim ? $$
latent $z$ is a denoised galaxy image
+ latent $z$ is a denoised galaxy image
$$ x \sim \mathcal{N}( \mathbf{P} z, \Sigma) \quad z \sim ?$$
latent $z$ is a super-resolved and denoised galaxy image
+ latent $z$ is a super-resolved and denoised galaxy image
$$ x \sim \mathcal{N}( \mathbf{P} (\Pi \ast z), \Sigma) \quad z \sim ? $$
latent $z$ is a deconvolved, super-resolved, and denoised galaxy image
+ latent $z$ is a deconvolved, super-resolved, and denoised galaxy image
$$ x \sim \mathcal{N}( \mathbf{P} (\Pi \ast g_\theta(z)), \Sigma) \quad z \sim \mathcal{N}(0, \mathbf{I}) $$
latent $z$ is a Gaussian sample
$\theta$ are parameters of the model
+ latent $z$ is a Gaussian sample
$\theta$ are parameters of the model
+
+
+
$\Longrightarrow$ Decouples the morphology model from the observing conditions.
+ Bayesian Inference a.k.a. Uncertainty Quantification
+
+
+
+
+ ![]()
+
+ 
+
+ The Bayesian view of the problem:
+ $$ p(z | x ) \propto p_\theta(x | z, \Sigma, \mathbf{\Pi}) p(z)$$
+ where:
+
+
+ +
-
+
- $p( z | x )$ is the posterior +
- $p( x | z )$ is the data likelihood, contains the physics +
- $p( z )$ is the prior +
+
+
+
+
+
+
+ ![]()
+ ![]()
+
+
+
+
+
+ Data
$x_n$
+ $x_n$
Truth
$x_0$
+ $x_0$
+
+
+
+
+
+
+
+ Posterior samples
$g_\theta(z)$
+ $g_\theta(z)$
+
+
+
+
+ ![]()
+
+
+ 
+
+
$\mathbf{P} (\Pi \ast g_\theta(z))$
+
+ $\mathbf{P} (\Pi \ast g_\theta(z))$
Median
+
+
+
+
+ ![]()
+
+
+ 
+
+
+ Data residuals
$x_n - \mathbf{P} (\Pi \ast g_\theta(z))$
+ $x_n - \mathbf{P} (\Pi \ast g_\theta(z))$
Standard Deviation
+ $\Longrightarrow$ Uncertainties are fully captured by the posterior.
+ The Scarlet algorithm: deblending as an optimization problem
+Melchior et al. 2018; Lanusse et al. 2019
+
+ $$ \mathcal{L} = \frac{1}{2} \parallel \mathbf{\Sigma}^{-1/2} (\ Y - P \ast A S \ ) \parallel_2^2 - \sum_{i=1}^K \log p_{\theta}(S_i) + \sum_{i=1}^K g_i(A_i) + \sum_{i=1}^K f_i(S_i)$$
+
+
+
+
+ $\Longrightarrow$ Explicit physical modeling of the observed sky
+
+ ![]()
+
+
+ 
+
+
+ Where for a $K$ component blend:
+
+
+ +
-
+
- $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... +
+
+
+
+
Scarlet in action
+ +
+
+
+
+ Input blend
+
+
+
+ ![]()
+
+ 
+
+ Solution
+
+
+
+ ![]()
+ ![]()
+
+ 
+ 
+
+ Residuals
+
+
+ ![]()
+ ![]()
+
+ 
+ 
+ -
+
- Classic priors (monotonicity, symmetry). +
- Deep Morphology prior. +
+ +
+
+
+ True Galaxy
+ ![]()
+
+
+ 
+
+ Deep Morphology Prior Solution
+
+ ![]()
+
+
+
+ 
+
+
+ Monotonicity + Symmetry Solution
+ ![]()
+
+ 
+ Beyond astrophysics: MRI Reconstruction
+
+ From left to right: raw reconstruction, posterior sample, ground truth +
+
+
+
-
+
- Super preliminary... +
- Same idea of learning a neural prior from example data and
using + it in a standard Bayesian inference scheme.
+
+
Takeaway message
++
+
-
+
- Generative models are a data-driven way of complementing
+ physical models
+
+ $\Longrightarrow$ This makes use of the vast amount of data. +
+
+ - Most of these tasks can be split into independent jobs.
+
+ $\Longrightarrow$ Requires High Throughput more than High Performance Computing. +
+
+
+ +
+ How do we manage when the size of the inverse problems become large?
+Reconstructing the Initial Conditions of the Universe
+
+
Going back to simpler times...
+
+
+ $$\arg\max_z \ \log p(x_{dm} = f(z)) \ + \ p(z) $$
+ where:
+
+ +
-
+
- $f$ is an N-body cosmological simulaion. + +
- $z$ are the initial conditions (early universe) + +
- $x_{dm}$ is the present day dark matter distribution + +
introducing FlowPM: Particle-Mesh Simulations in TensorFlow
+ +
+
+
+ ![]()
+
+
+
+ 
+
+
+ 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 + +
- Gradients readily available + +
+
+ ![]()
+
+ 
+ MAP optimization in action
+ $$\arg\max_z \ \log p(x_{dm} = f(z)) \ + \ p(z) $$ +credit: C. Modi
+ +
+
+
+ ![]()
+
True initial conditions
$z_0$ +
+
+ 
+ True initial conditions
$z_0$ +
+ ![]()
+
Reconstructed initial conditions $z$ +
+
+ 
+ Reconstructed initial conditions $z$ +
+ ![]()
+
Reconstructed dark matter distribution $x = f(z)$ +
+
+ 
+ Reconstructed dark matter distribution $x = f(z)$ +
+ ![]()
+
Data
$x_{DM} = f(z_0)$ +
+ 
+ Data
$x_{DM} = f(z_0)$ +
+
+ +
+ Check out this blogpost for more details
+ https://blog.tensorflow.org/2020/03/simulating-universe-in-tensorflow.html +
+ + https://blog.tensorflow.org/2020/03/simulating-universe-in-tensorflow.html +
The practical challenges
+ ++
+ +
+
+
+
+
+
+ -
+
- Simulations of scientifically interesting sizes do not fit on a single GPU RAM.
+ $\Longrightarrow$ We need for 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 + +
+ +
+ +
+ ![]()
+
+ 
+ +
+ +
+ $\Longrightarrow$ We have started investigating Mesh TensorFlow at NERSC.
+
+ Mesh TensorFlow in a few words
+ +
+
(Gholami et al. 2018)
+ Mesh FlowPM
++
+
+ -
+
- Successfully ran on multiple Cori-GPU nodes, and on TPUv3-32 slice + + +
- In general, efficient model parallelism on GPU clusters is NOT a solved problem. + +
Conclusion
+Conclusion
+ ++ +
-
+
- Upcoming surveys will probably rely on machine learning for part
+ of their low level processing.
+
-
+
- Deep Learning can be embedded in physical models. +
- Most likely only requires High Throughput computing. +
+ +
+ - The full scientific analysis of data will most likely evolve towards hybrid + physical and ML simulations, which require HPC infrastructures and associated software frameworks. +
+
+
Thank you !
++ + +