[Arxiv] Mask prior-guided denoising diffusion improves inverse protein folding

Mask Prior引导的去噪扩散模型提升逆向蛋白质折叠的性能

1. 背景介绍

蛋白质是生命活动中至关重要的生物大分子，由氨基酸线性序列折叠形成复杂的三维（3D）结构。理解蛋白质的结构与功能关系是生物学和医学领域的核心挑战。逆向蛋白质折叠 (Inverse Protein Folding, IPF)，也称为蛋白质序列设计，是一个 фундаментальный 的结构导向的蛋白质设计问题。它的目标是根据给定的目标 3D 骨架结构，生成能够折叠成该结构的有效氨基酸序列。这项技术在治疗性蛋白质工程、先导化合物优化和抗体设计等领域具有巨大的应用潜力。

传统的基于物理的方法将 IPF 视为一个能量优化问题，但计算成本高昂且精度有限。近年来，深度学习方法凭借其强大的非线性模式学习能力，已成为解决蛋白质结构问题的首选范式。早期的基于卷积神经网络的模型通常将每个蛋白质残基视为孤立单元或将整个蛋白质视为点云数据，对结构信息和残基间相互作用的考虑不足。最近，基于图的方法将 3D 蛋白质结构表示为邻近图，并使用图神经网络 (Graph Neural Networks, GNNs) 来建模残基表示并整合结构约束。GNNs 可以聚合和交换图结构数据内的局部信息，从而显著提升了基于图的方法的性能。

尽管基于图的方法取得了进展，但仅靠结构信息不足以确定某些挑战性区域（如环区和固有无序区）的残基类型。在这些不确定性高、置信度低的情况下，与其他准确预测的残基的相互作用可以为减轻这些区域的不确定性提供更可靠的指导。此外，现有的基于深度学习的 IPF 方法通常采用自回归解码或均匀随机解码来生成氨基酸序列，容易累积预测误差，并且难以捕捉蛋白质进化中的全局和长程依赖关系。最近，一些非自回归方法显示出在相关背景下超越自回归范式的潜力。

关键词解释:

• 逆向蛋白质折叠 (Inverse Protein Folding, IPF): 根据给定的蛋白质三维结构设计氨基酸序列的过程，也称为蛋白质序列设计。
• 图神经网络 (Graph Neural Networks, GNNs): 一种用于处理图结构数据的神经网络，能够学习节点和边之间的关系，并进行信息传递和聚合。
• 自回归解码 (Autoregressive Decoding): 一种序列生成方法，其中当前步的输出依赖于先前步骤的输出。在蛋白质序列设计中，意味着逐个氨基酸生成序列，当前氨基酸的生成依赖于之前生成的氨基酸。
• 非自回归解码 (Non-autoregressive Decoding): 一种序列生成方法，可以并行生成序列的每个部分，不依赖于先前步骤的输出。在蛋白质序列设计中，可以一次性预测整个蛋白质序列。
• 去噪扩散模型 (Denoising Diffusion Models): 一种生成模型，通过逐步向数据添加噪声，然后再反向逐步去噪来学习数据分布，从而生成新的数据样本。

2. 核心概述

本文提出了一种名为 Mask prior-guided denoising Diffusion (MapDiff) 的框架，用于改进逆向蛋白质折叠任务。MapDiff 采用离散去噪扩散概率模型，迭代地生成氨基酸序列，并在生成过程中结合了结构信息和残基相互作用。该框架使用图神经网络构建去噪网络，并引入掩码先验预训练策略来处理低置信度区域的预测不确定性。此外，在生成过程中，MapDiff 结合了去噪扩散隐式模型 (DDIM) 与 Monte-Carlo dropout，以提升不确定性估计和加速生成过程。在四个具有挑战性的序列设计基准数据集上的评估结果表明，MapDiff 显著优于现有最先进的方法。更重要的是，MapDiff 生成的序列在体外实验中展现出与天然蛋白质相似的物理化学和结构特征。

3. 方法论和实验细节

3.1 数据集

本文使用了以下几个数据集来评估 MapDiff 的性能：

• CATH 4.2 和 CATH 4.3: 两个常用的蛋白质结构域分类数据库。根据蛋白质拓扑结构进行划分，CATH 4.2 数据集包含 18,024 个训练蛋白质结构，608 个验证结构和 1,120 个测试结构；CATH 4.3 数据集包含 16,630 个训练结构，1,516 个验证结构和 1,864 个测试结构。训练集、验证集和测试集之间蛋白质拓扑结构代码无重叠。
• TS50: 一个常用的蛋白质序列设计基准数据集，包含 50 个独立的蛋白质链。
• PDB2022: 一个包含最近发布的来自蛋白质数据库 (Protein Data Bank, PDB) 的单链结构的集合，蛋白质长度小于 500 个氨基酸，分辨率优于 2.5 Å。该数据集包含 1,975 个蛋白质结构，与其他实验数据集无重叠。

对于 CATH 4.2 和 CATH 4.3 数据集，除了完整的测试集外，作者还研究了两个子类别：短蛋白（长度不超过 100 个残基）和单链蛋白（在 CATH 中标记为单链）。

3.2 算法和模型

MapDiff 框架的核心是一个 Mask prior-guided denoising network，其结构如图 1b 所示，包含以下三个关键组件：

1. Structure-based sequence predictor (基于结构的序列预测器): 使用等变图神经网络 (Equivariant Graph Neural Network, EGNN) 作为核心架构。EGNN 接收带噪氨基酸序列和 3D 骨架结构作为输入，通过消息传递机制学习残基之间的相互作用和结构信息，并预测去噪后的氨基酸序列。为了增强模型对全局信息的感知能力，EGNN 中引入了 global-aware module (全局感知模块)，在消息传递过程中整合全局池化向量。
2. Mask ratio adapter (掩码比例适配器): 动态调整掩码比例。掩码比例与去噪步骤 t 相关的噪声权重 βt 成正比，通过正弦函数和预定义的偏差和最小掩码比例参数进行计算。公式为：
   $$m_r^t = \sin\left(\frac{\pi}{2}\beta_t \cdot \sigma\right) + m$$
   其中，βt 是噪声调度中的噪声权重， σ 和 m 分别是预定义的偏差和最小掩码比例。
3. Masked sequence designer (掩码序列设计器): 使用 不变点注意力 (Invariant Point Attention, IPA) 网络，用于细化低置信度残基的预测。IPA 网络接收掩码后的氨基酸序列和 3D 骨架结构作为输入，利用几何感知注意力机制融合残基表示和空间关系，预测被掩码残基的类型。预训练 (Pre-training) 阶段，Masked sequence designer 通过 BERT 类似的掩码语言建模目标进行训练，学习结构和序列的先验知识。

损失函数:

• Base cross-entropy loss (基础交叉熵损失): 用于训练 Structure-based sequence predictor，计算预测的氨基酸概率分布与真实氨基酸类型的交叉熵损失。公式为:
  $$L_b = L_{CE}(p^b(X^{aa}), X^{aa})$$
• Mask cross-entropy loss (掩码交叉熵损失): 用于训练 Masked sequence designer，计算预测的掩码残基概率分布与真实掩码残基类型的交叉熵损失。公式为：
  $$L_m = L_{CE}(p^m({X^{aa}}{mask}), {X^{aa}}{mask})$$
• Total loss (总损失): 基础交叉熵损失和掩码交叉熵损失之和。公式为：
  $$L = L_b + L_m$$

扩散过程和去噪过程:

• 扩散过程 (Diffusion process): 如图 1c 所示，逐步向原始氨基酸序列 $X_0^{aa}$ 添加离散噪声，使其逐渐趋于均匀分布。噪声添加过程由转移概率矩阵 $Q_t$ 控制。
• 去噪过程 (Denoising process): 迭代地从带噪氨基酸序列 $X_t^{aa}$ 中去除噪声，逐步恢复原始序列 $X_0^{aa}$。去噪过程使用 Mask prior-guided denoising network $φ_θ$，学习从 $X_t^{aa}$ 预测 $X_{t-1}^{aa}$ 的分布。为了加速生成过程和提高不确定性估计，MapDiff 结合了 DDIM (Denoising Diffusion Implicit Model) 和 Monte-Carlo dropout。DDIM 允许跳过部分去噪步骤，Monte-Carlo dropout 通过多次前向传播并平均预测结果来提高不确定性估计。

3.3 训练和评估过程

训练过程:

1. Mask prior pre-training (掩码先验预训练): 首先对 Masked sequence designer 进行预训练。随机掩码氨基酸序列中的部分残基，然后使用 IPA 网络预测被掩码的残基类型。预训练目标是最小化掩码交叉熵损失。
2. Denoising diffusion model training (去噪扩散模型训练): 训练 Mask prior-guided denoising network。在每个训练步骤中，模型接收带噪氨基酸序列和对应的 3D 骨架结构作为输入，预测去噪后的氨基酸序列。训练目标是最小化总损失（基础交叉熵损失 + 掩码交叉熵损失）。

评估过程:

• 序列恢复性能评估: 使用以下指标评估生成序列的准确性：
  • Perplexity (困惑度): 衡量模型预测的氨基酸概率分布与真实氨基酸类型之间的对齐程度，困惑度越低表示模型预测越准确。
  • Recovery rate (回复率): 准确预测的氨基酸残基在蛋白质序列中所占的比例，回复率越高表示模型预测越准确。
  • Native Sequence Similarity Recovery (NSSR) (天然序列相似性回复): 使用 BLOSUM 矩阵评估预测序列与天然序列之间的相似性，考虑氨基酸的相似替换。本文使用了 BLOSUM42, BLOSUM62, BLOSUM80 和 BLOSUM90 四种不同的 cutoff level。
• Foldability evaluation (可折叠性评估): 使用 AlphaFold2 预测生成序列的 3D 结构，并与天然晶体结构进行比较，评估生成序列的可折叠性和预测结构的质量。评估指标包括：
  • predicted Local Distance Difference Test (pLDDT)
  • predicted Aligned Error (PAE)
  • predicted Template Modeling (pTM)
  • Template Modeling score (TM-score)
  • Root Mean Square Deviation (RMSD)
  • Global Distance Test-Total Score (GDT-TS)

实验在单个 Tesla A100 GPU 上进行。Adam 优化器和 one-cycle scheduler 用于参数优化，峰值学习率设置为 5e-4。

4. 研究过程和结论

文章首先在 CATH 4.2 和 CATH 4.3 数据集上评估了 MapDiff 的序列回复性能，并与多个最先进的基线模型进行了比较，包括 StructGNN, GraphTrans, GVP, AlphaDesign, ProteinMPNN, PiFold, LM-Design 和 GRADE-IF。结果表明，MapDiff 在不同的评估指标和数据集子集上均取得了最佳性能，在 CATH 4.2 和 CATH 4.3 完整测试集上，分别实现了 60.93% 和 60.68% 的回复率，显著优于现有方法。MapDiff 还展现出最低的困惑度，表明其生成了高置信度的概率分布，有利于准确预测。值得注意的是，MapDiff 在不依赖外部知识的情况下，依然取得了优异的性能，表明其模型架构和基于扩散的生成机制能够有效利用有限的训练数据捕捉相关模式，实现卓越的泛化能力。

为了进一步验证 MapDiff 的泛化能力，作者在 TS50 和 PDB2022 两个独立的测试数据集上进行了零样本迁移测试。结果表明，MapDiff 在这两个数据集上依然取得了最高的回复率和 NSSR 分数，表明其模型具有良好的零样本迁移能力。

为了评估生成蛋白质序列的可折叠性，作者使用 AlphaFold2 对 MapDiff 等模型生成的序列进行了结构预测，并将预测结构与天然晶体结构进行比较。结果表明，MapDiff 生成的蛋白质序列展现出卓越的可折叠性，预测结构具有最高的置信度和最小的偏差。尤其在 CATH 4.2 测试集上，MapDiff 的 pLDDT, PTM, TM-Score 和 GDT-TS 指标均为最高，RMSD 指标为最低，表明 MapDiff 生成的序列能够折叠成与天然结构高度相似的 3D 结构。

为了深入分析 MapDiff 各个模块的有效性，作者进行了消融研究，分别移除了 EGNN 中的 EdgeUpdate, CoordinateUpdate, GlobalContext 模块，以及 Refinement 模块中的 MaskAdapter 和 IPA network。结果表明，每个模块都对序列回复性能的提升有积极贡献，其中基于 IPA 的 Refinement 机制 (variant 5) 带来的提升最为显著，回复率提升了 7.9%。

为了评估 Monte-Carlo 采样和 DDIM 跳步对性能的影响，作者进行了敏感性分析。结果表明，增加 Monte-Carlo 样本数量和调整 DDIM 跳步步数可以进一步提升 MapDiff 的性能，表明结合 DDIM 和 Monte-Carlo dropout 能够有效加速生成过程并提高不确定性估计。

综上所述，实验结果充分证明了 MapDiff 框架在逆向蛋白质折叠任务中的有效性和优越性。MapDiff 不仅在序列回复性能上超越了现有最先进的方法，而且生成的蛋白质序列具有良好的可折叠性和结构质量。

5. 总结和客观评价

本文提出的 MapDiff 框架，利用掩码先验引导的去噪扩散模型，有效地提升了逆向蛋白质折叠的性能。该方法创新性地将 IPF 问题建模为离散去噪扩散过程，并设计了 Mask prior-guided denoising network 来捕捉结构信息和残基相互作用。通过结合 EGNN, IPA, DDIM 和 Monte-Carlo dropout 等技术，MapDiff 在生成准确率、泛化能力和不确定性估计方面都取得了显著的提升。实验结果表明，MapDiff 在多个基准数据集上超越了现有方法，并生成了高质量的蛋白质序列。

客观评价:

MapDiff 框架的创新性在于将掩码先验预训练策略与去噪扩散模型相结合，有效解决了逆向蛋白质折叠中低置信度区域预测不确定性的问题。模型设计精巧，充分利用了结构信息和残基相互作用，并在多个数据集上进行了全面的评估，实验结果扎实可靠。文章结构清晰，逻辑严谨，方法描述详细，结果分析深入。

优点:

• 提出了新颖的 Mask prior-guided denoising Diffusion 框架，有效提升了 IPF 性能。
• 模型设计巧妙，融合了 EGNN, IPA, DDIM 和 Monte-Carlo dropout 等先进技术。
• 实验评估全面，在多个数据集和评估指标上均取得了优异的性能。
• 消融研究和敏感性分析深入，验证了模型各组件的有效性。

不足:

• 模型参数量相对较大 (14.7M)，模型复杂度较高。
• 虽然使用了 DDIM 加速生成过程，但去噪扩散模型的迭代生成方式依然存在计算成本较高的局限性。

总体而言，MapDiff 是一项具有重要意义的工作，为逆向蛋白质折叠领域带来了新的突破，并为未来蛋白质设计方法的研究提供了新的思路。

6. 参考文献和链接

• 论文链接: arXiv:2412.07815v1 [q-bio.BM]
• 代码仓库: https://github.com/peizhenbai/MapDiff (根据论文描述，代码将在 GitHub 仓库提供)

请注意: 论文中没有明确提供数据集的下载链接，但提到了数据集来源和相关论文，可以根据描述在相关资源中查找。