OpenFold更多细节公开：重新训练AlphaFold2对其学习机制和泛化能力产生新见解

编辑 | 萝卜皮

AlphaFold2 彻底改变了结构生物学，能够以极高的准确性预测蛋白质结构。然而，它的实现缺乏训练新模型所需的代码和数据。这些对于解决新任务是必要的，例如蛋白质-配体复杂结构预测；可用于调查模型学习的过程，目前这一过程仍然知之甚少；另外，还可评估模型对折叠空间不可见区域的泛化能力。

在这里，哥伦比亚大学的 Mohammed AlQuraishi 团队联合哈佛大学、哈佛医学院等多个机构的研究人员，开发了 OpenFold（一种快速、内存高效且可训练的 AlphaFold2 实现）和 OpenProteinSet（最大的蛋白质多序列比对公共数据库）。该团队使用 OpenProteinSet 从头开始训练 OpenFold，完全匹配 AlphaFold2 的准确性。

建立奇偶校验后，研究人员通过使用精心设计的数据集对其进行再训练来评估 OpenFold 跨折叠空间进行泛化的能力。研究表明 OpenFold 在泛化方面非常稳健，尽管训练集的大小和多样性极度减少，包括二级结构元素类的近乎完全省略。

通过分析 OpenFold 在训练过程中产生的中间结构，进一步阐明了模型学习折叠蛋白质的方式，研究人员发现空间维度是按顺序学习的。

该研究以「OpenFold: Retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization」为题，于 2022 年 11 月 22 日发布在 BioRxiv 预印平台。

几十年来，根据序列预测蛋白质结构一直是生物学的一项决定性挑战。基于将深度学习应用于多序列比对 (MSA) 和同源结构中编码的共同进化信息的一系列工作，AlphaFold2 可以说已经解决了具有足够深 MSA 的天然蛋白质的问题。该模型已通过 DeepMind 的官方开源实现向公众开放，已用于预测数亿种蛋白质的结构。这一实施使研究人员能够优化 AlphaFold2 的预测程序和用户体验，并将其用作新算法中的一个模块，包括用于蛋白质复合物预测、肽-蛋白质相互作用、结构排名等的算法。

尽管具有出色的实用性，但官方的 AlphaFold2 实现省略了模型复杂训练过程的代码以及运行它所需的计算量大的训练数据。这使得研究 AlphaFold2 的学习行为和对数据组成和模型架构变化的敏感性以及创建模型变体来处理新任务变得困难。鉴于 AlphaFold2 的成功，它的许多新组件可能被证明对蛋白质结构预测以外的任务有用。例如，使用蛋白质-蛋白质复合物的数据集对 AlphaFold2 进行再训练，得到了 AlphaFold2-Multimer，这是预测蛋白质复合物结构的最先进模型。然而，直到最近，这种能力一直是 DeepMind 独有的。

为了解决这个缺点，哥伦比亚大学的 Mohammed AlQuraishi 团队联合哈佛大学、哈佛医学院等多个机构的研究人员开发了 OpenFold，AlphaFold2 的可训练开源实现；以及 OpenProteinSet，一个包含 500 万个深度和多样化 MSA 的数据库，它消除了最重要的计算障碍之一——数百万 CPU 小时——以 AlphaFold2 的规模训练新的蛋白质模型。

图示：OpenFold 与 AlphaFold2 的准确性相匹配。（来源：论文）

研究人员使用 OpenProteinSet 从头开始训练 OpenFold，在预测质量上与 AlphaFold2 相匹配。除了新的训练代码和数据之外，OpenFold 与 AlphaFold2 相比还有几个优势：(i) 它在大多数蛋白质上的运行速度快三倍， (ii) 它使用更少的内存，允许在单个 GPU 上预测极长的蛋白质和多蛋白质复合物，(iii) 它是在 PyTorch 中实现的，PyTorch 是使用最广泛的机器学习框架（AlphaFold2 使用 Google 的 JAX）。因此，OpenFold 可以很容易地被最广泛的开发人员社区使用，并与现有机器学习软件的丰富生态系统进行交互。

研究人员使用 OpenFold 来了解模型如何学习折叠蛋白质，重点关注训练中间阶段预测结构的几何特征，并确定多个不同的行为阶段。具体来说，通过分析多种分辨率下的预测结构并将它们分解为二级和三级元素，研究人员发现 OpenFold 以交错的方式学习空间维度、二级结构元素和三级尺度。

图示：OpenFold 预测的维度增长。（来源：论文）

接下来，由于 OpenFold 可以在约 3% 的训练时间内实现约 90% 的模型准确度，于是研究人员对 OpenFold 进行了多次特殊删减版本的训练集再训练，以量化其泛化到不可见蛋白质折叠的能力。

令人惊讶的是，该模型甚至对折叠空间的大省略也非常稳健，但其泛化能力因蛋白质片段和折叠的空间范围而异。在更多样化但更小的数据集上训练模型时，研究人员观察到了 OpenFold 更强大的性能，甚至有些数据集小到只有 1,000 个实验结构。总而言之，这些结果对 AlphaFold2 类型模型的学习行为产生了基本的新见解，并为生物分子建模算法的开发提供了新的概念和实用工具。

图示：OpenFold 在省略训练集上的泛化能力。（来源：论文）

OpenFold 使两大领域的进步成为可能：(i) 更深入地分析类似 AlphaFold2 模型的优势、劣势和学习行为；(ii) 开发利用 AlphaFold2 模块的新（生物）分子模型。在这项工作中，研究人员更专注于前者。

首先，通过分析部分训练模型的预测结构，研究人员发现类似 AlphaFold2 的模型按顺序学习空间维度。这种行为对模型架构和训练方案的设计有影响。

例如，将物理先验知识整合到机器学习模型中是一个具有突出科学意义的研究领域。此类综合的努力产生了不同的结果，事实上，AlphaFold2 是一个非常成功的模型的开创性例子，该模型几乎完全没有物理先验。它的学习行为说明了为什么合并这样的先验会很困难——在折叠的 1D 和 2D 学习阶段，所有预测的结构都表现出严重违反基本化学定律和大量空间冲突的情况。但是，禁止此类违规行为会极大地改变 AlphaFold2 的学习行为。

事实上，在最初的 AlphaFold2 论文中，未经进一步阐述就观察到，启用违规损失以惩罚空间冲突和非物理键长会破坏训练的稳定性。研究人员对空间折叠学习阶段的观察为这种观察提供了解释。

图示：二级结构类别是按顺序学习的。（来源：论文）

AlphaFold2 针对此问题采用的解决方案，即仅在训练的后期阶段对违反自然规律的行为进行惩罚，提出了一种更广泛的策略来解决物理先验的合并问题：这是一种课程学习方法，其中模型首先可以自由提取信息并从数据中学习，之后可以逐渐引入更复杂的物理先验，从而提高模型的泛化能力。正如，对 OpenFold 所做的那样，分析学习轨迹，为何时可以将此类先验注入训练过程，提供了具体的时间表。

其次，研究人员观察到空间坍缩阶段，对应于最终预测结构的不完美低维 PCA 投影。鉴于其他端到端可区分蛋白质结构模型不表现出相同的行为，所以这种情况无法从文献中得到解释。虽然没有直接证据，但研究人员怀疑 AlphaFold2 架构的某些方面——特别是 FAPE 损失函数——可能会导致这种现象。

研究人员推测，可能类似 PCA 的级数允许模型，通过首先解决对 FAPE 损失具有最大收益的问题来贪婪地最小化错误，根据定义，FAPE 损失位于地面实况结构的最大主成分上。解决后，该模型将继续处理与其他低维投影相关的较小问题。

图示：新 FAPE 钳位协议的稳定性。（来源：论文）

如果是这种情况，学习过程中空间维度的交错将取决于训练集中蛋白质的几何形状。完全由细长管状蛋白质组成的训练集的极端情况会产生更加显著的交错阶段。相反，由完美球形蛋白质组成的训练集将在所有空间维度上呈现均匀增长。这种行为是整个训练集的函数，不一定会反映在单个蛋白质中。无论如何，这些观察结果表明，有可能通过学习课程来故意简化分子建模中的其他难题，在该学习课程中，首先训练「玩具」模型以预测目标分子的低维投影，然后再负责预测其完全实现的实例化 。

图示：OpenFold 对比 AlphaFold2 预测速度。（来源：论文）

研究人员评估了 OpenFold 从大幅缩小的训练集中学习的能力。值得注意的是，研究人员发现即使数据集大小减少 100 倍，模型的性能也比第一版 AlphaFold 更高。换句话说，AlphaFold2 中引入的架构进步使其数据效率比其前身高出 100 倍，后者在推出时树立了新的技术水平。这些结果表明，与更大的数据集相比，架构创新可以对模型准确性产生更深远的影响，特别是在数据采集成本高昂或耗时的领域，如（生物）分子系统中经常出现的情况。

然而，值得注意的是，AlphaFold2 通常学习 MSA 结构关系，而不是序列结构关系。MSAs 隐含地编码了大量的结构知识，正如早期基于协同进化的结构预测方法所证明的那样，这些方法完全不受监督，不使用实验结构数据。因此，AlphaFold2 架构对不表现出共同进化信号的问题的适用性仍未得到证明。

该团队的数据省略结果可以根据最近对大型基于 transformer 的语言模型的研究进行解释，这些研究揭示了广泛适用的「缩放法则」，该法则将模型准确性预测为模型大小、使用的计算和训练集大小的简单函数。当不受这三个支柱中任何一个的限制时，模型会从对其他两个支柱的投资中受益。这些观察主要集中在基于 transformer 的架构上，AlphaFold2 就是一个例子，但最近的工作揭示了其他架构的类似行为。尽管确定 AlphaFold2 的精确缩放属性并不是该团队关注的重点，但他们的结果表明它几乎不受 PDB 的大小或多样性的限制，从而激发了对其架构的更大实例的潜在开发。

视频：蛋白折叠过程示例。（来源：论文）

OpenFold 为未来旨在改进 AlphaFold2 架构并将其重新用于新的分子建模问题的工作奠定了基础。自 2021 年 11 月该团队发布代码库以来，已经进行了多项构建和扩展 OpenFold 的工作。其中包括用于蛋白质结构预测的 ESMFold 方法，它用蛋白质语言模型取代了 MSA，以及 FastFold，这是一项社区工作，已经实现了包括快速模型并行训练和推理在内的重大改进。

研究人员希望未来的工作能够通过反汇编 OpenFold 来进一步解决蛋白质结构预测之外的问题。例如，evoformer 模块是用于推理进化相关序列的通用原语。DNA 和 RNA 序列也表现出共同进化信号，其努力旨在从快速实现的 MSA 预测 RNA 结构。进化生物学中更基本的问题（例如系统发育推断）可能证明适用于类似进化者的架构，这似乎是合理的。类似地，AlphaFold2 的结构模块，特别是不变点注意机制，为聚合物空间推理提供了一种通用方法，可以进一步扩展到任意分子。

研究人员预计，随着蛋白质结构和其他生物分子从要预测的输出转变为要使用的输入，依赖空间推理能力的下游任务将变得越来越重要。也许 OpenFold 将在促进这些发展方面发挥关键作用。

论文链接：https://www.biorxiv.org/content/10.1101/2022.11.20.517210v1

相关报道：https://twitter.com/gahdritz/status/1595420944880779266

人工智能 × [ 生物 神经科学 数学 物理 材料 ]

「ScienceAI」关注人工智能与其他前沿技术及基础科学的交叉研究与融合发展。

欢迎关注标星，并点击右下角点赞和在看。

点击阅读原文，加入专业从业者社区，以获得更多交流合作机会及服务。