Flow-based surrogate models for particle tracking
TL;DR
To capture distribution-dependent dynamics that vanilla CFM cannot represent, the model is extended with cross-attention over the initial particle ensemble (Cross-Attention-CFM) and it is demonstrated that this extension recovers performance on a space-charge benchmark in the PS, where vanilla CFM degrades.
Abstract
Particle tracking is a fundamental tool for particle-accelerator design and optimisation. Conventional tracking routines provide high accuracy but are computationally demanding, especially when simulating large particle ensembles or long time spans. As a result, optimising moderate- to high-dimensional parameter spaces is challenging, and real-time surrogate models remain out of reach for many applications. This contribution introduces a surrogate-modelling approach based on conditional flow matching (CFM). A CFM model is trained on tracking simulations of CERN's Proton Synchrotron (PS) over a 10-dimensional parameter space. The trained model reproduces final phase-space distributions with a median squared maximum mean discrepancy MMD$^2$ of $3\times 10^{-4}$ and mean inference time of 0.04 s, a speed-up of three orders of magnitude over conventional tracking. To capture distribution-dependent dynamics that vanilla CFM cannot represent, we extend the model with cross-attention over the initial particle ensemble (Cross-Attention-CFM) and demonstrate that this extension recovers performance on a space-charge benchmark in the PS, where vanilla CFM degrades. Finally, we introduce Hybrid-CFM, in which a small number of conventionally-tracked particles are used to inform the model. On the same 10-dimensional PS task, Hybrid-CFM with 100 auxiliary particles trained on 200 distributions matches the vanilla CFM trained on 1500, and improves the worst-case (90th-percentile) MMD$^2$ by roughly a factor of four, substantially reducing the upfront cost of building a surrogate.