Commentary of
Deep Learning to identify structures in 3D data

from Stella Offner at SCMA VII

Francois Lanusse (CNRS/CEA Saclay)

Back in my days...

My own "Deep Learning moment"

Lanusse et al. 2017

Metcalf et al. 2019

My main takeaways

  • U-net architecture
  • The challenges of 3D data

Ni et al. 2021

U-Nets
Everywhere

What is so special about this architecture?

Selected applications in Astrophysics

  • Semantic segmentation
     
  • Solvers for inverse problems (denoising, deconvolution)
     
  • Generative models

Semantic Segmentation: Morpheus
R. Hausen & B. Robertson (2020)

website, github, arXiv

Solving Inverse Problems: DeepMass
N. Jeffrey, F. Lanusse, O. Lahav, J.-L. Starck (2020)

 github, arXiv

\gamma = \mathbf{A} \kappa + n
\hat{\kappa} = \mathcal{F}_\theta(\gamma)
\mathcal{F}_{\theta^*} = \int \kappa p(\kappa | \gamma) d \kappa

Trained under an l2 loss:

Generative Models: Denoising Score Matching
Remy, Lanusse, Ramzi, Liu, Jeffrey, Starck (2020)

 github, arXiv

Li et al. 2019

See Benjamin's poster here

See Song & Ermon (2020)

Dealing with large scale
3D data

Excerpts from papers

Given the extra dimension in our model, we cannot maintain the same spatial resolution in our 3D model due to the limited memory on GPU. We reform the input data to an array shape of 64×64×32 in position position-velocity (PPV).

Application of Convolutional Neural Networks to Identify Stellar Feedback Bubbles in CO Emission, Xu et al. 2020

Limited by the size of GPU memory and the fact that 3D data consume more memory than lower dimensional tasks, we cannot feed the whole simulations into the GPU during training and testing.

AI-assisted super-resolution cosmological simulations, Li et al. 2020

My own struggle with 3D data

FlowPM, TensorFlow N-body solver
try me out here

CosmicRIM: Recurrent Inference Machine for initial condition reconstruction
(Modi, Lanusse, Seljak, Spergel, Perreault-Levasseur 2021)

Why is this still a problem today?

Pipeline parallelism

Towards a solution

#A simple 2 layer fully connected network
#y=Relu(xw+bias)v

#Define dimensions
batch = mtf.Dimension("batch", b)
io = mtf.Dimension("io", d_io)
hidden = mtf.Dimension("hidden", d_h)
# x.shape == [batch, io]
#Get tensors
w = mtf.get_variable("w", shape=[io, hidden])
bias = mtf.get_variable("bias", shape=[hidden])
v = mtf.get_variable("v", shape=[hidden, io])
#Define operations
h = mtf.relu(mtf.einsum(x, w, output_shape=[batch, hidden]) + bias)
y = mtf.einsum(h, v, output_shape=[batch, io])

 

 

 

 

 

 


 

Fresh from last week

  • 2048^3 simulation of the Universe, distributed on 256 GPUs

    => Runs in 3m
    => multiple TB of RAM

Main Takeaways

  • Deep Learning has revolutionized structure identification, Stella's talk is an excellent illustration of that.

     

  •  U-nets as an architecture are extremely versatile and successful.

     

  • Deep Learning with 3D data is still a challenge in 2021, this is because our needs in (astro-)physics are unique.