Prior-mean-assisted Bayesian optimization application on FRIB Front-End tunning

Kilean Hwang,Tomofumi Maruta,Alexander Plastun,Kei Fukushima,Tong Zhang,Qiang Zhao,Peter Ostroumov,Yue Hao

DOI: https://doi.org/10.48550/arXiv.2211.06400

2022-11-12

Abstract:Bayesian optimization~(BO) is often used for accelerator tuning due to its high sample efficiency. However, the computational scalability of training over large data-set can be problematic and the adoption of historical data in a computationally efficient way is not trivial. Here, we exploit a neural network model trained over historical data as a prior mean of BO for FRIB Front-End tuning.

Accelerator Physics,Machine Learning

What problem does this paper attempt to address?

This paper aims to solve several key problems in accelerator tuning, especially those encountered in the front - end tuning of FRIB (Facility for Rare Isotope Beams). Specifically:

1. **High sample efficiency**: Bayesian Optimization (BO) is often used in accelerator tuning due to its high sample efficiency. However, the computational scalability of training on large - scale data sets is a challenge, and how to utilize historical data in a computationally efficient manner is also a difficult problem.

2. **Distribution shift**: After each ion source type switch or restart, the machine state changes significantly (i.e., "distribution shift" or "machine drift"), causing past data to be a poor representation of the current machine state. Therefore, a strategy that can adapt to distribution shift with high sample efficiency is required.

3. **Rapid tuning**: FRIB front - end tuning requires rapid adjustment to maintain optimal beam optics performance. Traditional black - box optimization algorithms may not be able to meet this requirement, especially in the presence of distribution shift.

To solve the above problems, the authors propose a Prior - mean - assisted Bayesian Optimization (pmBO) method. Specifically, a neural network model is used to train historical data, and this model is used as the prior mean of the Gaussian Process (GP) surrogate model in Bayesian Optimization. This method can not only utilize the useful information provided by historical data, but also quickly find solutions close to the optimal solution with high sample efficiency.

Through experimental verification, pmBO performs excellently in virtual FE tuning tasks. Compared with traditional Bayesian optimization methods, pmBO can achieve better tuning results with fewer iterations. In addition, the actual FRIB front - end tuning results show that pmBO can find a nearly convergent solution in the second iteration, significantly improving the beam transmission rate and the position accuracy of the transverse beam center.