Sparse Fourier Backpropagation in Cryo-EM Reconstruction
Authors: Dari Kimanius
Kiarash Jamali
Sjors HW Scheres
Venue: NeurIPS 2022 Accept (2022)
Keywords: cryo-EM
NeRF
Fourier space reconstruction
VAE
generative modeling

Abstract: Electron cryo-microscopy (cryo-EM) is a powerful method for investigating the structures of protein molecules, with important implications for understanding the molecular processes of life and drug development. In this technique, many noisy, two-dimensional projection images of protein molecules in unknown poses are combined into one or more three-dimensional reconstructions. The presence of multiple structural states in the data represents a major bottleneck in existing processing pipelines, often requiring expert user supervision. Variational auto-encoders (VAEs) have recently been proposed as an attractive means for learning the data manifold of data sets with a large number of different states. These methods are based on a coordinate-based approach, similar to Neural Radiance Fields (NeRF), to make volumetric reconstructions from 2D image data in Fourier-space. Although NeRF is a powerful method for real-space reconstruction, many of the benefits of the method do not transfer to Fourier-space, e.g. inductive bias for spatial locality. We present an approach where the VAE reconstruction is expressed on a volumetric grid, and demonstrate how this model can be trained efficiently through a novel backpropagation method that exploits the sparsity of the projection operation in Fourier-space. We achieve improved results on a simulated data set and at least equivalent results on an experimental data set when compared to the coordinate-based approach, while also substantially lowering computational cost. Our approach is computationally more efficient, especially in inference, enabling interactive analysis of the latent space by the user.

Sparse Fourier Backpropagation in Cryo-EM Reconstruction
Dari Kimanius
Kiarash Jamali
Sjors H.W. Scheres
MRC Laboratory of Molecular Biology
{dari, kjamali, scheres) (aarress@rrr-lll @mrc-lmb. cam ac. uk
Abstract
Electron cryo-microscopy (cryo-EM) is a powerful method for investigating the structures of protein molecules, with important implications for understanding the molecular processes of life and drug development. In this technique, many noisy, (wo-dimensional projection images of protein molecules in unknown poses are combined into one or more three-dimensional reconstructions. The presence of multiple structural states in the data represents major bottleneck in existing processing pipelines, often requiring expert user supervision. Variational auto- encoders (VAEs) have recently been proposed the data manifold of data sets with large number of different states. These methods are based on coordinate-based approach, similar Neural Radiance Fields (NeRF), to make volumetric reconstructions from 2D image data in Fourier- space. Although NeRF is powerful method for real-space reconstruction, many of the benefits of the method do not transfer to Fourier-space space,, e.g. inductive bias for spatial locality. We present an approach where the VAE reconstruction expressed on a volumetric grid, and demonstrate how this model can be trained efficiently through novel backpropagation method that exploits the sparsity of the projection operation in Fourier-space. We achieve improved results on simulated data set and at least equivalent results on an experimental data set when compared to the coordinate- based approach, while also substantially lowering computational cost. Our approach computationally more efficient, especially in inference, enabling interactive analysis of the latent space by the user.
attractive means for learning
Introduction
Single-particle electron cryo-microscopy (single-particle cryo-EM) has become the fastest growing method for structure determination of biological macromolccules recent years, and is becoming an increasingly important cornerstone of molecular biology research and drug discovery 2 3 4 Briefly, a thin layer of an aqueous solution containing multiple copies of the target molecule in random orientations is flash frozen to cryogenic temperatures and imaged in a transmission electron micro- scope. Subsequently, wwo-dimensional projection images of the Coulomb scattering potential of the molecules in the sample can be acquired into micrograph. Micrographs typically contain information sub-nanometer details of up millions of copies of the target molecules that are called particles.
To limit damage to the radiation nedistion-eensitive particles, micrographs are acquired under low-dose condi- tions, which leads low lignal-to-noise ratios (of typically less than 1/10) which requires averaging over many images (5). If sufficient number of approximately identical particles are captured in distinet projection orientations (or poses), computational methods can be employed reconstruct a three-dimensional volume of the scattering potential of the particle. However, the poses of the individual particles are unknown and many particles adopt multiple conformational states, or assem- blies of multiple molecules may vary in their composition. The unknown poses and the potential
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

assignments to multiple different structural states renders the problem ill ill-posed. Combined with the low signal-to-noise ratio in the images, this makes single-particle cryo-EM reconstruction challenging task.
The reconstruction process starts with the extraction of large number of image patches (on the order of 10 - 10²) from the micrographs, each containing a single projection image of the particle, captured in an unknown pose. In order to perform a three-dimensional reconstruction, the poses of all particles need to be determined, and particles need to be separated into structurally homogeneous subsets. Typically, this task is performed in two stages. First, in classification stage, structurally homogeneous subsets are identified and undesirable particles are removed. Second, in a refinement stage, the selected subsets of particles are reconstructed to obtain high-resolution structure.
Although generally desirable, automation and high throughput are essential in several areas of
application, e.g. fragment-based drug discovery [ screening for disease-relevant structures [6]. Fast computational processing pipelines, hased on hardware acceleration 7 g 2 have proven critical keep up with the high data rates provided by automated image acquisition software (IO). However, classification remains a bottleneck in existing pipelines, often requiring expert user input and repeated trials. Therefore, unsupervised algorithms that can replace existing classification steps will be essential in further improving single-particle cryo-EM. Moreover, because conformational changes are often central the biological function of proteins, more powerful methods deal with structural heterogeneity have the potential to enable biological discoveries which are currently out of reach of the existing methods.
2 Background and Definitions
2.1 Cryo-EM Reconstruction
Let {y: E CM², isis N) be set of images of (Fourier transformed) individual particles, where M is the image size and N is the number particles. Each ys is modelled as projection of a three-dimensional signal. The projection operation H (fxx) is described by the pose of the particle, de = (R.: 44) E SFF3), where R3 E 3x33 is a rotation or reflection matrix and 1; E TR3 is translation. Through the Fourier-slice theorem, the projection operation is realized by taking central slice from the three-dimensional Fourier transform of the signal, V E CM The projected image is then convolved with the point spread function of the optical system, which is described in the Fourier domain by multiplication with the contrast transfer function (CTF), C (N²), with known parameters KG Finally, experimental noise is modelled as additive and independent Gaussian in the Fourier domain, with scale that is frequency-dependen becomes:
(O, oy)) Thereby, the forward model
V C (R:) (1.)
o c
If the data set is structurally homogeneous, all particles are projections of a common V, and the objective is to estimate this V from all yı with unknown o; and known Kg. If the data set is structurally heterogeneous, there is no longer a single V to solve for. Rather, the solution is some space V and each i corresponds to V e V. In either case, maximum-a-posterinn postterii estimate, where we marginalize over the nuisance variable Ø:
represent the task hand as the solution to
V argmax V
Σ log / (I P(ulli.))aa P(M)
In the above, P(V) represents a log prior over the solution space. A conventional approach to solving this problem is clustering the space V to discrete set of K structural states, which means that the integral in (2 becomes a sum over a fixed K. Popular software packages for sinle-partich cryo-EM, like RELION [I] and cryoSPARC (9. employ iterative algorithms, including expectation- maximization and gradient decent, to find solution for the problem with K distinct structural states. For computational reasons K is typically kept below 10. The expectation over & is also computationally expensive, which is why approximate approaches, including branch-and-bound methods [9, or variations thereof [12,, are regularly used.

---

2.2 Cryo-EM Generative Modelling
By parametrizing V differently, the cryo-EM reconstruction task can also be described with generative modelling. Using the autoencoder framework, let so be the encoder parameterized by and I be the decoder parameterized by w. Each particle image yi can be embedded in some latent encoding 24 4 fo().. Then, using the decoder, we reconstruct the original yı by minimizing a loss function L that will contain some priors and a log- ligeliidlhoodd term U/ (x²) o H. (Ø,) gw( To reduce the computational cost of generative models, the expectation in 2 can be replaced with the highest probability pose for each particle, di determined using homogeneous reconstruction approach.
The generative modelling approach enables replacing the fixed number of discrete classes, K, with a model that can account for a larger variability in V. This is especially relevant when dealing with structural heterogeneity that cannot be easily separated in distinet states, like continuous motions. Protein complexes often populate states that represent continuous transitions as part of their functional cycle. The continuous distribution of the learnt latent representation, expressed by Iw. provides a natural representation of the corresponding protein conformational landscape. Assuming that the experimental procedure does not perturb the distribution of structural states in the sample, the latent representation could provide insights into the dynamics and perhaps the biological function of the protein.
Moreover, although not strictly equivalent, the expansion of model capacity provided by the generative modelling approach can be approximated in the discrete setting by increasing K to large number. It has been shown that robustness of canonical non- non- parammeettic classifiers are critically dependent on large K [3]. Therefore, we expect generative models with larger model capacity improve robustness over discrete heterogeneous reconstruction methods with low K.
Related work
The generative modelling approach described in the previous section can be classified based on the space where the reconstruction takes place; either one decodes the volume in Fourier-space or in real-space. Both methods have their own benefits, as described below.
3.1 Fourier-space reconstruction
cryoDRGN 14 is currently the most popular approach to Fourier-space reconstruction using the generative modelling framework. In cryoDRGN, variational autoencoder (VAE) is used to encode particle image heterogeneity into a Gaussian parameterized latent space. In particular, cryoDRGN encodes each particle image with μ4; o3 (vt).. To obtain the latent space encoding, they sample zi ~ N(us; o:). Next, they independently reconstruct all Fourier components that correspond to the specific pose, &i,, of the particle image using coordinate-base describe the 3D coordinates. For this approach to work, the hidden variables, including the projection angle, are determined beforehand through global angular searches (see section 2. Similar neural radiance fields (NeRF) [15]. the decoder is typically a multilayer perceptron (MLP) that receives a positional embedding of the coordinates of each Fourier component of the image and its latent vector.
decoding function gu (=;, where
Another approach to Fourier-space reconstruction using the generative modelling framework is 3D variational analysis (3DVA) [16]. This method does not use deep learning gradient- gradient-based methods for learning its parameters, but starts from a reconstruction and formulates the generative process as a probabilistic principal components analysis to obtain a set of orthogonal basis (A) of 3D Fourier volumes representing the principal components that best fit the heterogeneity. This method optimizes the log-likelihood with an expectation expectation-maximine maximization approach.
Due to memory restrictions, cryoDRGN cannot reconstruct full volumes in a single batch, since the reconstruction of each coordinate in the volume requires a pass through the decoder. This prohibits backpropagation with respect to the full volume because the intermediate activations for (MM) coordinates have to be stored. This quickly becomes intractable. Although, this issue can be ameliorated with gradient checkpointing [17, it comes at the cost of extra training time. This hinders the application of real-space priors, e.g. solvent masks, mass preservation, and data-driven priors [8
3

---

3DVA does not have this issue since full volume reconstructions are more computationally efficient than they are with cry as volumes are weighted averages of the principal components and the original reconstruction. However, in contrast to 3DVA, VAE can be trained in a fully differentiable manner, which confers unparalleled flexibility to the method and allows easy modification of the loss function. For instance, one can easily add an auxiliary loss that imposes additional constraints to either the latent space or the output. It's noteworthy that gradient descent optimization has shown superior convergence properties over that of expectaion-mmxxiizattii Furthermore, since a VAE is trained to do amortized variational inference, the encoder can generalize to previously unseen data without the need for re-training. The approach presented in this paper is based on a VAE framework and hence enables both easy modification of the loss function and amortized variational inference with an encoder. On the other hand, we modify the decoder to gain some of the computational benefits of the 3DVA approach, which combined with the VAE framework enables fast and efficient full volume inference and differentiation.
for de novo reconstruction [9].
3.2 Real-space reconstruction
3DFlex 19 method implemented in the cryoSPARC software package that is based real-space reconstruction. In this approach, a consensus real-space reconstruction is learnt and then deformed by a flow field that is parameterized by an MLP, given latent representation for each image. Instead of viewing the problem as generating independent volumes, it deforms one consensus volume such that it maximizes the log-likelihood of the individual particle images.
Another real-space based method, e2gmm (20,. also learns flow fields, but these are applied to a Gaussian mixture model instead of volumetric representation. In this method, the encoding is done similar to cryoDRGN, where particle images are encoded a Gaussian latent space.
Both real-space approaches enforce pass-pres inductive bias, which can improve reconstruc- tion quality. Although not yet explored, both approaches would in principle enable computationally efficient application of other real-space priors, like solvent masking data-driven priors.
disadvantage of the real-space methods that the application of the CTF requires two Fourier transforms (forward and backward) to evaluate the log-likelihood loss. To avoid expensive 3D Fourier transforms, the implementation of both of the above mentioned methods first project along the orientation of each particle image and then apply a 2D Fourier transform. However, the projection operation is computationally expensive real-space, especially when dealing with volumetric representations, because the entire volume has to be accessed.
This contrasts with Fourier-space methods, where the CTF can be applied by multiplication and the projection operation can be performed by accessing a small fraction of the Fourier components in the volume, due to the Fourier-slice theorem. To exploit the latter, in what follows we present an approach where the decoder is not function of the 3D coordinates, but rather generates the entire volume in Fourier-space in a single pass, and present a novel approach for training such model in computationally efficient manner.
Method
Here, we present an approach for learning a volumetric Fourier reconstruction given a latent represen- tation in a VAE framework. Our decoder learns to generate an entire volume for each latent vector This volume consists of the model Fourier coefficients on a regular cubic grid with fixed coordinates. To enable efficient training, we employ sparse backpropagation method that limits evaluation and gradient calculations, in both the forward and backward pass, to the subset of model parameters that are affected by each projection image. This is implemented in PyTorch [21] with custom CUDA kernels. We have not preformed any runtime optimizations on these kernels. The code used in this paper is available at guihuu..omddkiimniuuusseeeye together with the instructions for compiling and running it.
4.1 Motivation
As previously noted, cryoDRGN uses a coordinate-based decoder, similar to NeRF, to parameterize the reconstructed volumes. However, unlike the real-space case to which NeRF is applied, cryoDRGN's

---

A
B
H(du) Vo
Bias
C(s)
Sw
H(dd) 900)
Figure 1: (A) shows the sparse backpropagation method. The projection image yi is input into the encoder fo which results in set of scalars (k) that weigh the learnable and interpretable structural basis set (kk) The weighted sum of the projection H (Ø:) is then CTF corrupted with C. (Ks). Finally, the input yi and the output are compared the loss function, L. (B) shows the bias parameter and
two structural bases of the model in real-space for the reconstruction in section 5.1
outputs are defined in Fourier-space, where the spatial smoothness assumption is less valid. In particular, this parameterization suffers from the spectral bias of neural networks [22] 23 That is, the inability of neural networks capture local, high-frequency changes. This makes Fourier-space representations more challenging to learn with a coordinate- -based approach. 'his issue can be ameliorated by incorporating high-frequency positional embedding [24]. but an alternative approach that avoids positional embedding altogether might be better suited for this particular problem. A potentially more natural approach learning discrete Fourier representations is on regular cubic grid, a so called voxel representation. Similar approaches in real-space have been shown to improve model capacity and training efficiency 225). In particular, we suggest to learn a set of volumes (V(K) E CM*) that are combined linearly into the final output.
While potentially better behaved, this method would ultimately result in the need to store B x K volumes for the backward pass during training, where B is the batch size. This computationally wasteful and has large memory footprint, which is why we propose novel training method. The following sections outline how we exploit the sparse property of the cryo-EM forward model in Fourier-space to more efficiently evaluate and train volumes (k) with backpropagation.
4.2 Structural Basis
We define the decoder as
K (b)) (z) + V())
Iu (2)
(3)
Here, ((k) C trainable volume-sized arrays that refer to as structural basis; s R are the latent-dependent weighting factors; and (0(0) is trainable bias term. K is strongly correlated with the expressive capacity of the model. Together with the volume size, K has a major impact on the total memory footprint. The volume size is determined by the data image size, i.c. by the image physical pixel size and the maximum resolution at which the model needs to be evaluated. For experiments in this work use K 16; s. is small three layer MLP with residual connections (26].
This approach has similarities to 3DVA, in that volumes are linearly combined generate the reconstruction results. However, in our approach the structural bases are learnt together with the representation an end-to-end differentiable setup, including an encoder parameterized by a neural network. The VAE is trained to perform amortized variational inference and we use gradient based optimization, which makes the computational burden comparable to eryoDRGN rather than 3DVA (see section 3.2
5

---

4.3 Training
For the loss the log-posterior expression used in the RELION-4.0 implementation 27 which is
(1 - A) C² (Ks) oy
111
C (me) o (o.) (z())) +XP(g(zz)).
(4)
Here, is a grequencc-depenndenn regularization parameter; P is a suitable log-prior; and By is the beqquene-deepenee noise amplitude of the data. In this work we set P = O.
During training, the Fourier slice belonging to the known projection plane of each data sample (particle) is extracted from the cubic grids of the structural bases, using tri-linear interpolation, similar to (28]. The extracted slice is then real-space shifted through phase-shifte in Fourier-space, and modulated with the CTF, in order to calculate its log- legliike loss. Although the projection angle, shift, and CTF parameters are given, it possible to jointly optimize them with the model parameters, since all the operations involving them are differentiable. However, in this paper we will keep them fixed.
Since the sparse backpropagation method generates gradients for a subset of the model parameters, is natural use the sparse version of the Adam optimizer (29], where the gradient moments are only updated for the affected parameters. This requires that the gradients are stored in sparse representation, which is less memory efficient for large batch sizes. However, since each batch contains particles from different orientations, when the batch size is sufficiently large, enough of the volume would receive gradients, which enables the use dense gradients. This makes the memory footprint of the sparse projection operation invariant to batch size.
Because of the low signal-to-noise ratio in cryo-EM data, can approximate the CTF corrected dg as the averaged spectral amplitude of the data divided by the average CTF amplitude, here denoted as oy. Let B = {E : E < [<] { + T). where 7 depends on the discretization used for the discrete Fourier transform. Then, for frequency & we get
IN SE (E) Ld
(5)
Em Jg C² (E, ki:) dg +
In the above, the integral is evaluated over the frequency shell with radius between and + r and the sum is over all the images in the data set. The value of is set to a small positive number for numerical stability.
Although the real-space particle images are standardized to zero-mean and standard deviation of one, the numerical values of their Fourier transforms typically have values spanning several orders of magnitude. To improve convergence of the model parameters, we perform a spectral normalization of the Fourier transforms of the data samples, where the normalization factor a function of Fourier frequency. In this work, we use o, as the spectral normalization factor, which makes the effective By in 4 approximately equal to one. To reverse the effect this has on the spectral weighting of the MSE-loss, we set À proportional to o²
We note that it is the sparse access of the projection operation H. (-) that enables an efficient evaluation and training of 9w. For the sparse backpropagation method, we define single fused operation that evaluates both gw and H (-) a single function call. In the forward pass, this operation disregards parameters in eu that are omitted by H (-). Similarly, in the backward pass, gradients belonging to omitted parameters are not calculated.
Additionally, to minimize the memory footprint, we only store Fourier components in ((k) in sphere with a radius corresponding to the Nyquist frequency, as defined by the user. This requires explicit coordinate arrays, which negatively impacts performance, but reduces the memory footprint by almost a factor of two, compared to storing Fourier components of the entire cube.
We evaluate our method one simulated data set and experimental cryo-EM data set, which was downloaded from the publicly available Electron Microscopy Public Image Archive (EMPIAR) 3D with CCO licence [I]. All caleulations were performed on NVIDIA RTX 3090.

---

Results
5.1
Simulated data set
UMAP
Max
CryoDRoN Small
Ours
CryoDRGNN
0.5
Large
A
UMAP
B
0.22 36.0
Resolution [A]
5.68
Dlosed
Intemedizatss Open
Closed
Intermediates
Open
0.85
UMAP
Ours
Mean FSC
F
Small
CryoDRGN Large
C
UMAP 2 D
0.70 10
Epoch
40
Figure 2: Reconstruction results of the simulated data set of SARS-CoV-2 spike protein after 50 epochs of training. (A) shows the learnt latent representation of our method versus (C) cryoDRGN- small. (B) shows the side view example reconstructions for our method and (D) shows the same for cryoDRGN-small. The outline of each ground truth structure is shown with a black contour and orange fill to highlight discrepancies between the reconstructions and the ground truth. (E) shows the average FSC to each ground truth for our method (green) versus the cryoDRGN-small (grey) and cryoDRGN- (orange). The dashed line shows the maximum possible FSC. (F) shows the average FSC per epoch. The average of the maximum FSC shown in (E) is 0.95.
The simulated data set consists of projection images of 3D volumes generated from frames of molecular dynamics simulation of the SARS-CoV-2 spike protein (DESRES PAERRES-ANTNN-LLE 10807500897850) [32. published under CC BY 4.0 license [33] Ten different conformational states were extracted out of the simulation and ten thousand projection images with random poses were generated for each state, amounting to 100,000 images in total. Three frames were selected of the so-called open state; three frames of the closed states 34; and four frames of the intermediate states between the open and the closed states. The image size is M = 96, with physical pixel of 2.7 Å. The simulated images were then distorted with additive Gaussian noise and CTF, The average signal to noise ratio is 1/10 and the defocus range is between 0.2 and 2 um. The pose for each particle image was then estimated by running a RELION consensus refinement, and kept fixed in the following steps. This data set, along with the PDB files used to create it and the refinement results, can be found at (1180looooo
ocnodo.org/recorr// 182156 (DOI
We trained our VAE with 8 latent dimensions for 50 epochs. After the end of training, we applied uniform manifold approximation and projection (UMAP) 35 to the distribution of each particle image's s latent vectors. Evaluating our decoder for points belonging to each cluster revealed that several distinct conformations where separated for each of the three states (see figure ( In total seven separated clusters, were observed, amongst which, four belong to single conformations and three are made up of mixtures of at most two conformations. CryoDRGN only fully separated one distinct conformation, belonging to the closed state. The rest remained in clusters mixed with other conformations (see figure 2C).
Running cryoDRGN on this data set with the default network architecture (cryoDRGN-small) and a larger one CryoDRGN-larrge results in worse reconstruction performance compared to our method, as can be seen in figure 2 Examples of reconstructions from the final epoch for our method and cryoDRGN-smen
are shown in figure2]pane B and D, respectivly. Since this simulated data, we can compare the results to the ground truth. Comparison to the outline of the ground truth, reveals that the cryoDRGN results have more missing structural features and overall lack the same level of details seen in the results of our method, particularly in the flexible region. Furthermore, we can calculate the Fourier shell correlation (FSC) between the generated and the ground truth structures. The FSC calculated here the normalized complex correlation between Fourier components of the solvent
7

---

masked reconstruction and the known ground truth for each frequency band [36], averaged over 30 samples (three out of each conformation). The results for this analysis are visualized in figure 2 At the Nyqvist frequency (5.4 Á resolution) the FSC drops due to aliasing [37]. This is also visible for the maximum FSC, which is calculated by reconstructing with the known conformational and pose assignments. Figure 2F shows the average FSC to ground truth for each epoch. Neither of the methods converged in 50 epochs. However, starting around epoch 10, cryo large diverges from ground truth, which is seen as the average FSC dropping, due to overfitiigg to noise. Additionally, increasing batch size gave worse results for cryoDRGN, as can be seen in the in supplementary table D and supplementary figure 5
took -ryoDRGN-smalll 213 minutes to train for 50 epochs, while our method takes 40.3 minutes. Evaluating the decoder takes ms on average for each 3D structure, which is on the same order of magnitude as the time for the inverse Fourier transform of ms. This can be compared to 250 ms for the cryoDRGN-smail architecture.
5.2 Ribosome data set
We also applied our method an experimental cryo-EM data set of 80S ribosome particles from the Plasmodium falciparum parasite (38). This data set has been used repeatedly as a test case for new algorithms [39] [4 40 41]. We used 108,72 7322 particle images, with M 160 and pixel size of 3.0A. We used 10 latent dimensions for this data set. To evaluate the ability to handle defective particle images, we also included 78,4 78,442 images that where classified by RELION as not containing ribosomes, containing partially (dis))assembled ribosomes, along with images classified as intact ribosomes, This brings N to 187.174.
After few epochs, the defective and empty particles were separated out the top principal component and the separation is clearly preserved in a UMAP of the latent space after 50 epochs (fi Additionally, we notice separation of different classes of particles in each of the two major classes. In the cluster of the intact ribosomes we observed several subtle continuous conformational transitions, including the "head swivel" 42 (not shown in figure) and the rotation of the small subunit, which is part of the "ratcheting" cycles [43] (figure BB) Although the figure shows the two endpoints of the rotation, continuous transition between the two states can be observed along the marked path in the latent representation. In the other cluster, containing the defective particles, we observed varying classes of particles. Images of partially assembled ribosomes are separated from images lacking any protein features (figure 3D).
Next, trained DryoDRGN- can be seen in figure 3 E-H. Unlike in [ we do not have access to the ground truth volumes in order to compute FSC. However, qualitatively,
this data set and achieved similar results our method, as
able to generate similar clustering and reconstruction performance compared to cryoDRGN, at lower computational cost. It should also be noted that training CryoDRGN-smail on this data set gives noticeably worse results, as shown in the appendix. Training took 210 minutes using our method, compared to 1638 minutes for cryoDRGN, Moreover, evaluating the decoder takes 10 ms on average for each 3D structure, compared to 480 for cryoDRGN.
We kept the batch size of cryoDRGN to the recommended value and what gave the best results for the simulated data set (see section
6 Discussion
The potential improvement in robustness make VAE approaches candidate for replacing the discrete classification methods that are commonly used in cryo-EM single-particle reconstruction pipelines. We identified the increased computational cost of deep learning based methods as an issue preventing this shift, and present an alternative method with potential to solve this issue.
Our results on simulated data illustrate that our method yields improved results compared to cryoDRGN with less than one-fifth the computational cost. We show that our method yields reconstructions with higher correlation to the known ground truth. For an experimental test case, our method converges onto a solution that at least as useful as the one provided by cryoDRGN, with one-cighth of the computational cost and a smaller memory requirement.
8

---

A
B
D
UMAP
UMAP
F
G
H
U
Figure Reconstruction results of the Plasmodium falciparum 80S ribosome data set after 50 epochs of training using our approach (A-D) or cryoDRGN-large (F-H). (A, E) Reconstruction of the ribosome, showing the small and large subunits. (B, F) Two structures (green and yellow) showing the two rotational states. (C, G) The learnt latent space, dimensionality-reducee Clusters containing particle images of intact ribosomes marked green and empty images or partially assembled ribosomes in grey. (D, H) Two reconstruction examples from the grey cluster, with ribosome lacking small subunit on the left and an empty particle containing mostly noise or low contrast features on the right.
with UMAP.
From the user's perspective, our approach also offers a practical advantage in allowing faster interactive analysis of latent space. Calculation volume for given point in latent space takes on the order of half second in cryoDRGN. Our decoder reconstructs entire volumes in one pass, which takes only a few milliseconds. This allows for a more convenient interactive analysis of latent space, where the user inspects multiple volumes on-the-fly. along interactively selected points of interest in latent space. Intuitive graphical tools enable inexperienced users to explore the heterogeneity in their data, and casily identify particle subsets to be used for subsequent refinements. This holds the potential to make cryo-EM accessible to a broader range of researchers.
This approach can also be combined with global angular searches because of the reduced overall computational cost and the fast full-volume generation (see section (3.T). This could enable recon- struction of data sets where the initial angular assignment based on homogeneous reconstruction fails. Additionally, this could accelerate replacement of the discrete 3D classification where global angular searches are integrated in all popular implementations.
Because of the low signal-to-non noise ratios in the experimental images, cryo-EM reconstruction is prone to overfitting, and explicit regularization of the output maps has played major role in the most popular software packages [II] []. Operations like (back-)prrojectio
and the application of the CTF are performed conveniently in Fourier-space. Most software packages thus perform reconstruction in Fourier-space, where they also regularize the power of the Fourier components of the reconstruction. Our approach, which also performs reconstruction in Fourier-space is well suited to apply similar regularization in Fourier-space, although we did not explore this here. Moreover, because reconstructed maps are generated at relatively low computational cost, one could also apply real-space priors, like solvent masks or data driven priors [8]. Because these priors are differentiable, one could backpropagate through them to improve this method's structural bases and network parameters. In this way, one could directly use the data-driven prior to improve the reconstruction in computationally efficient manner. Improved algorithmic robustness through the use of priors on the reconstructed set of volumes has the ability to increase the radius convergence of the reconstruction algorithm and allow more challenging data sets to be analyzed.
The structural basis, and the linear manner in which they are combined, also allow for better interpretability of what the model is learning. In addition to facilitating troubleshooting, the structural basis can potentially be utilized for two purposes. Firstly, they could be used to impose 3D priors on the reconstruction in a more direct manner than backpropagation. Real-space priors that are invariant to structural heterogeneity, e.g. solvent masks or structural symmetries, could be imposed to them directly rather than part of the gradient-bar based training loop. Secondly, real-space features of the

---

structural basis could be used at the analysis stage for automatic annotation of the latent space, e.g. through center of mass location or real-space PCA.
However, the use of the structural basis also implies a potential weakness in our method. the heterogeneity of the data set cannot be embedded into linear subspace of K dimensions, our method would not be able to correctly reconstruct the volume manifold. This limit is theoretically not present in cryoDRGN, because they parameterize their reconstruction with neural network. However, in practice, the capacity of the cryoDRGN network is also limited by computational cost, and we have so far not observed issues with the expressive capacity of our model. It should be noted that, due to the noise present in the cryo-EM reconstruction problem, one should expect that the dimension of the possible solution space V is much smaller than N. Therefore, the limit to the reconstruction of heterogeneity due to noise would dominate the limit due to the representation capacity of the structural basis set.
7 Broader Impact
The cost of data processing can be major limiting factor for research groups with limited access to computational resources. By improving robustmess and computational efficiency, we expand the accessibility of the method to include more diverse research community, which can in turn have an accelerating impact on the overall research quality of the field. Moreover, the potential for improved automation has implications for high-throughput structure determination. Along with improvements in microscope data acquisition, rapid processing can enable faster iterations in e-g. protein design or rragment-based drug discovery. Thereby, our method has the potential to impact medical research and the development of new therapeutics.
8 Acknowledgements
We are grateful to Johannes Schwab, Jasenko Zivanov and Daniel Mihaylov for helpful discussions, and to Jake Grimmett and Toby Darling for help with high-performance computing.
9 Funding
This work was funded by the UK Medical Research Council (MC_UP A0__0025_100 1013 to SHWS) and the European Union's Horizon 2020 research and innovation programme (under grant agreement no. 895412 to D.K.).
References
[1] Yifan Cheng. Single-particle cryo-em how did it get here and where will it go. Science, (6(())))) 2018. [Cited on page
[2] Rafael Fernandez-Leiro and Sjors H.W. Scheres. Unravelling biological macromolecules with cryo-electron microscopy. Nature, 537:339-346, 2016. [Cited on page []
[3] Taiana Maia de Oliveira, Lotte Beek, Fiona Shilliday, Judit É. Debreczeni, and Chris Phillips. Cryo-em: The resolution revolution and drug discovery. SLAS Discovery, 26: 17-31, 2021. [Cited page [
[4] Michael Saur, Michael Hartshorn, Jing Dong, Judith Reeks, Gabor Bunkoczi, Harren Jhoti, and Pamela A Williams. braseent-based drug discovery using cryo-em. Drug Discovery Today, 25(3):485 490, 2020. [Cited on pages and
[5] Richard Henderson. The potential and limitations of neutrons, electrons and x-rays for atomic resolution microscopy of unstained biological molecules. Quarterly Reviews of Biophysics, 28(2):171-193, 1995. [Cited on page [[]
[6] Soñia Lövestam and Sjors Scheres. High-throughput cryo-em structure determination of amy- loids. Faraday Discussions, 2022. [Cited on page 2
10

---

[7] Dari Kimanius, Björn O Forsberg, Sjors HW Scheres, and Erik Lindahl. Accelerated cryo-em structure determination with parallelisation using gpus in relion-2. elife, Sxe18722, 2016. [Cited on page 2
[8] Jasenko Zivanov, Takanori Nakane, Bjöm Forsberg, Dari Kimanius, Wim JH Hagen, Erik Lindahl, and Sjors HW Scheres. New tools for automated high-resolution cryo-em structure determination in relion-3. elife, 7:e42166, 2018. [Cited on page 2
[9] Ali Punjami, John L Rubinstein, David J Fleet, and Marcus A Brubaker. cryospare: algorithms for rapid unsupervised cryo-em structure determination. Nature methods, 14(3):290-266, 2017. [Cited on pages and [
[10] Philip R Baldwin, Yong Zi Tan, Edward T Eng, William J Rice, Alex J Noble, Carl J Negro, Michael A Cianfrocco, Clinton S Potter, and Bridget Carragher. Big data in cryoem: automated collection, processing and accessibility of em data. Current opinion in microbiology, 43:1-8,
2018. [Cited on page 2
[11] Sjors HW Scheres. Relion: implementation of bayesian approach cryo-em structure determination. Journal of structural biology, 180(3):519-3) 19-530, 2012. [Cited on pages 2 and 9
[12] Sjors HW Scheres, Mikel Valle, and José-Mardt Carazo. Fast maximun-likelhhod refinement of electron microscopy images. Bioinformatics, 21
2005. [Cited on
page 1
[13] Yizhen Wang, Somesh Jha, and Kamalika Chaudhuri. Analyzing the robustness of nearest neighbors adversarial examples. In International Conference Machine Learning, pages 5133-5142. PMLR, 2018. [Cited on page
[14] Ellen D. Zhong, Tristan Bepler, Bonnie Berger, and Joseph H. Davis. Cryodrgn: reconstruction of heterogeneous cryo-em structures using neural networks. Nature Methods, 18:176-185, 2021.
[Cited on pages 3 and [
[15] Ben Mildenhall, Pratul Srinivasan, Matthew Tancik, Jonathan T. Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes neural radiance fields for view synthesis. In ECCV, 2020. [Cited on page
[16] Ali Punjani and David J. Fleet. 3d variability analysis: Resolving continuous flexibility and dis- crete heterogeneity from single particle cryo-em. Journal of Structural Biology, 21702) 2021. [Cited page
[17] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. CoRR, abs/ /6/1106444....4. 2016. [Cited on page
[18] Dari Kimanius, Gustav Zickert, Takanori Nakane, Jonas Adler, Sebastian Lunz, Bibiane Schönlieb, Ozan Öktem, and Sjors H. W. Scheres. Exploiting prior knowledge about biological macromolecules in cryo-em structure determination. IUCrJ, 8(1):60-75, 2021. [Cited on pages
and
[19] Ali Punjani and David Fleet. 3d flexible refinement: Structure and motion of flexible proteins from cryo-em. bioRxiv, 2021. [Cited page
[20] Muyuan Chen and Steven J. Ludtke. Deep learning-based based model for characterizing variability in cryo-em. Nature Methods, 18(8):930 936,9999-93600-- 2021. [Cited on page 4
gaussian mixture
[21] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Bue, E. Fox, and Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8035- Curran Associates, Inc., 2019. [Cited on page

---

[22] Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 5301- 5310. PMLR, 09-15 Jun 2019. [Cited on page
[23] Ronen Basri, Meirav Galun, Amnon Geifman, David Jacobs. Yoni Kasten, and Shira Kritchman. Frequency bias neural networks for input of non-uniform density. International Conference on Machine Learning, pages 685-694, PMLR, 2020. [Cited on page 5
[24] Matthew Tancik, Pratul P. 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 in low dimensional domains. NeurlPS, 2020. [Cited on page
[25] Lingje Liu, Jiatao Gu, Kyaw Zaw Lin, Tat-Seng Chua, and Christian Theobalt. Neural sparse voxel fields. Advances in Neural Information Processing Systems, 33:15651 2020.
[Cited on page(5) 5]
[26] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016. [Cited on page 51
[27] Dari Kimanius, Liyi Dong, Grigory Sharov, Takanori Nakane, and Sjors HW Scheres. New tools for automated cryo-em single-particle analysis in relion-4.0. Biochemical Journal, 478(24):4169- 4185, 2021. [Cited on page [
[28] Karen Ulrrich, Rianne van den Berg, Marcus Brubaker, David Fleet, and Max Welling. Differ- entiable probabilistic models of scientific imaging with the fourier slice theorem. arXiv preprint 1906.0799777 2019. [Cited on page [6]
[29] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. [Cited on page 6
[30] Andrii Iudin, Paul K Korir, José Salavert-Torres, Gerard Kleywegt, and Ardan Patwardhan. Empiar: a public archive for raw electron microscopy image data. Nature Methods, 13(5):387- 388. 2016. [Cited on page 6]
[31] CcO license. Accessed: 2022-05-19, 19. [Cited on page 6
[32] Molecular Dynamics Simulations Related to SARS-CoV-2, D. E. Shaw Research Technical Data, 2020. http://wwww..ccl.eeeeeertt.emecs htt.:/ooooo...soco........ [Cited on page 7
[33] Cc by 4.0 license. http:///eeaatrrromoooooooooooeea 05- 19. [Cited on page
Accessed: 2022-
[34] Yuan Huang, Chan Yang, Xin-feng Xu, Wei Xu, and Shu-wen Liu. Structural and functional properties of sars-cov-2 spike protein: potential antivirus drug development for covid- 19. Acta Pharmacologica Sinica, 41(9):114-11144 2020. [Cited on page 7
[35] Leland McInnes, John Healy, Nathaniel Saul, and Lukas Großberger. Umap: Uniform manifold approximation and projection. Journal of Open Source Software, 3(29):661, 2018. [Cited on
[]
[36] Richard Henderson, Andrej Sali, Matthew I. Baker, Bridget Carragher, Batsal Devkota, Ken- neth Downing, Edward H Egelman, Zukang Feng. Joachim Frank, Nikolaus Grigorieff, et al. Outcome of the first electron microscopy validation task force meeting, Structure, 20(2):205-214, 2012. [Cited on page S]
[37] Steven W Smith et al. The scientist and engineer's guide to digital signal processing, 1997. [Cited on page 8
12

---

[38] Wilson Wong, Xiao-chen Bai, Alan Brown, Israel S Fernandez, Eric Hanssen, Melanie Condron, Yan Hong Tan, Jake Baum, and Sjors HW Scheres. Cryo-em structure of the Plasmodium falciparum 80s ribosome bound to the anti-protozoan drug emetine. elife,, 3:e03080, jun 2014.
[Cited on page []
[39] Takanori Nakane, Dari Kimanius, Erik Lindahl, and Sjors HW Scheres. Characterisation of molecular motions in cryo-em partieearticle data by multi-body refinement in relion. eLife, 7:e36661, jun 2018. [Cited on page(8) []
[40] Cyril F. Reboul, Frederic Bonnet, Dominika Elmlund, and Hans Elmlund. A stochastic hill climbing approach for simultancous 2d alignment and clustering of cryogenic electron mi- croscopy images. Structure, 24(6):988 996, 2016. [Cited on page
[41] Bernard Heymann, Roberto Marahini, Mohsen Kazemi, Carlos Oscar S. Sorzano, Maya Holmdahl, Joshua H. Mendez, Scott M. Stagg, Slavica Jonic, Eugene Palovcak, Jean-Paul Armache, Jianhua Zhao, Yifan Cheng, Grigore Pintilie, Wah Chiu, Ardan Patwardhan, and Jose- Maria Carazo. The first single particle analysis map challenge: A summary of the assessments. Journal of Structural Biology, 204(2):291-300 2018. [Cited on page
[42] Andreas H Ratje, Justus Loerke, Aleksandra Mikolajka, Matthias Brünner, Peter W Hildebrand, Agata I. Starosta, Alexandra Dönhöfer, Sean Connell, Paola Fucini, Thorsten Mielke, et al. Head swivel on the ribosome facilitates translocation by means of intra-subunit trna hybrid sites. Nature, (1167733)):7137716 2010. [Cited on page 8]
[43] Wen Zhang, Jack A Dunkle, and Jamie HD Cate. Structures of the ribosome in intermediate states of ratcheting. Science, 325(5543):00144-11 2009. [Cited page 8
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See
(c) Did you discuss any potential negative societal impacts of your work? [No] We don't 't believe to be relevant for this work
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A] Our paper does not contain theoretical results
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes] This will be
part of the supplemental material
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Most of the settings are clearly described in 5.1 and 5.2 The rest of this information will be available in the supplemental material.
(c) Did you report error bars (e.g., with respect to the random seed after running ex- periments multiple times)? [Yes] This will be provided as part of the supplemental material
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [Yes] This was a big portion of the discussion in 5.1 and 5.2
4. If you are using existing assets (c.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes] We cited all data sets and code
13

---

(b) Did you mention the license of the assets? [Yes] Please see 5.1 and 5.2
(c) Did you include any assets either in the supplemental material URL? [Yes] The data set used in will be made available as an anonymous URL for download
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A] We do not use personal data
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A] All of our data is either in publicly available databases 5.2 or was generated from such database 5.1
5. If you used crowdsoureing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A] We did not use crowdsourcing or conduct research with human subjects
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "bDmn2CPyhvi", "reviewer": "Reviewer_Vxds", "summary": "The paper proposes an approach for learning a volumetric Fourier reconstruction in a VAE framework for heterogeneous reconstruction of SPA cryoEM. The authors bring attention to the fact that the previous related work of coordinate-based reconstructions can be improved by learning Fourier representations on a regular cubic grid. The new approach is more efficient and reconstructs with better FSC than current SOTA but at one-seventh of the computational cost. ", "questions": "1. AFAIK, the paper or supplemental material makes no mention of the memory usage and the number of parameters. Moreover, you mention that the number of volumes for structural basis and volume size has major impact on memory footprint. Could you please expand on memory usage and at least approximately how it would scale?\n2. Did you encounter any cases where your method fails or which you expect to fail? For example, do you believe the method would still work with low SNR datasets or datasets with many classes?", "limitations": "The authors adequately addressed the limitations and potential negative societal impact of their work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Faster and more efficient method makes cryoEM more accessible, allowing to iterate faster.\n2. VAE framework, coupled with the previous point, allows interactive analysis of latent space (i.e. heterogeneouity), in contrast with discrete classification in the common cryoEM approaches.   \n3. The method is also better suited for cryoEM modality than previous coordinate-based networks.\n4. The method is validated on both simulated and experimental datasets and shows better FSC than the closest related work.\n\nWeaknesses:\n1. No comparisons with current defacto standard cryoEM software packages like RELION and cryoSPARC. For practical usage in cryoEM, it would be necessary to compare the proposed method with current standards. \n2. The method expects projection angle, shift and CTF parameters to be provided. The authors however mention that it should be doable to estimate them and leave it for the future work. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "RbMxbfyQVv", "reviewer": "Reviewer_4A7D", "summary": "This paper presents a method for heterogeneous cryo-EM reconstruction that fits 3D density volumes as a linear combination of voxel arrays where scale factors are inferred with an encoder network in a VAE framework. The model for volumes is represented in Fourier space and trained on projection images via the Fourier slice theorem (i.e. as 2D planar slices of 3D volumes) to reduce the memory requirement for backpropagation. The authors compare their method against cryoDRGN (a VAE-based model with a coordinate MLP decoder) and show slightly higher resolution reconstructions with their voxel-based generative model on a small synthetic dataset. The method yields qualitatively similar results to cryoDRGN on an experimental dataset of the Plasmodium falciparum 80S ribosome.", "questions": "* How does this method compare to 3DVA on the tested datasets? Can the authors discuss the advantages/disadvantages of this method compared to 3DVA?\n* How is the model initialized and how much does this impact the convergence/accuracy of the reconstruction?\n* On the synthetic dataset, the reconstruction accuracy of just a single volume is quantified. It would be good to verify that the accuracy is consistent across all the different states.\n* There are some architectural and training details that are missing, e.g. the size of the network for encoding images and decoding scale factors. Training details like the batch sizes are also missing, which has a huge impact on training times, especially when comparing methods by the number of epochs trained.\n", "limitations": "Yes", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* Novel formulation of cryo-EM heterogeneous reconstruction, which mixes and matches the generative model from 3DVA (linear combination of voxel arrays) and inference method from cryoDRGN (variational encoder MLP).\n* Rendering volumes is fast and interpretable, i.e. a linear combination of K 3D voxel arrays.\n* The paper is well written and easy to follow\n\nWeaknesses:\n\nWhile the method is interesting and could become a useful tool, the evaluation is flawed/limited.\n\n* The authors state that one of the main advantages is the computational cost for training (218 min for their method vs. 1638 min for cryoDRGN), but the training times are reported for a fixed number of training epochs without considering the convergence rate of the different methods. \n* Testing on synthetic data is useful for validating the method since there is ground truth, but the synthetic dataset used here seems a little too simple to be realistic (only 10k images of 10 discrete states at a very high signal-to-noise ratio of 2/3). \n* The authors mention a novel backpropagation method, but there is very little detail on what this method is. Is this just using the Fourier slice theorem to render 2D slices instead of 3D volumes (e.g. in cryoDRGN and in the voxel-based method described in Ullrich et all UAI 2019). Could the authors elaborate?\n\nA minor comment:\n* I don't follow some aspects of the motivation (Section 4.1). The authors point out that the cryo-EM signal is high-frequency in the Fourier domain and coordinate MLPs like cryoDRGN/NERF poorly capture high frequency signals unless a high-frequency positional encoding is provided in the input. But cryoDRGN already uses a positional encoding for this very reason, so what's the issue? Later in the results section, much of the results are instead focused on the training/render time of a voxel-based representation. The authors also discuss at length many advantages of a real space formulation, which further confuses the motivation of the work (a Fourier space representation). \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "-GZvGuaze_Z", "reviewer": "Reviewer_WYec", "summary": "This paper proposes to replace the coordinate-based Fourier space cryo-EM reconstruction approach in CryoDRGN with a volumetric Fourier space representation. The idea is that this volumetric representation allows the model to capture high-frequency variation typical of Fourier space data better than coordinate-based networks, which implicitly tend to captures smoother structures. While the volumetric representation increases memory requirements especially during backpropagation, the paper introduces a sparse back propagation strategy which takes advantage of the sparsity of the projection operation in the cryo-EM forward model. The paper demonstrates similar or better performance on reconstructing two datasets of cryo-EM images while requiring significantly shorter training and inference times, because forward passes do not require individual inference runs for each coordinate.", "questions": "- What aspects of the visual results lead to the claim that the reconstruction quality is improved with the proposed method (lines 256-257)?\n- How is the tri-linear interpolation able to handle the very high frequency changes in Fourier-space data which is ostensibly problematic for the coordinate-based approach?\n\n[My suggestion may change based on the responses to the above two questions, but as it currently stands, my suggestion is to remove claims about improved performance while focusing on the computational savings, which are more thoroughly justified. Or, if results using spatial priors capitalizing on the volumetric representation significantly improve the results, that could strengthen these claims -- but I don't think these experiments are necessary for publication of the paper as is.]", "limitations": "The authors describe a key limitation of their method in the last paragraph of Section 6 and address societal impacts in the \"Broader Impacts\" section. I believe that the proposed method has higher memory requirements than cryoDRGN, which I would encourage the authors to discuss more explicitly in the Discussion section as well.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n- The paper is quite clearly written and provides enough background for a reader outside of cryo-EM to easily understand the problem setup and categories of competing approaches, as well as some of their tradeoffs.\n- The reported savings in computational resources from the new method are significant, intuitively make sense (in that it is no longer necessary to run individual inference passes for each coordinate), and would be very useful for users of cryo-EM reconstruction tools.\n- The sparse backpropagation scheme is clever and makes use of the volumetric representation effectively.\n\n# Weaknesses\n- **It is not clear that the reconstruction quality of the proposed method is significantly higher than that of cryoDRGN.** This may be because I am not an expert in reading these kinds of images, and I am willing to change my mind on this. But as the manuscript currently stands, I can see differences between Figure 2B and 2D, but I’m not sure how to tell which is correct or why the proposed method reconstructions are better. Perhaps it would be helpful to point these out explicitly with arrows or more detailed text descriptions. Similarly, I am not sure what the salient takeaways are from the visual results in Figure 3. The FSC plot in 2E does show slightly higher FSC correlation than the cryoDRGN method, but this difference appears to decrease at high resolution, where I presume these methods would be most useful.\n- **It is not clear that the proposed method makes it easier for the networks to capture “local, high-frequency changes.”** While I agree that the coordinate-based approaches may suffer from this issue, it is not clear to me that the proposed approach is significantly better — in particular, I imagine the tri-linear interpolation (lines 204-205) might correspondingly struggle to handle very high frequency changes in the Fourier-space data. Perhaps the tri-linear interpolation does provide good results if the high-frequency changes are greater than ~1 pixel in width, but if so, I would expect the coordinate methods to also capture this reasonably well. This is further supported by my previous point that it is not clear that the proposed method results in significantly higher reconstruction quality, so at a minimum, some more discussion of this issue is needed.\n- (Minor) It may be helpful to report the real memory usage of both CryoDRGN and the proposed method so that readers can understand this tradeoff a bit more concretely. Relatedly, I see the parameter counts for cryoDRGN in Appendix C.2, but do not see these for the proposed method -- please provide these to compare.  This will give a rough sense for the expressive capacity of the models relative to each other (and, depending on the numbers, may make the point about capacity in lines 339-348 more concrete).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept