Differentiable Cosmological Simulation with Adjoint Method

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Li, Yin, Modi, Chirag, Jamieson, Drew, Zhang, Yucheng, Lu, Libin, Feng, Yu, Lanusse, François, Greengard, Leslie
Format: Preprint
Veröffentlicht: 2022
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866917587815759872
author Li, Yin
Modi, Chirag
Jamieson, Drew
Zhang, Yucheng
Lu, Libin
Feng, Yu
Lanusse, François
Greengard, Leslie
author_facet Li, Yin
Modi, Chirag
Jamieson, Drew
Zhang, Yucheng
Lu, Libin
Feng, Yu
Lanusse, François
Greengard, Leslie
contents 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. Based on analytic or automatic backpropagation, current differentiable cosmological simulations are limited by memory, and thus are subject to a trade-off between time and space/mass resolution, usually sacrificing both. 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 at the field level, and will improve gradient based optimization and inference. We implement it in an open-source 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.
format Preprint
id arxiv_https___arxiv_org_abs_2211_09815
institution arXiv
publishDate 2022
record_format arxiv
spellingShingle Differentiable Cosmological Simulation with Adjoint Method
Li, Yin
Modi, Chirag
Jamieson, Drew
Zhang, Yucheng
Lu, Libin
Feng, Yu
Lanusse, François
Greengard, Leslie
Instrumentation and Methods for Astrophysics
Cosmology and Nongalactic Astrophysics
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. Based on analytic or automatic backpropagation, current differentiable cosmological simulations are limited by memory, and thus are subject to a trade-off between time and space/mass resolution, usually sacrificing both. 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 at the field level, and will improve gradient based optimization and inference. We implement it in an open-source 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.
title Differentiable Cosmological Simulation with Adjoint Method
topic Instrumentation and Methods for Astrophysics
Cosmology and Nongalactic Astrophysics
url https://arxiv.org/abs/2211.09815