侧链条件化和建模用于全原子蛋白质序列设计 - FAMPNN

1. 背景介绍

蛋白质是生命活动的重要承担者，其三维结构决定了其生物学功能。蛋白质序列设计的目标是为给定的蛋白质骨架结构设计出能够稳定折叠并具有特定功能的氨基酸序列。传统的蛋白质序列设计方法，如基于物理的方法，通常依赖于能量函数来优化序列和侧链构象，以找到能量最低的配置。近年来，深度学习方法在蛋白质序列设计领域取得了显著的成功，尤其是基于固定骨架的序列设计。

关键术语解释:

• 蛋白质序列设计 (Protein Sequence Design): 在给定的蛋白质骨架结构下，设计出具有特定功能或稳定性的氨基酸序列的过程。
• 固定骨架序列设计 (Fixed-backbone Sequence Design): 蛋白质序列设计的一种特定形式，其中蛋白质的骨架结构是预先确定的，而目标是优化氨基酸序列。
• 侧链 (Sidechain): 氨基酸分子中，除了主链 (backbone) 之外的部分，也称为 R 基团。侧链的种类决定了氨基酸的类型，并且侧链原子的三维排列对于蛋白质的构象、稳定性和功能至关重要。
• 骨架 (Backbone): 蛋白质多肽链中重复的原子组分 (-N-Cα-C-)。骨架结构主要由 φ 和 ψ 二面角决定。
• 构象 (Conformation): 分子由于原子基团绕化学键旋转而呈现出的不同空间排列。对于蛋白质而言，构象通常指其三维结构。
• 能量函数 (Energy Function): 用于评估蛋白质构象稳定性的数学函数，通常基于物理化学原理，例如范德华力、静电相互作用等。
• 深度学习 (Deep Learning): 机器学习的一个分支，使用具有多层结构的神经网络进行数据建模和分析。在蛋白质科学中，深度学习被广泛应用于结构预测、序列设计、功能预测等任务。
• 图神经网络 (Graph Neural Network, GNN): 一种用于处理图结构数据的神经网络。在蛋白质领域，蛋白质可以被表示为图，其中氨基酸残基是节点，残基之间的相互作用是边。GNN 可以有效地学习蛋白质的结构特征。
• 掩码语言模型 (Masked Language Model): 一种自然语言处理中的模型训练方法，通过随机掩盖输入文本中的部分词语，并让模型预测被掩盖的词语，从而学习语言的上下文信息。在蛋白质序列设计中，可以借鉴掩码语言模型的思想来训练模型。
• 扩散模型 (Diffusion Model): 一类生成模型，通过逐步向数据中添加噪声，然后再学习逆向去噪过程来生成数据。在蛋白质领域，扩散模型可以用于生成蛋白质的构象和序列。
• 交叉熵损失 (Cross-entropy Loss): 一种用于分类任务的损失函数，衡量模型预测的概率分布与真实概率分布之间的差异。
• 扩散损失 (Diffusion Loss): 用于训练扩散模型的损失函数，通常基于均方误差 (MSE) 或其他距离度量，衡量模型去噪能力。

2. 核心概述

本文提出了 FAMPNN (Full-Atom MPNN)，一种用于蛋白质序列设计的新方法，该方法显式地建模了每个残基的序列身份和侧链构象。FAMPNN 模型学习残基离散的氨基酸身份和连续的侧链构象的联合分布，并使用结合了分类交叉熵损失和扩散损失的目标函数进行训练。论文证明了联合学习这两个分布是高度协同的任务，能够提高序列恢复率并实现最先进的侧链堆积性能。此外，显式全原子建模的优势可以推广到序列恢复之外的实际蛋白质设计应用，例如零样本预测实验结合亲和力和稳定性测量。该方法通过图神经网络 (GNN) 编码骨架结构，并迭代地生成序列和侧链，最终实现全原子蛋白质序列设计。

3. 方法论和实验细节

3.1. 数据集

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

• CATH 4.2 S40 数据集: 用于训练和评估序列恢复性能。该数据集是 CATH 数据库 4.2 版本的 S40 子集，包含冗余度低于 40% 的蛋白质结构域，并划分了训练集、验证集和测试集。数据集的划分方式与 Ingraham 等人 (2019) 的工作相同。
• PDB 数据集: 用于训练能够处理更复杂蛋白质结构的 FAMPNN 模型。该数据集基于整个蛋白质数据库 (PDB)，截至日期为 2021 年 9 月 30 日。数据集在链级别上进行聚类，保证了蛋白质结构的多样性，并移除了单链蛋白质在多链结构中的情况，确保模型学习到正确的上下文信息。
• CASP 数据集 (CASP13, CASP14, CASP15): 用于评估侧链堆积性能。论文使用了 MMseqs2 easy-search 工具从训练集和验证集中移除了与 CASP13-15 同源的序列（相似度cutoff为40%）。CASP13 和 CASP14 的测试集从 AttnPacker 的 GitHub 仓库获得，CASP15 的目标从 CASP 数据档案下载。
• RFdiffusion 生成的 de novo 骨架: 用于评估自洽性 (self-consistency)。论文使用了 RFdiffusion 生成了长度为 100 到 500 的 de novo 骨架，每个长度生成 100 个骨架。
• SKEMPIv2, Megascale, FireProtDB, S669, 抗体-抗原结合亲和力数据集: 用于评估蛋白质适应性 (fitness) 预测能力。这些数据集包含了实验测量的蛋白质突变体结合亲和力或稳定性变化的数据。详细的数据集处理过程在附录 G.2 中描述。

3.2. 算法和模型

FAMPNN 模型的核心架构是混合了 MPNN (Message Passing Neural Network) 和 GVP (Geometric Vector Perceptron) 的图神经网络。模型主要由以下三个组件构成：

1. 不变骨架编码器 (Invariant Backbone Encoder): 与 ProteinMPNN 的编码器相同，用于编码蛋白质的骨架结构。输入是蛋白质的骨架原子坐标，输出是节点的表示和边的表示。
2. 不变全原子编码器 (Invariant Full-Atom Encoder): 替换了 ProteinMPNN 的序列解码器，与骨架编码器结构相同，但扩展了特征表示，能够处理所有原子。输入除了骨架结构外，还包括序列信息。
3. 等变全原子编码器 (Equivariant Full-Atom Encoder): 使用几何向量感知器 (GVP) 层，学习向量特征和标量特征。GVP 层包含等变轨道 (equivariant track) 和不变轨道 (invariant track)。等变轨道用于学习向量特征，不变轨道用于学习标量特征。FAMPNN 使用等变轨道编码 Cα 到残基 i 中所有其他原子的单位向量，以及 Cαi 到残基 j 中所有原子的单位向量作为边特征。不变轨道则结合了全原子编码器中的距离信息。

训练损失函数:

FAMPNN 的训练目标是联合预测序列身份和侧链构象。因此，总损失函数由两部分组成：

• 序列预测的交叉熵损失 (L_MLM): 标准的掩码语言模型损失，用于训练模型预测被掩盖的氨基酸序列。
• 侧链构象预测的扩散损失 (L_diff): 使用欧几里得扩散损失，用于训练模型生成侧链构象。扩散过程基于方差爆炸 EDM 方案，模型学习一个噪声条件化的去噪器 D_θ，以最小化高斯噪声版本侧链坐标的 L2 误差。

总损失函数为两者的简单加和：

L<sub>total</sub> = L<sub>MLM</sub> + L<sub>diff</sub>

侧链扩散过程:

FAMPNN 使用扩散模型生成侧链坐标。扩散过程基于 EDM (Equilibrium Diffusion Model) 框架，并使用了方差爆炸 (variance-exploding) 策略。

• 噪声添加: 逐步向真实的侧链坐标 x<sub>0</sub> 添加高斯噪声，得到加噪后的坐标 x<sub>t</sub>。噪声水平由时间步长 t 控制， t 越大，噪声越大。
• 去噪: 训练一个神经网络 (去噪器 D_θ) 来预测给定加噪坐标 x<sub>t</sub> 和噪声水平 σ<sub>t</sub> 的原始坐标 x<sub>0</sub>。
• 损失函数: 使用均方误差 (MSE) 损失函数，衡量去噪器预测的坐标与真实坐标之间的差异。

采样过程:

采样过程是一个迭代的去噪过程。从完全噪声的侧链坐标开始，逐步使用训练好的去噪器 D_θ 进行去噪，最终得到生成的侧链坐标。FAMPNN 使用迭代掩码采样方法，在每个步骤中，模型并行预测所有残基的 token (序列和侧链)，然后解屏蔽一部分 token，重复此过程直到所有 token 都被解屏蔽。

3.3. 训练和评估过程

• 训练细节: 使用 AdamW 优化器，学习率设置为 1e-4。在 CATH 数据集上训练的模型，batch size 为 64，训练步数为 100k。在 PDB 数据集上训练的模型，使用 4 块 NVIDIA H100 GPU，effective batch size 为 128，训练步数为 300k。
• 自洽性评估 (Self-Consistency): 使用 RFdiffusion 生成的 de novo 骨架评估序列设计模型的自洽性。对于每个骨架，使用序列设计模型生成序列，然后使用 AlphaFold2 预测生成序列的结构，并计算预测结构与原始骨架之间的 TM-score 和 RMSD。
• 侧链堆积评估 (Sidechain Packing): 在 CASP13, CASP14, CASP15 数据集上评估侧链堆积性能。使用 RMSD (Root Mean Square Deviation) 和 chi 角的平均绝对误差 (MAE) 作为评估指标。
• 蛋白质适应性预测评估 (Protein Fitness Prediction): 在 SKEMPIv2, Megascale, FireProtDB, S669 和抗体-抗原结合亲和力数据集上评估蛋白质适应性预测能力。使用 Spearman 相关系数评估模型预测结果与实验测量值之间的相关性。

4. 研究过程和结论

论文的核心研究过程围绕着验证全原子建模和侧链条件化对于蛋白质序列设计和适应性预测的有效性展开。

研究过程:

1. 模型构建: 设计并实现了 FAMPNN 模型，该模型显式地建模了氨基酸序列和侧链构象，并采用联合训练策略。
2. 序列恢复和自洽性评估: 在 CATH 4.2 数据集和 de novo 骨架上评估 FAMPNN 的序列恢复率和自洽性。实验结果表明，FAMPNN 在序列恢复方面具有竞争力，在自洽性方面与 ProteinMPNN 相当。
3. 侧链堆积评估: 在 CASP 数据集上评估 FAMPNN 的侧链堆积性能。结果表明，FAMPNN 在侧链堆积方面取得了最先进的性能，尤其是在 RMSD 指标上优于其他方法。
4. 蛋白质适应性预测评估: 在多个蛋白质适应性数据集上评估 FAMPNN 的预测能力。结果表明，FAMPNN 在蛋白质稳定性预测和蛋白质-蛋白质结合亲和力预测方面优于其他无监督模型，甚至在某些情况下优于有监督模型。
5. 消融实验: 通过消融实验研究了全原子条件化和侧链堆积目标对模型性能的影响。结果表明，全原子条件化和侧链堆积目标都能够提高序列设计性能，并且全原子条件化在蛋白质适应性预测方面尤为重要。

结论:

论文的主要结论是：

• 显式地建模侧链构象能够提高蛋白质序列设计的性能，尤其是在侧链堆积方面。
• 全原子条件化能够显著提高蛋白质适应性预测的准确性，表明全原子结构信息对于理解蛋白质的功能至关重要。
• FAMPNN 模型通过联合学习序列和侧链分布，实现了在序列设计和侧链堆积方面的最先进性能，并展示了在蛋白质适应性预测方面的潜力。

5. 总结和客观评价

FAMPNN 是一种新颖的全原子蛋白质序列设计方法，它通过显式地建模侧链构象，并在训练过程中结合序列预测和侧链扩散，实现了优异的性能。论文的实验结果充分证明了全原子建模和侧链条件化对于蛋白质序列设计的重要性。FAMPNN 在序列恢复、侧链堆积和蛋白质适应性预测方面都取得了具有竞争力的结果，表明该方法具有广阔的应用前景。

客观评价:

• 优点:
  • 提出了新颖的全原子蛋白质序列设计方法 FAMPNN。
  • 显式建模侧链构象，提高了侧链堆积性能。
  • 全原子条件化显著提高了蛋白质适应性预测能力。
  • 模型设计和实验验证充分，结果可靠。
• 不足:
  • 模型在抗体-抗原结合亲和力预测方面的提升不明显，可能受限于数据集的质量和输入结构的准确性。
  • 模型训练和推理的计算成本较高，需要进一步优化。

总的来说，FAMPNN 是一项重要的研究工作，它为蛋白质序列设计领域提供了一种新的思路和方法，并为未来的研究方向提供了有益的启示。

6. 参考文献和链接

• 论文链接: https://doi.org/10.1101/2025.02.13.637498
• 代码仓库: https://github.com/richardshuai/fampnn

参考文献列表 (部分):

• Dauparas, J., et al. (2022). Robust deep learning-based protein sequence design using ProteinMPNN. Science, 378(6615), 49-56.
• Jing, B., et al. (2020). Learning from protein structure with geometric vector perceptrons. arXiv preprint arXiv:2009.01411.
• Karras, T., et al. (2022). Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35, 26565-26577.
• Lee, J. S., & Kim, P. M. (2024). Flowpacker: Protein side-chain packing with torsional flow matching. bioRxiv, pp. 2024-07.
• McPartlon, M., & Xu, J. (2023). An end-to-end deep learning method for protein side-chain packing and inverse folding. Proceedings of the National Academy of Sciences, 120(23), e2216438120.

注意: 由于 bioRxiv 是预印本平台，该论文尚未经过同行评审，请读者注意甄别。