Differentiable Cosmological Simulation with Adjoint Method - IRFU-AIM
Article Dans Une Revue The Astrophysical Journal Supplement Année : 2024

Differentiable Cosmological Simulation with Adjoint Method

Yin Li
  • Fonction : Auteur
Chirag Modi
  • Fonction : Auteur
Drew Jamieson
  • Fonction : Auteur
Yucheng Zhang
  • Fonction : Auteur
Libin Lu
  • Fonction : Auteur
Yu Feng
  • Fonction : Auteur
Leslie Greengard
  • Fonction : Auteur

Résumé

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.
Fichier principal
Vignette du fichier
Li_2024_ApJS_270_36.pdf (3.41 Mo) Télécharger le fichier
Origine Fichiers éditeurs autorisés sur une archive ouverte

Dates et versions

hal-03892303 , version 1 (20-04-2024)

Licence

Identifiants

Citer

Yin Li, Chirag Modi, Drew Jamieson, Yucheng Zhang, Libin Lu, et al.. Differentiable Cosmological Simulation with Adjoint Method. The Astrophysical Journal Supplement, 2024, 270 (2), pp.36. ⟨10.3847/1538-4365/ad0ce7⟩. ⟨hal-03892303⟩
215 Consultations
24 Téléchargements

Altmetric

Partager

More