[論文レビュー] Accelerating HEP simulations with Neural Importance Sampling

Nicolas Deutschmann, Niklas Götz|arXiv (Cornell University)|Jan 17, 2024

Nuclear reactor physics and engineering被引用数 6

ひとこと要約

本論文は、勾配ステップ間でサンプルを再利用してデータ効率を向上させるPyTorchベースのニューラル重要度サンプリングライブラリZüNISを提案し、VEGASと比較検証した。HEP様の積分のいくつかで顕著な分散削減とウェイト付けの改善を示す一方で、訓練コストは高い。

ABSTRACT

Many high-energy-physics (HEP) simulations for the LHC rely on Monte Carlo using importance sampling by means of the VEGAS algorithm. However, complex high-precision calculations have become a challenge for the standard toolbox, as this approach suffers from poor performance in complex cases. As a result, there has been keen interest in HEP for modern machine learning to power adaptive sampling. While previous studies have shown the potential of normalizing-flow-powered neural importance sampling (NIS) over VEGAS, there remains a gap in accessible tools tailored for non-experts. In response, we introduce ZüNIS, a fully automated NIS library designed to bridge this divide, while at the same time providing the infrastructure to customise the algorithm for dealing with challenging tasks. After a general introduction on NIS, we first show how to extend the original formulation of NIS to reuse samples over multiple gradient steps while guaranteeing a stable training, yielding a significant improvement for slow functions. Next, we introduce the structure of the library, which can be used by non-experts with minimal effort and is extensivly documented, which is crucial to become a mature tool for the wider HEP public. We present systematic benchmark results on both toy and physics examples, and stress the benefit of providing different survey strategies, which allows higher performance in challenging cases. We show that ZüNIS shows high performance on a range of problems with limited fine-tuning.

研究の動機と目的

• 複雑なHEPシミュレーションにおける手作業で調整された VEGAS への依存を減らす動機付け。
• 勾配ステップ間でサンプルを再利用するデータ効率の高いニューラル重要度サンプリング（NIS）訓練手法を提案する。
• 非専門家がNISを適用できるエンドツーエンドで使いやすいPyTorchベースのライブラリ（ZüNIS）を提供する。

提案手法

• 補助分布q(x)を用いて損失を推定しつつ、訓練用にp(x, θ)を評価する一方でqからサンプリングするデータ効率の高い訓練手法を用いて、Neural Importance Samplingを拡張する。
• 複数の勾配ステップにわたってサンプルを再利用する訓練アルゴリズム（バッファ付き/フォワードサンプリング）を導出・実装する。
• ３つのqサンプリング戦略：uniform（フラット）、frozen model、adaptive_varianceを導入し、柔軟な訓練を可能にする。
• モジュラーなコンポーネント（Flows、Trainers、Integrators）を備えたPyTorchベースのライブラリ（ZüNIS）を提供する。
• 低次元・高次元の積分、MadGraphのクロスセクション様問題を含む問題に対して、NIS（ZüNIS付き）をVEGASおよび uniform sampling と比較する。）

(a)

実験結果

リサーチクエスチョン

• RQ1 サンプル再利用を用いたNeural Importance Samplingは、難易度の高いHEP積分でVEGASを上回れるか？
• RQ2 データ効率訓練スキームはNISの収束速度とデータ使用量にどう影響するか？
• RQ3 異なる次元数や問題タイプ（低次元関数とMadGraph風積分）におけるZüNISの実用的な性能はどうか？

主な発見

関数	uniformに対する分散削減	VEGASに対する分散削減
Camel	1.8±0.4×10^3	7.0±1.4×10^2
Slashed Circle	8.9±0.9×10^1	8.8±0.9×10^1
Sinusoidal	2.0±0.5×10^2	1.6±0.5×10^2

• ZüNISは、2D–35D関数（キャメル、スラッシュドサークル、サイン波）に対して、uniform samplingおよびVEGASと比較して最大で1000倍の分散削減を達成。
• 複数の勾配ステップでサンプルを再利用することで訓練を著しく加速し、forward samplingで収束性が2–10倍改善。
• 一部の問題、特に評価コストが高いケースでは積分収束とunweighting効率でVEGASを上回ることがあるが、訓練とサンプリング全体は遅い。
• MadGraphクロスセクション試験では、過程ごとに改善が変動；滑らかな関数はKL損失、ピーク関数は分散損失が有利であり、適応的な損失選択が望ましいことを示唆。
• ZüNISは訓練とサンプリングでVEGASより遅いが、GPU加速を活用して関数評価がコスト高い場合に実用的な利点を提供。

(b)