CryoSPIN: Improving Ab-Initio Cryo-EM Reconstruction with Semi-Amortized Pose Inference
Authors: Shayan Shekarforoush
David B. Lindell
Marcus A Brubaker
David J. Fleet
Venue: NeurIPS 2024 poster (2024)
Keywords: Cryo-EM 3D reconstruction
Pose estimation
Semi-Amortization
Multi-choice learning
Primary area: machine_learning_for_other_sciences_and_fields
TL;DR: We propose a new approach to ab-initio cryo-EM 3D reconstruction using semi-amortization to accelerate pose convergence and multi-head pose encoder to handle pose uncertainty.

Abstract: Cryo-EM is an increasingly popular method for determining the atomic resolution 3D structure of macromolecular complexes (eg, proteins) from noisy 2D images captured by an electron microscope. The computational task is to reconstruct the 3D density of the particle, along with 3D pose of the particle in each 2D image, for which the posterior pose distribution is highly multi-modal. Recent developments in cryo-EM have focused on deep learning for which amortized inference has been used to predict pose. Here, we address key problems with this approach, and propose a new semi-amortized method, cryoSPIN, in which reconstruction begins with amortized inference and then switches to a form of auto-decoding to refine poses locally using stochastic gradient descent. Through evaluation on synthetic datasets, we demonstrate that cryoSPIN is able to handle multi-modal pose distributions during the amortized inference stage, while the later, more flexible stage of direct pose optimization yields faster and more accurate convergence of poses compared to baselines. On experimental data, we show that cryoSPIN outperforms the state-of-the-art cryoAI in speed and reconstruction quality.

CryoSPIN: Improving Ab-Initio Cryo-EM Reconstruction with Semi-Amortized Pose Inference
Shayan Shekarforoush 1.2 shayan@es. toronto.... edu
David B. Lindell lindell@cs toronto. edu
Marcus A. Brubaker 2,33 mabθeecs, yorku ca
David J. Fleet 2,4
fleet8cs. toronto edu
University of Toronto
2 Vector Institute
3 York University *Google Deep.Mind
Abstract
Cryo-EM is increasingly popular method for determining the atomic resolution 3D structure of macromolecular complexes (eg, proteins) from noisy 2D images captured by an electron microscope. The computational task is to reconstruct the 3D density of the particle, along with 3D pose of the particle in each 2D image, for which the posterior pose distribution highly multi-modal. Recent developments in cryo-EM have focused on deep learning for which amortized inference has been used to predict pose. Here, we address key problems with semi-amortized method, cryoSPIN, in which
this approach, and propose reconstruction begins with amortized inference and then switches to a form auto-decoding to refine poses locally using stochastic gradient descent. Through evaluation on synthetic datasets, we demonstrate that cryoSPIN able to handle multi-modal pose distributions during the amortized inference stage, while the later, more flexible stage of direct pose optimization yields faster and more accurate convergence of poses compared to baselines. On experimental data, we show that cryoSPIN outperforms the state-of-the-art cryoAl in speed and reconstruction quality. Project Webpage.
Introduction
Single particle electron cryo-microscopy (cryo-EM) has gained popularity among structural biologists as powerful experimental method for determining the 3D structure of macromolecular complexes, such as proteins and viruses, unlocking our understanding of biological function at the scale of the cell. Thanks to pioneering advances in hardware and data processing techniques, cryo-EM has enabled reconstruction of challenging structures at atomic or atar-itomic resolution [1].
During cryo-EM experiment, 10a-107 particle images, each containing instance of the target bio-molecule, are acquired using a transmission electron microscope. From that particle stack, the goal is to reconstruct the unknown 3D density map [2]. This ab-initio reconstruction task presents some challenges. First, the 3D pose (orientation and position) the particle in each image is unknown, and must be inferred along with the 3D structure. Second, low electron exposures are used to limit radiation damage to the particles (e.g., see Fig. 1). But this reduces the signal-to-noise ratio (SNR), obscuring high-resolution image detail and hindering pose and structure estimation. Third, bio-molecules are typically non-rigid and exhibit structural variations within a sample. Hence, for such heterogeneous datasets, it is crucial to account for the conformational variability in order to obtain high-resolution reconstruction [3-6].
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Stage 1: Auto-ing
L
H
Encoder
min
L
44
Stage 2: Auto-decoding
Decoder
Lock up Tadlo
8 (8
L
Hc
8Thh cxppt
6t (a yi
Figure 1: CryoSPIN consists of two stages: (i) auto-encoding stage where image encoder equipped with multiple heads maps the input image to the pose candidate set { Φx followed by computing projections by slicing through the volume decoder in Fourier space based the pose set. These projections are compared with the input image and the one with the minimum error is used. (ii) An auto-decoding stage where pose parameters in axis-angle representation are stored for all images. The same volume decoder is used to obtain projections, and the reconstruction loss is computed for single projection.
The development of methods in cryo-EM has recently focused on deep learning. CryoDRGN [5] proposed image encoder-volume decoder architecture model continuous heterogeneity, but the poses were assumed to be known. CryoDRGNv2 [7] introduced hicrarchical pose search, comprising grid search followed by a form of branch-and- and--ound (BnB) optimization, akin cryoSPARC [8]. More recently, cryoPoseNet [9] and cryoAI [10] introduced amortized pose inference, to avoid the expense of orientation matching for each particle image. They employ convolutional neural network (CNN) to map input image pose estimate. Nevertheless, although efficient, amortized inference may not accurately represent multi-modal posterior pose distributions, which occur frequently early in reconstruction before the structure is estimated well. As such, compared to previous methods like cryoSPARC [8] and cryoDRGNv2 [7] that consider many pose candidates, amortized methods sometimes fail to identify the correct mode. Moreover, amortization confines the pose search to global parametric function (the encoder) which can lead slow convergence.
Here, we introduce a new approach, called cryoSPIN, to ab-initio homogeneous reconstruction that is able handle multi-modal pose distributions with a tailored encoder, and accelerates pose optimization with comi-amortization []]]. First, we perform amortized pose inference through shared multi-choice encoder that maps each input image to multiple pose estimates (Fig. 1, top). In contrast to cryoPoseNet [9] and cryoAI [10], which use pose encoders that produce one or two estimates, we attach multiple pose predictor heads multiple plausible poses for each image. By design, our multi-output encoder is able to account for pose uncertainty and encourage exploration of pose space during initial stages of reconstruction. For the 3D decoder, unlike cryoDRGN and cryoAl that use MI.Ps with computationally expensive feed-forward implicit networks, we adopt an explicit parameterization reconstruction. To train the encoder-decoder architecture, apply wwinner-takes-aa loss in which the 3D decoder is queried to obtain a 3D-to-2D projection for each predicted pose. Inspired by the loss function in Multi-Choice Learning [12-14], select the projection with the lowest image reconstruction error to determine the loss.
shared CNN feature extractor to predict
further accelerate the
During the course of optimization, as higher resolution details are resolved, the pose posterior becomes predominantly uni-modal. Thus, the pose search can be narrowed down to a neighborhood containing the most likely mode. In this stage, we propose to transition from auto-encoding auto-decoding (Fig. 1, bottom). particular, for each image, we choose the pose with the lowest reconstruction loss, iteratively refine it with stochastic gradient descent (SGD), and continue pose optimization

---

and reconstruction in an alternating fashion. This direct, per-image optimization procedure achieves faster convergence to more accurate poses than amortized inference alone, which relies on potentially sub-optimal predictions of the encoder as a parametric function.
We evaluate cryoSPIN in comparison to state-of-the-art methods, eryoSPARC [8], cryoAl [10], and cryoDRGN [7] on both synthetic and real experimental datasets. Using synthetic data, we find that semi-amortized inference noticeably accelerates the convergence of poses, and yields 3D reconstructions at similar or higher resolutions than baseline methods. We also show that our multi- choice encoder empirically accounts for multiple modes in the pose posterior. Finally, we apply cryoSPIN to an experimental dataset and obtain reconstructions with higher resolution than those from cryoAI, while matching the resolution of cryoSPARC reconstructions.
To summarize, we make following contributions.
Building upon cryoAI, develop new encoder based on multi-head architecture to return multiple plausible candidates to further mitigate pose ambiguity.
We train the encoder coupled with explicit volume decoder using winner-take-all loss, which penalizes the pose hypothesis with least reconstruction error.
Beyond the architecture and objective, we introduce semi-amortized pose inference, which begins with feed-forward amortized pose inference, followed by direct, per-particle pose optimization.
In ab-initio reconstruction, our of cryoAl and achieves higher-resolution on synthetic and experimental datasets.
method, cryoSPIN, is faster than amortized method
2 Related work
Cryo-EM reconstruction. Methods for cryo-EM reconstruction can be categorized as either homogeneous or heterogeneous. Homogeneous techniques [8-10] assume a rigid structure while heterogeneous ones [5, 7, 6, 15] allow for conformational variation. We focus on homogencous reconstruction, but optimization framework could be extended to heterogeneous data as well.
Early reconstruction techniques rely on common- lines [16-18] or projection mapping [19, 20] to select optimal poses. Other works [21, 22] frame the reconstruction problem in the context of maximum a posteriori (MAP) estimation, and jointly reconstruct pose and structure via expectation maximization (EM). We compare our approach to cryoSPARC [8], state-of-theearr method that uses stochastic gradient descent and a branch-and -ouud search [23] for ab initio reconstruction and pose estimation. Like these methods, our auto-decoding stage directly optimizes pose of every image.
More recently, amortized inference techniques have been proposed for pose estimation [24, 9, 10, 15). These techniques avoid explicit per- ver-particl pose optimization; instead, they train an auto-encoder or variational one [25] associate each particle image with a predicted pose [9]. One challenge is that the auto-encoders can become stuck in local optima during training [9]. To address this issue, cryoAI [10] produces two pose estimates per image coupled with a symmetrized loss function that penalizes the best one. We build on this concept by adopting a multi-head neural architecture as the encoder to output multiple plausible pose candidates and avoid local optima.
Multi-choice learning (MCL). Inspired by scenarios where a set of hypotheses needs to be generated to account for uncertainty in the prediction task, MCL [12] was introduced in a supervised setup learn multiple stractured-outputs with SSVMs [26]. Their motivating question was: can we learn to produce a set of plausible hypotheses? To address this, they define an" an"oracle" loss in which only the most accurate output pays the penalty. This loss is minimized even if there only single accurate prediction in the set. The early follow-up work [27, 13] uses the same loss to learn a deep CNN ensemble composed of M heads with a shared backbone network. Importantly, they show that the ensemble-me mean loss hurts prediction diversity across different heads, while training with the "oracle" loss yields specialized heads. Variations have since been proposed to mitigate hypothesis collapse or overconfidence issues in MCL by modifying the loss applying learnable probabilistic scoring schemes [28-30, 14]. MCL has been used to mitigate the ambiguity in several tasks including image segmentation [31], optical-flow estimation [32], trajectory forecasting [33], human pose and shape estimation [14, 34]. In our work, we use the "oracle" loss in context of auto-encoder which supervises the pose encoder indirectly through projections provided by the decoder.
3

---

3 Problem Definition
Image formation in cryo-EM is well-approximated using the weak phase object model [35]. Under this model, the structure of interest is unknown 3D density map. V R + Rzo. represented in some canonical coordinate frame. Cryo-EM images, {4 are approximated orthographic projections of the 3D map under some unknown coordinate transformation. We denote the unknown orientation by rotation R E SO((3) and the unknown position in terms of an in-plane translation t2 = (tz; ty) E IR?. Formally,
(x,,)) = [g * (St. PR, ()] v) + (zz,))
where PR.() is the linear operator computing the integral along the optical axis, 2, over the input density map rotated by R,, and S is the shift operator. The projection is convolved with the image- specific sorntd function (PSF), 9s. and corrupted by additive noise n. It is common to assume that n follows zero-mean white (or colored) Gaussian distribution.
By the Fourier slice theorem [36], the Fourier transform of a projection is equal to a central slice through the density map's s 3D Fourier spectrum. That is,
(ww,, = S.S.,, (w,,,, + w(w)),y))
where i and V denote the 2D and 3D Fourier transforms of the image and the density map. The slice perpendicular to the projection is computed by (PR, V). The translation by St, becomes phase shift operator St,, and convolution with 9: is equivalent to element-wise multiplication with, 9t, the contrast transfer function (CTF). The noise ñ remains zero-mean Gaussian.
Under this model, given the structure V. the negative log-likelihood of observing image 7, with noise variance o² and pose (R; (:) is
L = 202
Σ 1,,S;,,Pu (Pr, - ((War Wan
Ab-initio reconstruction methods [8-10, 7] solve jointly for the unknown structure V and per per-particl poses (R,,t,) t.). They often follow an Expectation-Mxxiiiztion Maximization (EM) [37, 21] procedure in which the E-step aligns images with the structure yielding pose estimates (R,, t.). and then in the M-step the volume V is updated by minimizing the negative log-likelihood in Eq. 3. Since errors in pose estimates lead blurry reconstructions, accurate pose estimates are crucial finding high-resolution structures. As discussed above, poses are either optimized through search and projection matching [8, 7] or estimated by an encoder network [10, 9, 15].
Methodology
We introduce cryoSPIN, two-staged approach to ab-initio cryo-EM reconstruction, combining auto-encoding and auto-decoding. We start with amortized inference (Fig. 1, auto-encoding) where a shared encoder equipped with multiple heads provide a set of pose guesses. Multiple guesses enable multimodal inference in the initial stages of reconstruction when the 3D structure is poorly resolved. As reconstruction improves, the pose posterior becomes less uncertain, which point we switch to direct pose optimization (Fig. auto-decoding). The former enables handling the pose uncertainty carly in reconstruction while the latter circumvents sub-optimality encoder network leading to arguably more accurate poses and faster convergence. We couple pose estimation module with an explicit volumetric decoder representing the 3D structure in Hartley space [5]. Our explicit model enables faster evaluation of projections compared to multiple passes through implicit neural representations [38-40]. The following sections discuss these components in detail.
4.1 Multi-choice encoder
With randomly initialized or a low-resolution reconstruction, the pose posterior contains multiple modes. Also, due to high levels of noise in images and near-symmetries in biological structures, pose estimation inherently involves high uncertainty. As a result, there exist several plausible poses for each image, and naive optimization search is prone to local minima.

---

To account for uncertainty in pose estimation, we build upon cryo.I [10] and extend the encoder to return multiple candidate poses @. Formally, given the image 1; E R #xW MM pairs of rotations and translations are computed as,
Li = (Rug. (43) = (F)( i M
(4)
where g E SO(3) and 4,5 E ER² The image-specific intermediate features, F; = VGG(T)) E RC*H*W are extracted by a shared convolutional backbone based on the VGG16 architecture [41], which are then supplied to M separate pose predictor heads, He, 1 S S M. yielding poses ői.j The heads are composed of fully-connected layers with distinct learnable weights, while the backbone is shared across all heads. Using multiple heads facilitates pose exploration and reduces uncertainty in the pose estimation process. Moreover, as 02 are randomly initialized, different heads are free to specialize on subsets pose space so that they collectively produce set of likely poses . )..M)
Inspired by multi-choice learning [12-14], we optimize the encoder-decoder using "winner-take-all" loss. For each image I; and predicted pose di.j. the negative log-likelihood, Li, in (Eq. 3), is computed, and the minimum is selected as the final loss for the corresponding image, i.e.,
L, = min -,mmin Cr..
(5)
Minimizing this loss requires only one of the predicted poses to be accurate. Interestingly, we find that this approach produces heads that specialize to somewhat disjoint regions of pose space (see Supp. C), consistent with previous work showing that it improves diversity in prediction tasks [13, 27]. This "divide-and-comee
approach facilitates pose space exploration, which is crucial in order to avoid local minima during the early stages of reconstruction where uncertainty in the 3D structure is significant. CryoAI [10] can be viewed special case of this formulation; it assigns two poses to each image consequence of input augmentation, and selects the best one with a symmetrized loss. In contrast, our approach augments the output of the encoder network with multiple heads, each providing pose estimate.
4.2 Switching from auto-encoding to auto-decoding
As optimization progresses and fine-grained details of the 3D structure are resolved, the posterior pose variance tends to decrease, with the posterior approaching a unimodal distribution. At this point, the gap between the amortized and variational posterior distributions is mainly determined by the error the pose estimate (predicted mean), prioritizing accuracy over exploration. However, feed-for feed-forward network, as a globally parameterized function of the input image, may be limited in the accuracy of its predictions, and hence the approximations inherent in amortized inference become significant sources of error in the refinement of the 3D structure. In prior work [11, 42-44], a related issue, called the amortization gap, characterized by the KI--divergence between the true and predicted (amortized) variational posteriors.
To address this issue, we adopt a semi-amortized inference scheme [11] comprising two stages. First, the encoder predicts a set of pose candidates using multi-head architecture, In the second stage, rather than amortized inference, pose parameters are directly optimized for each image using stochastic gradient descent. To initialize poses for the i-th image, we choose the one with the lowest reconstruction loss from the set of candidates {
Φs.M namely Φ; Φi.s such that
arg min Ci.
(6)
Subsequently, the pose and structure are optimized by coordinate descent using the negative log- likelihood (Bq. 3) as the objective function. See Supp. A for more details on rotation optimization.
4.3 Explicit representation as volume decoder
Recent works [10, 5] use coordinate networks [39, 40, 38] to implicitly model the Fourier rep- resentation of the 3D structure. Instead, we couple the pose estimation module with an explicit parameterization (3D voxel array) of the structure in the Fourier domain. The explicit representation is less computationally expensive than an MI.P to evaluate and update. Also, this choice is motivated by the fact the implicit decoder needs be queried multiple times for each image with the multi-head encoder.
5

---

We parameterize the volume using the Hartley representation [5]. The Fourier and Hartley transforms, respectively denoted as F(w) and (w), are related as
(ww)
(7)
where w denotes the frequency coordinate and R and I are the real and imaginary part, respectively. The Hartley representation is real-valued, and so more memory efficient to use than storing complex- valued Fourier coefficients. To account for high dynamic range of the Hartley coefficients, assume the Hartley field is decomposed into mantissa, m(w) and exponent e(w) fields [10] as,
((w) = m(w) x exp(e(w)).
This decomposition restricts the range of values for m(w) and e(w) and makes the reconstruction less sensitive to the initialization of the field.
5 Experiments
In what follows, we consider both synthetic and real datasets to empirically compare our semi- amortized ab-initio reconstruction method with ofate-of-the-art methods, cryoAI [10] and cryo- DRGN [7] which use amortized inference for reconstruction, and cryoSPARC [8] which performs pose search separately for each particle image.
Synthetic data. We generate synthetic datasets by simulating the image formation process for- malized in Sec. 3 using atomic models deposited in the Protein Data Bank (PDB). We compute ground-truth density maps of size 1283 for heat shock protein (HSP) [45] (1.5 A), pre-catalytic spliceosome [46] (4.33 Å), and SARS-CoV-2 spike protein [47] (2.13 Å). Then, N = 50, 000 projec- tions of size L 128 are generated by randomly rotating and projecting the density map along the canonical z-axis. To simplify the analysis and interpretation of results (e.g. Figs. 4, 5), we consider synthetic partieles
centered, and therefore omit the estimation of 2D translation and allowing us to visualize the pose posterior more easily. Finally, a random CTF is applied in the Fourier space and Gaussian noise added yielding SNR 0.1.
Experimental data. As widely adopted experimental benchmark, we use the 80S ribosome
dataset TEMPIAR-10022 [48]), containing 105,247 particle images of linear box size L == 360 with pixel size 1.34 Á. Following cryoAl and cryoDRGN, we downsample the images to L = 128 (3.76 A) using cryoSPARC software [8], and randomly split the data into two halves and run the reconstruction methods independently on each (to enable Fourier Shell Correlation as quantitative measure of the resolution of the 3D reconstruction). Unlike synthetic data, particles are not well-centered so is required to estimate translation parameters as well rotations. This is achieved by augmenting each head to return translation parameters as well.
During auto-encoding, we use encoders with M = and M = 15 heads for reconstruction on synthetic and real datasets, respectively, with Adam [49] to optimize encoder and decoder with learning rates 0.0001 and 0.05. Once switched to direct optimization (after 7 epochs for synthetic and 15 epochs for real data), reduce the decoder learning rate 0.02 and allocate new optimizer for pose parameters with learning rate 0.05. To be consistent with cryoAl, we use a batch size of 64 and train for the same number of epochs (20 for synthetic, 30 for real data). We use the public cryoAI codebase, run cryoSPARC v4.4.0 [8] with default settings, and run cryoDRGN homogenous ab-initio job. Methods are implemented in Pytorch [50]. Experiments are run on a single NVIDIA A40 GPU.
Implementation details.
5.1 Results
We first qualitatively compare final reconstructions of cryoSPIN with cryoAI [10], cryoDRGN [7], and cryoSPARC [8] on synthetic and real datasets (Fig. 2, left). We observe that cryoAl gets stuck in local minima when evaluated on experimental data (eg, 80S ribosome). particular, we found that in the presence of unknown planar shift, cryoAl fails to estimate effective translation parameters, but after centering all particles via pre-processing. we could reproduce the cryoAl results. Hence, the provided results for cryoAl in Fig. 2 are obtained after this pre-processing step, whereas our method and others are given off-centered experimental images. Both cryoSPIN and cryoSPARC capture high-fin frequeegyynn details of the 3D structure on all datasets, whereas reconstructions by amortized

---

CryoSPARC
CryoDRGN
CryoAl
CryoSPIN (Ours)
Spliceosome
0.0
DrySSPAR (11.744
0.2
OtyoAl (1259))
OreDRGN (19304)
9.0
OcreSPiN (11.856
1285A 22.7A 1936A
TESIA 3.0A
1.0
0.5
Spike
F3C
ChyosPARC
Cryual (6.324)
0.2
Cryd220 (8/53)
0.0
CrecleNN (7.174)
20.7A 10:11A
137A
3.50A 3.TAA
HSP
PSC
0.0
CyysPARC (6 (8.8A))
CyyuBCX anb
0.0
Crydtrix (LU64)
10
14WA Y.8A 5.99
502A 3.18A
EMPIAR-10026 10008
0.8
CoydPARC (1,85))
OryAA (2and)
CruBBRNN (8.06)
CreaPIN (7,8A)
DXC 37.12A 19330 1804A 9.5A 7.01X
Figure 2: Qualitative and quantitative comparison of reconstructions obtained by our proposed semi-amortized method, cryoSPIN, with cryoAl [10], cryoDRGN [7] and cryoSPARC [8] We note that cryoAl often becomes stuck in local minima when particles are not centered, so for cryoAl translate the input images correct for the spatial offset. (Left) Final 3D reconstructions three synthetic datasets and one experimental dataset (EMPIAR- - 10028) are depicted using ChimeraX [51]. (Right) FSC curves are visualized for quantitative comparison. The red dashed lines show the standard threshold levels of 0.5 and 0.143 to report the resolution (in Angstrom) for synthetic and real data, respectively. CryoSPIN achieves higher resolution on the Spliceosome and HSP datasets, and it is competitive with the state of the art on the Spike and EMPIAR- 10028 datasets.
methods, cryoAl and cryoDRGN, are inferior on the HSP and 80S datasets. In particular, on HSP, cryoAl gets stuck in local minima as it fails to handle high uncertainty in poses caused by symmetries in this structure.
For quantitative comparison, we visualize the Fourier Shell Correlation (FSC) [52] between the reconstruction and the ground truth (Fig. 2, right). FSC is the gold-standard metric measuring the normalised cross-correlation coefficient between two 3D Fourier volumes along shells of increasing radius. Higher FSC implies more accurate reconstruction. On the Spliceosome, HSP and 80S experi- mental datasets, our reconstruction outperforms amortized methods of cryoAI and cryoDRGN. Also, our method outperforms cryoSPARC on Spliceosome, HSP and 80S, while achieving competitive FSC on Spike. We use the standard 0.5 and 0.143 cutoff thresholds to report the resolution for synthetic and real data, respectively. CryoSPIN achieves higher or competitive resolution compared to others on all datascts.

---

Spliceosome
HSP
Spike
22.5
CryoAl (1289A)
20.0
CryoAl (7.114) CrySSPIN (5 05.)
30
GrysAl (6.99A)
Resolition (1
-00
CiyuSP (11/79A)
Resolusion (2)
17.5
25
CryuSP (1-17A)
15.0
Resolution
20
12.5
30
10.0
50
20-
5.0
10
10
IF
10
Reconstruction Time (Seconds)
Reconstruction (Seconds)
Reconstruction Time (Seconds)
Figure 3: 31) resolution as function of log time for different methods. These plots show the
semi-amortized method is significantly faster than cryoAI.
FuliAmortieed Semimmorized
AuE-Amartioe mortAmortiaed
Seni-Ameiiaal
10
Grodesk
20
Epocis
150
Full-Amotiesd
Fullimortiasd
FuliAmotized mimmonnized
Fuli4morized
Scmi-Amerioo
45
Geslesi
Figure 4: Quantitative and qualitative comparison of fully- vs. semi-amortized methods in pose optimization for Spike (top) and Spliceosome (bottom) datasets. (Left) Mean geodesic distance between the predicted pose and ground-trute is visualized different epochs. By switching from amortized inference to direct optimization, semi-amortized method (blue) enjoys accelcrated pose convergence compared to amortized inference (red). (Right) To visualize pose inference, images depict the approximate log posterior for three particles (marginalized over in-plane rotations) as diewddirection distribution over HEALPix [53] uniform grid on unit sphere S After Gnomonic projection to 2D, show the neighborhood of the mode of interest. Black dots depict starting points. Blue and red dots are pose estimates from fully- and semi-amortized methods.
In Fig. visualize resolution-time plots, showing that cryoSPIN gets high-resolution reconstructions significantly faster than cryoAl. In fact, semi-amortization in cryoSPIN ac- celerates the improvement in the resolution and with our explicit structure decoder, we cir- cumvent expensive MLP evaluations and train faster without any degradation to the final re- construction quality as shown in Fig. 2. In Supp. B, through detailed comparison with cryoAI, we show cryoSPIN is ~ 6x
faster and uses ~ 5x less memory compared to cryoAI. Table Mean/median errors of estimated rotations in units of degrees evaluated on synthetic datasets.
Finally, report the mean and me- dian errors in estimated poses on syn- thetic datasets in Table 1, showing that our method outperforms cryoSPARC and cl High errors by cryoAl on HSP dataset shows that it gets stuck local minima and fails to accurately estimate poses.
Method
HSP
Spike
Spliccosome
CryoSPARC [8]
6,23 / 1.05
1.61 / 1.51
1.41 / 1.36
CryoAl [10]
45.83 / 61.86
2.52/ 2.52/2.00
2,85 / 2.61
CryoSPIN (Ours)
3.27 / 0.97
1.54 / 0.90
0.68 / 0.61

---

CryoAI
Multi-head
CrycAl (11)))
Mutti head (5.317)
Reso ution (3)
15
10
Proj
Proj
Input
20
Epochs
27.0%
25.5%
24 0%
23.5%
Input
Proj
Proj 2
nout
Head
Head 2
Head
Head
Head
Head
Head
Head 4
Figure 5: Comparing performance of our multi-head pose encoder (M 4) with cryoAl pose encoder on the challenging HSP [45] dataset. (Left) The approximate log posterior of view direction is visualized on the unit sphere with highlighted areas showing modes of the distribution. CryoAI and multi-head encoders provide two and four pose estimates, respectively, which are marked with colored dots on the sphere (the order of poses is arbitrary). Below the sphere, the corresponding projections are illustrated. CryoAI fails to find the correct mode while our method is able cover multiple modes. (Top, right) With multi-head encoder, the reconstruction converges to much higher resolution compared to cryoAI. (Bottom, right) Percentage of images assigned to each head visualized bar plot confirming that all heads participate in pose estimation.
5.2 Semi-Amortized vs. Fully-Amortized
To showcase the benefit of the auto-decoding stage, compare semi-amortized method with fully-amortized baseline on spike and spliceosome datasets. Starting with auto-encoding, we branch the running reconstruction into after epochs: the first continues using the encoder while the second switches to direct pose optimization. We plot the average pose error with respect to the ground-truie -ruth different epochs in Fig. 4 (left). When using the multi-head encoder, the pose with the least reconstruction error is selected among the candidates. Interestingly, as our method switches to direct pose optimization, the error in pose drops quickly, whereas the fully-amortized baseline exhibits slow convergence. This clearly shows the superiority of auto-decoding compared to auto auto-encoding during the later stages of optimization.
Next, we analyze the evolution of pose estimates over the optimization landscape depicted as heat map (Fig. 4, right). Qualitatively, poses obtained by direct optimization (blue dots) show stable convergence to the optimal point (highlighted area) while those inferred by the encoder (red dots) frequently oscillate. Indeed, the encoder in amortized inference which might be too restrictive, yielding sub-optimal pose parameters. Therefore, poses inferred in an amortized fashion might fail to consistently converge to the optimal point. On other hand, direct optimization is intuitively more flexible as it is performed separately and locally for each image, exhibiting more stable convergence. See Supp. for more examples.
globally parameterized function
5.3 Multi-Modal Pose Posterior
Lastly, we examine how well cryoSPIN multi-head encoder performs vs. cryoAI encoder in terms of handling the uncertainty in the pose on the challenging dataset of HSP [45]. To simplify visualization, we run our method with M = 4 heads in this experiment. We first inspect the behavior of encoders for two example cryo-EM images (Fig. 5, left, see Supp. for more examples). each row, the approximate posterior distribution over view direction visualized for both cryoAl and

---

multi-head encoders given the input image and reconstruction. Our multi-head encoder returns set of plausible candidates, while cryoAI obtains two pose estimates by data augmentation. both examples, the multi-head encoder identifies the correct mode while cryoAl selects the incorrect one. Note that the pose set predicted by the multi-heal encoder contains other posterior modes as well. Intuitively, cryoAl encoder with two pose predictions cannot explore the pose space much as our multi-head encoder. Thus, the pose posterior in cryoAl remains uncertain and multi-modal during reconstruction. This also contributes to slower convergence to higher-resolution reconstructions as depicted in the top left of Fig. 5. Finally, we visualize the percentage of images assigned to each head by our "winner-taker shows that the multi-head encoder does not suffer hypotheses collapse [30], namely all heads actively participate in the pose prediction. In Supp. C, we investigate how each head specializes in pose prediction for non-overlapping regions 50(3).
loss. The relatively even spread of assignments to different heads
Conclusion, Limitations and Future Work
In this paper, we introduce cryoSPIN, new semi-amortized approach ab initio cryo-EM recon- struction. We develop new multi-head encoder that estimates a set of plausible candidate pose to handle the uncertainty. As the uncertainty is reduced, we transition to an auto-decoding stage where poses are iteratively optimized using SGD for each image. Our results show that our multi-head encoder is able to capture multiple modes of the pose distribution, and our flexible direct optimization enables accelerated convergence of poses and reconstructions. CryoSPIN outperforms cryoAl on experimental data and achieves competitive results with cryoSPARC.
Limitations and future work. In this work, we assume that the 3D structure is rigid while many biomolecules are flexible in practice, exhibiting different conformations. Our results show that, even in this relatively constrained homogeneous case, there is a significant gap between amortized inference and our semi-amortized approach. We expect that the benefits of semi-amortized inference will also apply to such heterogeneous cases. In particular, our proposed multi-head encoder should be applicable to the estimation of the multimodal posterior well as pose. Yet, there are significant challenges in heterogeneous reconstruction, such as extra uncertainty due heterogeneity and conflation of latent spaces, making it outside the scope of our analysis on semi-amortized inference.
heterogeneity latent parameters as
The development of principled criteria for switching from amortized to auto-decoding is also an interesting direction for future research. Our experiments with a range of values show that after a sufficient number of epochs, the reconstruction error becomes insensitive to the switching time. Our intuition is that when switching too early, the pose posterior may have significant uncertainty, and the pose estimates may be too far from the correct basin of attraction to afford robust convergence. Therefore, it is better to switch late than early. However, a well-defined heuristic for the optimal switch time is lacking.
Structure determination with Cryo-EM for macromolecules one of the most exciting, fundamental areas in structural biology, and is already having with vast social impact. Cryo-EM was used to determine the structure of the SARS-COV2 spike protein, the discovery of its pre-fusion conformation, and the evaluation of medical counter-measures. It compliments advances in methods like alpha-fold for structure prediction, transforming our understanding the process of life at the scale of the cell, and the design and development of new therapeutics.
Societal impact.
Acknowledgments and Disclosure of Funding
This research was supported in part by the Province of Ontario, the Government of Canada, through NSERC, CIFAR, and the Canada First Research Excellence Fund for the Vision: Science to Applica- tions (VISTA) programme, and by companies sponsoring the Vector Institute.
References
[1] Werner Kühhbrannt. The resolution revolution. Science, 343(6178) 2014. 1
[2] Amit Singer and Fred J Sigworth. Computational methods for single-particn electron cryomi- croscopy. Annual Review of Biomedical Data Science, 3:163-190, 2020. 1
10

---

[3] Sjors HW Scheres, Haixiao Gao, Mikel Valle. Gabor T Herman, Paul PB Eggermont, Joachim Frank, and Jose-Maria Carazo. Disentangling conformational states of macromolecules in 3d-em through likelihood optimization. Nature Methods, 4(1):27-29, 2007.
[4] Roy R Lederman, Joakim Andén, and Amit Singer. Hyper-mone molecuooe:::ee: On the representation and recovery of dynamical structures, with application to flexible macro-molecular structures in cryo-em. arXiv preprint arXivv:907.0011 1907.0/59, 2019.
[5] Ellen D Zhong, Tristan Bepler, Bonnie Berger, and Joseph H Davis. CryoDRGN: reconstruction of heterogeneous cryo-EM structures using neural networks. Nature Methods, 18(2):176-185, 2021. 2, 3, 6
[6] Ali Punjani and David J Fleet. 3DFlex: Determining structure and motion of flexible proteins from cryo-EM. Nature Mthods, 20(6):860-870, 2023. 1,3
[7] Ellen Zhong, Adam Lerer, Joseph H Davis, and Bonnie Berger. CryoDRGN2: Ab initio neural reconstruction of 3d protein structures from real cryo-EM images. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4066-4075, 2021. 2, 3, 4, 6, 7
[8] Ali Punjani, John L Rubinstein, David J Fleet, and Marcus A Brubaker. cryoSPARC: algorithms for rapid unsupervised cryo-EM structure determination. Nature Methods, (4(3):290-296 2017. 2, 2, 4, 6, 7,
[9] Youssef SG Nashed, Frédéric Poitevin, Harshit Gupta, Geoffrey Woollard, Michael Kagan, Chun Hong Yoon, and Daniel Ratner. CryoPoseNet: End-to-end simultaneous learning of single-particle orientation and 3D map reconstruction from cryo-electron microscopy data. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages (06664076.. 2021. 2, 3,4
[10] Axel Levy, Frédéric Poitevin, Julien Martel, Youssef Nashed, Ariana Peck, Nina Miolane, Daniel Ratner, Mike Dunne, and Gordon Wetzstein. CryoAI: Amortized inference of poses for ab initio reconstruction of 3d molecular volumes from real cryo-EM images. In European Conference Computer Vision, pages 540-557. Springer, 2022. 2, 4, 5, 6, 7, 8, 18
[11] Yoon Kim, Sam Wiseman, Andrew Miller, David Sontag, and Alexander Rush. Semi-amortized variational autoencoders. In International Conference on Machine Learming, pages 2678-2687. PMLR, 2018. 2,5
[12] Abner Bouzman-Rivers Dhruv Batra, and Pushmeet Kohli. Multiple choice learning: Learning to produce multiple structured outputs. Advances in Neural Information Processing Systems, 25, 2012.22222 3, 5
[13] Stefan Lee, Senthil Purushwalkam Shiva Prakash. Michael Cogswell, Viresh Ranjan, David Crandall, and Dhruv Batra. Stochastic multiple choice learning for training diverse deep ensembles. Advances in Neural Information Processing Systems, 29, 2016. 16
[14] Christian Rupprecht, Iro Laina, Robert DiPietro, Maximilian Baust, Federico Tombari, Nassir Navab, and Gregory D Hager. Learning in an uncertain world: Representing ambiguity through multiple hypotheses. In Proceedings of the IEEE International Conference on Computer Vision, pages 3591-3600, 2017. 2, 3,
[15] Axel Levy, Gordon Wetzstein, Julien NP Martel, Frederic Poitevin, and Ellen Zhong, Amortized inference for heterogeneous reconstruction in cryo-EM. Advances in Neural Information Processing Systems, 35:130388 1303--13049, 2022. 3,
[16] Amit Singer, Ronald R Coifman, Fred J Sigworth, David W Chester, and Yoel Shkolnisky. Detecting consistent common lines in cryo-EM by voting. Journal of Structural Biology, 169(3):312-322, 2010. 3
[17] Ido Greenberg and Yoel Shkolnisky. Common lines modeling for reference free ab-initio reconstruction in cryo-EM. Journal of Structural Biology, 200(2):106-117) 2017.
[18] Gabi Pragier and Yoel Shkolnisky, common lines approach for ab initio modeling of eyclically symmetric molecules. Inverse Problems, 35(12):124005, 2019.

---

[19] Pawel A Penczek, Robert A Grassucci, and Joachim Frank. The ribosome at improved resolution: new techniques for merging and orientation refinement in 3D cryo-electron microscopy of biological particles. Ultramccoscopyy 53(3):251-270, 1994. 3
[20] Timothy S Baker and R Holland Cheng. A model-based approach for determining orientations of biological macromolecules imaged by cryoelectron microscopy. Journal of Structural Biology, 1 116(1):120-13)) 1996. 3
[21] Sjors HW Scheres. A Bayesian view on cryo-EM structure determination. Journal of Molecular Biology, 415(2):406411,, 2012. 3,
[22] Sjors HW Scheres. RELION: implementation of a Bayesian approach to cryo-EM structure determination. Journal of Structural Biology, 180(3):51--535 2012. 3
[23] Eugene I. Lawler and David Wood. Branch-and-bound methods: A survey. Operations
Research, 14(4):699-719, 1966. 3
[24] Dan Rosenbaum, Marta Gamelo,, Michal Zielinski, Charlic Beattie, Ellen Clancy, Andrea Huber, Pushmeet Kohli, Andrew W Senior, John Jumper, Carl Doersch, et al. Inferring a continuous distribution of atom coordinates from eryo-EM images using VAFs. arXiv preprint arXiv:2106. /4/08, 2021. 3
[25] Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv preprint 11Xiv:1312.6114, 2013. 3
[26] Ioannis Tsochantaridis, Thorsten Joachims, Thomas Hofmann, Yasemin Altun, and Yoram Singer. Large margin methods for structured and interdependent output variables. Journal of Machine Learning Research, 6(9), 2005.
[27] Stefan Lee, Sentil Purushwalkam, Michael Cogswell, David Crandall, and Dhruy Batra. Why heads are better than one: Training a diverse ensemble of deep networks. arXiv preprint 111.v:151..06311 2015. 3, 5, 16
[28] Kimin Lee, Changho Hwang, KyoungSoo Park, and Jinwoo Shin. Confident multiple choice learning. In International Conference on Machine Learning, pages 2014-2023 PMLR, 2017.
[29] Kai Tian, Yi Xu, Shuigeng Zhou, and Jihong Guan. Versatile multiple choice learning and its application to vision computing. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pottern Recognition, pages 6349-6333 60999-6355550 2019.
[30] Victor Letzelter, Mathicu Fontaine, Mickaël Chen, Patrick Pérez, Slim Essid, and Gael Richard. Resilient multiple choice learning: A learned scoring scheme with application to audio scene analysis. Advances in Neural Information Processing Systems, 36, 2024. 3, 10
[31] Simon Kohl, Bernardino Comera-Paredess, Clemens Meyer, Jeffrey De Fauw, Joseph R Ledsam, Klaus Maier-Hein, SM Eslami, Danilo Jimenez Rezende, and Olaf Ronneberger. A probabilistic U-Net for segmentation of ambiguous images. Advances in Neural Information Processing Systems, 31, 2018. 3
[32] Eddy Ilg, Ozgun Cicek, Silvio Galesso, Aaron Klein, Osama Makansi, Frank Hutter, and Thomas Brox. Uncertainty estimates and multi-hypotheses networks for optical flow. In European Conference on Computer Vision, pages 652 2018.
[33] Ye Yuan and Kris Kitani. Diverse trajectory forecasting with determinantal point processes. arXiv preprint 107.47790. 2019.
[34] Benjamin Biggs, David Novotny, Sebastien Ehrhardt, Hanbyul Joo, Ben Graham, and Andrea Vedaldi. 3D Multi-bodies: Fitting sets of plausible 3d human models to ambiguous image data. Advances in Neural Information Processing Systems, 33:20496-20507, 2020. 3
[35] Earl J Kirkland. Advanced computing in electron microscopy, volume 12. Springer, 1998. 4
[36] Ronald N Bracewell. Strip integration radio astronomy. Australian Journal of Physics, 9(2):198-217 1956. 4
12

---

[37] Todd K Moon. The expectation maximizationnnn algorithm. IEEE Signal processing magazine, 13(6):47-60, 1996. 4
[38] Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. Advances in Neural Information Processing Systems, 33:7462-7473, 2020. 4, 5
[39] Matthew Taneik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions low dimensional domains. Advances in Neural Information Processing Systems, 33:7537-7547, 2020. 5
[40] Ben Mildenhall, Pratul Srinivasan, Matthew Tancik, Jonathan Barron, Ravi Ramamoor- thi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1)999~106, 2021. 5,5
[41] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint 1409. 2014.
[42] Devon Hjelm, Russ R Salakhutdinov, Kyunghyun Cho, Nebojsa Jojic, Vince Calhoun, and Junyoung Chung. Iterative refinement of the approximate posterior for directed belief networks. Advances in Neural Information Processing Systems, 29, 2016. 5
[43] Rahul Krishnan, Dawen Liang. and Matthew Hoffman. On the challenges learning with inference networks sparse, high-dimensional data. In International Conference Artificial Intelligence and Statistics, pages 1511. PMLR, 2018
[44] Joe Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. In International Conference on Machine Learning, pages 3403-3412. PMLR, 2018. 5
[45] D Goodsell. PDB-101 Molecule of the Month: Hsp90. 2008. 6, 9
[46] Clemens Plaschka, Pei-Chun Lin, and Kiyoshi Nagai. Structure of a pre-catalytic spliceosome. Nature, 546676660)667-662,, 2017. 6
[47] Alexandra Walls, Young-Jun Park, M Alejandra Tortorici, Abigail Wall, Andrew McGuire, and David Veesler. Structure, function, and antigenicity of the SARS-CoV-2 spike glycoprotein. Cell, 181(2):811-992, 2020. 6
[48] Wilson Wong, Xiao-chen Bai, Alan Brown, Israel Fernandez, Eric Hanssen, Melanie Condron, Yan Hong Tan. Jake Baum, and Sjors HW Scheres. Cryo-EM structure of the Plasmodium faleiparum 80S ribosome bound to the anti-protozoan drug emetine. Elife, 3:e03080, 2014.
[49] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 2014.
15
[50] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 32, 2019. 6
[51] Thomas D Goddard, Conrad C Huang, Elaine C Meng, Eric F Pettersen, Gregory S Couch, John H Morris, and Thomas Ferrin. UCSF ChimeraX: Meeting modern challenges in visualization and analysis. Protein Science, 27(1):14-25, 2018 7
[52] Marin Van Heel and Michael Schatz. Fourier shell correlation threshold criteria. Journal of Structural Biology, 151(3)2550-66. 2005. 7
[53] Krzysztof Gorski, Eric Hivon, Anthony Banday, Benjamin D Wandelt, Frode K Hansen, Mstvos Reinecke, and Matthia Bartelmann. HEALPix: A framework for high-resolution discretization and fast analysis of data distributed on the sphere. The Astrophysical Journal, 622(2):759, 2005. 8, 16, 17
13

---

[54] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Viti,o, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017. 15
[55] Julian Panetta. Optimizing over SO(3). Technical report, University of California, Davis, 2018.
15
[56] Sebastian Grassia. Practical parameterization rotations using the exponential map. Journal of Graphics Tools, 3(3):29 48, 1998. 15
14

---

Supplementary Material
CryoSPIN: Improving Ab-Initio Cryo-EM Reconstruction with Semi-Amortized Pose Inference
A Pose Optimization
For the auto-encoding stage, we follow cryoAl and design the convolutional backbone to output the six-dimensional representation commonly referred to as S2S2. To compute the rotation matrix from this representation, the 6D vector is split into two 3D vectors and normalized, denoted as U1, 22 R We then compute the cross product between them, fty = 2/1 X 122, yielding a new unit vector. Selecting 2) and v/ the first and third columns of the target rotation matrix, compute 82 24 11 which is another unit vector orthogonal to both Uj and V3. Then, R = [1), 02, va].
Once switched to direct optimization, we change parameterization to axis-angle representation. During auto-decoding, we alternate between five iterations of pose SGD updates and one iteration of volume update. To update poses, we keep the volume fixed and optimize for the negative log- likelihood (Eq. 3) with respect to the pose parameters. We define the new pose estimate based on the current one as follows:
R4+1 = R4
where Rs an infinitesimal rotation matrix perturbing the current estimate. The perturbation matrix R3 is parameterized by axis-angle representation. By a single vector w E RR one can represent both the axis ||w|| and the angle 0 < ||ا|| < π for any given rotation. Using Rodrigues formula, the perturbation matrix Ro(w) can be parameterized as a function of w. To find the optimal w, one can initialize it with zero vector, and then use automatic differentiation [54] in pytorch to compute the gradient with respect to w and make updates using Adam [49]. However, a naive implementation of the function Rs(w) would lead to numerically unstable calculations of the partial derivative "w
w
0R. In
fact, there is singularity the zero vector, and the partial derivative involves terms that are unstable around the origin. Formally, the derivative of i-th column of the rotation matrix Rs with respect to the vector w is [55, 56],
2R(")
( @ w [e(0 sin([w[ll)
(L-cos([as)) cos (
+ I(w e + w @ e0
dw
اil
||ا |
+ (w 8 w) (w e()) 2 cos(||wo|)
+ ||اا|| (in(|اull
|
+ ((w x ))) @ W/ |W cos( (|]ll
sin ([wl|)
where @ and x are tensor and cross products, respectively. (+)) is the i-th standard basis in 3D and [v]x denotes the cross product matrix for the vector v. In all four terms, there are scalars such as sin([[ell)
0 ||||| or 1-cos([[w]]) that evaluate to o at zero angle = 0. Similar to [55], for ||w||| < 1, we ||]]| substitute these terms with their numerically robust Taylor expansion, for instance,
sin(lww]] |
|
+ O(Ilw))))
6
1 cos( [
1
||||
+ +[[[w])))
|
2
24
We implement differentiable and numerically stable version of the function Ro(w) in pytorch and use in our pose estimation module.
B Ablation Study on Decoder
We perform an ablation study on the decoder, detailed in Table 2, as well as comparison with baselines in terms of reconstruction time per epoch, GPU memory usage, and number of parameters. As our method consists of two stages, report numbers for each stage separately. In the first
15

---

Table The reconstruction time per epoch, GPU memory, and number parameters for different methods averaged across three runs. GPU memory is recorded separately for the encoding and decoding modules. Also, two numbers provided for semi-amortized method corresponds to auto- encoding and auto-decoding stages, respectively. In CryoAl-explicit, the implicit decoder of CryoAl replaced with an explicit decoder. Since we store rotations in 3D representation during direct optimization, the number of parameters of pose module is N x 3 with N denoting number of images in millions.
GPU Mem. (GB) Decoding
Model
Time (s)
# Params (M)
Encoding
Semi-Amortized
19.76, 81.61
1.17,
2.22, 0.35
13.01, 4.29 + N x
Fully-Amortized
99.75
1.17
2.22
13.01
CryoAl-explicit
142.40
2.32
0.75
9.05
CryoAl
622.39
2.78
14.05
5.09
stage, the encoder has H=7 heads, while in the second stage, it is replaced with pose module with size depending on the number of particles (N). To facilitate comparison, we include another baseline, cryoAl-explicit, in which the implicit decoder is replaced by an explicit one as in our method. The explicit and implicit decoders have 4.29 and 0.33 million parameters, respectively. Despite larger decoder, our method is 6x faster and uses 5x less memory compared to cryoAI. Importantly, swapping the implicit decoder with an explicit one (cryoAl-explicit) significantly drops time and memory, indicating that the implicit decoder is major computational and memory bottleneck. Yet, cryoAl-explicit uses more GPU memory for encoding and is 1.5x slower than our method as it performs early input augmentation and runs the entire encoder twice per image. Our method, by augmenting the encoder head, saves memory and time during pose encoding. Finally, the amortized baseline, which uses explicit decoder coupled with a multi-head encoder (H=7), uses more memory in decoding (negligible vs implicit decoder) and runs slower than the direct optimization stage the semi-amortized method.
C Specialization of Encoder Heads
natural question about the multi-head encoder is: how each head does take part in pose encoding process? To address this, using the synthetic datasets, we conduct experiment with multi- architecture (M = 4) and visualize the performance of each head on different regions of SO(3) space. In particular, as before, we define uniform grid on the unit sphere using HEALPix [53] and assign images to their corresponding cells based on the view-direction. Now, for all images end up in the same cell, shown in Fig. 6, all heads actively participate in pose estimation and they are able to specialize in prediction of poses for images with certain view-direction. A similar result has been provided in prior work on MCL [27, 13], to show that minimizing the eгror made by the best prediction ("oracle"" loss) encourages diversity in deep ensembles. In our problem, by optimizing "winner-takes-all" loss, the whole burden of pose estimation is no longer on single network but it gets divided between multiple heads as separate predictors.
compute the average rotation error and visualize it separately for each head. As
D More on Semi-Amortized vs. Fully-Amortized
To validate the advantages of direct pose optimization in our comi-amortized method, we further show more qualitative examples of paths taken by pose estimates over the optimization landscape during reconstruction in Fig. 7. For both methods, optimization start from the same point marked by black dot in the vicinity of the distribution mode. will then continue paths colored in blue and red for semi-amortized and full-amortized methods, respectively. We observe in all examples that iterative updates by stochastic gradient descent demonstrate stable convergence toward the optima while poses obtained by amortized inference show unstable behavior around the mode.
16

---

Head
Head
Head
Head
Spliceosome
Spike
HSP
Figure 6: Average rotation error is visualized over the sphere for different heads of the multi-head pose encoder (M 4). The sphere is uniformly divided into cells using HEALPix [53]; images are assigned to cells based on ground-truth view-direction. For each cell, the average rotation error is visualized, showing diverse behavior of different heads across the space. Blue and red colors show low and high error regions, respectively. Error ranges from zero 180 degrees.
NulAmrotiaed Semi-Amorters
Fulommationl
Fillmmation SemHomontre
Senbimottee
NatAmmottes
Nattrrotice
Figure 7: We compare the behavior of mullt-amottized and semi-amortized pose inference four examples per dataset. Two rows correspond Spike and Spliceosome datasets, respectively. Each plot shows the approximate log pose posterior, marginalized over in- in-plane rotations represented as a heat map on uniform grid over the unit sphere S2. Gnomonic projection to 2D is also applied, followed by zooming the proximity of the mode of interest. Black cross is the starting point while blue dots and red dots show poses estimated by fully- and cemi-amortzeed methods, respectively.
E
More on Multi-Modal Pose Posterior
Through more examples (Fig. 8). we demonstrate that cryoAl fails to handle ambiguity in pose estimation HSP dataset. The visualization shows that pose estimates by cryoAI become stuck in incorrect modes whereas our pose encoder with multi-head architecture is able to return a pose candidate that captures the correct mode.
17

---

CryoAl
Multi-head
2
Input
Proj. 1
Proj. 2
Input
Head 1
Head 2
Head 3
Head 4
Input
Proj. 1
Proj. 2
Input
Head
Head 2
Head 3
Head 4
Figure 8: The approximate log posterior of view-direction visualized on the unit sphere with high- lighted areas showing modes of the distribution. CryoAI [10] and our multi-head encoders provide two and four pose estimates, respectively, which are marked with colored dots the sphere (the order of poses is arbitrary). The corresponding projections are also illustrated. CryoAl cannot identify the correct mode of pose distribution.
18

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justifiicton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We clearly state the claims in the abstract, discuss them through the introduction section, and summarize them list at the end of introduction section.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
19

---

We discuss limitations in the conclusion section. We acknowledge that ours assumptions about homogeneity and centered images could be limiting some experimental setups. However, our framework could be extended to account for heterogeneity and translation parameters.
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-texi system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
We provide no theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
20

---

Justification: In a subsection called implementation details, we included all details needed to reproduce reconstruction results our method, cryoAI and cryoSPARC. Steps generate synthetic datasets and preprocessing steps to prepare experimental data are also provided.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (c.g., with an open-soure dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm ca/ sublic/guides wttps///oooSumiiiiiooooooooo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: //nips wtips.ccpppiii wwtpc//ullcc/eeeecceeeceeeeeeele wttgggudeee///ee /ewiiidee/ooodeiiiiiccccc SotooeeSubmsssiooool.... ssionoolccy) for more details. wttpc//pcccggui/e//e/eeeee/ce/ee
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
21

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have discussed all the necessary information about experimental settings in implementation details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: When reporting the resolution, we follow the standard to compute FSC by splitting data into two halves and make independent reconstructions and compare resulting half-maps.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
22

---

Answer: [Yes]
Justification: We have mentioned that we are using A40 Nvidia GPU to prodcue the results. Also, in Fig. we compare the execution time of our method vs. cryoAl.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https
Answer: [Yes]
Justification:
We conform with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Societal impacts of the work is dicussed in the conclusion section.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeet of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and funetioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
23

---

Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: For cryoAI, are using the available code provided in github, For cryoSPARC, we have mentioned which version the software used. Synthetic datasets are generated based on atomic models on PDB which are free of all copyright restrictions and freely available for both non-commercial and commercial use. For experimental data, EMPIAR database which is freely and publicly available to the global community under the CCO licence.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, saperswithcone com/ datasets has curated licenses for some datascts. Their licensing guide can help determine the license of dataset.
For existing datasets that are re se-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We don't 't release any new assets.
Guidelines:
The answer NA means that the paper does not release new assets.

---

Researchers should communicate the details of the dataset/codeemmodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing or research with human subject.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: No crowdsourcing or research with human subject.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
25


Reviews:
[{"review_id": "OVQ8kj6QdO", "reviewer": "Reviewer_diCb", "summary": "The submission addresses ab-initio cryo-EM reconstruction, where both image poses and the 3D structure are estimated. The authors adopt a multihead architecture to estimate multiple poses for each image, encouraging the exploration of pose space early in the reconstruction process. They then refine poses in an auto-decoding fashion using SGD. Experiments on synthetic and experimental datasets demonstrate acceleration and resolution improvement over baseline approaches.", "questions": "1. The results presented for cryoSPARC are significantly worse than my experience. What is the setting for cryoSPARC? Are there any important assumptions?\n2. Will the multi-head structure predict very similar poses, thereby undermining the assumption that this design encourages exploration of the pose space?\n3. Can the authors provide a comparison of the training time and memory consumption between the proposed method and the baselines?", "limitations": "n/a", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Compared to cryoAI, mapping the input image to multiple pose candidates in the auto-encoding stage can account for pose uncertainty and encourage exploration of the pose space during the initial stages of reconstruction.\n2. The experimental results provided in the paper include both synthetic and real experimental data, as well as video results.\n3. The proposed method offers a speed advantage over cryoAI.", "weaknesses": "1. The paper's technical contribution is primarily limited to a multi-head architecture for estimating multiple plausible poses.\n2. The results presented for cryoSPARC are significantly worse than my experience suggests (i.e., on Splicesome or EMPIAR-10028). CryoSPARC's results should not be inferior to those obtained by the proposed approach.\n3. The paper cannot handle heterogeneous reconstruction, which limits its general applicability in structural biology research.\n4. Import baselines are missing: RELION, cryoFIRE, cryoDRGN-BNB and cryoDRGN2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review and questions. Below, we address the concerns raised in the review.\n\n**\"technical contribution limited to a multi-head architecture\".**\nOur contributions include the design of a new encoder equipped with multiple heads to mitigate uncertainty in pose auto-encoding, and with it, a new “winner-takes-all” loss that encourages diversity in pose predictions. Beyond the architecture and objective, we propose semi-amortized pose inference, combining amortized pose inference and auto-decoding, which outperforms the predominant fully-amortized alternative. In particular, semi-amortized inference accelerates and stabilizes pose convergence (as illustrated in Fig. 4) leading to improved reconstruction.\n\n**CryoSPARC results.** In our experiments, for fair comparison to cryoAI (which learns a $128^3$ volume), the particle images are downsampled to D=128 before reconstruction. The deposited images in EMPIAR-10028 have a box size of D=360, with 1.34 A/pixel. The downsampling increases the pixel size to 3.77 A, and hence the best possible resolution (at the Nyquist rate) is 7.54 A. Both cryoSPARC and our method reach this limit as shown by the FSC plot in Fig. 2. We agree this is a mid-level resolution and hence atomic resolution details are not determined.\n\n**Heterogeneous reconstruction.** We intentionally focus on the homogeneous case instead of heterogeneous reconstruction; we show that even for the simpler homogeneous reconstruction problem there is a significant disparity between applying amortization in this setting (e.g., CryoAI) and semi-amortization (our proposed method, see Fig. 4). We expect that the benefits of semi-amortized inference will also apply to the heterogeneous reconstruction case. In particular, our proposed multi-head encoder should be applicable to the estimation of the multimodal posterior over heterogeneity latent parameters as well as pose. Yet, there are significant challenges in heterogeneous reconstruction, making it outside the scope of our analysis on semi-amortized inference—for example, (1) heterogeneity adds significant uncertainty and multimodality in the posterior over the latent variables, (2) changes in heterogeneous structure during optimization complicate pose estimation using amortized inference, and (3) existing methods for heterogeneous reconstruction conflate pose and conformational state, and disentangling these properties remains a challenge. We plan to investigate these lines in future work.\n\n**Other baselines.** We primarily compare with CryoAI and CryoSPARC as they are the state-of-the-art methods for ab-initio homogeneous reconstruction, and our focus is on investigating the tradeoff between amortized and semi-amortized inference. The deep-learning-based methods cryoFIRE and cryoDRGN (which adopt the auto-encoding approach similar to cryoAI) are specially designed for reconstruction of heterogeneous datasets. Still, we provide an additional comparison to cryoDRGN with the same setup as in the paper (see rebuttal PDF, Fig. 1-C). We run a homogenous ab initio reconstruction with cryoDRGN on synthetic and experimental datasets. The results show that our semi-amortized method outperforms cryoDRGN and reaches higher resolution reconstruction in most cases. We will add these results in the revision.\n\n**“Will the multi-head structure predict very similar poses?”.** The multi-head encoder does not predict similar poses. Figures 5 and 8, for example, show that different heads return different poses. Moreover, in the supplement, Fig. 6, we investigate the behavior of each head across all images by mapping their accuracy across different regions of pose space. As shown in the figure, the multi-head encoder adopts a divide-and-conquer approach where each head effectively specializes in pose estimation over localized regions with minimal overlap with other heads. This specialization behavior is a result of the winner-takes-all loss, which has been demonstrated to improve diversity in other prediction tasks as well [12, 13, 27]. Thus, the encoder is able to explore the pose space to help avoid local minima in the early stages of reconstruction.\n\n**Comparison on time and memory.** In the rebuttal PDF, we provide detailed comparison with baselines in terms of time, memory, and number of parameters (Table 1). As the semi-amortized method consists of two stages, we report numbers for each stage separately. In the first stage, the encoder has H=7 heads, while in the second stage, it is replaced with a pose module with size depending on the number of particles (N). For comparison, we include another baseline, cryoAI-explicit, in which the implicit decoder is replaced by an explicit one as in our method. The explicit and implicit decoders have 4.29 and 0.33 million parameters, respectively. Despite a larger decoder, our method is ~6x faster and uses ~5x less memory compared to cryoAI. Importantly, swapping the implicit decoder with an explicit one (cryoAI-explicit) significantly drops time and memory, indicating that the implicit decoder is a major computational and memory bottleneck. Yet, cryoAI-explicit uses ~2x more GPU memory for encoding and is ~1.5x slower than our method as it performs early input augmentation and runs the entire encoder twice per image. Our method, by augmenting the encoder head, saves memory and time during pose encoding. Finally, the amortized baseline, which uses an explicit decoder coupled with a multi-head encoder (H=7), uses more memory in decoding (negligible vs implicit decoder) and runs slower than the direct optimization stage of the semi-amortized method. We will add these performance comparisons in the revision.", "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": "", "extra_scores": {}}, {"review_id": "52BaykmS15", "reviewer": "Reviewer_Kgto", "summary": "This work introduces a new method for 3D ab initio estimation in cryo-EM using amortized inference. It relies on an existing technique which predicts the 3D rotation of a cryo-EM image using a convolutional neural network. That estimated rotation is then used together with a neural representation of the 3D volume to regenerate the projected image. The network weights (both of the pose estimation network and the 3D neural representation) are then optimized to reduce this auto-encoding loss.\n\nThe first important difference between the proposed framework and previous work is that this method produces several estimates using multiple heads and computes the loss with respect to the rotation estimates that yields the closest matching image. The second important difference is that after a few epochs of training, the pose estimation network is discarded, and the best pose estimates are refined using alternating minimization together with the 3D neural representation. The resulting algorithm is shown to perform well on both simulated and experimental datasets.", "questions": "– On like 54, what is meant by “we adopt an explicit paremeterization to further accelerate the reconstruction”? How does the parametrization affect convergence here?\n– The idea of “semi-amortized inference” is never actually explicitly described. Presumably this means switching from amortized estimation to pose refinement at some point.\n– How is L_{i,j} defined in eq. (6)?\n– Why can the proposed method not handle shifts? Should this not be just a matter of adding another output to the pose estimation heads?\n– For the synthetic and real datasets, the method switches from amortized inference to pose refinement after 7 and 15 epochs, respectively. How was this switching point determined and is there some way to automate this?\n– The labels (Left) and (Right) appear to be switched in the caption for Figure 4.\n– On line 227, “experiment” should be “experiments”.", "limitations": "The authors have largely addressed the limitations of the work in the manuscript.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The writing of the manuscript is clear, in particular with respect to the introduction of the problem and its context as well as the presentation and interpretation of numerical results. The method seems to be well thought through and successfully mitigates the important pitfalls of the 3D ab initio reconstruction problem.", "weaknesses": "The exposition of the pose inference network and the 3D neural representation are quite superficial and deserve a more detailed explanation. For example, how are the 3D rotations obtained from the convolutional backbone output? Are they encoded in axis–angle parametrization? If so, how are the outputs properly constrained? Similarly, the 3D neural representation is described as leveraging the Hartley transform and somehow uses a decomposition of this transform into its mantissa and exponent. How is this relevant to the actual neural network architecture?\n\nThe work would also benefit from more extensive testing on experimental datasets beside EMPIAR-10028. This is a relatively high-SNR dataset depicting a very large molecule (an 80S ribosome) and is not representative of experimental datasets in general.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your thoughtful feedback. In what follows, we address the main concerns and questions individually.\n\n**3D rotation parameterization.** In the supplement, Sec. A, we discuss the rotation parameterization and its optimization using PyTorch Autodiff.  Specifically, we follow cryoAI and design the convolutional backbone to output the six-dimensional representation commonly referred to as S2S2. To compute the rotation matrix from this representation, the 6D vector is split into two 3D vectors and normalized, resulting in the  unit vectors, $v_1, v_2$. We then compute the cross product ($v_1\\times v_2$) to get another unit vector ($v_3$), which becomes the last column of the rotation matrix. Selecting v1 as the first column, we compute $\\tilde{v_2}=v_3\\times v_1$, which is a new unit vector orthogonal to $v_1$ and $v_3$. Then, $R = [v_1, \\tilde{v_2}, v_3]$. During the auto-decoding stage, we switch to using a 3D axis-angle representation which requires fewer parameters. \n\n**\"What is meant by “we adopt an explicit parameterization to further accelerate the reconstruction”?** We meant that the decoder uses an explicit volumetric representation of density values to model the 3D structure rather than a multi-layer perceptron (MLP). Compared to MLPs, using the explicit representation yields faster reconstruction, which is especially helpful in our case because the multi-head architecture requires querying the decoder multiple times for each input image. We will revise the text to better clarify the distinction between explicit and implicit representations, and the approach taken in our work.\n\n**On Hartley transform and decomposition.** We use the Hartley transform to save memory and reduce redundancy of the transformation of a real signal; compared to the Fourier transform, it represents frequency information using real values instead of complex values. One notable issue with the Fourier/Hartley coefficients is their high dynamic range across different frequencies. To account for this, we parameterize each coefficient using a  mantissa and exponent such that we actually store two values per voxel.\n\n**Extensive testing on experimental datasets.** To ensure fair comparison, we benchmark on EMPIAR-10028 which is the same dataset used in cryoAI. We believe this is an important experimental benchmark as it is widely adopted in methodology papers (e.g. cryoDRGN). Also, we acknowledge that further study on more recent experimental datasets with smaller particles and a range of SNRs is an important direction to explore in future work.\n\n**“The idea of semi-amortized inference is never actually explicitly described”** Thank you for raising this issue. Semi-amortized inference is described in the introduction (L58-65) and depicted in Figure 1. We provide more detail in Sec. 4.2 where “semi-amortized inference” is explicitly mentioned with references to prior work. In the revised paper we will provide a more detailed description of “semi-amortization” and justify its use in our case. For example, as higher frequency details are resolved, the variance of the pose posterior tends to decrease, becoming unimodal. At this point, the gap between the amortized and variational posterior is mainly determined by the error in the pose estimate (predicted mean). The encoder might be too restrictive as a globally parameterized function leading to limited prediction accuracy hindering further refinement of the 3D structure. This motivates us to stop amortization and switch to direction optimization, which is the main idea of semi-amortized inference. We will clarify these points in the revision. \n\n**Loss function.** We define $L_{i,j}$ as the negative log likelihood for the i-th image given its j-th predicted pose. This is formalized in Eq. 3 and is computed in the Fourier space in practice.\n\n**Estimating in-plane shifts.**  As discussed in the ‘global response’ to all reviewers, we show in the Rebuttal PDF (Fig 1-A) that our semi-amortized approach extends in a straightforward fashion to estimation of rotation and translation while still outperforming the baselines. \n\n\n**Switching point.** We considered the switching point as a hyper-parameter and experimented with a range of values. We found that after a sufficient amount of optimization (i.e., 7 epochs for synthetic data and 15 epochs for real data) the reconstruction error is insensitive to when the switch occurs. Our intuition is that when switching too early, the pose posterior may have significant uncertainty, and the pose estimates may be too far from the correct basin of attraction to afford robust convergence. Therefore, it is better to switch late than early. Exploring how to identify the optimal switching point is a promising direction that we will explore in future work.\n\nFinally, thank you for bringing the typos to our attention. We will incorporate all feedback into the revision.", "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": "", "extra_scores": {}}, {"review_id": "n1qFtbMGSd", "reviewer": "Reviewer_xz2Y", "summary": "This paper introduces an approach to ab-initio homogeneous reconstruction that handles multi-modal pose distributions with a tailored encoder and accelerates pose optimization with semi-amortization. The approach uses a shared CNN feature extractor with multiple pose predictor heads, predicting several plausible poses for each image to account for pose uncertainty. Unlike the computationally expensive implicit networks used by cryoDRGN and cryoAI, this method employs an explicit 3D decoder, speeding up reconstruction. The encoder-decoder architecture is trained with a \"winner-takes-all\" loss, where the 3D decoder generates multiple 3D-to-2D projections, and the one with the lowest reconstruction error determines the loss. This method achieves faster convergence and more accurate poses compared to relying solely on the encoder's predictions. Evaluations show that the approach outperforms cryoAI and is competitive with cryoSPARC on both synthetic and experimental datasets.", "questions": "Overall, I am not convinced that the proposed semi-amortized approach can fundamentally improve pose prediction results without any significant improvements on the image encoder side. All experiments conducted in this paper show only minor performance improvements or even no improvements compared to cryoSPARC, which has no deep learning but only traditional optimization, does it imply that the actual bottleneck is the weak image encoder? The image encoder (VGG in this paper) is not pre-trained on a diverse cryo-EM dataset but trained from scratch on a per-scene dataset in a self-supervised manner. How can it be ensured that the extracted features can be applied for accurate pose prediction? I believe the image encoder is also a performance bottleneck. A PCA visualization of the feature space of the encoder (or calculating the cosine similarity) would be very helpful to validate if images with similar poses can be clustered and images with different poses can be far from each other. It is important for this work to conduct a performance analysis for both the encoder and decoder to identify the actual problem. Unfortunately, I see no analysis in the current version of this paper. I am happy to improve my ratings if my main concerns (including weaknesses) can be addressed by the authors.", "limitations": "Yes, there is a discussion of the limitations in this paper, but I believe the limitation of the image encoder used in this work should be also discussed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper proposes a two-stage ab-initio homogeneous reconstruction algorithm based on a novel multi-head transformer architecture to output multiple proposals of predicted poses. As a result, the performance on both synthetic and real datasets outperforms learning-based baselines.\n- A novel winner-takes-all loss has been introduced to specialize multi-heads of the decoder. Each head of the decoder can account for the pose estimation in a local region, mitigating the burden on the decoder.\n- This paper is well-structured and easy to follow, with appropriate references.", "weaknesses": "Despite its strengths, I see some major weaknesses in this paper:\n1. It is unfair to claim \"our semi-amortized method is faster than the amortized method of cryoAI\" (Line 79) as there are some minor modifications, like changing the decoder from an implicit representation to an explicit volume to reduce the computational cost. An explicit density volume or a feature volume can easily replace all coordinate-based neural networks. In this paper, there is no ablation study about the representation of the decoder or report on how much training and inference time has been reduced by this replacement. Explicit volume with multi-head predictions leads to higher memory costs during training. Advanced representations that can achieve a better trade-off between computational and memory costs, like TensoRF, triplanes, and hash encodings, are not well discussed or explored.\n\n2. The improvement of results is incremental. In Figure 2, the performance improvements on synthetic datasets are very subtle. On the experimental results, it seems cryoSPARC, which uses a traditional ab-initio method, achieves the best performance. Additionally, why do the resolutions of EMPIAR-10028's FSC curve all meet the resolution limit?\n\n3. Translation estimation is important for a complete ab-initio model. Why does this method not support translation estimation while cryoAI is capable?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the insightful comments and questions. Below we address the concerns.\n\n**speed comparison and decoder ablation.** To clarify, we perform an additional study, detailed in the rebuttal PDF Tab.1. We report reconstruction time per epoch, GPU memory usage, and number of parameters for various methods. As our method consists of two stages, we report numbers for each stage individually: the first stage uses an image encoder with H=7 heads, while, in the second stage, the encoder is replaced with a pose module with size depending on the number of particles (N). For comparison, we include another baseline, cryoAI-explicit, in which the implicit decoder is replaced by an explicit one as in our method. The explicit and implicit decoders have 4.29 and 0.33 million parameters, respectively. Despite a larger decoder, our method is ~6x faster and uses ~5x less memory compared to cryoAI. Importantly, swapping the implicit decoder with an explicit one (cryoAI-explicit) significantly drops time and memory, indicating that the implicit decoder is a major computational and memory bottleneck. Yet, cryoAI-explicit uses ~2x more GPU for encoding and is ~1.5x slower than our method as it performs early input augmentation and runs the entire encoder twice per image. Our method, by augmenting the encoder head, saves memory and time during pose encoding. Moreover, the amortized baseline, which uses an explicit decoder with multi-head encoder (H=7), uses ~7x more memory in decoding (negligible vs implicit decoder) and runs slower than the direct optimization stage of our method.  We will include these results and clarify this comparison.\n\n**Hybrid representations.** We share our findings on more recent hybrid representations in Fig.1-B of the rebuttal PDF. In a simplified setting with known poses, we optimize the structure represented by TensoRF which yields spurious blobs of density. We conjecture this is due to the low-rank assumption made by these representations which assume a 3D signal can be factorized into 2D axis-aligned planar signals. While this assumption may be adequate for structured, real-world scenes, it tends to break down for cryo-EM structures. Thus factorization might reduce the capacity leading to artifacts of wrong density masses. We believe effective adoption of hybrid representation in cryo-EM context is non-trivial and requires further investigation. We will add these failure cases in the supplement.\n\n**Significance of improvement.** Our primary goal is to show that amortized inference methods can be enhanced by combining with more traditional optimization approaches to achieve the best of both worlds: (1) enabling pose space exploration in early stages with amortized inference empowered by deep-learning architectures, and (2) boost accuracy of DL-based models to obtain results competitive with cryoSPARC. As shown in Fig. 2, on synthetic datasets of Spliceosome and HSP, our method outperforms cryoAI by ~1A and ~2A in resolution, respectively, while achieving the same or better resolution to cryoSPARC. Similarly, for experimental data (see last row of Fig. 2), our approach provides substantial qualitative and quantitative improvements compared to cryoAI. Our reconstruction is also comparable with that obtained by cryoSPARC. Moreover, our results with shift estimation (rebuttal PDF, Fig.1A) shows that our method is able to outperform cryoSPARC on experimental data. We believe our work is a step towards developing DL-based methods for cryo-EM reconstruction and reducing their gap with traditional methods like cryoSPARC.\n\n**Resolution for EMPIAR-10028.** The deposited images in EMPIAR-10028 have a box size of D=360, with 1.34 A/pixel. In our experiments, for fair comparison to cryoAI (which learns a $128^3$ volume), the particle images are downsampled to D=128 before reconstruction. With downsampled images, the best possible resolution to obtain (at the Nyquist rate) is 7.54A. Both cryoSPARC and our method reach this limit as shown by the FSC plot in Fig. 2. We should note that this is a mid-level resolution.\n\n**Estimating in-plane shifts.**  As discussed in the ‘global response’ to all reviewers, we show in the Rebuttal PDF (Fig 1-A) that our semi-amortized approach extends in a straightforward fashion to estimation of rotation and translation while still outperforming the baselines. \n\n**Image encoder bottleneck.** There is evidence in the literature that image encoders for pose prediction are a bottleneck in amortized inference for cryoEM [1, 2]. Indeed, while improvements to image encoder may have benefits, accurate pose prediction from cryo-EM images is non-trivial because (1) input images are extremely noisy, resulting in pose ambiguities, and (2) the job of the encoder is to invert the decoder and predict pose—intuitively, this requires knowledge of the 3D structure and is difficult to accomplish from the input image alone as the 3D structure changes during optimization. Our work pursues an orthogonal direction: we take advantage of the encoder’s ability to rapidly converge to sufficiently accurate poses such that direct optimization can be performed. Our approach significantly improves pose prediction accuracy (see Table 1 and Fig. 4). For instance, on the Spliceosome our method reaches <1 degree error while cryoSPARC and cryoAI reach ~1.5 and ~2.5 degree error. Also, on the HSP dataset, cryoAI fails to handle pose uncertainty (see Fig. 5) leading to high error: ~45 degrees on average. Our method reaches ~3 degree error vs ~6 degrees for cryoSPARC.\n\n**References**\\\n[1] Edelberg, et al. \"Using VAEs to learn latent variables: Observations on applications in cryo-EM\" (2023)\\\n[2] Klindt, et al. \"Towards interpretable Cryo-EM: disentangling latent spaces of molecular conformations\" (2024)", "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": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their thoughtful feedback. In the attached rebuttal PDF, we include results of additional experiments to clarify and address concerns raised by reviews, including additional baselines and datasets, and including empirical results using semi-amortized inference to estimate both particle rotation and translation, as requested by reviewers Kgto and xz2Y.

**Estimation of in-plane translation.**  Our paper originally focused mainly on the inference of 3D orientation as this is widely regarded to be more challenging than the inference of planar shifts per se. As such, we assumed that particles have been centered. Nevertheless it is straightforward to extend our method to estimate rotation and translation: We did so simply by allocating additional encoder heads to estimate translation parameters, as suggested in the reviews. We demonstrate the resulting model through application to experimental data with non-centered input particles. Fig. 1-A in the rebuttal PDF shows the reconstructed 3D density maps along with FSC curves on an experimental data (EMPIAR-10028), allowing comparison to cryoSPARC and cryoAI. A direct comparison of the FSC curves in Fig. 1-A shows that our semi-amortized approach outperforms cryoSPARC, and in Figure 4 of the original cryoAI paper clearly shows that cryoSPARC outperforms cryoAI on this same dataset. (Using the publicly available code for cryoAI by the authors, we also tried to replicate these results in order to show all FSC curves on the same plot.  However, in the presence of unknown planar translation, we found that cryoAI becomes trapped in local minima and hence we could not replicate the published FSC curves in the original paper.) We will add these results to the camera-ready version of the paper, along with the other changes or clarifications requested in the individual reviews.

Meta-review: This paper introduces a method for 3D ab initio reconstruction in cryo-EM using amortized inference. The paper makes several important contributions to state-of-the-art deep learning approaches, and it is clear and well-written.
Decision: Accept (poster)