[論文レビュー] PC-Droid: Faster diffusion and improved quality for particle cloud generation

Matthew Leigh, Debajyoti Sengupta+3 more|arXiv (Cornell University)|Jul 13, 2023

Advanced Neural Network Applications参考文献 59被引用数 5

ひとこと要約

PC-Droid は、高エネルギー物理学における粒子クラウド生成のための最先端の拡散モデルを導入し、改良された EDM ノイズスケジューラー、高度な ODE スolver、および一貫性蒸留を活用することで、PC-JeDi よりも最大2桁、Delphes よりも3桁速い生成速度を達成しながら、すべてのジェットタイプと指標において品質を維持または向上させた。

ABSTRACT

Building on the success of PC-JeDi we introduce PC-Droid, a substantially improved diffusion model for the generation of jet particle clouds. By leveraging a new diffusion formulation, studying more recent integration solvers, and training on all jet types simultaneously, we are able to achieve state-of-the-art performance for all types of jets across all evaluation metrics. We study the trade-off between generation speed and quality by comparing two attention based architectures, as well as the potential of consistency distillation to reduce the number of diffusion steps. Both the faster architecture and consistency models demonstrate performance surpassing many competing models, with generation time up to two orders of magnitude faster than PC-JeDi and three orders of magnitude faster than Delphes.

研究の動機と目的

• 高エネルギー物理学のシミュレーションにおける粒子ジェットのポイントクラウドとしての生成のための、より高速かつ高精度な拡散モデルの開発。
• PC-JeDi や Delphes などの既存のサーヴィレートモデルと比較して、品質を維持または向上させながら推論時間を短縮すること。
• アテンションベースのアーキテクチャと一貫性モデルを用いて、生成速度と品質のトレードオフを調査すること。
• すべてのジェットタイプにわたる共同学習を実施し、一般化性能と無条件生成能力を向上させること。
• 標準的な指標と実世界のハードウェア評価を用いて、粒子クラウド生成のための新しいベンチマークを確立すること。

提案手法

• SDE と ODE を簡素化した EDM（ノイズ除去スコアベース）拡散定式化を採用し、$\gamma(t) = 1$ および $\sigma(t) = t$ を使用することで、より直線的な逆方向軌道と少ない切り捨て誤差を実現した。
• 性能向上が最も顕著に現れる中間ノイズレベルに焦点を当てるために、ノイズ率 $\sigma \sim \mathcal{N}(-1.2, 1.2)$ の対数正規分布サンプリングを採用した。
• ノイズレベルが低い場合にネットワークをバイパス可能にするために、$\sigma$-依存のスケーリングおよびスキップ接続（$c_{\text{in}}, c_{\text{out}}, c_{\text{skip}}$）を導入し、単位分散を維持した。
• ノイズ予測ではなく再構成誤差を直接最小化するため、損失関数を $\mathcal{L} = \|D_\theta(\boldsymbol{x}; \sigma, \boldsymbol{y}) - \boldsymbol{x}_0\|^2_2$ に再定式化した。
• k-diffusion ライブラリの高度な ODE 結合スケジューラーを活用し、一貫性モデル学習において DPM2 が優れた性能を示した。
• PC-Droid を教師モデルとして用いて、一貫性蒸留モデル（CD1, CD5）を訓練し、許容可能な忠実度で1ステップ生成が可能になった。5ステップで性能が飽和した。

Figure 1 : PC-Droid network architecture and training setup.

実験結果

リサーチクエスチョン

• RQ1 再考された拡散定式化は、粒子クラウド生成の推論時間を著しく短縮するとともに、品質を向上させることができるか？
• RQ2 現代の ODE スolver と改良されたネットワーク正規化技術は、ジェット粒子クラウド生成の忠実度と速度にどのように影響を与えるか？
• RQ3 一貫性モデルは、標準的な拡散モデルと比較して、競争力のある品質で1ステップ生成を可能にするか、その程度はどの程度か？
• RQ4 さまざまなアーキテクチャと蒸留戦略において、生成速度と品質の性能トレードオフはどのように変化するか？
• RQ5 すべてのジェットタイプにわたる共同学習は、一般化性能を向上させ、性能劣化を伴わずに無条件生成を可能にするか？

主な発見

• PC-Droid は、すべてのジェットタイプと評価指標において最先端のパフォーマンスを達成し、PC-JeDi や Delphes を上回る忠実度と速度を実現した。
• 最も高速なモデル（CD1）は、EPiC-GAN の約3倍の時間でトップジェットを生成したが、FPND スコアは理想的な Delphes ベンチマークに1桁分近い。
• CD5（5ステップの一貫性モデル）は、8ステップの FPCD（PD8）モデルより2.5倍速く、$\mathrm{W}_1^{\mathrm{M}}$ と $\mathrm{W}_1^{\tau_{32}}$ を含むすべての指標を向上させた。
• フル PC-Droid モデルは EPiC-GAN よりも250倍遅いが、$\mathrm{W}_1^{\tau_{32}}$ を除き、すべての指標でほぼ理想に近いパフォーマンスを達成した。
• 一貫性モデルは、5ステップで性能が飽和することを示した。CD1 は質量スメアを示したが、特に $\tau_{32}$ においては適切な部分構造再現が可能だった。
• 改良されたノイズスケジューラーとネットワーク正規化により、切り捨て誤差が低減され、より少ない逆方向ステップで安定した学習が可能になり、品質を損なわず高速な推論が実現した。

Figure 2 : A single cross-attention encoder block updating both an input point cloud $x$ and global tokens $g$ , using layer-norm (LN), multi-headed attention (MHA). Converging arrows represent vector addition and contextual information $c$ is injected into the network by concatenating to the inputs