Autonomous Electron Tomography Reconstruction Using Bayesian Optimization

For electron tomography [1,2], the challenge is to obtain 3D reconstructions faithful to the true morphology of inorganic and biological materials. Experimentally, resolution and reconstruction quality are constrained by an insufficient number of projections with mixed Poisson-Gaussian noise. Modern tomographic algorithms utilize compressed sensing (CS) methods that maximize sparsity in the gradient domain to produce superior 3D reconstructions [3–6]. CS tomography recovers the 3D structure of specimens by solving an optimization problem that finds a solution which (1) strongly correlates with the experimental measurements (e.g. ADF) and (2) is maximally sparse in the gradient magnitude domain (minimal total variation): argminx‖Ax−b‖22+λ‖x‖TVwhere b is the measured ADF, x is the reconstructed object, A is the measurement matrix, and λ is a tunable parameter balancing the two terms[5]. CS tomography is greatly influenced by a scientist’s decision on the appropriate weight, λ, between the data consistency and sparsity. This sensitivity to tunable hyperparameters limits CS tomography. Domain experts need to explore several tens or hundreds of possible parameters, resulting in days of delay until a reliable 3D reconstruction is available. We present a framework for autonomous compressed sensing parameter tuning of 3D electron tomography reconstructions by leveraging Bayesian optimization with Gaussian processes (BO-GP) [7]. Bayesian optimization is a machine learning process that autonomously and efficiently finds minima of unknown functional landscapes, allowing for efficient CS hyperparameter selection. The framework consists of two core components: (1) develop a posterior probability distribution of the parameter space with Gaussian Process (GP) regression and (2) control the exploration of future measurements with an acquisition function [8]. When applied to electron tomography simulations, BO-GP can reduce parameter exploration by approximately 50%, efficiently modeling and methodically sampling the parameter space. Fig. 1 shows the Bayesian optimization process for determining optimal tomography hyperparameters from simulated measurements of a simulated strontium titanate (SrTiO3) nanoparticle. Bayesian optimization uses Gaussian processes to model the root mean square error (RMSE) landscape as a family of functions and quantifies uncertainty in unsampled regions (Fig. 1a). As more parameters are sampled, its confidence in the RMSE landscape improves until the optimization completes (Fig. 1b). For the SrTiO3 dataset, Bayesian optimization confidently identifies the solution within approximately 13 iterations. Pareto fronts quantify the tradeoff between maximizing sparsity and preserving consistency with measurements. The Pareto front plot (Fig. 2a) is generated from many reconstructions across a comprehensive range of λ values and highlights the optimal tradeoff between each term in the CS optimization problem [9]. Here, the resulting TV and least squares are the x and y coordinates, respectively, and each point is assigned a color value based on RMSE. We see that in regions with maximal curvature (within the leftmost corner of the curve) the λ values produce reconstructions with marginally higher data distance but minimal TV; this compressed sensing approach results in minimal RMSE (Fig. 2c-e). Bayesian Optimization with Gaussian Processes offers efficient compressed sensing electron tomography without the need for human intervention. Determining optimal parameters for simulated compressed sensing electron tomography with Bayesian Optimization. a The temporal convergence of BO performed on a synthetic SrTiO3 dataset. BO strategically determines the next point to sample after measuring the reconstruction's RMSE -- typically converging within 12 to 20 estimates. b The final estimated landscape of the CS tomography. The blue star represents the Bayesian optimization’s minimum RMSE solution, and the yellow and red represent DD and TV minima, respectively. c-e The resulting reconstructions with no TV weight (λ = 0), minimal RMSE, and excessive TV weight (large λ), respectively. f The ground truth depicts the true morphology of the nanoparticle. Exploring the relationship between reconstruction quality and lambda parameter with Pareto fronts. a Depicted are the tradeoffs from three evaluation metrics: data distance, TV, and root mean square error (ground truth dependent). The ground truth (bottom left) is the true morphology of the nanoparticle. b-e The minima of three metrics, TV, DD, and RMSE, are shown alongside their visual tradeoffs. The white circle highlights reconstructions preserving finer details of object morphology.