Sparse Transformer：Stride 和 Fixed Factorized Attention

研究人员@OpenAI 提出了跨步和固定注意 在名为“使用稀疏变换器生成长序列”的论文中。他们认为 Transformer 是一种强大的架构，但是，它具有与序列长度相关的二次计算时间和空间。因此，这抑制了使用大序列的能力。这就是为什么他们还提出了一种新的 Sparse Transformer，通过对 Transformer 架构进行一些更改，将时间复杂度降低到 O(n ^(1+1/p) ) 其中 p>1 (p ~=2)，

• 他们还对残差块和权重初始化方法进行了一些更改，以改进深度网络的训练。
• 为了以更快更有效的方式计算注意力矩阵，他们还引入了一组稀疏内核，可以有效地计算它的子集。

标准变压器

在上图中，标准转换器处理 6×6 和 2 个注意力头的图像。上面一行表示一个 6×6 的图像，当计算给定的输出时，它会定位两个注意力头接收作为输入。底行表示所有这些输出（行）和输入（列）的连接矩阵 b/w。我们也将在以后的图像中遵循此约定。

建筑学：

让我们考虑一个自回归序列生成问题，其中分布 x ={x ₁ ,x ₂ ,..x _n } 的联合概率是条件概率的乘积，参数化为$\theta$

$p(x) = \prod_{i=1}^{n} p(x_i | x_1,x_2,x_3 ...x_{i-1} ; \theta)$

在这里，图像、音频、视频也可以被视为原始字节序列。

该网络接受序列输入并输出 v 上的分类分布，其中 v 是在网络 \theta 的输出上进行 softmax 后获得的词汇量大小。

此类模型的一个不错选择是 Transformer。然而，Transformer 需要计算需要 O(n ² ) 可计算性的自注意力矩阵。因此，本文的作者对变压器架构进行了一些修改。让我们详细讨论这些修改。

分解自注意力：

自注意力矩阵将输入嵌入映射到输出矩阵，并通过连接模式对其进行参数化

$S = \left \S_1,S_2, ..., S_n{ \right \}$

其中，S_i 表示第 i^个输出向量所涉及的输入向量的索引集。 Self-attention 可以用下面的表达式来表达：

$Attend \left ( X,S \right ) = a\left ( a(x_i, S_i)\right )_{\left \{ i \epsilon i_1, i_2 ....i_n \right \}} \\ \\ a(x_i,S_i) = softmax(\frac{(W_q x_i) K_{S_i}^{T} }{\sqrt(d)}V_{S_i}) \\ \\ K_{S_i} = \left ( W_k x_j \right )_{j \epsilon S_i} \\ \\V_{S_i} = \left ( W_V x_j \right )_{j \epsilon S_i}$

其中，W _q ，W _k W _v是权重，

作者没有在预测当前值之前关注所有先前的值，而是通过将所有注意力头划分为 p 个单独的注意力头来更有效地执行它，其中^第m 个头定义了索引 A_i^(m) A_i^( m) A_i^{(m)} \subset {j: j \leq i} 并让 S_i =A_i^{(m)}。作者的目标是将注意力头划分为有效的子集，以便

$\left |A_i^{(m)} \right | \propto \sqrt[p]{n}$

此外，作者寻找 A 的选择，以便所有输入都连接到p 个注意力步骤中的所有未来输出位置。

对于每对 j<=i，作者设置每个 A，使得i可以通过最大长度为p+1的位置路径来处理j 。

$j \epsilon A_{a}^{(1)} ; a \epsilon A_{b}^{(2)}; b \epsilon A_{c}^{(3)}$

二维分解注意力：

跨步注意：

二维预测当前像素的一种方法是，一个头关注l个前一个像素，另一个头关注每^第l^个位置，其中l称为步幅，其值选择为\sqrt{n} .在这种模式中，第一个负责人出席

步幅注意力可以用以下等式表示：

$A_i^{1} = \left \{ t, t+1, ... i \right \} ; for \, t = max\left ( 0,i-l \right ) \\ A_i^{2} = \left \{ j: \left ( i-j \right ) mod \, \, l =0 \right \}$

当数据本身具有某种与步幅一致的模式（如图像数据和音乐数据）时，这种注意力模式就会起作用。这不适用于不包含任何周期性模式并且无法以黑白注意力方式传达信息的数据集。

稀疏变压器（跨步）

固定注意：

在数据缺乏结构和周期性的情况下，作者提出了一种称为固定注意力的二维分解方法，其中特定单元格总结了先前的位置并将信息传播到所有未来的单元格。这可以用以下等式在数学上表示： $attention(X) = W_p \cdot attend(X,S)$

$A_{(i)}^{1} = \left \{ j: \lfloor{\frac{j}{l}} \rfloor = \lfloor{\frac{i}{l}} \rfloor \right \}$

$A_{(i)}^{2} = \left \{ j: j mod \,\, l \, \epsilon \left \{ t, t+1 , .. l-1 \right \} \right \}$

其中，t=lc，c 是超参数。

稀疏变压器（固定）

稀疏变压器：

分解注意力

标准密集注意力执行出席函数的线性变换：

$attention(X) = W_p \cdot attend(X,S)$

其中，W _p是后注意权重矩阵。但是，从上面的方法，我们知道我们有不止一个注意力头，那么我们如何整合它们。作者提出了三种方法来执行此操作：

• 集成factorized self-attention最基本的技术是在每个残差块上使用一种attention，并以一个可以由超参数确定的比例顺序交错它们。此类积分的数学表达式如下：

$attention(X) = W_p \cdot attend(X,A^{(r\, mod \, p)})$

这里，r 是当前残差块，p 是注意力头的数量。

• 另一种方法是让单个头部关注两个分解的头部都会关注的像素位置，作者合并了头部。这种方法的计算成本稍高。

$attention(X) = W_p \cdot attend(X,\bigcup_{m=1}^{P} \, A^{(m)})$

• 在多头注意力中使用的第三种方法也被作者大部分时间使用。在该方法中，n _{h 个}注意力产品并行计算并沿特征维度连接

$attention(X) = W_p\left ( attend(X, A)^{(i)} \right )_{i \epsilon \left \{ 1,2...,n_h \right \}} \\$

缩放到 100 层

这篇论文的作者在架构上做了一些重要的改变，以便更容易地训练多层的变换器，首先他们使用了类似于 ResNet 论文中使用的残差块的预激活块

$H_0 = embed(X ,W_e) \\ \\ H_k = H_{k-1} + resblock(H_{k-1}) \\ \\ Y = softmax(norm(H_N)W_{out})$

其中 W _out是权重矩阵， resblock(h) 对输入进行归一化，并按以下方式将注意力前馈：

$a(H) = dropout(attention(norm(H))) \\ b(H) = dropout(ff(norm(H + a(H)))) \\ resblock(H) = a(H) + b(H)$

这里，范数函数表示层归一化，ff是激活函数，根据作者的说法，ff的正确选择是高斯误差线性单元。

建模不同的数据类型

为了对数据的空间关系进行编码，Transformer 主要使用基于位置的嵌入和其他特定于位置的架构。作者还得出结论，位置明智的嵌入最适合该模型。作者计算这些嵌入如下：

$embed (X, W_e) = \left ( x_i*W_e + \sum_{j=1}^{n_{emb}}o_{i}^{(j)}*W_j \right )_{x_i \epsilon X}$

其中 n _emb = d_数据（_数据维度）或 n _emb = d _attn （注意维度）。对于图像数据，n _emb =3（d_数据=3 宽*高*通道），对于文本和音频数据，n _emb =2 因为d _attn =2 并且这里的索引对应每个位置的行和列索引在宽度等于步幅的矩阵中。这里，x _i是序列中的 one-hot 编码的^第i^个元素，o _i ^(j)表示 x _i在^第j 维的 one-hot 编码位置。

通过重新计算注意力权重来节省内存

梯度检查点是 OpenAI 研究人员设计的一种技术，可以让用户以 20% 的额外计算为代价将 10 倍大的模型装入 GPU。当处理长序列时，它对训练自注意力层特别有效，因为它们需要大量内存使用。

仅使用重新计算，作者能够将长度为 16,384 (2^14) 的序列与 100 层的层拟合，否则即使使用现代架构也是不可行的。

高效的块稀疏注意力内核：

strided和fixed的稀疏注意力掩码 可以通过切出查询、键和值矩阵的一部分并以块为单位计算乘积来计算注意力。此外，从不计算注意力矩阵的上三角，将操作次数减少了一半。

混合精度训练

网络权重以单精度浮点数存储，但网络激活和梯度以半精度计算。这有助于我们在 Nvidia V100 GPU 上加速模型训练。

参考

• 稀疏变压器纸