Protenix: AlphaFold3的再现 – ByteDance AML AI4Science团队

DRUGAI

蛋白质结构预测取得了显著进展。然而，预测涉及多种生物分子的复杂结构（如蛋白质-配体、蛋白质-核酸和抗体-抗原相互作用）仍面临额外的挑战。最近的研究已专注于这些复杂的预测任务。其中，AlphaFold 3 (AF3) 已设立了新的里程碑，代表了该领域的重大进步。尽管AF3取得了显著的进展，其有限的可访问性限制了在研究界的广泛应用。AF3论文中部分描述存在模糊及印刷错误，也为希望复现或改进模型的机器学习和计算生物学研究者带来了挑战。开源项目已为使这些先进模型的普及做出了重要贡献，但缺乏完整的训练代码和预处理数据仍然是研究人员复现和充分利用这些模型的障碍。为应对这些挑战，ByteDance AML AI4Science团队推出了Protenix，旨在降低这些门槛，并更好地支持跨学科研究社区。

以下是ByteDance AML AI4Science团队的主要贡献概要：

• 模型性能。将Protenix与现有的模型进行了基准测试。Protenix在不同分子类型的结构预测中表现出强大的性能。作为一个完全开源的模型，它使研究人员能够生成新的预测并对模型进行微调，以满足特定的应用需求。
• 方法。在复现过程中，依据AF3的描述实现了Protenix，并优化了一些模糊步骤，纠正了排版错误，并根据模型行为进行了有针对性的调整。通过分享复现经验，希望支持社区在这些改进的基础上进一步推动该领域的发展。
• 可访问性。已将Protenix开源，提供了模型权重、推理代码和可训练代码供研究用途。
• https://github.com/bytedance/Protenix

方法

数据管道

纳入不同类型的生物分子为数据整理和特征化带来了巨大挑战。研究人员详细分析了AF3的补充材料以复现数据管道。以下列出了在实现中的主要差异。

解析器。在选择替代位置时，研究人员使用第一个占据位置，而非最大占据位置。因为使用最大占据位置可能导致一些相邻残基采用不同构象，阻碍共价键的形成，导致链断裂。

多序列比对 (MSA)。研究人员使用MMSEQS2和ColabFold MSA管道进行MSA搜索，并从Uniref100数据库中提取用于配对的MSA（物种通过分类ID进行识别）。对于核酸链，不使用MSA。

模板。研究人员不使用模板。

特征化。

• 参考特征。研究人员使用CCD的2024-06-08版本，并通过pdbeccdutils v0.8.5生成RDKit构象。当RDKit构象生成失败时，使用CCD理想坐标，避免使用代表性坐标以防止潜在的数据泄漏。
• 键特征。键特征只包含配体内部的键、配体与聚合物之间的键，以及非标准残基内部的键。

裁剪。对裁剪方法的实现考虑了以下方面：

• 在连续裁剪方法中，研究人员剔除了金属和离子，因为该方法无法确保裁剪后的原子彼此接近。当金属或离子与其他链隔离裁剪时，模型可能无法准确预测其位置，导致较大的训练损失。因此，在连续裁剪中排除金属和离子，但在空间裁剪中保留它们。
• 由于配体和非标准氨基酸由每个标记一个原子表示，它们可能会被裁剪成片段。为防止这种情况发生，研究人员确保配体和非标准氨基酸保持完整，不被裁剪成不完整的片段。
• 在训练过程中，研究人员考虑了裁剪对链排列的影响。在空间裁剪中，将排列限制在裁剪区域内的链上，以确保预测集中在空间上接近的结构。相比之下，在连续裁剪中，裁剪的链可能较远，但等价的链可能更接近彼此，因此允许等价的排列并根据模型预测调整裁剪区域。

模型与训练

纠正与调整。在对原始论文中算法的分析中，研究人员发现了若干错误和模糊之处。部分调整对性能至关重要。此外，研究人员对置信度头进行了轻微修改，加入了LayerNorm并增加了一些线性层（详见发布的代码）。发现当严格按照原始论文中的描述实现时，置信度损失收敛效果不佳。

参数初始化。研究人员未对不同初始化策略的影响进行广泛分析，但对若干模块采用了零初始化。具体而言，为确保每个残差层在初始化时表现为恒等操作，部分线性层采用零初始化。参照最近的条件扩散模型研究，研究人员在AdaptiveLayerNorm中对两个线性层进行了零初始化。此外，对Pairformer中的N循环连接块也应用了零初始化。经验表明，这些零初始化有助于防止网络权重的维度崩溃，并缓解隐藏值的爆炸问题。

可扩展性。尽管AF3仅包含约3.86亿个参数，但它需要大量计算资源。尤其是其Pairformer块会产生大量中间激活。此外，在训练期间，扩散模块具有较大的扩散批量大小，并对每个原子对应用激活，从而为训练和推理带来了内存瓶颈。为应对这些挑战并提升训练效率，研究人员引入了以下方法：

• BF16混合精度训练。与FP32训练相比，BF16训练显著降低了峰值内存使用量，并几乎使实际训练速度翻倍。然而，并非所有GPU架构都支持BF16，仅限于NVIDIA的Ampere和Hopper等特定型号。
• 自定义CUDA内核。在PyTorch中直接实现LayerNorm常常面临内存限制。借鉴FastFold和OneFlow的优化，研究人员开发了一种自定义的LayerNorm实现，能够自适应调整块和网格大小，同时利用共享内存以优化GPU利用率。相比PyTorch的原生LayerNorm，该方法在各训练阶段实现了30%-50%的端到端训练速度提升。
• DS4Sci_EvoformerAttention。通过利用来自DeepSpeed4Science的DS4Sci_EvoformerAttention内核，在各训练阶段实现了10%-20%的端到端训练速度提升。
• 其他考虑。按照AF3论文直接实现模型更新步骤会导致创建不必要的大张量。为避免这一情况，研究人员采用了替代实现方式以防止这些大张量的生成。此外，为解决训练和推理过程中的内存瓶颈，研究人员采用了多种优化技术，包括梯度检查点、就地操作、张量卸载和分块，类似于OpenFold中使用的方法。

训练。Protenix在192块GPU上训练了约两周。研究人员还显著简化了超参数，因为训练仅使用了Protein Monomer Distillation数据集，未包括其他蒸馏集。尽管实现了与AF3类似的多阶段训练设置，但模型的训练步数显著减少。在初始阶段训练了75K步，在微调阶段1训练了15K步。微调阶段2仅进行了4K步，可能不足以充分训练置信度头。

结果

Protenix使用从PDB数据库整理的实验结构进行训练，数据截止日期为2021年9月30日，同时结合AlphaFold2和OpenFold预测的蛋白单体结构。研究人员在不同的评估数据集上对Protenix进行基准测试，以评估其在不同类型分子上的表现。由于商业实体对性能基准测试的限制，研究人员无法将Protenix与HelixFold3和Chai-1进行比较。因此，研究人员选择将Protenix的性能与AF3、AF2.3和RF2NA进行对比。

与AF3不同，AF3使用了两个不同数据截止日期（2019年9月30日用于Posebusters，2021年9月30日用于其他数据集）训练了两个独立的模型，而研究人员由于资源限制，选择训练单个模型。为避免数据泄漏，PoseBusters Benchmark Set版本2（PoseBusters V2）中的目标被排除。

对于每个PDB条目，研究人员遵循AF3的推理设置，使用5个模型种子生成25个预测，每个种子生成5个扩散样本。通过置信度评分对预测进行排序

配体

研究人员在PoseBusters Version 2基准测试集上评估了Protenix的性能，并将其与AF3进行比较。为确保对比的一致性和公平性，研究人员遵循了AF3论文中描述的评估程序。

RMSD成功率。成功率定义为预测中与真实值的口袋对齐配体的均方根偏差（RMSD）不超过2 Å的比例。图2[A]显示，Protenix在成功率方面优于AF3-2019，在RMSD和PB-Valid指标上都略有提高。这表明Protenix代表了蛋白-配体共折叠任务的最新技术水平（SOTA）。图2[B]展示了不同RMSD阈值下的成功率，提供了模型性能的更详细视图。Protenix - oracle配置（红色虚线）在所有RMSD阈值下均实现了最高成功率，始终优于AF3-2019的置信度排名（黑色实线）。Protenix - 置信度排名和Protenix - 所有样本配置与AF3的性能相当，显示了Protenix生成高质量蛋白-配体共折叠预测的能力。值得注意的是，所选结果仍然落后于所有样本中的最佳候选项，表明通过更好的样本排序器可能还有改进空间。

相似性分析。Protenix使用2021年9月30日截止的PDB数据训练，因此性能提升可能归因于额外的训练数据。为避免潜在的数据泄漏问题，研究人员从训练集中删除了相同的PDB ID。此外，研究人员进行了相似性分析，如图2[C]和[D]所示。这些图表明，与AF3-2019相比，Protenix在常见配体上的表现更好，可能得益于更多的训练数据。然而，由于在非常见配体上两者表现相同，可以得出结论，Protenix的泛化能力与AF3-2019相当。在目标端，Protenix在高同源组上优于AF3-2019，但在其他目标上略逊一筹。

蛋白质

研究人员将Protenix与当前表现最好的开源模型AF2.3进行基准测试。两个模型均预测每个结构的第一个生物组装体。对于AF2.3，非蛋白链被排除，非标准残基被映射为标准残基，因为该模型不支持这些成分。Protenix预测完整的复合物，但研究人员未评估包含额外链在输入中是否影响性能。结果总结在图1[B]中。

研究人员报告了不同接口类型的DockQ成功率，包括所有蛋白-蛋白界面和蛋白-抗体界面。研究人员报告了Protenix在10次循环后的结果，以及AF2.3 (5x5)在20次循环后的结果，分别提供了最佳（oracle）、中位数和排名第一的预测结果的分数。对于Protenix，蛋白质界面的排序器使用“链对ipTM”置信度。总体而言，Protenix的DockQ成功率高于AF2.3，表明预测准确性有所提高。唯一的例外是蛋白-抗体的情况，在该场景下，Protenix的最高排名预测与AF2.3表现相当，显示出样本排序器方面的改进潜力。

核酸

研究人员将评估范围扩展到RNA和DNA目标，结果显示Protenix在表现上与AF3相当，同时准确性高于RF2NA。Protenix未使用核酸链的MSA。

研究人员在CASP15 RNA目标上评估Protenix，遵循AF3论文中的方法，聚焦于截至2023年12月1日公开的8个目标。研究人员将Protenix的结果与AF3、AIchemy_RNA2和RF2NA进行比较。AIchemy_RNA2的预测结构来自CASP网站，选取LDDT分数最高的结构进行分析。RF2NA的预测结构从Zenodo存储库下载，与原论文一同提供。AF3的指标来自Bernard等人的预印本，其中系统地评估了AF3在RNA预测上的表现，并使用OpenStructure计算LDDT和TM分数。由于AF3的评估基于单个种子，研究人员也报告了Protenix在单个种子下的性能，而非五个种子。每个结构的最高排名样本基于pLDDT进行选择，用于评估。对于RF2NA、AIchemy_RNA2和Protenix，研究人员一致使用OpenStructure重新计算LDDT和TM分数。

如图1[C]和图6所示，Protenix的平均LDDT和TM分数与AF3相似，显著优于RF2NA，但仍落后于AIchemy_RNA2，后者受益于人为输入。

讨论

• 置信度分数。研究人员对置信度头的架构进行了轻微修改，与AF3略有不同。在训练置信度头后，观察到预测分数可以用于在低同源性的recentPDB数据集中选择界面。由于PAE头仅训练了4K步，仍有巨大的进一步改进空间。
• 局限性。研究人员审查了AF3论文中展示的示例，以评估Protenix在这些案例上的表现。发现模型在这些结构上的预测精度也很高。在观察到类似结果后，研究人员进一步调查了训练集，发现其中一些复合物与训练样本存在显著相似性。这表明Protenix和AF3的高准确性部分可能源于记忆效应，这凸显了需要更多分布外（OOD）测试集以更好地评估其泛化能力。

编译 | WJM

参考资料

• https://github.com/bytedance/Protenix
• https://github.com/bytedance/Protenix/blob/main/Protenix_Technical_Report.pdf