阅读笔记-蛋白质序列预训练ESM

概要

          \,\,\,\,\,\,\,\,\, 数十年来，公共数据库中蛋白质序列数量的增长呈指数趋势，从而深入了解了整个生命中蛋白质序列的广度和多样性。 该数据为使用人工智能研究生物学的预测和生成模型提供了有希望的基础。 本文的重点是使单个模型适合整个进化过程中的许多不同序列。 因此，本文研究了高容量的神经网络，研究了从大规模进化数据建模中可以学到的关于蛋白质生物学的知识。NLP中有一个假设：单词的语义可以从其出现的上下文中得出。最近，基于自我监督的技术已经被证明可以实现可以在自然语言任务中提升的词义表示，并且借助更大的训练数据集，模型学习此类表示的能力会显着提高。蛋白质序列产生于一个与自然语言极为不同的过程。这些模型和目标函数是否能有效地促进自然语言跨领域的迁移还不确定。最近，自我监督已成为人工智能研究的核心方向。 与需要对每个数据点进行手动注释的监督学习不同，自我监督方法使用未标记的数据集，因此可以利用大量数据。 自我监督学习使用代理任务进行训练，例如在给定所有先前单词的情况下预测句子中的下一个单词或预测从上下文中被掩盖的单词。
          \,\,\,\,\,\,\,\,\, 本文探索了自我监督的语言建模方法，这些方法在一系列自然语言处理任务上表现出了最先进的性能，并将其应用于未标记氨基酸序列形式的蛋白质数据。 由于蛋白质序列使用的词汇量很少，只有二十个规范元素，因此建模问题与字符级语言模型相比，更类似于字符级语言模型。

数据与模型

          \,\,\,\,\,\,\,\,\, 作者采用了UniProt数据集（250M条蛋白质序列），以MLM为目标来预训练模型。并且没有采用BPE算法来处理蛋白质序列。作者随机选取了1M条序列作为验证集，作者从Uniref50里随机选取了10%聚类保留集是所含的代表性序列组成测试集，并从训练集中删掉了这些聚类保留集。作者探讨了潜在的序列多样性对预训练数据的影响。UniParc聚类显示了簇大小的幂律分布，这意味着大多数序列属于簇的一小部分。使用序列的聚类训练的结果是对掩蔽语言建模损失的重新加权，朝向更多样化的序列集。作者使用UniRef 创建三个具有不同多样性水平的预训练数据集:(i)低多样性数据集(UR100)使用UniRef100代表性序列;(ii)高多样性稀疏数据集(UR50/S)使用UniRef50代表性序列;(iii)高多样性密集数据集(UR50/D)在UniRef50集群中均匀采样UniRef100序列。
作者训练的模型如下所示：

同源性建模

          \,\,\,\,\,\,\,\,\, 进化是受约束的;如果对蛋白质的修饰导致其结构被打乱，生物体将失去这种功能并受损。因此，我们希望从数据库中提取的表亲集告诉我们，进化在哪里有自由统治，哪里有回旋的空间，哪里被完全制约。这些蛋白质系列的表示，被恰如其分地称为多序列对齐，是蛋白质结构预测管道的关键输入。直观地讲，在 3D 空间中接近的位置将共同演化，即一个位置的突变通常与另一个位置的容纳性突变同时发生。尽管直系同源基因的序列存在差异，但其结构和功能仍可能保留下来。每个点代表一个基因，每个基因都由其所属的直系同源基团着色（t-SNE降低了维数）。 直系同源基因群密集地聚集在训练的表示空间中。 相比之下，未经训练的表示空间和会标表示不能通过进化关系反映强大的组织。（b）在训练的表征空间中，与常见生物学变异相对应的基因线性相关。 基因由其直系同源基团着色，其种类由字符标签指示。 PCA在训练的表示空间中恢复物种轴（水平）和正交轴（垂直），但在未训练的或唯一字母空间中恢复。 表示形式来自在UniParc上训练的36层Transformer模型。
          \,\,\,\,\,\,\,\,\, 作者对PFAM数据集中的PF01010蛋白质家族用训练前后的模型进行相似性分析，发现经过预训练后的模型能够捕捉到相同蛋白质家族内序列之间的相似性。如下图所示：

• 深蓝色：利用训练后的模型来提取的相似性较高的蛋白质序列
• 浅蓝色：利用训练后的模型未对齐的蛋白质序列（可看作相似性较低）
• 深红色：未训练的模型来提取相似性较高的蛋白质序列
• 浅红色：未训练的模型来提取相似性较低的蛋白质序列
• 图b是利用与训练前后的模型来提取相似性极高的蛋白质序列集的表示，并用来计算序列之间的余弦相似度。可以看出经过预训练的模型更能提取到相似性较高的序列之间的相似性。而图c是利用预训练前后的模型来提取相似性较低的蛋白质序列集的表示，并用来计算余弦相似度。可以看出预训练后的模型不但让相似性高的序列的表示向量相似，更能区分出相似性较低蛋白质序列。而未预训练的模型则不具备这种能力。

蛋白质中氨基酸在特定结构或功能背景下的可互换性取决于它们的生化特性。自监督可以捕捉这些模式，建立反映生化知识的表示空间。作者将各种氨基酸的表示向量利用t-SNE算法投影到二维空间（模型为36层的transformer），如下图所示：

氨基酸的生化特性在Transformer模型的输出嵌入中表示，这里用t-SNE表示。通过无监督学习，残基被聚集成疏水基团、极性基团和芳香基团，并通过分子量和电荷反映整体组织。

蛋白质嵌入将整个蛋白质表示为高维空间中的点。蛋白质的嵌入可以通过对序列中所有氨基酸的嵌入求平均值得到，接下来作者利用t-SNE算法将同源蛋白质的嵌入投影到二维空间，来研究同源基因在这个空间中是如何表现的。

上图a显示，预训练后模型提取到的蛋白质序列的嵌入，同源的蛋白质往往能聚集到一起，而训练前模型并不具备这种能力。

之后作者进行了蛋白质远程同源性检测，发现34层transformer的与训练模型在该任务上与基于比对的HMM模型的检测水平非常接近，Fold检测这项指标上甚至超过了HMM模型。

结构预测

为了预测蛋白质的三级结构信息，作者将两个单独的线性投影拟合到序列中位置对的隐藏表示中，使它们的点积回归到一个二进制变量，该二进制变量指示这两个氨基酸在空间中是否足够的接近。对于二级结构预测，作者使用来自预训练模型的表示向量通过全连接层去预测每个氨基酸对应的8种二级结构标签。

SSP:

Contact Pre：

上图表明了，ECE越低模型越能从蛋白质序列捕获到它的空间结构信息，并且模型容量越大，捕获到蛋白质序列的空间结构的能力上限越高。

与比对特征的结合

对于蛋白质序列x,作者探讨了三种整合表征学习信息的方法：

• direct：直接将transformer的输出作为每个氨基酸的表示向量；
• avg：通过MSA将得到的多序列比对中相同位置处每个氨基酸的嵌入的平均值作为这个位置的最终表示；
• cov：在用PCA将输出嵌入降维后，使用来自x的MSA的序列之间的非中心协方差，为每对位置生成特征(感觉比较适合蛋白质三级结构预测的任务)。

结果显示在蛋白质二级结构预测的任务上，avg方法效果最好，成为了该任务的SOTA。

在蛋白质三级结构预测的任务上，cov成为了该任务的SOTA

方法与数据

数据集

训练集：

• UR100：124.0M条UniRef序列；
• UR50/S：27.1M条UniRef50序列；
• UR50/D：从UniRef50集群中平均抽样了124.9M条序列；

验证集：1M，删除了序列长度超过1024个氨基酸的蛋白质序列。

下游任务

• 远程同源性检测；
• 二级结构预测SSP；
• 三级结构预测Contace Prediction；
• 突变效应预测

Transformer

• layer normalization
• no dropout
• 除了输出层与输入层，其余全连接层均包含偏差
• q,k,v均为d维的，feed forward层是4d维的

UniParc 模型

UniRef模型

ESM-1b 模型

更大的学习率，使用了dropout，可以学习到的位置表示向量，输出层添加了LayerNorm，作者发现这些均有助于提升模型的性能。

预训练任务

Masked Language Modeling
15%token被mask，80%被mask替代，10%随机替代，10%保持不变
序列长度最大为1024，对于序列长度大于1024的序列，随机裁剪出1024个token。