论文笔记 | Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers

介绍

常规Transoformer的计算复杂度是$O(L^{2}d)$，空间复杂度是$O(L^2)$，在长序列的情况下，比如基因序列的处理，序列长度$L$太大会导致计算成本大到难以接受。

本文提出Fast Attention Via Orthogonal Random features (FAVOR)方法，在不改变原先transformer模型架构的前提下，高效地估计出attention矩阵，时间复杂度为$O(L{d}^2\log d)$，空间复杂度为$O(Ld\log d)$。

相比于其他改进transformer在长序列情况下复杂度的方法，本文方法并没有提出新的模型结构，可以方便地应用于已经预训练好的模型。

方法

常规transformer的attention的计算有两种形式，区别在于时序依赖关系的处理。

第一种称为双向点乘的attention（Bidirectional or non-directional dot-product attention），形式如下：

第二种称为单向点乘attention（unidirectional dot-product attention），形式如下：

类比cosine相似度，attention中的点乘实际上可以看作相应向量的相似性的度量，于是作者在文中对其进行泛化，基于核函数的表示给出了Generalized Attention (GA):

在GA中，核函数$K$负责query $Q$和key $K$中元素的交互，选择不同的核函数，对应不同的attention的形式。原先的点乘attention对应的GA形式为（对应着高斯核）：

如果知道了核函数的random feature map函数，核函数的计算可以表示为:

对于大多数的核函数来说，相应的random feature map都相似的形式：

特别的，对于高斯核来说，有如下的random feature map：

表示成矩阵的形式有：

现在，attention矩阵$A$表示成两个低秩矩阵的乘积，计算的时间复杂度下降了。但是如果显示计算出$A$的话，空间复杂度仍然是$L^2$的，因此在attention与值矩阵相乘的时候，先用低秩矩阵与值矩阵相乘，避免A矩阵的出现。

对于双向点乘attention来说，与值矩阵相乘的计算方式如下，注意括号决定的乘法优先级：

其中，期望的估计使用了一个样本。

对于单向点乘attention来说，要计算的目标具有时序依赖关系，可以通过前缀和的形式得到：

正交随机特征

估计核函数需要采样$W$，文中提出$W$具有一定正交性的话，效果会好。推荐的有Gaussian orthogonal matrices，random Hadamard/Givens matrices。

理论收敛性分析

$M_{opt}$最优的随机特征的个数不依赖于序列长度$L$，而依赖于特征维度$d$。$M_{opt}=\Theta (d\log d)$。

实验

上图表明，序列越长，加速越明显。在序列长1000以下，加速不太明显。

上图，使用了随机正交特征，在特征数够大的时候，近似误差较小，不使用随机特征的情况下，考虑到attention score的scale本身就是0.0x左右，有一定的误差。网络深度会加重误差传播，应用的时候不能直接迁移，需要finetune。

上图，在多层transformer下，训练效果能达到甚至超过原先的transformer，比Reformer要好。

注意

• M大小如何选取
  量级是$\Theta (d\log d)$，实际中100多近似效果较好。
• 在预训练好的transformer上，进行finetune再用，不然准确率很差。