[論文レビュー] 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のガウス事前分布を置き換え、cryo-EMデータ内の複雑で多モーダルなコンformational分布をよりよく捉えることを目的としている。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.

研究の動機と目的

• VAEフレームワークにおける事前分布の穴問題に取り組むこと。これは、タンパク質コンformationの複雑で多モーダルな事後分布を正しく捉えることができないガウス事前分布の限界に起因する。
• VAEの事前分布をより表現力のある拡散モデルに置き換えることで、cryo-EMデータから直接に分子コンformationの真の生成的モデリングを可能にする。
• 学習された拡散事前分布を活用し、異なるコンformational状態間の潜在空間内探索と補間を高速化する。
• 将来的な条件付き生成、原子モデルとの連携、および分子動力学との統合を支援するフレームワークの開発。

提案手法

• cryo-EM画像から分子コンformationの連続的潜在表現を学習するために、標準的なVAEの目的関数を用いてcryoDRGNを訓練する。
• 訓練済みのcryoDRGN VAEから潜在埋め込みを抽出し、それらを用いて潜在空間でノイズ除去拡散確率的モデル（DDPM）を事前学習する。
• 訓練済みの拡散モデルを柔軟で表現力のある事前分布として用い、真のコンformational状態のデータ分布に従う新しい潜在コードをサンプリングする。
• Langevinダイナミクスなどの反復的サンプリング手法を潜在空間で適用し、異なる分子状態間の高速な探索と補間を可能にする。
• 拡散モデルが学習した滑らかで連続的な潜在多様体を活用し、コンformational状態の最適化のための勾配ベースの最適化を支援する。
• 将来的な拡張として、エンドツーエンド学習や階層的拡散モデルを提案し、表現力の向上とコンformational特徴の分離を図る。

実験結果

リサーチクエスチョン

• RQ1 拡散モデル事前分布は、VAEのガウス事前分布の限界を克服し、cryo-EMデータ内のタンパク質コンformationの複雑で多モーダルな分布を効果的に捉えることができるか？
• RQ2 拡散モデルは、真のデータ分布を正確に反映する高精度な生成的サンプリングを可能にするか？
• RQ3 学習された拡散事前分布は、潜在空間内で異なるコンformational状態間の効率的かつ意味のある探索と補間を可能にするか？
• RQ4 拡散モデルは、条件付き生成や所望の性質を持つコンformational状態の最適化を支援するようにcryoDRGNに統合可能か？
• RQ5 本フレームワークは、エンドツーエンド学習や階層的モデリングをサポートするように拡張可能か？これにより、大規模で不均一なタンパク質アンサンブルのモデリングが向上するか？

主な発見

• 潜在空間拡散モデルは、cryo-EMデータにおけるコンformational状態の複雑で多モーダルな分布を効果的に捉えており、VAEのガウス事前分布を上回る生成的サンプリング性能を示した。
• 合成および実際のcryo-EMデータセットの両方で、拡散事前分布からのサンプルは真のデータ分布に正確に従っており、VAEベースのサンプリングに比べて優れた忠実度を示した。
• 拡散モデルに従うLangevinダイナミクスを用いた潜在空間内探索により、遠く離れたコンformational状態間への高速なジャンプが可能となり、分子多様体の探索が加速した。
• 関心のある状態間の補間が可能で滑らかであり、連続的かつ現実的な形でコンformational遷移をシミュレートできるようになった。
• 本手法により、勾配ベースの最適化が潜在空間で可能となり、例えば結合親和性などの所望の性質に向かってコンformational状態の最適化が可能になった。
• 本フレームワークは、条件付き生成、原子モデル再構築、および分子動力学との統合へと拡張可能であり、バイオ分子モデリングの新たな道を開いた。