[논문 리뷰] Latent Space Diffusion Models of Cryo-EM Structures

Karsten Kreis, Tim Dockhorn외 2명|arXiv (Cornell University)|2022. 11. 25.

Advanced Electron Microscopy Techniques and Applications인용 수 12

한 줄 요약

이 논문은 코어드라이드라곤(cryoDRGN)에서 학습 가능한 사전으로 잠재 공간 확산 모델을 제안하며, 기존 VAE의 가우시안 사전을 대체하여 냉동전자현미경(cryo-EM) 데이터에서 복잡하고 다중모달적인 구조적 분포를 더 잘 포착한다. 코어드라이드라곤이 학습한 잠재 임bedding에 기반해 확산 모델을 훈련시킴으로써, 정확한 생성적 샘플링, 효율적인 잠재공간 탐색, 그리고 분자 상태 간의 보간을 가능하게 하여 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 기반 복원에서의 사전 구멍 문제를 해결한다.
• VAE 사전을 더 표현력 있는 확산 모델로 대체하여 냉동전자현미경(cryo-EM) 데이터에서 직접 분자 구조적 분포를 진정으로 생성 모델링할 수 있도록 한다.
• 학습된 확산 사전을 활용해 잠재공간 탐색과 서로 다른 구조적 상태 간의 보간을 가속화한다.
• 향후 조건부 생성, 원자 모델 결합, 분자 동역학와의 통합을 지원하는 프레임워크를 개발한다.

제안 방법

• 냉동전자현미경(cryo-EM) 영상에서 단백질 구조적 분포의 연속적인 잠재 표현을 학습하기 위해 표준 VAE 목표로 코어드라이드라곤을 훈련시킨다.
• 훈련된 코어드라이드라곤 VAE의 잠재 임베딩을 추출하고, 이를 기반으로 잠재공간에서 노이즈 제거 확산 확률 모델(DDPM)을 사전 훈련한다.
• 학습된 확산 모델을 유연하고 표현력 있는 사전으로 사용하여 진정한 데이터 분포를 따르는 새로운 잠재 코드를 샘플링한다.
• 잠재공간에서 랭글레인 역동법과 같은 반복적 샘플링 방법을 적용하여 서로 다른 분자 상태 간의 빠른 탐색과 보간을 가능하게 한다.
• 확산 모델이 학습한 매끄럽고 연속적인 잠재 다양체를 활용해 구조적 상태 정밀 조정을 위한 기울기 기반 최적화를 지원한다.
• 표현력 향상과 구조적 특징의 분리에 기여하기 위해 엔드 투 엔드 훈련 및 계층적 확산 모델과 같은 향후 확장 방안을 제안한다.

실험 결과

연구 질문

• RQ1 확산 모델 사전은 냉동전자현미경(cryo-EM) 데이터에서 단백질의 복잡하고 다중모달적인 분포를 효과적으로 포착할 수 있는가, 기존 VAE의 가우시안 사전의 한계를 극복할 수 있는가?
• RQ2 확산 모델은 진정한 데이터 분포를 정확히 반영하는 고해상도의 분자 구조 생성 샘플링을 가능하게 하는가?
• RQ3 학습된 확산 사전은 잠재공간에서 서로 다른 구조적 상태 간의 효율적이고 의미 있는 탐색과 보간을 가능하게 하는가?
• RQ4 확산 모델은 코어드라이드라곤에 어떻게 통합되어 조건부 생성 또는 원하는 특성을 지닌 구조적 상태의 최적화를 지원할 수 있는가?
• RQ5 이 프레임워크는 엔드 투 엔드 훈련이나 계층적 모델링을 지원하도록 확장되어 대규모 이질적인 단백질 군집의 모델링을 향상시킬 수 있는가?

주요 결과

• 잠재공간 확산 모델은 냉동전자현미경(cryo-EM) 데이터에서 복잡하고 다중모달적인 구조적 상태 분포를 성공적으로 포착하였으며, 생성 샘플링에서 VAE의 가우시안 사전보다 뛰어난 성능을 보였다.
• 확산 사전에서 생성된 샘플은 합성 및 실제 냉동전자현미경(cryo-EM) 데이터셋 모두에서 진정한 데이터 분포를 정확히 따르며, VAE 기반 샘플링 대비 개선된 정밀도를 보였다.
• 확산 모델에 의해 안내된 랭글레인 역동법을 활용한 잠재공간 탐색은 먼 구조적 상태 간의 빠른 점프를 가능하게 하여 분자 다양체 탐색을 가속화했다.
• 관심 있는 상태 간의 보간은 가능하고 매끄럽게 이루어져 연속적이고 현실적인 구조적 전이 시뮬레이션을 가능하게 하였다.
• 이 방법은 기울기 기반 최적화를 통해 잠재공간에서 구조적 상태를 원하는 성질(예: 결합 친화도) 향한 정밀 조정이 가능하게 하였다.
• 이 프레임워크는 조건부 생성, 원자 모델 복원, 분자 동역학와의 통합을 지원하도록 확장 가능하여 생물분자 모델링의 새로운 길을 열었다.