表演者

作者: Valar_Morghulis | 来源:发表于2021-01-08 09:47 被阅读0次

https://zhuanlan.zhihu.com/p/319866371

进阶版：Slim-Performer ，https://hub.fastgit.org/google-research/google-research/tree/master/performer/models/slim_performer/pytorch

摘要：我们提出表演者，这是一种变换器，它以可验证的准确率可近似常规的满秩注意力的变换器，但时间和空间复杂度是线性的，而无需任何先验知识，比如稀疏性和低秩。为了近似softmax注意力机制，表演者使用了一种新的基于正值正交随机特征的快速注意力方法(FAVOR+)，这对于可伸缩的核方法来说可能是独立的。FAVOR+还可用于有效地近似softmax以外的可核化注意力机制。这种表示能力对于在常规变换器无法完成的大规模任务中将softmax与其它核做准确率的比较，以研究比softmax更好的注意力，是关键的。表演者是线性结构，完全兼容常规变换器，具有很强的理论保证：注意力矩阵的无偏或近似无偏估计、一致性收敛和低估计方差。我们做了一系列实验，从像素预测，到文本建模，到蛋白质序列建模。我们和其它有效的稀疏或密集的注意力方法对比，取得有竞争力的结果。证明了表演者的新的注意力范式的有效性。

1.引言及相关工作

变换器(Vaswani et al.，2017；Dehghani et al.，2019)是强大的神经网络体系结构，已成为机器学习多个领域的SOTA。

变换器依赖于一个可训练的注意力机制，其是要确认输入序列的元素之间的复杂的依赖关系。不幸的是，常规变换器的复杂度和序列的tokens的长度 $L$ 成二次关系，当 $L$ 很大时，常规变换器的计算量就很昂贵，因此它无法在算力有限的设备上使用， $L$ 为中等大小时也是这样。已有一些工作解决该问题，大多数是使用局部注意力或利用关于结构的先验，例如稀疏性，基于池化的压缩、或者是使用聚类/分组/卷积的方法，其使用k-means聚类来学习动态的稀疏的关注区域，或者使用位置敏感的哈希算法将嵌入相似的tokens分到一组、或者使用滑动窗口或截断目标。使用密集注意力的方法也有很多的研究工作，这些工作是使用低秩的核来代替softmax，这些方法主要依赖于核函数，核函数可以将表示显式地表示为有限个正特征向量的点积。

这些方法的目的不是为了近似常规注意力，而是提出更简单、更易处理的注意力机制，通常通过加入额外的约束（例如Kitaev et al2020将Q集和K集设为相同），或者堆叠很多的稀疏注意力层来近似密集注意力。不幸的是，这些方法产生的表示能否奏效缺乏严格的保证，有时稀疏的注意力方法的有效性只能通过C++去写CUDA核或者TVMs不断地试错来验证。其它的降低变换器空间复杂度的技术包括，使用可逆残差层，这样就可以一次激活存储，以及共享的注意力权重。增加这些约束可能会导致无法用到长序列上，因为它们对注意力机制的近似是不足的。基于截断反向传播的近似无法捕获长程相关性，因为梯度仅在局部窗口内传播。其它方法提出对常规注意力的有偏估计，但只适用于非因果任务，或者有较大的均方误差。

我们提出表演者，是第一个既能保持准确度又能实用的近似常规softmax满秩注意力的变换器结构，空间复杂度和时间复杂度是线性，不依赖像稀疏或低秩这样的任何先验。表演者通过正值正交随机特征(FAVOR+)来实现快速的注意力机制，使用我们提出的新方法来近似softmax和高斯核。我们相信这些方法可作为单独的研究点，有助于可伸缩核方法的研究。因此，表演者是第一个与常规变换器完全兼容的线性结构（只需要少量的微调），而且它具有强有力的理论保证：注意力矩阵的无偏估计或近似无偏估计、一致性收敛和较低的近似方差。

FAVOR+还可以有效地近似softmax以外的其他可核化注意力机制。这种表征能力对于首次在常规变换器无法完成的大规模任务中准确地将softmax与其他内核进行比较，并为它们找到最佳的注意力内核是至关重要的。FAVOR+也可以应用于Transformer以外的领域，作为常规注意力的更具伸缩性的替代品，它本身在计算机视觉(Fu et al.，2019)、强化学习(Zambaldi et al.，2019)、softmax交叉熵损失训练甚至组合优化(Vinyals et al.，2015)中有着广泛的用途。

我们在一系列任务上实验，从像素预测，到文本模型，到蛋白质序列建模。实验表明，我们的方法和其它稀疏的或者是密集的注意力方法的性能相比是由竞争力的，表明表演者利用新的注意力学习范式的有效性。我们强调，FAVOR+理论上还可与其它方法结合实验，例如可逆层(Kitaev等人，2020)或基于蔟的注意力(Roy等人，2020)。

2. FAVOR+及正值正交随机特征

下面我们将描述表演者的核心部分FAVOR+。我们提出一种新的方法来估计softmax（或高斯）核的方法，使用正值正交随机特征对常规softmax注意力做鲁棒和无偏的估计，并且介绍FAVOR+如何应用到其它注意力核上。

2.1 偏好 - 常规注意力机 $1_{L}$ 制

令 $L$ 表示序列长度，即tokens的数量。由Vawani提出的常规点积注意力是以 $Q,K,V\in R^{L\times d}$ 为输入矩阵的映射，其中 $d$ 是中间表示的维度。 $Q,K,V$ 是输入的中间表示，它们的行可以看做是像字典里的queries、keys和values。双向或者无向的点积注意力的形式如下，其中 $A\in R^{L\times L}$ 被称为是注意力矩阵：

$A_{\leftrightarrow } (Q,K,V)=D^{-1}AV$

$A=exp(QK^T/\sqrt{d} )$

$D=diag(A1_{L} )$ $(1)$

其中 $exp(\cdot )$ 是按元素计算的， $1_{L}$ 是长度为 $L$ 的全1向量， $diag(\cdot )$ 是一个以输入向量为对角线的对角矩阵。公式(1)的时间复杂度和空间复杂度分别是 $O(L^2 d)$ 和 $O(L^2 +Ld)$ ，因为 $A$ 必须被显式地存储。因此，公式(1)的点积注意力无法用于长序列的端到端处理。双向注意力应用于Seq2Seq架构中的编码器自注意力和编码-解码注意力。

另外一个重要的类型是单方向的点积注意力，公式是：

$A_{\rightarrow } (Q,K,V)=\tilde{D} ^{-1}\tilde{A} V$

$\tilde{A} = tril(A)$

$\tilde{D} =diag(\tilde{A} 1_{L} )$ $(2)$

其中 $tril(\cdot )$ 返回参数矩阵的下三角部分，包括对角线。如（Vaswani等人，2017）所述，单向注意用于自回归生成建模，例如，作为生成式变换器中的自注意力，以及Seq2Seq变换器中的解码器部分。

我们将展示注意力矩阵 $A$ 在时间复杂度 $O(Ld^2log(d) )$ 内以任意地精度近似上近似到任何精度。作为比较，常用的方法Reformer，通过局部敏感散列技术利用稀疏性，具有 $O(Ld^2log(L) )$ 的时间复杂度。在本文主体部分，我们将描述双向注意力的FAVOR+。单向注意力在附录B.1部分，通过prefix-sums机制得到。

2.2 广义的可核化的注意力

对于注意力块 $A\in R^{L\times L}$ ，FAVOR+使用形式 $A(i,j)=K(q_{i}^T ,k_{j}^T)$ 来处理注意力块，其中 $q_{i}$ 表示 $Q$ 的第 $i$ 个query行向量， $k_{j}$ 表示 $K$ 的第 $j$ 个key行向量。核 $K：R^d \times R^d\rightarrow R_{+}$ 定义（通常是随机化的）映射 $\phi ：R^d \rightarrow R_{+}^r$ （一般地 $r>0$ ）：

$K(x,y)=E({\phi (x)}^T \phi (y))$ $(3)$

我们将 $\phi (u)$ 称为 $u\in R^d$ 的随机特征映射。令 $Q’,K’\in R^{L\times r}$ 的行分别为 ${\phi (q_{i}^T)}^T,{\phi (k_{i}^T)}^T$ ，从公式(3)可得到高效的注意力计算方式：

$\hat{Att_{\leftrightarrow } } (Q,K,V)=\hat{D} ^{-1}(Q’({(K’)^T }V) )$

$\hat{D} =diag(Q’({(K’)^T1_{L} } ))$ $(4)$

这里 $\hat{Att_{\leftrightarrow } }$ 表示近似的注意力。括号表示了计算顺序。很容易看出，这样计算的特点是空间复杂度为 $O(Lr+Ld+rd)$ 和时间复杂度 $O(Lrd)$ ，而不是常规注意力的 $O(L^2+Ld )$ 和 $O(L^2d)$ ，如图1所示。

图1 使用随机特征映射来近似常规注意力。虚线框表示计算顺序及对应的时间复杂度。

上述转换构成了FAVOR+机制的FA部分。剩下的OR+部分要回答的问题是：

——公式(3)中定义的注意力模型的表达能力如何？特别是，我们能否用它从原则上近似常规softmax注意力？

——在实践中我们如何鲁棒地实现它？特别的，我们能否选择 $r<<L$ $for$ $L>>d$ 来获得想要的空间和时间复杂度上的增益？我们在下一节回答这些问题。

2.3 怎样近似，以及，怎样不近似注意力的softmax核

事实证明，通过取 $\phi$ 为函数 $f_{1} ,...f_{l} :R\rightarrow R$ 、函数 $g:R^d \rightarrow R$ 以及符合某种分布 $D\in P(R)^d$ 的 $\omega _{i}$ $or$ $\omega _{1} ,...\omega _{m}$ $iid$ $D$ ：

我们可以对实践中使用的大多数核进行建模。此外，在大多数情况下， $D$ 是各向同性的（即在球体上具有pdf函数常数），通常是高斯的。例如，通过取 $h(x)=1$ ， $l=1$ ， $D=N(0,I_{d} )$ ，我们就得到了所谓PNG核(Choromanski et al., 2017)的估计，例如 $f_{1} =sgn$ 对应角核。如果令 $h(x)=1$ ， $l=2$ ， $f_{1} =sin$ ， $f_{2} =cos$ ，则对应的移位不变核，特别的是，当 $D=N(0,I_{d} )$ 就会导致高斯核(Rahimi & Recht, 2007)。常规注意力矩阵 $A$ 定义的softmax核是：

$SM(x,y)$ $define:$ $exp(x^T y)$ $(6)$

在上式中，不失通用性，我们省略了 $\sqrt{d}$ 的再归一化，因为我们可以对输入的keys和queries的再归一化来等效这个计算。因为： $SM(x,y)=exp(\frac{\vert \vert x \vert \vert ^2 }{2} )K_{gauss} (x,y)exp(\frac{\vert \vert y \vert \vert ^2 }{2} )$ ，根据我们所述的，我们通过以下配置获得 $SM(x,y)$ 的随机特征映射无偏近似：

$h(x)=exp(\frac{\vert \vert x \vert \vert ^2 }{2} )$ ， $l=2$ ， $f_{1} =sin$ ， $f_{2} =cos$ 。我们称之为 $\hat{SM} _{m}^{trig}(x,y)$ 。

但是这里有一个风险。公式(1)的注意力模块为每个token构造一个值向量和系数的凸组合，其系数是相应的再归一化的核分数。这就是为什么核要产生非负的分数。而使用的随机特征映射会产生潜在的负值(negative dimension-values)，会导致不稳定的行为，特别是当核分数接近0的时候（这是对应于 $A$ 中大量的元素，即token和token不相关的情况），在这些区域会被近似为较大的方差。这会导致异常行为，例如，负对角值的再归一化器 $D^{-1}$ ，从而导致训练中断，或者产生一个次优模型。我们从经验上证明 $\hat{SM} _{m}^{trig}$ 会发生这样的情况，并提供了详细的理论解释，表明当近似值趋于0时， $\hat{SM} _{m}^{trig}$ 的方差会很大，见第3节。这就是为什么鲁棒随机特征映射机制来近似常规softmax注意力从未被提出的一个主要原因。

本文提出一种鲁棒机制。此外，我们新的无偏正值随机特征映射近似，当近似值为0 时，方差也趋近于0。见第3节。

引理1 用于Softmax的正值随机特征PRFs 对于 $x,y\in R^d$ ， $z=x+y$ ，我们有：

其中 $\Lambda =exp(-\frac{\vert \vert x \vert \vert ^2 +\vert \vert y\vert \vert ^2 }{2} )$ ， $cosh$ 是双曲余弦。因此，softmax核可以这样地由正值随机特征映射来无偏近似：

$h(x)=exp(-\frac{\vert \vert x \vert \vert ^2 }{2} )$ ， $l=1$ ， $f_{1} =exp$ ， $D=N(0,I_{d} )$

或者

$h(x)=\frac{1}{\sqrt{2} } exp(-\frac{\vert \vert x \vert \vert ^2 }{2} )$ ， $l=2$ ， $f_{1} (u)=exp(u)$ ， $f_{2} (u)=exp(-u)$ ， $D=N(0,I_{d} )$ 。

后者是为了进一步地减小方差。我们称之为相关估计量： $\hat{SM} _{m}^{+}$ 和 $\hat{SM} _{m}^{hyp+}$ 。

左边：在原点周围对称的效用函数 $r$ 定义为基于三角函数和正值随机特征的估计器的均方误差MSEs的比率。它是角度 $\phi$ 的函数， $\phi$ 的单位是弧度， $\phi$ 是输入特征向量和它们长度 $l$ 之间的角度。较大的值表示空间 $(\phi ,l)$ 区域有更好的正值随机特征性能。我们可以看到，在 $\phi$ 足够大的临界区域（softmax核值足够小），我们的方法比三角随机特征更准确。图中画出了 $[-\pi ,\pi ]\times [-2,2]$ 的区域。右边：函数 $r$ 在固定 $l=1$ ，改变角度 $\phi$ 时的切片。右上角：对比两个近似方法在低softmax核值区域的均方误差MSEs。

在图2中，我们看到了正值相比标准三角随机特征的优势。在核值很小的关键区域，需要仔细的逼近，我们的方法要优于它的对手。第4节我们进一步用实验验证了我们方法的优势，使用正值特征来高效地训练基于softmax的线性变换器。如果我们将公式(7)中的 $\omega$ 替换为 $\sqrt{d} \frac{\omega }{\vert \vert \omega \vert \vert }$ ，我们就得到了称为正则化的softmax核 SMREG，这个我们可以用类似的方法去逼近，只要简单地将 $D=N(0,I_{d} )$ 替换成 $D=Unif(\sqrt{d} S^{d-1} )$ ，在 $R^d$ 空间中，半径 $\sqrt{d}$ 球面上与Haar测度相对应的分布，得到估计器 $\hat{SMREG} _{m}^+$ 。如我们在第3节所展示的，这种随机特征也能够准确地近似常规softmax核。

2.4 正交随机特征orthogonal random features(ORFs)

以上构成了FAVOR+方法的R+部分。还有待解释O部分。为了进一步减少估计量的方差（这样我们就可以使用更少的随机特征 $r$ ），我们使(entangle)不同的随机样本 $\omega _{1} ,...\omega _{m}$ 是完全正交的。这个在实现的时候，各向同性分布 $D$ （尤其是在我们目前所考虑的所有核中）无论何时被标准的Gram-Schmidt再归一化程序（详见Choromanski et al., 2017）使用时，它都可以保持无偏性。ORFs是一个知名的方法，事实证明，它与我们为softmax引入的PRFs配合得特别好。这导致了第一个结论：ORFs可以用于降低任何维数 $d$ 的softmax/高斯核估计量的方差，而不仅仅是对足够大的 $d$ 的渐近估计（就像以前的方法一样，请参阅：下一节），并导致第一个大偏差概率的指数小界，该概率严格小于非正交方法。随机特征是正值在这些边界中起着关键作用。ORF机制需要 $m\leq d$ ，而在我们的所有实验设置中都是这样。整个FAVOR+算法的伪代码在附录B中给出。

我们的理论结果与实验结果非常吻合。我们在第4节中展示了PRFs+ORFs极大地提高了注意力矩阵近似的精度，并使我们能够减少 $r$ ，这导致高准确度的同时，时间复杂度和空间复杂度大大降低，这个机制我们称之为FAVOR+。

3 理论结果

本文提出了用正值正交随机特征来估计softmax核的理论。所有这些也可以应用于高斯核，因为如前一节所述，一个可以通过再归一化从另一个获得（见：第2.3节）。附录中给出了所有的证明和更一般的理论结果，并进行了讨论。

引理2 （正值（双曲）vs. 三角随机特征）。对于独立的随机样本 $\omega _{i}$ ，以下是正确的，其中MSE表示均方误差：

因此，当 $SM(x,y)\rightarrow 0$ 的时候，有 $MSE({\hat{SM} }_{m}^{trig} (x,y))\rightarrow \infty$ ，以及 $MSE({\hat{SM} }_{m}^{+} (x,y))\rightarrow 0$

此外，双曲型估计器能够带来更多的准确度提升，要比使用2倍数量随机特征的估计器 ${\hat{SM} }_{2m}^{+} (x,y)$ 。后一个结果表明，正则化的softmax核实际上是注意力机制中softmax核的精确代理。

定理 1 （正则化核 vs. softmax核）假定softmax核的注意力矩阵的 $L_{\infty}$ 范数满足：对于某个常数 $C\geq 1$ ，有 $\vert \vert A \vert \vert_{ \infty} \leq C$ 。用 $A^{reg}$ 表示正则化的softmax核对应的注意力矩阵。以下是成立的：

此外，即使 $L_{\infty}$ 的条件不满足，当 $d\geq 2$ 时，后面一个公式也是成立的，也就是，正则化的softmax核是softmax核的通用下界。

因此，SMREG（即上文的正则化的SM，上上文有定义）的正值随机特征可以用来逼近softmax核。我们的下一个结果表明，正交性被证明能够降低在任意维度 $d>0$ 下正值随机特征估计器的均方误差，我们会提供显式的差距。

定理2 如果用 ${\hat{SM} }_{m}^{ort+} (x,y)$ 来表示用正交随机特征来修改 ${\hat{SM} }_{m}^{+} (x,y)$ （and thus for $m\leq d$ ），对于任意的 $d>0$ ，以下都是成立的：

此外，对于正则化的softmax核SMREG，完全类似的结果成立。

对于正则化的softmax核，正交特征提供了额外的集中结果：估计器的尾部概率的第一个指数小界要严格优于任何 $d>0$ 的非正交版本。我们的下一个结果使我们能够显示地估计差距。

定理3 令 $x,y\in R^d$ ，对于任意的 $a>SMREG(x,y)$ 与 $m\leq d$ ，以下都是成立的：

其中 ${\hat{SMREG} }_{m}^{ort+} (x,y)$ 表示用ORFs修改 ${\hat{SMREG} }_{m}^{+} (x,y)$ 的版本， $X=\Lambda exp(\sqrt{d}\frac{\omega ^T }{\vert \vert \omega \vert \vert _{2} } (x+y))，\omega \sim N(0,I_{d} )$ ， $\Lambda$ 和引理1中的含义相同， $L_{Z}$ 是 $Z$ 的勒让德变换(Legendre Transform)，定义为 $L_{Z} (a)=sup_{\theta >0} log(\frac{e^{\theta a} }{M_{Z}(\theta ) } )$ ，其中 $M_{Z}$ 是 $Z$ 的矩母函数。

我们可以看到，ORFs为softmax核很小的临界区域(关键区域)提供了指数级的小而尖锐的边界。下面我们说明，即使对于具有ORFs的 ${SM}^{trig}$ 机制，采用 $m=\Theta (d\cdot log(d))$ 的随机映射也足以精确地逼近注意力矩阵（因此，如果不进行注意力的再归一化，就不需要PRFs）。一般来说， $m$ 取决于嵌入的维数 $d$ 、所有查询/键所在球的半径R和精度参数 $\epsilon$ （更多讨论见附录F.6），但不取决于输入序列长度 $L$ 。

定理4 （注意力逼近的一致收敛性）令 $h(x)=exp(\frac{{\vert \vert x \vert \vert }^2 }{2} )$ 。假定queries和keys的L2范数有上界 $R>0$ 。定义 $l=Rd^{-\frac{1}{4} }$ ，令 $h^{\ast }=max_{x\in B(l)} \vert h(x) \vert$ ，其中 $B(l)$ 是一个中心在0、半径为 $l$ 的球体。于是，对于任何的 $\epsilon >0$ 、 $\delta =\frac{\epsilon }{(h^{\ast } )^2 }$ 以及随机投影的数目 $m=\Theta (\frac{d}{\delta ^2 } log(\frac{4d^{\frac{3}{4} }R }{\delta } ))$ ，以下的注意力近似机制是成立的，该机制利用具有任意常数概率的 $ORFs:{\vert \vert {\hat{A}-A } \vert \vert }_{1} \leq \epsilon$ 的估计器 ${\hat{SM}}^{trig}$ ，其中 $\hat{A}$ 是注意力矩阵 $A$ 的近似。

4. 实验

dd和随机投影的数目mm以下为注意近似机制，该机制利用具有任意常数概率的orf的估计器sm，其中Aaa是注意矩阵AA的近似

我们在Jax（Frostig et al.，2018）中预先存在的Transformer培训代码的基础上实现了我们的设置，并使用just-in-time进行了优化(jax.jit文件)汇编，并用经验证据补充我们的理论，以证明FAVOR+在多种环境下的实用性。除非明确说明，否则执行者仅用我们的方法替换注意组件，而所有其他组件与常规转换器完全相同。对于速记法，我们将单向/因果建模表示为（U），将双向/屏蔽语言建模表示为（B）。

就基线而言，我们使用其他变压器模型进行比较，尽管其中一些仅限于一种情况，例如，重整器（Kitaev et al.，2020）仅为（U），Linformer（Wang et al.，2020）仅为（B）。此外，我们使用PG-19（Rae et al.，2020）作为替代（B）训练前基准，因为与BERT（Devlin et al.，2018）和Linformer中使用的（目前公开不可用的）BookCorpus（Zhu et al.，2015）+Wikipedia数据集相比，PG-19用于长序列训练。所有模型和标记化超参数如附录A所示。

图3：变压器和表演者在向前和向后传球速度以及允许的最大L方面的比较。”X“（OPT）表示当注意力简单地返回V矩阵时，可以达到的最大可能加速。当模型在16GB的V100 GPU上产生内存不足错误时显示的绘图。使用的词汇量是256。颜色最好。

4.1计算成本

当使用常规默认大小（n个头部；n个层；d ff；d）=（8；6；2048；512）时，我们比较了（B）设置中变压器和执行者的向后传递速度，因为这是训练期间的主要计算瓶颈之一，其中d ff表示MLP层的宽度。

我们观察到（图3）在L方面，表演者达到了几乎线性的时间和次二次的内存消耗（因为显式的O（L 2）注意矩阵没有存储）。事实上，当注意力被简单地返回V矩阵的“身份函数”所取代时，表演者实现了几乎最优的加速和记忆效率。对于大L，内存和向后传递效率的结合分别允许大批量训练和较低的每个梯度步的挂钟时间。附录E中通过不同的层、原始注意和体系结构大小展示了大量的附加结果

4.2 SOFTMAX注意近似误差

我们通过图4中的FAVOR+进一步检查了近似误差。我们证明了这一点1。正交特征比非结构特征误差小，正特征比三角正余弦特征误差小。这两个实验验证了PORF机制。

图4：比较正交与IID特征和三角正弦/余弦与正特征时近似输出的均方误差。我们取L=4096；d=16，并改变随机样本数m。标准偏差显示了15个样本的适当标准化随机矩阵输入数据

为了进一步提高注意块在多次迭代中的整体近似性，从而进一步提高训练效果，随机样本应该定期重画（图5，右）。这是一个廉价的程序，但可以进一步优化（附录B.2）。

4.3 SOFTMAX近似变换器

即使注意机制的近似值很小，小错误也很容易在多个变压器层（例如MLP、多个磁头）中传播，如图14（附录）所示。换句话说，模型的Lipschitz常数可以很容易地放大较小的注意近似误差，这意味着有时可能需要非常紧密的近似。因此，当对变压器模型（即“Performer-X-softmax”）应用FAVOR（+）的softmax近似时，我们证明：

与预训练模型的向后兼容性是softmax近似的一个优点，通过小的微调（误差传播所需），甚至对于LM1B数据集上的三角特征（图5，左）（Chelba等人，2014）。然而，在更大的数据集PG-19上，2.正（POS）softmax特性（带重绘）对于实现与常规变压器匹配的性能至关重要（图5，右）。

图5：我们将原始的预训练变压器的权重转移到Performer中，它产生了一个初始的非零0.07精度（橙色虚线），但是很快就恢复了原始梯度步数的一小部分的精度。然而，在PG-19上，三角（TRIG）softmax近似变得非常不稳定（附录D.2中的完整曲线），而正特征（POS）（无重绘）和Linformer（也近似softmax），即使重绘投影，也处于相同的困惑。带特征重绘的正softmax是匹配变压器所必需的，带有SMREG（来自Sec的正则化）。3）允许更快的收敛。附录D.3中给出了对许多注意核的额外消融研究，也表明三角随机特征甚至导致训练中的NaN值。

4.4蛋白质的多层训练

我们通过使用2019年1月发布的TrEMBL（Consortium，2019）中的蛋白质序列训练36层模型（Madani et al.，2020），进一步对（U）和（B）病例的执行者进行基准测试。在图6中，重整器和Linformer在蛋白质数据集上的精确度显著下降。此外，在（U）和（B）两种情况下，表演者RELU（在等式5中取f=RELU）都达到了最高的精确度，这证明了广义注意的有用性。我们提出的softmax近似也被证明是严密的，实现了与精确的softmax变压器相同的精度，并证实了我们在第3节中的理论主张

图6：序列=虚线，验证=实线。对于TrEMBL，我们使用了（Madani et al.，2020）中完全相同的模型参数（n头；n层；d ff；d）=（8；36；1024；512）。为了公平起见，所有TrEMBL实验都使用16x16 TPU-v2。在给定计算约束的情况下，每次单独运行的批大小都最大化。超参数可以在附录A中找到。扩展结果包括数据集统计、分布外评估和可视化可以在附录C中找到

4.5大长度训练-通用数据集

根据（Parmar et al.，2018）的标准（U）ImageNet64基准，L=12288，这对于常规变压器是不可行的，我们将所有模型设置为使用相同（n个磁头；d ff；d）但不同的n层。执行者/6层与重整器/12层匹配，执行者/12层与重整器/24层匹配（图7：左）。根据硬件（TPU或GPU），我们还发现，通过Jax优化（U）设置，执行者可以比改革者快2倍。

为了进行原理性研究，我们还创建了一个初始蛋白质基准，通过将蛋白质序列连接到TrEMBL的长度L=8192，来预测蛋白质组之间的相互作用，长度足以模拟蛋白质相互作用网络，而无需现有方法所需的大序列比对（Cong等人，2019）。在这种设置下，一个常规的转换器甚至以每片1的批量大小使内存过载，幅度很大。因此，作为基线，我们被迫使用一个明显更小的变量，减少到（nheads；n layers；d ff；d）=（8；f1；2；3g；256；256）。同时，Per-former使用标准（8；6；2048；512）架构以每芯片8的批量大小高效地训练。我们在图7（右图）中看到，较小的变压器（n层=3）很快被限制在 19%，而表演者能够连续训练到24%。

图7：序列=虚线，验证=实线。对于ImageNet64，所有型号都使用标准（n个磁头；d ff；d）=（8；2048；512）。我们进一步证明，我们的正softmax近似实现了与附录D.2中的ReLU相同的性能。对于串联的TrEMBL，我们改变了n层2f1；2；3gf，用于较小的变压器。超参数见附录A。

5结论

提出了一种新型的变换器Performer，通过正正交随机特征（FAVOR+）机制快速注意，显著提高了规则变换器的时空复杂度。我们的机制为我们的知识提供了第一个有效的基于softmax的线性时空复杂度变压器的无偏估计，为变压器和非稀疏注意机制的作用的研究开辟了新的途径。

6更广泛的影响

我们相信，所提出的算法可以以各种方式产生影响：

生物学和医学：

我们的方法有可能直接影响生物序列分析的研究，使变压器可以应用于更长的序列，而不受注意矩阵结构的限制。我们考虑的最初应用是在蛋白质组尺度上预测蛋白质间的相互作用。最近发表的方法需要大量的进化序列比对，这是哺乳动物基因组应用的瓶颈（Cong等人，2019）。将这些方法应用于生物序列的潜在广泛翻译影响是这项工作的主要动机之一。我们相信，现代生物信息学可以极大地受益于新的机器学习技术与变压器是其中最有前途的。扩展这些方法来训练更快更精确的语言模型，为设计具有预先指定的相互作用特性的分子集合打开了大门。这些方法可用于增强现有的基于物理的设计策略，这些策略在开发新的纳米疫苗等方面至关重要（Marcandalli等人，2019年）。

环境：

如我们所示，受欢迎的表演者的特点是计算成本低得多，空间复杂度低得多，可以直接转化为二氧化碳减排（Strubell等人，2019年）和低能耗（You等人，2020年），因为常规变压器需要非常大的计算资源。

变压器研究：

我们相信，我们的研究结果可以形成高效变压器结构的研究，指导该领域的方法具有强大的数学基础。我们的研究也有望将变形金刚扩展到它们的标准范围之外（例如，通过考虑广义注意机制和与内核的联系）。探索可扩展的转换器架构，能够处理数千个或更多数量级的L，同时保持基线的准确性，是生物信息学新突破的一个门户，如我们在论文中解释的蛋白质语言建模。我们提出的方法可能是第一步。

向后兼容性：

我们的表演者可以在一个常规的预先培训的变压器顶部使用，而不是其他变压器的变种。即使不需要向上训练，FAVOR仍然可以用于快速推理而不损失准确性。我们认为这种向后兼容性是所介绍的技术的一个非常重要的附加特性，可能对实践者特别有吸引力

关注超越变压器：

最后，赞成可以适用于近似准确的注意范围外也变压器。这开启了大量新的潜在应用，包括：层次注意网络（HANS）（Yang等人，2016）、图形注意网络（Velickovic等人，2018）、图像处理（Fu等人，2019）和强化学习/机器人学（Tang等人，2020）。

附录

实验用超参数

我们为执行者使用的最佳设置（包括与近似softmax的比较）在广义注意（第。A、 4），除非特别提及（例如，使用名称“Performer SOFTMAX”），“Performer”指的是使用此广义注意设置。

A、 1指标

我们报告了以下评估指标：

准确度：对于单向模型，我们测量下一个标记预测的准确度，在数据集中所有序列位置的平均值。对于双向模型，我们用15%的概率屏蔽每个标记（与（Devlin et al.，2018）相同），并测量屏蔽位置的精度。

困惑：对于单向模型，我们测量数据集中所有序列位置的困惑。对于双向模型，类似于准确度的情况，我们测量蒙面位置的困惑度。

三。

每维/字符位数（BPD/BPC）：由损耗除以ln（2）计算得出。

我们在主要部分的图中使用了完整的TrEMBL评估数据集，而对于其他数据集，如ImageNet64和PG-19，其评估数据集大小非常大，我们使用随机批（>2048个样本）绘制曲线。

A、 1.1 PG-19预处理

PG-19数据集（Rae等人，2020）是一项具有挑战性的远程文本建模任务。它包括1919年以前出版的古腾堡图书版权项目。它没有固定的词汇表大小，而是选择任何可以模拟任意文本字符串的标记化。我们使用了一个单格句段词汇表（Kudo&Richardson，2018），包含32768个标记，它保留了空格，并且与原始图书文本完全可逆。困惑度的计算方法是每个标记的平均对数似然，乘以句子片段标记化与原始数据集中标记数量的比率。每次拆分的原始数据集令牌计数为：train=1973136207，validation=3007061，test=6966499。我们的sentencepiece标记化在每次分割中产生以下标记计数：train=3084760726，valid=4656945，test=10699704。在计算复杂度之前，这给出了序列的对数似然乘子=1.5634，有效值=1.5487，测试=1.5359，这等于exp（对数似然乘子*损失）

TrEMBL的预处理在附录C中有详细说明

A、 2个培训高参数

除非特别说明，否则默认情况下，所有Performer+Transformer都使用0:5grad clip、0:1 weight decay、0:1 dropout、10 3固定学习率和Adam hyperparameters运行，并使特定型号的批处理大小最大化（直到TPU内存过载）。

所有36层蛋白质实验使用相同的计算量（即16x16tpu-v2，每个芯片8GB）。对于串联实验，16x16 TPU-v2也用于表演者，而8x8用于1-3层（d=256）变压器模型（使用16x16不会造成精度上的差异）。

请注意，表演者使用与变形金刚相同的训练超参数，但却取得了竞争性的结果——这表明，青睐可以作为一个简单的下降，而不需要太多的调整。

A、 3个MaxAttentionDefaultValue的近似值

设置为默认参数的最佳值为：renormalize\u attention=True，numerical stabilizer=106，number of features=256，ortho\u features=True，ortho\u scaling=0.0

A、 4个GeneralizedAttentionDefaultValue

设置为默认参数2的最佳值为：renormalize\u attention=True，numerical stabilizer=0.0，number of features=256，kernel=ReLU，kernel\u epsilon=103。

A、 5个ReformerDefaultValue

对于重整器，我们使用了与蛋白质实验相同的超参数，没有梯度剪裁，而使用ImageNet-64的默认值3（使用学习率衰减）。在这两种情况下，改革者使用相同的默认LSH注意参数。

A、 6个默认值

使用我们上面提到的标准管道，我们通过Jax用Linformer变量替换了注意函数，d=10-6；k=600（本文中使用的符号相同（Wang等人，2020）），其中d是重整化过程中的指数，使用e-d作为乘数，以近似softmax，而k是Q和k矩阵投影的维数。作为健全性检查，我们发现我们在Jax中的Linformer实现正确地近似于softmax的输出，对于所有条目都没有0:02的错误。

请注意，为了进行严格的比较，我们的Linformer超参数甚至比（Wang等人，2020）中的默认值更强，如下所示：

我们使用k=600，这比论文中默认的k=256多了一倍，也比默认的m=256多了一倍。

我们还使用重绘，这避免了在Q和K上的“不幸”投影。

B类

我们概述了FAVOR+的主要算法：

B、 1单向用例和前缀

我们将解释如何将第2.2节中的分析扩展到本节中的单向机制。注意，此时注意力矩阵A被掩蔽，即其不在下三角部分（包含对角线）的所有条目被置零（也参见图8）

图8：单向注意前缀和算法的视觉表示。为了清楚起见，我们省略了这个可视化中的注意规范化。该算法保留了前缀和，前缀和是由键对应的随机特征的外积与值向量求和得到的矩阵。在前缀和算法的每个给定迭代中，与查询相对应的随机特征向量乘以最近的前缀和（通过将与前面标记相对应的所有外积相加获得），以获得由注意机制输出的矩阵av的新行

B、 2正交随机特征-扩展

如正文中所述，对于各向同性（适用于大多数实际应用，包括常规注意），而不是采样！i独立地，我们可以使用正交随机特征（ORF）（Yu等人，2016；Choromanski等人，2017；2018b）：这些保持样本wi的边际分布，同时强制不同样本是正交的。如果我们需要m>d，ORF仍然可以在W的每个d*d区块内局部使用（Yu et al.，2016）。

引入ORF以减少蒙特卡罗估计量的方差（Yu等人，2016；Choromanski等人，2017；2018b；2019a；Rowland等人，2019；Choromanski等人，2018a；2019b），我们在主体的理论和实验部分表明，它们确实导致了更精确的近似和更好的下游结果。ORF机制有几种变体，在正文中我们只讨论了基本的一种（这里我们称之为常规的）。下面我们简要回顾一下最有效的ORF机制（基于它们的优势和成本），以呈现最完整的情况。

（1）正则ORF[R-ORF]：应用高斯正交矩阵（Yu等人，2016）。在O（md）空间中对W样本的矩阵W进行编码（不同的行对应不同的样本）。提供用于计算任意x2Rd的Wx in O（md）时间的算法。给出无偏估计。需要一次性的O（md2）预处理（Gram-Schmidt正交化）。

（2） Hadamard/Givens ORFs[H/G-ORFs]：应用随机Hadamard（Choromanski et al.，2017）或Givens矩阵（Choromanski et al.，2019b）。在O（m）或O（mlog（d））空间中对矩阵W进行编码。提供用于计算任意x2Rd的Wx in O（mlog（d））时间的算法。给出较小的偏差（d趋于0！1).

B、 3时间和空间复杂性-详细分析

我们看到，使用iid样本或R-ORFs的双向FAVOR+的变体具有O（md+Ld+mL）空间复杂度，而不是基线的（L2+Ld）空间复杂度。并行使用快速前缀和预计算（Ladner&Fischer，1980；Cormen et al.，2009）的单向FAVOR+具有O（mLd）空间复杂度来存储GPS，其可通过运行GPS i；：；：的简单（尽管在L中非并行）聚合而减少到O（md+Ld+mL），而无需将整个张量GPS存储在内存中。从第。B、 2.我们知道，如果使用G-ORFs，则空间复杂度降低到O（mlog（d）+Ld+mL），如果使用H-ORFs机制，则空间复杂度进一步降低到O（m+Ld+mL）=O（Ld+mL）。因此，对于m；d<<L，我们的所有变体都提供了大量的空间复杂性改进，因为它们不需要显式地存储注意矩阵。

算法1的时间复杂度为O（Lmd）（注意，q0和k0的构造可以在时间O（Lmd）内完成）。注意，我们方法的时间复杂度远低于L>>m基线的O（l2d）。

如第。B、 2，R-ORF机制产生额外的一次性O（md 2）成本（与L>>d的O（Lmd）项相比可以忽略不计）。H-ORF或G-ORF没有这个成本，当FAVOR+使用它们时，计算Q 0和K 0可以在时间O（Llog（m）d）中进行，而不是在时间O（Lmd）中进行（见第。B、（2）。因此，即使H/G-ORFs不改变渐近时间复杂度，它们也改进了前导项的常数因子。这可能在训练非常大的模型时起到重要作用

随机特征m的数量允许在计算复杂性和近似水平之间进行权衡：m越大，计算成本越高，但a的估计方差也越小。在主体部分的理论部分，我们表明，在实践中，我们可以取m=（dlog（d））。

观察到FAVOR+算法是高度并行的，并且受益于快速矩阵乘法和gpu或tpu上的广播操作。

C蛋白质建模任务的实验细节

C、 1个TREMBL数据集

我们使用了TrEMBL数据集，它包含139394261个序列，其中106030080个是唯一的。虽然训练数据集似乎比Madani等人（Madani等人，2020）中使用的数据集小，但我们认为它包含了大多数相关序列。具体地说，TrEMBL数据集由UniProtKB序列的子集组成，这些序列已经过计算分析，但没有手动管理，并且占UniProtKB数据集中序列总数的99:5%。

按照Madani et al.（Madani et al.，2020）中描述的方法，我们使用了OOD测试集（其中Pfam家族的选定子集用于评估）和IID分裂（其中剩余的蛋白质序列随机分裂为序列、有效和测试测试）。我们发现了以下蛋白质家族（PF18369、PF04680、PF17988、PF12325、PF03272、PF03938、PF17724、PF10696、PF11968、PF04153、PF06173、PF12378、PF04420、PF10841、PF06917、PF03492、PF06905、PF15340、PF17055、PF05318），它们产生了29696个OOD序列。我们注意到，由于重复数据消除和潜在的TrEMBL版本不匹配，我们的OOD测试集与Madani等人（Madani等人，2020）中的测试集不完全匹配。我们还注意到，这种OOD测试选择方法不能保证评估序列与训练过程中使用的序列之间的最小距离。在未来的工作中，我们将包括严格的基于距离的分割。

结果数据集拆分的统计数据如表1所示。在标准序列建模任务中，给定表中报告的长度统计信息，我们将单个序列的最大长度L=1024，这导致少数序列被显著截断。

在长序列任务中，训练集和验证集是通过将序列串联起来，由序列结束标记分隔，并将得到的链分组为长度=8192的非重叠序列来获得的。

C、 2经验基线

一个随机基线，在每个位置的所有词汇标记上具有统一的概率，其准确度为5%（仅包括20个标准氨基酸时）和4%（同时包括5个异常氨基酸时（Consortium，2019））。然而，我们的数据集中各种氨基酸的经验频率可能并不一致，因此我们还考虑了一个经验基线，其中氨基酸概率与其在训练集中的经验频率成正比。

图9显示了估计的经验分布。我们同时使用标准氨基酸和异常氨基酸，并将序列裁剪为1024的长度，以匹配对Transformer模型执行的数据处理。图中仅显示了20种标准氨基酸，按类别着色，以便与TrEMBL网页6上的可视化结果进行比较

C、 3表格结果

表2包含单个蛋白质序列建模任务的结果（L=1024）。我们报告附录A中定义的准确性和复杂性：

表2：单蛋白序列建模结果（L=1024）。我们注意到，经验基线结果适用于单向（UNI）和双向（BID）模型。

C、 4注意事项基质净化

在本节中，我们将说明由执行者模型生成的注意矩阵。我们专注于双向的情况，并选择一个执行者模型训练的标准单序列TrEMBL任务超过500K步骤。同样的分析也适用于单向执行者。

我们注意到，虽然Transformer模型实例化了注意矩阵，以便计算包含（查询Q，键K，值V）三元组的注意输出（见主要论文中的等式1），但偏好机制直接返回注意输出（见算法1）。为了解释这种差异，我们通过两次应用每个注意机制来提取注意矩阵：一次在每个原始（Q；K；V）三元组上获取注意输出，另一次在修改的（Q；K；V）三元组上获取注意矩阵，其中V为每个位置索引包含一个热指标。V的选择确保注意输出的维度等于序列长度，并且维度i上的非零输出只能从到第i序列位置的非零注意权重产生。实际上，在Transformer的情况下，当将这个过程的输出与实例化的注意矩阵进行比较时，输出是匹配的。

注意矩阵示例。

我们首先通过可视化单个蛋白质序列的注意矩阵。我们使用了BPT1\u-BOVIN蛋白序列，它是研究最广泛的球状蛋白之一，含有100个氨基酸。在图10中，我们展示了前4层的注意矩阵。请注意，许多头部显示为径向模式，其中每个节点注意其相邻节点，一些头部显示为垂直模式，其中每个头部注意相同的固定位置。这些模式与自然语言训练的变压器模型中发现的模式一致（Kovaleva et al.，2019）。在图12中，我们通过关注前25个标记来突出这些注意模式，在图11中，我们更详细地说明了两个注意头。

氨基酸相似性。

此外，如Vig等人（Vig等人，2020）所述，我们还分析了根据执行者模型产生的注意矩阵估计的氨基酸相似性矩阵。我们将注意力矩阵聚合到800个序列上。得到的相似矩阵如图13所示。请注意，执行者识别高度相似的氨基酸对，如（D，E）和（F，Y）。

图10：

我们展示了前4层和所有8个头部的注意矩阵（每行是一层，每列是头部索引，每个细胞包含整个BPT1\u-BOVIN蛋白序列的注意矩阵）。请注意，许多头部显示为径向模式，其中每个节点注意其相邻节点，一些头部显示为垂直模式，其中每个头部注意相同的固定位置。

图11：

我们更详细地说明了两个注意头。子图分别对应于：（1）头1-2（第二层，第三头），（2）头4-1（第五层，第二头）。注意头1-2中的块注意和头4-1中的垂直注意（开始标记（'M'）和第85标记（'C'））。

图12：

我们通过将注意力限制在前25个标记来突出注意模式（注意，我们不会将注意力重新规范化到这些标记）。插图基于Vig等人（Vig，2019；Vig&Belinkov，2019）。注意，与先前关于蛋白质转换器的研究（Madani等人，2020年）类似，注意矩阵包括局部和全局模式。

图13：

氨基酸相似性矩阵是由聚集在一小部分序列上的注意矩阵估计的，如Vig等人（Vig等人，2020）所述。这些子图分别对应于：（1）标准化BLOSUM矩阵，（2）通过训练的执行者模型估计的氨基酸相似性。请注意，执行者识别高度相似的氨基酸对，如（D，E）和（F，Y）。

D扩展逼近结果

D、 1向后兼容性-错误传播

尽管前面提到过（第。4.2）具有额外微调的执行者与变压器向后兼容，我们在下面的图14中证明，由于变压器的非注意组件导致的误差传播是预训练变压器权重不能立即用于推断相应执行者的主要原因之一。

图14：不同层数的vanilla Transformer和执行者（具有正交特征）之间的输出近似误差

D、 2近似SOFTMAX-扩展属性

我们在图15中展示了softmax近似的以下特性：

重绘：

而重绘特征的好处在第。4.3在论文的主体部分，我们还介绍了多层次大规模（16x16tpu-v2）训练的好处。

单向的：

虽然我们已经在TrEMBL上展示了具有广义ReLU注意力的执行者比softmax表现更好，但是我们也展示了近似softmax注意力仍然是一个可靠的选择，例如在ImageNet64（U）上。经过10万步的训练，Performer ReLU、Performer Softmax和Performer Softmax（SMREG）变体分别达到3.67、3.69和3.67 BPD。

三角函数特征的不稳定性：

在使用三角softmax时，我们看到了不稳定训练曲线的全貌

图15：

最佳观看缩放在左边：重绘特征的重要性。如果不使用重绘，一组“不幸运”的随机特征可能会导致训练降级，如种子1的提前停止曲线所示，而一组“幸运”的随机特征可能不会导致问题，如种子2的曲线所示。重绘允许训练自我修正，如黑色垂直线所示线。中间：使用相同的8x8 TPU-v2计算和相同的6层标准模型，具有正特征的近似softmax获得与广义ReLU相同的结果注意。好的：图5右侧子图的缩小视图，表明三角softmax会导致非常不稳定的训练行为。

D、 3.注意事项

我们研究了广义注意机制（见第。2.2）对于各种核函数，当L=512时，在TrEMBL上。这与（Tsai et al.，2019）相似，Tsai et al.，2019还对自然语言的各种注意力内核进行了实验。使用多个变量的超参数扫描，我们比较了几个核和重整化开/关（图16和图17），其中重整化对应于注意应用D1算子，至于标准机制，尽管我们注意到，禁用它并不一定会损害准确性），以便为执行者生成最佳的训练配置。我们注意到，有效批量大小略微影响排名（如2x2和4x4 TPU运行之间的差异所示）-默认情况下，我们使用广义ReLU内核和附录A中所示的其他默认超参数，因为我们观察到，它们在经验上对于大批量运行（即8x8或16x16 TPU）是最佳的。

图16：为了强调最高精度的运行，同时也显示某些内核导致运行提前停止的NaN问题，我们将x轴和y轴都设置为对数刻度。我们测试了由不同函数f定义的内核（参见：第。2.2）：sigmoid、指数、ReLU、绝对、gelu、余弦（原始softmax近似）、tanh和恒等式。所有培训运行均在2x2 TPU-v2上进行，每台设备128批次大小。

E计算成本-扩展结果

在本小节中，我们根据经验测量了图18、19中三种情况下，向前和向后通行证的挂钟时间计算成本：

不同层数的表演者。我们表明，我们的方法可以扩展到（但不一定限于）甚至20层。

比较标准注意（来自变压器）和偏好（来自表演者）时的注意时间复杂性。请注意，这里的最大内存大小并不反映实际模型（如下所示）中的最大内存大小，因为这个基准测试需要在Jax中计算显式张量（导致内存增加），而模型则不需要。比较Transformer和Performer模型时的时间复杂性。”“X”（OPT）表示当注意力简单地返回V向量时所能达到的最大可能的加速，表明表演者几乎是最优的。我们看到，使用常规尺寸时，V100GPU（16GB）上允许的最大功率（2长度）为2 15=32768

由于转换器中的一些计算瓶颈可能来自额外的前馈层（Kitaev等人，2020），因此我们还对“小”版本进行了基准测试，即（n个头；n层；d ff；d）=（1；6；64；64），而注意力组件是计算和记忆的主要来源。我们提醒读者，“常规”版本包括（n个头；n层；d ff；d）=（8；6；2048；512）。

表演者

相关文章

网友评论

延伸阅读

深度阅读

栏目导航

热点阅读