[Paper Review] Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch

Aojun Zhou, Yukun Ma+6 more|arXiv (Cornell University)|Feb 8, 2021

Advanced Neural Network Applications43 references74 citations

TL;DR

This work trains N:M fine-grained structured sparse networks from scratch using SR-STE, enabling hardware-friendly sparsity and up to 2x speed-up on Nvidia A100 while preserving accuracy, and introduces the SAD metric to analyze sparse topology changes.

ABSTRACT

Sparsity in Deep Neural Networks (DNNs) has been widely studied to compress and accelerate the models on resource-constrained environments. It can be generally categorized into unstructured fine-grained sparsity that zeroes out multiple individual weights distributed across the neural network, and structured coarse-grained sparsity which prunes blocks of sub-networks of a neural network. Fine-grained sparsity can achieve a high compression ratio but is not hardware friendly and hence receives limited speed gains. On the other hand, coarse-grained sparsity cannot concurrently achieve both apparent acceleration on modern GPUs and decent performance. In this paper, we are the first to study training from scratch an N:M fine-grained structured sparse network, which can maintain the advantages of both unstructured fine-grained sparsity and structured coarse-grained sparsity simultaneously on specifically designed GPUs. Specifically, a 2:4 sparse network could achieve 2x speed-up without performance drop on Nvidia A100 GPUs. Furthermore, we propose a novel and effective ingredient, sparse-refined straight-through estimator (SR-STE), to alleviate the negative influence of the approximated gradients computed by vanilla STE during optimization. We also define a metric, Sparse Architecture Divergence (SAD), to measure the sparse network's topology change during the training process. Finally, We justify SR-STE's advantages with SAD and demonstrate the effectiveness of SR-STE by performing comprehensive experiments on various tasks. Source codes and models are available at https://github.com/NM-sparsity/NM-sparsity.

Motivation & Objective

• Motivate combining unstructured and structured sparsity to accelerate DNNs on GPUs.
• Propose a framework to train N:M sparse networks from scratch without large performance loss.
• Introduce SR-STE to mitigate gradient-induced architecture perturbations during training.
• Define Sparse Architecture Divergence (SAD) to quantify topology changes during training.
• Demonstrate effectiveness across vision tasks and machine translation.

Proposed method

• Define N:M sparsity where in every group of M consecutive weights at most N are non-zero.
• Extend Straight-through Estimator (STE) to enable back-propagation for online pruning during training.
• Introduce Sparse Architecture Divergence (SAD) to measure topology changes during training.
• Propose Sparse-refined STE (SR-STE) with a regularization term that penalizes pruned weights to stabilize architecture during training.
• Evaluate on image classification, object detection, instance segmentation, optical flow, and machine translation; compare with ASP, STE, and other sparsity methods.

Experimental results

Research questions

• RQ1 Can we train N:M sparse networks from scratch without sacrificing performance?
• RQ2 Does SR-STE reduce pruned-weight gradient mismatch and stabilize sparse architectures during training?
• RQ3 How do different N:M patterns (e.g., 2:4, 4:8, 1:4, 2:8) affect accuracy and speedups across tasks?
• RQ4 Does the proposed method retain transferability of sparse models to downstream tasks?

Key findings

• 2:4 sparse networks can achieve about 2x speed-up on Nvidia A100 GPUs with negligible accuracy loss on ImageNet for ResNet-50 when compared to dense baselines.
• 4:8 sparsity (same 50% sparsity) can outperform 2:4 at similar FLOPs on ImageNet with ResNet-50.
• SR-STE consistently improves Top-1 accuracy over STE and ASP baselines on ImageNet across several patterns (e.g., 2:4, 4:8).
• In COCO object detection, 2:8 sparsity yields mAP close to dense baselines, and 4:8 can even surpass dense performance for Faster R-CNN with ResNet-50.
• In optical flow (RAFT) and neural machine translation (Transformer), SR-STE achieves comparable performance to dense models with substantially fewer parameters and FLOPs.
• The SAD metric correlates with performance and decreases when SR-STE stabilizes sparse architectures.