BERTology遇上生物学：深入理解蛋白质语言模型中的注意力机制

今天给大家介绍Salesforce Research投稿在ICLR2021上的一项工作。在很多蛋白质任务中，Transformer模型被用来学习蛋白质的有效表示形式，但是这些表示形式在可解释性方面提出了挑战。针对这个问题，作者通过分析Transformer的内部工作原理，探索了Transformer模型如何识别蛋白质的结构和功能特性，并展示了注意力机制和蛋白质结构之间相互作用的三维可视化。实验结果表明，作者的发现与已知的生物学过程相吻合，并提供了一个工具用以辅助蛋白质工程和合成生物学。

1

介绍

蛋白质是控制生物和生命本身的基本大分子，对蛋白质的研究有助于理解人类健康和发展疾病疗法。测序技术成本的下降使我们能够形成大量天然蛋白质数据库，这些数据库中的信息十分丰富，已经可以开发出优秀的基于序列的深度学习方法。由氨基酸序列组成的蛋白质某种程度上可以视为一种语言，可以使用针对自然语言开发的神经网络模型对它们进行建模。Transformer 模型引发了文本无监督学习的革命，并在蛋白质序列建模领域有着相似的作用。然而，Transformer的强大性能是以可解释性为代价的，这种缺乏透明度的现象可能掩盖了诸如模型偏差和虚假相关性之类的潜在问题。

为此，作者分析了Transformer蛋白质模型中的核心架构-注意力机制，并提出了一系列捕获蛋白质序列独特功能和结构特征的分析方法。作者的分析表明，注意力机制捕获了蛋白质的高级结构特性，连接了在三维空间中接近但在基础序列中相距甚远的氨基酸，并聚焦于蛋白质的关键功能组成部分——结合位点。此外，作者指出注意力机制捕获了氨基酸的取代特性，并随着网络深度的增加逐步构造了蛋白质结构和功能的更高级的表示。

2

实验

作者使用来自TAPE的BERT-Base模型，该模型已在3100万个蛋白质序列的数据集上进行了预训练。该模型接受氨基酸x =（x1，...，xL）的序列作为输入，并输出连续嵌入向量z =（z1，...，zL）。该模型包括一系列编码层，每层都包含多个multi-head attention。注意力权重定义了每个token对当前token下一层表示的影响,BERT-Base模型共有12层和12个头部，总共产生144种不同的注意力机制。作者用-表示一个特定的层-头对，例如 head3-7代表第三层的第七头。

注意力分析

作者分析了在token级别（二级结构，结合位点）和token对级别（接触图）上，注意力如何与各种蛋白质特性保持一致。针对token对情况，作者为属性f定义了一个指标函数f（i，j），如果该属性存在于token对（i，j）中（即，如果氨基酸i和j处于接触状态），则返回1。否则为零。然后，计算在数据集X上与f对应的注意力比例

αij（x）表示从i到j对输入序列x的注意力，pα（f）表示针对属性f的注意力比例。

探测任务

除了分析注意力之外，作者还使用诊断分类器来探测输出，以确定它们包含有关上述属性的哪些信息。对于分类器，左使用单个线性层，后跟softmax，并冻结原始模型的权重。对于token级别的探测任务（绑定站点，二级结构），我们将每个token的输出向量直接输入分类器。对于token对探测任务（联系图），按照TAPE的实现，通过两个token输出向量的乘积来构建成对特征向量.对于不同的任务，作者采用不同的评价标准。

数据集

作者使用TAPE中的两个蛋白质序列数据集进行分析：ProteinNet数据集和二级结构数据集。两个数据集均包含氨基酸序列，其中ProteinNet还标有每个氨基酸的空间坐标，用于生成接触图，二级结构数据集在每个序列位置标有二级结构。对于氨基酸和接触图的分析，我们使用ProteinNet；对于二级结构和结合位点的分析，使用二级结构数据集。对于结合位点分析，作者采用从Protein Data Bank中获得了可用的token结合位点注释。为了分析注意力，作者使用了来自各个训练分组的5000个序列的随机子集，为了训练探测分类器，使用全部训练集来训练模型，并使用验证集来进行评估。

3

实验分析

注意力头专注研究某些类型的氨基酸

作者计算了每个头集中在特定类型氨基酸上的注意力比例，该比例是在5000个序列的数据集上平均得到的，这些序列的总长度为1,067,712个氨基酸。作者发现，对于20种氨基酸中的14种，存在一个关注头，将超过25％的注意力集中在该氨基酸上。下图显示head 1-11将其总注意力的78％集中在氨基酸Pro上，head 12-3将27％的注意力集中在Phe上。

注意力与替代关系一致

为了进一步探究注意力与氨基酸的关系，作者分析了注意力如何与氨基酸现有的结构和功能特性度量相关。通过计算每对氨基酸之间的关注相似度，然后将该指标与基于取代矩阵的成对相似度进行比较，来评估关注是否跟踪相似属性。为了衡量注意力的相似性，作者计算了每个氨基酸在注意力头中获得的注意力比例之间的皮尔森相关性。例如，为了测量Pro和Phe之间的注意力相似度，采用两个热图的Pearson相关性。所有这种成对相关性的值如图3a所示。作者将这些分数与图b中的BLOSUM分数进行比较，发现皮尔森相关系数为0.80，这表明注意力在很大程度上与替代关系保持一致。

注意力与一个注意头中的联系图紧密对齐

下图显示了每个头的注意力与联系图对齐的百分比。12-4的单个头部比其他任何一个头部更紧密地与接触图对齐。考虑到模型是使用掩码语言建模(Masked Language Model, MLM)目标进行训练的，其输入或训练标签中没有空间信息，因此可以识别联系是很令人惊讶的。作者认为，这种定位行为的一个潜在原因可能是接触更可能发生生化相互作用，从而限制了可能占据这些位置的氨基酸。因此，在语言模型中，了解屏蔽token的联系可以为token预测提供有价值的上下文。对于捕获这些联系关系，作者指出可能是该模型学习了蛋白质结构的基本几何形状，或者该模型可以利用由接触关系产生的统计规律性。

注意力是连接图的良好预测器

通过比较上文讨论的head 12-4中的注意力权重以及两个氨基酸接触的可能性，作者针对接触图镜像测试。通过将数据集中所有氨基酸对（i，j）的注意力权重αi，j进行分箱，并计算每个仓中所接触的成对氨基酸的比例，来估算接触概率。结果如下图所示。估计的概率与注意权重（基于每个bin的中点）之间的Pearson相关性是0.97。这表明在这种情况下，注意力权重是一个经过良好校准的估计器，可以提供对注意力的原则性解释，作为置信度的度量。

注意力集中在结合位点，尤其是在更深的层次上

下图显示了每个头部专注于结合位点的注意力比例。在大多数层中，跨头部的平均百分比显着高于结合位点的本底频率（4.8％）。该效果在模型的最后6层中最强，其中包括15个头部，每个头部将超过20％的注意力集中在结合位点上，Head 7-1最关注的是结合位点（34％）。下图中还显示了该头部靶向结合位点的估计概率，它是注意力权重的函数。作者指出token通常在序列中较远处靶向结合位点。过去的工作表明，结合位点可以揭示蛋白质之间的进化关系，并且结合位点中的特定结构基序主要限于蛋白质的特定家族或超家族。因此，结合位点提供了蛋白质的高级表征，可能与整个序列中的模型有关。

注意力集中在更深层次的更高层次的属性

如下图所示，较深的层将更多的注意力集中在结合位点和连接（高层概念）上，而二级结构（低层到中层概念）则更均匀地针对各层。探测分析类似地显示，该模型在完全编码接触图和结合位点之前首先形成二级结构的表示。此结果显示了模型必须先了解局部结构，然后才能形成高阶结构和功能的表示形式。

4

总结

本文通过将NLP的可解释性方法改编并扩展到蛋白质序列建模来建立NLP与计算生物学之间的协同作用。作者展示了Transformer语言模型如何恢复蛋白质的结构和功能特性，并将这种信息直接整合到其注意力机制中。作者分析表明，注意力不仅捕获了单个氨基酸的特性，而且还发现了更多的全局特性，例如结合位点和三级结构。在某些情况下，注意力机制还可以很好地校准模型对这些特性的预测。

代码

https://github.com/salesforce/provis

参考资料

论文链接：https://arxiv.org/pdf/2006.15222v1.pdf