[论文解读] Latent Space Diffusion Models of Cryo-EM Structures

Karsten Kreis, Tim Dockhorn+2 more|arXiv (Cornell University)|Nov 25, 2022

Advanced Electron Microscopy Techniques and Applications被引用 12

一句话总结

本文提出一种潜在空间扩散模型作为 cryoDRGN 中的可学习先验，取代标准 VAE 的高斯先验，以更好地捕捉冷冻电镜数据中复杂的多模态构象分布。通过在 cryoDRGN 学习到的潜在嵌入上训练扩散模型，该方法实现了精确的生成采样、高效的潜在空间遍历以及分子状态之间的插值，显著改善了基于 VAE 的重构中的先验洞问题。

ABSTRACT

Cryo-electron microscopy (cryo-EM) is unique among tools in structural biology in its ability to image large, dynamic protein complexes. Key to this ability is image processing algorithms for heterogeneous cryo-EM reconstruction, including recent deep learning-based approaches. The state-of-the-art method cryoDRGN uses a Variational Autoencoder (VAE) framework to learn a continuous distribution of protein structures from single particle cryo-EM imaging data. While cryoDRGN can model complex structural motions, the Gaussian prior distribution of the VAE fails to match the aggregate approximate posterior, which prevents generative sampling of structures especially for multi-modal distributions (e.g. compositional heterogeneity). Here, we train a diffusion model as an expressive, learnable prior in the cryoDRGN framework. Our approach learns a high-quality generative model over molecular conformations directly from cryo-EM imaging data. We show the ability to sample from the model on two synthetic and two real datasets, where samples accurately follow the data distribution unlike samples from the VAE prior distribution. We also demonstrate how the diffusion model prior can be leveraged for fast latent space traversal and interpolation between states of interest. By learning an accurate model of the data distribution, our method unlocks tools in generative modeling, sampling, and distribution analysis for heterogeneous cryo-EM ensembles.

研究动机与目标

• 解决 cryoDRGN 的 VAE 框架中因高斯先验无法匹配蛋白质构象复杂且多模态的后验分布而导致的先验洞问题。
• 通过用更具表达能力的扩散模型替代 VAE 先验，实现在冷冻电镜数据上对分子构象的真正生成建模。
• 利用学习到的扩散先验加速潜在空间中不同构象状态之间的遍历与插值。
• 构建一个支持未来条件生成、原子模型耦合以及与分子动力学集成的框架，以增强生物分子模拟能力。

提出的方法

• 使用标准 VAE 目标训练 cryoDRGN，以从冷冻电镜图像中学习分子构象的连续潜在表征。
• 从训练好的 cryoDRGN VAE 中提取潜在嵌入，并在潜在空间中预训练去噪扩散概率模型（DDPM）。
• 将训练好的扩散模型用作灵活且表达能力强的先验，以采样符合构象状态真实数据分布的新潜在代码。
• 在潜在空间中应用如拉普拉斯动力学等迭代采样方法，以实现不同分子状态间快速遍历与插值。
• 利用扩散模型学习到的平滑连续潜在流形，支持基于梯度的优化，以实现构象状态的精细化调整。
• 提出未来扩展方向，如端到端训练和分层扩散模型，以提升表达能力并解耦构象特征。

实验结果

研究问题

• RQ1 扩散模型先验能否有效捕捉冷冻电镜数据中蛋白质构象的复杂多模态分布，从而克服 VAE 高斯先验的局限性？
• RQ2 扩散模型能否实现高保真度的分子结构生成采样，使其准确反映底层数据分布？
• RQ3 学习到的扩散先验能否促进潜在空间中不同构象状态之间高效且有意义的遍历与插值？
• RQ4 如何将扩散模型集成到 cryoDRGN 中，以支持条件生成或针对特定性质优化构象状态？
• RQ5 该框架能否扩展以支持端到端训练或分层建模，从而改善大规模异质蛋白复合物的建模？

主要发现

• 潜在扩散模型成功捕捉了冷冻电镜数据中构象状态的复杂多模态分布，在生成采样方面优于 VAE 的高斯先验。
• 从扩散先验生成的样本在合成和真实冷冻电镜数据集上均准确遵循真实数据分布，其保真度显著优于基于 VAE 的采样方法。
• 借助扩散模型引导的拉普拉斯动力学进行潜在空间遍历，可实现从远距离构象状态间的高效跳跃，加速分子流形的探索。
• 在感兴趣状态之间实现可行且平滑的插值，能够以连续且真实的方式模拟构象转变过程。
• 该方法支持潜在空间中的基于梯度的优化，可将构象状态优化至期望属性（如结合亲和力）方向。
• 该框架可扩展至条件生成、原子模型重建以及与分子动力学的集成，为生物分子建模开辟新途径。