Cheetah: Bridging the Gap Between Machine Learning and Particle Accelerator Physics with High-Speed, Differentiable Simulations

TL;DR

Cheetah tackles data and computation bottlenecks in accelerator ML by delivering a PyTorch-based, differentiable, linear-beam dynamics simulator that runs orders of magnitude faster than existing codes. By supporting two beam representations, dynamic transfer-map reduction, and seamless integration with ML workflows, it enables efficient gradient-based optimization, system identification, and RL-driven control. The paper demonstrates five concrete use cases, including zero-shot RL transfer to real accelerators, gradient-based tuning and identification, Bayesian optimisation priors, and modular NN surrogates for space-charge effects, highlighting both speed gains and practical applicability. The work argues that fast differentiable simulations can accelerate ML development for accelerators and accelerate real-world adoption, while outlining future directions like more modular surrogates and JAX-based backends to push speeds further.

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.

Figures (9)

• Figure 1: Overview of where Cheetah fits into the proposed applications, with the Cheetah logo marking its use as a component of these applications. (a) Cheetah is used as a physical prior for BO. (b) Cheetah provides a differentiable beam dynamics model which can be used for accelerator tuning and system identification. (c) Cheetah enables the implementation of fast beam dynamics environments for training RL agents. (d) Cheetah provides the infrastructure to seamlessly integrate modular NN surrogate models with physical beam dynamics simulations.
• Figure 2: Visualisation of a simple example for transfer map reduction. The tracking function of the screen is denoted by $f_S$. It cannot be reduced along with other transfer maps. The transfer maps drift sections and magnets upstream of the screen $\left\{ R_{D1}, R_{M1}, R_{D2} \right\}$ and downstream of the screen $\left\{ R_{D3}, R_{M2}, R_{D4} \right\}$ can be reduced to two transfer maps $R_{A1}$ and $R_{A2}$, one on each end of the screen.
• Figure 3: Flowchart of the RL loop for the ARES EA transverse tuning task. The environment -- during training defined in terms of Cheetah -- outputs an observation $\bm{o}_t$ and a $r_t$ based on the previous action $\bm{a}_{t-1}$. The agent then computes a new action $\bm{a}_t$ using the neural network policy. The new action is applied to the environment and results in the next observation $\bm{o}_{t+1}$ and reward $r_{t+1}$.
• Figure 4: An NN policy trained with RL tuning of the transverse beam parameters in the ARES EA. (a) The MSE loss development over parameter update iterations. (b) Beam parameters on the diagnostics screen over parameter update iterations. (c) and (d) quadrupole and dipole magnet settings over parameter update iterations.
• Figure 5: Gradient-based tuning example of the transverse beam parameters in the ARES EA. (a) The loss development over parameter update iterations. (b) Beam parameters on the diagnostics screen over parameter update iterations. (c) and (d) quadrupole and dipole magnet settings over parameter update iterations.