[論文レビュー] Fast Point Cloud Generation with Diffusion Models in High Energy Physics

V. M. Mikuni, Benjamin Nachman+1 more|arXiv (Cornell University)|Apr 3, 2023

Generative Adversarial Networks and Image Synthesis被引用数 4

ひとこと要約

本稿では、高エネルギー物理学のジェットシミュレーションに特化した、可変長の点群を高精度に生成する拡散モデルであるFast Point Cloud Diffusion (FPCD) を紹介する。プログレッシブ・ディスティルレーションを用いることで、ベースライン比で推論時間を450倍短縮し、1ステップ推論で最先端の性能を達成しながら、トップクォーク、Wボソン、グルーオンを含む複数のジェットタイプにおいても正確性を維持する。

ABSTRACT

Many particle physics datasets like those generated at colliders are described by continuous coordinates (in contrast to grid points like in an image), respect a number of symmetries (like permutation invariance), and have a stochastic dimensionality. For this reason, standard deep generative models that produce images or at least a fixed set of features are limiting. We introduce a new neural network simulation based on a diffusion model that addresses these limitations named Fast Point Cloud Diffusion (FPCD). We show that our approach can reproduce the complex properties of hadronic jets from proton-proton collisions with competitive precision to other recently proposed models. Additionally, we use a procedure called progressive distillation to accelerate the generation time of our method, which is typically a significant challenge for diffusion models despite their state-of-the-art precision.

研究の動機と目的

• 高エネルギー物理学のデータシミュレーションにおいて、グリッドベースおよび固定長の生成モデルの限界を克服し、特にハドロンジェットのような可変長・順列不変性を持つ点群を生成すること。
• ピxls化や固定サイズの制約なしに、ジェットデータの連続的・確率的・対称的な性質を保持する拡散モデルの開発。
• 通常は遅い拡散モデルの推論を、プログレッシブ・ディスティルレーションによって高速化し、大規模データセットにおけるリアルタイムまたはニアリアルタイムのシミュレーションを可能にすること。
• 個別のモデルを訓練せずに、1つの条件付きモデルで複数の粒子種類（例：トップクォーク、グルーオン、Wボソン）のジェット生成を統合すること。

提案手法

• FPCDは2段階の拡散モデルを採用：1つはジェットの運動量（例：エネルギー、横運動量）を生成し、もう1つはジェット種別および運動量を条件として、粒子レベルの運動量（4元運動量）を生成する。
• スコアベースの拡散フレームワークを用い、対数確率密度の勾配を学習することで、可換最適化が可能になり、逆写像の流れを必要としない安定な学習が可能になる。
• マスキング戦略により、トレーニング時にマスクされた入力を条件として用いることで、30〜150個の粒子を有するジェットを可変粒子多重度で生成可能になる。
• プログレッシブ・ディスティルレーションを適用し、512ステップの完全な拡散プロセスを8ステップ、さらには1ステップで生成可能なモデルに縮小することで、推論を顕著に高速化する。
• 粒子種別（例：トップクォーク、グルーオン）を条件として用いることで、1つの統合アーキテクチャで多様なジェット種別を生成可能になり、個別のモデルを訓練する必要がなくなる。
• グラフニューラルネットワークを活用することで、点群データの順列不変性を処理し、複雑なトポロジカル関係を捉える。

実験結果

リサーチクエスチョン

• RQ1 可変長・順列不変性を持つ点群を、高エネルギー物理学におけるハドロンジェットの代表的特性として効果的に生成できるか？
• RQ2 標準的な拡散モデルの高い計算コストを、ジェットシミュレーションにおける物理的忠実性を損なわず、どのように低減できるか？
• RQ3 1つの統合モデルが粒子種別を条件として用いることで、トップクォーク、Wボソン、グルーオンなどの複数のジェット種別を高精度に生成できるか？
• RQ4 512ステップから1ステップに減少させた拡散推論ステップにおいて、プログレッシブ・ディスティルレーションが生成品質をどの程度保持できるか？
• RQ5 FPCDは、EPiC-GAN や MP-GAN といった既存の生成モデルと比較して、ジェット生成タスクにおいて忠実度と速度の両面で優れているか？

主な発見

• FPCDは物理的指標において最先端の性能を達成し、150粒子のデータセットにおいてグルーオンでは粒子多重度誤差が 0.157 ± 0.036、Zボソンでは 0.241 ± 0.090 を記録した。
• 1ステップ推論で実行されるFPCD 1モデルは高い忠実度を維持し、グルーオンでは粒子多重度誤差が 0.157 ± 0.035、Zボソンでは 0.239 ± 0.090 を達成した。
• 150粒子のFPCDにおける生成時間は、ベースラインの 31 × 10³ μs から FPCD 1の 66 μs に短縮され、全モデル比で450倍の高速化が達成された。
• FPCD 1はMP-GAN（150粒子で11 μs 対 12 μs）よりも高速であり、1ステップの拡散モデルが複数ステップのGANを上回る速度性能を示した。同時に、忠実度は同等またはそれを上回った。
• 1つのアーキテクチャでトップクォーク、Wボソン、Zボソン、グルーオン、軽いクォークの5種類のジェットを生成可能であり、個別モデルの訓練を回避できた。
• 付録Bのヒストグラム比較により、トップクォーク発生ジェットのすべての運動量分布において、蒸留モデルが完全なシミュレーションと密接に一致していることが確認された。