[Paper Review] Dynamic Sparse Training: Find Efficient Sparse Network From Scratch With Trainable Masked Layers

Junjie Liu, Zhe Xu+3 more|arXiv (Cornell University)|May 14, 2020

Advanced Neural Network Applications31 references34 citations

TL;DR

The paper introduces Dynamic Sparse Training (DST), an end-to-end method that jointly learns sparse network structure and weights via trainable pruning thresholds, enabling fine-grained, step-wise pruning and recovery during training.

ABSTRACT

We present a novel network pruning algorithm called Dynamic Sparse Training that can jointly find the optimal network parameters and sparse network structure in a unified optimization process with trainable pruning thresholds. These thresholds can have fine-grained layer-wise adjustments dynamically via backpropagation. We demonstrate that our dynamic sparse training algorithm can easily train very sparse neural network models with little performance loss using the same number of training epochs as dense models. Dynamic Sparse Training achieves the state of the art performance compared with other sparse training algorithms on various network architectures. Additionally, we have several surprising observations that provide strong evidence for the effectiveness and efficiency of our algorithm. These observations reveal the underlying problems of traditional three-stage pruning algorithms and present the potential guidance provided by our algorithm to the design of more compact network architectures.

Motivation & Objective

• Motivate the need for efficient sparse networks to reduce memory and computation in inference.
• Propose an end-to-end sparse training framework that learns both weights and layer-wise pruning masks.
• Enable fine-grained, step-wise pruning and recovery within each training step rather than between epochs.
• Automatically adjust layer-wise pruning rates via back-propagation and a trainable threshold mechanism.
• Demonstrate state-of-the-art performance on MNIST, CIFAR-10, and ImageNet across multiple architectures.

Proposed method

• Represent pruning as trainable, layer-wise thresholds per layer for neuron/filters.
• Use a binary mask M derived from a unit step function S applied to |W| - t to obtain sparse W ∘ M.
• Introduce a straight-through estimator-based derivative for training the threshold vector t.
• Incorporate a sparse regularization term Ls = sum exp(-ti) to encourage higher sparsity.
• Replace dense layers with trainable masked layers to enable back-propagation of both performance and structure gradients.
• Show that thresholds and masks can be updated at each training step, enabling fine-grained pruning and recovery.

Experimental results

Research questions

• RQ1 Can weights and sparse structures be learned jointly in an end-to-end framework with trainable pruning thresholds?
• RQ2 Does step-wise pruning and recovery during training outperform predefined pruning schedules in sparse training?
• RQ3 How do layer-wise, trainable thresholds affect final sparsity patterns and model performance across architectures?
• RQ4 What guidance does DST provide for designing compact architectures via observed sparse patterns?

Key findings

• DST can prune almost 98% of parameters on MNIST Lenet-300-100 with little performance loss (Sparse results: remaining 2.48%).
• On MNIST Lenet-5-Caffe, sparse training yields remaining 1.64% with near-dense accuracy (Sparse: 99.11%).
• LSTM models for sequential MNIST achieve comparable or better sparse accuracy with over 99% parameter pruning.
• On CIFAR-10 with VGG-16 and WideResNet, DST outperforms Sparse Momentum and Dynamic Sparse Reparameterization at high sparsities (e.g., VGG-16: DST 8.82% remaining with 93.93% sparse accuracy vs 10% remaining for others).
• ImageNet (ResNet-50): DST achieves higher top-1/top-5 accuracy with slightly higher sparsity (e.g., remaining ~9.87-19.24% depending on setting) compared to baselines.
• DST reveals a consistent sparse pattern across different α values, indicating layer-wise redundancy and providing architectural design guidance.