CryoDRGN-AI: Notes on ab initio heterogeneous reconstruction and its implications for spliceosome dynamics

Structural biology is increasingly moving beyond the idea of a single static structure. Many macromolecular machines do not exist in one dominant conformation. Instead, they occupy an ensemble of states. Ribosomes rotate, spliceosomes bend and remodel, molecular motors step through cycles, and membrane complexes assemble, disassemble, or shift between functional conformations.

This makes cryo-EM and cryo-ET reconstruction a natural machine learning problem. The question is no longer only:

What is the 3D structure?

but rather:

What is the distribution of 3D structures that explains the observed particle images?

The CryoDRGN series is important because it reformulates heterogeneous cryo-EM reconstruction as a generative modeling problem. Instead of assigning particles into a small number of discrete classes, CryoDRGN learns a continuous latent space where each point corresponds to a possible molecular state.

The 2025 CryoDRGN-AI work pushes this idea further. It moves from heterogeneous reconstruction with pre-estimated poses toward ab initio heterogeneous reconstruction, where both particle conformations and particle poses are optimized within a unified neural framework.

This post helps to learn the CryoDRGN series together with readers from a machine learning perspective, with emphasis on CryoDRGN-AI: its network architecture, forward model, loss design, and why these choices are well matched to structural heterogeneity. At the end, I discuss the spliceosome as an example of how latent structural landscapes can help us reason about dynamic molecular machines and also might help with RNA design.

1. Cryo-EM reconstruction as a Generative Modeling Problem

In single-particle cryo-EM, each observed particle image can be modeled as a noisy 2D projection of a 3D molecular volume:

where:

• $I_i$ is the observed 2D particle image;
• $V_i$ is the 3D density volume corresponding to particle $i$;
• $R_i$ is the 3D orientation;
• $t_i$ is the in-plane translation;
• $\mathcal{P}_{R_i,t_i}$ is the projection operator under pose $(R_i,t_i)$;
• $\mathrm{CTF}_i$ is the contrast transfer function;
• $\epsilon_i$ is imaging noise.

For homogeneous reconstruction, all particles are assumed to come from the same structure:

The goal is then to estimate one consensus 3D density map.

However, this assumption is often too restrictive. Many biological complexes are heterogeneous. Different particles may correspond to different conformational or compositional states. In this case, it is more natural to write:

Here, $z_i$ is a latent variable representing the structural state of particle $i$.

This is the key conceptual shift behind CryoDRGN:

Each particle is not merely a noisy view of one fixed volume.
Each particle is a noisy view of one point on a structural manifold.

2. CryoDRGN: A neural implicit representation of structural heterogeneity

CryoDRGN represents the heterogeneous 3D density field using a coordinate-based neural network. Instead of storing each volume as a voxel grid, it learns a function:

or, in Fourier space:

where:

• $\mathbf{x}$ is a real-space coordinate;
• $\mathbf{k}$ is a Fourier-space coordinate;
• $\mathbf{z}$ is a low-dimensional latent vector;
• $f_\theta$ is a neural decoder, typically implemented as a multilayer perceptron;
• $\theta$ denotes the shared decoder parameters.

This representation means that a 3D volume is not a fixed array. A volume is generated by querying the decoder at many spatial or Fourier coordinates under a given latent code:

Different latent codes generate different structures:

This design has several important implications.

First, the representation is continuous. The decoder can be queried at arbitrary coordinates, so the volume is represented as a continuous field rather than a fixed voxel array.

Second, the representation is generative. Sampling or interpolating latent variables generates new density maps along the learned structural landscape.

Third, the model directly couples representation learning with cryo-EM physics. The network does not simply reconstruct 2D images; it must explain 2D images through 3D density fields and projection operators.

3. Original CryoDRGN as a VAE-Like Model

The original CryoDRGN model can be understood as a variational autoencoder-like architecture.

An encoder maps each particle image to a distribution over latent conformational states:

A latent code is sampled from this approximate posterior and passed to the neural decoder. The decoder generates a 3D density volume conditioned on $\mathbf{z}_i$. This volume is then projected according to the particle pose and modulated by the CTF:

The training objective is based on the evidence lower bound:

Equivalently, as a minimization objective:

The reconstruction term comes from a Gaussian noise assumption. If the image noise is modeled as independent Gaussian noise, then maximizing the image likelihood is equivalent to minimizing squared reconstruction error.

The KL term regularizes the latent space. It encourages the posterior latent distribution to remain close to a simple prior, usually a standard Gaussian. This prevents the latent representation from becoming an arbitrary lookup table with no smooth global structure.

The parameter $\beta$ controls the trade-off between reconstruction fidelity and latent smoothness.

4. Why CryoDRGN Is Not a Standard Autoencoder

A standard image autoencoder would learn:

2D image → encoder → latent code → decoder → reconstructed 2D image

CryoDRGN is different. Its decoder must generate a 3D molecular density:

2D image → latent code → 3D density → projection → reconstructed 2D image

The image reconstruction loss is therefore not merely a pixel-level autoencoding objective. It is a physics-constrained likelihood.

The model must explain each observed image through:

latent conformation
    ↓
3D density field
    ↓
pose-dependent projection
    ↓
CTF modulation
    ↓
predicted 2D particle image

This provides a crucial inductive bias. The network cannot freely memorize image-level statistics. It must explain all particles through a shared family of 3D structures.

This is why CryoDRGN is more than a neural image reconstruction method. It is a generative model constrained by cryo-EM image formation physics.

5. From CryoDRGN to CryoDRGN-AI

The original CryoDRGN framework usually assumes that particle poses are already known or estimated from an upstream pipeline such as RELION or cryoSPARC. In mathematical terms, it optimizes over the decoder and latent conformations while treating $R_i$ and $t_i$ as given:

Here, $A_i$ denotes the cryo-EM forward operator, including projection, CTF modulation, and image-space or Fourier-space sampling.

This assumption can become a bottleneck. If pose estimates are wrong, especially in highly heterogeneous or low-SNR datasets, the learned latent landscape may be distorted. Pose errors can be absorbed into the latent space, causing the model to confuse orientation uncertainty with structural variation.

CryoDRGN-AI addresses this problem by optimizing pose and conformation together:

This turns the problem into ab initio heterogeneous reconstruction. The model no longer relies on fixed upstream pose estimates. Instead, it learns:

global neural decoder parameters
particle-specific conformational latent codes
particle-specific orientations
particle-specific translations

within a unified optimization problem.

6. CryoDRGN-AI Uses an Autodecoder Architecture

A major architectural change in CryoDRGN-AI is the use of an autodecoder.

In the original VAE-style CryoDRGN, the latent code is inferred by an encoder:

In an autodecoder, there is no encoder. Each particle has its own learnable latent code:

These latent codes are stored in a particle-indexed lookup table and optimized directly during training.

The model therefore contains two types of learnable variables.

Global parameters:

• $\theta$: parameters of the neural decoder

Particle-specific parameters:

• $z_i$: conformational or compositional latent code
• $R_i$: particle orientation
• $t_i$: in-plane particle translation

The forward pass can be summarized as:

particle index i
    ↓
retrieve z_i, R_i, and t_i
    ↓
query neural decoder f_θ at Fourier coordinates conditioned on z_i
    ↓
generate pose-dependent 2D projection
    ↓
apply CTF and translation phase shift
    ↓
compare predicted image with observed image

In formula form:

where:

This autodecoder design is important because cryo-EM particle images are extremely noisy. Asking an encoder to infer a reliable latent code from a single low-SNR 2D image is difficult. Directly optimizing particle-specific latent codes can be more stable for reconstruction, even though it sacrifices the fast amortized inference provided by an encoder.

This is conceptually similar to neural implicit shape models such as DeepSDF, where each object has an optimized latent code and a shared decoder maps coordinates plus latent codes to shape values.

7. Fourier-Space Forward Model

CryoDRGN-style methods often operate naturally in Fourier space because of the Fourier slice theorem. The 2D Fourier transform of a projection corresponds to a central slice through the 3D Fourier transform of the volume.

For a 2D Fourier coordinate $\mathbf{k}_{2D}$, the predicted image can be written as:

where $C_i$ is the CTF for particle $i$.

Translation in real space corresponds to a phase shift in Fourier space:

This formulation is powerful because the entire image formation process becomes differentiable with respect to:

decoder parameters θ
latent code z_i
orientation R_i
translation t_i

The model can therefore refine both structure and pose through gradient-based optimization once a sufficiently good pose initialization is available.

8. Loss Function: Why this objective makes sense

Assume the observed image is generated as:

with Gaussian noise:

Then the negative log likelihood is:

This gives the basic reconstruction loss:

However, reconstruction loss alone is not enough. If every particle has a free latent code, the model may use $\mathbf{z}_i$ to memorize noise, junk features, or pose errors. The latent space would then lose biological meaning.

A common solution is to place a Gaussian prior on the latent codes:

The negative log prior gives an L2 penalty:

Thus, a simplified CryoDRGN-AI objective can be written as:

9. Why Latent Regularization Is Essential

Without regularization, an autodecoder can become too flexible. Each particle could receive a latent code that explains not only true conformational variability but also noise, contamination, or reconstruction artifacts.

Regularization constrains the information stored in each latent code:

This encourages the model to encode only variation that is shared across many particles. In other words:

A structural feature should appear in latent space only if it helps explain population-level variation.

This is especially important for biological interpretation. A useful latent space should separate or organize molecular states, not merely index individual particle images.

The low-dimensional bottleneck plays a similar role. If $z_i$ has too many dimensions, the model may memorize particle-specific noise. If it has too few dimensions, the model may fail to capture important biological variation. The dimension of $z$ therefore controls the representational capacity of the structural manifold.

10. Pose Estimation: The key optimization challenge

The most difficult part of ab initio cryo-EM reconstruction is pose estimation.

The pose variable is:

where $R_i$ lies on the rotation group $SO(3)$, and $\mathbf{t}_i$ is the in-plane shift.

Directly optimizing pose from random initialization is difficult because the loss landscape over rotations is highly non-convex. Many orientations can produce partially similar projections, especially when the data are noisy or the molecule has approximate symmetry.

CryoDRGN-AI therefore combines two strategies:

Stage 1: hierarchical pose search
Stage 2: stochastic gradient descent refinement

10.1 Hierarchical Pose Search

Hierarchical pose search performs a coarse-to-fine search over a discretized pose grid:

This stage is robust because it does not require a good gradient-based initialization. It can search broadly over orientation space.

However, pose search is computationally expensive, and its precision is limited by the grid resolution.

10.2 SGD-Based Pose Refinement

After hierarchical search provides a reasonable initialization, pose can be refined as a continuous variable:

This gradient-based stage is more precise and efficient once the optimization has entered a good basin of attraction.

The overall strategy is therefore:

use search for robustness
use gradients for precision

This is one of the most important engineering and algorithmic choices in CryoDRGN-AI.

11. Why CryoDRGN-AI Is More Unified Than a Traditional Pipeline

A conventional cryo-EM workflow often separates the problem into multiple stages:

particle picking
    ↓
2D classification
    ↓
ab initio reconstruction
    ↓
homogeneous refinement
    ↓
heterogeneous refinement
    ↓
3D classification
    ↓
downstream analysis

Each stage compresses uncertainty into intermediate results. For example, once poses are estimated, downstream heterogeneity analysis may treat them as fixed.

CryoDRGN-AI instead places several unknowns into one optimization problem:

From a probabilistic perspective, this resembles MAP estimation:

It is not full Bayesian posterior inference, but it is more unified than a pipeline that fixes poses before modeling heterogeneity.

The model jointly learns:

what the structure looks like
how the structure varies
how each particle is oriented
which particles occupy which structural states

This is why CryoDRGN-AI is best viewed as an ab initio heterogeneous reconstruction framework rather than merely a downstream heterogeneity visualization tool.

12. Conformational and Compositional Heterogeneity

CryoDRGN-style models can represent both conformational and compositional heterogeneity.

Conformational heterogeneity refers to continuous structural motion, such as:

• domain bending
• subunit rotation
• opening and closing
• local flexibility

Compositional heterogeneity refers to the presence or absence of density, such as:

• factor bound or unbound
• ligand present or absent
• subunit assembled or missing
• RNA or protein segment ordered or disordered

Because the decoder generates a full density field conditioned on $\mathbf{z}$, both types of variation can be represented:

In latent space, these patterns may appear differently:

• continuous trajectory: conformational motion
• separate cluster: compositional state
• outlier cluster: broken or junk particles
• branching manifold: multiple assembly or transition pathways

This is one of the main advantages of the method. The user does not need to predefine the number of structural classes. The model can discover continuous and discrete variation from the data.

13. The Inductive Biases Behind CryoDRGN-AI

CryoDRGN-AI works not simply because it uses a neural network, but because it combines several strong inductive biases.

13.1 Three-Dimensional Consistency

Every predicted 2D image must come from a 3D density field.

The model cannot explain images without constructing 3D molecular structure.

13.2 Shared Decoder

All particles share the same neural decoder:

different particles have different latent codes, but the mapping from latent space to density is shared.

This encourages a global structural manifold rather than unrelated per-particle reconstructions.

13.3 Low-Dimensional Bottleneck

Each particle state is represented by a low-dimensional latent code:

the latent variable should capture structural variation, not pixel-level noise.

13.4 Physics-Based Forward Model

Projection, rotation, translation, Fourier slicing, and CTF modulation are explicitly included.

the neural network is constrained by microscope physics.

13.5 Search Plus Differentiable Optimization

Hierarchical search provides robust pose initialization. SGD provides continuous refinement.

global search improves robustness, local gradient optimization improves precision.

Together, these biases make the model suitable for noisy, heterogeneous structural biology data.

14. Why the Loss Should Not Be Interpreted as Ordinary Image Reconstruction

At first glance, the reconstruction term looks like ordinary mean squared error:

But this interpretation is incomplete.

A generic image model could reconstruct each 2D particle without ever learning a molecule. CryoDRGN-AI cannot do this, because the predicted image must be generated through:

Thus, image reconstruction is only the observable consequence of a deeper constraint: the model must explain all images as projections of a shared family of 3D structures.

This is the essential difference between a physics-informed generative model and a standard image autoencoder.

15. CryoDRGN, CryoDRGN-ET, and CryoDRGN-AI

The CryoDRGN series can be viewed as a progression.

15.1 CryoDRGN

Core problem:

given particle images and estimated poses, learn continuous structural heterogeneity.

Core model:

VAE-like architecture with a neural implicit 3D decoder

15.2 CryoDRGN-ET

Core problem:

extend heterogeneous reconstruction to cryo-electron tomography.

Core model change:

use subtomogram tilt-series observations to learn heterogeneous 3D density maps in situ.

CryoDRGN-ET is important because it brings neural heterogeneous reconstruction closer to the cellular environment, where molecules are observed in crowded and native-like contexts.

15.3 CryoDRGN-AI

Core problem:

perform ab initio heterogeneous reconstruction without relying on fixed upstream pose estimates.

Core model change:

autodecoder architecture
particle-specific latent codes
particle-specific pose optimization
hierarchical pose search
SGD pose refinement
physics-informed image likelihood

CryoDRGN-AI is therefore not merely another post-processing method. It is a more complete reconstruction model that jointly estimates structure, heterogeneity, and pose.

16. Spliceosome as a Case Study

The spliceosome is an ideal example for CryoDRGN-style modeling.

Splicing is not a single structural event. It is a sequence of RNA-protein rearrangements involving snRNP assembly, branch point recognition, splice-site positioning, catalytic activation, exon ligation, and disassembly.

The spliceosome therefore does not naturally fit into a small number of rigid structural classes. Even within one nominal biochemical state, such as the pre-catalytic spliceosome, particles may show bending, rotation, subcomplex rearrangement, and local density variation.

Traditional 3D classification faces a trade-off:

too few classes: continuous motion is oversimplified;
too many classes: each class has lower signal and noisier maps.

CryoDRGN-AI offers a different representation. Instead of asking which class each spliceosome particle belongs to, it asks where each particle lies in a learned structural latent space.

This latent space can reveal:

continuous conformational trajectories
separated compositional states
outlier particle populations
broken or junk particles
rare intermediate conformations

For pre-catalytic spliceosome data, CryoDRGN-AI demonstrates how a learned latent representation can recover continuous bending motion and separate abnormal particles into outlier regions of latent space. This is biologically meaningful because spliceosome function depends on coordinated structural rearrangements rather than a single static architecture.

17. Interpreting a Spliceosome Latent Landscape

Suppose each spliceosome particle receives a latent code:

Collecting them gives a latent matrix:

We can visualize this latent distribution using PCA, UMAP, diffusion maps, or graph-based trajectory analysis.

Then we can select representative latent points:

and decode them into density maps:

If moving along one latent direction generates maps where the spliceosome gradually bends, then that direction may correspond to a major conformational coordinate.

In this sense, the latent space becomes a candidate reaction coordinate for structural dynamics:

• latent direction
• decoded density maps
• interpretable molecular motion

However, interpretation must be cautious. A latent direction may mix several factors:

• true conformational variation
• compositional heterogeneity
• pose uncertainty
• particle quality
• CTF-related artifacts
• local flexibility
• junk particles

18. From Spliceosome Dynamics to RNA Design

The long-term value of CryoDRGN-AI is not only reconstructing structures. It may help define design objectives for dynamic molecular systems.

For a static protein, the design objective is often to stabilize a target structure.

For a dynamic molecular machine such as the spliceosome, a more meaningful objective may be to shift the population from one conformational basin to another.

In latent-space terms, a perturbation changes the structural distribution:

This provides a quantitative way to describe how a mutation, small molecule, antisense oligonucleotide, or RNA sequence change alters a molecular ensemble.

18.1 Stabilizing Specific Spliceosome States

If a latent region corresponds to a stalled, inactive, or pre-catalytic spliceosome state, one could imagine designing perturbations that increase the occupancy of that region:

This design goal is ensemble-based rather than structure-based. The objective is not merely to bind a pocket, but to reshape the conformational population.

18.2 Connecting RNA Sequence to Structural State Distribution

Splicing regulation is strongly sequence-dependent. Mutations near splice sites, branch points, polypyrimidine tracts, or regulatory motifs may alter spliceosome assembly and activation.

A future experimental-computational pipeline could look like:

RNA sequence perturbation
    ↓
cryo-EM or cryo-ET data collection
    ↓
CryoDRGN-AI latent landscape
    ↓
population shift analysis
    ↓
connection to splicing outcome

This would make it possible to ask whether a sequence perturbation changes not only the final splicing product, but also the distribution of structural intermediates.

18.3 Designing Dynamic Function Instead of Static Stability

For molecular machines, flexibility is often functional. A useful intervention may need to preserve one motion while blocking another.

Possible design objectives include:

preserve productive motion
block harmful transition
stabilize a transient intermediate
destabilize a nonproductive basin
increase occupancy of an active state

CryoDRGN-AI provides a way to measure whether such objectives are achieved at the structural ensemble level.

19. Limitations and Cautions

CryoDRGN-AI is powerful, but its outputs require careful interpretation.

First, the learned latent space is not unique. Different initializations, masks, regularization strengths, preprocessing choices, and pose parameterizations may affect the geometry of the latent representation.

Second, decoded intermediate maps need particle support. If a region of latent space contains few particles, the generated structure may reflect interpolation rather than a well-supported biological state.

Third, pose and conformation can be entangled. Even though CryoDRGN-AI jointly optimizes pose and conformation, some pose uncertainty may still leak into latent variables, especially in low-SNR or preferred-orientation datasets.

Fourth, latent outlier clusters may represent biologically meaningful rare states, but they may also represent junk particles, broken complexes, ice contamination, or other data-quality issues.

Fifth, density-level interpretation should eventually be connected to atomic modeling, biochemical perturbation, mutagenesis, functional assays, or orthogonal structural evidence.

20. Conclusion

The CryoDRGN series reframes heterogeneous cryo-EM and cryo-ET reconstruction as a problem in physics-informed generative modeling.

Original CryoDRGN introduced the core idea:

image encoder
    ↓
latent conformation
    ↓
neural 3D decoder
    ↓
projection and CTF
    ↓
image likelihood

CryoDRGN-AI extends this into ab initio heterogeneous reconstruction:

particle-specific latent code
    +
particle-specific pose
    ↓
shared neural implicit 3D decoder
    ↓
differentiable cryo-EM forward model
    ↓
physics-based reconstruction loss

Its main machine learning contributions are the combination of:

neural implicit representation
autodecoder latent optimization
learnable particle poses
hierarchical pose search
SGD pose refinement
Gaussian image likelihood
latent regularization
physics-informed forward modeling

For dynamic systems such as the spliceosome, this framework is especially compelling. It does not force particles into a small number of static classes. Instead, it learns a structural landscape where continuous motion, compositional variation, abnormal particles, and possible intermediate states can be studied together.

In the long term, this type of model may shift structural biology from determining individual structures toward measuring and designing structural distributions. For biological machines, that may be the more natural level of description: not one structure, but a landscape of possible states and transitions. This is especially relevant for RNA design, where the goal is often not only to design a sequence that folds into one target structure, but to design an RNA molecule whose ensemble, accessibility, interactions, and conformational switching behavior produce the desired function. In this view, CryoDRGN-style latent landscapes could provide structural feedback for designing RNAs that stabilize productive spliceosomal states, avoid nonproductive conformations, or tune the population balance between competing molecular states.

References

1. Zhong, E. D., Bepler, T., Berger, B., and Davis, J. H. CryoDRGN: reconstruction of heterogeneous cryo-EM structures using neural networks. Nature Methods, 2021.

2. Kinman, L. et al. Uncovering structural ensembles from single-particle cryo-EM data using cryoDRGN. Nature Protocols, 2023.

3. Rangan, R. et al. CryoDRGN-ET: deep reconstructing generative networks for visualizing dynamic biomolecules inside cells. Nature Methods, 2024.

4. Levy, A. et al. CryoDRGN-AI: neural ab initio reconstruction of challenging cryo-EM and cryo-ET datasets. Nature Methods, 2025.

5. Plaschka, C., Lin, P. C., and Nagai, K. Structure of a pre-catalytic spliceosome. Nature, 2017.

6. Nakane, T. et al. Characterisation of molecular motions in cryo-EM single-particle data by multi-body refinement in RELION. eLife, 2018.