Nat Methods | OpenFold：对AlphaFold2重新训练提供了关于其学习机制和泛化能力的新见解

今天为大家介绍的是来自Mohammed AlQuraishi团队的一篇论文。AlphaFold2凭借极高的蛋白质结构预测准确性，彻底改变了结构生物学。然而，其实现方式缺乏训练新模型所需的代码和数据。这些对于解决新任务（如蛋白质-配体复合物结构预测）、研究模型的学习过程以及评估模型在未见过的折叠空间区域的泛化能力都是必要的。在此，作者介绍了OpenFold，这是一种快速、高效且可训练的AlphaFold2实现。作者从零开始训练OpenFold，达到了与AlphaFold2相同的准确性。在此基础上，作者发现OpenFold在泛化能力上表现出色，即使训练集的规模和多样性被刻意限制，包括几乎完全排除某些二级结构元素类别的情况。通过分析训练过程中产生的中间结构，作者还获得了OpenFold在折叠学习过程中层次化方式的见解。总之，作者的研究展示了OpenFold的强大和实用性，并相信其将成为蛋白质建模领域的重要资源。

根据序列来预测蛋白质结构，这一直是生物学领域数十年来的一个重大挑战。基于共进化信息以及同源结构的一系列深度学习研究，AlphaFold2可以说已经解决了对于具有足够深度MSA的天然蛋白质的结构预测问题。尽管AlphaFold2很有用，但是其官方实现中省略了模型复杂的训练过程代码以及运行所需的计算量巨大的数据。这使得研究人员难以(1) 研究AlphaFold2的学习行为及其对数据组成和模型架构变化的敏感性，(2) 创建模型变体以解决新任务。鉴于AlphaFold2的成功，它的许多新组件可能在蛋白质结构预测之外的任务中也会有用。例如，使用蛋白质复合物数据集重新训练AlphaFold2产生了AlphaFold2-Multimer，这是预测蛋白质复合物结构的最先进模型。然而，直到最近，这一能力还仅限于DeepMind。为了解决这一缺陷，作者开发了OpenFold，一个可训练的开源AlphaFold2实现。作者使用OpenProteinSet（作者对AlphaFold2训练集的开源复现）从头开始训练OpenFold，其预测质量与AlphaFold2相匹配。除了新的训练代码和数据，OpenFold相比AlphaFold2有几个优势：(1) 对于大多数蛋白质，运行速度快了三到五倍，(2) 使用更少的内存，允许在单个图形处理单元（GPU）上预测极长的蛋白质和多蛋白复合物，(3) 使用最广泛使用的机器学习框架PyTorch实现（AlphaFold2使用谷歌的JAX）。因此，OpenFold可以被最广泛的开发者社区轻松使用，并与现有的机器学习软件生态系统进行接口。

作者的模型在约3%的训练时间内可以达到约90%的准确度的优势，作者多次在特意省略部分训练集的版本上重新训练OpenFold，以量化其对未见过的蛋白质折叠的泛化能力。令人惊讶的是，模型即使在折叠空间大规模省略的情况下也相当稳健，但其泛化能力会根据蛋白质片段和折叠的空间范围而变化。作者在使用更多样化但更小的数据集进行训练时，观察到了更强的性能，有些数据集仅包含1000个实验结构。接下来，作者利用OpenFold来理解模型如何学习折叠蛋白质，重点关注训练中间阶段预测结构的几何特征。总之，这些结果提供了关于AlphaFold2类模型学习行为的基本新见解，并为生物分子建模算法的发展提供了新的概念和实用工具。

OpenFold完整重现了AlphaFold2的模型架构，没有任何会改变其内部数学计算的修改。这确保了OpenFold与AlphaFold2之间的完美互操作性，使得在OpenFold中使用原始AlphaFold2模型参数成为可能，反之亦然。为了验证OpenFold实现是否重现了AlphaFold2训练的所有方面，作者使用它从头开始训练一个新模型。OpenFold和AlphaFold2的训练需要一组蛋白质序列、多序列比对（MSAs）和结构。由于AlphaFold2的MSA数据库尚未公开，作者使用了OpenProteinSet，这是AlphaFold2训练数据集的复制版，在可用时替换了更新版本的序列数据库。作者从大约1500万个Uniclust MSAs中，选择了大约27万个多样且深度的MSAs，形成一个“自蒸馏”集；这种集合用于通过高质量预测来增强实验训练数据。作者使用AlphaFold2预测了该集合中所有MSAs的蛋白质结构，并将它们与来自蛋白质数据库（PDB）的大约13.2万个唯一（64万个非唯一）实验结构结合，形成了OpenFold训练数据集。在训练自蒸馏蛋白质时，AlphaFold2置信度评分低于0.5的残基（<0.5 pLDDT，预测局部距离差异测试）被屏蔽。作者的验证集由来自CAMEO的近200个结构组成，这是一个用于持续质量评估蛋白质结构预测模型的在线仓库，数据涵盖了截至2022年1月16日的3个月期间。

在主要的训练过程中，作者选择了七个快照模型，形成一组不同但相关的模型集合。在预测时，这些模型可以为同一蛋白质生成不同的结构假设。为了进一步增加这一集合的多样性，作者对从主模型分支出来的第二组模型进行了微调。在这个第二分支中，作者禁用了模型的“template”流程，类似于AlphaFold2使用的过程。从这个分支中选择的快照被添加到最终模型池中，总共形成了十个不同的模型。

图1

作者在图1中总结了训练实验的主要结果。使用OpenFold和AlphaFold2对CAMEO验证集进行的预测，使用lDDT-Cα（相对于α碳的局部距离差异测试）指标进行评估，显示了OpenFold和AlphaFold2之间的高度一致性，表明OpenFold成功重现了AlphaFold2。尽管OpenFold在CASP15竞赛的大部分时间内仍在训练中，但作者的回顾性评估显示，最终模型在CASP15领域也达到了等同水平。图1c提供了这种一致性的视觉说明。跟踪预测准确性与训练阶段的关系（图1d）揭示了一个显著的事实，即OpenFold在仅1500 GPU小时（约3%的训练时间）内达到了其最终准确度的约90%，在2500 GPU小时内达到了约95%；总训练时间约为50000 GPU小时。尽管作者会展示扩展的第二阶段训练在学习更详细的物理相互作用中起着重要作用，但这种准确度的快速提升表明，训练高性能的OpenFold变体所需的计算远少于完整模型训练，从而促进了模型架构的快速探索。作者在数据省略实验中利用了这一事实。

AlphaFold2的训练大致分为两个阶段：初始训练阶段和计算密集的微调阶段。在后者中，用于训练的蛋白质片段的大小增加到384个残基，并启用了额外的损失函数以惩罚结构违例（例如，空间碰撞）。通过比较初始阶段和微调阶段预测的结构，作者发现第二阶段对整体结构质量指标的影响仅为适中，即使仅考虑长度超过500个残基的长蛋白质。相反，微调的主要用途似乎是解决已知化学约束的违例。在作者的训练实验中，这一现象在微调开始后不久便发生，表明省略的微调运行可以在对预测质量影响最小的情况下使用。

除了预测准确性，作者还跟踪了pLDDT随训练阶段的变化。pLDDT是模型对预测结构的lDDT-Cα的估计，并作为其主要的置信度指标。作者发现，pLDDT在训练初期与真实的lDDT高度相关，尽管最初在自我评估中过于自信，后来进入了一个自信不足的阶段（图1b）。值得注意的是，模型在训练初期，即其整体预测能力仍然非常有限时，就能够评估自身预测的质量。

在确立了AlphaFold2和OpenFold的等效性之后，作者开始研究该架构的特性，从其数据效率入手。AlphaFold2使用了大约132,000个来自PDB的蛋白质结构进行训练，这是几十年艰苦且昂贵的实验结构确定工作成果。对于其他可能开发AlphaFold2风格模型的分子系统，数据要稀缺得多；例如，PDB中仅包含1,664个RNA结构。作者想知道AlphaFold2实现的高精度是否确实依赖于其相对较大的训练集，还是可以在使用更少数据的情况下实现可比的性能。如果是后者，这将表明AlphaFold2范式在分子问题上的广泛适用性。为了探讨这种可能性，作者进行了多次OpenFold训练，在每次训练中逐步减少训练数据，评估模型精度与训练集大小的关系。

图 2

第一组测试随机抽取原始训练数据至17,000、10,000、5,000、2,500、2,000和1,000条蛋白质链。作者使用每个抽取的集合训练OpenFold至少7,000步，经历初期快速提升阶段到早期收敛。为了避免来自完整训练集的信息泄漏，作者没有使用自蒸馏，这使新训练的模型相对于原始OpenFold处于劣势。作者训练了有结构模板和无结构模板的模型。在其他方面，训练与标准OpenFold模型完全相同。模型精度（使用lDDT-Cα评估）随训练步数的变化如图2a所示，颜色表示所使用的训练集大小。

作者发现，仅10,000条蛋白质链（约占全部训练数据的7.6%（黄色曲线））就足以达到与在完整训练集上训练的模型（粉色曲线）基本相同的初始lDDT-Cα值。经过20,000步（未在图中显示），完整数据模型达到峰值lDDT-Cα为0.83，而在7,000步后，10,000样本模型已超过0.81 lDDT-Cα。尽管随着训练集大小的进一步减少性能逐渐下降，且收敛速度在随机种子之间似乎相当可变（如图2a右侧的10,000次消融实验所示），作者发现所有模型的性能都令人惊讶地优异，即使是训练集中最小的1,000条蛋白质链的模型，仅占全部训练集的0.76%。作者强调，重要的注意事项是lDDT并不是衡量模型成功的唯一指标；在如此短时间内训练的模型无法学习罕见的二级结构元素（SSEs）（图4）或完全训练模型所获得的隐式生物物理能量函数。尽管如此，这些结果显然将AlphaFold2与先前的蛋白质结构预测模型区分开来。1,000链消融模型达到了0.64的lDDT-Cα，超过了CASP13中首个AlphaFold达到的0.62的中位lDDT-Cα，该模型在当时是表现最好的模型。

比较有模板和无模板模型的精度，作者发现即使在低数据环境下，模板对早期预测质量的平均贡献也很小。这与原始AlphaFold2的消融研究结果一致，后者显示模板除非在MSAs较浅或完全不存在时，影响才最小。

如前面的分析所示，随机抽样OpenFold训练集减少了所使用的训练数据量，但不一定会降低整体多样性。为了评估OpenFold对分布外数据的泛化能力，作者以结构分层的方式对训练集进行抽样，这样整个折叠空间的区域被排除在训练之外，但保留用于模型评估。蛋白质的多种结构分类体系存在，包括蛋白质结构域的层次分类（CATH）和蛋白质结构的分类（SCOP）系统。对于这项任务，作者使用CATH，它将蛋白质结构域按特异性递增的顺序分配到类（C）、架构（A）、拓扑（T）和同源超家族（H）。同一同源超家族分类的结构域可能表面上有所不同，但其结构核心高度相似。作者的前期分析可以视为在同源超家族层面对数据进行结构分层。在本次分析中，作者进一步分层数据，排除整个拓扑、架构和类。

首先，作者过滤掉未被CATH分类的蛋白质结构域，剩下大约440,000个结构域，涵盖了1,385个拓扑、42个架构和4个类。对于拓扑（T）层级，作者随机抽取100个拓扑，并从训练集中移除所有相关链。作者从这些排除的拓扑中抽取一个代表链，构建验证集。作者还从剩余的拓扑中逐步减少份额，构建更小的训练集，包括一个涵盖所有拓扑的训练集。对于架构（A）分层，作者采用类似的程序，但在这种情况下，验证集由五个架构中随机选出的100条链组成（每个架构20条链）。对于类（C）层级，验证集包含既不属于主要α类也不属于主要β类的结构域，因此富含同时具有高比例SSE的结构域。对于训练，作者构建了两个集合，一个仅包含主要α类，另一个仅包含主要β类；这使得作者能够确定主要在α螺旋或β折叠上训练的模型对同时包含两者的蛋白质的泛化能力。对于所有分层，作者从零开始训练OpenFold至早期收敛（约7,000步）。为了防止结构信息从排除类别泄露，所有训练都在没有模板的情况下进行。作者将模型精度随训练步数的变化绘制在图2b中，颜色表示每个CATH层次级别在训练中保留的类别比例。

如预期的那样，移除整个折叠空间区域对模型性能的影响比仅减少训练集大小更显著。例如，保留10%的拓扑进行训练（图2b中的绿色曲线，拓扑分层），对应约6,400条唯一链，导致模型性能低于包含5,000条随机选择链的模型（图2a中的绿色曲线，无模板）。然而，即使在最严重的训练集多样性省略中，绝对精度仍然出人意料地高。例如，包含5%拓扑的训练集（2,000条链）仍然达到了接近0.6的lDDT-Cα，类似于首个AlphaFold的结果，该模型在超过100,000条蛋白质链上进行训练。同样，基于最小构架层级的训练集仅包含一个构架（覆盖PDB几乎全部的42个构架之一）的结构域，然而其lDDT-Cα仍然接近0.6。最令人惊讶的是，类分层模型中几乎完全没有α螺旋或β折叠的训练集在包含α螺旋和β折叠的结构域上达到了超过0.7的lDDT-Cα分数。这些模型可能受益于其训练集中相对较多的唯一链数量：主要α螺旋集合包含15,400条链，而主要β折叠集合包含21,100条链。还应注意，主要α和主要β类别中分别包含小部分β折叠和α螺旋。总之，这些结果表明AlphaFold2架构具有显著的泛化能力。

为了更好地理解类分层模型的行为，作者分析了两个蛋白质结构域的结构，一个几乎完全由α螺旋组成的Lsi1结构域和另一个由β折叠组成的结构域TMED1，它们由主要α或主要β数据集训练的模型预测。在图2c的上排中，作者展示了Lsi1的实验结构（橙色）以及主要α训练模型（黄色）和主要β训练模型（红色）的预测结果。在下排，作者展示了TMED1的类似图像。可以预测到，主要α训练模型准确预测了Lsi1的α螺旋，但未能正确形成TMED1的β折叠，并在部分结构中错误地采用了一个小α螺旋。主要β训练模型则遇到了相反的问题：其Lsi1预测包含对齐不佳的螺旋和一个错误的β折叠，但TMED1预测得相当好。然而，值得注意的是，两者都没有出现灾难性错误。Lsi1的β折叠区域由主要α模型预测出大致正确的形状，只是其原子坐标不够精确，无法使DSSP将其分类为β折叠。

表 1

由于拓扑消融在排除的拓扑上进行了评估，构架消融在排除的构架上进行了评估，以此类推，因此不同数据省略实验的验证分数不能直接比较。为了更一致地了解每组数据省略实验的相对最终准确度，作者在表1中重新评估了每个模型在标准CAMEO验证集上的最终checkpoint。

OpenFold在折叠空间中排除区域的惊人泛化能力表明，它在全局折叠水平上对训练集多样性不太敏感。相反，模型似乎学会了从 MSA 和/或 序列-结构相关性的局部模式（片段、SSEs、单个残基等）中预测蛋白质结构，而不是从CATH的全局折叠模式。

尽管作者在上述实验中使用的OpenFold模型完全匹配AlphaFold2的计算逻辑，作者还做了一些最小化改变模型特性但在训练新模型和进行大规模预测时提高易用性和性能的改进。

图 3

首先，作者对数据预处理和训练过程进行了多项改进，包括一种低精度（‘FP16’）训练模式，这种模式能够在商业上可用的过去几代GPU（如广泛使用的NVIDIA V100）上更方便地进行模型训练。

其次，作者对主要结构损失FAPE（帧对齐点误差）进行了更改，以增强训练稳定性。在原始模型中，FAPE在大量训练批次中被限制在一个固定的最大值内。作者发现，在训练的动态早期阶段，这种策略过于激进，限制了有用训练信号的批次数量，并且经常阻止及时收敛。与其这样钳制整个批次，作者改为钳制每个批次中等效比例的样本，确保每个批次至少包含一些未钳制的链。这样，作者能够显著提高训练稳定性并加快模型收敛速度（见图3a）。重要的是，这一变化似乎并未影响完全训练模型的最终精度。

第三，作者在训练和推理过程中优化了内存效率。优化操作包括神经网络注意力机制的高级实现，这些机制对于异常短和长的序列具有有利特性，模块重构以降低内存使用，某些计算的可选近似以减少内存负担，以及针对GPU硬件定制的专业低级代码。

总之，这些优化使OpenFold的实现比AlphaFold2高效得多。作者在图3b中报告了OpenFold的运行时间。在推理过程中，OpenFold比AlphaFold2快达四倍。OpenFold在推理时比AlphaFold2更高效。当残基超过2,500时，AlphaFold2由于内存限制在单个GPU上崩溃。OpenFold能够成功运行超过4,000个残基长度的长蛋白质和复合物。OpenFold的训练速度与AlphaFold2相匹配或优于AlphaFold2，正如其他使用OpenFold的研究人员报告的那样，并且可以在较长的长度（最多可达1,200）上进行训练。

图 4

前面的分析表明，二级结构元素（SSEs）的学习是在三级结构之后进行的。接下来，作者正式确认了这一观察结果，并记录了不同SSEs的学习顺序。对于验证集中的每个蛋白质以及训练的每一步，作者使用DSSP识别出匹配八种已知SSE状态的残基。作者将实验结构中残基的DSSP分配作为真实值，并计算每种SSE在不同训练步骤中的召回率和精确度的F1分数（图4a）。

作者观察到SSEs的学习顺序非常清晰：首先学习的是α螺旋，随后是β折叠，然后是较不常见的SSE。这一顺序大致对应于蛋白质中SSEs的相对频率（图4b），除了不常见的螺旋变体。正如之前所见，模型对SSEs的发现滞后于准确的全局结构。例如，β折叠（‘E’）的F1分数在全局结构准确度达到顶峰数百步之后才达到平稳期，这可以通过GDT-TS（一种衡量蛋白质三级结构相似性的指标）来测量。这在逐步训练预测的动画中也清晰可见；对于每个蛋白质，二级结构在全局几何基本完成后才被识别和正确呈现。

为了调查OpenFold是否通过逐渐学习小片段螺旋来达到高α螺旋F1分数，作者将预测的螺旋按其恢复的真实螺旋的最长连续片段分类，并绘制了随训练步骤变化的直方图（图4c）。显然，几乎没有概率质量积累在0.0（完全未恢复的螺旋）和1.0（完全恢复的螺旋）之间。这表明，从DSSP的角度来看，大多数螺旋几乎同时被正确预测。这一突变与早期训练阶段螺旋DSSP F1分数的大部分提升相吻合。

基于上述观察，作者推测，随着训练的进行，OpenFold可能首先学习预测较小的结构片段，然后再学习较大的片段，这在三级和二级结构层次上都可能是显而易见的。

图 5

首先关注三级结构，作者在验证集中评估了每个训练步骤中所有非重叠片段（10、20和50个残基长）的预测质量。平均GDT-TS值如图5a所示。与全局GDT-TS（粉色）相比，训练的前300-400步内全局GDT-TS的改进最小，而片段GDT-TS在此阶段显著改善，短片段的提升更大。到1,000步时，模型达到了一个暂时的高原期，它已经学会了比全局结构更好地预测局部结构（10个残基片段的GDT-TS > 50，而整个蛋白质的GDT-TS < 10）。不久之后，在1,800步时，所有片段长度（包括全局结构）的准确度开始迅速提高。然而，短片段的提升小于长片段，以至于在3,000步时10个残基片段和整个蛋白质之间的差距比1,800步时要小（10个残基片段的GDT-TS ≈ 90，而整个蛋白质的GDT-TS ≈ 70）。这种趋势持续到模型完全训练完成时，10个残基片段和整个蛋白质之间的差距缩小到仅10个GDT-TS点。因此，尽管模型最终学会了几乎与局部结构一样好地预测全局结构，但它首先学会的是后者。

转向二级结构，作者调查了在检查SSEs时是否可以检测到同样的多尺度学习行为。与之前一样，作者将验证集中实验结构的DSSP分类作为真实值，仅关注α螺旋和β折叠。作者根据大小将两种SSE分类，定义为α螺旋的残基数和β折叠的链数。由于统一分类会导致高度不平衡的分类，作者选择了一种动态分类过程。首先，将每个SSE分配到与其大小相对应的（可能不平衡的）分类中。然后，在不超过最大大小（本例中，螺旋为200，折叠为30）的条件下，分类被迭代合并，直到不能再合并。最后，小于最小分类大小（螺旋和折叠均为20）的分类无条件与相邻分类合并。作者独立计算每个分类的平均指标（图5b）。在扩展数据图8中，作者通过区分小（6 Å）、中（12 Å）和长（24 Å）范围接触，提供了图5b中折叠数据的另一视图。

与三级结构的观察结果类似，短螺旋和窄折叠在训练的早期阶段比长和宽的对应物预测得更好。SSE准确度的提高与三级结构准确度的快速增加相吻合，尽管有偏移，如图4a所示。值得注意的是，尽管初始预测准确度差距较大，但预测的SSE的最终质量基本上与长度和宽度无关，这表明OpenFold最终在预测能力上变得独立于尺度，至少对于二级结构是这样。作者指出，SSE的识别是由DSSP进行的，DSSP对其氢键网络的细节很敏感。在训练的早期阶段，OpenFold可能已经恢复了DSSP由于原子定位不精确而未识别的二级结构方面。

编译 | 黄海涛

审稿 | 王建民

Ahdritz, G., Bouatta, N., Floristean, C., Kadyan, S., Xia, Q., Gerecke, W., ... & AlQuraishi, M. (2024). OpenFold: Retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization. Nature Methods, 1-11.