[论文解读] 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 的高速、可微分线性束动力学模拟器，可在实现多数量级加速的同时，通过自动微分实现高效的梯度计算。在一项基准测试中，它将强化学习训练时间从超过 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.

研究动机与目标

• 解决由于束流时间有限和模拟成本高昂而导致的粒子加速器机器学习中的数据稀缺瓶颈。
• 克服传统加速器模拟代码在大规模数据生成方面计算不可行的问题。
• 在高维参数空间中实现基于梯度的优化，以实现束流调谐和系统识别。
• 促进基于物理的模拟与神经网络代理模型的无缝集成，实现可重用、可微分的加速器建模。
• 提供一个用户友好且可扩展的框架，弥合机器学习工具与加速器物理工作流程之间的鸿沟。

提出的方法

• 在 PyTorch 中实现一个可微分、张量化的束流动力学引擎，用于高速模拟线性束流动力学。
• 利用 PyTorch 的自动微分功能，高效计算束流可观测量对所有系统参数的梯度。
• 通过向量化运算和 GPU 加速优化计算，将模拟时间减少多个数量级。
• 支持与神经网络的模块化集成，用于对空间电荷等复杂效应进行代理建模。
• 实现从神经网络输出到可微分物理模拟器的端到端梯度传播。
• 设计框架以兼容标准机器学习工具，如强化学习环境（例如 Gymnasium）和贝叶斯优化。

实验结果

研究问题

• RQ1 可微分、高速的束流动力学模拟器能否将加速器调谐中强化学习的训练时间从数天缩短至数小时？
• RQ2 可微分模拟在多大程度上能够实现在高维加速器参数空间中的基于梯度的优化？
• RQ3 可微分物理模型在贝叶斯优化中作为先验，对加速器系统识别任务的收敛性提升效果如何？
• RQ4 模块化的神经网络代理模型能否被训练并无缝集成到可微分物理模拟器中，以对空间电荷等复杂效应进行建模？
• RQ5 可微分模拟器与机器学习框架的集成，如何提升代理模型的可重用性和可扩展性？

主要发现

• Cheetah 将横向束流调谐任务的模拟时间从使用 Ocelot 时的超过 12 天缩短至 1 小时以内，使大规模强化学习训练成为可能。
• 该框架支持对所有束流线参数的高效梯度计算，使高维优化成为可能，突破了传统数值方法的限制。
• Cheetah 通过提供可微分先验，实现了物理信息引导的贝叶斯优化，显著提升了系统识别任务的收敛性。
• 针对空间电荷效应的模块化神经网络代理模型可被训练并无缝集成到 Cheetah 中，实现了快速、可微分且可重用的建模。
• Cheetah 与 Gymnasium 和 PyTorch 等机器学习生态系统的集成，实现了结合可微分物理与神经网络的端到端训练流程。
• 基准测试证实，Cheetah 相较于 Ocelot 等传统模拟代码实现了多个数量级的加速，使大规模数据生成成为现实。