基于Transformer预训练的蛋白语言模型是无监督的蛋白质结构学习器

作者 | 何文嘉 编辑 | 李仲深

• 摘要
• 1 介绍
• 2 背景
  • 2.1 蛋白质语言模型
  • 2.2 多序列比对 (Multiple Sequence Alignments, MSAs)
  • 2.3 评估指标
• 3 模型分析
  • 3.1 ESM-1b
  • 3.2 Gremlin
• 4 结果分析
  • 4.1 普通情况
  • 4.2 限制监督
  • 4.3 MSA集成
• 5 总结
• 6 参考文献

‍

‍摘要

无监督接触预测 (Unsupervised Contact Prediction) 是在蛋白质结构测定和设计过程中揭示蛋白质物理、结构和功能约束的核心。几十年来，主要的方法是从一组相关序列中推断进化约束。在过去的一年里，蛋白质语言模型已经成为一种潜在的替代方法，但目前性能还没有达到生物信息学中最先进的方法。本文证明了Transformer的注意图 (Attention Map) 能够从无监督语言建模目标中学习蛋白质序列中各个残基之间的接触距离。我们发现，迄今为止已经训练过的最高容量模型已经优于当前最先进的无监督接触预测的工作流程，这表明过去这些繁琐的工作流程可以用端到端模型的单向传递工作流程来代替。

关键词：无监督接触预测，Transformer，预训练，蛋白质语言模型

1 介绍

无监督接触预测任务是当前用于预测蛋白质结构的基本步骤，通常先预测蛋白质残基接触图，再将蛋白质残基接触图用于蛋白质结构预测。过去的在生物信息学中无监督接触预测任务的标准的做法通常会包括各种特定的工具和数据库，整个任务的工作流程相当繁琐且复杂，而本篇工作则提出了一种用于取代多阶段工作流的单向工作流，这种工作流是基于端到端的预训练蛋白质语言模型的。

端到端的优势很明显：

1. 取代了资源和时间消耗昂贵的查询，对齐和训练等步骤，加速特征提取过程。
2. 对所有的蛋白质家族共享参数，通过捕获数百万个进化多样性和无关序列的共性来实现泛化。

• 过去的方法通常是通过对相关和对齐的序列的分析来利用蛋白质序列的进化信息。
• 目前的方法是使用类似于NLP中的语言模型来捕捉序列中的内在隐含性质。

文章主要是将用Transformer训练的蛋白质语言模型ESM-1b和传统的对数线性模型Gremlin来进行比较和分析，从而得出相关的结论，提高Transformer模型的可解释性。

本文属于预训练模型 (PTMs) 的一个应用，不了解PTMs的读者建议先去查阅相关资料，了解一下预训练任务如LM，MLM等和经典的预训练模型如Transformer，BERT等。本文需要用到的先行知识有预训练任务MLM，注意力机制Self-Attention和经典模型Transformer。

2 背景

2.1 蛋白质语言模型

文中的蛋白质语言模型ESM-1b的预训练任务是MLM (Masked Language Model)，所以本文中PLM (Protein Language Model) 默认指代MLM。既然是MLM，其实训练方式就和BERT很相似，按BERT去理解就好，只不过是把输入序列从自然语言文本转换为氨基酸字母序列而已。

2.2 多序列比对 (Multiple Sequence Alignments, MSAs)

多序列比对由一组进化相关的蛋白质序列组成。由于真正的蛋白质序列很可能有插入、缺失和替换，因此通过在所有序列上最小化编辑距离 (Levenshtein Distance) 的距离类度量来对序列进行对齐。在实践中，使用了启发式对齐方案。潜在的误差来源包括

1. 预滤波的启发式方法失效
2. 相关序列未被检测到
3. 未能找到最佳对齐
4. 替代矩阵和间隙惩罚的次优性，即找不到全局最优解

2.3 评估指标

对于长度为

的蛋白质，我们评估了长度为前

、

和

的序列的接触，在评估上考虑了三种接触范围中的预测精度，分别是短程接触范围 （

）、中程接触范围 （

）、远程接触范围 （

），其中

指蛋白质序列中的位置索引。一般来说，所有接触都提供了关于蛋白质结构相互作用的信息，短程接触对二级和局部结构的帮助较大，而较长距离接触对确定全局结构的帮助较大。

3 模型分析

3.1 ESM-1b

研究动机

使用Transformer来预测残基接触是因为Transformer模型内部产生的注意力图天然地对应了序列中各个残基之间的信息。形式上，注意力图和残基接触图都可以表示为一个

的方阵。

训练任务

因为ESM-1b采用的预训练任务是掩码语言建模 (MLM)，所以损失为：

即给定训练数据集

，在数据预处理时，每个序列样本

上的多个位置对应的氨基酸残基会被特殊符号

替换掉，利用非

位置对应的氨基酸残基的信息来预测每个序列样本

中所有

位置对应的氨基酸残基字母。

模型结构

ESM-1b是先前工作ESM-1的改进版本，关于ESM-1的内容读者可参考论文[1]

ESM-1b的框架如下图所示：

Transformer首先通过掩码语言建模 (MLM) 对来自大型数据库 (Uniref50) 的序列进行预训练。一旦完成训练，注意力图就能够被提取出来，先后进行对称化 (Symmetrization) 和平均乘积 (Average Product Correction, APC) 校正两个操作将注意力图转换为需要的形式，然后进行回归任务，在每个氨基酸对

上独立应用带有

正则化的

回归。

回归是在少量的

个蛋白质 (

) 上训练的，以确定哪些注意力头 (Attention Head) 包含了丰富的信息。在训练时，我们只训练

回归的权重，并不会反向传播整个模型，也就是在逻辑回归的分支上，梯度值反向传播至注意力图处，到上图中的 “Stop Gradient”部分便停止梯度的反向传播。

在测试时，输入序列的接触预测可以完全在GPU上通过一个前向传播进行。整个预测工作流程可以在单个前向传播中进行，为蛋白质接触预测提供端到端的工作流程，而不需要从序列数据库中进行任何检索步骤。

3.2 Gremlin

对于一组训练序列

，Gremlin优化了以下伪似然损失：

其中单个位置被掩盖并从其上下文中预测其真实的氨基酸残基，其实就是利用上下文预测中心词，本质上和Word2Vec中的连续词袋模型CBOW一样。由于对Gremlin的输入是对齐的，所以它们有统一的长度

。

Gremlin以MSA作为输入。输出预测的质量在很大程度上取决于MSA的构建。

4 结果分析

4.1 普通情况

多种方法的平均精度对比如 Table 1 所示：

作者在具有15051个蛋白质的trRosetta训练数据集中评估的预训练的蛋白质语言模型，其中在trRosetta数据集中去除了43个序列长度大于1024的蛋白质，因为ESM-1b的训练上下文大小设置为1024了。在这些序列中，当我们试图使用ESM训练集构造MSAs时，Jackhmer在126个蛋白质中失效，也就是构造失败，所以最终只留下了14882个蛋白质序列。在最后，作者保留20个序列用于训练，20个序列用于验证，剩余的14842个序列用于测试。

在 Table 1 和 Table 2 中，所有Transformer模型的接触预测器都是在预训练完成后用Logistic回归对20种蛋白质进行训练得到的。作者发现，虽然ESM-1b模型只用了20个蛋白质序列进行训练，但其在短、中、远距离范围的接触预测方面均比Gremlin具有更高的精度。

4.2 限制监督

虽然语言建模目标MLM是完全无监督的，但

回归是用少量的监督例子来训练的。在本节中，作者研究了结果对这种监督的依赖性，提供了在无监督阶段确实学习到了接触预测所需要的信息，而

回归只需要提取相关的特征。限制监督后的结果对比如 Table 2 所示：

Top Head

在这里，使用

回归来确定最重要的头部。一旦它们被选择，我们从

回归中丢弃权重，并简单地平均与top-k权重值对应的注意头。通过从ESM-1b中获取单个最佳头部，在相同的数据下，已经接近Gremlin的性能；平均前5名的注意力头便可以使预测性能优于Gremlin；平均前10名的注意力头便可以使预测性能优于所有其他Transformer模型 +

回归，并接近Gremlin在给定最优MSAs的情况下的表现。

总体来看，这说明了ESM-1b比Gremlin的最佳性能还要好。

Low-N

我们考虑的第二个变化是限制提供给

回归的监督样本的数量。作者发现，即便只有一个单一的训练例子，ESM-1b在长度为L的序列的远距离范围预测中，平均精度已经达到39.2，接近SOTA，相同的情况下，这与Gremlin在统计学上已经是无法区分孰优孰劣了（假设检验，p>0.05）。仅使用10个训练示例，该模型在所有度量上都优于Gremlin。

总体来看，这说明了在预训练后，模型只需要几个样本，很快就可以适应残基接触预测任务，一方面验证了模型动机的合理性，另一方面说明了注意力图中真的蕴含了不少有助于残疾接触预测的信息。

MSA Only

最后，我们考虑只用MSAs而不是ESM-1b中的Transformer预训练模型来对

回归进行监督学习，而且和Gremlin基线使用的训练数据相同，结果是预测精度和Gremlin差不多。

说明让ESM-1b表现出优秀的性能的关键因素在于Transformer产生的注意力图而不是

回归模型本身。

4.3 MSA集成

Transformer模型从根本上说是单序列模型，但可以通过将来自多个序列的预测进行集成并对齐来进一步提高性能。Table 2 的最后三行就是实验结果，整体上性能有一点提高。

5 总结

本文用无监督目标在Transformer上训练的蛋白质语言模型在它们的注意图中学习到了蛋白质序列三级结构的相关信息。通过稀疏（L1正则化）的

回归可以从注意力图中提取残基接触的有用信息。另外，也发现了不同的注意力头部专门负责不同类型的接触。限制监督部分的分析证实了接触预测所需要的信息是在无监督阶段学习到的，

回归只需要提取模型中代表接触的那部分信息即可。

原文中还对模型进行了很多细致的说明和分析，比如参数设置，鲁棒性分析，语言模型的困惑度分析等，但鉴于篇幅，小编只能对论文的整体内进行阐述，希望了解更多的读者可以参考论文原文[2]。

这些结果对蛋白质结构的确定和设计有一定的意义。这个工作第一次表明，蛋白质语言模型可以优于目前性能最好的传统的无监督接触预测方法，这个意义更加重大。而且很多研究已经表明预训练提取的特征很多时候比传统的手工特征要好得多，相信以后用预训练自动提取序列特征很有前景。

6 参考文献

[1]Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Demi Guo, Myle Ott, C. Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological Structure and Function Emerge from Scaling Unsupervised Learning to 250 Million Protein Sequences. bioRxiv, pp. 622803, 4 2019. doi: 10.1101/622803. URL https://www.biorxiv.org/content/10.1101/622803v3.

[2]https://openreview.net/forum?id=fylclEqgvgd. (Under Review)‍