[論文レビュー] pmwd: A Differentiable Cosmological Particle-Mesh $N$-body Library

Yin Li, Libin Lu+8 more|arXiv (Cornell University)|Nov 18, 2022

Galaxies: Formation, Evolution, Phenomena被引用数 14

ひとこと要約

pmwd は、随伴法と JAX の自動微分を活用して、メモリ効率が良く、GPU で加速された大規模構造形成の微分可能（differentiable）な粒子・メッシュ $N$-体シミュレーションライブラリです。時間反転統合によりシミュレーション履歴を再構築することで、従来の時間遡及バックプロパゲーション（backpropagation-through-time）手法と比較して、顕著に低いメモリコストで完全な微分可能性を達成しています。

ABSTRACT

The formation of the large-scale structure, the evolution and distribution of galaxies, quasars, and dark matter on cosmological scales, requires numerical simulations. Differentiable simulations provide gradients of the cosmological parameters, that can accelerate the extraction of physical information from statistical analyses of observational data. The deep learning revolution has brought not only myriad powerful neural networks, but also breakthroughs including automatic differentiation (AD) tools and computational accelerators like GPUs, facilitating forward modeling of the Universe with differentiable simulations. Because AD needs to save the whole forward evolution history to backpropagate gradients, current differentiable cosmological simulations are limited by memory. Using the adjoint method, with reverse time integration to reconstruct the evolution history, we develop a differentiable cosmological particle-mesh (PM) simulation library pmwd (particle-mesh with derivatives) with a low memory cost. Based on the powerful AD library JAX, pmwd is fully differentiable, and is highly performant on GPUs.

研究の動機と目的

• 宇宙論的パrameter制約のための勾配ベース推論をサポートする微分可能な宇宙論的シミュレーションライブラリの開発を目的とする。
• 従来の時間遡及バックプロパゲーションにおけるメモリボトルネックを克服するため、随伴法を用いる。
• フィールドレベルの前向きモデリングに適した、高解像度で GPU で加速された完全な微分可能性を備えたシミュレーションを実現する。
• 深層学習や半解析的モデルと統合し、宇宙論的パrameterと初期条件の統合推論を可能にする。
• 次世代の宇宙論的データ解析を支援するスケーラブルでパフォーマンスに優れ、オープンソースなツールを提供する。

提案手法

• 時間逆転統合により勾配を計算し、前向きな進化履歴の全保存を回避することで、随伴法を用いる。
• JAX を用いて自動微分と GPU 加速を実現し、高性能な微分可能シミュレーションを可能にする。
• 「離散化してから最適化する（discretize-then-optimize）」アプローチを適用し、後向き勾配が前向き統合と同じ離散的経路に従うことを保証する。
• 時間逆転統合が可能な粒子・メッシュ（PM）重力ソルバーを実装し、一貫性のある随伴計算を実現する。
• 設計に C2 対称性を活用し、時間反転性を反映するとともに、メモリ効率の良い勾配計算を支援する。
• エミュレーターや時空間最適化器などの微分可能コンポONENTとの統合をサポートする。

実験結果

リサーチクエスチョン

• RQ1 随伴法は、宇宙論的 $N$-体シミュレーションに効果的に適用可能であり、メモリコストを削減しながら微分可能性を維持できるか？
• RQ2 pmwd は、全前向き状態を保存する従来の微分可能シミュレーションコードと比較して、どの程度メモリ効率が優れているか？
• RQ3 pmwd は、微分可能宇宙論的シミュレーションにおいて、より高い時間分解能および質量分解能にどの程度スケーリングできるか？
• RQ4 pmwd は、初期条件と物理的パrameterを含む、宇宙論的観測量のエンドツーエンド微分可能モデリングを可能にするか？
• RQ5 pmwd は、統合推論パイプライン内で深層学習や他の微分可能コンポONENTとどの程度良好に統合されるか？

主な発見

• pmwd は、全前向き状態を保存するのではなく随伴法を用いることで、顕著に低いメモリ使用量で完全な微分可能性を達成している。
• ライブラリは JAX と完全に互換性があり、GPU でも効率的に実行され、高性能なシミュレーションを可能にする。
• 随伴法により、従来の AD ベースのシミュレーションと比較して、より多くの時間ステップと高解像度にスケーリング可能である。
• pmwd はフィールドレベルの前向きモデリングをサポートしており、宇宙論的パrameterと初期条件の統合推論が可能である。
• ライブラリは微分可能な深層学習および半解析的モデルと相互運用可能であり、複雑な推論パイプラインの構築を可能にする。
• pmwd はオープンソースであり、公開アーカイブも提供されており、再現性とコミュニティによる拡張を支援する。