[Paper Review] 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 dynamicsEngineering24 references3 citations

TL;DR

Cheetah is a PyTorch-based, high-speed, differentiable linear-beam dynamics simulator that accelerates beam dynamics simulations by multiple orders of magnitude while enabling efficient gradient computation via automatic differentiation. It reduces reinforcement learning training time from over 12 days to under 1 hour in one benchmark, enabling scalable machine learning applications in particle accelerator physics.

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.

Motivation & Objective

• Address the data scarcity bottleneck in machine learning for particle accelerators due to limited beam time and expensive simulations.
• Overcome the computational infeasibility of large-scale data generation using traditional accelerator simulation codes.
• Enable gradient-based optimization in high-dimensional parameter spaces for beam tuning and system identification.
• Facilitate seamless integration of physics-based simulations with neural network surrogate models for reusable, differentiable accelerator modeling.
• Provide a user-friendly, extensible framework that bridges machine learning tools with accelerator physics workflows.

Proposed method

• Implement a differentiable, tensorized beam dynamics engine in PyTorch for high-speed simulation of linear beam dynamics.
• Leverage PyTorch's automatic differentiation to compute gradients of beam observables with respect to all system parameters efficiently.
• Optimize computation using vectorized operations and GPU acceleration to reduce simulation time by multiple orders of magnitude.
• Support modular integration with neural networks for surrogate modeling of complex effects like space charge.
• Enable end-to-end gradient flow from neural network outputs back through the differentiable physics simulator.
• Design the framework to be compatible with standard ML tools such as reinforcement learning environments (e.g., Gymnasium) and Bayesian optimization.

Experimental results

Research questions

• RQ1 Can a differentiable, high-speed beam dynamics simulator reduce training time for reinforcement learning in accelerator tuning from days to hours?
• RQ2 To what extent can differentiable simulations enable gradient-based optimization in high-dimensional accelerator parameter spaces?
• RQ3 How effectively can differentiable physics models serve as priors in Bayesian optimization for accelerator system identification?
• RQ4 Can modular neural network surrogates be trained and integrated with differentiable physics simulators to model complex effects like space charge?
• RQ5 How does the integration of differentiable simulators with ML frameworks improve the reusability and scalability of surrogate models?

Key findings

• Cheetah reduces simulation time for a transverse beam tuning task from over 12 days (using Ocelot) to under 1 hour, enabling practical large-scale RL training.
• The framework supports efficient gradient computation across all beamline parameters, making high-dimensional optimization feasible beyond the limits of traditional numerical methods.
• Cheetah enables physics-informed Bayesian optimization by providing differentiable priors that improve convergence in system identification tasks.
• Modular neural network surrogate models for space charge effects can be trained and seamlessly integrated with Cheetah, enabling fast, differentiable, and reusable modeling.
• The integration of Cheetah with ML ecosystems like Gymnasium and PyTorch allows for end-to-end training pipelines combining differentiable physics and neural networks.
• Benchmarking confirms that Cheetah achieves speedups of multiple orders of magnitude over legacy simulation codes like Ocelot, making large-scale data generation practical.