[論文レビュー] CONNECT: A neural network based framework for emulating cosmological observables and cosmological parameter inference

Andreas Nygaard, Emil Brinch Holm+2 more|arXiv (Cornell University)|May 30, 2022

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

ひとこと要約

CONNECT は、CLASS ボルツマンソルバーからの宇宙論的観測量をエミュレートする PyTorch ベースのニューラルネットワークフレームワークであり、高速かつ高精度な宇宙論的パラメータ推論を可能にする。反復的に最適化されたラテンハイパーキューブサンプリングを用いて効率的なトレーニングデータを生成することで、CLASS の評価回数を2桁削減し、MCMC推論を桁単位で高速化する。また、完全な CLASS 計算と比較して、事後分布の差は 0.1 標準偏差未満にとどまる。

ABSTRACT

Bayesian parameter inference is an essential tool in modern cosmology, and typically requires the calculation of $10^5$--$10^6$ theoretical models for each inference of model parameters for a given dataset combination. Computing these models by solving the linearised Einstein-Boltzmann system usually takes tens of CPU core-seconds per model, making the entire process very computationally expensive. In this paper we present extsc{connect}, a neural network framework emulating extsc{class} computations as an easy-to-use plug-in for the popular sampler extsc{MontePython}. extsc{connect} uses an iteratively trained neural network which emulates the observables usually computed by extsc{class}. The training data is generated using extsc{class}, but using a novel algorithm for generating favourable points in parameter space for training data, the required number of extsc{class}-evaluations can be reduced by two orders of magnitude compared to a traditional inference run. Once extsc{connect} has been trained for a given model, no additional training is required for different dataset combinations, making extsc{connect} many orders of magnitude faster than extsc{class} (and making the inference process entirely dominated by the speed of the likelihood calculation). For the models investigated in this paper we find that cosmological parameter inference run with extsc{connect} produces posteriors which differ from the posteriors derived using extsc{class} by typically less than $0.01$--$0.1$ standard deviations for all parameters. We also stress that the training data can be produced in parallel, making efficient use of all available compute resources. The extsc{connect} code is publicly available for download at \url{https://github.com/AarhusCosmology}.

研究の動機と目的

• 宇宙論におけるベイズ的パラメータ推論の高い計算コストに対処すること。これは通常、アインシュタイン＝ボルツマンソルバーの 10^5–10^6 回の評価を必要とする。
• 新しい反復的サンプリング戦略の開発により、ニューラルエミュレータのトレーニングに必要な高価な CLASS 評価回数を削減すること。
• 新しいデータセットに対して再トレーニングなしで、訓練済みのニューラルネットワークを用いて高速かつ高精度なパラメータ推論を可能にする、即席型のフレームワークを構築すること。
• 標準的な MCMC サンプラー（例：MontePython や Cobaya）と互換性を持たせ、既存の宇宙論ワークフローへの広範な適用を可能にすること。
• カスタム宇宙論的出力と GPU 加速の可能性を備えた、完全に再現可能でオープンソースのソリューションを提供すること。

提案手法

• CLASS が計算する宇宙論的観測量（例：CMB スペクトル、物質パワー スペクトル）をエミュレートするために、完全に接続されたフィードフォワードニューラルネットワークを用いる。
• 高品質なトレーニングデータを生成するために、反復的ラテンハイパーキューブサンプリング戦略を採用し、必要な CLASS 評価回数を2桁削減する。
• 宇宙論的パラメータとそれに対応する CLASS の出力のデータセット上でニューラルネットワークをトレーニングし、入力構成ファイルを介して出力変数をカスタマイズ可能にしている。
• CLASS の呼び出しを高速なニューラルネットワーク推論に置き換えるプラグインを介して、MontePython や Cobaya とシームレスに統合する。
• モジュラーなコード構造により、ユーザー定義の損失関数やネットワークアーキテクチャのサポートが可能で、特定の宇宙論的モデルに合わせたカスタマイズが可能である。
• 並列処理によりトレーニングデータを生成し、高性能コンピューティングリソースを効率的に活用する。

実験結果

リサーチクエスチョン

• RQ1 少ない CLASS 評価回数でトレーニングされたニューラルネットワークエミュレータが、宇宙論的観測量を高精度にエミュレートできるか？
• RQ2 反復的サンプリング戦略は、エミュレータの精度を維持しつつ、必要な CLASS 評価回数をどの程度削減できるか？
• RQ3 エミュレータを用いて得られる宇宙論的パラメータの事後分布は、完全な CLASS 計算と比較してどの程度正確か？
• RQ4 再トレーニングなしで、異なるデータセットの組み合わせに対してエミュレータを再利用可能か？また、これにより推論速度と精度にどのような影響があるか？
• RQ5 MontePython などの既存の MCMC サンプリングパイプラインに、最小限のコード変更でフレームワークを統合できるか？

主な発見

• CONNECT は、標準的な推論実行と比較して、トレーニングに必要な CLASS 評価回数を2桁削減し、データ生成コストを顕著に低減する。
• CONNECT を用いた宇宙論的パラメータ推論では、すべてのパラメータについて、完全な CLASS 計算と比較して、事後分布の差が 0.01–0.1 標準偏差未満にとどまる。
• フレームワークにより、MCMC 推論が数時間で完了するようになり、数日から数週間かかっていた従来の方法と比較して著しく高速化される。推論時間はモデル計算ではなく、尤度評価に支配されるようになる。
• 訓練済みのエミュレータは、再トレーニングなしで、さまざまなデータセットの組み合わせに完全に再利用可能であり、新しいデータが追加された際の迅速な再分析を可能にする。
• フレームワークは MontePython や Cobaya と互換性があり、プラグインインターフェースは元のコードベースの変更を一切必要としない。
• トレーニングデータは並列で生成可能であり、利用可能な HPC リソースを最適に活用し、スケーラブルなデータ生成を実現する。