介绍
常规Transoformer的计算复杂度是 O ( L 2 d ) O(L^2d) O(L2d),空间复杂度是 O ( L 2 ) O(L^2) O(L2),在长序列的情况下,比如基因序列的处理,序列长度 L L L太大会导致计算成本大到难以接受。
本文提出Fast Attention Via Orthogonal Random features (FAVOR)方法,在不改变原先transformer模型架构的前提下,高效地估计出attention矩阵,时间复杂度为 O ( L d 2 log d ) O(Ld^2\log d) O(Ld2logd),空间复杂度为 O ( L d log d ) O(Ld\log d) O(Ldlogd)。
相比于其他改进transformer在长序列情况下复杂度的方法,本文方法并没有提出新的模型结构,可以方便地应用于已经预训练好的模型。
方法
常规transformer的attention的计算有两种形式,区别在于时序依赖关系的处理。
第一种称为双向点乘的attention(Bidirectional or non-directional dot-product attention),形式如下:
第二种称为单向点乘attention(unidirectional dot-product attention),形式如下:
类比cosine相似度,attention中的点乘实际上可以看作相应向量的相似性的度量,于是作者在文中对其进行泛化,基于核函数的表示给出了Generalized Attention (GA):
在GA中,核函数 K K K负责query Q Q Q和key K K K中元素的交互,选择不同的核函数,对应不同的attention的形式。原先的点乘attention对应的GA形式为(对应着高斯核):
如果知道了核函数的random feature map函数,核函数的计算可以表示为:
对于大多数的核函数来说,相应的random feature map都相似的形式:
特别的,对于高斯核来说,有如下的random feature map:
表示成矩阵的形式有:
现在,attention矩阵 A A A表示成两个低秩矩阵的乘积,计算的时间复杂度下降了。但是如果显示计算出 A A A的话,空间复杂度仍然是 L 2 L^2 L2的,因此在attention与值矩阵相乘的时候,先用低秩矩阵与值矩阵相乘,避免A矩阵的出现。
对于双向点乘attention来说,与值矩阵相乘的计算方式如下,注意括号决定的乘法优先级:
其中,期望的估计使用了一个样本。
对于单向点乘attention来说,要计算的目标具有时序依赖关系,可以通过前缀和的形式得到:
正交随机特征
估计核函数需要采样 W W W,文中提出 W W W具有一定正交性的话,效果会好。推荐的有Gaussian orthogonal matrices,random Hadamard/Givens matrices。
理论收敛性分析
M o p t M_{opt} Mopt最优的随机特征的个数不依赖于序列长度 L L L,而依赖于特征维度 d d d。 M o p t = Θ ( d log d ) M_{opt}=\Theta(d\log d) Mopt=Θ(dlogd)。
实验
上图表明,序列越长,加速越明显。在序列长1000以下,加速不太明显。
上图,使用了随机正交特征,在特征数够大的时候,近似误差较小,不使用随机特征的情况下,考虑到attention score的scale本身就是0.0x左右,有一定的误差。网络深度会加重误差传播,应用的时候不能直接迁移,需要finetune。
上图,在多层transformer下,训练效果能达到甚至超过原先的transformer,比Reformer要好。
注意
- M大小如何选取
量级是 Θ ( d log d ) \Theta(d\log d) Θ(dlogd),实际中100多近似效果较好。 - 在预训练好的transformer上,进行finetune再用,不然准确率很差。