[논문 리뷰] Improving AlphaFlow for Efficient Protein Ensembles Generation

Shaoning Li, Mingyu Li외 5명|arXiv (Cornell University)|2024. 07. 08.

Cell Image Analysis Techniques인용 수 5

한 줄 요약

Alpha Flow -Lit는 특징으로 조건화하여 단백질 앙상블 생성 속도를 높이고, Alpha Flow와 유사한 정확도를 유지하며 약 47x 더 빠른 샘플링을 달성합니다.

ABSTRACT

Investigating conformational landscapes of proteins is a crucial way to understand their biological functions and properties. AlphaFlow stands out as a sequence-conditioned generative model that introduces flexibility into structure prediction models by fine-tuning AlphaFold under the flow-matching framework. Despite the advantages of efficient sampling afforded by flow-matching, AlphaFlow still requires multiple runs of AlphaFold to finally generate one single conformation. Due to the heavy consumption of AlphaFold, its applicability is limited in sampling larger set of protein ensembles or the longer chains within a constrained timeframe. In this work, we propose a feature-conditioned generative model called AlphaFlow-Lit to realize efficient protein ensembles generation. In contrast to the full fine-tuning on the entire structure, we focus solely on the light-weight structure module to reconstruct the conformation. AlphaFlow-Lit performs on-par with AlphaFlow and surpasses its distilled version without pretraining, all while achieving a significant sampling acceleration of around 47 times. The advancement in efficiency showcases the potential of AlphaFlow-Lit in enabling faster and more scalable generation of protein ensembles.

연구 동기 및 목표

• Alpha Flow의 단백질 앙상블 샘플링 효율성의 한계를 이해한다.
• 계산 부하를 줄이기 위한 가볍고 특징 조건화된 변형을 개발한다.
• MD 궤적 데이터에서 Alpha Flow -Lit를 Alpha Flow 및 그 증류 버전과 비교 평가한다.
• 생성된 앙상블이 단백질 다이내믹스, 국부 잔류 배열, 장거리 상관에 대해 어떤 것을 보이는지 분석한다.

제안 방법

• 홈으로 유지된 AlphaFold 임베더 및 Evoformer를 고정 상태로 두어 단일 특징과 쌍 특징을 사전 계산한다.
• MSA 대신 계산된 단일/쌍 특징(비틀림 각도와 접촉 맵)을 사용해 경량 StructureModule을 조건화한다.
• 사전 학습 가중치를 최소로 교란하는 ControlNet에서 영감을 받은 사전 계산된 특징 조건 입력 임베딩을 사용한다.
• PDB 사전학습 없이 ALTAS MD 궤적에서 Alpha Flow -Lit를 전체 디노이징 스텝과 경량 디노이저를 사용해 학습한다.
• MD 기반 지표(RMSD, RMSF, JSD, DCCM 등)를 사용해 Alpha Flow -Full 및 Alpha Flow -Distilled와 런타임 및 앙상블 품질을 비교한다.
• x_t를 특징으로 하고 학습된 벡터 필드를 갖춘 AlphaFlow 기반 CNF/flow 매칭 프레임워크를 유지한다。

Figure 1: Model architecture of sequence-conditioned Alpha Flow ( left ) and feature-conditioned Alpha Flow -Lit ( right ). $T$ : Denoising steps; $x_{t}$ : Noisy structure; $\tilde{x}_{0}$ : Predicted structure.

실험 결과

연구 질문

• RQ1 특징 조건화된 더 가벼운 구조 모듈이 전체 MSAs 기반 조건화를 대체해 Alpha Flow를 가속화하면서 앙상블 품질을 손상시키지 않는가?
• RQ2 Alpha Flow와 그 증류 버전에 비해 Alpha Flow -Lit가 단백질 길이에 따라 어떻게 규모를 확장하는가?
• RQ3 Alpha Flow -Lit 앙상블이 ground-truth MD 앙상블에 비해 필수 다이내믹스, 국부 잔류 동작, 장거리 상관을 보존하는가?
• RQ4 Alpha Flow -Lit와 Alpha Flow -Full/Distilled 간의 다양성과 정확성 사이의 트레이드오프는 무엇인가?

주요 결과

• Alpha Flow -Lit는 평균적으로 Alpha Flow -Full에 비해 약 47x 더 빠른 샘플링 속도를 달성한다.
• Alpha Flow -Lit는 ATLAS MD 궤적에서 앙상블 품질이 Alpha Flow와 일치하며, 여러 지표에서 Alpha Flow -Distilled를 능가한다.
• Alpha Flow -Lit는 일반 다이내믹스(RMSD 기반)에서 ground-truth MD 앙상블과 가장 강한 Pearson 상관 관계를 보이며 Alpha Flow -Full와 유사한 다양성을 유지한다.
• PCA 기반 필수 다이내믹스에서 Alpha Flow -Lit는 Alpha Flow -Full과 유사하게 작동하고 Alpha Flow -Distilled보다 낫다.
• 잔류에 대한 로컬 배열 분석(RMSF, 접촉 확률, 이면 각 분포)은 Alpha Flow -Lit가 ground-truth 패턴을 밀접하게 추적하며 Distilled에 비해 의미 있는 개선을 보인다.
• 장거리 결합 분석(DCCM)에서 Alpha Flow -Lit가 증류 버전보다 잔류 상관을 더 잘 포착한다.

Figure 2: Visualization of MD evaluation from MD, Alpha Flow -Lit and Alpha Flow . ( A ) Runtime comparison corresponding to the sequence length and their fitted curves. ( B ) Principal components analysis (PCA) for 6q9c_A ensembles. The representative structures are pointed out. ( C, D ) Ensembles