[논문 리뷰] Fast protein backbone generation with SE(3) flow matching

Jason Yim, Andrew M. Campbell외 9명|arXiv (Cornell University)|2023. 10. 08.

Protein Structure and Dynamics인용 수 21

한 줄 요약

FrameFlow는 SE(3) 흐름 매칭을 단백질 백본 생성에 적용하여 FrameDiff에 비해 훨씬 적은 샘플링 타임스텝으로 비슷한 설계가능성을 달성하고, 설계가능성에서 GENIE를 능가합니다. 다양성 및 참신성 특성을 유지하면서 샘플링 속도를 크게 높입니다.

ABSTRACT

We present FrameFlow, a method for fast protein backbone generation using SE(3) flow matching. Specifically, we adapt FrameDiff, a state-of-the-art diffusion model, to the flow-matching generative modeling paradigm. We show how flow matching can be applied on SE(3) and propose modifications during training to effectively learn the vector field. Compared to FrameDiff, FrameFlow requires five times fewer sampling timesteps while achieving two fold better designability. The ability to generate high quality protein samples at a fraction of the cost of previous methods paves the way towards more efficient generative models in de novo protein design.

연구 동기 및 목표

• SE(3)에서의 흐름 매칭을 이용한 de novo 단백질 설계를 위한 더 빠르고 효율적인 생성 모델링을 제시한다.
• 백본 생성에 SE(3) 흐름 매칭을 적용하고 깨끗한 프레임을 회귀하도록 등변 벡터장(equivariant vector field)을 학습한다.
• FrameFlow가 확산 기반 기준선 대비 더 적은 샘플링 단계로 비슷하거나 더 우수한 설계가능성을 달성함을 보여준다.
• 성능과 안정성을 극대화하기 위한 아키텍처 및 학습 조정(등변성, 프라이어, 사전 정렬)을 탐구한다.

제안 방법

• 덧셈 메트릭스를 갖는 곱공간 SE(3)^N을 사용하여 단백질 백본에 SE(3) 흐름 매칭을 적용한다.
• 조건부 벡터장 학습을 사용하여 조건부 흐름을 맞추고 ODE 적분을 통한 CNF를 가능하게 한다.
• 이동은 선형 보간(translation), 회전은 Rodrigues 기반의 exp/log를 사용하여 적절한 지오데시스로 번역 및 회전을 매개화한다.
• SE(3)-등변 FramePred 아키텍처와 SE(3) 등가성 확보를 위한 IPA 업데이트로 학습한다.
• 학습 효율성을 높이기 위해 Kabsch로 사전 정렬(pre-alignment)을 사용하고 회전 샘플링을 개선하기 위해 IGSO3 프라이어를 적용한다.
• 샘플링 중 지수적 램프를 갖는 SO(3) 추론 일정표를 사용하고 학습 안정화를 위해 손실 가중치를 조정한다.
• SCOPe 데이터에서 평가하고 DESIGNability, 다양성, 참신성을 GENIE 및 FrameDiff와 비교하며 길이당 10개의 샘플을 사용하고 ProteinMPNN/ESMFold 설계가능성 점검을 수행한다.

실험 결과

연구 질문

• RQ1 SE(3) 흐름 매칭으로 확산 모델보다 훨씬 적은 샘플링 단계로 고품질의 단백질 백본을 생성할 수 있는가?
• RQ2 FrameFlow가 FrameDiff 및 GENIE에 비해 다양성이나 참신성을 희생하지 않고 향상된 설계가능성을 제공하는가?
• RQ3 단백질 백본을 위한 SE(3) 흐름 매칭을 가장 효과적으로 안정화하는 아키텍처 및 학습 조정은 무엇인가?
• RQ4 다른 프라이어와 샘플링 스케줄러가 단백질의 SE(3) 흐름 매칭에서 성능과 효율에 어떤 영향을 미치는가?

주요 결과

• FrameFlow는 500 타임스텝에서 0.81의 설계가능성을, 100 타임스텝에서 0.77의 설계가능성을 달성하여 해당 설정에서 FrameDiff 및 GENIE를 능가한다.
• FrameFlow는 여러 타임스텝에서 0.22–0.28(클러스터) 수준의 다양성 점수를 얻어, 여러 설정에서 FrameDiff보다 더 많은 설계 가능한 샘플과 더 많은 모드를 보인다.
• FrameFlow는 0.69–0.67의 참신도(낮을수록 좋음)에 도달하여 높은 설계가능성을 유지하면서 알려진 구조에 대한 유사성을 경쟁적으로 나타낸다.
• GENIE와 비교했을 때, 100 타임스텝의 FrameFlow는 설계가능성이 더 높다(0.77 대 GENIE SDE의 1000 타임스텝 0.22) 및 더 적은 단계에서도 비슷하거나 더 나은 설계가능성을 보인다.
• 샘플링 시간은 크게 단축된다: 길이-100 단백질에서 1000 타임스텝의 GENIE는 128초가 걸리는 반면, 100 타임스텝의 FrameDiff/FrameFlow는 유사한 하드웨어에서 5.7초에 수행된다.
• FrameFlow와 FrameDiff는 동일한 백본 아키텍처를 사용한다; FrameFlow의 ODE 기반 샘플링은 FrameDiff의 SDE 기반 샘플링보다 궤적이 더 직선적이고 수렴이 더 빠르다.