力闻｜力文所参与的蛋白质语言模型解释性工作于PNAS发表

背景

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

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

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

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

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

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

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

研究方法

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

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

为了探索大模型中共进化信息的提取机制以及实现可解释性，文章采用了研究梯度的方式，即分析对输入的信息微小扰动。这种方法对于大模型来说，梯度的步长十分重要，小步长难以对于ESM2的输出产生影响，而增加步长可提高预测接触图的准确度，改变氨基酸类型则可获得最佳效果。

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

如图所示，这种方法本质上利用正向传播来提取模型中的二阶梯度依赖关系，并通过APC（Average Product Correction）方法去除偏差，进而将这些依赖关系转化为接触图信号。此方法与语言模型自身额外参数实现的接触预测（即ESM contact-regression head）、以及从MSA（多重序列比对）构建的马尔可夫随机场（MRF）或多元高斯模型中提取的共进化权重相比，属于一种完全无监督的提取方式。

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

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

   
   

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

   
   

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

   
   

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

探索结果

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

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

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

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

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

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

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

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

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

为了理解共进化信息如何被提取出来的以及这个过程中什么信息最重要，本文对不同间隔的已知二级结构片段配对之外的序列信息使用了3种策略进行掩码处理，包括解开目标配对的周围氨基酸，随机解开掩码，以及在目标配对周围以外随机随机解开掩码。

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

• 逐步解开接触点附近的掩码比起随机解开掩码，更有利于模型提取出接触的共进化信息。无论是对于单个片段还是对于相隔较远的二级结构片段，当给予其接触位点附近残基时，均可有效提回复接触预测。

  
  

• 当上下文长度提升至某个阈值时，回复率会迎来突变式提升。以淀粉结合蛋白 SusD为例，模型无法通过周围13个残基预测位置225-421的𝛼-螺旋之间的接触，但是在周围14个残基时就可以。该现象具有普适性，对于不同间隔的配对片段，至少有64%在增加周围一个残基后，回复率突然提升了 0.5 以上。

  
  

• 另外上图C的蓝色方框中出现了一些额外的接触信息，这些区域位于掩码区，本身是不包含氨基酸类型信息的，这也许额外挖掘出了蛋白中的基序信息。

  
  

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

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

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

力评

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

计算机科学中的可解释性研究方法，如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公众号  标星

在文末右下角点击  在看

给本文作者  点赞