[Paper Review] Neural Network Prior Mean for Particle Accelerator Injector Tuning

Connie Xu, Ryan Roussel+1 more|arXiv (Cornell University)|Nov 16, 2022

Machine Learning and Data Classification4 citations

TL;DR

This paper proposes using a neural network surrogate model as a prior mean in Gaussian processes for Bayesian optimization in particle accelerator injector tuning. By incorporating a physics-informed, fast-executing neural network prior, the method accelerates initial optimization convergence—especially when the surrogate model correlates well with ground truth—demonstrating that even moderately accurate surrogates improve performance over constant priors, with the most accurate model matching the performance of the true surrogate.

ABSTRACT

Bayesian optimization has been shown to be a powerful tool for solving black box problems during online accelerator optimization. The major advantage of Bayesian based optimization techniques is the ability to include prior information about the problem to speed up optimization, even if that information is not perfectly correlated with experimental measurements. In parallel, neural network surrogate system models of accelerator facilities are increasingly being made available, but at present they are not widely used in online optimization. In this work, we demonstrate the use of an approximate neural network surrogate model as a prior mean for Gaussian processes used in Bayesian optimization in a realistic setting. We show that the initial performance of Bayesian optimization is improved by using neural network surrogate models, even when surrogate models make erroneous predictions. Finally, we quantify requirements on surrogate prediction accuracy to achieve optimization performance when solving problems in high dimensional input spaces.

Motivation & Objective

• To improve the sample efficiency of Bayesian optimization in high-dimensional accelerator tuning by incorporating prior physics knowledge via surrogate models.
• To investigate whether neural network-based surrogate models can serve as effective prior mean functions in Gaussian processes for online accelerator optimization.
• To quantify the required accuracy of surrogate models to achieve measurable improvements in Bayesian optimization performance.
• To compare the performance of fixed neural network priors against constant and ground truth prior means in real-world accelerator tuning scenarios.
• To evaluate the trade-offs between surrogate model complexity and optimization convergence speed in high-dimensional input spaces.

Proposed method

• Trained two neural network surrogate models on simulated data from the LCLS injector, with varying complexity and training data density (3^9 vs. 4^9 samples).
• Used the more accurate neural network as a custom prior mean function in a Gaussian Process model for Bayesian optimization.
• Modified standard Bayesian optimization to incorporate the neural network prior mean while retraining kernel hyperparameters at each step.
• Employed the Upper Confidence Bound (UCB) acquisition function with β = 2.0 and β = 0.1 to guide the selection of next evaluation points.
• Initialized optimization with three random data points from the injector model and tracked the maximum objective value over 50 iterations.
• Compared optimization performance across four priors: constant mean, ground truth surrogate mean, and two neural network priors (Model1: less accurate, Model2: more accurate).

Experimental results

Research questions

• RQ1 Can a neural network surrogate model serve as an effective prior mean function in Gaussian processes for Bayesian optimization in accelerator tuning?
• RQ2 How does the accuracy of a neural network surrogate model affect the convergence rate of Bayesian optimization in high-dimensional input spaces?
• RQ3 Does using a fixed neural network prior improve optimization performance compared to a constant prior mean function?
• RQ4 At what level of surrogate model accuracy does the optimization performance begin to approach that of the true ground-truth surrogate?
• RQ5 Does the performance of a fixed neural network prior degrade over time due to model inaccuracies, and can this be mitigated by retraining?

Key findings

• The more accurate neural network prior (Model2) consistently outperformed both the constant prior and the less accurate neural network prior (Model1) after 50 optimization iterations.
• Model1 initially outperformed the constant prior but stagnated around iteration 30, eventually being surpassed by the constant prior, indicating performance degradation due to model inaccuracies.
• Model2’s optimization performance was comparable to that of the ground truth surrogate mean, demonstrating that a well-trained neural network prior can closely approximate the true prior.
• The results show that even moderately accurate surrogate models (with high correlation to ground truth) can significantly improve initial optimization speed and convergence.
• The study confirms that fixed prior means, even when inaccurate, can degrade performance over time, suggesting the need for retraining or adaptive priors in long-running optimization.
• The findings support the use of neural network surrogates as effective prior means in Bayesian optimization, especially when they correlate well with experimental or simulation-based ground truth.