深度学习进阶篇[7]：Transformer模型长输出序列、狭义注意力、FAVOR+疾速注意力、蛋白质序列建模实操。

基于Transformer模型在泛滥畛域已获得卓越成绩，包含自然语言、图像甚至是音乐。然而，Transformer架构始终以来为人所诟病的是其注意力模块的低效，即长度二次依赖限度问题。随着输出序列长度的减少，注意力模块的问题也越来越突出，算力和内存耗费是输出序列长度的平方。

1.基于Transformer模型应答长输出序列

比拟晚期的是一方面从内存方面动手：Transformer-XL和内存压缩的Transformer，而另一种办法是各种稠密注意力机制。稠密注意力机制不再暴力地计算所有token之间的注意力，而仅仅计算出无限token对之间的类似度得分。这些须要重点attend的稠密注意力项能够是人工指定，也能够是通过优化办法找到的，还能够是通过学习学到的，甚至能够是随机的。

因为稠密矩阵也能够用图和边来示意，所以稠密化办法也是由图神经网络文献所启发的，具体关系在图注意力网络中概述。这种基于稠密性的架构通常须要额定的层来隐含产生残缺的注意力机制。

规范的稠密化技术。左：稠密模式的例子，其中tokens只关注其余左近的tokens。右：在图注意力网络中，tokens只关注图中的街坊，这些街坊应该比其余节点有更高的相关性。

目前的这些稠密注意力办法依然有以下局限：

  • 须要高效的稠密矩阵乘法运算，而这些运算并不是所有加速器都能提供。比方Longformer的实现采纳TVM技术将计算自注意力的代码编译为CUDA核。
  • 通常不足严格的实践保障。这点Big Bird是有的。
  • 次要针对Transformer模型和生成式预训练进行优化。
  • 通常须要重叠更多的注意力层来弥补稠密表征，这使得它们很难与其余预训练模型一起应用，因而须要进行再训练，这就很费时费力。

除了上述这些毛病之外，稠密留神机制往往还无奈解决惯例注意力办法所利用的全副问题，比方如 Pointer Networks。此外，还有一些操作是无奈进行稠密化的，比方softmax操作。

为了解决这些问题，Performer提出了一个具备线性复杂度的注意力，其注意力机制可线性扩大，从而实现更快的训练，同时容许模型解决较长的长度，这对于某些图像数据集（如ImageNet64）和文本数据集（如PG-19）是必须的。Performer 应用一个高效的（线性）狭义注意力框架（generalized attention framework），容许基于不同相似性度量（核）的一类宽泛的注意力机制。该框架通过谷歌的新算法 FAVOR+（ Fast Attention Via Positive Orthogonal Random Features）来实现，后者可能提供注意力机制的可扩大低方差、无偏预计，这能够通过随机特色图合成（惯例 softmax-attention）来表白。该办法在放弃线性空间和工夫复杂度的同时准确率也很有保障，也能够利用到独立的 softmax 运算。此外，该办法还能够和可逆层等其余技术进行互操作。

2.两种惯例的注意力机制回顾

用L示意输出的token序列的长度，惯例的点乘注意力就是一个映射，把矩阵Q，K，$V \in \mathbb{R}^{L \times d}$作为输出，其中d代表的是隐含维度。矩阵Q，K，V是输出的两头示意，他们的行别离示意为间断字典数据结构的queries，keys，values。双向的点乘注意力有如下的模式，其中$A\in \mathbb{R}^{L \times L}$ 叫做注意力矩阵：

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

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

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

其中这里的$exp(\cdot)$ 是元素级别的，$1_{L}$是长度为L的全1向量，$diag(\cdot)$是对角矩阵。工夫复杂度为$O(L^2d)$，空间复杂度为 $O(L^2+Ld)$ ，因为A须要显式的存储。原则上，点乘注意力的类型和端到端的长序列解决是不相容的。在Seq2Seq构造的encoder自注意力和encoder-decoder注意力利用的是双向注意力。

另一种重要的注意力类型是无向点积注意力有着如下的模式：

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

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

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

其中$tril(\cdot)$返回参数矩阵的下三角局部，包含对角线。无向的注意力用于自回归生成式模型，例如生成式Transformers中的自注意力，和Seq2Seq Transformers中的decoder局部。

3. 狭义注意力

在以往的注意力机制中，别离对应矩阵行与列的 query 和 key 输出相乘，通过 softmax 计算造成一个注意力矩阵，以存储类似度系数。值得注意的是，这种办法不能将 query-key 生成后果传递给非线性 softmax 计算之后，再将其合成为原始的 query 和 key。然而，将注意力矩阵合成为原始 query 和 key 的随机非线性函数的乘积是能够的，即所谓的随机特色（random feature），这样就能够更加高效地对类似度信息进行编码。

LHS：规范注意力矩阵，由 query(示意为q) 和 key(示意为k) 上的softmax计算组成，示意两两之间的类似得分。RHS：规范注意力矩阵能够通过低阶随机化矩阵Q′和K′来近似，行编码由随机化非线性函数对原始query/key进行编码而失去。对于惯例的softmax-attention，变换十分紧凑，波及指数函数以及随机高斯投影。

惯例的 softmax 注意力能够看作是由指数函数和高斯投影定义的非线性函数的一个特例。在这里咱们也能够反向推理，首先实现一些更狭义的非线性函数，隐式定义 query-key 后果中其余类型的相似性度量或核函数。研究者基于晚期的核办法（kernel method），将其定义为狭义注意力（generalized attention）。只管对于大多核函数来说，闭式解并不存在，但这一机制依然能够利用，因为它并不依赖于闭式解。

4. FAVOR+：通过矩阵相关性实现疾速注意力

上文形容的合成容许咱们以线性而非二次内存复杂度的形式存储隐式注意力矩阵。咱们还能够通过合成取得一个线性工夫注意力机制。尽管在合成注意力矩阵之后，原始注意力机制与具备值输出的存储注意力矩阵相乘以取得最终后果，咱们能够重新排列矩阵乘法以近似惯例注意力机制的后果，并且不须要显式地构建二次方大小的注意力矩阵。最终生成了新算法 FAVOR+。

4.1 FAVOR+实践

FAVOR+应用矩阵$A \in \mathbb{R}^{L \times L}$ ，$A(i,j)=K(q_{i}^T,k_{j}^T)$，其中$q_{i}/k_{j}$代表的是Q/K中第$i^{th}/j^{th}$ 的query/key 行向量，核K：$\mathbb{d}^d \times \mathbb{d}^d \rightarrow R_{+}$ 定义为映射$\phi$:$\mathbb{R}^d \rightarrow \mathbb{R}^r_{+}$，公式为：

$$K(x,y)=\mathbb{E}[(\phi(x)^T \phi(y))]$$

咱们把$\phi(u)$认为是一个随机特色图（random feature map），其中$u \in \mathbb{R}^d$。$Q^{'},K^{'} \in \mathbb{R}^{L \times r}$，行别离用$\phi(q_{i}^T)^T$和$\phi(k_{i}^T)^T$示意，上述核函数间接促成了高效注意力机制：

$$\widehat{Att\longleftrightarrow}(Q,K,V)=\widehat{D}^{-1}(Q^{'}((K^{'})^{T}V))$$

$$\widehat{D}=diag(Q^{'}((K^{'})1_{L}))$$

$\widehat{Att\longleftrightarrow}$示意的是近似注意力，括号示意的是计算的程序。很容易失去这个机制的空间复杂度为$O(Lr+Ld+rd)$，工夫复杂度为$O(Lrd)$，而惯例的注意力机制的空间复杂度为$O(L^2+Ld)$，工夫复杂度为$O(L^2d)$

4.2 双向注意力近似

左：规范注意力模块计算，其中通过执行带有矩阵 A 和值张量 V 的矩阵乘法来计算最终的预期后果；右：通过解耦低秩合成 A 中应用的矩阵 Q′和 K′以及依照虚线框中批示的程序执行矩阵乘法，研究者取得了一个线性注意力矩阵，同时不必显式地构建 A 或其近似。

4.3 单向注意力近似

上述剖析与双向注意力（即非因果注意力）相干，其中没有过来和将来的概念。对于输出序列中没有留神前后 token 的单向（即因果）注意力而言，研究者略微批改办法以应用前缀和计算（prefix-sum computation），它们只存储矩阵计算的运行总数，而不存储显式的下三角惯例注意力矩阵。

左：规范单向注意力须要 mask 注意力矩阵以取得其下三角局部；右：LHS 上的无偏近似能够通过前缀和取得，其中用于 key 和值向量的随机特色图的外积（outer-product）前缀和实现动静构建，并通过 query 随机特征向量进行左乘计算，以在最终矩阵中取得新行。

4.4 softmax-kernels近似softmax

把$\phi$用上面的模式示意，函数为$f_{1},...f_{l}:\mathbb{R} \rightarrow \mathbb{R}$， 函数$g:\mathbb{R}^d \rightarrow \mathbb{R}$和

$$\phi(x)=\frac{h(x)}{\sqrt(m)}(f_{1}(w_{1}^Tx),...,f_{1}(w_{m}^Tx),...,f_{l}(w_{m}^Tx))$$

这能够对大多数的核办法进行建模。大多数状况下，D各向同性的，通常是高斯分布。例如

$$h(x)=1,l=1, D=N(0,I_{d})$$

能够失去PNG核，配置：

$$h(x)=1,l=2,f_{1}=sin,f_{2}=cos$$

对应的是平移不变性核，特地地

$$D=N(0,I_{d})$$

就变成了高斯核 $K_{gauss}$。惯例的注意力矩阵的softmax-kernel如下：

$$SM(x,y)=exp(x^Ty)$$

不失一般性，省去了$\sqrt{d}$，因为咱们能够对keys核queries进行规范化。

$$SM(x,y)=exp(\frac{||y||^2}{2})K_{gauss}(x,y)exp(\frac{||y||^2}{2})$$

通过应用三角函数，咱们失去了$SM(x,y)$随机特色图无偏近似：

$$h(x)=exp(\frac{||x||^2}{2}),l=2,f_{1}=sin,f_{2}=cos$$

咱们把它叫做$\hat{SM}_{m}^{trig}(x,y)$，前面还有一堆实践的剖析，有趣味的能够翻翻论文。

4.5 正交随机特色（ORF）¶

下面的局部组成了FAVOR+的R+局部，还差O局部要解释。为了进一步缩小预计的反差，咱们引入了不通的随机样本$w_{1},...,w_{m}$是正交的(orthogonal)。ORF是一个家喻户晓的办法，但事实证明它是无效的，特地是咱们为softmax引入的PRF，咱们的第一个实践试验结果表明ORFs能够利用于任何维数d的softmax/高斯核预计的方差的升高，而不仅仅是对足够大的d的渐近预计，这导致了大偏差上的第一个指数小界线严格小于非正交办法的概率。随机特色的正性在这些界线中起着关键作用。

5. 试验后果

对 Performer 的空间和工夫复杂度进行基准测试，试验结果表明，注意力机制的减速比和内存缩小在实证的角度上近乎最优，也就是说，这十分靠近在模型中基本不应用注意力机制的状况。

在以工夫(T)和长度(L)为度量的双对数坐标轴中，惯例 Transformer 模型的双向 timing。行方向的起点是GPU内存的极限。黑线(X)示意应用;注意力块时可能的最大内存压缩和减速，它基本上绕过了注意力计算，代表着模型可能的最大效率。Performer模型在注意力局部简直可能达到这种最佳性能。

文章进一步证实应用无偏 softmax 近似的Performer 模型，在微调之后能够向后兼容预训练 Transformer 模型，从而在晋升推理速度的同时升高能耗，且无需齐全从新训练已有的模型。

应用 One Billion Word Benchmark（LM1B）数据集，将原始预训练的Transformer权重间接迁徙到Performer模型中，这个初始后果的精度是0.07（橙色虚线）。通过微调，Performer 的精度在很少的梯度步数之后就迅速复原了。

6. 案例：蛋白质序列建模

蛋白质是具备简单三维构造和特定性能的大分子，对生命来说至关重要。与单词一样，蛋白质被指定为线性序列，其中每个字符是20个氨基酸构建块中的一个。

将 Transformers 利用于大型未标记的蛋白质序列产生的模型可用于对折叠的功能性大分子进行精确的预测。

Performer-ReLU (应用基于 relu 的注意力，这是一个不同于 softmax 的狭义注意力)在蛋白质序列数据建模方面有很强的体现，而 Performer-Softmax 与 Transformer 的性能相匹配，正如实践所预测的后果那样。

上面，咱们可视化一个蛋白质Performer模型，应用基于 relu 的近似注意力机制进行训练，应用 Performer 来预计氨基酸之间的相似性，从序列比对中剖析进化替换模式失去的替换矩阵中复原相似的构造。

更一般地说，咱们发现部分和全局注意力机制与用蛋白质数据训练的Transformer模型统一。Dense Attention的近似Performer有可能捕获逾越多个蛋白质序列的全局相互作用。


作为概念的验证，对长串联蛋白质序列进行模型训练，会使得惯例 Transformer 模型的内存过载，但 Performer模型的内存不会过载，因为它的空间利用很高效。

7. 总结

Google AI的这项工作有助于改良基于非稠密的办法和基于Kernel的Transformer，这种办法也能够与其余技术互操作，钻研人员甚至还将 FAVOR 与Reformer的代码集成在一起。同时钻研人员还提供了论文、 Performer的代码和蛋白质语言模型的代码链接。

该钻研首次证实了，任意注意力矩阵都能够通过随机特色在上游 Transformer 利用中实现无效地近似。实现这一点的的新机制是应用正随机特色，即原始 query 和 key 的耿直非线性函数，这对于防止训练过程中的不稳定性至关重要，并实现了对惯例 softmax 注意力的更精确近似。Performer的钻研开拓了一种对于Attention、Transformer架构甚至Kernel的全新的思维形式，对于进一步的改良有微小的启发作用。

  • 参考文献

Rethinking attention with performers