R ECONSTRUCTING CONTINUOUS DISTRIBUTIONS OF 3D PROTEIN STRUCTURE FROM CRYO -EM IMAGES
Tristan
Davis jhdavis Bonnie Berger bab
A BSTRACT
Cryo-electron microscopy (cryo-EM) is a powerful technique for determining the structure of proteins and other macromolecular complexes at near-atomic resolution.
In single particle cryo-EM, the central problem is to reconstruct the 3D structure of a macromolecule from 10 4 − 7 noisy and randomly oriented 2D projection images.
However, the imaged protein complexes may exhibit structural variability, which complicates reconstruction and is typically addressed using discrete clustering approaches that fail to capture the full range of protein dynamics.
Here, we introduce a novel method for cryo-EM reconstruction that extends naturally to modeling continuous generative factors of structural heterogeneity.
This method encodes structures in Fourier space using coordinate-based deep neural networks, and trains these networks from unlabeled 2D cryo-EM images by combining exact inference over image orientation with variational inference for structural heterogeneity.
We demonstrate that the proposed method, termed cryoDRGN, can perform ab initio reconstruction of 3D protein complexes from simulated and real 2D cryo-EM image data.
To our knowledge, cryoDRGN is the ﬁrst neural networkbased approach for cryo-EM reconstruction and the ﬁrst end-to-end method for directly reconstructing continuous ensembles of protein structures from cryo-EM images.
I NTRODUCTION
Cryo-electron microscopy (cryo-EM) is a Nobel Prize-winning technique capable of determining the structure of proteins and macromolecular complexes at near-atomic resolution.
In a single particle cryo-EM experiment, a puriﬁed solution of the target protein or biomolecular complex is frozen in a thin layer of vitreous ice and imaged at sub-nanometer resolution using an electron microscope.
After initial preprocessing and segmentation of the raw data, the dataset typically comprises 10 4 −noisy projection images.
Each image contains a separate instance of the molecule, recorded as the molecule’s electron density integrated along the imaging axis (Figure 1).
A major bottleneck in cryo-EM structure determination is the computational task of 3D reconstruction, where the goal is to solve the inverse problem of learning the structure, i.e. the 3D electron density volume, which gave rise to the projection images.
Unlike classic tomographic reconstruction (e.g. MRI), cryoEM reconstruction is complicated by the unknown orientation of each copy of the molecule in the ice.
Furthermore, cryo-EM reconstruction algorithms must handle challenges such as an extremely low signal to noise ratio (SNR), unknown in-plane translations, imperfect signal transfer due to microscope optics, and discretization of the measurements.
Despite these challenges, continuing advances in hardware and software have enabled structure determination at near-atomic resolution for rigid proteins (Kühlbrandt (2014);
Scheres (2012b); Renaud et al.
; Li et al.
Many proteins and other biomolecules are intrinsically ﬂexible and undergo large conformational changes to perform their function.
Since each cryo-EM image contains a unique instance of the molecule of interest, cryo-EM has the potential to resolve structural heterogeneity, which is experimentally infeasible with other structural biology techniques such as X-ray crystallography.
However, this heterogeneity poses a substantial challenge for reconstruction as each image is no longer of the same structure.
Traditional reconstruction algorithms address heterogeneity with discrete clustering approaches, however, protein conformations are continuous and may be poorly approximated with discrete clusters (Malhotra & Udgaonkar (2016);
Nakane et al.
Here, we introduce a neural network-based reconstruction algorithm that learns a continuous low-dimensional manifold over a protein’s conformational states from unlabeled 2D cryoEM images.
We present an end-to-end learning framework for a generative model over 3D volumes using an image encoder-volume decoder neural network architecture.
Extending spatialVAE, we formulate our decoder as a function of 3D Cartesian coordinates and unconstrained latent variables representing factors of image variation that we expect to result from protein structural heterogeneity (Bepler et al.
All inference is performed in Fourier space, which allows us to efﬁciently relate 2D projections to 3D volumes via the Fourier slice theorem.
By formulating our decoder as a function of Cartesian coordinates, we can explicitly model the imaging operation to disentangle the orientation of the molecule during imaging from intrinsic protein structural heterogeneity.
Our learning framework avoids errant local minima in image orientation by optimizing with exact inference over a discretization of SO (3) × R 2 using a branch and bound algorithm.
The unconstrained latent variables are trained in the standard variational autoencoder approach.
We present results on both real and simulated cryo-EM data.
B ACKGROUND AND N OTATION 2. I MAGE F ORMATION M ODEL
Cryo-EM aims to recover a structure of interest V : R 3 → R consisting of an electron density at each point in space based on a collection of noisy images X 1 , ..., X N produced by projecting (i.e. integrating) the volume in an unknown orientation along the imaging axis.
Formally, the generation of image X can be modeled as:
[Equation]
where V is the electron density ( volume ),
R ∈ SO (3) , the 3D rotation group, is an unknown orientation of the volume, and t = ( tx, ty, 0) is an unknown in-plane translation, corresponding to imperfect centering of the volume within the image.
The image signal is convolved with g , the point spread function for the microscope before being corrupted with frequency-dependent noise and registered on a discrete grid of size DxD, where D is the size of the image along one dimension.
The reconstruction problem is simpliﬁed by the observation that the Fourier transform of a 2D projection of V is a 2D slice through the origin of V in the Fourier domain, where the slice is perpendicular to the projection direction.
This correspondence is known as the Fourier slice theorem (Bracewell (1956)).
In the Fourier domain, the generative process for image ˆ X from volume ˆ V can thus be written:
[Equation]
where ˆ g = F g is the contrast transfer function (CTF) of the microscope, S ( t ) is a phase shift operator corresponding to image translation by t in real space, and
A ( R ) ˆ V =
ˆ V ( R T ( · , · , 0) T ) is a linear slice operator corresponding to rotation by R and linear projection along the z-axis in real space.
The frequency-dependent noise (cid:15) is typically modelled as independent, zero-centered Gaussian noise in
Fourier space.
Under this model, the probability of of observing an image ˆ X with pose φ
=
( R, t ) from volume ˆ V is thus:
[Equation]
[Equation]
(cid:12)(cid:12)(cid:12)
[Equation]
(cid:19)
[Equation]
where l is a two-component index over Fourier coefﬁcients for the image, σ l is the width of the
Gaussian noise expected at each frequency, and Z is a normalization constant.
2. T RADITIONAL CRYO -EM RECONSTRUCTION
To recover the desired structure, cryo-EM reconstruction methods must jointly solve for the unknown volume V and image poses φ i
=
( R i , t i ) .
Expectation maximization (Scheres (2012a)) and simpler variants of coordinate ascent are typically employed to ﬁnd a maximum a posteriori estimate of V marginalizing over the posterior distribution of φ i ’s, i.e.:
[Equation]
Intuitively, given V ( n ) , the estimate of the volume at iteration n , images are ﬁrst aligned with V ( n ) (Estep), then with the updated alignments, the images are backprojected to yield V ( n +1)
(M-step).
This iterative reﬁnement procedure is sensitive to the initial estimate of V as the optimization objective is highly nonconvex; stochastic gradient descent is commonly used for ab initio reconstruction 1 to provide an initial estimate V (0) (Punjani et al.
Given sample heterogeneity, the standard approach in the cryo-EM ﬁeld is to simultaneously reconstruct K independent volumes.
Termed multiclass reﬁnement , the image formation model is extended to assume images are generated from V 1 , ..., V K independent volumes, with inference now requiring marginalization over φ i ’s and class assignment probabilities π j ’s:
[Equation]
While this formulation is sufﬁciently descriptive when the structural heterogeneity consists of a small number of discrete conformations, it suffers when the heterogeneity is complex or when conformations lie along a continuum of states.
In practice, resolving such heterogeneity is handled through a hierarchical approach reﬁning subsets of the imaging dataset with manual choices for the number of classes and the initial models for reﬁnement.
Because the number and nature of the underlying structural states are unknown, multiclass reﬁnement is error-prone, and in general, the identiﬁcation and analysis of heterogeneity is an open problem in single particle cryo-EM.
M ETHODS
We propose a neural network-based reconstruction method, cryoDRGN (Deep Reconstructing Generative Networks), that can perform ab-initio unsupervised reconstruction of a continuous distribution over 3D volumes from unlabeled 2D images (Figure 2).
We formulate an image encoder-volume decoder architecture based on the variational autoencoder (VAE) (Kingma & Welling (2013)), where protein structural heterogeneity is modeled in the latent variable.
While a standard VAE assumes all sources of image heterogeneity are entangled in the latent variable, we propose an architecture that enables modelling the intrinsic heterogeneity of the volume separately from the extrinsic orientation of the volume during imaging.
Our end-to-end training framework explicitly models the forward image formation process to relate 2D views to 3D volumes and employs two separate strategies for inference: a variational approach for the unconstrained latent variables and a global search over
SO (3) × R 2 for the unknown pose of each image.
These elements are described in further detail below.
3. G ENERATIVE MODEL
We design a deep generative model to approximate a single function, ˆ V : R 3+ n → R , representing a n-dimensional manifold of 3D electron densities in the Fourier domain.
Speciﬁcally, the volume ˆ V is modelled as a probabilistic decoder p θ ( ˆ V | k, z ) , where θ are parameters of a multilayer perceptron (MLP).
Given Cartesian coordinates k ∈ R 3 and continuous latent variable z , the decoder outputs distribution parameters for a Gaussian distribution over ˆ V ( k, z ) , i.e. the electron density of volume ˆ V z at frequency k in Fourier space.
Unlike a standard deconvolutional decoder which produces a separate distribution for each voxel of a D 3 lattice given the latent variable, following spatial-VAE, we model a function over Cartesian coordinates (Bepler et al.
Here, these coordinates are explicitly treated as each pixel’s location in 3D Fourier space and thus enforce the topological constraints between 2D views in 3D via the Fourier slice theorem.
By the image formation model, each image corresponds to an oriented central slice of the 3D volume in the Fourier domain (Section 2).
During training, the 3D coordinates of an image’s pixels can be explicitly represented by the rotation of a DxD lattice initially on the x-y plane.
Under this model, the log probability of an image, ˆ X , represented as a vector of size DxD, given the current MLP, latent pose variables R ∈ SO (3) and t ∈ R 2 , and unconstrained latent variable, z , is:
[Equation]
where i indexes over the coordinates of a ﬁxed lattice c 0 .
Note that ˆ X (cid:48)
=
S ( − t )
ˆ X is the centered image, where S is the phase shift operator corresponding to image translation in real space.
We deﬁne c 0 as a vector of 3D coordinates of a ﬁxed lattice spanning [ − 0 .
5 , 0 . 5] 2 on the x-y plane to represent the unoriented coordinates of an image’s pixels.
Instead of directly supplying k , a ﬁxed positional encoding of k is supplied to the decoder, consisting of sine and cosine waves of varying frequency:
[Equation]
Without loss of generality, we assume a length scale by our deﬁnition of c 0 which restricts the support of the volume to a sphere of radius 0.5.
The wavelengths of the positional encoding thus follow a geometric series spanning the Fourier basis from wavelength 1 to the Nyquist limit ( 2 /D ) of the image data.
While this encoding empirically works well for noiseless data, we obtain better results with a slightly modiﬁed featurization for noisy datasets consisting of a geometric series which excludes the top 10 percentile of highest frequency components of the noiseless positional encoding.
3. I NFERENCE
We employ a standard VAE for approximate inference of the latent variable z , but use a global search to infer the pose φ
= ( R, t ) using a branch and bound algorithm.
Variational encoder:
As each cryo-EM image is a noisy projection of an instance of the volume at a random, unknown pose (viewing direction), the image encoder aims to learn a pose-invariant representation of the protein’s structural heterogeneity.
Following the standard VAE framework, the probabilistic encoder q ξ ( z | ˆ X ) is a MLP with variational parameters ξ and Gaussian output with diagonal covariance.
Given an input cryo-EM image ˆ X , represented as a DxD vector, the encoder
MLP outputs µ z | ˆ X and Σ z | ˆ X , statistics that parameterize an approximate posterior to the intractable true posterior p ( z | ˆ X ) .
The prior on z is a standard normal, N (0 , I ) .
Pose inference: We perform a global search over SO (3) × R 2 for the maximum-likelihood pose for each image given the current decoder MLP and a sampled value of z from the approximate posterior.
Two techniques are used to improve the efﬁciency of the search over poses: (1) discretizing the search space on a uniform grid and sub-dividing grid points after pruning candidate poses with [Figure] branch and bound (BNB), and (2) band pass limiting the objective to low frequency components and incrementally increasing the k-space limit at each iteration ( frequency marching ).
The pose inference procedure encodes the intuition that low-frequency components dominate pose estimation, and is fully described in Appendix A.
In summary, for a given image ˆ X i , the image encoder produces µ z | ˆ X i and Σ z | ˆ X i .
A sampled value of the latent z i ∼ N ( µ z | ˆ X i , Σ z
| ˆ X i ) is broadcast to all pixels.
Given z i and the current decoder, BNB orientational search identiﬁes the maximum likelihood rotation R i and translation t i for ˆ X i .
The decoder p θ then reconstructs the image pixel by pixel given the positional encoding of R Ti c 0 and z i .
The phase shift corresponding to t i and optionally the microscope CTF ˆ
g i is then applied on the reconstructed pixel intensities.
Following the standard VAE framework, the optimization objective is the variational lower bound of the model evidence:
[Equation]
where the expectation of the log likelihood is estimated with one Monte Carlo sample.
By comparing many 2D slices from the imaging dataset, the volume can be learned through feedback from these single views.
Furthermore, this learning process is denoising as overﬁtting to noise from a single image would lead to higher reconstruction error for other views.
We note that the distribution of 3D volumes models heterogeneity within a single imaging dataset, capturing structural variation for a particular protein or biomolecular complex, and that a separate network is trained per experimental dataset.
Unless otherwise speciﬁed, the encoder and decoder networks are both MLPs containing 10 hidden layers of dimension 128 with ReLU activations.
Further architecture and implementation details are given in Appendix A.
R ELATED W ORK
Homogeneous cryo-EM reconstruction: Cryo-EM reconstruction is typically accomplished in two stages: 1) generation of an initial low-resolution model followed by 2) iterative reﬁnement of the initial model with a coordinate ascent procedure alternating between projection matching and reﬁnement of the structure.
In practice, initial structures can be obtained experimentally (Leschziner & Nogales (2006)), inferred based on homology to complexes with known structure, or via ab-initio reconstruction with stochastic gradient descent (Punjani et al.
Once an initial model is generated, there are many tools for iterative reﬁnement of the model (Scheres (2012b); Punjani et al. (2017); Hohn et al. ; Lyumkis, Dmitry et al. ; Tang et al.
For example, Scheres (2012a) presents a Bayesian approach based on a probabilistic model of the image formation process and reﬁnes the structure via Expectation Maximization.
Frequency marching is used extensively in existing tools to speed up the search for the optimal pose for each image (Scheres (2012b); Barnett et al.
; Punjani et al.
CryoSPARC implements a branch and bound optimization scheme, where their bound is a probabilistic lower bound based on the noise characteristics from the image formation model (Punjani et al.
Ullrich et al.
(2019) propose a differentiable voxelbased representation for the volume and introduce a variational inference algorithm for homogeneous reconstruction with known poses.
Heterogeneous cryo-EM reconstruction: In the cryo-EM literature, standard approaches for addressing structural heterogeneity use mixture models of discrete, independent volumes, termed multiclass reﬁnement (Scheres (2010); Lyumkis, Dmitry et al.
These mixture models assume that the clusters are independent and homogeneous, and in practice require many rounds of expertguided hierarchical clustering from appropriate initial volumes and manual choices for number of clusters.
More recently, Nakane et al.
(2018) extend the image generative model to model the protein as a sum of rigid bodies (determined from a homogeneous reconstruction), thus imposing structural assumptions on the type of heterogeneity.
Frank & Ourmazd (2016) aim to build a continuous manifold of the images, however their approach requires pose supervision and ﬁnal structures are obtained by clustering the images along the manifold and reconstructing with traditional tools.
Recent theoretical work for continuous heterogeneous reconstruction includes expansion of discrete 3D volumes in a basis of Laplacian eigenvectors (Moscovich et al.
and a general framework for modelling hyper-volumes (Lederman et al.
e.g. as a tensor product of spatial and temporal basis functions (Lederman & Singer (2017)).
To our knowledge, our work is the ﬁrst to apply deep neural networks to cryo-EM reconstruction, and in doing so, is the ﬁrst that can learn a continually heterogeneous volume from real cryo-EM data.
Neural network 3D reconstruction in computer vision: There is a large body of work in computer vision on 3D object reconstruction from 2D viewpoints.
While these general approaches have elements in common with single particle cryo-EM reconstruction, the problem in the context of computer vision differs substantially in that 2D viewpoints are not projections and viewing directions are typically known.
For example, Yan et al.
(2016) propose a neural network that can predict a 3D volume from a single 2D viewpoint using only 2D image supervision.
Gadelha et al.
(2017) learn a generative model over 3D object shapes based on 2D images of the objects thereby disentangling variation in shape and pose.
Tulsiani et al.
(2018) also reconstruct and disentangle the shape and pose of 3D objects from 2D images by enforcing geometric consistency.
These works attempt to encode the viewpoint ‘projection’ operation 2 explicitly in the model in a manner similar to our use of the
Fourier slice theorem.
Coordinate-based neural networks in computer vision: Using spatial (i.e. pixel) coordinates as features to a convolutional decoder to improve generative modeling has been proposed many times, with recent work computing each image as a function of a ﬁxed coordinate lattice and latent variables (Watters et al.
However, directly modeling a function that maps spatial coordinates to values is less extensively explored.
In CocoNet, the authors present a deep neural network that maps 2D pixel coordinates to RBG color values.
CocoNet learns an image model for single images, using the capacity of the network to memorize the image, which can then be used for various tasks such as denoising and upsampling (Bricman & Ionescu (2018)).
Similarly, Spatial-VAE proposes a similar coordinate-based image model to enforce geometric consistency between rotated 2D images in order to learn latent image factors and disentangle positional information from image content (Bepler et al.
Our method extends many of these ideas from simpler 2D image modelling to enable 3D cryo-EM reconstruction in the Fourier domain.
R ESULTS
Here, we present both qualitative and quantitative results for 1) homogeneous cryo-EM reconstruction, validating that cryoDRGN reconstructed volumes match those from existing tools; 2) heterogeneous cryo-EM reconstruction with pose supervision, demonstrating automatic learning of the latent manifold that previously required many expert-guided rounds of multiclass reﬁnement; and 3) fully unsupervised reconstruction of continuous distributions of 3D protein structures, a capability not provided by any existing tool.
5. U NSUPERVISED HOMOGENEOUS RECONSTRUCTION
We ﬁrst evaluate cryoDRGN on homogeneous datasets, where existing tools are capable of reconstruction.
We create two synthetic datasets following the cryo-EM image formation model (image size D=128, 50k projections, with and without noise), and use one real dataset from EMPIARconsisting of 105,247 images of the 80S ribosome downsampled to image size D=90.
The encoder network is not used in homogeneous reconstruction.
As a baseline for comparison, we perform homogeneous ab-initio reconstruction followed by iterative reﬁnement in cryoSPARC (Punjani et al.
We compare against cryoSPARC as a representative of traditional state-of-the-art tools, which all implement variants of the same algorithm (Section 2).
Further dataset preprocessing and training details are given in Appendix B.
[Table]
We ﬁnd that cryoDRGN inferred poses and reconstructed volumes match those from state ofthe-art tools.
The similarity of the volumes to the ground truth can be quantiﬁed with the with the Fourier shell correlation (FSC) curve 3 .
Reconstructed volumes and quantitative comparison with the FSC curve is given in Figure S5.
Pose error to the ground truth image poses are given in Table 1.
For the real cryoEM dataset (no ground truth), the median pose difference between cryoDRGN and cryoSPARC reconstructions is 0.002 for rotations and 1.0 pixels for translations, and the resulting volumes are correlated above a FSC cutoff of 0.5 across all frequencies.
5. H ETEROGENEOUS RECONSTRUCTION WITH POSE SUPERVISION
Next, we evaluate cryoDRGN for heterogeneous cryo-EM reconstruction on EMPIAR-10076, a real dataset of the E. coli large ribosomal subunit (LSU) undergoing assembly (131,899 images, downsampled to D=90)
(Davis et al.
Here, poses are obtained through alignment to an existing structure of the LSU and treated as known during training.
In the original analysis of this dataset, multiple rounds of discrete multiclass reﬁnement with varying number of classes followed by human comparison of similar volumes were used to identify 4 major structural states of the
LSU.
We train cryoDRGN with a 1-D latent variable treating image pose as ﬁxed to skip BNB pose inference.
As a baseline, we reproduce the published structures originally obtained through multiclass reﬁnement with cryoSPARC.
Further baseline and training details are given in Appendix C.
We ﬁnd that CryoDRGN automatically identiﬁes all 4 major states of the LSU (Figure 3a).
Quantitative comparison with FSC curves 3 and additional volumes along the latent space are shown in Figure S7.
We compare the cryoDRGN latent encoding µ z | X for each image to the MAP cluster assignment in cryoSPARC and ﬁnd that the learned latent manifold aligns with cryoSPARC clusters (Figure 3b).
CryoDRGN identiﬁes subpopulations in some of the cryoSPARC clusters (e.g. Class D), which is partitioned by a subsequent round of cryoSPARC multiclass reﬁnement (Figure S8).
Published structures A and F correspond to impurities in the sample.
CryoDRGN correctly assigns images from these impurities to distinct clusters, but does not learn their correct structure since the poses inferred from aligning to the LSU template structure are [Table]
5. U NSUPERVISED HETEROGENEOUS RECONSTRUCTION
We test the ability of cryoDRGN to perform fully unsupervised heterogeneous reconstruction from datasets with different latent structure.
We generate four datasets (each 50k projections, D=64) from an atomic model of a protein complex, containing either a 1D continuous motion, 2D continuous motion, 1D continuous circular motion
[Table]
or a mixture of 10 discrete conformations (Figure S7).
We train cryoDRGN with a 1D latent variable for the linear 1D dataset and a 10D latent variable for the other 3 datasets.
As a baseline, we perform multiclass reconstruction in cryoSPARC sweeping K=2classes.
We compare against K=3, which had the best qualitative results.
We also propose a modiﬁcation to cryoDRGN in order to train on tilt series pairs datasets.
Tilt series pairs is a variant of cryo-EM in which, for each image
X i , a corresponding image X (cid:48)
i is acquired after tilting the imaging stage by a known angle.
This technique was originally employed to identify the chirality of molecules (Belnap et al.
(1997)), which is lost in the projection from 3D to 2D.
We propose using tilt series pairs to encourage invariance of q ξ with respect to pose transformations for a given ˆ V z (and incidentally to identify the chirality of ˆ V z ).
We make minor modiﬁcations to the architecture as described in Appendix D.
In Figure 4, we show that cryoDRGN reconstructed volumes for the circular 1D dataset qualitatively match the ground truth structures.
Note that while we only visualize 10 structures sampled along the latent space, the volume decoder can reconstruct the full continuum of states.
In contrast, cryoSPARC multiclass reconstruction, a discrete mixture model of independent structures, is only able to reconstruct 2 (originally unaligned) structures which resemble the ground truth.
Volumes contain blurring artifacts from clustering images from different conformations into the assumedhomogeneous clusters in the mixture model.
Results for the remaining datasets are given in Figures
S10-13.
We quantitatively measure performance on this task with an FSC resolution metric computed between the MAP volume for each image V z i | ˆ X i and the ground truth volume which generated each image, averaged across images in the dataset (Table S4).
We ﬁnd that cryoDRGN reconstruction accuracy is much higher than state-of-the-art discrete multiclass reconstruction in cryoSPARC, with further improvement achieved by training on tilt series pairs.
C ONCLUSIONS
We present a novel neural network-based reconstruction method for single particle cryo-EM that learns continuous variation in protein structure.
We applied cryoDRGN on a real dataset of highly heterogeneous ribosome assembly intermediates and demonstrate automatic partitioning of structural states.
In the presence of simulated continuous heterogeneity, we show that cryoDRGN learns a continuous representation of structure along the true reaction coordinate, effectively disentangling imaging orientation from intrinsic structural heterogeneity.
The techniques described here may also have broader applicability to image and volume generative modelling in other domains of computer vision and 3D shape reconstruction.
A CKNOWLEDGMENTS
We thank Ben Demeo, Ashwin Narayan, Adam Lerer, Roy Lederman and Kotaro Kelley for helpful discussions and feedback.
This work was funded by the National Science Foundation Graduate
Research Fellowship Program, NIH grant R01-GM081871, NIH grant R00-AG050749, and the MIT
J-Clinic for Machine Learning and Health
A A PPENDIX - M ETHODS A. B RANCH AND BOUND IMPLEMENTATION DETAILS
We perform a global search over SO (3) × R 2 for the maximum-likelihood pose for each image given the current decoder MLP.
Two techniques are used to improve the efﬁciency of the search over poses: (1) discretizing the search space on a uniform grid and sub-dividing grid points after pruning candidate poses with branch and bound , and (2) band pass limiting the objective to low frequency components and incrementally increasing the k-space limit at each iteration ( frequency marching ).
Our branch and bound algorithm for pose optimization is given in Algorithm 1.
Brieﬂy, we discretize
SO (3) uniformly using the Hopf ﬁbration Yershova et al.
(2010) at a predeﬁned base resolution of the grid and incrementally increase the grid resolution by sub-dividing grid points.
At each resolution of the grid, the set of candidate poses is pruned using a branch and bound (BNB) optimization scheme, which alternates between a computationally inexpensive lower bound on the objective function evaluated at all grid points and an upper bound consisting of the true objective evaluated on the best lower-bound candidate.
Grid points whose lower bound is higher than this value are excluded for subsequent iterations.
In our case, the loss is evaluated on low-frequency components of the image; speciﬁcally, Fourier components with | k | < k max is an effective lower bound, as it is both inexpensive to compute and captures most of the power (and thus the error).
This bound encodes the intuition that low-frequency components dominate pose estimation.
We concomitantly increase k max at each iteration of grid subdivision.
At each iteration, some poses are excluded by BNB, and the remaining poses are further discretized.
Although BNB is risk-free in the sense that the optimal pose at a given resolution will not be pruned, our application of it is not risk-free as a candidate pose with high loss at a given resolution doesn’t guarantee that its neighbor in the next iteration will not have a lower loss.
Irrespective, in practice, we ﬁnd that at a sufﬁciently ﬁne base resolution, we obtain good results on a tractable timescale (hours on a single GPU).
We reimplement the uniform multiresolution grids on SO (3) based on Yershova et al.
(2010), using the Healpix Gorski et al.
(2005) grid for the sphere and the Hopf ﬁbration to uniformly lift the grid to
SO (3) .
The base grid on SO (3) contains 576 orientations.
We use the ordinary grid for translations containing 7 2 points with an extent of 20 pixels for D=128 datasets.
We subdivide the grid 5 times for a ﬁnal resolution of 0.92 degrees for the orientation and 0.08 pixels for the translation.
For D= datasets, we use a translational grid with extent of 10 pixels.
Algorithm
1: procedure OPTPHI ( ˆ X, ˆ V z ) (cid:46) Find the optimal image pose given the current decoder 2: k min ← 12 , k max ← D/ 2 , N iter ←3:
Φ ← SO (3) × R 2 grid at base resolution 4: k ← k min 5: for iter = 1 . . .
N iter do 6: for φ i ∈
Φ do (cid:46) Compute lower bound at all grid points 7: lb ( φ i ) ← loss between ˆ X and SLICE ( ˆ V z , φ i ) at k < k 8: φ ∗ ← arg min( lb ) 9:
ub ← loss between ˆ X and SLICE ( ˆ V z , φ ∗ ) at k < k max (cid:46)
Compute upper bound 10: Φ new ← {} 11: for φ i ∈
Φ do (cid:46) Subdivide grid points below the upper bound 12: if lb ( φ i ) < ub then 13: Φ new ← Φ new ∪ SUBDIVIDE ( φ i ) 14: Φ ← Φ new 15: k ← k + ( k max − k min ) /
( N iter − 1) (cid:46) Increase frequency band limit 16: return φ ∗ 4 The difference in
A. T RAINING DETAILS
Given an imaging dataset, ˆ X 1 , ... ˆ X N , we summarize three training paradigms of cryoDRGN. 1) For homogeneous reconstruction, we only train the volume decoder p θ and perform BNB pose inference for the unknown φ i ’s for each image.
As an intermediate task, we can perform heterogeneous reconstruction training the image encoder q ξ and the volume decoder p θ with known φ i ’s to skip
BNB pose inference.
For fully unsupervised heterogeneous reconstruction, we jointly train q ξ and p θ to learn a continuous latent representation, performing BNB pose inference for the unknown pose of each image.
Unless otherwise speciﬁed, the encoder and decoder networks are both MLPs containing 10 hidden layers of dimension 128 with ReLU activations.
A fully connected architecture is used instead of a convolutional architecture because the images are not represented in real space.
Instead of representing both the real and imaginary components of each image, we use the closelyrelated Hartley space representation (Hartley (1942)).
The Hartley transform of real-valued functions is equivalent to the real minus imaginary component of the FT, and thus is real valued.
The Fourier slice theorem still holds and the error model is equivalent.
In this work, we simplify the image generation model to Gaussian white noise.
Therefore, for a given image, the negative log likelihood for a reconstructed slice from the decoder corresponds to the mean squared error between the phase-shifted image and the oriented slice from the volume decoder.
We leave the implementation of a colored noise model to future work.
We use the Adam optimizer (Kingma & Ba (2014)) with learning rate of 5e-4 for experiments involving noiseless, homogeneous datasets, and 1e-4 for all other experiments.
All models are implemented in Pytorch (Paszke et al.
B H OMOGENEOUS RECONSTRUCTION B. D ATASET PREPARATION
Simulated datasets: From a ground truth 3D volume, we simulated datasets following the cryo-EM image formation model by 1) rotating the 3D volume in real space by R , where R ∈ SO (3) is sampled uniformly, 2) projecting (integrating) the volume along the z-axis, 3) shifting the resulting 2D image by t , where t is sampled uniformly from [ − 10 , 10] 2 pixels, and 4) optionally adding noise to an SNR of 0.1, a typical value for cryo-EM data (Baxter et al.
Following convention in the cryo-EM ﬁeld, we deﬁne SNR as the ratio of the variance of the signal to the variance of the noise.
We deﬁne the noise-free signal images to be the entire DxD image.
50k projections were generated for each dataset with image size of D=128.
Real dataset: To generate the real cryo-EM dataset for homogeneous reconstruction, images from
EMPIAR-10028 (Wong et al.
(2014)) were downsampled by a factor of 4 by clipping in Fourier space.
The images were then ’phase ﬂipped’ in Fourier space by their contrast transfer function, a given real-valued function with range [-1,1] determined by the microscopy conditions, i.e. the Fourier components are negated where the CTF is negative.
B. T RAINING
For each dataset, we train the volume decoder (10 hidden layers of dimension 128) in minibatches of 10 images with random orientations for the ﬁrst epoch to learn a volume with roughly correct spatial extent, followed by 4 epochs with branch and bound (BNB) pose inference (30 min/epoch noiseless, 80 min/noisy datasets).
Since BNB pose inference is the bottleneck during training, we employ a multiscale training protocol, where after 4 epochs with BNB pose inference, the latent pose is ﬁxed, and we train a separate, larger volume decoder (10 hidden layers of dimension 500) forepochs with ﬁxed poses to "reﬁne" the structure to high resolution (20 min/epoch).
Training times are reported for 50k, D=128 image datasets trained on a Nvidia Titan V GPU.
[Figure]
C H ETEROGENEOUS RIBOSOME RECONSTRUCTION WITH POSE SUPERVISION
Dataset preparation: We used the dataset from EMPIAR-10076 which contains 131,899 images of the E. coli large ribosomal subunit (LSU) in various stages of assembly (Davis et al.
Images were downsampled to D=128 by clipping in Fourier space.
Poses were determined by aligning the images to a mature LSU structure obtained from a homogeneous reconstruction of the full resolution dataset in cryoSPARC, i.e. "a consensus reconstruction".
Baseline:
In the original analysis of this dataset, multiple rounds of multiclass reﬁnement in sweeps of varying number of classes followed by expert manual alignment and clustering of similar volumes were used to identify 6 classes, labeled A-F consisting of 4 major structural states of the LSU (classes
B-E) and 2 additional structures of the 70S and 30S ribosome, class A and F, respectively.
Since the published dataset did not contain the corresponding image cluster assignments, we perform multiclass reﬁnement in cryoSPARC using the published structures of the 6 major states, low pass ﬁltered to 25Å as initial models, to reproduce the results and obtain image cluster assignments.
Aside from class A and F (low population impurities in the sample), the remaining structures correlate well with the published volumes (Figure S6).
cryoDRGN training: We train cryoDRGN with a 1-D latent variable in minibatches of 10 images for 200 epochs, treating image pose as ﬁxed (11 min/epoch on a Nvidia Titan V GPU).
To simplify representation learning for q ξ , we center and phase ﬂip images before inputting to the encoder.
We encode and decode a circle of pixels with diameter D=128 instead of the full 128x128 image.
C. S UPPLEMENTARY RESULTS
D F ULLY UNSUPERVISED HETEROGENEOUS RECONSTRUCTION D. D ATASET PREPARATION
Linear 1D motion: We generated a dataset containing one continuous degree of freedom as follows:
From an atomic model of a protein complex, a single bond in the atomic model was rotated while keeping the remaining structure ﬁxed, and 50 atomic models were sampled along this reaction coordinate.
1000 projections with random rotations and in-plane translations were generated for each model, yielding a total of 50k images, approximating a uniform distribution along a continuous reaction coordinate.
Linear 2D motion: We extended the linear 1D motion dataset by introducing a second degree of freedom from rotating a bond in the atomic model that connected a different protein in the complex.
Similar to the 1D motion dataset, from a starting conﬁguration, the original bond was rotated +/-
N degrees, and 50 models were sampled along this reaction coordinate.
Then from the starting conformation, the second bond was rotated +/- 90 degrees, and 50 additional models were sampled along the second reaction coordination.
500 projections were generated from each model, yielding a total of 50k images.
Circular 1D motion: For this dataset, we rotated a bond a full 360 degrees and sample 100 models along this circular reaction coordinate.
500 projections were generated from each model, yielding a total of 50k images.
Discrete 10 class: For this dataset, we sampled 10 random conﬁgurations for the proteins in the complex.
5000 projection images were generated from each model, yielding a dataset containing a mixture of 10 discrete states.
For all four datasets, random rotations were generated uniformly from SO (3) , and translations were sampled uniformly from [ − 5 , 5] pixels.
The image size was D=64 with absolute spatial extent of 720Åand Nyquist limit of 22.5Å. Schematics of the simulated motions are given in Figure S9.
[Figure]
Tilt series pairs is a variant of cryo-EM in which, for each image
X i , a corresponding image X (cid:48)
i is acquired after tilting the imaging stage by a known angle.
This technique was originally employed to identify the chirality of molecules (Belnap et al.
(1997)), which is lost in the projection from 3D to 2D and therefore cannot be inferred from standard cryo-EM.
Inferential procedures such as expectation maximization converge to one handedness or the other depending on their initialization.
In multiclass reconstruction, different classes are not guaranteed to possess the same handedness even if there is a high relatedness between structures.
We remark on this experimental technique as we propose using tilt series pairs to encourage invariance of q ξ with respect to pose transformations for a given ˆ V z (and incidentally also to identify the chirality of ˆ V z ).
To train on tilt series pairs, the encoder is split into two MLPs, the ﬁrst learning an intermediate encoding of each image, and the second mapping the concatenation of the two encodings to the latent space.
We use an 8 layer MLP with output dimension 128 for the former and a 2 layer MLP with input dimension 256 for the latter.
All hidden layers have dimension 128.
For branch and bound, the combined loss over both images is evaluated for each grid point of SO (3) × R 2 .
To generate the image X tilt,i associated with X i , prior to rotating the volume by R i , we rotate the volume by a constant 45 degrees around the x-axis.
D. T RAINING
We trained cryoDRGN in minibatches of 5 images for 40 epochs without tilt series pairs andepochs with tilt series pairs.
We trained a 1-D latent variable for the linear 1D motion dataset and 10-D latent variables for the remaining datasets.
Random angles were used for the ﬁrst epoch of training to learn roughly the correct spatial extent of the volume and BNB pose inference was used for the remaining epochs.
The runtime was 120 min/epoch vs 2 min/epoch with and without BNB pose inference, respectively, on a Nvidia Titan V GPU.
D. S UPPLEMENTARY RESULTS
cryoSPARC
Dataset
K=
K= K=
K=
Linear 1D motion 5.11(3.82) 3.60(2.27) 7.40(4.16) 7.59(4.58)
Linear 2D motion 6.89(2.21) 6.90(3.77) 5.98(2.10) 6.76(4.47)
Circular 1D motion 5.16(2.70) 4.87(2.17) 7.50(3.32) 4.62(1.93)
cryoDRGN cryoDRGN+tilt
Dataset z-D= z-D= z-D= z-D= z-D= z-D=
Linear 1D motion 2.50(0.62) 2.34(0.12) – 2.35(0.36) 2.43(0.26) –
Linear 2D motion 7.16(4.69) 4.38(3.15) 4.44(2.50) 3.38(1.
[Table]
Circular 1D motion
[Table]
[Table]