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