Sequence-Augmented SE(3)-Flow Matching For Conditional Protein Backbone Generation

📄 arXiv: 2405.20313v2 📥 PDF

作者: Guillaume Huguet, James Vuckovic, Kilian Fatras, Eric Thibodeau-Laufer, Pablo Lemos, Riashat Islam, Cheng-Hao Liu, Jarrid Rector-Brooks, Tara Akhound-Sadegh, Michael Bronstein, Alexander Tong, Avishek Joey Bose

分类: cs.LG, q-bio.BM

发布日期: 2024-05-30 (更新: 2024-12-11)

备注: Presented at NeurIPS 2024

💡 一句话要点

FoldFlow-2：序列增强的SE(3)-Flow Matching用于条件蛋白质骨架生成

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 蛋白质结构生成 SE(3)-Flow Matching 蛋白质大型语言模型 多模态融合 几何Transformer 条件生成 药物设计

📋 核心要点

  1. 现有蛋白质结构生成模型在多样性和新颖性方面存在不足，限制了其在药物设计等领域的应用。
  2. FoldFlow-2利用蛋白质序列信息，结合SE(3)-等变Flow Matching，提升蛋白质结构生成的多样性和设计性。
  3. FoldFlow-2在设计性、多样性和新颖性等指标上超越RFDiffusion，并在纳米抗体支架设计等任务中取得进展。

📝 摘要（中文）

本文提出FoldFlow-2，一种新型序列条件SE(3)-等变Flow Matching模型，用于蛋白质结构生成，充分利用了氨基酸序列丰富的生物学归纳偏置。FoldFlow-2在之前的FoldFlow模型基础上引入了显著的新架构特征，包括用于编码序列的蛋白质大型语言模型、结合结构和序列表示的新型多模态融合主干网络，以及基于几何Transformer的解码器。为了增加生成样本的多样性和新颖性（这对于从头药物设计至关重要），我们在一个比以往PDB数据集大一个数量级的新数据集上大规模训练FoldFlow-2，该数据集包含PDB中的已知蛋白质和通过过滤获得的高质量合成结构。我们进一步通过引入强化微调(ReFT)目标，展示了FoldFlow-2与任意奖励对齐的能力，例如增加二级结构多样性。实验结果表明，FoldFlow-2在无条件生成方面优于之前的最先进的基于蛋白质结构的生成模型RFDiffusion，在所有蛋白质长度上的设计性、多样性和新颖性等所有指标上均有所提高，并且在平衡构象采样任务中表现出泛化能力。最后，我们证明了微调后的FoldFlow-2在具有挑战性的条件设计任务（如设计VHH纳米抗体的支架）方面取得了进展。

🔬 方法详解

问题定义：论文旨在解决蛋白质结构生成问题，特别是提高生成蛋白质结构的多样性和新颖性，以满足从头药物设计等应用的需求。现有方法，如RFDiffusion，在生成蛋白质结构的多样性和设计性方面存在局限性，无法充分利用蛋白质序列信息。

核心思路：论文的核心思路是将蛋白质序列信息融入到SE(3)-等变Flow Matching模型中，利用蛋白质序列的生物学归纳偏置来指导蛋白质结构的生成。通过结合蛋白质大型语言模型、多模态融合主干网络和几何Transformer解码器，FoldFlow-2能够更好地利用序列信息，生成更多样、更具设计性的蛋白质结构。

技术框架：FoldFlow-2的整体架构包括以下几个主要模块：1) 蛋白质大型语言模型：用于编码蛋白质序列信息。2) 多模态融合主干网络：将蛋白质结构和序列表示进行融合。3) 几何Transformer解码器：基于融合后的表示生成蛋白质结构。整个流程基于SE(3)-等变Flow Matching框架，保证生成的蛋白质结构在三维空间中的旋转和平移不变性。

关键创新：FoldFlow-2的关键创新在于：1) 引入蛋白质大型语言模型来编码序列信息，充分利用了蛋白质序列的生物学归纳偏置。2) 设计了新的多模态融合主干网络，能够有效地结合结构和序列表示。3) 提出了强化微调(ReFT)目标，能够将模型与任意奖励对齐，例如增加二级结构多样性。

关键设计：FoldFlow-2的关键设计包括：1) 使用预训练的蛋白质大型语言模型来初始化序列编码器。2) 多模态融合主干网络采用注意力机制来融合结构和序列表示。3) 几何Transformer解码器采用SE(3)-等变层来保证旋转和平移不变性。4) ReFT目标使用策略梯度方法来优化模型，使其能够生成具有更高奖励的蛋白质结构。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

FoldFlow-2在无条件生成方面优于RFDiffusion，在设计性、多样性和新颖性等指标上均有所提高。例如，在蛋白质长度范围内，FoldFlow-2在设计性方面取得了显著提升。此外，FoldFlow-2在平衡构象采样任务中表现出泛化能力，并在纳米抗体支架设计等具有挑战性的条件设计任务中取得了进展。这些结果表明FoldFlow-2是一种有效的蛋白质结构生成模型。

🎯 应用场景

FoldFlow-2在药物设计、蛋白质工程等领域具有广泛的应用前景。它可以用于从头设计具有特定功能的蛋白质，例如设计新型药物靶点或优化现有蛋白质的性能。通过与奖励函数对齐，FoldFlow-2还可以用于设计具有特定性质的蛋白质，例如具有高稳定性的蛋白质或具有特定结合能力的蛋白质。该研究的成果将加速蛋白质设计和药物发现的进程。

📄 摘要（原文）

Proteins are essential for almost all biological processes and derive their diverse functions from complex 3D structures, which are in turn determined by their amino acid sequences. In this paper, we exploit the rich biological inductive bias of amino acid sequences and introduce FoldFlow-2, a novel sequence-conditioned SE(3)-equivariant flow matching model for protein structure generation. FoldFlow-2 presents substantial new architectural features over the previous FoldFlow family of models including a protein large language model to encode sequence, a new multi-modal fusion trunk that combines structure and sequence representations, and a geometric transformer based decoder. To increase diversity and novelty of generated samples -- crucial for de-novo drug design -- we train FoldFlow-2 at scale on a new dataset that is an order of magnitude larger than PDB datasets of prior works, containing both known proteins in PDB and high-quality synthetic structures achieved through filtering. We further demonstrate the ability to align FoldFlow-2 to arbitrary rewards, e.g. increasing secondary structures diversity, by introducing a Reinforced Finetuning (ReFT) objective. We empirically observe that FoldFlow-2 outperforms previous state-of-the-art protein structure-based generative models, improving over RFDiffusion in terms of unconditional generation across all metrics including designability, diversity, and novelty across all protein lengths, as well as exhibiting generalization on the task of equilibrium conformation sampling. Finally, we demonstrate that a fine-tuned FoldFlow-2 makes progress on challenging conditional design tasks such as designing scaffolds for the VHH nanobody.