End-to-end Training of Latent Space Diffusion Models for Conformational Heterogeneity in Cryo-EM Reconstruction

Biological macromolecules are dynamic and undergo conformational changes to perform their function. An understanding of the conformational landscape is therefore important to gain insights into how biomolecules transition between conformations. Cryo-electron microscopy (cryo-EM) has emerged as a powerful tool for visualizing biological macromolecular complexes. While it allows the reconstruction of heterogeneous structures concurrently, exploring the structure dynamics, particularly the conformational path between structures, remains challenging. The CryoDRGN method, based on variational autoencoders (VAEs), provides an effective way to generate the embeddings of the structures in latent spaces. However, the mismatch of its Gaussian prior and the actual latent distribution limits its ability to generate the conformation path. While previous works suggested training a standalone diffusion model to model the latent distribution of pretrained CryoDRGN VAEs, we introduce an innovative end-to-end approach that trains VAE and the latent diffusion model jointly. We test our method's ability on three datasets, demonstrating its ability to model the latent embedding landscapes and generate plausible structures. The ability to generate transition states and pathways consistent with the data distribution will allow for the integration of generated models with techniques such as molecular-dynamics and mechanistic calculations for the exploration of free-energy landscapes.