Dynamic Sparse Training with Structured Sparsity

CPAL| Dynamic Sparse Training with Structured Sparsity

• 论文链接：https://arxiv.org/pdf/2305.02299.pdf
• 源码链接：https://github.com/calgaryml/condensed-sparsity

简介

动态稀疏训练（Dynamic Sparse Training，DST）如RigL是SOTA稀疏训练方法，学习满足参数比密集网络少85%-95%权重的非结构化稀疏神经网络（SNN），同时维持相近的泛化性能。稀疏学习方法在训练和推断时都利用稀疏度，不像剪枝算法Learning n:m fine-grained structured sparse neural networks from scratch只在推断时利用稀疏度。

尽管使用DST方法训练的模型是高度稀疏的，且在FLOPs上有高约简率，但在硬件上实现加速是有挑战性的当系数凡是是非结构化的。即使考虑目前在加速非结构化SNN的先进研究，结构化稀疏度实现了在现实视觉硬件上更强的加速。另一方面，结构化稀疏剪枝经常移除显著的权重，导致比相当的相同稀疏度下非结构化SNN更差的泛化能力。本文研究提供了一种两全其美的方法：利用DST框架来学习高度稀疏和结构化表示，同时保持泛化性能。

本文方法

大部分已有的稀疏-稀疏DST方法，包括SOTA算法RigL，学习非结构化稀疏掩码，但是结构化稀疏度实现了实际上更好的加速。本文目标是在RigL学习的稀疏掩码上引入结构化约束，使得在训练和推断时更顺应加速，同时不影响RigL泛化性能。恒定扇入约束（Constant fan-in Constraint）展现了特殊的N:M稀疏度。

恒定扇入N:M结构化稀疏度

恒定扇入约束是N:M稀疏度的一种特殊例子，N是每个神经元的非0权重数量。M是给定层内每个神经元的密集扇入。虽然目前在Nvidia的Ampere和更高版本的架构上，商品硬件加速仅存在2:4稀疏度。但正如Towards Memory-Efficient Training for Extremely Large Output Spaces – Learning with 500k Labels on a Single Commodity GPU证明的那样，恒定的扇入约束可以利用N:M稀疏性带来的高效内存访问和吞吐量增加。该研究解决了具有大输出标签空间的DNN倒数第二层的缩放问题，并提出使用恒定扇入稀疏性（或如他们表示的均匀）稀疏性，以允许在商品硬件上输出大约670000标签的标签空间，在GPU硬件上使用恒定扇入稀疏性的高效定制CUDA实现。

恒定扇入稀疏性有几个属性，这些属性将其与N:M稀疏性区分开来：

• 恒定扇入稀疏比N:M稀疏性更灵活，使得能够将任意全局稀疏性值应用于该模式，而N:M稀疏性被限制为特定的稀疏率。
• 在恒定扇入约束下，模型可以应用ERK等逐层稀疏性分布。通过将参数重新分配到参数较少的层，ERK已被证明优于均匀稀疏性分布。相反，N:M稀疏性只能应用于均匀的稀疏性分布。
• 对N:M稀疏性的硬件支持仅限于2:4稀疏性。虽然这确实提供了一些x2级别的适度加速，但高稀疏模型（>=90%稀疏性）的潜在前景是比等效密集模型快x10。

结构化RigL

本文提出在稀疏到稀疏的DST方法中强制执行恒定扇入约束，以从头开始学习结构化稀疏连接。具体地，本文使用RigL，它可以获得高度稀疏的网络，其泛化性能与密集baseline相当。

简洁地，RigL的方法是训练期间通过剪枝每个层中具有最小幅度的权重并重新生成具有最大相应梯度幅度的权重来更新SNN连接。在大多数训练中，这种情况发生在周期性但相对不频繁的掩码更新步骤中。在SRigL中，权重显著性必须在神经元级别（卷积层中每个滤波器级别）确定，因为本文强制每个神经元（输出通道）具有相同数量的未掩码传入权重，从而满足恒定扇入约束。

然而在非常高稀疏性下，仅此方法就明显落后于RigL的推广（>90%）和Transformer架构。这是因为恒定的扇入约束有一个重要的副作用：在严格的恒定扇入约束下，神经元永远不会被完全掩码（消融）。在非常高的稀疏性水平下，这可能导致许多神经元只有1-2个权重，限制了学习复杂特征的能力，从而降低泛化性能。事实上，根据经验观察到RigL会消融大量神经元。有效地，RigL在高稀疏度下减小模型宽度，以保证泛化性能；本文相信本文是第一个DST方法中明确识别这种行为的研究。为了解决SRigL中的这个问题，本文实现了一种神经元消融方法，使SRigL既能保持恒定的扇入约束，又能在高稀疏度下减小层宽度。本文引入了一个新的超参数${\gamma }_{sal}$，它定义了每个神经元所需的显著权重最小的百分比。给定一个具有恒定扇入k的神经元，如果根据下降或生长标准${\gamma }_{sal}\ast k$的权重被认为是显著的，则该神经元被消融，其权重重新分配给同一层内其他神经元。值得注意的是，这种神经元消融方法允许SRigL以比本文在RigL中发现的低得多的稀疏性水平利用神经元消融结构的稀疏性，同时保持良好的泛化能力。

下面的步骤概述了本文最后的神经元消融SRigL方法。在以下程序中，前两个步骤与RigL中步骤相同，而其他步骤是SRigL特有的，包含修改，包括恒定扇入约束和动态神经元消融。首先设置消融消融${\gamma }_{sal}$。对于每一层执行如下步骤：

1. 获得激活权重大小和剪枝后的权重梯度大小，这些将分别作为剪枝和生长的标准，
2. 计算K，即在该层的当前步骤中要增长和剪枝的权重数量。本文算法中总是在剪枝时增长相同数量的连接。
3. 计算每个神经元的显著性权重数量。如果权重在最大幅度权重或最大幅度梯度的TOP-K时，则该权重被认为是显著的。
4. 消融显著权重小于${\gamma }_{sal}\ast k$的神经元，其中k是扇入。消融是通过剪枝所有传入的权重完成的。这些剪枝后权重在以下步骤中被重新分配到剩下的神经元。
5. 基于消融神经元数量计算新的恒定扇入约束$k^{\prime }$。
6. 剪枝当前层中K个最小幅度值权重。该剪枝标准考虑了层内所有权重。而不是仅剪枝每个神经元的最小权重。
7. 对于每个激活神经元，根据需要重新生成尽可能多的权重，按梯度大小递减顺序进行，直到达到目标扇入$k^{\prime }$。