[Paper Review] Dynamic Sparse Training with Structured Sparsity

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

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Structured RigL (SRigL), a dynamic sparse training method that learns fine-grained structured N:M sparsity with a constant fan-in constraint, enabling high model sparsity while maintaining generalization comparable to unstructured dynamic sparse training. It achieves real-world inference speedups of up to 13.0× on GPU and 3.4× on CPU by combining structured sparsity with neuron ablation, outperforming both dense and unstructured sparse baselines.

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.

Motivation & Objective

• To address the challenge of achieving real-world inference acceleration with unstructured sparse neural networks on standard hardware.
• To develop a sparse-to-sparse dynamic training method that maintains generalization performance comparable to unstructured dynamic sparse training at high sparsity levels.
• To enable structured sparsity with constant fan-in and neuron ablation, combining parameter efficiency with hardware-friendly sparsity patterns.
• To demonstrate that structured sparsity, when learned via dynamic training, enables both high model efficiency and practical acceleration on CPUs and GPUs.
• To show that structured sparsity is complementary to neuron ablation in improving inference efficiency without sacrificing model accuracy.

Proposed method

• SRigL extends the RigL dynamic sparse training framework by imposing a constant fan-in constraint, ensuring each neuron maintains a fixed number of non-zero incoming weights.
• The method learns structured N:M sparsity patterns where N out of M consecutive weights are non-zero, enabling compact and hardware-optimized representations.
• Neuron ablation is incorporated across all sparsity regimes, allowing the model to dynamically remove entire neurons when their contributions are low.
• A novel empirical analysis reveals that unstructured DST methods like RigL ablate entire neurons at sparsity levels above 90%, motivating the inclusion of neuron ablation in SRigL.
• The structured sparsity pattern is learned end-to-end during training, with pruning and regrowth guided by gradient magnitude and fan-in constraints.
• The method uses a condensed representation of the sparse weight matrix, enabling efficient storage and computation on both CPU and GPU.

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

Experimental results

Research questions

• RQ1 Can dynamic sparse training be adapted to learn structured sparsity patterns that enable real-world hardware acceleration while preserving generalization performance?
• RQ2 How does neuron ablation affect model performance and efficiency in high-sparsity regimes when combined with structured sparsity?
• RQ3 Does enforcing a constant fan-in constraint in structured sparsity lead to more stable and predictable output-norm variance compared to unstructured sparsity?
• RQ4 To what extent can structured sparsity, learned via dynamic training, outperform unstructured sparsity in real-world inference benchmarks on CPUs and GPUs?
• RQ5 Is structured sparsity with constant fan-in and neuron ablation complementary in improving inference speed without sacrificing model accuracy?

Key findings

• SRigL achieves state-of-the-art performance in sparse-to-sparse structured dynamic training, matching the generalization of unstructured RigL up to 99% sparsity across multiple architectures.
• On CPU, SRigL achieves a 3.4× speedup for online inference and 2.5× for batched inference at 90% sparsity compared to equivalent dense layers.
• On GPU, SRigL achieves a 1.7× speedup for online inference and a 13.0× speedup for batched inference (batch size 256) compared to unstructured CSR sparse layers.
• The constant fan-in constraint reduces output-norm variance, leading to more stable and predictable model behavior under high sparsity.
• Neuron ablation in SRigL enables better performance at extreme sparsity (e.g., 99%) by preserving salient neurons and avoiding the removal of critical weights.
• The combination of structured sparsity and neuron ablation results in a compact, memory-efficient representation that is both parameter-efficient and highly accelerated on real hardware.

(b) Output-norm variance analysis.