[論文レビュー] Bridging the gap between machine learning and particle accelerator physics with high-speed, differentiable simulations

Jan Kaiser, Chenran Xu+2 more|arXiv (Cornell University)|Jan 1, 2024

Particle accelerators and beam dynamics参考文献 24被引用数 3

ひとこと要約

Cheetah は、PyTorch に基づく高速で微分可能な線形ビームダイナミクスシミュレータであり、複数のオーダーの速度向上を実現するとともに、自動微分を用いた効率的な勾配計算を可能にする。1つのベンチマークでは、強化学習の訓練時間を 12 日以上から 1 時間未塔に短縮し、素粒子加速器物理学におけるスケーラブルな機械学習応用を可能にする。

ABSTRACT

Machine learning has emerged as a powerful solution to the modern challenges in accelerator physics. However, the limited availability of beam time, the computational cost of simulations, and the high-dimensionality of optimisation problems pose significant challenges in generating the required data for training state-of-the-art machine learning models. In this work, we introduce Cheetah, a PyTorch-based high-speed differentiable linear-beam dynamics code. Cheetah enables the fast collection of large data sets by reducing computation times by multiple orders of magnitude and facilitates efficient gradient-based optimisation for accelerator tuning and system identification. This positions Cheetah as a user-friendly, readily extensible tool that integrates seamlessly with widely adopted machine learning tools. We showcase the utility of Cheetah through five examples, including reinforcement learning training, gradient-based beamline tuning, gradient-based system identification, physics-informed Bayesian optimisation priors, and modular neural network surrogate modelling of space charge effects. The use of such a high-speed differentiable simulation code will simplify the development of machine learning-based methods for particle accelerators and fast-track their integration into everyday operations of accelerator facilities.

研究の動機と目的

• 限られたビーム時間と高コストなシミュレーションによる、素粒子加速器における機械学習のデータ不足のボトル neck を解決する。
• 従来の加速器シミュレーションコードでは大規模なデータ生成が計算的に非現実的であるという課題を克服する。
• 高次元パrameter空間における勾配ベース最適化を、ビームチューニングおよびシステム同定に可能にする。
• 物理ベースのシミュレーションとニューラルネットワークの代替モデルをシームレスに統合し、再利用可能で微分可能な加速器モデリングを実現する。
• 機械学習ツールと加速器物理学ワークフローを橋渡しする、使いやすく拡張可能なフレームワークを提供する。

提案手法

• PyTorch を用いて、高速な線形ビームダイナミクスのための微分可能でテンソル化されたビームダイナミクスエンジンを実装する。
• PyTorch の自動微分を活用し、すべてのシステムパrameter に対してビーム観測量の勾配を効率的に計算する。
• ベクトル化された演算と GPU アクセeleration を最適化して、シミュレーション時間を複数オーダー短縮する。
• 空間電荷のような複雑な効果の代替モデリングのため、ニューラルネットワークとのモジュラーフォーマット統合をサポートする。
• ニューラルネットワークの出力から、微分可能な物理シミュレータを経由して逆方向に勾配が流れることを可能にする。
• 強化学習環境（例：Gymnasium）やベイズ最適化など、標準的な機械学習ツールと互換性を持つようにフレームワークを設計する。

実験結果

リサーチクエスチョン

• RQ1 微分可能で高速なビームダイナミクスシミュレータは、加速器チューニングの強化学習における訓練時間を数日から数時間に短縮できるか？
• RQ2 微分可能なシミュレーションは、高次元加速器パrameter空間における勾配ベース最適化をどの程度可能にするか？
• RQ3 微分可能な物理モデルは、加速器システム同定のためのベイズ最適化において、どの程度有効な事前分布として機能するか？
• RQ4 モジュラーなニューラルネットワーク代替モデルは、微分可能な物理シミュレータと統合され、空間電荷のような複雑な効果を効果的にモデル化できるか？
• RQ5 微分可能なシミュレータと機械学習フレームワークの統合は、代替モデルの再利用性とスケーラビリティをどのように向上させるか？

主な発見

• Cheetah は、Ocelot を用いた場合に 12 日以上かかっていた横方向ビームチューニングタスクのシミュレーション時間を 1 時間未塔に短縮し、大規模な強化学習訓練を実用可能にした。
• フレームワークは、すべてのビームラインパrameter における効率的な勾配計算をサポートしており、従来の数値的手法の限界を超えた高次元最適化を可能にした。
• Cheetah は、微分可能な事前分布を提供することで、物理的制約を組み込んだベイズ最適化を可能にし、システム同定タスクでの収束性を向上させた。
• 空間電荷効果のためのモジュラーなニューラルネットワーク代替モデルは、Cheetah とシームレスに統合可能であり、高速で微分可能かつ再利用可能なモデリングを実現した。
• Cheetah と Gymnasium や PyTorch などの機械学習エコシステムの統合により、微分可能な物理とニューラルネットワークを組み合わせたエンドツーエンドの訓練パイプラインが可能になった。
• ベンチマークの結果、Cheetah は Ocelot のようなレガシーシミュレーションコードと比較して、複数オーダーの高速化を達成しており、大規模なデータ生成を現実可能にした。