Physics-informed Bayesian Optimization of an Electron Microscope

Achieving sub-Angstrom resolution by aberration correction in an electron microscope requires the precise control of a large set of multipole electron-optical elements whose multidimensional optimization is comparable to the complexity of a synchrotron light source. We draw on this experience by realizing the equivalence between aberration correction and phase space beam emittance minimization. We then show a convolutional neural network (CNN) can be trained to predict beam emittance from single Ronchigrams, providing a bounded metric that enable Bayesian optimization (BO) for the autonomous aberration correction of the microscope. We show that the description in phase space via Wigner-Weyl transform can translate the aberration function of an electron beam to a single valued beam quality metric, beam emittance [1]. With such, we train a CNN on simulated Ronchigrams to capture the phase space variances of the beam from aberrations. Then the aberration correction task reduces to the minimization of beam emittance growth via this surrogate model by finding the optimal Ronchigram [2][3]. We test the validity of the CNN on real data by varying different aberration coefficients on a Spectra 300 (Fig. 1 (b-c)). We apply BO to efficiently optimize the emittance by modelling the full posterior of the response over control parameters to account for uncertainties at each query. The workflow of this scheme is demonstrated in Fig. 1 (a). We then validate our method with a simulation study using the General Particle Tracer (GPT) model which simulates the trajectory of electrons in the STEM column with a hexapole corrector (2 hexapoles and 4 transfer lenses) [4]. Fig. 2 (a-f) shows an example of the correction of A2 aberration. In the initial uncorrected case (left panel), the distribution is dominated by 3-fold astigmatism. On contrary, in the corrected case (right panel), the center is almost perfectly round while only at the periphery does the distribution show uncorrected rotational symmetry of order of 6. This indicates the two hexapoles have been optimized to counteract each other to correct the aberrations. To better capture complex couplings between optimization parameters, we adapted a deep kernel (DKBO), which integrates a deep neural network (DNN) to the kernel to outperform standard kernels by learning finer representations in the higher dimensional space projections [5]. Fig. 2 (g) benchmarks the performance of our method. Bayesian optimization always outperforms the Simplex method regardless of the choice of kernel. We observe DKBO (green solid line) outperforms in the long term but underperforms in the beginning. This can be explained by the overfitting of the DNN having too few observations to train on at the early stage. As observations accumulate, the model starts to benefit from the expressive power of the DNN, which help us reach a lower emittance at higher convergence rate. We have implemented the optimization method on both Nion UltraSTEM and ThermoFisher Spectra 300 [5]. (a). Workflow of the online Bayesian optimization of an electron microscope. (b). CNN response to change in single aberration functions averaged from five measured scans on the Spectra 300. (i): First-order aberrations. (ii): Second-order aberrations. The CNN accurately captures the emittance growth with increased aberrations. (c). Measured Ronchigrams, at different tuning stages. (a)(b). GPT-simulated electron Ronchigrams, and (c)(d). Transverse distribution of electrons at the center of the objective lens, both before and after optimization of the hexapole elements. The comparison indicates that optimization of aberration correctors by minimization of beam emittance effectively minimizes aberrations. (e)(f).Phase space distribution of the electrons at the center of the objective lens in polar coordinates, before after optimization of the hexapole elements. The purple area represents beam emittance. (g), Benchmark results of deep kernel Bayesian optimization v.s. Simplex method on the GPT-6D simulation.