[論文レビュー] Branches of a Tree: Taking Derivatives of Programs with Discrete and Branching Randomness in High Energy Physics

M. Kagan, L. Heinrich|arXiv (Cornell University)|Aug 31, 2023

Distributed and Parallel Computing Systems被引用数 5

ひとこと要約

本論文は、素粒子物理学（HEP）における離散的かつ分岐的な確率性を有するプログラムの微分可能化に向け、勾配推定手法——特にベースライン付きスコア関数推定と確率的AD——を導入する。本研究では、最初の完全に微分可能な分岐プログラムを提示し、確率的ADおよびベースライン付きスコア関数推定が、小さなバッチサイズでも効果的な勾配ベースの検出器設計最適化を可能にすることを示している。

ABSTRACT

We propose to apply several gradient estimation techniques to enable the differentiation of programs with discrete randomness in High Energy Physics. Such programs are common in High Energy Physics due to the presence of branching processes and clustering-based analysis. Thus differentiating such programs can open the way for gradient based optimization in the context of detector design optimization, simulator tuning, or data analysis and reconstruction optimization. We discuss several possible gradient estimation strategies, including the recent Stochastic AD method, and compare them in simplified detector design experiments. In doing so we develop, to the best of our knowledge, the first fully differentiable branching program.

研究の動機と目的

• 素粒子物理学（HEP）における離散的かつ分岐的な確率性を有するプログラム、例えば粒子ショワーおよび検出器シミュレーションにおいて、勾配ベースの最適化を可能にすること。
• HEPソフトウェアにおけるパラメータ依存の離散的確率的性質を扱えない標準的な自動微分（AD）の限界を解消すること。
• 粒子ショワーのシミュレーションに適した完全に微分可能な分岐プログラムを設計・検証し、HEPにおける微分可能プログラミング分野における新規貢献を果たすこと。
• 簡略化された検出器設計最適化タスクにおいて、勾配推定手法（スコア関数、数値勾配、確率的AD）を比較すること。
• 検出器設計、シミュレータチューニング、再構築最適化などの応用を想定し、エンドツーエンド微分可能なHEPシミュレータの実現可能性を示すこと。

提案手法

• 離散的確率性を有する確率的プログラムにおける勾配推定に、制御変数（ベースライン）を用いたスコア関数勾配推定法を適用する。
• 近年開発された離散確率変数を有するプログラムの微分に適した手法である、確率的AD（Arya et al., 2022）を実装し、HEPワークロードに適用する。
• 検出器半径を調整可能なパラメータとする分岐的ショワーモデルを備えたおもちゃ検出器シミュレータを設計し、損失関数をショワーの径方向深さとして定義する。
• ミニバッチ勾配推定（バッチサイズ = 2）を用いたAdam最適化法を用い、検出器半径パラメータに対して500ステップの最適化を実行する。
• 各手法について10回の繰り返し最適化を実施し、最適化ステップごとの平均損失と標準偏差を用いて性能を評価する。
• 確率的AD推定器のプライマルプログラムと代替プログラムにおいて、確率変数の再利用を実装し、推定器の分散を低減する。再利用が適用された際、標準偏差が1.5倍以上低減されるのを観測した。

(a) Single particle energy loss

実験結果

リサーチクエスチョン

• RQ1 離散的確率性を有するHEP検出器設計最適化に、ベースライン付きスコア関数勾配推定法を効果的に適用できるか？
• RQ2 確率的ADは、部分子ショワーなどのHEPにおける分岐プロセスに成功裏に適用可能か、微分可能プログラミングを可能にするか？
• RQ3 低サンプルサイズのHEP最適化環境において、勾配推定手法（スコア関数、数値勾配、確率的AD）の分散と収束速度はどのように比較されるか？
• RQ4 確率的ADにおける確率変数の再利用は、HEP応用における勾配推定器の分散にどのような影響を与えるか？
• RQ5 勾配推定器の分散が著しく高い状況下でも、1ステップあたり2回のシミュレーションで勾配ベースの最適化が可能か？

主な発見

• 確率的ADおよびベースライン付きスコア関数推定法は、全推定手法の中で最小の標準偏差を示し、1ステップあたり2回のシミュレーションでも安定した最適化が可能である。
• ベースラインなしのスコア関数推定法および数値勾配は、分散が著しく高いため、低サンプル環境下では最適化が非現実的である。
• 確率的ADおよびベースライン付きスコア関数推定法は、損失最小値への収束進捗が類似しており、今後のHEP微分可能プログラミング応用において両者とも有効であることが示唆された。
• 確率的ADの代替プログラムにおける確率変数の再利用は、勾配推定器の分散を1.5倍以上低減させ、カップリング戦略の重要性を強調した。
• 本研究では、HEPにおける離散的分岐プロセスを通過する勾配推定が可能な、最初の既知の微分可能な分岐プログラムを提示した。これにより、粒子ショワーのエンドツーエンド微分可能なシミュレーションが可能になった。
• 確率的ADおよびベースライン付きスコア関数推定法が、最小限のデータでも勾配ベースの検出器設計最適化に有効であることが示され、広範なHEP応用への強い可能性を示唆した。

(b) Particle shower with splitting