[论文解读] SE(3)-Stochastic Flow Matching for Protein Backbone Generation

Avishek Joey Bose, Tara Akhound-Sadegh+8 more|arXiv (Cornell University)|Oct 3, 2023

Protein purification and stability被引用 14

一句话总结

FoldFlow 引入了 SE(3) 不变的正则化流，用于蛋白质骨架生成，结合基础流、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) 上的无仿真条件流匹配，具备用于训练的闭式对数映射。
• 扩展为 FoldFlow-OT，结合黎曼 OT 以获得更短、更直的测地线插值。
• 引入 FoldFlow-SFM，利用对 SO(3) 上布朗运动/热桥的无仿真近似来学习随机 SE(3) 流。
• 将 SE(3)^N 分解为 SO(3) 与 R^3 分量，并对每个残基训练具有 SE(3) 不变性的流。
• 通过每个残基的类似 AF2 的框架对蛋白质骨架进行参数化，并训练速度预测器 v_theta，将从 x_t 推进到 x_0。
• 将 SO(3) 与 R^3 上的流匹配损失与原子级辅助损失结合，以实现稳健的骨架设计。

实验结果

研究问题

• RQ1 SE(3)^N 不变流是否能够在蛋白质骨架的任意源分布和目标分布之间映射？
• RQ2 基于黎曼 OT 的和随机流变体是否比基础流在训练稳定性和设计质量上有改进？
• RQ3 与未预训练的最先进方法相比，FoldFlow 变体在设计性、多样性和新颖性方面的表现如何？
• RQ4 在为 SE(3) 骨架引入随机桥时，无仿真训练目标是否足够？
• RQ5 在 FoldFlow-Base、FoldFlow-OT 和 FoldFlow-SFM 之间，精度与速度的计算权衡是什么？

主要发现

• FoldFlow-OT 与 FoldFlow-SFM 在合成的 SO(3) 数据上达到低于 FoldFlow-Base 的 Wasserstein 距离。
• 在 PDB 骨架设计任务中，所有 FoldFlow 变体在设计性、多样性和新颖性方面都优于未预训练的 FrameDiff-Improved 基线。
• FoldFlow-SFM 提供一个无仿真训练框架，具有随机 SE(3) 流，通过与无条件损失的梯度等价性得到验证。
• FoldFlow-OT 由于基于 OT 的插值而实现更稳定和更快的训练。
• 在最多 300 残基的骨架设计中，FoldFlow 变体产生高质量、可设计、具有多样性和新颖性的样本。