力文所携手MIT发布突破性蛋白质大模型解释性研究

稿件来源

雅克

目前力文所与MIT Ovchinnikov课题组的合作研究取得重大进展，发表了一篇关于蛋白质大模型学习机制的论文。

此项研究，曾在著名蛋白质设计社区Rosetta Common进行过内部分享，主要探讨了解释蛋白质语言模型预测蛋白质结构的机理，特别是采用了创新的无监督Jacobian共进化提取方法。

该方法在与传统线性模型和监督学习大模型的对比中展示了卓越性能。文章还通过掩码分析揭示了ESM2模型在进行接触预测时，如何通过识别片段配对来学习序列间的共进化依赖。

这一方法不仅强调了理解蛋白质大模型基本机制的重要性，而且通过可视化大模型的无监督共进化信息，为评估和设计大模型提供了更高的可靠性，对未来大模型的分析应用具有潜在的重要价值。

01

背景

至今，蛋白质结构预测领域已经经历了显著的进步，特别是以AlphaFold2为典型的模型，它采用多重序列比对（MSA）作为输入，在同源序列信息丰富的情况下，能够实现高精度的蛋白质结构预测。

为了解决依赖MSA的限制，一系列基于大模型的预测方法相继被提出，这些方法能够仅从单一蛋白质序列中预测其结构，例如依赖于语言模型ESM2的ESMFold。

然而，尽管这些蛋白质语言模型在结构预测方面取得了成果，其具体的预测机制还未完全明确。目前，关于这些模型是真正理解了蛋白质折叠的内在物理机制，还是主要依靠序列间的共进化信息进行预测的问题，在学术界仍然存在广泛的讨论和争议。

在探讨蛋白质结构预测的现有问题时，本研究提出了三个主要假设：首先，ESM2可能在物理层面上学习到了蛋白质折叠的能量函数；其次，语言模型可能记忆了特定的蛋白质结构模式，并依此进行预测；最后，ESM2可能已经掌握了基于序列motif和序列间关系的配对依赖关系。

针对第一个假设，Ovchinnikov团队在AF2Rank论文中曾有过对AlphaFold模型学习局部势能能力的探讨，AlphaFold显示出在不依赖MSA的情况下对蛋白质构象进行有效打分的能力。然而，对于ESM2是否具备相似能力，本研究持保留态度。

这种怀疑主要基于两个观察：首先，ESMFold论文中指出，模型性能与其训练集中序列近邻数量呈现高度相关性；其次，本研究发现ESM2在预测同源异构体序列结构时出现误差，推测这可能是由于模型错误地应用了共进化信息。

基于以上观察，本研究倾向于认为，大模型在蛋白质结构预测中的学习模式可能实际上更接近第三个假设，即基于序列motif和序列间关系的配对依赖关系。这一点对于理解和优化蛋白质结构预测模型的机制至关重要。

02

研究方法

本研究首先系统地构建了一个蛋白质异构体数据集，该数据集包括从已发表的文献中收集的18个异构体序列及其对应的全长蛋白质序列。

接着，作者采用了多种结构预测方法对这些异构体序列进行了结构预测。

为了探索大模型中共进化信息的提取机制，本研究采用了Categorical Jacobian方法。该方法通过对每个氨基酸残基进行突变为其他氨基酸的操作，计算这些变化对模型输出logits的影响，从而提取共进化信息。

如图所示，这种方法本质上利用正向传播来提取模型中的二阶梯度依赖关系，并通过APC（Average Product Correction）方法去除偏差，进而将这些依赖关系转化为接触图信号。

此方法与语言模型自身额外参数实现的接触预测、以及从MSA（多重序列比对）构建的马尔可夫随机场（MRF）或多元高斯模型中提取的共进化权重相比，属于一种完全无监督的提取方式。

在上述分析方法的基础上，本文进一步进行了多个掩码实验，以深入探究大模型如何提取共进化信息：

1. 对Jacobian预测为有接触的位点附近的氨基酸残基进行序列掩码，观察模型在逐步解开附近序列时接触预测的恢复情况。

2. 针对二级结构相关分析，选择距离较远的二级结构片段对，进行相似的序列掩码实验，观察模型在逐步解开两侧序列时接触预测的恢复情况。

3. 调整语言模型的编码，包括起始结束编码、位置编码等，以观察整体信息对于局部共进化信息提取的影响。

通过这些实验，本研究旨在揭示大模型在蛋白质结构预测中的工作机制，特别是它们如何处理和利用共进化信息。

03

探索结果

本研究首先分析了蛋白质语言模型在对蛋白异构体的结构预测问题，这些模型倾向于将异构体预测为完整蛋白质结构中的某个特定片段，从而导致不合理的预测结果。

如图所示的肌红蛋白异构体例子中，AF2、OmegaFold和ESMFold等三种结构预测模型均能很好地将异构体与全长蛋白质中相应的片段对齐。

然而，这些预测中多个疏水残基被暴露在一个裂隙中，而在完整形式的肌红蛋白结构中，这样的裂隙是不存在的。

这种现象在其他的异构体案例中也有所体现。基于这些观察，作者推断ESM2等模型在结构预测时主要依赖共进化信息，而不是基于整体的势能函数进行预测。

这一发现对于理解蛋白质语言模型在结构预测中的工作机制具有重要意义，同时也提示我们在处理异构体结构预测时需要考虑到这种模型的局限性。

本研究中采用了Categorical Jacobian方法来提取蛋白质语言模型中的共进化信息，并发现该方法在仅使用单序列作为输入的情况下表现出色。

在ESM2模型中，利用Jacobian方法提取的共进化权重在预测蛋白质间接触点的精度上优于传统的线性模型，并且与在ESM2上专门训练的接触预测头所预测的接触精度相当。此外，Jacobian方法的性能与模型的参数量密切相关。

如上图所示，文章对比了使用Jacobian方法和线性模型提取的接触图，结果表明两种方法提取的共进化信息在一致性上有较高的相似度。

这种比较为我们提供了一种无监督的方式来评估语言模型在提取共进化信息方面的效能。总体而言，Jacobian方法为理解和评估蛋白质语言模型在共进化信息处理上提供了一个有效的工具。

通过序列掩码实验，本文发现：

• 逐步解开接触点附近的掩码比起随机解开掩码，更有利于模型提取出接触的共进化信息。如上图所示的2个例子，无论是对于单个片段还是对于相隔较远的二级结构片段，当给予其接触位点附近残基时，均可有效提回复接触预测，并且当上下文长度提升至某个阈值时，回复率会迎来突变式提升。

• 修改位置编码比如删掉掩码所在的位置信息，实际上对模型表现影响有限，作者归因于语言模型进行接触预测时不需要考虑完整fold信息，也不需要类似inpainting任务一样预留位置。

• Beta片之间比起螺旋之间接触预测所需的上下文序列信息更少。作者归因于Beta片需要更少的残基形成。这一点与扩散模型中往往更容易形成螺旋的表现异曲同工，螺旋的连续程度更高，上下文自由生成的情况下比起离散的Beta片片段更容易连接起来形成。

总体来说，本文的实验结果表明ESM2模型在预测接触时依赖于局部序列上下文的信息，而不是依赖于完整的蛋白质fold信息。

语言模型学习到的主要是基序片段的配对信息，其中的序列长度大约在20-40个残基范围内。尽管如此，本文认为不能完全排除语言模型可能学习到了一些整体fold信息的可能性。

另外本文基于前文提到的假设3，即大模型学习到了配对依赖关系，展望了未来蛋白大模型的开发，想要学到全部共进化信息，大概需要 20000（UniProt的家族数量）* C(256，2) （序列平均长度中的氨基酸对数量）* 20*20 （氨基酸对的可能性） = 2610亿的参数，这个数量和目前最大的蛋白大模型只有2.6倍之差，ESM2-3B模型可以被粗略认为学到了平均4对的相互作用信息，蛋白大模型开发未来可期。

04

力评

本研究深入探讨了蛋白质大模型的可解释性，这是理解模型如何学习和预测蛋白质结构的关键。

计算机科学中的可解释性研究方法，如Grad-CAM、特征分析、注意力分析等，已经在揭示模型内部工作机制方面取得了显著进展。特别是早期的研究已经开始将注意力机制与蛋白质的接触图、结合位点等结构特征联系起来。

Jacobian分析方法则从蛋白质的二阶信息和梯度角度出发，提供了另一种解释模型工作机制的视角。与传统的计算机方法相比，经过去除偏差的Jacobian方法提供了更高的可读性和直观性。

此外，蛋白质大模型的可解释性分析不仅揭示了模型的当前短板，也为未来的研究指明了方向。

例如，大模型中包含的片段信息如何服务于整体蛋白质fold，同一家族中不同构象是如何通过片段信息进行学习的，如何更高效地利用大模型进行蛋白质评估和设计等，以及如何更有效改进大模型架构。

综上所述，本研究不仅对当前的蛋白质模型提供了深入的理解，也为未来在蛋白质设计和模型优化方面的研究提供了新的视角和方法。

文献链接

Protein language models learn evolutionary statistics of interacting sequence motifs, Zhidian Zhang, Hannah K Wayment-Steele, Garyk Brixi, Haobo Wang, Matteo Dal Peraro, Dorothee Kern, Sergey Ovchinnikov, bioRxiv 2024.01.30.577970;

2024丨力文所

EXPLORE EVOLUTION

DECIPHER LIFE

©️ 力文所原创内容，未经许可转载必究。

欢迎给力文所LEVINTHAL公众号  标星

在文末右下角点击  在看

给本文作者  点赞