论文解读：基于结构条件类别扩散的蛋白质反向折叠强化学习

论文标题: REINFORCEMENT LEARNING ON STRUCTURE-CONDITIONED CATEGORICAL DIFFUSION FOR PROTEIN INVERSE FOLDING

论文链接: arXiv:2410.17173v1 [cs.AI]

代码链接: https://github.com/flagshippioneering/pi-rldif

1. 背景介绍

蛋白质反向折叠 (Protein Inverse Folding, IF) 是生物信息学和蛋白质工程领域中的一个核心问题。它的目标是：给定一个目标蛋白质的三维 (3D) 结构，预测能够折叠成该结构的氨基酸序列。 这与蛋白质结构预测问题相反，结构预测是从氨基酸序列预测蛋白质结构。

专业术语解释:

传统的蛋白质反向折叠方法通常以序列回复率为优化目标，即模型试图尽可能还原天然蛋白质的氨基酸序列。然而，反向折叠问题本质上是“一对多”的映射，即一个给定的蛋白质结构可以由多个不同的氨基酸序列折叠而成。在许多实际应用中，例如药物设计、生物材料和合成生物学等领域，拥有多样化的、能够折叠成目标结构的序列集合是非常有益的。这样可以为后续的优化步骤 (如提高稳定性、防止聚集、降低免疫原性等) 提供更广泛的选择空间。

然而，提高序列多样性的方法 (例如，提高自回归模型的采样温度) 可能会降低生成序列的结构一致性。因此，一个有用的蛋白质反向折叠模型的关键特性是其 “可折叠多样性”，即在保证结构一致性的前提下，生成多样化序列的能力。

2. 核心概述

本文提出了一种新的蛋白质反向折叠模型 RL-DIF (Reinforcement Learning on Structure-Conditioned Categorical Diffusion)。该模型基于类别扩散模型，首先通过序列回复率进行预训练，然后通过强化学习 (Reinforcement Learning, RL) 进行微调，以优化结构一致性。实验结果表明，RL-DIF 在保持与现有最佳模型相当的序列回复率和结构一致性的同时，显著提高了可折叠多样性。在 CATH 4.2 数据集上的实验显示，RL-DIF 的可折叠多样性达到了 29%，而基于相同数据集训练的其他模型仅为 23%。该研究强调了可折叠多样性作为评估反向折叠模型性能的新指标的重要性，并证明了通过强化学习微调扩散模型可以有效提高蛋白质序列设计的质量和多样性。

3. 方法与实验细节

3.1 数据集

本文使用了以下数据集进行模型训练和评估：

文章在附录 A.2 中分析了这些基准数据集与 CATH 4.2 训练集的结构和序列相似性，发现 TS50, TS500 和 CASP15 与训练集存在一定程度的重叠 (42%-84%)。尽管如此，为了与先前的工作保持一致，论文仍然使用了这些数据集进行评估，并将重点放在 CATH 4.2 数据集上的比较。

3.2 算法和模型

3.2.1 RL-DIF 模型架构

RL-DIF 模型的核心是类别扩散模型，其灵感来源于 GradeIF 和 PiFold 模型。 模型架构基于改进的 PiFold 结构，并添加了多层感知器 (MLP) 来处理部分加噪的氨基酸序列和扩散时间步长。

具体来说，给定蛋白质骨架坐标 $X \in R^{4N \times 3}$，首先构建 kNN 图 (k=30)，然后使用 PiFold 的特征提取器提取节点和边的特征 $h_V$ 和 $h_E$。 这些特征包括原子之间的距离、二面角和方向向量。去噪模型是一个关于 $h_V$, $h_E$，部分去噪的序列 $s_t$ 和时间步长 $t$ 的函数。 模型架构如下:

其中 $p(S_{t+1}|S_t)$ 代表在给定当前去噪步骤 $S_t$ 的情况下，预测下一步 $S_{t+1}$ 的概率分布。 [a, b] 表示连接操作， PiGNN 是 PiFold 模型中引入的图神经网络层。

3.2.2 训练损失函数和强化学习

预训练阶段 (Diffusion Model Pre-training):

RL-DIF 模型首先使用离散去噪扩散概率模型 (D3PM) 进行预训练，目标是序列回复率。 论文使用了 GradeIF 提出的条件离散去噪扩散模型，将蛋白质反向折叠过程建模为马尔可夫扩散过程。前向扩散过程定义如下：

$S_t \sim q(S_t|S_{t-1}, S_0) = Cat(S_t; p = S_{t-1}Q_t)$ (公式 1)

其中 $Q_1, ..., Q_T$ 是一系列 $|V| \times |V|$ 的转移矩阵， $S_0 = S$ 是天然序列， $Cat$ 表示类别分布。 反向扩散过程需要学习一个模型 $p_\theta(S_{t-1}|S_t; X)$ 来逐步去噪，从而生成新的氨基酸序列。 论文使用了 D3PM 混合损失函数进行训练，并使用了均匀转移矩阵。

强化学习微调阶段 (Reinforcement Learning Fine-tuning):

预训练后的扩散模型通过强化学习进行微调，目标是优化结构一致性。 论文使用了去噪扩散策略优化 (DDPO) 算法。DDPO 将反向去噪过程视为一个 T 步马尔可夫决策过程，并定义策略梯度来最大化期望奖励 $J(\theta)$:

$J(\theta) = E_{X \sim p(X), \hat{S} \sim p_\theta(\hat{S}|X)} [R(\hat{S})]$ (公式 3)

其中 $R(\hat{S})$ 是奖励函数，本文中使用的是自洽 TM-score (sc-TM)。策略梯度 $\nabla_\theta J(\theta)$ 的计算公式如下：

$\nabla_\theta J(\theta) = E_{X \sim p(X), \hat{S}0, ..., \hat{S}T \sim p{old}} [\sum{t=1}^T \frac{p_\theta(S_{t-1}|S_t, X)}{p_{old}(S_{t-1}|S_t, X)} \nabla_\theta log p_\theta(S_{t-1}|S_t, X) R(\hat{S}_0, X)]$ (公式 4)

其中 $\frac{p_\theta(S_{t-1}|S_t, X)}{p_{old}(S_{t-1}|S_t, X)}$ 是重要性采样比率，允许每个样本进行多次优化迭代。

3.3 训练和评估过程

预训练阶段:

强化学习微调阶段:

蛋白质结构预测:

在强化学习微调阶段，以及模型评估过程中，需要预测氨基酸序列的 3D 结构。为了平衡速度和精度，论文使用了 ESMFold 而不是 AlphaFold2 进行结构预测。 使用 Huggingface Transformers 库中的 ESMFold 实现，并在 Kubernetes 集群上部署，包含 20 个 Nvidia A10 GPU 和负载均衡器，以实现高效的在线策略训练。

模型评估指标:

模型采样策略:

4. 研究过程与结论

本文的研究过程主要包括以下几个方面：

  1. 可折叠多样性指标的提出: 论文首先分析了现有蛋白质反向折叠模型在序列多样性和结构一致性之间的权衡，并提出了 “可折叠多样性” (Foldable Diversity, FD) 这一新的评估指标。 FD 指标能够更有效地衡量在保证结构正确性的前提下，模型生成多样化序列的能力，更符合实际应用的需求。

  2. RL-DIF 模型的构建与训练: 为了提高可折叠多样性，论文提出了 RL-DIF 模型。该模型结合了类别扩散模型和强化学习，通过两阶段训练策略：

    • 预训练阶段: 使用扩散模型进行序列回复率预训练，学习生成氨基酸序列的基本能力。
    • 强化学习微调阶段: 使用 DDPO 算法，以结构一致性 (自洽 TM-score) 为奖励信号，对预训练的扩散模型进行微调，优化模型生成具有正确结构的序列的能力。

  3. 实验评估与结果分析: 论文在 CATH 4.2, TS50, TS500 和 CASP15 数据集上对 RL-DIF 模型进行了全面的实验评估，并与现有的 SOTA 模型 (ProteinMPNN, PiFold, KWDesign, DIF-Only, ESM-IF) 进行了比较。 实验结果主要体现在 Table 1, Table 2 和 Figure 2 中:

    • 可折叠多样性显著提升: 实验结果表明，RL-DIF 模型在所有基准数据集上都实现了最高的或接近最高的可折叠多样性。尤其在 CATH-all 数据集 (包含多蛋白复合物) 上，RL-DIF 的可折叠多样性达到了 29%，显著高于其他模型 (最高 23%)。 这表明 RL-DIF 模型能够有效生成更多样化的、同时保持结构正确的蛋白质序列。

    • 结构一致性保持竞争力: 尽管侧重于提高可折叠多样性，RL-DIF 模型在结构一致性 (sc-TM) 方面仍然保持了与 SOTA 模型相当的水平。在大多数数据集上，RL-DIF 的 sc-TM 略低于一些模型，但差距不大，且可折叠多样性的提升幅度远大于结构一致性的轻微下降。

    • 序列回复率适中: RL-DIF 的序列回复率 相对较低，但作者认为序列回复率并非评估反向折叠模型性能的最佳指标，因为它倾向于生成与天然序列高度相似的序列，限制了设计空间的多样性。论文更强调可折叠多样性和结构一致性。

    • 消融实验: Table 3 展示了强化学习微调对 RL-DIF 性能的影响。结果表明，RL 微调能够有效提高结构一致性，但在一定程度上会降低可折叠多样性。论文选择 RL 训练 1000 步，以平衡结构一致性和可折叠多样性。 Figure 2 展示了可折叠多样性指标对 TMmin 阈值的敏感性分析，结果表明 RL-DIF 在不同的 TMmin 阈值下都表现出最佳的可折叠多样性。

    • 与 ESM-IF 的对比: Table 2 对比了 RL-DIF-100K (使用更大规模数据集预训练的 RL-DIF) 和 ESM-IF 的性能。 结果表明，RL-DIF-100K 在 TS50 数据集上超越了 ESM-IF，在其他数据集上也接近 ESM-IF 的性能，尽管 RL-DIF-100K 使用的训练数据量和参数量远小于 ESM-IF。 这表明 RL-DIF 模型具有更高的数据效率。

  4. 示例序列展示: Figure 1 展示了 RL-DIF, ProteinMPNN 和 PiFold 模型在相同蛋白质骨架上生成的示例序列。 颜色编码显示了每个位置氨基酸的多样性，结果表明 RL-DIF 生成的序列具有最高的氨基酸多样性，同时保持了良好的结构一致性。

结论: 本文提出的 RL-DIF 模型通过结合类别扩散模型和强化学习，有效地提高了蛋白质反向折叠模型的可折叠多样性，同时保持了良好的结构一致性。 可折叠多样性被证明是一个更符合实际应用需求的评估指标。 强化学习微调是提高结构一致性的有效手段，但需要在多样性和一致性之间进行权衡。 RL-DIF 模型为蛋白质序列设计提供了一种新的有效方法，能够生成更多样化的、高质量的候选序列。

5. 总结与客观评价

总结: 本文创新性地提出了 RL-DIF 模型，该模型巧妙地融合了类别扩散模型和强化学习，用于解决蛋白质反向折叠问题。通过两阶段训练策略，RL-DIF 模型在保证结构一致性的前提下，显著提升了生成蛋白质序列的可折叠多样性。 论文还强调了可折叠多样性作为评估反向折叠模型的新指标的价值，并进行了全面的实验验证，证明了 RL-DIF 模型在性能上的优势。

客观评价:

总而言之，本文工作具有重要的学术价值和应用前景。RL-DIF 模型代表了蛋白质反向折叠领域的一个重要进展，为生物医药、生物材料等领域的蛋白质设计提供了更有力的工具。 可折叠多样性指标的提出也为未来评估和改进反向折叠模型提供了新的视角。

6. 参考文献与链接