[論文レビュー] CosmicRIM : Reconstructing Early Universe by Combining Differentiable Simulations with Recurrent Inference Machines

Chirag Modi, François Lanusse+3 more|arXiv (Cornell University)|Apr 26, 2021

Computational Physics and Python Applications参考文献 17被引用数 5

ひとこと要約

CosmicRIMは、微分可能N体シミュレーション（FlowPM）と学習された最適化を組み合わせた、微分可能で再帰的な推論マシン（RIM）フレームワークを提案する。このフレームワークは、宇宙論における初期原始密度場の再構成を実現し、ADAM や LBFGS に加えてアニーリングを行う手法と比較して40倍速い収束を達成するとともに、大スケールでは再構成品質を同等または上回る。これは、物理的アニーリングスケジュールを暗黙的に学習することで実現される。

ABSTRACT

Reconstructing the Gaussian initial conditions at the beginning of the Universe from the survey data in a forward modeling framework is a major challenge in cosmology. This requires solving a high dimensional inverse problem with an expensive, non-linear forward model: a cosmological N-body simulation. While intractable until recently, we propose to solve this inference problem using an automatically differentiable N-body solver, combined with a recurrent networks to learn the inference scheme and obtain the maximum-a-posteriori (MAP) estimate of the initial conditions of the Universe. We demonstrate using realistic cosmological observables that learnt inference is 40 times faster than traditional algorithms such as ADAM and LBFGS, which require specialized annealing schemes, and obtains solution of higher quality.

研究の動機と目的

• 宇宙論的調査データから原始密度場を再構成する高次元的・非線形的逆問題に対処すること。
• ADAM や LBFGS といった従来の最適化手法が、高価なN体シミュレーションを数百回も必要とする計算的に非現実的な問題を克服すること。
• 手動でのチューニングなしに、物理的アニーリングに類似したスケジュールを暗黙的に学習する効率的な推論スキームを学習すること。
• 次世代の宇宙論的調査で繰り返し再構成が必要となるような、高速でスケーラブルな推論を可能にすること。

提案手法

• 前方モデルの初期条件に関して勾配を計算するために、微分可能N体シミュレータ（FlowPM）を用いる。
• 隠れ状態、現在の再構成、および対数後確率の勾配を組み合わせた再帰的更新ルールを持つ再帰的推論マシン（RIM）を採用する。
• RIMは対数後確率の負の対数を最小化するように訓練され、物理的アニーリングを模倣する暗黙の最適化経路を学習する。
• メモリ状態と再構成および隠れ状態のための別々の更新関数を含むアーキテクチャにより、反復的で改善可能な再構成が可能になる。
• 初期条件にガウス事前分布を、尤度モデルには加法的ガウスノイズを用いて後確率を定義する。
• 固定されたフィducialパワースペクトルを用いたシミュレーテッドデータ上で訓練することで、推論ネットワークのエンドツーエンド最適化を可能にする。

実験結果

リサーチクエスチョン

• RQ1 学習された最適化スキームは、ADAM や LBFGS といった従来の勾配ベース手法よりも、原始密度場の再構成において優れているか？
• RQ2 RIMフレームワークは、明示的な設計なしに、物理的アニーリングに類似したスケジュールを暗黙的に学習しているか？
• RQ3 収束速度と再構成精度の観点から、RIMの性能は、専用のアニーリングスキームと比べてどうか？
• RQ4 RIMフレームワークは、異なる初期条件や調査設定に一般化して訓練可能か？

主な発見

• CosmicRIMは、ADAM や LBFGS にアニーリングを施した手法と同等またはそれ以上の再構成品質を達成し、特に大スケールでのトランスファ関数の向上が見られた。
• この手法はたった10回の推論ステップで十分であり、アニーリング付きADAM や LBFGS が必要とする400ステップと比較して40倍速い。
• 標準的な再構成で初期化された場合でも、CosmicRIMはさらなる性能向上を示しており、初期化に頼らない強靭性を示している。
• RIMは、大スケール構造を最初に再構成する最適化経路を暗黙的に学習しており、物理的アニーリングスキームの挙動を模倣している。
• RIMの最終出力を出発点としてADAM や LBFGS でさらなる最適化を試みても改善が得られず、RIMが高品質なMAP推定値に到達していることが示された。
• このフレームワークはスケーラブルな推論を可能にし、メモリコストを削減するためのサブネットワークの逐次訓練による拡張も可能である。