Proteins, which are large and intricate molecules, are fundamental to all forms of life and perform a plethora of functions within organisms. Historically, scientists have relied on experimental methods such as nuclear magnetic resonance and X-ray crystallography to ascertain the structures of proteins. However, these techniques are notoriously labor-intensive, requiring extensive setups and significant trial and error. Cryo-electron microscopy (cryo-EM) has emerged as a promising alternative, becoming the preferred technique for determining 3D protein structures with atomic-level resolution (Nakane et al., 2020). A key advantage of cryo-EM is its ability to analyze conformational mixtures as the molecules are imaged in their near-native states. This capability has proven invaluable since the outbreak of COVID-19 in January 2020, enabling the creation of the first dynamic visualization of the 2019-nCoV Spike trimer structure (Wrapp et al., 2020) and facilitating the first molecular-level structural analysis of the Omicron variant’s spike protein. This analysis has provided crucial insights into how the heavily mutated Omicron variant binds to and infects human cells (Mannar et al., 2022). Nevertheless, data obtained through cryo-EM are characterized by significant noise, vast dimensions, a large volume of unlabeled data, and high heterogeneity coupled with unknown orientations. These factors complicate the achievement of reliable computational conclusions (Singer and Sigworth, 2020).

It is noted that classifying and refining data collected in a single day can take weeks due to the complexity involved in estimating 3D orientations and the necessity for human intervention in the clean-up process. In this study, we build on the concepts presented in (Banjac et al., 2021), which employed contrastive learning (Le-Khac et al., 2020) and neural networks to estimate the distances of orientations between projections. We believe incorporating pairwise relationships can add regularization to orientation learning, thus enabling us to alleviate the issues posed by high noise levels. Furthermore, utilizing neural networks that carry out amortized inference (Donnat et al., 2022) on orientation can substantially reduce the processing time. Lastly, we suggest a novel uncertainty measure and reconstruction pipeline to assess the reliability of our estimates and speed up the entire image analysis process. The contributions of this paper can be summarized as follows:

Uncertainty estimation for orientation estimation: Evaluating the reliability of the network’s predictions is critical, particularly given the substantial presence of outliers and contaminants in the cryo-EM dataset. Uncertainty estimation is vital in quantifying the dependability of the network’s predictions and can aid us in filtering particles during data cleaning. In this study, we introduce uncertainty measures that can serve as proxies for estimating testing errors. Additionally, we propose a strategy, based on the uncertainty measure, to clean up the dataset directly at the 3D level, which, as we demonstrate, can lead to a more accurate 3D reconstruction. This advancement could potentially decrease the time spent on dataset clean-up using traditional 2D classification methods.

Model design and generalization capability: Another essential concern is the model’s ability to generalize effectively. In this study, we systematically assess the network’s generalization capabilities. Specifically, we suggest using distance learning as an auxiliary loss to regularize the learning process and explore the potential of different components in the neural network. This comprehensive study can guide us in understanding the design choices made when leveraging a neural network for amortized inference. Furthermore, it addresses the current gap in the design of the encoder network in the generative model, an emerging framework in 3D reconstruction (Donnat et al., 2022). Ultimately, this analysis may pave the way to obtain a pre-trained model, thereby further accelerating the process of orientation estimation.

By addressing these critical issues, this paper aims to make significant contributions to the ongoing efforts to leverage neural networks for orientation estimation in cryo-EM (Banjac et al., 2021, Lian et al., 2022) and to improve the design of encoders in current generative models (Donnat et al., 2022, Nashed et al., 2021, Levy et al., 2022). The rest of this paper is organized as follows. In Section 2, we first review the related work on orientation estimation, followed by an examination of the related work on generative models that use 3D orientation in their autoencoder frameworks. In Section 3, we present the design of the proposed framework and provide insights into each component. In Section 4, we offer numerical results to compare different design choices in the methods and demonstrate the superior performance of the proposed framework. Finally, in Section 5, we discuss the potential of the methodology and draw conclusions. For readers not mathematically inclined, Appendix A offers brief definitions of key terms.