[논문 리뷰] SE(3)-Stochastic Flow Matching for Protein Backbone Generation

Avishek Joey Bose, Tara Akhound-Sadegh외 8명|arXiv (Cornell University)|2023. 10. 03.

Protein purification and stability인용 수 14

한 줄 요약

FoldFlow는 SE(3)–불변 정규화 흐름을 단백질 백본 생성에 도입하며, 기본(base), OT, 확률적 흐름 매칭을 사용해 SE(3)^N의 원천 분포와 목표 분포를 매핑하고, 빠르고 안정적인 학습과 높은 품질의 설계를 달성한다.

ABSTRACT

The computational design of novel protein structures has the potential to impact numerous scientific disciplines greatly. Toward this goal, we introduce FoldFlow, a series of novel generative models of increasing modeling power based on the flow-matching paradigm over $3\mathrm{D}$ rigid motions -- i.e. the group $ ext{SE}(3)$ -- enabling accurate modeling of protein backbones. We first introduce FoldFlow-Base, a simulation-free approach to learning deterministic continuous-time dynamics and matching invariant target distributions on $ ext{SE}(3)$. We next accelerate training by incorporating Riemannian optimal transport to create FoldFlow-OT, leading to the construction of both more simple and stable flows. Finally, we design FoldFlow-SFM, coupling both Riemannian OT and simulation-free training to learn stochastic continuous-time dynamics over $ ext{SE}(3)$. Our family of FoldFlow, generative models offers several key advantages over previous approaches to the generative modeling of proteins: they are more stable and faster to train than diffusion-based approaches, and our models enjoy the ability to map any invariant source distribution to any invariant target distribution over $ ext{SE}(3)$. Empirically, we validate FoldFlow, on protein backbone generation of up to $300$ amino acids leading to high-quality designable, diverse, and novel samples.

연구 동기 및 목표

• 단백질 설계를 SE(3)^N 분포에서 샘플링하는 것으로 동기를 부여하고 불변 생성이 가능하도록 한다.
• SE(3)에서 시뮬레이션 없이 동작하는 흐름 모델을 개발하여 임의의 불변 소스에서 타깃 분포로 매핑한다.
• 리만 계 OT 및 확률적 브리지로 학습 안정성과 속도를 향상시킨다.
• 백본 생성을 위한 신경망 백본을 포함하는 엔드 투 엔드 FoldFlow 아키텍처를 제공한다.
• 최대 300 잔기로 구성된 백본에서 최신 방법과 비교하여 경험적으로 검증한다.

제안 방법

• FoldFlow-Base 개발: 학습을 위한 닫힌 형태의 로그 맵을 갖는 SO(3)에서의 시뮬레이션 없는 조건부 흐름 매칭.
• 리만 OT를 도입하여 더 짧고 더 곧은 측지 보간을 얻도록 FoldFlow-OT로 확장한다.
• FoldFlow-SFM을 도입하여 SO(3)에서 Brownian/열 다리의 시뮬레이션 없이 근사를 사용해 확률적 SE(3) 흐름을 학습한다.
• SE(3)^N을 SO(3)과 R^3 성분으로 분해하고, SE(3) 불변성을 갖춘 잔기별 흐름을 학습한다.
• 각 잔당에 대해 AF2와 같은 프레임으로 단백질 백본을 매개화하고, x_t에서 x_0로 이동시키는 속도 예측기 v_theta를 학습한다.
• SO(3)와 R^3에서의 흐름 매칭 손실을 보조 원자 수준 손실과 함께 결합하여 강건한 백본 설계를 달성한다.

실험 결과

연구 질문

• RQ1 SE(3)^N 불변 흐름이 단백질 백본의 임의의 소스와 타깃 분포 간 매핑이 가능한가?
• RQ2 리만 OT 기반 및 확률적 흐름 변형이 기본 흐름 대비 학습 안정성 및 설계 품질을 개선하는가?
• RQ3 FoldFlow 변형들이 설계성, 다양성, 새로움 측면에서 비사전학습(state-of-the-art) 방법과 비교하여 어떤 성능을 보이는가?
• RQ4 SE(3) 백본에 확률적 브리지를 도입할 때 시뮬레이션 없는 학습 목표가 충분한가?
• RQ5 FoldFlow-Base, FoldFlow-OT, FoldFlow-SFM 간의 정확도와 속도 간 컴퓨팅 트레이드오프는 무엇인가?

주요 결과

• FoldFlow-OT 및 FoldFlow-SFM은 합성 SO(3) 데이터에서 Wasserstein 거리 측정에서 FoldFlow-Base보다 더 낮은 값을 보인다.
• 모든 FoldFlow 변형은 프레임디프-임프로브드(FrameDiff-Improved) 비사전학습 baselines보다 디자인성, 다양성, 새로움 면에서 우수하다.
• FoldFlow-SFM은 시뮬레이션 없이 학습 가능한 프레임워크를 제공하며, 무조건부 손실에 대한 그래디언트 등가성으로 확률적 SE(3) 흐름을 검증한다.
• FoldFlow-OT는 OT 기반 보간에 의해 더 안정적이고 빠른 학습을 제공한다.
• 최대 300 잔기로 구성된 백본 설계에서 FoldFlow 변형들은 고품질, 설계 가능하고, 다양하며, 새로움이 있는 샘플을 생성한다.