CosmicRIM : Reconstructing Early Universe by Combining Differentiable Simulations with Recurrent Inference Machines
A side effect of the Deep Learning revolution has been the advent of general purpose GPU-accelerated automatic differentiation frameworks, such as TensorFlow. This opens the door to a range of applications of the generic
concept of automatically differentiable physics. In this talk, I will first present the FlowPM tool, a TensorFlow re-implementation of the well-known FastPM N-body solver, leading to distributed, GPU accelerated, and automatically differentiable cosmological simulations.One of the direct applications of such a framework is the reconstruction of Gaussian initial conditions from observed survey data as a Maximum A Posteriori (MAP) solution obtained by gradient-descent. However, optimization through the cosmological forward model is typically difficult as the objective is non-convex, and has up until now typically required combining ad-hoc annealing schemes with traditional optimizers.
I will present our solution to this difficult optimization problem, building on a "learning-to-learn" approach developed in the machine learning community. We combine our automatically differentiable N-body solver, with a recurrent networks to learn the inference scheme and obtain the MAP estimate of the initial conditions of the Universe. Learnt inference is 40 times faster than traditional algorithms such as ADAM and LBFGS and obtains solution of higher quality.