PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch

Aojun Zhou, Yukun Ma, Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, Hongsheng Li

arXiv:2102.04010v2cs.CVcs.AR

TL;DR

DNN sparsity must balance compression, hardware acceleration, and model performance, while existing sparse-training recipes can require costly dense pretraining or suffer performance losses. The paper trains N:M fine-grained structured networks from scratch with SR-STE and SAD, reporting negligible performance degradation across diverse tasks and sometimes better performance than dense models.

  • Problem

    Existing fine-grained sparsity offers high compression but limited hardware acceleration, while structured sparsity does not consistently combine modern-GPU acceleration with decent performance.

  • Method

    The paper trains N:M sparse networks from scratch using magnitude-based forward pruning, STE-based back-propagation, the SR-STE refinement, and SAD to analyze topology changes.

  • Results

    Across image classification, detection, segmentation, optical flow estimation, and machine translation, the proposed sparse models achieve negligible performance drops and sometimes outperform dense models.

  • Takeaways & Limitations

    N:M fine-grained structured sparsity provides a paper-supported route to train sparse networks from scratch while retaining strong performance across varied tasks.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper addresses the tension between fine-grained sparsity’s compression benefits and structured sparsity’s hardware friendliness by training N:M sparse networks from scratch. It introduces a sparse-refined estimator and evaluates the approach across diverse tasks.

  • Problem and setting: Unstructured sparsity can achieve high compression but incurs irregular memory-access overhead, whereas structured sparsity is more hardware-friendly but may limit compression or performance.The paper frames combining both properties as a challenging problem for accelerating DNNs on modern hardware.
  • Problem and setting: N:M sparsity keeps N non-zero weights in every continuous group of M weights, combining fine-grained pruning with structured hardware support.For 2:4 sparsity, only two weights remain non-zero in each group of four; Nvidia A100 Sparse Tensor Cores can process the corresponding operation in one cycle instead of two.
  • Research gap: A simple recipe for learning N:M sparse neural networks from scratch is needed because existing two-stage pruning requires dense training followed by additional retraining.The paper explicitly asks whether an efficient, simple, and universal recipe can learn N:M sparse networks from scratch.
  • Method: The proposed framework applies magnitude-based pruning during the forward pass and uses a straight-through estimator to enable back-propagation through the non-differentiable pruning operation.The authors analyze perturbations from approximated gradients and introduce Sparse Architecture Divergence to study sparse-topology updates.
  • Results: Across image classification, detection, segmentation, optical flow estimation, and machine translation, the proposed sparse models show negligible performance drops and sometimes outperform dense models.The experiments cover 1:4, 2:4, 2:8, and 4:8 fine-grained structured sparsity patterns.
  • Method: SR-STE adds a sparse-refined term to reduce the influence of approximated gradients and stabilize sparse architecture updates during training.The contribution extends STE for training N:M sparse networks from scratch.

2 RELATED WORK

Prior work includes unstructured and coarse-grained structured sparsity, along with one-stage and two-stage procedures for obtaining sparse networks. The paper positions N:M fine-grained structured sparsity as a distinct hardware-oriented pattern.

  • Unstructured and structured sparsity: Unstructured sparsity prunes individual weights and can provide high compression, but irregular non-zero indices limit acceleration and add storage overhead.Reported examples include 13× and 108× compression ratios.
  • Unstructured and structured sparsity: Structured sparsity removes larger organized units such as filters, channels, or filter shapes to improve hardware efficiency on GPUs.The paper distinguishes these existing coarse-grained patterns from its N:M fine-grained structured sparsity.
  • One-stage and two-stage methods: Two-stage methods prune a trained dense network and retrain the resulting fixed sparse network, while one-stage methods jointly adapt parameters and sparse connections.One-stage methods save training time and cost compared with two-stage methods but usually obtain lower performance.
  • One-stage and two-stage methods: The lottery ticket hypothesis showed that sparse subnetworks can be trained from scratch after their structures are discovered through dense training.Other one-stage approaches include dynamic rewiring, evolutionary training, and momentum-based connection updates.

3 METHOD

The method trains N:M fine-grained structured sparse networks from scratch by pruning weights during the forward pass and using STE-based updates. It introduces SAD to track topology changes and SR-STE to stabilize pruned connections during optimization.

  • N:M Sparse Network Training: N:M sparsity keeps at most N non-zero weights in each group of M consecutive network weights.The objective is to train sparse parameters S(W,N,M) using observed data D and loss L.
  • N:M Sparse Network Training: Magnitude-based pruning projects dense weights into an N:M sparse sub-network during the forward pass.Within each M-parameter group, the least significant absolute values are set to zero while the others are retained.
  • Dynamic Pruning with STE: Vanilla STE applies gradients computed from the pruned sub-network directly to dense parameters, but produces a significant performance drop versus dense training.The paper attributes this drop to unstable architecture updates caused by the mismatch between sparse forward and approximated backward gradients.
  • Sparse Architecture Divergence: SAD measures sparse-topology change as the L1 difference between binary masks at two training iterations; smaller SAD indicates more similar architectures.Experiments report lower SAD for dense-forward training than sparse-forward STE training, and a positive correlation between high SAD and poor sparse-network performance.
  • Sparse-Refined STE: SR-STE adds a regularization term on pruned weights to encourage them to remain pruned and stabilize the sparse architecture.Positive λW constrains only pruned parameters, preventing them from becoming unpruned and reducing architecture alternation during training.

4 EXPERIMENTS

Experiments evaluate N:M sparse patterns and training methods across image classification, detection, segmentation, optical flow, and machine translation. SR-STE generally preserves or improves accuracy while reducing training cost and model parameters.

  • Image Classification: 4:8 outperforms 2:4, and 2:8 outperforms 1:4 despite equal sparsity within each comparison.The reported explanation is that larger M provides more abundant convolution-kernel shapes, while larger N provides more parameters and computational cost at fixed M.
  • Image Classification: SR-STE outperforms NVIDIA ASP and STE with fewer training epochs and better accuracy.Table 2 compares these training methods for N:M sparse networks.
  • Image Classification: SR-STE outperforms state-of-the-art sparsity methods on ResNet50 at both 80% and 95% sparsity.The comparison includes methods using unstructured sparsity.
  • Object Detection and Instance Segmentation: 2:8 (25%) structured sparsity achieves comparable COCO detection and instance-segmentation results with dense baselines, while 4:8 (50%) performs better.The results also indicate similar or better feature-transfer ability for N:M-sparse pretrained models.
  • Cross-task Results: Across image classification, detection, segmentation, optical flow, and machine translation, N:M models show negligible performance drops and sometimes outperform dense models.The experiments cover multiple neural-network architectures and N:M sparsity levels.
  • Optical Flow and Machine Translation: Comparable accuracy with half of the parameters is achieved for RAFT optical-flow prediction relative to the dense model.The evaluation uses end-point-error, for which smaller values represent better performance.

5 DISCUSSION AND CONCLUSION

The paper introduces SR-STE and SAD for training and analyzing N:M fine-grained structured sparse networks from scratch. Experiments support SR-STE’s effectiveness and report a strong relationship between SAD and pruned-network performance.

  • Discussion and Conclusion: SR-STE extends STE with a regularization term to reduce ineffective sparse-architecture updates caused by coarse gradients.The method targets N:M fine-grained structured sparse networks trained from scratch.
  • Discussion and Conclusion: SAD measures sparse-network topology changes during training and correlates strongly with pruned-network performance.The authors propose SAD as a basis for further theoretical and empirical studies of sparse-network training.

A.1 ALGORITHM

Algorithm 1 trains an N:M sparse neural network from scratch by initializing weights, applying sparse pruning during the forward pass, and updating weights during backpropagation.

  • The procedure takes N, M, λW, and dataset D as inputs for SR-STE training.These parameters define the sparsity pattern, regularization strength, and training data.
  • Each iteration samples a mini-batch (X, Y) from D after random model initialization and learning-rate setup.
  • The forward pass computes fW by applying the N:M sparsification operator S(W, N, M) and obtaining its mask E.
  • The update uses learning rate γt and gradient g(fW) with an additional λW(Ē ⊙ W) term in the weight update.

A.2 KERNEL SHAPE

The paper illustrates learned convolution-kernel shapes under 2:8 sparsity and contrasts them with the shapes available under 1:4 sparsity.

  • Figure 5 shows six convolution kernels selected from a trained ResNet50 model using 2:8 sparsity.
  • Under 2:8 constraints, the six kernels exhibit non-zero-element shapes that cannot be acquired or learned with 1:4 sparsity.

A.3 REGNETXS ON IMAGENET-1K

On RegNetX compact models, SR-STE improves performance over STE, while 2:4 structured sparsity can exceed dense models at matched FLOPs.

  • SR-STE significantly improves RegNetX002 performance compared with STE.
  • The illustrated 2:8 ResNet50 kernel is identified as layer1.1.conv2, with coordinates denoting input-channel and output-channel indices.
  • Table 8 reports ImageNet validation accuracy for RegNet models across different N:M sparse patterns.

A.4.1 CLASSIFICATION

The classification experiments use ImageNet-1K, train models with a specified 120-epoch cosine schedule, and evaluate them using Top-1 accuracy.

  • ImageNet-1K contains about 1.2 million training images, 50 thousand validation images, and 1000 object classes.
  • Reported classification results use standard Top-1 accuracy.

A.4.2 OBJECT DETECTION AND INSTANCE SEGMENTATION

The section reports evaluation settings and compares sparse-refined formulations using accuracy and Sparse Architecture Divergence (SAD) curves. The sign-constant formulation improves performance relative to the other refined formulation and STE baseline.

  • Object Detection and Instance Segmentation: COCO-style Average Precision measures object detection and instance segmentation results on MS COCO 2017.The dataset contains 115K training images and 5K validation images, with mAP computed over IoUs from 0.5 to 0.95.
  • Optical Flow Prediction: Endpoint error evaluates optical-flow predictions on the FlyingChairs validation set.FlyingChairs contains 22,232 training samples and 640 validation samples; EPE is averaged Euclidean distance between predicted and ground-truth flow vectors.
  • Sparse-Refined Formulations: 69.4% is achieved by the refined gradient term, versus 69.9% for STE.The refined term has lower SAD early in training but higher SAD later, accompanying worse performance than STE.
  • Sparse-Refined Formulations: 0.9% improvement comes from the sign-constant formulation, whose SAD converges similarly to Eq. 4.Figure 6(b) compares SAD curves for Eq. 4, Eq. 5, and Eq. 6 across training epochs.
Loading 2102.04010v2…