Events
Astronomy Colloquium (01.03/2023)： Differentiable Cosmological Simulation with Adjoint Method
Description
Abstract: Rapid advances in deep learning have brought not only myriad powerful neural networks, but also breakthroughs that benefit established scientific research. In particular, automatic differentiation (AD) tools and computational accelerators like GPUs have facilitated forward modeling of the Universe with differentiable simulations. Current differentiable cosmological simulations are limited by memory, thus are subject to a trade-off between time and space/mass resolution. They typically integrate for only tens of time steps, unlike the standard non-differentiable simulations. We present a new approach free of such constraints, using the adjoint method and reverse time integration. It enables larger and more accurate forward modeling, and will improve gradient based optimization and inference. We implement it in a particle-mesh (PM) N-body library pmwd (particle-mesh with derivatives). Based on the powerful AD system JAX, pmwd is fully differentiable, and is highly performant on GPUs.
Bio:
Yin Li is a research scientist at the Peng Cheng Laboratory in Shenzhen, China, working at the interface between machine learning, statistics, cosmology and astrophysics. Previously, he was a Flatiron Research Fellow with a joint appointment at the Center for Computational Astrophysics (CCA) and the Center for Computational Mathematics (CCM) in the Simons Foundation. Before that, he was a joint postdoctoral cosmologist at the Berkeley Center for Cosmological Physics (BCCP), University of California, Berkeley, and the Kavli Institute for the Physics and Mathematics of the Universe (IPMU), University of Tokyo. He did his Ph.D. and B.S. in physics, respectively at the University of Chicago and Peking University.
