[论文解读] Fast protein backbone generation with SE(3) flow matching

Jason Yim, Andrew M. Campbell+9 more|arXiv (Cornell University)|Oct 8, 2023

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) 上的流匹配进行新生蛋白质设计的更快、更高效的生成建模。
• 将 SE(3) 流匹配适应于主链生成，并训练一个等变向量场来回归干净的框架。
• 证明 FrameFlow 在较少采样步骤的情况下实现与基于扩散的基线相当或更优的可设计性。
• 探索架构与训练方面的调整（等变性、先验、预对齐）以最大化性能和稳定性。

提出的方法

• 在蛋白质主链上应用 SE(3) 流匹配，使用具有加法度量的乘积空间 SE(3)^N。
• 使用条件向量场学习来匹配条件流，从而通过微分方程积分实现一个 CNF。
• 使用适当的测地线对平移和旋转进行参数化（平移采用线性插值；旋转采用罗德里格斯基的指数/对数）。
• 使用 SE(3) 等变 FramePred 架构和 IPA 更新进行训练，以确保 SE(3) 等变性。
• 使用 Kabsch 的预对齐以提升训练效率，并引入 IGSO3 先验以改善旋转采样。
• 在采样时使用带指数斜坡的 SO(3) 推断调度器，并调整损失权重以稳定训练。
• 在 SCOPe 数据上进行评估，比较与 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 相比，FrameFlow 使用 100 时步在设计性方面更高（0.77 对 GENIE SDE 的 1000 时步的 0.22），在更少的步数下实现类似或更好的设计性。
• 采样时间显著减少：在长度为 100 的蛋白质上，GENIE 使用 1000 时步需要 128 秒，而 FrameDiff/FrameFlow 在 100 时步仅需 5.7 秒，且硬件相近。
• FrameFlow 与 FrameDiff 使用相同的主干架构；FrameFlow 的基于 ODE 的采样显示出更直的轨迹和比 FrameDiff 的基于 SDE 的采样更快的收敛。