Abstract: Plasma wake field accelerator achieves high accelerating gradient on trailing electron (or positron) beams by exciting wake fields in a plasma through high-energy particle beams or high intensity laser pulses. Both underloaded and overloaded beam can significantly affect the accelerated beam quality, including energy spread and emittance, as well as the acceleration efficiency. Such beam loading effect involves complex nonlinear physical processes, which makes its optimization difficult with the traditional parameter scan method. In this work, we employs the quasi-static particle-in-cell simulation code QuickPIC together with the Bayesian optimization algorithm to find the optimal beam profile that can flattened longitudinal electric field inside the plasma wake. We also improve this Bayesian optimization algorithm by introducing an efficient exploration method for high-dimensional parameter spaces. Numerical simulations demonstrate that the optimal beam parameters can significantly improve the wake field’s quality. The enhanced algorithm achieves a 7.6 times increase in optimization efficiency and a 4.65 times decrease in the absolute objective function value. This work provides a reliable beam optimization solution for the plasma wake field accelerator applications in the compact accelerators and future high-energy colliders as well as shows the availability and high efficiency of the improved Bayesian optimization algorithm.