Sparse Fourier Backpropagation in Cryo-EM Reconstruction
Abstract
Electron cryo-microscopy (cryo-EM) is a powerful method for investigating thestructures of protein molecules, with important implications for understandingthe molecular processes of life and drug development.
In this technique, manynoisy, two-dimensional projection images of protein molecules in unknown posesare combined into one or more three-dimensional reconstructions.
The presenceof multiple structural states in the data represents a major bottleneck in existingprocessing pipelines, often requiring expert user supervision.
Variational autoencoders (VAEs) have recently been proposed as an attractive means for learningthe data manifold of data sets with a large number of different states.
Thesemethods are based on a coordinate-based approach, similar to Neural Radiance
Fields (NeRF), to make volumetric reconstructions from 2D image data in Fourierspace.
Although NeRF is a powerful method for real-space reconstruction, many ofthe benefits of the method do not transfer to Fourier-space, e.g. inductive bias forspatial locality.
We present an approach where the VAE reconstruction is expressedon a volumetric grid, and demonstrate how this model can be trained efficientlythrough a novel backpropagation method that exploits the sparsity of the projectionoperation in Fourier-space.
We achieve improved results on a simulated data setand at least equivalent results on an experimental data set when compared to thecoordinate-based approach while also substantially lowering computational cost.
Our approach is computationally more efficient, especially in inference, enablinginteractive analysis of the latent space by the user.
Introduction
Single-particle electron cryo-microscopy (single-particle cryo-EM) has become the fastest growingmethod for structure determination of biological macromolecules in recent years and is becomingan increasingly important cornerstone of molecular biology research and drug discovery [1, 2, 3, 4].
Briefly, a thin layer of an aqueous solution containing multiple copies of the target molecule in randomorientations is flash frozen to cryogenic temperatures and imaged in a transmission electron microscope.
Subsequently, two-dimensional projection images of the Coulomb scattering potential of themolecules in the sample can be acquired into a micrograph.
Micrographs typically contain informationto sub-nanometer details of up to millions of copies of the target molecules that are called particles.
To limit damage to the radiation-sensitive particles, micrographs are acquired under low-dose conditions, which leads to low signal-to-noise ratios (of typically less than 1/10) which requires averagingover many images [5].
If a sufficient number of approximately identical particles are captured indistinct projection orientations (or poses), computational methods can be employed to reconstructa three-dimensional volume of the scattering potential of the particle.
However, the poses of theindividual particles are unknown and many particles adopt multiple conformational states, or assemblies of multiple molecules may vary in their composition.
The unknown poses and the potential assignments to multiple different structural states renders the problem ill-posed.
Combined withthe low signal-to-noise ratio in the images, this makes single-particle cryo-EM reconstruction achallenging task.
The reconstruction process starts with the extraction of a large number of image patches (on theorder of 10 4 − 10 7 ) 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 allparticles need to be determined, and particles need to be separated into structurally homogeneoussubsets.
Typically, this task is performed in two stages.
First, in a classification stage, structurallyhomogeneous subsets are identified and undesirable particles are removed.
Second, in a refinementstage, the selected subsets of particles are reconstructed to obtain a high-resolution structure.
Although generally desirable, automation and high throughput are essential in several areas ofapplication, e.g. fragment-based drug discovery [6] or screening for disease-relevant structures [7].
Fast computational processing pipelines, based on hardware acceleration [8, 9, 10], have provencritical to keep up with the high data rates provided by automated image acquisition software [11].
However, classification remains a bottleneck in existing pipelines, often requiring expert user inputand repeated trials.
Therefore, unsupervised algorithms that can replace existing classification stepswill be essential in further improving single-particle cryo-EM.
Moreover, because conformationalchanges are often central to the biological function of proteins, more powerful methods to deal withstructural heterogeneity have the potential to enable biological discoveries which are currently out ofreach of the existing methods.
Background and Definitions2.1 Cryo-EM Reconstruction
Let { y i ∈
C M 2 : 1 ≤ i ≤ N } be a set of images of (Fourier transformed) individual particles,where M is the image size and N is the number of particles.
Each y i is modelled as a projectionof a three-dimensional signal.
The projection operation H ( ϕ i ) is described by the pose of theparticle, ϕ i
=
( R i , t i ) ∈ SE(3) , where R i ∈
R 3 × 3 is a rotation or reflection matrix and t i ∈ R 3 isa translation.
Through the Fourier-slice theorem, the projection operation is realized by taking acentral slice from the three-dimensional Fourier transform of the signal, V ∈ C M 3 .
The projectedimage is then convolved with the point spread function of the optical system, which is described inthe Fourier domain by a multiplication with the contrast transfer function (CTF), C ( κ i ) , with knownparameters κ i .
Finally, experimental noise is modelled as additive and independent Gaussian in the
Fourier domain, with a scale that is frequency-dependent, ε ∼ N (0 , σ y ) .
Thereby, the forward modelbecomes:
[Equation]
If the data set is structurally homogeneous, all particles are projections of a common V , and theobjective is to estimate this V from all y i , with unknown ϕ
i and known κ i .
If the data set is structurallyheterogeneous, there is no longer a single V to solve for.
Rather, the solution is some space V andeach y
i corresponds to a V ∈ V .
In either case, we can represent the task at hand as the solution to amaximum-a-posteriori estimate, where we marginalize over the nuisance variable ϕ :
[Equation]
In the above, P ( V ) represents a log prior over the solution space.
A conventional approach to solvingthis problem is clustering the space V to a discrete set of K structural states, which means thatthe integral in (2) becomes a sum over a fixed K .
Popular software packages for single-particlecryo-EM, like RELION [12] and cryoSPARC [10], employ iterative algorithms, including expectationmaximization and gradient decent, to find a solution for the problem with K distinct structuralstates.
For computational reasons K is typically kept below 10 .
The expectation over ϕ is alsocomputationally expensive, which is why approximate approaches, including branch-and-boundmethods [10], or variations thereof [13], are regularly used.
2.2 Cryo-EM Generative Modelling
By parametrizing V differently, the cryo-EM reconstruction task can also be described with generativemodelling.
Using the autoencoder framework, let f θ be the encoder parameterized by θ , and g ω bethe decoder parameterized by ω .
Each particle image y i can be embedded in some latent encodingz i ← f θ ( y i ) .
Then, using the decoder, we reconstruct the original y i by minimizing a loss function Lthat will contain some priors and a log-likelihood term (cid:12)(cid:12)(cid:12)(cid:12)
y i − C
( κ i ) ◦ H ( ϕ i )
◦
g ω ( z ) (cid:12)(cid:12)(cid:12)(cid:12) 2 .
To reducethe computational cost of generative models, the expectation in (2) can be replaced with the highestprobability pose for each particle, ˆ ϕ
i , as determined using a homogeneous reconstruction approach.
The generative modelling approach enables replacing the fixed number of discrete classes, K , with amodel that can account for a larger variability in V .
This is especially relevant when dealing withstructural heterogeneity that cannot be easily separated in distinct states, like continuous motions.
Protein complexes often populate states that represent continuous transitions as part of their functionalcycle.
The continuous distribution of the learnt latent representation, expressed by g ω , provides anatural representation of the corresponding protein conformational landscape.
Assuming that theexperimental procedure does not perturb the distribution of structural states in the sample, the latentrepresentation could provide insights into the dynamics and perhaps the biological function of theprotein.
Moreover, although not strictly equivalent, the expansion of model capacity provided by the generativemodelling approach can be approximated in the discrete setting by increasing K to a large number.
It has been shown that robustness of canonical non-parametric classifiers are critically dependenton a large K [14].
Therefore, we expect generative models with larger model capacity to improverobustness over discrete heterogeneous reconstruction methods with low K .
Related work
The generative modelling approach described in the previous section can be classified based on thespace where the reconstruction takes place; either one decodes the volume in Fourier-space or inreal-space.
Both methods have their own benefits, as described below.
3.1 Fourier-space reconstruction
cryoDRGN [15] is currently the most popular approach to Fourier-space reconstruction using thegenerative modelling framework.
In cryoDRGN, a variational autoencoder (VAE) is used to encodeparticle image heterogeneity into a Gaussian parameterized latent space.
In particular, cryoDRGNencodes each particle image with µ i , σ i ← f θ ( y i ) .
To obtain the latent space encoding, they samplez i ∼ N ( µ i , σ i ) .
Next, they independently reconstruct all Fourier components that correspond to thespecific pose, ϕ i , of the particle image using a coordinate-based decoding function g ω
( z i , ξ ) , where ξdescribe the 3D coordinates.
For this approach to work, the hidden variables, including the projectionangle, are determined beforehand through global angular searches (see section 2.1).
Similar to neuralradiance fields (NeRF) [16], the decoder is typically a multilayer perceptron (MLP) that receives apositional embedding of the coordinates of each Fourier component of the image and its latent vector.
Another approach to Fourier-space reconstruction using the generative modelling framework is 3DVA[17].
This method does not use deep learning or gradient-based methods for learning its parameters,but starts from a reconstruction and formulates the generative process as a probabilistic principalcomponents analysis to obtain a set of orthogonal basis V ( k ) of 3D Fourier volumes representing theprincipal components that best fit the heterogeneity.
This method optimizes the log-likelihood withan expectation maximization approach.
Due to memory restrictions, cryoDRGN cannot reconstruct full volumes in a single batch, sincethe reconstruction of each coordinate in the volume requires a pass through the decoder.
Thisprohibits backpropagation with respect to the full volume because the intermediate activations for
O ( M 3 ) coordinates have to be stored.
This quickly becomes intractable.
Although, this issue can beameliorated with gradient checkpointing [18], it comes at the cost of extra training time.
This hindersthe application of real-space priors, e.g. solvent masks, mass preservation, and data-driven priors[19].
3DVA does not have this issue since full volume reconstructions are more computationally efficient than cryoDRGN, as volumes are weighted averages of the principal components and theoriginal reconstruction.
3.2 Real-space reconstruction
3DFlex [20] is a method implemented in the cryoSPARC software package that is based on real-spacereconstruction.
In this approach, a consensus real-space reconstruction is learnt and then deformedby a flow field that is parameterized by an MLP, given a latent representation for each image.
Insteadof viewing the problem as generating independent volumes, it deforms one consensus volume suchthat it maximizes the log-likelihood of the individual particle images.
Another real-space based method, e2gmm [21], also learns flow fields, but these are applied to a
Gaussian mixture model instead of a volumetric representation.
In this method, the encoding is donesimilar to cryoDRGN, where particle images are encoded in a Gaussian latent space.
Both real-space approaches enforce a mass-preserving inductive bias, which can improve reconstruction quality.
Although not yet explored, both approaches would in principle enable a computationallyefficient application of other real-space priors, like solvent masking or data-driven priors.
A disadvantage of the real-space methods is that the application of the CTF requires two Fouriertransforms (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 alongthe orientation of each particle image and then apply a 2D Fourier transform.
However, the projectionoperation is computationally expensive in real-space, especially when dealing with volumetricrepresentations, because the entire volume has to be accessed.
This contrasts with Fourier-space methods, where the CTF can be applied by multiplication and theprojection operation can be performed by accessing a small fraction of the Fourier components inthe volume, due to the Fourier-slice theorem.
To exploit the latter, in what follows we present anapproach where the decoder is not a function of the 3D coordinates, but rather generates the entirevolume in Fourier-space in a single pass, and present a novel approach for training such a model in acomputationally efficient manner.
Method
Figure method projection image
y i is the encoder f
θ which results in a set of scalars s ( k ) i that weigh the learnable and interpretable structural basis set V ( k ) .
The weighted sum of the projection H ( ϕ i ) is then CTF corrupted with C ( κ i ) .
Finally, the input y i and the output are compared in 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.
Here, we present an approach for learning a volumetric Fourier reconstruction given a latent representation 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 a sparse backpropagation method that limits evaluation andgradient 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 [22] with a custom CUDAkernel.
We have not preformed any runtime optimizations on these kernels.
4.1 Motivation
As previously noted, cryoDRGN uses a coordinate-based decoder, similar to NeRF, to parameterize thereconstructed volumes.
However, unlike the real-space case to which NeRF is applied, cryoDRGN’soutputs are defined in Fourier-space, where the spatial smoothness assumption is less valid.
Inparticular, this parameterization suffers from the spectral bias of neural networks [23, 24].
That is,the inability of neural networks to capture local, high-frequency changes.
This makes Fourier-spacerepresentations more challenging to learn with a coordinate-based approach.
This issue can beameliorated by incorporating a high-frequency positional embedding [25], but an alternative approachthat avoids positional embedding altogether might be better suited for this particular problem.
Apotentially more natural approach to learning discrete Fourier representations is on a regular cubicgrid, a so called voxel representation.
Similar approaches in real-space have been shown to improvemodel capacity and training efficiency [26].
In particular, we suggest to learn a set of volumes{ V ( k ) ∈ C M 3 } that are combined linearly into the final output.
While potentially better behaved, this method would ultimately result in the need to store B × Kvolumes for the backward pass during training, where B is the batch size.
This is computationallywasteful and has a large memory footprint, which is why we propose a novel training method.
Thefollowing sections outline how we exploit the sparse property of the cryo-EM forward model in
Fourier-space to more efficiently evaluate and train volumes V ( k ) with backpropagation.
4.2 Structural Basis
We define the decoder as
[Equation]
Here, V ( k ) ∈ C M 3 are trainable volume-sized arrays that we refer to as structural basis ; s ( k ) ω ∈ R arethe latent-dependent weighting factors; and V (0) is a trainable bias term.
K is strongly correlated withthe expressive capacity of the model.
Together with the volume size, K has a major impact on the totalmemory footprint.
The volume size is determined by the data image size, i.e. by the image physicalpixel size and the maximum resolution at which the model needs to be evaluated.
For all experimentsin this work we use K = 16 ; s ω is a small three layer MLP with residual connections [27].
This approach has similarities to 3DVA (see section 3.2), in that volumes are linearly combined togenerate the reconstruction results.
However, in our approach the structural bases are learnt togetherwith the representation in an end-to-end differentiable setup, including an encoder parameterized bya neural network.
4.3 Training
For the loss we use the log-posterior expression used in the RELION-4.0 implementation [28], whichis
[Equation]
Here, λ is a frequency-dependent regularization parameter; P is a suitable log-prior; and σ y is thefrequency-dependent noise amplitude of the data.
In this work we set P = 0 .
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.
Theextracted slice is then real-space shifted through phase-shifts in Fourier-space, and modulated withthe CTF, in order to calculate its log-likelihood loss.
Although the projection angle, shift, and CTFparameters are given, it is possible to jointly optimize them with the model parameters, since all theoperations 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,it is natural to 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 a sparserepresentation, which is less memory efficient for large batch sizes.
However, since each batchcontains particles from different orientations, when the batch size is sufficiently large, enough of thevolume would receive gradients, which enables the use of dense gradients.
This makes the memoryfootprint of the sparse projection operation invariant to batch size.
Because of the low signal-to-noise ratio in cryo-EM data, we can approximate σ y as the averagedspectral amplitude of the data, here denoted as ¯ σ y .
Let B = { ξ : ξ < | ξ | < ξ + τ } , where τ dependson the discretization used for the discrete Fourier transform.
Then, for a frequency ξ we get
[Equation]
In the above, the integral is evaluated over the frequency shell with radius between ξ and ξ + τ andthe sum is over all the images in the data set.
The value of ϵ is set to a small positive number fornumerical 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 ofmagnitude.
To improve convergence of the model parameters, we perform a spectral normalization ofthe Fourier transforms of the data samples, where the normalization factor is a function of Fourierfrequency.
In this work, we use ¯ σ y as the spectral normalization factor, which makes the effective σ yin (4) equal to one.
To reverse the effect this has on the spectral weighting of the MSE-loss, we set1 − λ proportional to ¯ σ 2 .
We note that it is the sparse access of the projection operation H ( · ) that enables an efficientevaluation and training of g ω .
For the sparse backpropagation method, we define a single fusedoperation that evaluates both g ω and H ( · ) in a single function call.
In the forward pass, this operationdisregards parameters in g ω that are omitted by H ( · ) .
Similarly, in the backward pass, gradientsbelonging to omitted parameters are not calculated.
Additionally, to minimize the memory footprint, we only store Fourier components in V ( k ) in asphere with a radius corresponding to the Nyquist frequency, as defined by the user.
This requiresexplicit coordinate arrays, which negatively impacts performance, but reduces the memory footprintby almost a factor of two, compared to storing Fourier components of the entire cube.
We evaluate our method on one simulated data set and an experimental cryo-EM data set, which wasdownloaded from the publicly available Electron Microscopy Public Image Archive (EMPIAR)
[30],with a CC0 licence [31].
All calculations were performed on an NVIDIA RTX 3090.
Results
5.1 Simulated data set
The simulated data set consists of projection images of 3D volumes generated from frames of amolecular dynamics simulation of the SARS-CoV-2 spike protein (DESRES-ANTON-10897850)[32], published under a CC BY 4.0 license [33].
Ten different conformational states were extractedout of the simulation and ten thousand projection images with random poses were generated for eachstate, 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 openand the closed states.
The image size is M = 96 , with a physical pixel of 2.7 Å. The simulatedimages were then distorted with additive Gaussian noise and CTF.
The average signal to noise ratiois 1/10 and the defocus range is between 0.2 and 2 µm.
The pose for each particle image was thenestimated by running a RELION consensus refinement, and kept fixed in the following steps.
We trained our VAE with 8 latent dimensions for 50 epochs.
After the end of training, we applieduniform manifold approximation and projection (UMAP) [35] to the distribution of each particleimage’s latent vectors.
Evaluating our decoder for points belonging to each cluster revealed thatseveral distinct conformations where separated for each of the three states (see figure 2A).
In totalseven separated clusters, were observed, amongst which, four belong to single conformations andthree 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 otherconformations (see figure 2A).
Running cryoDRGN on this data set with the default network architecture (cryoDRGN-small) and alarger one (cryoDRGN-large) results in worse reconstruction performance compared to our method,as can be seen in figure 2.
The reconstructions for cryoDRGN-small are shown in figure 2D and itslatent space is shown in figure 2C. Since this is simulated data, we can compare the results to theground truth. Comparison to the outline of the ground truth, reveals that the cryoDRGN results havemore missing structural features and overall lack the same level of details seen in the results of ourmethod, particularly in the flexible region.
Furthermore, we can calculate the Fourier shell correlation(FSC) between the generated and the ground truth structures.
The FSC is calculated here as thenormalized complex correlation between Fourier components of the solvent masked reconstructionand the known ground truth for each frequency band [36], averaged over 30 samples (three out ofeach conformation).
The results for this analysis are visualized in figure 2E. 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. Basedon the MSE-loss, neither of the methods converged in 50 epochs. The cryoDRGN-large run gaveworse results also compared to cryoDRGN-small due to overfitting to noise.
Additionally, increasingbatch size gave worse results for cryoDRGN, see mean FSC in supplementary table 1.
It took cryoDRGN-small 213 minutes to train for 50 epochs, while our method takes 40.3 minutes. Neither method converged within 50 epochs, see supplementary figure
6A. Evaluating the decodertakes 5 ms on average for each 3D structure, which is on the same order of magnitude as the time forthe inverse Fourier transform of 1 ms.
This can be compared to 250 ms for the cryoDRGN-smallarchitecture.
5.2 Ribosome data set
We also applied our method to an experimental cryo-EM data set of 80S ribosome particles fromthe Plasmodium falciparum parasite [38].
This data set has been used repeatedly as a test case fornew algorithms [15, 39, 40, 41].
We used 108,732 particle images, with M = 160 and a pixel sizeof 3 .
Å. We used 10 latent dimensions for this data set.
To evaluate the ability to handle defectiveparticle images, we also included 78,442 images that where classified by RELION as not containingribosomes, or containing partially (dis-)assembled ribosomes, along with images classified as intactribosomes.
This brings N to 187,174.
After a few epochs, the defective and empty particles were separated out as the top principalcomponent and the separation is clearly preserved in a UMAP of the latent space after 50 epochs(figure 3C).
Additionally, we notice separation of different classes of particles in each of the two majorclasses.
In the cluster of the intact ribosomes we observed several subtle continuous conformationaltransitions, including the "head swivel" [42] (not shown in figure) and the rotation of the smallsubunit, which is part of the "ratcheting" cycles [43] (figure 3B).
Although the figure shows the twoendpoints of the rotation, a continuous transition between the two states can be observed along themarked path in the latent representation.
In the other cluster, containing the defective particles, weobserved varying classes of particles.
Images of partially assembled ribosomes are separated fromimages lacking any protein features (figure 3D).
Next, we trained cryoDRGN-large on this data set and achieved similar results to our method, ascan be seen in figure 3E-H. Unlike in 5.1, we do not have access to the ground truth volumes inorder to compute an FSC.
However, qualitatively, we are able to generate similar clustering andreconstruction performance compared to cryoDRGN, at lower computational cost.
It should also benoted that training cryoDRGN-small on this data set gives noticeably worse results, as shown in theappendix.
Training took 210 minutes using our method, compared to 1638 minutes for cryoDRGN.
Neither method converged within 50 epochs, see supplementary figure 6B. Moreover, evaluating thedecoder takes 10 ms on average for each 3D structure, compared to 480 ms for cryoDRGN.
We kept the batch size of cryoDRGN to the recommended value and what gave the best results forthe simulated data set (see section 5.1).
Discussion
The potential improvement in robustness make VAE approaches a candidate for replacing the discreteclassification 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 preventingthis 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 tocryoDRGN with less than one-fifth the computational cost.
We show that our method yieldsreconstructions with higher correlation to the known ground truth.
For an experimental test case,our method converges onto a solution that is at least as useful as the one provided by cryoDRGN,with one-eighth of the computational cost and a smaller memory requirement.
From the user’s perspective, our approach also offers a practical advantage in allowing a fasterinteractive analysis of latent space.
Calculation of a volume for a given point in latent space takes on the order of half a second in cryoDRGN.
Our decoder reconstructs entire volumes in one pass, whichtakes 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 interestin latent space.
Intuitive graphical tools enable inexperienced users to explore the heterogeneity intheir data, and easily identify particle subsets to be used for subsequent refinements.
This holds thepotential 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 overallcomputational cost and the fast full-volume generation (see section 3.1).
This could enable reconstruction of data sets where the initial angular assignment based on homogeneous reconstruction fails.
Additionally, this could accelerate the replacement of the discrete 3D classification where globalangular searches are integrated in all popular implementations.
Because of the low signal-to-noise ratios in the experimental images, cryo-EM reconstruction isprone to overfitting, and explicit regularization of the output maps has played a major role in themost popular software packages [10, 12].
Operations like (back-)projection and the applicationof the CTF are performed conveniently in Fourier-space.
Most software packages thus performreconstruction in Fourier-space, where they also regularize the power of the Fourier components ofthe reconstruction.
Our approach, which also performs reconstruction in Fourier-space, is well suitedto 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 alsoapply real-space priors, like solvent masks or data driven priors [19].
Because these priors aredifferentiable, one could backpropagate through them to improve this method’s structural basesand network parameters.
In this way, one could directly use the data-driven prior to improve thereconstruction in a computationally efficient manner.
Improved algorithmic robustness through theuse of priors on the reconstructed set of volumes has the ability to increase the radius of convergenceof 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 betterinterpretability of what the model is learning.
In addition to facilitating troubleshooting, the structuralbasis can potentially be utilized for two purposes.
Firstly, they could be used to impose 3D priors onthe reconstruction in a more direct manner than backpropagation.
Real-space priors that are invariantto structural heterogeneity, e.g. solvent masks or structural symmetries, could be imposed to themdirectly rather than as part of the gradient-based training loop.
Secondly, real-space features of thestructural 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.
If theheterogeneity of the data set cannot be embedded into a linear subspace of K dimensions, our methodwould not be able to correctly reconstruct the volume manifold.
This limit is theoretically not presentin cryoDRGN, because they parameterize their reconstruction with a neural network.
However, inpractice, the capacity of the cryoDRGN network is also limited by computational cost, and we haveso far not observed issues with the expressive capacity of our model.
It should be noted that, dueto the noise present in the cryo-EM reconstruction problem, one should expect that the dimensionof the possible solution space ˆ V is much smaller than N .
Therefore, the limit to the reconstructionof heterogeneity due to noise would dominate the limit due to the representation capacity of thestructural basis set.
Broader Impact
The cost of data processing can be a major limiting factor for research groups with limited accessto computational resources.
By improving robustness and computational efficiency, we expand theaccessibility of the method to include a more diverse research community, which can in turn have anaccelerating impact on the overall research quality of the field.
Moreover, the potential for improvedautomation has implications for high-throughput structure determination.
Along with improvementsin microscope data acquisition, rapid processing can enable faster iterations in e.g. protein design orfragment-based drug discovery.
Thereby, our method has the potential to impact medical researchand the development of new therapeutics.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See
(c) Did you discuss any potential negative societal impacts of your work?
[No] We don’t
believe it to be relevant for this work
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
[Table]
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
Our paper
[Table]
does not contain theoretical results
[Table]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] This will bepart of the supplemental material [Table]
(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
[Table]
of this information
will be available in the supplemental material.
[Table] (c) Did you report error bars (e.g., with respect to the random seed after running ex- [Table]
periments multiple times)?
[Yes] This will be provided as part of the supplemental [Table]
material
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] This was a big portion of the
discussion in 5.1 and 5.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] We cited all data sets and code [Table]
(b) Did you mention the license of the assets?
[Yes] Please see 5.1 and 5.
[Table] (c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
[Table]
The data set used in 5.1 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 a database 5.
If you used crowdsourcing or 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 to 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]