# Exploring generative atomic models in
cryo-EM reconstruction

Ellen D. Zhong
MIT
<EMAIL_ADDRESS>
&Adam Lerer
Facebook AI Research
<EMAIL_ADDRESS>
&Joseph H. Davis
MIT
<EMAIL_ADDRESS>
&Bonnie Berger
MIT
<EMAIL_ADDRESS>

###### Abstract

Cryo-EM reconstruction algorithms seek to determine a molecule’s 3D density
map from a series of noisy, unlabeled 2D projection images captured with an
electron microscope. Although reconstruction algorithms typically model the 3D
volume as a generic function parameterized as a voxel array or neural network,
the underlying atomic structure of the protein of interest places well-defined
physical constraints on the reconstructed structure. In this work, we exploit
prior information provided by an atomic model to reconstruct distributions of
3D structures from a cryo-EM dataset. We propose Cryofold, a generative model
for a continuous distribution of 3D volumes based on a coarse-grained model of
the protein’s atomic structure, with radial basis functions used to model atom
locations and their physics-based constraints. Although the reconstruction
objective is highly non-convex when formulated in terms of atomic coordinates
(similar to the protein folding problem), we show that gradient descent-based
methods can reconstruct a continuous distribution of atomic structures when
initialized from a structure within the underlying distribution. This approach
is a promising direction for integrating biophysical simulation, learned
neural models, and experimental data for 3D protein structure determination.

##  1 Introduction

Single particle cryo-electron microscopy (cryo-EM) is a powerful experimental
technique for structure determination of proteins and macromolecular complexes
at near-atomic resolution [1]. In this technique, an electron microscope is
used to image a purified sample of the molecule of interest suspended in
vitreous ice. Initial processing of the resulting micrograph produces a
dataset of $10^{4-7}$ 2D projection images, where each image captures a unique
molecule suspended in a random, unknown orientation. Reconstruction algorithms
are then used to computationally infer the 3D protein structure from the
dataset. Unlike structure determination through x-ray crystallography, cryo-EM
is able to determine the structure of molecules that are not amenable to
crystallization, including those that adopt an ensemble of heterogeneous
conformational states. Thus, cryo-EM is in principle capable of modeling the
entire ensemble of conformational states in a sample and probing the dynamic
conformational landscape of proteins. To date, this promise has been limited
by a lack of computational techniques for dealing with heterogeneous cryo-EM
data, but promising new tools have recently been developed to reconstruct
heterogeneous structures [2, 3, 4].

Computational processing of cryo-EM datasets consists of several distinct
stages (Figure 1). First, the noisy micrograph image is preprocessed and
segmented, where bounding boxes containing individual molecules (i.e.
particles) are identified and extracted (i.e. particle picking) [5]. Next, a
3D cryo-EM density volume (or distribution of volumes) that maps the
molecule’s electron scattering potential is reconstructed from the 2D
projections [2, 6, 7, 8, 9]. Finally, an atomic model is built into the
resulting 3D volume either de novo or fit starting from a related model,
typically with the help of automated tools [10, 11, 12, 13, 14, 15]. The
reader is referred to [16] for further details on the single particle cryo-EM
image processing pipeline.

Figure 1: Structure determination via cryo-EM. Schematic of cryo-EM
reconstruction (A $\rightarrow$ B) and atomic model fitting (B $\rightarrow$
C). In this work, we investigate the possibility of fitting the atomic protein
structure directly during the reconstruction process. Dataset from Walls et
al. [17].

In this work, we investigate the possibility of fitting the atomic protein
structure directly during the reconstruction process. We have two motivations
for this. First, atomic fitting is a labor-intensive step in cryo-EM post-
processing, and in particular it is not clear how to perform atomic fitting
for models of distributions of protein conformations learned during
reconstruction with new tools such as cryoDRGN [2]. Second, modeling the
atomic structure of the 3D volume provides a strong prior over structures. In
fact, in many cases the protein sequence and an approximate reference
structure are known beforehand, strongly constraining the space of feasible 3D
volumes. These structural priors are especially important for models of
heterogeneous distributions of molecules, because they constrain the
conformational dynamics to those that approximate realistic protein motions.
Without such priors, it is common to observe artifacts in the motion of the
protein, e.g. mass appearing and disappearing between two distinct
conformations with rarely-sampled transition states.

To this end, we propose a reconstruction process based on a coarse-grained
atomic model for the cryo-EM volume. The model fits parameters including
atomic coordinates, to maximize the likelihood of the dataset under a
generative model that maps atomic structures to cryo-EM images. When
initialized appropriately, this approach is able to learn both homogeneous
structures and heterogeneous ensembles from synthetic cryo-EM images.

##  2 Background

A cryo-EM experiment produces a dataset of $10^{4-7}$ noisy 2D projection
images, each containing a unique molecule captured in a random, unknown
orientation. The goal of cryo-EM reconstruction is to infer the 3D density
volume $V:\mathbb{R}^{3}\rightarrow\mathbb{R}$ that gave rise to the imaging
dataset $X_{1},...,X_{N}$. As cryo-EM images are integral projections of the
molecule in this imaging modality, 2D images can be related to the 3D volume
by the Fourier slice theorem [18], which states that the Fourier transform of
a 2D projection is a central slice from the 3D Fourier transform of the
volume. Traditional methods approximate the volume as a voxel array
$\hat{V}(\bf k)$ in Fourier space [16].

To recover the desired structure, cryo-EM reconstruction methods must jointly
solve for the unknown volume $\hat{V}$ and image poses
$\phi_{i}=(R_{i},t_{i})$, where $R_{i}\in SO(3)$ and
$t_{i}\in{\mathbb{R}}^{2}$ are the 3D orientation of the molecule and in-plane
image translation, respectively. Expectation maximization and simpler variants
of coordinate ascent are typically employed to find a maximum a posteriori
estimate of $\hat{V}$ marginalizing over the posterior distribution of
$\phi_{i}$’s [19].

A unique advantage of cryo-EM is its ability to image heterogeneous molecules.
Heterogeneous reconstruction algorithms aim to reconstruct a distribution of
structures from the dataset. A standard approach involves extending the
generative model to assume that images are generated from a mixture model of
$K$ volumes $V_{1},...,V_{K}$ [20, 21]. More recently, cryoDRGN proposed a
neural model to reconstruct heterogeneous ensembles of particles from cryo-EM
data [22]. CryoDRGN represents a continuous $n$-dimensional distribution over
volumes as a function $\hat{V}:{\mathbb{R}}^{3+n}\rightarrow{\mathbb{R}}$
approximated by a multi-layer perceptron (MLP) with positional encoding of
Cartesian coordinate inputs [22]. To simplify reconstruction, cryoDRGN and
other advanced reconstruction methods [2, 3, 23, 4] find it sufficient to use
poses $\phi$ computed using a traditional reconstruction method, and focus on
the volume reconstruction.

##  3 Related Work

A large body of work has investigated continuous deformations of protein
structures produced from normal mode analysis of atomic, coarse-grained, or
pseudo-atomic models [24, 25, 26, 27, 28]. The top N normal modes of a system
where (pseudo-)atomic bonds are approximated by harmonic springs can summarize
a molecule’s flexing motions, which can then be used to model cryo-EM data
during reconstruction [25, 26, 27, 28]. However, it is unclear how accurate
the hypothetical motions that are generated from the underlying harmonic
spring approximation are. In BioEM and other approaches, likelihood-based
analysis of cryo-EM images or reconstructed maps based on an ensemble of
atomic models (e.g. generated from molecular dynamics) has been used to
investigate conformational heterogeneity [29, 30, 31].

Recently, deep learning has been used to incorporate prior beliefs from an
existing dataset of atomic structures into the cryo-EM reconstruction process.
Kimanius et al. train a convolutional denoising model on PDB structures that
is integrated into the cryo-EM reconstruction process in the form of updates
to regularization parameters during iterative refinement [32]. They show
higher resolution reconstructions of synthetic datasets but observe artifacts
from the learned general model in some cases. DeepEMhancer trains a network to
perform volume post-processing (e.g. high frequency sharpening and
solvent/background masking) from pairs of raw and postprocessed experimental
density maps, where the postprocessed maps have been refined by the fitted
atomic structure [33].

In concurrent work, Rosenbaum et al. propose a VAE with a similar RBF-based
generative model over atomic coordinates [34]. While we model proteins at the
residue level, Rosenbaum et al. model them at the atomic level [34] and also
jointly infer image pose. They evaluate this model on synthetic data from a
molecular dynamics simulation.

##  4 Method

###  4.1 Cryofold volume model

The central contribution of this work is a cryo-EM density model that is
parameterized in terms of coordinates for a coarse-grained atomic model given
a known atomic reference structure (Figure 2A). In this coarse-grained model,
each amino acid is represented by two Gaussian radial basis functions (RBFs),
one representing the backbone and the second representing the sidechain. Each
RBF is parameterized by $(\mu_{i},\sigma_{i},a_{i})$, where $\mu_{i}$ is the
position of the $i$th RBF; $a_{i}$ is the amplitude of the $i$th RBF; and
$\sigma_{i}$ is the width of the $i$th RBF. We tie $a_{i}=a_{0}Z_{i}$ where
$a_{0}$ is a global learned amplitude constant, and $Z_{i}$ is the total
number of electrons in the fragment represented by RBF $i$. Furthermore, we
tie all $\sigma_{i}$ to the same value $\sigma$. Thus, the full RBF model has
$3K+2$ parameters, where $K$ is the total number of amino acids in the protein
complex. The cryo-EM density can be computed as a function of these parameters
as:

$V(\mathbf{r})=\sum_{i}(2\pi\sigma_{i})^{-3/2}a_{i}\exp\left({\frac{-||\mathbf{r}-\mathbf{\mu_{i}}||^{2}}{2\sigma_{i}^{2}}}\right)$
(1)

As described in Section 2, reconstruction is performed in the Fourier domain.
This makes the choice of Gaussian RBFs convenient, as $V$ can be computed
efficiently in Fourier space111Incidentally, projections of Gaussian kernels
can also be computed analytically in real space, obviating the need for the
Fourier slice theorem altogether. This real space formulation allows for
algorithms that scale better with the number of RBFs since the RBFs are
localized in real space, allowing for spatial decomposition via gridding, KD-
trees, etc.:

$\hat{V}(\mathbf{k})=\sum_{i}{a_{i}e^{-2\pi i\mathbf{\mu_{i}\cdot
k}}\exp\left(\frac{-\pi^{2}\mathbf{k}^{2}\sigma_{i}^{2}}{2}\right)}$ (2)

To impose physical constraints on the RBF model, we add a set of fixed
harmonic bond terms $\mathcal{B}$ between consecutive backbone RBF centers.
Each bond term $(i,j,k,l)\in\mathcal{B}$ is specified by a pair of RBF indices
$i,j$ and a bond strength $k$, which we set to 0.1. The bond length $l$ is set
to 3.8 Å, the distance between protein $C_{\alpha}$ backbone carbons.

We constrain the side-chain RBFs to be located close to their backbone RBF
using one-sided harmonic restraints $\mathcal{C}$. These are similar to the
bond terms, but only induce a loss when the distance between RBF centers
exceeds the max length $l$. We set $l$ to the maximum distance between a
backbone C-alpha carbon and its side chain center of mass observed in the
reference structure.

For homogeneous reconstruction, we fit $\\{\mu_{i}\\},\sigma,a_{0}$ directly
using stochastic gradient descent (SGD). The overall loss function, given a
set of $N$ images $\mathbf{X}$ with poses $\mathbf{\phi}$, bond terms
$\mathcal{B}$ and side-chain constraints $\mathcal{C}$ is:

$\begin{split}\mathcal{L}(\mu,\sigma,a_{0}|\mathbf{X})=\frac{1}{N}\sum_{(x,\phi)\in
X}||\hat{X}(\phi|\mu,\sigma,a_{0})-X||^{2}&+\sum_{(i,j,k,l)\in\mathcal{B}}k(||\mathbf{\mu}_{i}-\mathbf{\mu}_{j}||-l)^{2}\\\
&+\sum_{(i,j,k,l)\in\mathcal{C}}k~{}\max\left(||\mathbf{\mu}_{i}-\mathbf{\mu}_{j}||-l,0\right)^{2}\end{split}$
(3)

###  4.2 Heterogeneous reconstruction

For heterogeneous reconstruction, we learn a continuous latent variable model
of conformational heterogeneity expressed through motion of RBF centers
(Figure 2B). Unlike standard heterogeneous cryo-EM reconstruction algorithms
that use an unconstrained volume representation (e.g. voxel arrays or
positionally encoded MLPs), the RBF model constrains the effect of the latent
degrees of freedom to motion of the underlying atomic structure.

An image encoder $E$ with parameters $\theta_{E}$ predicts a latent $z\sim
E(\hat{X}|\theta_{E})$; A decoder network $D$ with parameters $\theta_{D}$
predicts a $z$-dependent translation of the RBF centers
$\mu_{het}(z)=\mu+D(z|\theta_{D})$. We optimize
$\mu,\sigma,a_{0},\theta_{D},\theta_{E}$ together end-to-end with SGD.

Figure 2: Cryofold model and architecture. A) The physics-inspired cryo-EM
density model consists of set of Gaussian RBFs, two for each amino acid in the
reference structure. One RBF represents the backbone (blue) and one represents
the sidechain (red). Backbone RBFs are connected with harmonic bond terms
$\mathcal{B}$, and sidechain RBFs are connected to their backbone with max-
distance harmonic constraints $\mathcal{C}$. B) Architecture for heterogeneous
reconstruction. We use a VAE to learn $z$-dependent offsets of RBF centers
$\mu_{i}$, given an unlabelled imaging dataset of image, pose pairs
$(\hat{X_{i}},\phi_{i})$.

###  4.3 Local minima

A major shortcoming of this atomic reconstruction approach is the existence of
many local minima of the loss function that do not approximate the true atomic
coordinates. This is in contrast to voxel-based models, where SGD converges to
the global minimum of the convex loss given the poses, and in contrast to
large-scale generic neural models (e.g. [22, 2]) which typically do not suffer
from problematic local minima in practice [35]. We address the local minima
problem primarily by initializing RBF centers $\mu_{i}$ from a reference
structure that is a close approximation of the imaged structure (homogeneous)
or some point in the distribution of structures (heterogeneous). Such
reference structures are often available when studying a variant of a known
structure, or heterogeneous protein dynamics.

##  5 Results

Here, we present results for reconstructing coarse-grained atomic structures
from synthetic cryo-EM image data with Cryofold. We first explore the effect
of reference structure initialization in homogeneous reconstruction, when
initialized from either an approximate reference structure (fitting), the
exact structure (cheating), or from randomly initialized locations (folding).
We then turn to heterogeneous cryo-EM datasets, where we evaluate the ability
of Cryofold to reconstruct continuous distributions of structures and their
atomic coordinates. Finally, we assess choices in the design of the RBF model
by ablating key components in the homogeneous setting.

Datasets. We generate homogeneous and heterogeneous cryo-EM datasets using a
141-residue atomic model (PDB 5NI1) as the ground truth structure. To generate
homogeneous data, we simulate 50,000 noisy projection images based on the
cryo-EM image formation model. To model heterogeneity, we introduce a bond
rotation in the backbone of 5ni1 to create a continuous 1D motion, and
generate cryo-EM images sampling along the ground truth reaction coordinate.
Further details on dataset generation are given in the Appendix.

Architecture and training. For all experiments, we train for 10 epochs using
the Adam optimizer in minibatches of 8 images and a learning rate of 1e-4. We
initialize $\sigma$ and $a_{0}$ to $3.71\AA$ and $0.1$, and perform one epoch
of ‘warm-up’ to refine these global parameters after which we reset the atomic
coordinates to their initial values. For homogeneous reconstruction, we
directly optimize all RBF parameters. For heterogeneous reconstruction, we use
a VAE to predict the $z$-dependent offsets of the RBF centers from their
reference values. Both the encoder and decoder are 3-layer MLPs of width 256
and residual connections, and a 1-dimensional latent variable. We use ground-
truth poses $\phi$ for training. In real applications, the poses would be
inferred from traditional cryo-EM tools [6, 7, 8]. The model is implemented in
PyTorch [36].

###  5.1 Homogeneous reconstruction

To explore the effect of reference structure initialization, we compare the
reconstruction accuracy of Cryofold when initialized from the ground truth
coordinates ("Exact 5NI1") and from three alternate starting configurations:
1) an approximate reference structure generated by evolving the system under a
molecular dynamics simulation ("Approximate 5NI1"), 2) the ground truth
structure with 6 Å uniform noise added to the ground truth coordinates of each
$C_{\alpha}$ ("5NI1 + Uniform[6 Å]"), and 3) a random initialization of each
$C_{\alpha}$ RBF in the $64\AA^{3}$ region in the center of the box
("Random"). Side-chain RBFs are initialized to their corresponding
$C_{\alpha}$ RBF centers.

We report the root-mean-square-error (RMSE) of model backbone coordinates to
the $C_{\alpha}$ of the true structure, the percent of $C_{\alpha}$ backbone
atoms predicted within 3 Å of the true structure, and the normalized mean-
square-error (NMSE) of the reconstructed volume to the true structure (Table
1). We find that our atomic model is quite sensitive to initialization. While
the model performs adequately when initialized at nearby reference structures,
it makes some mistakes due to local minima, and performs much better when
initialized with the ground truth coordinates. When initialized from random
coordinates, SGD is unable to recover the ground truth atomic coordinates
whatsoever. Instances of atomic local minima include ‘mismatched buttoning’ of
the amino acid backbone where multiple backbone RBFs are trapped in the same
location, as well as incorrect tracing of the protein backbone through the
volume (Figure 3).

Initialization | Initial RMSE (Å) |  $C_{\alpha}$ RMSE (Å) | % within 3Å | Volume NMSE
---|---|---|---|---
Exact 5NI1 | 0.00 | 0.843 | 99.29% | 0.28
Approximate 5NI1 | 5.15 | 3.81 | 77.30% | 0.32
5NI1 + Uniform[6 Å] | 6.50 | 3.19 | 53.19% | 0.35
Random | 34.87 | 17.16 | 2.12% | 0.43
Table 1: Comparison of different choices of initial atomic coordinates when
training the model. Figure 3: Reconstructed atomic structures with RBF centers
initialized either 1) at the exact ground truth values of 5NI1, 2) at an
approximate structure generated from evolution by molecular dynamics, 3) at
5NI1 coordinates randomly perturbed by Uniform[6 Å] noise, or 4) from random
initial values. Structures are colored by $C_{\alpha}$ RMSE to the ground
truth structure (top left).

###  5.2 Heterogeneous reconstruction

The main motivation behind proposing this model is to provide inductive bias
from known atomic structures when reconstructing heterogeneous protein
conformations. As opposed to the previous section where we reconstruct a
single static structure, here we attempt to reconstruct a continuous manifold
of structures from a heterogeneous cryo-EM dataset.

We consider a synthetic dataset of noisy projection images, where the
underlying protein structure possesses a 1D continuous motion of the 5NI1
protein. A bond in the center of the protein is rotated, leading to a large-
scale global conformational change (Figure S2). In our experiments, we
initialize the RBF model to the structure at one end of the reaction
coordinate, and train a heterogeneous RBF model with a 1-dimensional latent
variable on the imaging dataset with ground truth poses.

The RBF model is able to correctly reconstruct the full distribution of 3D
structures containing this large conformational change. The latent encodings
of the images are well correlated with the true reaction coordinate (Figure 4,
left), and the RBF atomic coordinates from traversing the latent space nearly
exactly reconstruct the underlying protein motion (Figure 4, right).

As a baseline, we perform heterogeneous reconstruction with cryoDRGN [2],
which learns an unconstrained neural representation of cryo-EM volumes.
Similarly, we provide the ground truth poses and train a 1-D latent variable
model with identical encoder architecture and training settings. While the
latent space is well correlated with the ground truth motion similar to the
RBF model, the reconstructed volumes from cryoDRGN contain noise and blurring
artifacts in the mobile region whereas Cryofold volumes are regularized using
structural priors from the underlying atomic model (Figure 4, bottom).

We also measure reconstruction accuracy across the reaction coordinate for
four distributions of images across the reaction coordinate: a uniform
distribution, two non-uniform distributions corresponding to an energy barrier
of different heights between the end states; and two discrete clusters with no
images in the middle (Figure 5). For each value of the reaction coordinate, we
approximate the atomic coordinates using the median latent from the images at
that coordinate, and measure its $C_{\alpha}$ RMSE with the ground truth. We
see that when there is even a small probability mass across the reaction
coordinate (Figure 5, top row), the atomic model learns the full distribution
of conformations with high accuracy. However, if transition states are nearly
or completely unobserved in the image distribution (Figure 5, bottom row),
reconstruction accuracy is poor.

Figure 4: Heterogeneous reconstruction results of an unlabeled dataset
containing a uniform distribution of images across the ground truth reaction
coordinate. Predicted 1D latent encoding $z$ plotted against the ground truth
reaction coordinate (left), and reconstructed structures at the specified
values of $z$. Cryofold directly reconstructs atomic coordinates (top). The
unconstrained neural volume representation (cryoDRGN) contains noise and
blurring of moving atoms in the reconstructed volumes (bottom) and does not
produce atomic coordinates.

Figure 5: $C_{\alpha}$ RMSE (blue) for heterogeneous reconstruction of a large
synthetic conformational change of 5ni1, with different distributions of
images (red) across the reaction coordinate. The reference structure used for
initialization is the ground truth atomic structure at reaction coordinate 0.

###  5.3 Model ablations

Using the homogeneous dataset and an initialization from the approximate 5NI1
structure, we explore various choices in design of the RBF model by ablating
the side chain RBF, bond constraints, and the cryo-EM supervision (Table 2).
We find that removing the sidechain RBFs and modeling each amino acid as a
single RBF slightly degrades quality, whereas removing the internal bond terms
leads to a dramatic degradation. Atomic accuracy is substantially worse but
not completely degraded even when ignoring the cryo-EM images due to the
initialization; however, as expected there is low overlap with the true volume
in this case. We expect the volume NMSE to be higher than unconstrained
approaches (e.g. cryoDRGN), as the coarse-grained RBF model does not exactly
match the underlying all-atom generative model even with correct coordinates
(Figure S3).

Model |  $C_{\alpha}$ RMSE (Å) | % within 3Å | Volume NMSE
---|---|---|---
Full Model | 3.81 | 77.30% | 0.32
No Sidechain RBFs | 4.10 | 70.21% | 0.49
No Bonds | 10.42 | 37.59% | 0.36
No Cryo-EM Loss | 5.15 | 74.47% | 1.82
CryoDRGN | N/A | N/A | 0.08
Table 2: Ablations of model components. We remove the sidechain RBFs, the bond
terms between RBFs, and the cryo-EM reconstruction loss; each degrades the
quality of reconstruction.

##  6 Discussion

This work presents Cryofold, a simple framework for incorporating prior
information provided from an atomic model into cryo-EM reconstruction. By
constraining the generative model to act on a (coarse-grained) atomic model,
we constrain the model output to a submanifold of coordinate positions, which
both provides strong regularization to reduce artifacts seen in unconstrained
methods and directly yields interpretable atomic coordinates. Our experiments
suggest that such models are a promising direction for incorporating
structural priors into cryo-EM reconstruction, especially for heterogeneous
structures. However the experimental validation is preliminary: we worked
entirely with synthetic cryo-EM datasets of a small protein using exact poses.
Follow-up work is required to understand how these techniques behave with real
cryo-EM images and volumes, using realistic protein complexes of interest, and
using approximate poses generated by existing tools. For larger protein
complexes, follow-up work will investigate whether a coarser model granularity
may be more appropriate, especially when modeling large-scale heterogeneous
dynamics.

Results on homogeneous datasets suggest that local minima in optimization
space are a major problem for this class of methods if coordinates are not
initialized near the ground truth structure. These local minima problems could
potentially be ameliorated with a combination of improved modeling of
structural priors such as additional bonded terms and steric interactions, and
optimization methods such as Hamiltonian dynamics or Markov Chain Monte Carlo
that can escape local minima. There is a rich literature on these topics in
the domain of molecular dynamics simulation which could carry over to cryo-EM
reconstruction [37, 38].

However, results on heterogeneous datasets suggest that we can correctly learn
distributions with large continuous conformational changes when initialized
from some structure in the distribution. Even structural changes that are too
large to be modeled correctly when only the endpoint structures are observed
(as in homogeneous reconstructions) were successfully reconstructed from
heterogeneous datasets. Additional work is required to more fully characterize
exactly when neural models of heterogeneous structures converge to the correct
distribution.

It should be noted, however, that prior information biases the model output.
In particular, this approach will not be able to model other sources of image
heterogeneity that can not be expressed as motion of the RBFs (e.g.
compositional variation) whose presence are not known a priori. Importantly,
the separation of atomic modelling from volume reconstruction in existing
cryo-EM pipelines provides a crucial mechanism for validation of the volume
reconstruction accuracy. This and related approaches that incorporate atomic
priors should be carefully employed only in settings where prior beliefs are
strongly held, and we believe validation of these methods is an open problem.
Nevertheless, future extensions of this method that can flexibly integrate
biophysical simulation, learning on existing structures, and experimental
single particle cryo-EM data is a promising direction for improved protein
structure determination.

##  7 Acknowledgements

We thank the MIT-IBM Satori team for GPU computing resources and support. This
work was funded by the National Science Foundation Graduate Research
Fellowship Program to E.D.Z., NIH grant R01-GM081871 to B.B., NIH grant
R00-AG050749 to J.H.D., and a grant from the MIT J-Clinic for Machine Learning
and Health to J.H.D. and B.B.

## References

  * [1] Eva Nogales.  The development of cryo-EM into a mainstream structural biology technique.  Nature Methods, 13(1):24–27, jan 2016.
  * [2] Ellen D Zhong, Tristan Bepler, Bonnie Berger, and Joseph H Davis.  CryoDRGN: reconstruction of heterogeneous cryo-EM structures using neural networks.  Nature methods, 18(2):176–185, 2021.
  * [3] Ali Punjani and David J Fleet.  3d variability analysis: Resolving continuous flexibility and discrete heterogeneity from single particle cryo-em.  Journal of Structural Biology, 213(2):107702, 2021.
  * [4] Takanori Nakane, Dari Kimanius, Erik Lindahl, and Sjors Hw Scheres.  Characterisation of molecular motions in cryo-EM single-particle data by multi-body refinement in RELION.  eLife, 7:e36861, June 2018.
  * [5] T Bepler, A Morin, M Rapp, J Brasch, and L Shapiro.  Positive-unlabeled convolutional neural networks for particle picking in cryo-electron micrographs.  Nature Methods, 2019.
  * [6] Ali Punjani, John L Rubinstein, David J Fleet, and Marcus A Brubaker.  cryoSPARC: algorithms for rapid unsupervised cryo-EM structure determination.  Nature Methods, 14(3):290–296, March 2017.
  * [7] Jasenko Zivanov, Takanori Nakane, Björn O. Forsberg, Dari Kimanius, Wim J.H. Hagen, Erik Lindahl, and Sjors H.W. Scheres.  New tools for automated high-resolution cryo-EM structure determination in RELION-3.  eLife, 7, nov 2018.
  * [8] Timothy Grant, Alexis Rohou, and Nikolaus Grigorieff.  cisTEM, user-friendly software for single-particle image processing.  eLife, 7:e14874, mar 2018.
  * [9] Guang Tang, Liwei Peng, Philip R Baldwin, Deepinder S Mann, Wen Jiang, Ian Rees, and Steven J Ludtke.  EMAN2: an extensible image processing suite for electron microscopy.  Journal of structural biology, 157(1):38–46, January 2007.
  * [10] P Emsley, B Lohkamp, W G Scott, and K Cowtan.  Features and development of Coot.  Acta Crystallographica Section D: Biological Crystallography, 66(4):486–501, April 2010.
  * [11] Dorothee Liebschner, Pavel V. Afonine, Matthew L. Baker, Gábor Bunkoczi, Vincent B. Chen, Tristan I. Croll, Bradley Hintze, Li Wei Hung, Swati Jain, Airlie J. McCoy, Nigel W. Moriarty, Robert D. Oeffner, Billy K. Poon, Michael G. Prisant, Randy J. Read, Jane S. Richardson, David C. Richardson, Massimo D. Sammito, Oleg V. Sobolev, Duncan H. Stockwell, Thomas C. Terwilliger, Alexandre G. Urzhumtsev, Lizbeth L. Videau, Christopher J. Williams, and Paul D. Adams.  Macromolecular structure determination using X-rays, neutrons and electrons: Recent developments in Phenix.  Acta Crystallographica Section D: Structural Biology, 75(10):861–877, oct 2019.
  * [12] Leonardo G Trabuco, Elizabeth Villa, Kakoli Mitra, Joachim Frank, and Klaus Schulten.  Flexible fitting of atomic structures into electron microscopy maps using molecular dynamics.  Structure (London, England : 1993), 16(5):673–683, May 2008.
  * [13] Daisuke Kihara, Genki Terashi, and Sai Raghavendra Maddhuri Venkata Subramaniya.  De Novo Computational Protein Tertiary Structure Modeling Pipeline for Cryo-EM Maps of Intermediate Resolution.  Biophysical Journal, 118(3):292a, feb 2020.
  * [14] Dong Si, Spencer A Moritz, Jonas Pfab, Jie Hou, Renzhi Cao, Liguo Wang, Tianqi Wu, and Jianlin Cheng.  Deep Learning to Predict Protein Backbone Structure from High-Resolution Cryo-EM Density Maps.  Scientific Reports, 10(1):1–22, mar 2020.
  * [15] Tristan Ian Croll.  ISOLDE: A physically realistic environment for model building into low-resolution electron-density maps.  Acta Crystallographica Section D: Structural Biology, 2018.
  * [16] Amit Singer and Fred J. Sigworth.  Computational Methods for Single-Particle Electron Cryomicroscopy.  Annual Review of Biomedical Data Science, 3(1):163–190, jul 2020\.
  * [17] Alexandra C Walls, Young-Jun Park, M Alejandra Tortorici, Abigail Wall, Andrew T McGuire, and David Veesler.  Structure, function, and antigenicity of the sars-cov-2 spike glycoprotein.  Cell, 2020.
  * [18] Ronald N Bracewell.  Strip integration in radio astronomy.  Australian Journal of Physics, 9(2):198–217, 1956.
  * [19] Sjors H W Scheres.  A Bayesian view on cryo-EM structure determination.  Journal of Molecular Biology, 415(2):406–418, January 2012.
  * [20] Sjors HW Scheres, Mikel Valle, Rafael Nuñez, Carlos OS Sorzano, Roberto Marabini, Gabor T Herman, and Jose-Maria Carazo.  Maximum-likelihood multi-reference refinement for electron microscopy images.  Journal of Molecular Biology, 348(1):139–149, 2005.
  * [21] Lyumkis, Dmitry, Brilot, Axel F, Theobald, Douglas L, and Grigorieff, Nikolaus.  Likelihood-based classification of cryo-EM images using FREALIGN.  Journal of Structural Biology, 183(3):377–388, September 2013.
  * [22] Ellen D Zhong, Tristan Bepler, Joseph H Davis, and Bonnie Berger.  Reconstructing continuous distributions of 3D protein structure from cryo-EM images.  ICLR, 2020.
  * [23] Joachim Frank and Abbas Ourmazd.  Continuous changes in structure mapped by manifold embedding of single-particle data in cryo-EM.  Methods, 100:61–67, May 2016.
  * [24] Ivet Bahar and A J Rader.  Coarse-grained normal mode analysis in structural biology.  Curr. Opin. Struct. Biol., 15(5):586–592, October 2005.
  * [25] Konrad Hinsen, Nathalie Reuter, Jorge Navaza, David L Stokes, and Jean-Jacques Lacapère.  Normal mode-based fitting of atomic structure into electron density maps: application to sarcoplasmic reticulum ca-atpase.  Biophysical journal, 88(2):818–827, 2005.
  * [26] Qiyu Jin, Carlos Oscar S Sorzano, José Miguel De La Rosa-Trevín, José Román Bilbao-Castro, Rafael Núñez-Ramírez, Oscar Llorca, Florence Tama, and Slavica Jonić.  Iterative elastic 3d-to-2d alignment method using normal modes for studying structural dynamics of large macromolecular complexes.  Structure, 22(3):496–506, 2014.
  * [27] Sandra Schilbach, Merle Hantsche, Dmitry Tegunov, Christian Dienemann, Cristoph Wigge, Henning Urlaub, and Patrick Cramer.  Structures of transcription pre-initiation complex with tfiih and mediator.  Nature, 551(7679):204–209, 2017.
  * [28] Mohamad Harastani, Carlos Oscar S Sorzano, and Slavica Jonić.  Hybrid electron microscopy normal mode analysis with scipion.  Protein Science, 29(1):223–236, 2020.
  * [29] Pilar Cossio and Gerhard Hummer.  Bayesian analysis of individual electron microscopy images: Towards structures of dynamic and heterogeneous biomolecular assemblies.  Journal of structural biology, 184(3):427–437, 2013.
  * [30] Massimiliano Bonomi, Riccardo Pellarin, and Michele Vendruscolo.  Simultaneous determination of protein structure and dynamics using cryo-electron microscopy.  Biophysical journal, 114(7):1604–1613, 2018.
  * [31] Pilar Cossio and Gerhard Hummer.  Likelihood-based structural analysis of electron microscopy images.  Current opinion in structural biology, 49:162–168, 2018.
  * [32] Dari Kimanius, Gustav Zickert, Takanori Nakane, Jonas Adler, Sebastian Lunz, C-B Schönlieb, Ozan Öktem, and Sjors HW Scheres.  Exploiting prior knowledge about biological macromolecules in cryo-em structure determination.  IUCrJ, 8(1), 2021.
  * [33] R Sánchez-García, J Gomez-Blanco, A Cuervo, J M Carazo, COS Sorzano, and J Vargas.  DeepEMhancer: a deep learning solution for cryo-EM volume post-processing.  bioRxiv, 2020.
  * [34] Dan Rosenbaum, Marta Garnelo, Michal Zielinski, Charlie Beattie, Ellen Clancy, Andrea Huber, Pushmeet Kohli, Andrew W. Senior, John Jumper, Carl Doersch, S. M. Ali Eslami, Olaf Ronneberger, and Jonas Adler.  Inferring a continuous distribution of atom coordinates from cryo-em images using vaes, 2021.
  * [35] Anna Choromanska, Mikael Henaff, Michael Mathieu, Gérard Ben Arous, and Yann LeCun.  The loss surfaces of multilayer networks.  2015\.
  * [36] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach 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.  NeurIPS, 2019.
  * [37] Rafael C Bernardi, Marcelo CR Melo, and Klaus Schulten.  Enhanced sampling techniques in molecular dynamics simulations of biological systems.  Biochimica et Biophysica Acta (BBA)-General Subjects, 1850(5):872–877, 2015.
  * [38] Yi Isaac Yang, Qiang Shao, Jun Zhang, Lijiang Yang, and Yi Qin Gao.  Enhanced sampling in molecular dynamics.  The Journal of chemical physics, 151(7):070902, 2019.
  * [39] Eric F Pettersen, Thomas D Goddard, Conrad C Huang, Gregory S Couch, Daniel M Greenblatt, Elaine C Meng, and Thomas E Ferrin.  UCSF Chimera: A visualization system for exploratory research and analysis.  Journal of Computational Chemistry, 25(13):1605–1612, oct 2004\.
  * [40] Ellen D Zhong and Michael R Shirts.  Thermodynamics of Coupled Protein Adsorption and Stability Using Hybrid Monte Carlo Simulations.  Langmuir, 30(17):4952–4961, 2014.

##  Appendix A Appendix

###  A.1 Dataset generation

Synthetic cryo-EM datasets were generated from an atomic model of PDB 5NI1
according to the image formation model as follows: starting from the deposited
atomic model of 5NI1, the 141-residue A-chain subunit of the complex was
extracted. A cryo-EM volume was generated with the ‘molmap‘ command in Chimera
[39] at 3 Å resolution and grid spacing of 1 Å. The volume was zero-padded to
a cubic dimension of $128^{3}$. 50,000 projection images were generated at
random orientations of the volume uniformly from SO(3). Images were then
translated in-plane by $t$ uniformly sampled from $[-10,10]^{2}$ pixels. We
omit the application of the CTF for simplicity in this synthetic dataset.
Gaussian noise was added leading to a signal to noise ratio (SNR) of 0.1,
where we define the whole image as the signal (Figure S1). As real cryo-EM
datasets have variable resolution and thus resolvability of the atomic
structure, we investigate the effect of volume resolution and the included
atoms in generating the dataset’s ground truth volume/images (Table S1). We
find that our atomic modeling is quite robust to the resolution and which
atoms we model in the synthetic data, which suggests that it may transfer well
to real cryo-EM images.

To generate the "Approximate 5NI1" reference structure, we evolve a coarse-
grained $C_{\alpha}$-only model of the 5NI1 A chain under a short Hamiltonian
Monte Carlo simulation near the system’s melting temperature. We use the
implementation from [40], and run 10k Monte Carlo moves where each move
consists of 2.25 ps of molecular dynamics at 320 K. The final frame of the
trajectory is used as the approximate reference structure.

To create synthetic heterogeneous datasets, a dihedral angle in the backbone
of the atomic model was rotated through 0.25 radians with a structure
generated every 0.05 radians along the motion (Figure S2). The resulting 50
structures are used as the ground truth structures for approximating the
continuous motion. For simplicity of heterogeneous dataset generation, we use
the 5NI1 atomic models with $C_{\alpha}$ backbone atoms only. We generate
multiple datasets of 50k total images following the homogeneous dataset image
formation process, each with a different distribution of images along the
reaction coordinate as shown in Figure 5.

Figure S1: Example projection images. Figure S2: Ground truth structures of the heterogeneous datasets simulating a 1D continuous motion transitioning from left (5NI1) to right. All generated structures are shown in the center. Figure S3: Ground truth and reconstructed volumes with a neural network representation of density (cyroDRGN) and with our RBF model initialized from exact coordinates. The RBF volume reconstruction is somewhat worse than that of cryoDRGN because even with exact coordinates, it cannot match the all-atom generative model. Modeled Atoms | Resolution |  $C_{\alpha}$ RMSE (Å) | % within 3Å |
---|---|---|---|---
C$\alpha$ | 3Å | 3.21 | 80.8% |
C$\alpha$ & C$\beta$ | 3Å | 3.29 | 77.3% |
All Atoms | 3Å | 3.75 | 78.0% |
C$\alpha$ | 5Å | 3.20 | 80.9% |
C$\alpha$ & C$\beta$ | 5Å | 3.87 | 73.1% |
All Atoms | 5Å | 3.80 | 74.5% |
C$\alpha$ | 8Å | 3.80 | 74.5% |
C$\alpha$ & C$\beta$ | 8Å | 3.87 | 73.8% |
All Atoms | 8Å | 4.20 | 78.7% |
Table S1: Homogeneous reconstruction model accuracy for different settings of
generating the dataset of cryo-EM images (resolution and which atoms are
included). The model performs similarly across these choices. For other
homogeneous experiments, we use all atoms at 3Å. For heterogeneous
experiments, we use $C_{\alpha}$ at 3Å.

###  A.2 cryoDRGN baseline

For homogeneous reconstruction, we train a cryoDRGN positionally-encoded
3-layer MLP of width 256 for 20 epochs. For heterogeneous reconstruction, both
the encoder and decoder networks are 3-layer MLPs of width 256, and are
trained for 20 epochs.