AlphaFold Meets Flow Matching for Generating Protein Ensembles

The biological functions of proteins often depend on dynamic structural ensembles. In this work, we develop a flow-based generative modeling approach for learning and sampling the conformational landscapes of proteins. We repurpose highly accurate single-state predictors such as AlphaFold and ESMFold and fine-tune them under a custom flow matching framework to obtain sequence-conditoned generative models of protein structure called AlphaFlow and ESMFlow. When trained and evaluated on the PDB, our method provides a superior combination of precision and diversity compared to AlphaFold with MSA subsampling. When further trained on ensembles from all-atom MD, our method accurately captures conformational flexibility, positional distributions, and higher-order ensemble observables for unseen proteins. Moreover, our method can diversify a static PDB structure with faster wall-clock convergence to certain equilibrium properties than replicate MD trajectories, demonstrating its potential as a proxy for expensive physics-based simulations. Code is available at https://github.com/bjing2016/alphaflow.

한 줄 요약

AlphaFlow와 ESMFlow는 PDB 및 MD 앙상블 데이터를 기반으로 단백질 구조 다양성을 생성하는 flow matching 기반 생성 모델이다.

핵심 기여도

AlphaFold와 ESMFold를 flow matching 프레임워크로 재구성하여 AlphaFlow와 ESMFlow라는 시퀀스 조건부 생성 모델을 제안.
PDB에서 학습했을 때 AlphaFold의 MSA subsampling 기반 방법보다 정밀도와 다양도의 균형이 우수함.
MD 앙상블 데이터를 사용한 학습 시, 유연성, 위치 분포, 고차 관측량을 정확히 재현.
정적 PDB 구조에서 MD 시뮬레이션보다 빠르게 평형 속성을 수렴.

핵심 아이디어

기존 AlphaFold와 ESMFold는 단일 구조 예측을 위한 회귀 모델이지만, 이 연구는 flow matching을 통해 이를 생성 모델로 변환하는 전략을 제안한다. 핵심 아이디어는 반복적 디노이징 프레임워크가 회귀 모델을 생성 모델로 전환할 수 있다는 통찰에 기반한다. 이는 MSA subsampling과 같은 추론 시 조작과 달리, 학습 시에 구조 다양성을 학습할 수 있는 원칙적인 접근법이다. 특히, flow matching은 이미지 분야에서 확립된 기법이지만, 단백질 구조에 적용된 것은 최근이며, 이 연구는 AlphaFold 및 ESMFold의 아키텍처와 훈련 방식에 맞춘 커스텀 flow matching 프레임워크를 개발하였다. 이 프레임워크는 harmonic diffusion의 polymer-structured prior를 활용하면서도, 결측치나 자른 residue에 강한 scale-invariant noising 과정을 도입하여 성능을 향상시켰다.

기술적 접근법

**모델**: AlphaFlow와 ESMFlow는 AlphaFold와 ESMFold를 기반으로 flow matching을 적용한 생성 모델.
**프레임워크**: 커스텀 flow matching 프레임워크를 사용.
**데이터**: PDB 구조와 ATLAS MD 앙상블 데이터셋을 사용.
**하이퍼파라미터**: flow matching의 noising 과정에서 scale-invariant 특성을 도입.
**학습 전략**: 기존 AlphaFold의 회귀 손실(FAPE)을 기반으로 flow matching의 vector field를 학습.

주요 결과

**PDB 데이터셋**: AlphaFlow는 MSA subsampling 기반 AlphaFold 대비 정밀도와 다양도의 균형이 우수함.
**ATLAS MD 앙상블**: AlphaFlow는 유연성, 위치 분포, 고차 관측량(예: 간헐적 접촉, 용매 노출)을 정확히 재현.
**정적 PDB 구조**: AlphaFlow는 MD 시뮬레이션 대비 더 빠르게 평형 속성에 수렴 (wall-clock convergence).

의의 및 한계

이 연구는 단백질 구조 예측 모델을 생성 모델로 변환하는 새로운 패러다임을 제시하며, MSA subsampling과 같은 추론 시 조작에 의존하지 않는 학습 기반 접근법을 도입하였다. 특히, PLM 기반 모델(EsmFold)에도 적용 가능하다는 점에서 유연성이 높다. 또한, MD 앙상블 데이터를 학습에 활용할 수 있어, 물리 기반 시뮬레이션 대체 모델로서의 잠재력을 보여준다. 그러나, flow matching 기반 모델은 아직 단백질 구조 생성 분야에서 초기 단계에 있으며, 대규모 앙상블 데이터셋에 대한 일반화 능력이나 계산 효율성에 대한 추가 연구가 필요하다.

실용적 활용

AlphaFlow와 ESMFlow는 단백질 기능 연구, 분자 역학 시뮬레이션 대체, 구조 기반 약물 설계 등 구조 생물학 분야에서 널리 활용될 수 있다. 특히, 고해상도 cryo-EM 데이터와 구조 이질성을 분석하는 알고리즘과 결합될 경우, 더 정밀한 앙상블 생성이 가능할 것으로 기대된다.