[论文解读] Dynamic Sparse Training with Structured Sparsity

Mike Lasby, А. В. Голубева+3 more|arXiv (Cornell University)|May 3, 2023

Advanced Neural Network Applications被引用 4

一句话总结

本文提出结构化RigL（SRigL），一种动态稀疏训练方法，通过恒定输入连接数约束学习细粒度的结构化N:M稀疏模式，实现在保持与非结构化动态稀疏训练相当泛化能力的同时实现高模型稀疏度。通过结合结构化稀疏性与神经元消融，SRigL在GPU上实现了高达13.0×的实时推理加速，在CPU上实现了3.4×的加速，优于所有密集模型和非结构化稀疏基线模型。

ABSTRACT

Dynamic Sparse Training (DST) methods achieve state-of-the-art results in sparse neural network training, matching the generalization of dense models while enabling sparse training and inference. Although the resulting models are highly sparse and theoretically less computationally expensive, achieving speedups with unstructured sparsity on real-world hardware is challenging. In this work, we propose a sparse-to-sparse DST method, Structured RigL (SRigL), to learn a variant of fine-grained structured N:M sparsity by imposing a constant fan-in constraint. Using our empirical analysis of existing DST methods at high sparsity, we additionally employ a neuron ablation method which enables SRigL to achieve state-of-the-art sparse-to-sparse structured DST performance on a variety of Neural Network (NN) architectures. Using a 90% sparse linear layer, we demonstrate a real-world acceleration of 3.4x/2.5x on CPU for online inference and 1.7x/13.0x on GPU for inference with a batch size of 256 when compared to equivalent dense/unstructured (CSR) sparse layers, respectively.

研究动机与目标

• 为解决在标准硬件上使用非结构化稀疏神经网络实现真实世界推理加速的挑战。
• 开发一种稀疏到稀疏的动态训练方法，在高稀疏度下保持与非结构化动态稀疏训练相当的泛化性能。
• 实现具有恒定输入连接数的结构化稀疏性，并结合神经元消融，兼顾参数效率与硬件友好的稀疏模式。
• 证明通过动态训练学习的结构化稀疏性，可在CPU和GPU上实现高模型效率与实际加速。
• 表明结构化稀疏性与神经元消融在提升推理效率方面具有互补性，且不牺牲模型精度。

提出的方法

• SRigL通过施加恒定输入连接数约束，扩展了RigL动态稀疏训练框架，确保每个神经元维持固定数量的非零输入权重。
• 该方法学习结构化N:M稀疏模式，其中M个连续权重中有N个为非零，从而实现紧凑且硬件优化的表示。
• 在所有稀疏度范围内均引入神经元消融，使模型能动态移除贡献较低的完整神经元。
• 一项新颖的实证分析表明，如RigL等非结构化DST方法在稀疏度超过90%时会完全消融神经元，从而促使SRigL中引入神经元消融。
• 结构化稀疏模式在训练过程中端到端学习，剪枝与再生由梯度幅值和输入连接数约束共同引导。
• 该方法使用稀疏权重矩阵的紧凑表示，实现在CPU和GPU上高效存储与计算。

(a) Constant fan-in pruning v.s. unstructured pruning.

实验结果

研究问题

• RQ1 能否将动态稀疏训练方法适配以学习结构化稀疏模式，实现在保持泛化性能的同时实现真实世界硬件加速？
• RQ2 在高稀疏度下，结合结构化稀疏性时，神经元消融对模型性能与效率有何影响？
• RQ3 在结构化稀疏性中强制实施恒定输入连接数约束，是否能带来比非结构化稀疏性更稳定、更可预测的输出范数方差？
• RQ4 通过动态训练学习的结构化稀疏性，在CPU和GPU上的真实世界推理基准测试中，能在多大程度上超越非结构化稀疏性？
• RQ5 具有恒定输入连接数和神经元消融的结构化稀疏性是否在提升推理速度方面具有互补性，且不牺牲模型精度？

主要发现

• SRigL在稀疏到稀疏的结构化动态训练中达到最先进性能，在多种架构上实现高达99%稀疏度时，泛化能力与非结构化RigL相当。
• 在CPU上，SRigL在90%稀疏度下实现在线推理3.4×加速，批量推理2.5×加速，优于等效的密集层。
• 在GPU上，SRigL在在线推理中实现1.7×加速，在批量推理（批量大小256）中实现13.0×加速，优于非结构化CSR稀疏层。
• 恒定输入连接数约束降低了输出范数方差，使模型在高稀疏度下表现出更稳定、更可预测的行为。
• SRigL中的神经元消融通过保留关键神经元，避免移除关键权重，在极端稀疏度（如99%）下实现了更优性能。
• 结构化稀疏性与神经元消融的结合，产生紧凑、内存高效的表示，兼具参数效率与真实硬件上的高度加速。

(b) Output-norm variance analysis.