NeurIPS 2024｜AlphaFold结合流匹配生成蛋白质系综

DRUGAI

今天为大家介绍的是来自Tommi Jaakkola团队的一篇论文。蛋白质的生物学功能通常依赖于动态结构系综。在这项工作中，作者开发了一种基于流的生成模型方法，用于学习和采样蛋白质的构象景观。作者将AlphaFold和ESMFold等高精度的单态预测器重新利用，并在自定义流匹配（Flow Matching）框架下对其进行微调，以获得序列条件的蛋白质结构生成模型，称为AlphaFLOW和ESMFLOW。在PDB上训练和评估时，该方法在精度和多样性上比AlphaFold的MSA子采样方法有显著优势。在对全原子MD的集合进行进一步训练后，该方法能够准确捕捉未见蛋白质的构象灵活性、位置分布和更高阶的系综观测值。此外，该方法可以通过更快的时间收敛于某些平衡特性，将静态PDB结构多样化，展示了其作为昂贵物理模拟代理的潜力。

代码可在https://github.com/bjing2016/alphaflow获取。

尽管AlphaFold等深度学习方法在单态实验蛋白质结构建模方面表现出色，但它们未能考虑到这种构象异质性。因此，一种在单结构预测器的精确度基础上揭示潜在结构集合的方法，对结构生物学家来说将具有重要价值。

现有的机器学习方法在生成结构集合时，主要集中于对AlphaFold进行推理时的干预，修改多序列比对（MSA）输入，从而对每个版本的MSA产生不同的结构预测。虽然这些方法取得了一些成功，但它们存在两个主要限制。首先，通过操作MSA，它们无法推广到基于蛋白质语言模型（PLM）的结构预测器，如ESMFold或OmegaFold，这些模型因其快速运行时间和易用性而越来越受欢迎。其次，这些推理时的干预无法对PDB以外的蛋白质集合进行训练，例如，分子动力学产生的系综，这些集合具有重要的科学价值，但模拟成本极高。

为了解决这些限制，作者在这项工作中将AlphaFold和ESMFold与流匹配（Flow Matching）相结合，提出了一种采样蛋白质构象景观的系统方法。流匹配是一种最近发展的生成建模框架。虽然AlphaFold和ESMFold最初是作为回归模型开发和训练的，旨在为给定的MSA或序列输入预测单一最佳蛋白质结构，但作者开发了一种策略，将其重新利用为（序列条件的）蛋白质结构生成模型。这种合成依赖于一个关键见解，即迭代去噪框架（如扩散和流匹配）提供了一种通用方法，可以将回归模型转换为生成模型，而对架构和训练目标的修改相对较少。与推理时的MSA消融不同，这种策略同样适用于基于PLM的预测器，并且可以用于任意集合的训练或微调。

式 1

模型部分

给定一个由氨基酸序列A构成的蛋白质序列，作者的目标是对其3D坐标x上的分布p(x | A) 进行建模，该分布表示该蛋白质序列的结构系综。考虑到开发一个确定性序列到结构模型（即AlphaFold）所需的大量智力投入，开发一个具有同等精度和泛化能力的分布模型似乎是一个巨大的挑战。作者的解决方案是利用生成建模的最新概念进展，几乎无需大幅修改地将AlphaFold重新利用为一个生成模型。

图 2

例如，考虑原型文本到图像扩散模型的（简化）架构，这些模型旨在对图像x的条件分布p(x | s)进行建模，其中图像x是基于文本提示s生成的。在这些模型的核心是一个去噪神经网络（例如，UNet），它接收一个噪声图像和一个文本提示，预测出一个干净的图像。基于这些输入，这类模型通常通过简单的类似回归的均方误差（MSE）目标进行训练。类似地，一个基于类似回归损失训练的蛋白质结构预测器——如AlphaFold或ESMFold——可以通过提供一个额外的噪声结构输入转换为去噪模型（图2）。这并非巧合，这让人联想到某些AlphaFold工作流中使用的模板结构的理念。因此，作者开发了一个非常类似于AlphaFold模板嵌入堆栈的输入嵌入模块，并将其预先添加到AlphaFold和ESMFold的成对折叠主干之前。通过这样做，作者获得了结构去噪架构，这些架构是在经过良好验证的单结构预测器基础上的简单包装。通过这些架构修改，作者准备将AlphaFold和ESMFold接入任何基于迭代去噪的生成建模框架。接下来，将具体介绍这如何应用于蛋白质系综的流匹配。

设计一个流匹配生成框架相当于选择一个条件概率路径及其对应的向量场。受基于插值的流匹配观点的启发，作者通过从先验中采样噪声并与数据点进行线性插值来定义条件概率路径（式2）。概率路径与向量场相关（式3）。而作者随后学习一个神经网络来根据式1逼近边缘向量场。然而，作者定义一个神经网络并通过式4重新参数化。

式 2

式 3

式 4

通过对式1和式4的重新排列显示，可以通过式5等效地学习的期望值。

式 5

为了将流匹配应用于蛋白质结构，作者通过其β-碳（甘氨酸的α-碳）的3D坐标来描述一个结构x。然后，作者定义这些β-碳位置的先验分布q(x)为谐波先验（式6）。

式 6

对学习条件期望值的参数化（方程5）表明，神经网络应通过均方误差（MSE）损失进行训练。然而，这种直接方法存在一些问题。为了将这些问题与流匹配框架调和，作者将蛋白质结构空间重新定义为商空间，并将先验分布投影到该空间。作者重新定义了在此空间中两点之间的插值。此外，因为商空间不再是向量空间，所以不再存在分布“期望”的概念；相反，作者旨在学习条件分布 的更一般的弗雷歇平均值（式7）。

式 7

作者利用FAPE作为有效度量来定义弗雷歇平均值。为了学习这一目标，作者使用与原始FAPE相同的训练损失，只是现在将其平方。训练和推理过程的最终结果在算法1和算法2中提供。这个修改框架的重要意义在于，虽然作者的模型在所有原子坐标上进行监督，但从技术上讲，它仅学习β-碳坐标上的分布。

算法 1

算法 2

PDB系综的生成能力

作者使用模型的流匹配框架微调了AlphaFold和ESMFold的所有权重，基于它们公开可用的预训练权重进行操作。为了实现架构和训练管道，作者使用了OpenFold，并且用OpenProteinSet来训练MSA。按照原始工作，作者分别使用了2018年5月1日和2020年5月1日作为AlphaFold和ESMFold的训练截止日期。在这一阶段训练结束时（分别为1.28M和720k例子），作者获得了AlphaFold和ESMFold的流匹配变体，分别命名为AlphaFLOW和ESMFLOW。

作者首先检查了AlphaFLOW和ESMFLOW在采样蛋白质数据库（PDB）中蛋白质多样构象的能力。为此，作者构建了一个测试集，包含在AlphaFold训练截止日期（2018年5月1日）之后存入的100个蛋白质，这些蛋白质具有多个链并且存在构象异质性的证据。对于每个蛋白质，作者进行了50次预测，分别使用（1）未修改的AlphaFold/ESMFold，（2）不同程度的MSA子采样的AlphaFold，以及（3）不同程度的流截断的AlphaFLOW/ESMFLOW，以调节多样性。每组预测通过三个指标进行评估：精度——每个预测到最接近的晶体结构的平均lDDTCα；召回率——每个晶体结构到最接近的预测的平均lDDTCα；以及多样性——预测结构对之间的平均相异度（1-lDDTCα）。

图 3

100个测试目标的中位数结果如图3所示。AlphaFLOW类似于MSA子采样，增加了相对于未修改的AlphaFold的预测多样性，但精度有所降低。然而，AlphaFLOW的变体相对于MSA子采样展现出显著优越的帕累托前沿。在某些情况下，真实结构和预测系综的PCA分析可以解释这一结果：在MSA子采样中，随着输入信号的削弱，预测系综偏离了真实结构，而AlphaFLOW的预测则保持在真实构象周围，同时达到了相同或更高的多样性。在精度和召回率方面，AlphaFLOW表现与MSA子采样非常相似。有些令人意外的是，这两种方法都无法相对于基线AlphaFold显著提高总体召回率，这表明它们通常未能增加实验确定的PDB结构的覆盖率，或者（更乐观地）预测的构象变化尚未被实验观察到。

分子动力学系综的生成能力

图 4

表 1

接下来，作者评估了AlphaFLOW和ESMFLOW为ATLAS数据库中的82个测试蛋白质生成代理MD系综的能力。这些测试蛋白质与训练系综的结构重叠最小，提供了严格的泛化测试。对于每个目标，作者使用每种方法采样250次预测，并通过一系列评估探测它们与MD系综的相似性，这些评估分为三大类，难度逐渐增加：（1）预测灵活性，（2）分布准确性，以及（3）系综可观察性。除非另有说明，作者主要关注仅使用MSA输入（即没有PDB模板）生成的AlphaFLOW系综。主要结果在表1和图4中展示。

编译 | 黄海涛

审稿 | 曾全晨

参考资料

Jing, B., Berger, B., & Jaakkola, T. (2024). AlphaFold Meets Flow Matching for Generating Protein Ensembles. arXiv preprint arXiv:2402.04845.