Amortized Inference for Heterogeneous Reconstruction in Cryo-EM
Authors: Axel Levy
Gordon Wetzstein
Julien N. P. Martel
FREDERIC P POITEVIN
Ellen D Zhong
Venue: NeurIPS 2022 Accept (2022)
Keywords: Cryo-electron microscopy
3D Reconstruction
Neural Scene Representation

Abstract: Cryo-electron microscopy (cryo-EM) is an imaging modality that provides unique insights into the dynamics of proteins and other building blocks of life. The algorithmic challenge of jointly estimating the poses, 3D structure, and conformational heterogeneity of a biomolecule from millions of noisy and randomly oriented 2D projections in a computationally efficient manner, however, remains unsolved. Our method, cryoFIRE, performs ab initio heterogeneous reconstruction with unknown poses in an amortized framework, thereby avoiding the computationally expensive step of pose search while enabling the analysis of conformational heterogeneity. Poses and conformation are jointly estimated by an encoder while a physics-based decoder aggregates the images into an implicit neural representation of the conformational space. We show that our method can provide one order of magnitude speedup on datasets containing millions of images, without any loss of accuracy. We validate that the joint estimation of poses and conformations can be amortized over the size of the dataset. For the first time, we prove that an amortized method can extract interpretable dynamic information from experimental datasets.

Amortized Inference for Heterogeneous Reconstruction in Cryo-EM
Axel Levy Stanford University
Gordon Wetzstein Stanford University
Julien Marte Stanford University
Frédéric Poitevir C
Ellen D. Zhong
SLAC National Accelcrator Laboratory
Princeton University
Abstract
Cryo-electron microscopy (cryo-EM) is an imaging modality that provides unique insights into the dynamies of proteins and other building blocks of life. The algo- rithmic challenge of jointly estimating the poses, 3D structure, and conformational heterogeneity of biomolecule from millions of noisy and randomly oriented 2D projections in computationally efficient manner, however, remains unsolved. Our method, cryoFIRE, performs ab initio heterogeneous reconstruction with unknown poses in amortized framework, thereby avoiding the computationally expensive step pose search while enabling the analysis of conformational heterogeneity. Poses and conformation jointly estimated by encoder while physics-based decoder aggregates the images into an implicit neural representation of the con- formational space. We show that method can provide one order of magnitude speedup on datasets containing millions of images without any loss of accuracy. We validate that the joint estimation of poses and conformations over the size the dataset. For the first time, we prove that an amortized method can extract interpretable dynamic information from experimental datasets.
be amortized
Introduction
Proteins and other biological macromolecules in the cell function through finely-tuned choreography of transitions between metastable conformational states. Analyzing the structural heterogeneity of biomolecule is therefore critical for applications such as drug design and, more generally, for understanding these essential building blocks of life.
In a single particle cryo-electron microscopy (cryo-EM) experiment, an aqueous solution of pu- rified biomolecules is Bash-frozen in a thin layer of vitreous ice and imaged with a transmission electron microscope (Fig. [ (a)). A cryo-EM experiment outputs large set of unlabeled images, each containing a 2D projection of a unique molecule, whose 3D structure sampled from some thermodynamic distribution (i.e. conformation) and viewed from an unknown orientation (i.e. pose) di RR,, t,) (3(3) R (Fig. (b)). While homogeneous reconstruction methods focus on estimating the average 3D electron scattering potential of the studied molecule (the consensus volume), heterogeneous methods take into account structural variability and introduce a variable z; for the conformational state that characterizes the electron scattering potential V(., =;) R3 IR associated with observation (5). Given the image formation model (detailed in Section 3.T). hetero- geneous reconstruction can be seen inverse problem where V, 2 and d; must be inferred from the observed images Y, graphical formulation the problem is given in Fig. (c).
*Corresonddecee to: zhonge@princeton..e
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

(a)
[b]
V(,))
R² -
(c)
Zi
d, 50(3) x IR?
Discriminctive Model (encoder)
Generative Model (decoder)
Figure I: (a) Illustration of a cryo-EM experiment. Molecules are frozen in random orientation d; and in a random conformation 2,. Their electron scattering potential V.,, z4) interacts with an electron beam resulting in noisy projected images Y; on the detector. (b) (Top) Example scattering potentials visualized as an isosurface. (Bottom) Poses are characterized by rotation in SO(3) and translation in I2 (not shown). (c) Graphical model for our method. The encoder, parameterized by v. a discriminative model that predicts Φ and 2 from Y; while the decoder is generative model that outputs noise-firee estimation of the image.
Driven by recent advances in data collection capabilities, the number of images collected per cryo- EM experiment has been steadily increasing, now reaching millions Established methods for heterogeneity analysis break down 3D reconstruction into two alternating iterative refinement steps: 1) the latent pose Φ; and conformational states 24 are first estimated from the images and the current estimate of V, and 2) the estimation of V is then updated using the current estimates of the latent variables. The primary computational bottleneck of this approach is pose estimation, which is done by rendering images from view points distributed on the the 5-dimensional space (3 x R Although previous methods came up with accelerated branch-and- bound strategies [14] 24 38 state-of-the-art methods for pose search raise two problems: (1) solving this problem for single image is time consuming, especially when the images are rendered using a neural model (39] and (2) the pose needs to be solved independently for each image in the dataset and does not use the fact that similar images are likely associated with the similar variables. Furthermore, most heterogeneous reconstruction algorithms limit their applicability by relying homogeneous reconstruction, treating poses as known to simplify the optimization problem for 2; or due to the sheer computational cost of pose search.
of millions [ [9].
upstream
In this work, we focus on ab initio heterogeneous reconstruction, meaning that reconstruction occurs de novo without any upstream estimation of the latent pose or conformation variables. We present a method that amortizes (over the size of the dataset) the cost of inference by learning a function that directly maps observed images estimates of the latent variables. Our method, cryoFIRE (Fast heterogeneous ab Initio Reconstruction for cryo-EM), leverages an amortized approach to address the increasing size of datasets, while tackling the opportunity to learn complex conformational spaces of dynamic proteins. Our contributions include:
An autoencoder-batt based architecture with a tailored loss function that processes images 10 times faster than existing methods and enables amortization of the runtime over the size of the dataset,
A demonstration that our model accurately learns the structure of low-dimensional mani- fold in the conformational space on benchmark synthetic datasets with both discrete and continuous heterogeneity, and
To our knowledge, the first instance of amortized inference for ab initio heterogeneous reconstruction of experimental cryo-EM datasets.
2
Related Work
Heterogeneous reconstruction methods in cryo-EM can be differentiated according to the way the latent pose di and conformation state z; are estimated [5). The vast majority of approaches assume that o; are previously estimated. We describe related work in terms of how o; is estimated, how 2; is

---

estimated (without pose inference), and progress towards joint amortization of all latent, unknown variables in cryo-EM.
(EM) algo- rithm, posterior distributions over the conformational states computed (or approximated) for each image at expectation step. In cryo-EM, this method was first popularized by RELION with 3D Classi- fication [29] which implemented a discrete class indicator 2; E {l..... . K}, where K is the number of classes defined by the user. More recently, a continuous reconstruction method called multi-body refinement was made available in RELION T8 In this approach, a homogeneous reconstruction is segmented by the user into a set of rigid bodies, where each rigid body is free to move relative to the others. The relative pose (orientation and translation) of each body is subsequently estimated through EM for each particle, and further reduced through Principal Component Analysis (PCA). Various methods have been developed to estimate the conformational space as a linear deformation around a known reference. HEMNMA [9] uses the normal modes, or eigenvolumes, of the reference to model its deformation and proceeds through a projection- wwt--mmatacchhaheeeeeee elastic alignment of each lingleppaticcle image with the reference structure to yield reduced representation of the variability in the dataset. 3DVA (23) relaxes the dependence on known normal modes and instead implements a variant of the EM algorithm for Probabilistic PCA which iteratively updates the eigenvolumes and the projection of the particles on them, until convergence, following [3]. Other approaches explicitly estimate the deformation field mapping the reference structure to the best-fi fit structure for each image. In Herreros et al. [8, the deformation field is defined as a linear expansion over Zermike polynomials, with 2, as their coefficients. In 3DFlex (2)), the deformation field is generated from 2; through neural flow generator, alleviating the linear constraint. Formally, each 2; can be seen as a point estimate that maximizes the posterior distribution over conformational states. An in-depth review of EM-based reconstruction method can be found in [31]. Other computational approaches like Markov Chain Monte Carlo have been explored for ab initio heterogeneous reconstruction in [15], introducing a mathematical framework for representing deformable molecules ("hyper-molecules)) but remain a prototype implementation.
Expectation-Maxiinaatioo over Conformations With an Sopectation-Maxiiizaan
Amortization over Conformations Instead of optimizing each variable independently, amortized inference learns a function qe- parameterized by v. that maps images Y; to probability (posterior) distributions over the space of latent variables [T]. If N represents the size of the dataset, amortized inference therefore replaces the estimation of N latent variables with the learning of th, which complexity (number of dimensions) does not scale with N. Previous methods explored the possibility of using amortized inference estimate the conformational state 2,, in a setting where the poses were known. CryoDRGN (37] introduced amortized variational inference to estimate the confor- mational state z; in the setting where poses are known. In cryoDRGN (37), distributions over the conformational states are predicted by the encoder of a Variational Autoencoder (VAE) [T]] 12 and the conformational space parameterized with an implicit neural representation. E2GMM
and cryoFold (40) both reconstruct a deformable atomic representation of the molecule and train an encoder associate each image with a low-dimensional 2; that characterizes the deformation. These two methods require the initialization of the backbone structure of the molecule. Although these methods can be used to analyze the structural heterogeneity in a dataset, they assume the poses to be given by another reconstruction method. In cryoDRGN-BM BNB 38) and cryoDRGN2 [39] the poses are not given anymore but instead are estimated using an exhaustive search strategy. The estimation of the poses therefore does not amortize over the size of the dataset and, in spite of branch-and- bound approach 14), the 5D pose search remains the most computationally expensive step in the pipeline.
Amortization over Poses In the context of homogeneous reconstruction, previous methods used amortized inference to predict the latent variables. In cryo [17. Miolane et al. showed that the in-plane rotation and the contrast transfer function (CTF) parameters could be jointly estimated in the latent space of an encoder. CryoGAN 6 showed that homogeneous reconstruction could be achieved with a generative framework using discriminative loss to avoid explicitly recovering the poses, considered as "nuisance" variables. SpatialVAE (2) showed that the translations and the in in-plane rotations in 2D images could be estimated with VAE-based architecture, and was later generalized to other transformations in [34]. CryoPoseNet (20) first demonstrated the possibility of using autoencoder architecture with synthetic datasets. CryoAl AI [66] later introduced the symmetric loss to help the model avoid local minima and showed that homogeneous reconstruction could be done experimental datasets. All of these methods, however, reconstruct 3D consensus volume and do not address the question of conformational hetcrogeneity.
3

---

Joint Amortization Rosenbaum et al. [26] demonstrated heterogeneous reconstruction from unknown poses in a jointly amortized framework on simulated data and with strong priors on the structure given by initial atomic model. To the best of our knowledge, no amortized inference technique has been shown a more general setting on experimental datasets with unknown poses. Here, cryo uses implicit neural representation for the conformational space that can reconstruct experimental datasets with accuracy comparable to the state of the art at a fraction of the compute time.
Methods
3.1 Image Formation Model
In single particle cryo-EM, probing electrons interact with the electrostatic potential created by the molecules embedded in thin layer of vitreous ice (Figure T (a)). During reconstruction, assume that this potential can be broken down into spatially bounded independent potentials (volumes) created by individual molecules. Each volume Vi be by E .I.....N. We assume that the volumes (Vili=t,...., NN are drawn independently from a probability distribution Py supported on oow-dimensional manifold (the conformational space) in the space F(RS, R) of all possible 3D potentials (F (A, B) is the set of functions from A to B). More specifically, we assume there exist d and ; R3 x Rd + R such that Iv is supported on the conformational space (V, 2 E Rd).
mapping from IR to R and is indexed
In the sample, each molecule is in an unknown orientation R, E SO(3) 3x3 in the frame of the observer. The probing electron beam interacts with the electrostatic potential and its projections,
(x,z) + J. (ysss v.
24 dt
(1)
are considered mappings from IR2 to IR. The interaction between the beam and the lens is modeled by the Point Spread Function (PSF) P.- Imperfect centering of the molecule in the image is characterized by small translations t. E IR3. Finally, taking into account signal arising from the vitreous ice into which the molecules are embedded as well as the non-idealities of the lens and the detector, each image Y is generally modeled as
Y, Te, * P. Q2 + 7
(2)
where is the convolution operator, T. the t-translation kernel and 77, white Gaussian noise on 2 35 28
The Fourier-Slice Theorem [3] (FST) avoids the computation of integrals and convolutions in Eq. and is also satisfied with the Hartley transform:
H2D [Q1] = S, Ha 2))],
(3)
where HoD and HoD are the 2D and 3D Hartley transform operators 7 (real minus imaginary parts of the Fourier transform). The "slice" operator S. is defined such that for any V R
+ R.
S. S. (liz, Myy) + V (Ra (kay ky of²)
(4)
Finally, if 8; = Han[Y] and ((,, 2:) = HaD (.., ); the image formation model Hartley space can be expressed as
= C. + S. =4]] + nn
where @ indicates clement-wiss multiplication, C. = H2D [P.] is the Contrast Transfer Function (CTF), St the t-translation operator phase-shift in Fourier space and ń): represents complex white Gaussian noise on R²2 See Supplement E for a discussion on the discretization step.
3.2 Overview of cryOFIRE
Fig. summarizes the architecture of cryoFIRE. Images Y; are fed into an encoder, parameterized by
k/, that predicts pose & (R;, t;) and conformational state 2,. R; is used to rotate a grid of D 3D-coordinates in Hartley space. These coordinates are concatenated with the conformational state before being fed into neural representation of the function Va (in Hartley space), parameterized by 0.

---

[Kez Ky. 0
Go
Ri
k
MLP
Encoder
Zi
ii
w
t;
tc
H2D
Laym
Figure 2: Architecture of cryoFIRE. An encoder predicts the latent variable 2; (conformational state), Rd (rotation) and t; (translation) for each image. These variables are used by physics- decoder that reproduces the image formation model (in Hartley space) and contains a neural representation V of the conformational space. Reconstructed and measured images compared with the symmetric loss.
Vo is a mapping from 33+ to IR and can be seen as a parametric representation of the conformational space in Hartley domain (a manifold of dimension d). Queried for all k in the grid, the neural representation outputs set of real values corresponding to discrete sampling of the slice defined in Eq. 4 Based on the estimated translation t, and given CTF parameters C, the rest of the image formation model described in Eq. (5) simulated to obtain X,, a noise-free estimation of Y. The whole forward model is refered to as To.e. such that X, = Teee(())) Pairs measured and reconstructed images compared using the symmetric loss (see Section (355) and gradients backpropagated throughout the differentiable model in order to optimize both the encoder and the neural representation.
3.3 Discriminative Model
cryoFIRE, an encoder acts as discriminative model by mapping images to estimates Φ; and 2,. The encoder is structured sequentially with the following components:
1. A Convolutional Neural Network (CNN) containing 7 convolutional layers extracts high- level visual features from images and divides the width and height of images by 32. The architecture the CNN is inspired from the first layers of VGG16 [30], known to perform well on visual tasks.
2. A shared Multi-Layer Perceptron (MLP) with 2 hidden layers that outputs a feature y of dimension 256.
3. A conformation MLP that maps y: to
4. Rotation and translation MLPs that map the concatenation of ys and 2; respectively to R and t,. The pose prediction is conditioned 2; since rotations and translations are only defined for a given conformation of the molecule and 2; contains, with small number of dimensions, all the required information to determine the conformational state. The rotation is represented in the 6-dimensional space S x S2 [41 as it was shown to lead to the best results for the rotation prediction in 20
We experimented with both variational (predicting an approximate posterior distribution over latent variables) and non-variational (predicting a point estimate) approaches and did not observe any signif- icant performance difference (although the variational approach is more computationally expensive). Full details about the architecture of the encoder given in Supplement A.
3.4 Generative Model
The generative model is a simulation of the image formation model Hartley space, order to make use of the FST (Section 3.1). The forward pass is differentiable with respect to R2 and t; For each grid of D2 coordinates on the x-y plane Hartley space are rotated by R2 (mapped from S x 2 to 3x3 using Gram-Schmide orthogonalization) via matrix multiplication. Rotated
5

---

points k are positionally encoded with bank of D/2 random frequencies B3 and independently concatenated with the conformational state 2; before being fed into an MLP mapping TR3+ to IR. This neural network is an implicit representation of the function V, where Vi,, =) = HaD [V =)].
Additional implementation details about the neural network are given in Supplement
The CTF is defined by the defocus parameters and the astigmatism angle (provided by the simulator or by external software like CTFFIND [25] for experimental datasets). It is then applied to the slice of size D outputted by the neural representation. Finally, the slice is element-wiss multiplied by Tt.. the Hartley transform of the t.-translation kernel.
3.5 Training Procedure
In the setting where poses are unknown, one cannot distinguish, given a of 2D projections, two molecules with different handedness. This is called the "handedness ambiguity" [27]. Ref. [[6] showed that this leads amortized inference techniques get stuck in local minima where the predicted molecule contains spurious planar symmetries. A solution suggested to alleviate this problem is to use the symmetric loss
Csym Σ Emin min - (Y)) (1 T (Rx (2))))) EEB
(6)
where is a batch of indices and R₂ applies an in-plane rotation of on images (we refer the reader to [16] for an ablation study on the symmetric loss and its instrumental role in the amortized inference of poses). In order to enable the model to converge to consensus volume, we disable the optimization of the conformation MLP at the start of the training. During this "pose-only phase", 2; are randomly sampled from a normalized Gaussian distribution.
4 Results
We evaluate cryoFIRE for ab initio heterogeneous reconstruction and compare it with the state-of- the-art method cryoDRGN2 (39. We first validate that using an encoder to predict poses, instead of performing an exhaustive pose search, enables us to reduce the runtime of heterogeneous recon- struction a synthetic dataset. We show that the encoder is able to accurately predict Φ; and 2; for images it has never processed during training, thereby validating the ability of an encoder-like architecture to amortize the runtime over the size of the dataset. Secondly, we show that cryoFIRE enables the detection of discrete heterogeneity in a dataset, by clustering images in the 2, space Finally, show that our method can perform heterogeneous reconstruction of real data, a first for method that jointly amortizes the estimation poses and conformations.
4.1 Runtime Improvement and Amortization
Experimental Setup We prepare three synthetic datasets of varying sizes small (50k images), medium (500k images) and large (5M images) (Supplement C). Images are generated using a simulation of the image formation model ground truth volumes of the 80S ribosome. Each dataset contains a mix of projections with 90% sampled from the one the volumes (major) and the rest from the other volume (minor). CTFs are drawn randomly with a log-normal distribution over the defocus range. We add Gaussian noise with variance a² = 10 to the normalized projections (SNR=- 10dB). With cryoFIRE, seen 1.5M images. We compare our method with cryoDRGN2 (39 with default parameters. Pixels' intensities are set 0 outside of circle of radius 32 pixels for pose search. With cryoDRGN2. pose search is done every 5 epochs for the small datasets and every epoch for the medium and large datasets. We train the models on a single NVIDIA A100 SXM4 40GB GPU. Images of size D 128 are fed by batches of maximum sizes (128 for cryoFIRE, 32 for cryoDRGN2) and loaded on the fly using
fix d = 8 and activate the conformation MLP after the model has
on 16 CPUs. The model optimized with the ADAM optimizer [10] and learning rate of 10-4 Convergence is assessed via visual inspection of the pose prediction accuracy. Once convergence reached on the train set, latent variables can be estimated on lest set of 10k images in order to validate that the encoder is not just memorizing the training set. With cryoDRGN2, poses from the test set are estimated with a randomly-initialzzz pose search on the model obtained at train time.

---

Table 1: Heterogeneous reconstruction on datasets of varying sizes. Each dataset contains mix of images, 90% from the major volume and 10% from the minor volume. After convergence on a train set, each method estimates the latent variables (R;, t, 21) on a test set. We report the resolution (Res. in pixels, }), the confusion error (4), the rotation accuracy (Rot. in degrees, }) and the translation accuracy (Trans. in pixels, 4). Convergence is detected on the pose accuracy (ep. epochs).
Dataset / Method
Time
Confusion
Res. (major/mnizrr)
Rot. (Med/MSE)
Trans. (Med/MS))
Small (Trin: 50k / Test: 10k)
cryoDRGN2 (train)
1:21h (20 ep.)
0.00005
2.4/2.8
0.8 0.8
0.007 0.01
cryoFIRE (train)
1:33b (70 ep.)
0.0004
2.6/
2.3/2.6
0.09/ 0.1
cryoDRGiNZ (lest)
3 min. (1 ep.)
0
0,8 0.8
0.006 0.01
cryoFTRE (test)
11 sec. (1 ep.)
0.001
2.6/
0.2/0.3
Medium (Train:
10k)
cryDRRGN2 (train)
5:10 (2 ep.)
0.002
2.5/3.
0.8
0.007 0.01
cryoFIRE (train)
1:28h (7 ep.)
0.0008
2.7/3.2
2.7
/
cryoDRGN2 (lest)
3 min. (1 cp.)
0.0001
0,9
0,007 0,01
cryoFIRE (test)
11 sec. (1
0.0001
2.5/
0.1 0.2
Large (Train: SM/ /Test: 10k)
cryoDRGN2 (train)
21:37h (1 ep.)
0.002
2.3/2.6
0.81 1.6
0.01 1.2
cryoFTRE (train)
1:55h (1 op.)
0.0002
2.3/2.7
1.5/ 1.7
0.1 1.0
cryoDRGN2 (lest)
min. (1 ep.)
1.0/
0,007 0.1
cryoFIRE (test)
sec. (1 ep.)
1.2/14
0.1/
Metries We compute the median and the mean the angular error on the view direction, in degrees, and the median and mean square error of the predicted t,, in pixels. Before computing pose errors, rigid 6D-body alignment is applied on the set of predicted poses to align the model into the same global reference frame the ground truth volume. We perform Principal Component Analysis (PCA) on the set of predicted 2 's and, looking at the first principal component, we group images in two classes, according the cluster they fall into. The confusion error is defined as the ratio of misclassified images over the total number images. After convergence, we generate two volumes by taking the centroids of the clusters in space and using them as inputs for the neural network Vi, on them. We report the resolution (in pixels) obtained on the major (90%) and on the minor (10%) volumes. The reported resolution is the inverse of the maximum frequency for which the Fourier Shell Correlation above 0.5.
At train time, the average time for processing one image is 19 ms (5:10h / (2 500 000)) with cryoDRGN2 vs. 1.5 ms (1:28h / (7 x 500 000)) for cryoFIRE. (The runtime cryoDRGN2 is lower with the small dataset because pose search is only done every 5 epochs). At test time, only the encoder of cryoFIRE is used, which further decreases the runtime per image (1.1 ms). The accuracy obtained at test time on Φ; and z; validates that method effectively amortizes the estimation of latent variables, and the errors on the test sets decrease with the size of the dataset provided for training. The accuracy of predicted poses is slightly worse with cryoFIRE than with an exhaustive pose search strategy (see Supplement E for a comparison with fully non-an an-ommottzzed method (24)). The quantitative resolution of the reconstructed volumes remains however very similar with both methods (see Supplement E for a qualitative comparison).
Results We report the runtime and quantitative metrics in Table I
4.2 Heterogeneous Reconstruction of Synthetic Datasets
Experimental Setup We prepare three heterogencous datasets of 100k images of the Plasmodium falciparum 80S ribosome (Supplement C). The three datasets contain an equal mix of projections from K E {2, 3, 10} different ground truth volumes, generated sequentially along reaction path that corresponds to a rotation of the small ribosomal subunit relative to the large ribosomal subunit identified in (37. The image formation model is simulated with the same parameters as in Section 4.1
We run FIyO with 8 and perform a PCA on the set of predicted 2, 's after 95 epochs.
Results In Fig. we plot the first principal component (PC1) of the set of z;'s vs. the ground truth index of the volume each image was generated from. With the dataset containing K conformations, images can be clustered in K groups, based the value of PCI (K e {2, (2. Looking for the
7

---

(a) 2 classes
(b)
3 closses
(c)
10 classes
(d)
Confratoo
etroct
Conviattio:
Cualering trrmbald Consitional mears Conditiona uld
Conditional
Conditional
2123433 Ground Buth Conformeiinn
Ground Truth Conformtiinn
Ground Conformatioo
Figure 3: Conformational heterogeneity analysis on synthetic datasets. On (a-b-c), the y-axis shows first principal component (PC1) among the 2; s. We indicate means + standard deviations of the PCIs, conditioned on the ground truth conformation. On (a-b) dashed lines are the threshold used to compute the confusion error. Spearman correlation [13] is indicated in (c). (d) is a superposition of two volumes obtained by sampling two points along the PCI axis for the experiment 2 classes.
optimal thresholds on PC1, we compute confusion error of 1 x 10-3 (resp. 2x 1x10) the 2 classes (resp. classes) dataset. On the classes dataset, the conformations cannot be separated but the value of PCI embeds information about the movement of the ribosome, as indicated by the Spearman correlation [ between the index of the ground truth conformation and PCI. Sampling different values for PCI (and setting the other components to zero), we can use the neural representation V to generate set of volumes from set of conformational states 2;. Qualitative results are given in Fig. 3 (d), showing an accurate reconstruction of the movement of a subunit of the ribosome.
4.3 Heterogeneous Reconstruction of an Experimental Dataset
We use the publicly available dataset EMPIAR- 10180 21 of pre-catalytic spliccosome (Supplement C). We run cryoFIRE with d = 8, activate the conformation MLP after 30 epochs and train for a total of 100 epochs. Results from cryoDRGN2 39 and cryoDRGN- on the same dataset are obtained from [39].
Experimental Setup
38
Results The first two components of PCA on the of predicted 2,'s is shown in Fig 4 By traversing the conformational space along the direction of PC1 and generating volumes on set of points, cryoFIRE generates a trajectory, showing large flexing motion of the spliceosome. CryoFIRE qualitatively recovers the non-uniform distribution of viewing directions, avoids the local minima cryoDRGN-BNN BNB falls into, and reconstructs volumes which qualitatively match the state of the art. See Supplement D for more results and reconstruction on an experimental dataset of the 80S ribosome (EMPIAR-10028 (36)).
Discussion
As the resolving power of sampling method scales with the size of the dataset, an increasing number of images needs to be collected in cryo-EM in order to extract meaningful structural information. CryoFIRE answers the need for reconstruction methods that scale appropriately with the dataset size: heterogeneous ab initio reconstruction on dataset of 5M images can be performed within 2 hours. CryoFIRE is also the first reconstruction method to perform joint amortization poses and conformations on experimental cryo-EM dataset, opening the door to fast analysis of structural heterogeneity on real datasets.
The interpretability of the conformational space remains an open question. Since distances in this space are not meaningful (z, conditions the volume via nonlinear decoder), the probability distribution cannot be straightforwardly interpreted in that space. In the case of discrete heterogeneity. general quantities like the number of states or the relative populations of the states can be reliably inferred from the conformational space by clustering the predicted conformational states. In the case of continuous heterogeneity, dimensionality reduction methods like PCA can highlight directions of maximal variance and, by visual inspection of the reconstructed volumes, cryoFIRE can help understanding the main degrees of freedom of continuously deformable molecule. However, "small" and "large" deformations can stem from similar changes in the conformational space and no physically

---

(1)
12)
(3)
(4)
(5)
(1)
[3]
15)
PC-1
GyoDRGN-BNB
oyoDRGN2
ciyoTIRE
Raía ance
8
Figure 4: Ab initio heterogeneous reconstruction of an experimental dataset of the pre-catalytic spliceosome (EMPIAR--01S1 10180 (21)). (Top left) Distribution of the first principal components of & and sampling points along the first component. (Top right) Structures generated by traversing along PC1. The outline indicates the edges of volume (5). (Bottom) Distribution of the viewing directions. The "reference" shows the viewing directions published in (2I)
interpretable notion of distance is currently associated with the conformational space: providing this space with physically interpretable metric is an interesting avenue for future work.
is worth highlighting that the parameterization proposed here is not uniquely defined: a rotation of the molecule equivalently be represented by R3 or by change of conformation state 2;- Empirically, cryoFIRR decorrelates poses and conformation by relying on a "pose-only phase" at the beginning of training, but nothing explicitly prevents 2; from containing information about R;- Future work could explore improvements in the optimization procedure to further enforce minimal mutual information between poses and conformation.
Finally, we note that accuracy of pose estimation (especially on translations) lower with cryoFIRE than with exhaustive pose search methods. Inaccurate translation prediction lead to "blurry" reconstructions and therefore limits the resolution of reconstructed volumes. This is a direct reflection of the fact that amortization maximizes non-tight lower-bound of the likelihood, allowing for faster inference at the cost of accuracy (5]. Future work could investigate a hybrid approach that would estimate poses with encoder the beginning of training and switch, the end, local pose search initialized from encoder-estimate poses.
Broader Impacts Statement The cryo-EM field is generating data at rapidly increasing pace, potentially leading to wasteful storage and inefficient computing. The design more efficient methods such as CryoFIRE provides path to mitigate sustainability risks. In order to better compare modern reconstruction methods, careful attention must be paid toward designing common datasets and evaluation metrics. By providing open-source implementation of CryoFIRE upon publication, together with benchmark metrics, we hope to make cryo-EM research accessible to a broader class of researchers. We also acknowledge that. by pushing further the boundaries of biological research, our method may facilitate the development of harmful biologics, however this is outweighed by its significant societal benefits. We encourage the cryo-EM community to play wisely with cryoFIRE.
Acknowledgements
This work was supported by the U.S. Department of Energy, under DOE Contract No. DE-AC02- 76SF00515 and the SLAA L.RDD program. We acknowledge the use the computational resources at the SLAC Shared Scientific Data Facility (SDF) and Princeton Research Computing.

---

References
[1] Baldwin, P.R.. Tan, Y.Z.,, Eng. E.T., Rice, W.J., Noble, A.J, Negro, C.J., Cianfrocco, M.A., Pot- C.S., Carragher, B.: Big data in cryoem: automated collection, processing and accessibility of em data. Current Opinion in Microbiology 43, 1-8 (2018)
[2] Bepler, T., Zhong, E., Kelley, K., Brignole, E.,, Berger, B.: Explicitly disentangling image content from translation and rotation with spatial-vae. Advances in Neural Information Processing Systems 32 (2019)
[3] Bracewell, R.N.: Strip integration in radio astronomy. Australian Journal of Physics 198-217 (1956)
[4] Chen, M., Ludtke, S.J.:
Deep learning-based mixed-dimensional gaussian mixture model for characterizing variability in cryo-em. Nature Methods 18(8), 930-936 (2021)
[5] Donnat, C.,, Levy, A., Poitevin, F., Miolane, N.: Deep generative modeling for volume recon- struction in cryo-electron microscop. arXiv preprint arXiv:2201 .02867 (2022)
[6] Gupta, H., McCann, M.T., Donati, L., Unser, M.: Cryogan: reconstruction paradigm
for single-particle cryo-em via deep adversarial learning. IEEE Transactions on Computational Imaging 7, 759-774 (2021)
[7] Hartley, R.V.: A more symmetrical fourier analysis applied to transmission problems. Proceed- ings of the IRE 30(3), 144-150 (1942)
[8] Herreros, D., Lederman, R.R., Krieger, J., Jiménez-Moreno, A., Martínez, M., Myška, D.. Strelak, D., Filipovic, J., Bahar, I., Carazo, J.M., et al.: Approximating deformation fields for the analysis of continuous heterogeneity of biological macromolecules by 3d zernike polynomials. IUCrJ 8(6) (2021)
[9] Jin, Q., Sorzano, C.O.S., De La Rosa-Trevin, J.M., Bilbao-Castro, J.R., Jüüiez-Ramrez, R. Llorca, O., Tama, F,, Jonić, S.: Iterative clastic 3d-to-2d alignment method using normal modes for studying structural dynamics of large macromolecular complexes. Structure 22(3), 496-506 (2014)
[10] Kingma, D.P., Ba, J.: Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6988 (2014)
[11] Kingma, D.P., Welling, M.: Auto-encoding variational bayes. arXiv preprint arXiv:1312.611 (2013)
[12] Kingma, D.P., Welling, M.: An introduction to variational autoencoders. arXiv preprint arXiv:1906.0269 02699 (2019)
[13] Kokoska, S., Zwillinger, D.: CRC standard probability and statistics tables and formulae. Cre Press (2000)
[14] Lawler, E.L., Wood, D.E.: Branch-and-bound methods: A survey, Operations research 14(4), 699-719 (1966)
[15] Lederman. R.R., Andén, J., Singer, A.: Hyper-molecules: on the representation and recovery of dynamical structures for applications in flexible macro-molecules in cryo-em. Inverse Problems 36(4), 044005 (2020)
[16] Levy, A., Poitevin, F., Martel, J., Nashed, Y., Peck, A., Miolane, N., Ratner, D., Dunne, M., Wetzstein, G.: Cryoai: Amortized inference poses for ab initio reconstruction of molecular volumes from real cryo-em images. arXiv preprint arXiv:2203.088 001808888888888 (2022)
[17] Miolane, N., Poitevin, F, Li, Y.T., Holmes, S.: Estimation of orientation and camera parameters from cryo-electron microscopy images with variational autoencoders and generative adversarial networks. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops. pp. 970-971 (2020)
10

---

[18] Nakane, T., Kimanius, D., Lindahl, E., Scheres, S.H.: Characterisation of molecular motions in cryo-em sangle-partice data by multi-body refinement in relion. elife 7, e36861 (2018)
[19] Namba, K., Makino, F: Recent progress and future perspective of electron cryomicroscopy for structural life sciences. Microscopy 71, i3-il4 (2022)
[20] Nashed, Y.S., Poitevin, F., Gupta, H., Woollard, G., Kagan, M., Yoon, C.H., Ratner, D.: Cryoposenet: End-to-end simultaneous learning of single-particle orientation and 3d map recon- struction from cryo-electron microscopy data. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 4066-4076 (2021)
[21] Plaschka, C., Lin, P.C., Nagai, K.: Structure of a pre-catalytic spliceosome. Nature 546(7660), 617-621 (2017)
[22] Punjani, A., Fleet, DJJ.: 3d flexible refinement: structure and motion of flexible proteins from cryo-em. BioRxiv (2021)
[23] Punjani, A., Fleet, DJJ.: 3d variability analysis: Resolving continuous flexibility and discrete heterogeneity from single particle cryo-em. Journal of Structural Biology 213(2), 107702 (2021)
[24] Punjani, A., Rubinstein, J.L., Fleet, DJJ., Brubaker, M.A.: cryosparc: algorithms for rapid unsupervised cryo-em structure determination. Nature methods 14(3), 290-296 (2017)
[25] Rohou, A., Grigorieff, N.: Ctffind:: Fast and accurate defocus estimation from electron micrographs. Journal of structural biology 192, 216-22 221 (2015)
[26] Rosenbaum, D., Garnelo, M., Ziclinski, M., Beattic, C., Clancy, E., Huber, A., Kohli, P., Senior, A. W., Jumper, J., Doersch, C., et al.: Inferring a continuous distribution of atom coordinates from cryo-em images using vaes. arXiv preprint arXiv:2106. 14108 (2021)
[27] Rosenthal, P.B., Henderson, R.: Optimal determination of particle orientation, absolute hand, and contrast loss in single-particle electron cryomicroscopy. Journal of molecular biology 333, 721-745 (2003)
[28] Scheres, S.H.: RELION: Implementation of Bayesian approach to cryo-EM structure determi- nation. Journal of Structural Biology 180, 519-530 (2012)
[29] Scheres, S.H., Gao, H., Valle, M., Herman, G.T., Eggermont, P.P.. Frank, J., Carazo, J.M.: Dis- entangling conformational states macromolecules 3d-em through likelihood optimization. In: Single-Partiche Cryo-Electron Microscopy: The Path Toward Atomic Resolution: Selected Papers of Joachim Frank with Commentaries, pp. 413-415. World Scientific (2018)
[30] Simonyan, K., Zisserman, A.; Very deep convolutional networks for acrge-scale image recogni- tion. arXiv preprint arvi: 1409.1556 (2014)
[31] Singer, A., Sigworth, F.J.: Computational methods for single-particle electron cryomicroscopy. Annual Review of Biomedical Data Science 3, 163-190 (2020)
[32] Tagare, H.D., Kucukelbir, A., Sigworth, F.J., Wang, H., Rao, M.: Directly reconstructing principal components of heterogeneous particles from cryo-em images. Journal of structural biology 191(2), 245-262 (2015)
[33] Tancik, M., Srinivasan, P., Mildenhall, B.. Fridovich- Keil, S., Raghavan, N., Singhal, U., Ramamoorthi, R., Barron, J., Ng, R.: Fourier features let networks learn high frequency functions in low dimensional domains. Advances in Neural Information Processing Systems 33, 7537-7547 (2020)
[34] Uddin, M.R., Howe, G., Zeng. X., Xu, M.: Harmony: A generic unsupervised approach for disentangling semantic content from parameterized transformations. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 20646-20655 (2022)
[35] Vulović, M., Ravelli, R.B.G., van Vliet, L.J., Koster, A.J., Lazić, I., Lücken, U., Rullgård, H., Öktem, O., Rieger, B.: Image formation modeling in cryo-electron microscopy. Journal of Structural Biology 183, 19-32 (2013)

---

[36] Wong, W., Bai, X.c., Brown, A., Fernandez, I.S., Hanssen, E., Condron, M., Tan, Y.H., Baum, J., Scheres, S.H.: Cryo-em structure of the plasmodium falciparum 80s ribosome bound the anti-protozoan drug emetine. Elife 3, e03080 (2014)
[37] Zhong, E.D., Bepler, T., Berger, B., Davis, J.H.: Cryodrgn: reconstruction of heterogeneous cryo-em structures using neural networks. Nature methods 18(2), 176-185 (2021)
[38] Zhong, E.D., Bepler, T., Davis, J.H., Berger, B.: Reconstructing continuous distributions of 3d protein structure from cryo-em images. arXiv preprint arXiv:1909.05211 (2019)
[39] Zhong, E.D., Lerer, A., Davis, J.H., Berger, B.: Cryodrgn2: Ab imitio neural reconstruction of
3d protein structures from real cryo-em images. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 4066-4075 (2021)
[40] Zhong, E.D., Lerer, A., Davis, J.H., Berger, B.: Exploring generative atomic models in cryo-em reconstruction. arXiv preprint arXiv: 107 .01331 (2021)
[41] Zhou, Y., Barnes, C., Lu, J., Yang. J., Li, H.: On the Continuity of Rotation Representations in Neural Networks. arXiv: 1812.07035 (2020)
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] see Sec.
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or URL)? [No] We plan on making code available.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [No]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes] see Sec. 4
4. If you are using existing assets (c.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [No]
(c) Did you include any new assets either in the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're
using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [No]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review
Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
12


Reviews:
[{"review_id": "vHQf0C3KCAl", "reviewer": "Reviewer_hGi9", "summary": "The paper proposes a novel method to infer both pose and conformational states of protein structures from raw cryo-EM images in a way that is amortized over the size of the dataset. The authors compared the proposed method with cryoDRGN2 on several synthetic and experimental datasets. They demonstrated that their method requires less time compared to cryoDRGN2 by sacrificing little accuracy. \n", "questions": "1. How good is cryoFIRE in conformation estimation compared to other methods?\n2. Why Rosembaum et al.’s method can not be applicable to a real dataset?\n3. I would suggest taking a look at methods like SpatialVAE [1] and Harmony [2] that disentangles the rotation and translation from cryo-EM images. How are those methods related to yours? The setup of the discriminative model (Section 3.3) looks very similar. \n\n[1] Bepler, Tristan, et al. \"Explicitly disentangling image content from translation and rotation with spatial-VAE.\" Advances in Neural Information Processing Systems 32 (2019).\n\n[2] Uddin, Mostofa Rafid, et al. \"Harmony: A Generic Unsupervised Approach for Disentangling Semantic Content From Parameterized Transformations.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.\n", "limitations": "The authors discussed the limitations and possible negative social impact of the work.\n", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is well-written and overall easy to follow. The motivation and importance of the problem are well discussed. The paper provides a thorough overview of the literature. The proposed method is novel to the best of my knowledge.\n\nHowever, there are several issues with the experiments, baselines, and results. The authors mentioned that Rosenbaum et. al.’s method was the only one to perform amortization for both pose and conformation estimation. In such a case, that particular method should have been kept as a baseline, not cryoDRGN2 since it amortizes for conformation estimation only. Even the comparison with cryoDRGN2 is not very impressive. It is quite usual that methods with an exhaustive search for pose would take more time with more accuracy. If the proposed method, cryoFIRE could achieve better or similar accuracy with less time then it would have been much more impressive. But cryoFIRE performs much worse than cryoDRGN2 even in large datasets.\n\nNo comparison has been showed for conformation estimation. Is cryoFIRE as good as cryoDRGN2 for conformation estimation as well? Can it recover some conformations that can not be recovered by cryoDRGN2 or 3DVA or E2GMM?\n\nThe paper mentions Rosembaum et al’s method only worked on the simulated dataset, but it did not discuss why it can not be applied to real datasets. What makes it inapplicable for real datasets but applicable for realistically simulated datasets? Or is it the case that the authors are the first to apply the method to the experimental dataset, this is clearly not a big contribution. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "W519kzeyg0A", "reviewer": "Reviewer_7QvN", "summary": "The paper proposes an ab initio heterogeneous reconstruction method for SPA cryoEM, named cryoFIRE. There are multiple advantages to the proposed approach and most of them come from the ability to jointly estimate both poses and conformations, allowing to skip computationally expensive step of pose search. The authors show the effectiveness of the method on multiple datasets (including one experimental) and compare the approach with SOTA of cryoEM NN-based heterogeneous reconstruction (cryoDRGN/2). The authors promise to provide an open-source implementation of cryoFire upon publication. ", "questions": "1. The datasets used for the experiments have quite high SNR (especially synthetic ones), at least judging by the samples in the figure S3. Did you explore whether there is a limit of SNR at which the method does not work?\n2. Even though Table 1 and Table S2 show that cryoFIRE is slightly bit less precise than cryoDRGN2 and cryoSPARC, the speedup in training and especially testing is more than worth the tradeoff. However, what do you think is the reason for decreased precision? What is the number of parameters in the models and have you experimented with changing number of parameters/layers?\n", "limitations": "The authors adequately addressed the limitations and potential negative societal impact of their work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper is written cleanly and I especially enjoyed the succinct formal description of cryoEM image formation.\n2. The method is novel and presents the first instance of amortised inference for ab initio cryoEM heterogeneous reconstruction.\n3. Joint estimation of poses and conformations brings one order of magnitude speedup, scaling much better with the ever-increasing amount of data produced by cryoEM.\n\nWeaknesses:\n1. As the authors themselves point, parameterisation calculated with cryoFIRE is not uniquely defined (\"a rotation of the molecule can be equivalently represented by R_i or by a change of conformation state z_i\"). The proposed workaround in the paper is to use training schedule and have \"pose-only\" phases at the start of training and then every N epochs. Additional work is required to be able to decouple rotation from conformation without training schedules.\n2. cryoFIRE is directly compared with cryoDRGN and cryoDRGN2 (Table 1), as well as cryoSPARC (Table S2). I don't think that this represents software frameworks currently used by cryoEM groups.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "IJsKGv2yc03", "reviewer": "Reviewer_fKm9", "summary": "This manuscript introduces a new method for the reconstruction of continuous variability in cryogenic electron-microscopy (cryo-EM) based on amortized inference and neural network autoencoders. While amortized inference and autoencoders have previously been used for different aspects of this problem, the proposed method provides a more comprehensive approach. Specifically, previous methods have relied on a consensus model for pose estimation or performed expensive pose estimation as part of an alternating refinement process where the main focus has been on estimating the latent states of the molecule. The proposed method performs the pose estimation at the same time as the latent state estimation, achieving significant speedups with respect to the state of the art. The method is validated on numerous synthetic datasets and one experimental dataset.\n", "questions": "– Please provide a formal definition for amortized inference early on (Section 1 or 2) for context.\n– What does “disentanglement” mean in the context of cryoVAEGAN (line 108)?\n– Another approach for joint estimation of pose and latent state (similar to 3DVA of cryoSPARC) is presented by the hypervolume model of Lederman et al. (2020). It would be good to compare this approach with the proposed model.\n– What is calligraphic F in line 134?\n– Section 3.1 describes an image formation model entirely in the continuous domain. Please discuss issues of discretization (since we are dealing with discrete data, after all). How is this taken into account?\n– Where does D come from in line 156?\n– How and why is the rotation “represented in the 6-dimensional space S² × S²”?\n– “defoci” on line 196 should be “defocus” (plural is on “parameters”).\n– What is meant by major and minor on line 224? (This is explained in the supplement, but needs to be briefly described in the main text.)\n– What is meant by “images are masked with a circle of radius 32 pixels for pose search” on line 228?\n– Similarly, what is meant by “pose search is done every 5 epochs”? Given that there is no search per se, does this mean that the inferred pose parameters are held fixed (i.e., not recomputed using the corresponding MLP) for 5 epochs in a row before the corresponding weights are updated?\n– Why split the data into training and testing? Is this merely for validation purposes or part of the overall workflow? If it is part of the overall workflow, how is a user supposed to perform this split and how is the final reconstruction obtained?\n– On line 243, the error is the “ratio of misclassified images” to what quantity?\n– In Table 1, why is resolution not reported for the test sets?\n– What is meant by “conditioning” on line 245?\n– What is K on line 267?\n– The Spearman correlation on line 271 is calculated between what variables?\n– Why is only PCA of dimension 1 considered in Section 4.2.\n", "limitations": "As stated above, the high SNRs of the synthetic data, limited experimental validation, and lack of validation for translational pose estimation in the experimental data detract from the overall claims of the manuscript and should either be resolved or acknowledged in the text.\n\nThe authors discuss potential negative societal impact of their work (and the field in general) quite well.\n", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The main components of the proposed method have been proposed earlier: amortized inference for pose estimation has been introduced by Miolane et al. [15] and for latent state estimation by Zhong et al. [33]. Joint estimation of both pose estimation has been performed by Rosenbaum et al. [24], but only for synthetic data. These previous results are accurately described in the manuscript. As such, the original work in this manuscript consists of combining these existing components into a functioning whole, which is no easy feat. It also does this while achieving significantly reduced running times compared to the state of the art (albeit with a slight loss of accuracy).\n\nThe method is promising and the authors provide numerical results to support their claims. That being said, these numerical results do not completely validate the method as they currently stand. First, the synthetic data generated has a very high signal-to-noise ratio (SNR) by cryo-EM standards. This can be seen most clearly by comparing the sample projections in Figure S3, where two synthetic projection images are compared to an experimental projection image. The noise level of the former is much lower compared to the latter. In the text, the authors simply mention that the variance of the noise applied is one without providing a measure of the variance of the clean images. Judging by eye, I would guess that the SNR for the synthetic images is close to one, while realistic cryo-EM data typically has SNRs ranging from one hundredth to one tenth. It is important to note that despite these high SNRs, the proposed method performs worse than the state-of-the-art, with rotation errors on the order of 2× and translation errors typically on the order of 10×.\n\nThe results on the experimental data is also concerning. In the supplementary material, the authors mention that “particle images are shifted by their published poses, since the particles in this dataset are significantly out of center”. In other words, the pose estimation is not fully validated on experimental data (specifically, while the rotational component of pose estimation is validated, the translational component is not). While it may well be the case that the particles in this dataset are significantly out of center, this is a common occurrence in many cryo-EM datasets, and it is something that a robust method should be able to handle. Given the fact that an important contribution of the proposed method is its validation on experimental data, this discrepancy significantly detracts from that contribution. Furthermore, the fact that only one experimental dataset is tested also reduces the scope of validity (compared to the cryoDRGN and Relion multi-body refinement papers, which present results on three or more experimental datasets).\n\nThe overall presentation of the work is clear, with most concepts being well defined. Given the importance put on amortized inference, it would be good to specify early on what is meant in a formal manner to reduce any possible confusion about the term. In a similar vein, a definition of the Hartley transform (and its relative advantages over the Fourier transform) would be a good addition.o\n\nAs stated above, the significance of this paper is the joint amortized inference of pose and latent space applied to experimental data. As such, it provides an important methodological milestone for the processing of cryo-EM data. However, the numerical results do not completely back up these claims. Additionally, the authors also do not make clear what it is about their setup that makes the method successful. In other words, what has been missing from previous attempts at joint inference (such as the work of Rosenbaum et al. [24]) that is addressed by the current work?\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "8Cl4k30AB9r", "reviewer": "Reviewer_ypJe", "summary": "The paper describes a method, CryoFIRE, to recapitulate the 3D structure of a biomolecule from cryo-EM images using coordinate-based neural model that enables continuous representation. This is inspired by recent advancements of CryoDRGN/CryoDRGN2 for ab initio 3D molecular reconstruction. Different from existing methods, the authors propose to use an encoder to jointly estimate the particle poses and conformation state (a latent vector) from a set of 2D projections. The predicted poses and latent vector are then fed into a feedforward neural network to aggregate each particle image into an implicit representation of the scattering potential volume. The encoder and decoder are trained in an end-to-end fashion using L2 distance. This method achieves comparable results on both simulated and experimental datasets to SOTA. In addition, the paper shows that such autoencoder-based method can significantly speed up the reconstruction time compared to CryoDRGN2 when using a large number of 2D projections.", "questions": "Here are my detailed review comments.\n\n1), One of the claimed contributions of this work is confusing, and some clarifications will be helpful. In specific, this work claims that “the first instance of amortized inference for ab initio heterogeneous reconstruction of experimental cryo-EM datasets.” Seems that CryoFIRE uses the same publicly available dataset as cryoDRGN2, and directly takes the results from cryoDRGN2.\n\n2), The discussions around the differences and contributions over existing CryoAI is missing. Both methods use an encoder to capture the high-level poses of particle images and symmetrized loss.  In addition, CryoAI proposes to use two MLPs to represent the complex signals as phase and magnitude separately, while this work seems not adopt this strategy. What results in this different?\n\n3), The rational and advantages of using the symmetrized loss during training in Sec 3.5 is not very instructive at its current state. A deeper understanding and intuition with some ablation studies will be helpful. In addition, how is the symmetrized loss generalizable? Would other training losses such as the statistical distances used in CryoDRGN benefit from this in-plane constrain?\n\n4), While explicitly estimating the poses and amortized variations for parametrizing 3D structure is an interesting direction, recent work in deep learning such as CryoGAN [R1], for single particle cryo-EM has attempted to recover the 3D structure without explicitly recovering the poses using adversarial networks. It may be worth to discuss the relations and differences in the revision. \n\n5), Consider explaining CryoDRGN2 in detail given that is the state-of-the-art technique and to which all the CryoFIRE results are compared.\n\n6), What would runtime comparisons be for cryoDGRN, for an unsupervised heteregeneous reconstruction?\n\n7), Consider adding more description of the neural network architecture that is the main contribution of this work (more than what is reported on page 5 and page S1 in the supplement).\n\n[R1] Gupta et al., \"CryoGAN: A New Reconstruction Paradigm for Single-Particle Cryo-EM Via Deep Adversarial Learning,\" in IEEE TCI, 2021.\n\nMinor comments:\n\nFigure 1 (c) is not very informative; it is hard to be linked to Figure 2.\n", "limitations": "The limitations of current method are missing. Considering single particle cryo-EM is an important but challenging problem, some discussions will be helpful (future work).", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThe paper presents a novel approach for single particle cryo-EM that quantitatively and qualitatively works well. In addition, the paper is overall well written, and the method is well evaluated on experimental dataset. \n\nWeaknesses:\n\nSome claimed contributions and technical details might need to be clarified and discussed.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept