阅读笔记--Protein PreTrain

1、预备知识

1.1 Longformer

Longformer是一种可高效处理长文本的模型，传统Transformer-based模型在处理长文本时存在一些问题，因为它们均采用"我全都要看"型的attention机制，即每一个token都要与其他所有token进行交互，无论是空间还是时间复杂度都很高。Longformer则改进了Transformer传统的self-attention机制。具体来说，每一个token只对固定窗口大小附近的token进行local attention（局部注意力）。并且Longformer针对具体任务，在原有local attention的基础上增加了一种global attention（全局注意力）
三种注意力机制：

• 滑窗机制：对于每一个token，只对其附近$n$个token计算attention。作者认为根据应用任务的不同，可以对Transformer每一层施以不同的窗口大小。
  
• 空洞滑窗机制：对每一个token进行编码时，普通的滑窗机制只能考虑到长度为$n$的上下文。作者进一步提出空洞滑窗机制（实际上空洞滑窗是CV领域中很早就有的一项技术,比如空洞卷积），在不增加计算负荷的前提下，拓宽视野范围。在滑动窗口中，被attented到的两个相邻token之间会存在大小为$d$的间隙，因此每个token的视野范围便增大了。实验表明，由于考虑了更加全面的上下文信息，空洞滑窗机制比普通的滑窗机制表现更佳。
• 融合全局信息的滑窗机制 (Global+sliding window)
  我们知道BERT类的语言模型在fine-tune时，实现方式略有不同。比如，对于文本分类任务，我们会在整个输入的前面加上[CLS]这个token；而对于QA任务，我们则会将问题与文本进行拼接后进行输入。在Longformer中，作者也希望能够根据具体任务的不同，在原本local attention的基础上添加少量的global attention。例如，在分类任务中会在[CLS]初添加一个global attention（对应下图第一行第一列全蓝）；而在QA任务上会对question中的所有token添加global attention。如下图所示，对于添加了global attention的token，我们对其编码时要对整个序列做attention，并且编码其它token时，也都要attend到它
  

1.2 BPE介绍

BPE(byte pair encoder)字节对编码，也可以叫做digram coding双字母组合编码，主要目的是为了数据压缩。

作者准备了两个来自STRING的蛋白质序列对的预训练数据集，其中一半是强结合对（以0-1,000的置信度表示，结合相互作用评分至少为700），另一半是随机对。

2、相关工作

• $UnirRep$是第一个应用深度学习从未标记序列中获取蛋白质表示的研究。它以自回归的方式对两千两百多万个Pfam序列进行了乘法LSTM的预训练，来学习每个氨基酸1900维的表示。
• Bepler等人的工作，根据蛋白质序列对预测蛋白质的二级结构，同时通过预测这些序列对内蛋白质序列的相似性以及单个序列内氨基酸之间是否接触的关系来共同预训练双向LSTM模型。
• 最近的蛋白质语言模型致力于利用上下文相关的语言模型架构：ELMO与Bert。与ELMO相关的工作有：P-ELMO,SeqVec,PLUS-RNN。与BERT相关的工作是ESM,TAPE,ProtTrans。
• 与本文的工作类似，ProBerta也是以BPE子词表示来进行训练的。并且BPE子词表的大小都为10K。但本文所用于预训练的数据集以及模型都更大。

3、方法

3.1 方法概述

作者以Transformer为架构，训练了两种蛋白质序列预训练的策略：

• 单个蛋白质序列（对应的数据集：Pfam，String，String LF，SwissPort）
• 成对的蛋白质序列（这两个序列要么具有较强的结合性要么是随机匹配的）（对应的数据集：String2Seq，StringLF2Seq）

作者仅以掩码语言模型（MLM）来预训练：随机选取输入序列中15%的tokens做掩码，并以对应的预测的二进制交叉熵的值来报告。所有的预训练模型在Pfam holdout数据集进行测试。（Pfam holdout数据集：根据蛋白质家族对数据集内蛋白质进行分割，一共包括六个Pfam家族：- PF01112，PF03417， PF03577，PF14604， PF18346，PF18697）

• CE：Cross Entropy loss（交叉熵损失函数）

模型架构如下图所示：

• $E_i与T_i$分别代表着token$i$的输入和上下文嵌入。
• $[CLS]$使用以分类任务输出的特殊$token$，而$SEP$则是将两个非连续的序列分开的特殊$token$。
• 与字符级嵌入相比，使用10k子词词汇表（a）进行字节对编码（BPE）可使输入的蛋白质序列的平均长度缩短为字符集嵌入的时蛋白质序列平均长度的36％。
• 该模型是由STRING DB数据集中的结合与非结合蛋白序列的混合序列仅适用MLM目标函数来预训练的。
  
  不同数据集对应的模型的参数如下：
  
  为了容纳2048个标记的蛋白质序列，作者采用了Longformer架构，滑窗的大小为512对于序列分类任务，作者让$[CLS]token$执行全序列的自注意机制计算。

3.2 BPE压缩表示

首先介绍蛋白质序列的字符级表示：蛋白质序列通常表示为字母串，列出了氨基末端开始至羧基末端的氨基酸。三个字母代码或单个字母代码可以用于表示20种天然存在的氨基酸，以及混合物或不确定的氨基酸（类似于核酸符号）。

一个长度为2521的蛋白质序列，如下图所示：

与TAPE中将氨基酸直接视为token（总共有25种token）不同的是作者使用了BPE对输入的序列进行编码，并将subword词表的大小设为10K，即总共有10K个子词。这种方法非常有效的压缩了蛋白质序列的长度。和采用字符级表示的蛋白质序列相比，其平均长度从360缩短为130，对于SwissPort中的蛋白质序列，仅有1.58%的蛋白质序列长度超过了512。如果观察UniProKB_2020_06数据集，我们可以发现最长的蛋白质序列由超过3万个氨基酸组成，长度超过1000的蛋白质也有很多，现有的工作大多数会将过长的蛋白质序列（比如超过500或1000）丢弃，从而损失了有价值的信息，其次，对于序列对分类任务，例如PPB预测，在模型长度约束内拟合两个序列至关重要。下图为两种表示方法下，蛋白质序列的长度：

3.3 数据集

作者在STwissProt,Pfam,STRING三个数据集上进行了预训练。并且对于成对输入的预训练模型，作者将STRING数据集分为两个数据集：String2Seq，StringLF2Seq
其中String2Seq数据集每对中蛋白质序列的长度之和小于512，而StringLF2Seq数据集中每对中蛋白质序列的长度之和小于2048。
不同预训练数据集的超参数和语言模型预训练结果如下图所示：

3.4 Fine-tuning任务

预训练任务的评价标准是针对MLM的交叉熵损失值，但预训练模型的最终评价则来自于下游蛋白质预测任务的表现。作者尝试了三种类型的下游任务：

• 单个的蛋白质序列分类 ：首先，对于单个蛋白质序列分类，作者将个任务特定的输入送入经过预训练的模型中，并通过对补充输出层的单次修改来获取来自$[CLS]$的表示，从而得出输入的蛋白质序列对应到各种类别的概率。对应的任务有：
• 成对蛋白质序列分类 ：两个被$token[SEP]$分隔开的蛋白质序列的关系也是根据$token[CLS]$的表示来判断。对应的任务有：蛋白质-蛋白质结合预测，T细胞受体结合预测。
• token级分类（序列标注）：除了$[CLS]$，所有的token的表示都要经过一个SoftMax层以给出token属于各类的概率。

3.4 Fine-tuning数据集

• 蛋白质-蛋白质结合预测：2M的短蛋白质序列对（<512 tokens），0.667M的长蛋白质序列对(>512tokens and <2048tokens)，且数据集中结合对于非结合对的数量比为$1:1$
• T细胞受体结合预测：该任务为给定的一个T细胞受体（TCR）预测与TCR结合的最可能的肽片段。数据集包括124486个TCR与蛋白质序列的序列对，验证集与测试集分别含有13832对，10K对序列。
• 远程同源性预测：预测蛋白质的折叠类别（对应有1995种折叠类型）
• 水溶性预测：预测蛋白质是否可溶的二元任务- 训练集中有28972个可溶序列以及40448个不可溶序列
• 稳定性预测：类似于TAPE，是一个回归任务，其中每个输入蛋白x被映射到一个标签y∈R,测量在最极端的情况下，蛋白x将其折叠保持在浓度阈值以上（代表内在稳定性）。
• 二级结构预测：

4、实验结果

4.1 具有蛋白质结合对的MLM预训练优于单个序列

4.2 序列对预训练模型在微调任务方面优于单个序列模型

预训练完成后的交叉熵损失小并不意味着微调任务的效果会更好。微调任务的结果如下图所示：

未来的工作

• 训练一个含有20-25个氨基酸词汇表的蛋白质模型是很简单的，我们把每个字符作为一个标记，并在测试集中预测新的标签。但是，如果我们有来自BPE的10k令牌，并且假设一组更长的、多字符的BPE令牌将携带更多的结构信息来预测二级结构。为了充分利用10k词汇表，我们必须克服一个问题：输入序列标记（在BPE编码之后）将破坏标签的连续性。二级结构标签通常形成2-7个氨基酸长序列，我们称为标签岛。作者建议探索使用BPE标记生成器对所有训练数据进行预标记，并创建一个新的具有标签岛意识的词汇表，并通过删除新词汇表中未出现的标记来修剪BPE词汇表。然后，在评估过程中，我们希望测试序列能够正确解析，并且在使用新的修剪词汇进行标记后，可以保留标签的连续性。 另外，可以将序列分成三元组。
  
  

• 补充的内在预训练目标（如掩蔽语言建模，MLM），什么对二级结构和PPI(蛋白质交互)预测微调任务有利。有的工作提出蛋白质家族预测，但另有工作表明这一任务较为简单，准确率很容易达到90%以上。

• 研究Longformer预训练提供的巨大模型输入空间（2,048个令牌）的新用途；总的来说，如果我们假设自我注意能够完全理解蛋白质的二级和三级结构，那么我们可以预测两个（1022个）、四个（511个）或八个（255个）蛋白质之间的结合。但首先，我们仍然需要在所有下游任务上验证Longformer。

• 将TCR结合预测作为一个问题回答任务：T细胞受体结合预测是非常困难的模型。我们建议在NLP中引入一个类似于问答的新的蛋白质序列任务。我们可以训练完整的TCR序列，而不仅仅是CDR3，还有一个额外的目标，即确定短CDR3在长TCR序列中的确切范围

• 我们假设，通过对蛋白质序列对的训练，自我注意的大脑会跨层学习哪些序列模式导致了两序列结合，并识别出哪些模式构成了结合位点。通过可视化层和自我注意头的组合，我们试图检索正确的结合位点。本文利用RXNMapp来发现结合位点，并给出预测结合位点的概率，而不是简单结合位点可视化。

• 新的子词编码策略。作者提出了一种方法：随机的删去子词表中的子词，来以不同的方式表示同一个蛋白质序列。