为内存塞不下Transformer犯愁？OpenAI应用AI研究负责人写了份指南（2）

为了推动 N:M 结构稀疏化，需要将一个矩阵的列拆分为 M 列的多个 slide（也称为 stripe），这样可以很容易地观察到每个 stripe 中的列顺序和 stripe 的顺序对 N:M 稀疏化产生的限制。

Pool 和 Yu 提出了一种迭代式的贪心算法来寻找最优排列，使 N:M 稀疏化的权重幅度最大化。所有通道对都被推测性地交换，并且只采用幅度增加最大的交换，然后生成新的排列并结束单次迭代。贪心算法可能只会找到局部极小值，因此他们引入了两种技术来逃避局部极小值：

1. 有界回归：在实践中，两个随机通道的最大交换次数是固定的。每次搜索只有一个通道可以进行交换，以保持搜索空间宽而浅；2. 窄且深的搜索：选择多个 stripe 并同时优化它们。

图 9. 贪心算法实现迭代地寻找 N:M 稀疏化最佳排列的算法。

与按默认通道顺序对网络进行剪枝相比，如果在剪枝之前对网络进行置换，可以获得更好的性能。

为了从头开始训练具有 N:M 稀疏化的模型，Zhou & Ma 扩展了常用于模型量化中的反向传播更新的 STE，用于幅度剪枝和稀疏参数更新。

STE 计算剪枝后的网络 的密集参数的梯度 ，并将其作为近似值应用于稠密网络 W：

STE 的扩展版本 SR-STE（稀疏精化 STE）通过以下方式更新稠密权重 W：

其中 是 的掩码矩阵，⊙是元素对应位置相乘。SR-STE 通过（1）限制 中对权重的剪枝，以及（2）维持 中未被剪枝的权重，来防止二进制掩码剧烈变化。

图 10. STE 和 SR-STE 的对比。⊙的比较是元素乘积；⊗是矩阵乘法。

与 STE 或 SR-STE 不同，Top-KAST 方法可以在前向和反向传播的整个训练过程中保持恒定的稀疏性，还不需要使用具有稠密参数或梯度的前向传播。

在训练到第 t 步时，Top-KAST 过程如下：

稀疏前向传递：选择参数 的一个子集，包含每层按大小排列的前 K 个参数，限制为权重的前 D 比例。如果时间 t 的参数化 α^t 不在 A^t（活动权重）中，则参数化为零。

其中 TopK (θ,x) 是根据大小排序后从 θ 中的前 x 个权重。

稀疏向后传递：然后将梯度应用于更大的参数子集 , 其中 B 包含 (D+M)， A⊂B。扩大需要更新的权重比例可以更有效地探索不同的剪枝掩码，从而更有可能将前 D% 的激活权重排列好。

训练分为两个阶段，集合 B∖A 中的附加坐标控制引入的探索量。探索量会在训练过程中逐渐减少，最终掩码会稳定下来。

图 11. Top-KAST 的剪枝掩码会随时间稳定下来。

为了防止马太效应，Top-KAST 通过 L2 正则化损失来惩罚激活权重，以鼓励产生更多新的探索。在更新期间，B∖A 中的参数比 A 受到更多的惩罚以稳定掩码。

稀疏 Transformer

稀疏 Transformer 将 Transformer 架构中的自注意力层和 FFN 层稀疏化，使单个样本推理的速度提高了 37 倍。

图 12. 当在不同网络层上应用稀疏化时，Transformer 模型解码单个 token（非批量推理）的速度。

稀疏 FFN 层：每个 FFN 层包含 2 个 MLP 和中间的一个 ReLU。因为 ReLU 会引入很多零值，所以该方法在激活函数上设计了一个固定结构，来强制要求在一个包含 N 个元素的块中只包含 1 个非零值。稀疏模式是动态的，每个 token 都不同。

其中 Y_(sparse ) 中的每个激活函数结果对应于 W_1 中的一列和 W_2 中的一行。控制器是一个低秩的 bottleneck 全连接层，其中 、 在训练期间使用 argmax 进行推理以选择哪些列应为非零和，以及 Gumbel-softmax 技巧 。因为可以在加载 FFN 权重矩阵之前计算 Controller (x)，所以可以知道哪些列将被清零，因此选择不将它们加载到内存中以加快推理速度。

图 13. (a) 稀疏 FFN 层；红色列未加载到内存中以进行更快的推理。(b) 1:4 稀疏度的稀疏 FFN 控制器。

稀疏注意力层：在注意力层中，维度 d_(model) 被划分为 S 个模块，每个模块的大小为 M=d_(model)/S。为了确保每个细分都可以访问嵌入的任何部分，Scaling Transformer 引入了一个乘法层（即，一个乘法层将来自多个神经网络层的输入按元素相乘），它可以表示任意排列，但包含的参数少于全连接层。

给定输入向量 ，乘法层输出 ：

乘法层的输出是一个大小为 的张量。然后由二维卷积层对其进行处理，其中 length 和 S 被视为图像的高度和宽度。这样的卷积层进一步减少了注意力层的参数数量和计算时间。

图 14. (a) 引入乘法层以使分区能够访问嵌入的任何部分。(b) 乘法全连接层和二维卷积层的结合减少了注意力层的参数数量和计算时间。