[Paper Review] Inferring a Continuous Distribution of Atom Coordinates from Cryo-EM Images using VAEs

Dan Rosenbaum|arXiv (Cornell University)|Jun 26, 2021

Advanced Electron Microscopy Techniques and Applications24 references31 citations

TL;DR

The paper introduces a fully end-to-end variational autoencoder framework that reconstructs a continuous distribution of atomic protein structures directly from cryo-EM particle images by encoding conformation and pose in separate latent spaces and rendering atom coordinates through a differentiable cryo-EM renderer.

ABSTRACT

Cryo-electron microscopy (cryo-EM) has revolutionized experimental protein structure determination. Despite advances in high resolution reconstruction, a majority of cryo-EM experiments provide either a single state of the studied macromolecule, or a relatively small number of its conformations. This reduces the effectiveness of the technique for proteins with flexible regions, which are known to play a key role in protein function. Recent methods for capturing conformational heterogeneity in cryo-EM data model it in volume space, making recovery of continuous atomic structures challenging. Here we present a fully deep-learning-based approach using variational auto-encoders (VAEs) to recover a continuous distribution of atomic protein structures and poses directly from picked particle images and demonstrate its efficacy on realistic simulated data. We hope that methods built on this work will allow incorporation of stronger prior information about protein structure and enable better understanding of non-rigid protein structures.

Motivation & Objective

• Motivate recovering conformational heterogeneity in cryo-EM data beyond discrete states.
• Propose an end-to-end deep learning model that infers a distribution of atomic structures from noisy 2D projections.
• Leverage a differentiable renderer to map atom coordinates to simulated EM images for training.
• Incorporate explicit atom-space priors (bond lengths, geometry) to regularize reconstructions.
• Evaluate on simulated data to demonstrate multi-modality recovery and continuous latent representations.

Proposed method

• Model the image formation with a differentiable renderer from atom coordinates to EM images.
• Use a variational autoencoder with two latent spaces: conformation (32-d) and pose (32-d) to disentangle structure from orientation.
• Decode latent samples into delta backbone frames (per-residue pose) relative to a base conformation.
• Render decoded atom configurations to mean images and model observation noise.
• Train with an IWAE-like objective by sampling multiple posterior z's per image to encourage multimodal coverage.
• Incorporate auxiliary losses: backbone continuity and centering losses; use Gram-Schmidt to obtain rotation matrices; optimize with Adam on batches.

Experimental results

Research questions

• RQ1 Can a VAE-based model recover a continuous distribution of atomic protein conformations directly from 2D cryo-EM images?
• RQ2 Does disentangling pose and conformation in latent space improve recovery of multimodal structural states compared to a single latent space?
• RQ3 Can a differentiable renderer and atom-space priors enable end-to-end learning from images to atom coordinates?
• RQ4 How well does the method capture multimodal conformational distributions and what are its limitations on simulated data?

Key findings

• The method can train neural networks to represent diverse states of a protein from 2D particle images (AurA kinase) with a true distribution approximated within 3.35 Angstrom by the EMD-RMSD metric.
• Latent space interpolation shows continuous transitions between states, demonstrating a continuous distribution of conformations.
• Disentangling pose and conformation latent spaces and pre-training pose prediction are crucial to avoid degenerate solutions and enable non-degenerate structural reconstructions.
• Conditional sampling from the posterior yields weaker correlations to the true distributions under high-noise 2D data, indicating the necessity of shared capacity across many images.
• The evaluation on simulated data acknowledges limitations, with the model performing well in unconditional sampling and partial success in conditional predictions.