PaperlayerView exact revision on arXiv

Source-linked AI summary

Dynamic Sparse Training: Find Efficient Sparse Network From Scratch With Trainable Masked Layers

Junjie Liu, Zhe Xu, Runbin Shi, Ray C. C. Cheung, Hayden K. H. So

arXiv:2005.06870v1cs.LGstat.ML

TL;DR

Pruning can reduce the memory and computation demands of over-parameterized neural networks, but existing sparse-training methods face challenges with pruning schedules, weight recovery, and layer-wise pruning rates. Dynamic Sparse Training jointly learns network parameters and sparse structure using trainable thresholds. It achieves state-of-the-art performance against other sparse-training algorithms and produces highly sparse models, including on MNIST benchmarks.

  • Problem

    Neural-network inference is memory- and computation-intensive, while existing sparse-training approaches have limitations in pruning schedules, weight recovery, and layer-wise pruning rates.

    Source: §1 · pp. 1–2

  • Method

    Dynamic Sparse Training jointly learns network parameters and sparse structure using trainable pruning thresholds that adjust through backpropagation.

    Source: §1 · p. 2

  • Results

    Dynamic Sparse Training achieves state-of-the-art performance against other sparse-training algorithms; on MNIST, it prunes almost 98% of parameters on Lenet-300-100 and Lenet-5-Caffe with little performance loss.

    Source: §1 · p. 2, §4.1 · p. 5

  • Takeaways & Limitations

    DST can produce highly sparse or better-performing-than-dense models and reveal consistent layer-wise sparse patterns that may guide compact network design.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper targets inefficiencies in pruning approaches and proposes end-to-end sparse training with dynamic, trainable thresholds. Its design addresses pruning schedules, weight recovery, and layer-wise pruning rates.

1 INTRODUCTION

  • Traditional pruning trains a dense network before pruning and fine-tuning, while sparse-training methods prune during training to avoid those costly iterations.The paper identifies three remaining problems with sparse training methods.Source: §1 · p. 1
  • Fixed, epoch-wise pruning schedules require difficult hyperparameter choices and apply the same schedule even to architectures with very different complexity.No pruning occurs within an epoch under these schedules.Source: §1 · p. 1
  • Hard pruning sets weights to zero and loses their historical importance, making it difficult to decide which weights to recover and how to initialize them.Existing recovery approaches instead restore a predefined portion of weights, randomly or with the same value.Source: §1 · pp. 1–2
  • Uniform layer pruning can be suboptimal, while global thresholds and greedy layer-by-layer pruning can produce inconsistent or amplified errors.Layers differ in redundancy; early-layer pruning can affect later-layer responses.Source: §1 · p. 2
  • The proposed end-to-end method uses one additional hyperparameter for final sparsity, dynamically prunes and recovers weights, and automatically adjusts layer-wise rates.It reports state-of-the-art performance against other sparse-training algorithms.Source: §1 · p. 2

2 RELATED WORK

The related work traces pruning from dense-train-prune-fine-tune methods to approaches that seek sparse networks during training. It highlights trade-offs in how sparse methods choose and adjust weights and thresholds.

2 RELATED WORK

  • Traditional pruning uses second-order derivatives, parameter magnitude, Taylor approximations, or L1 regularization to identify weights for removal.The cited approaches include LeCun et al. (1990), Han et al. (2015), Molchanov et al. (2016), and Li et al. (2016).Source: §2 · p. 2
  • Sparse Evolutionary Training (SET) regrows connections randomly, while DEEP-R uses Bayesian sampling to choose pruning and regrowth configurations.The paper notes that SET requires manual sparsity settings and that DEEP-R is computationally expensive.Source: §2 · pp. 2–3
  • Other sparse-training methods use coarse threshold adjustments, fixed thresholds, or fixed pruning percentages, with additional settings to tune.The discussion includes Dynamic Sparse Reparameterization, Dynamic Network Surgery, and sparse momentum.Source: §2 · p. 3

3 DYNAMIC SPARSE TRAINING

Dynamic Sparse Training represents pruning with trainable masks and layer-wise thresholds, then optimizes network loss and sparsity together. Its masked layers support backpropagation-based threshold updates and step-wise pruning and recovery.

3.1 NOTATION

  • Each network layer has a parameter matrix, and pruning seeks a binary mask for each matrix to encode its sparse structure.The paper also expresses convolution kernels as flattened parameter matrices.Source: §3.1 · p. 3

3.2 THRESHOLD VECTOR AND DYNAMIC PARAMETER MASK

  • A trainable threshold vector and unit step function create masks from weight magnitudes; masked weights retain their underlying values and historical importance.Thresholds can be neuron-wise for fully connected and recurrent layers, or filter-wise for convolutional layers.Source: §3.2 · p. 3

3.3 TRAINABLE MASKED LAYERS

  • Trainable masked layers use W⊙M in layer computation; a long-tailed derivative estimator supports backpropagation through the step function.The estimator has a non-zero gradient over [-1, 1] and a tighter approximation near zero than STE.Source: §3.3 · p. 4, Figure 1 · p. 4
  • The derivative estimator lets thresholds and masked weights update via backpropagation, enabling automatic step-wise pruning and recovery.The paper describes performance and structure gradients for network parameters.Source: §3.3 · p. 4

3.4 SPARSE REGULARIZATION TERM

  • A sparse regularization term penalizes low thresholds, using exp(−x), which approaches zero as the threshold increases.The term is defined for the thresholds of trainable masked layers.Source: §3.4 · p. 5

3.5 DYNAMIC SPARSE TRAINING

  • Dynamic Sparse Training replaces standard layers with trainable masked layers and jointly optimizes task loss and sparse regularization.The sparse regularization term favors higher thresholds and sparsity, while the task loss counteracts increased sparsity to balance structure and performance.Source: §3.5 · p. 5

4 EXPERIMENTS

The experiments evaluate DST on MNIST, CIFAR-10, and ImageNet across fully connected, convolutional, and recurrent architectures. Results show strong pruning performance across these settings, while varying α controls the tradeoff between model size and accuracy.

4 EXPERIMENTS

  • The evaluation spans MNIST, CIFAR-10, and ImageNet with fully connected, convolutional, and recurrent architectures; thresholds are initialized to zero.The zero initialization reflects the assumption that the original network is dense.Source: §4 · p. 5

4.1 PRUNING PERFORMANCE ON VARIOUS DATASETS

  • On MNIST, DST prunes almost 98% of parameters in LeNet-300-100 and LeNet-5-Caffe with little performance loss, and over 99% in LSTMs that outperform dense baselines.The LSTM models are adapted for sequential MNIST classification.Source: §4.1 · p. 5, Table 1 · p. 6
  • On CIFAR-10, DST outperforms DSR and Sparse momentum in almost all settings.The comparison covers VGG and WideResNet architectures.Source: Table 2 · p. 6, §4.1 · p. 6
  • For ImageNet ResNet-50, DST achieves better top-1 and top-5 accuracy with slightly higher model sparsity.The dense ResNet-50 model has 25.56 M parameters.Source: Table 3 · p. 6, §4.1 · p. 6

4.2 PRUNING PERFORMANCE ON VARIOUS REMAINING RATIO

  • Increasing α decreases the model remaining ratio; moderate values can retain comparable or higher accuracy than dense models, while values yielding under 5% remaining parameters cause a noticeable accuracy drop.The reported α range is 10^-9 to 10^-4.Source: §4.2 · p. 6, Figure 3 · p. 7

5 DISCUSSION

DST changes layer sparsity throughout training, adapts its pruning schedule to learning-rate settings, and reveals recurring layer-wise sparsity patterns that the authors present as guidance for compact architecture design.

5.1 FINE-GRAINED STEP-WISE PRUNING

  • Layer remaining ratios change smoothly, fluctuate during the first 100 steps, and layer 1 falls below 10% while layer 3 returns to 1; sparse accuracy remains almost the same as dense accuracy.The reported LeNet-300-100 patterns show layer-wise pruning and recovery during training rather than fixed-size pruning steps.Source: Figure 4 · p. 7, §5.1 · p. 7

5.2 DYNAMIC PRUNING SCHEDULE

  • At initial LR 0.1, VGG-16 ended with 8.82% remaining and 93.93% sparse accuracy versus 93.75% dense; at 0.01, 12.45% remaining and 92.74% versus 92.54%.Across training, sparse accuracy stayed above dense accuracy at LR 0.01, while its running average did so at LR 0.1.Source: Figure 5 · p. 8, Figure 6 · p. 8, §5.2 · p. 9
  • For initial LR 0.1, VGG-16's FC 2 remaining ratio rises after epoch 80; at LR 0.01, it reaches 1 around epoch 40, with accuracy increasing about 2% after decay.The authors attribute the LR 0.1 change to learning-rate decay making previously unimportant output-layer connections useful as preceding layers extract helpful features.Source: Figure 5 · p. 8, Figure 6 · p. 8, §5.2 · pp. 8–9

5.3 INFORMATION REVEALED FROM CONSISTENT SPARSE PATTERN

  • Across α settings, Conv 10–13 and FC 1 are highly sparse while Conv 3–7 and FC 2 retain many parameters, a consistent pattern used to identify layer redundancy for compact design.The authors report similar patterns in other architectures and suggest applying DST to newly designed architectures to assess layer-wise redundancy.Source: §5.3 · p. 9, Figure 7 · p. 10

6 CONCLUSION

The conclusion presents DST as direct sparse-model training with trainable pruning thresholds, and as a source of layer-wise redundancy guidance for compact network design.

6 CONCLUSION

  • Dynamic Sparse Training uses trainable masked layers and pruning thresholds to train sparse models directly, and its consistent sparse patterns can guide compact network design.The conclusion says DST can be applied across neural-network layer types and architectures to obtain highly sparse or better-performing models than dense counterparts.Source: §6 · p. 10

A.1 EXPERIMENTAL SETUP

The appendix specifies training configurations for MNIST, CIFAR-10, and ImageNet, along with a safeguard for layers whose masks become nearly all zero.

A.1 EXPERIMENTAL SETUP

  • MNIST LeNet models use SGD with momentum 0.9, batch size 64, 20 epochs, learning rate 0.01 without decay, and α = 0.0005; LSTMs use Adam for 20 epochs with batch size 100, learning rate 0.001, and α = 0.001.The passage says the LSTM Adam hyperparameters use the default setting.Source: §A.1 · p. 12
  • CIFAR-10 models use SGD with momentum 0.9, batch size 64, and 160 epochs; the learning rate starts at 0.1 and decays by 0.1 at epochs 80 and 120, with α = 5 × 10−6.Source: §A.1 · p. 12
  • ImageNet ResNet-50 models use SGD with momentum 0.9, batch size 1024, and 90 epochs, starting at learning rate 0.1 and decaying by 10 at epochs 30, 60, and 80.Source: §A.1 · p. 12
  • Trainable thresholds start at zero; if over 99% of a layer's mask elements are zero, its pruning threshold is reset to zero, except for LSTM models.The authors say this reset makes training smoother and permits a wider range of α choices.Source: §A.1 · p. 12

A.2 ANALYSIS OF DIFFERENT THRESHOLD CHOICES

The appendix compares scalar, vector, and matrix trainable thresholds, reporting similar pruning for vector and matrix forms but different robustness and storage overhead.

A.2 ANALYSIS OF DIFFERENT THRESHOLD CHOICES

  • Matrix thresholds yield almost the same remaining ratio as vector thresholds, while scalar thresholds are less robust than either vector or matrix thresholds.Source: §A.2 · p. 12
  • Vector thresholds add less than 1% network parameters, whereas matrix thresholds almost double parameter storage during training.Source: §A.2 · p. 12
  • The authors choose vector thresholds because their added element-wise computations are O(n^2), compared with O(n^3) batched matrix multiplication, and deployment stores only masked weights.No deployment overhead is introduced by the thresholds because only W ⊙ M needs to be stored.Source: §A.2 · p. 12

A.3 MORE DETAILS ABOUT THE TRAINABLE MASKED LAYER

Trainable masked layers use masked weights in forward computation and pass separate performance and structure gradients during back-propagation, allowing masked weights to receive structure gradients.

A.3 MORE DETAILS ABOUT THE TRAINABLE MASKED LAYER

  • Forward computation uses the sparse product W ⊙ M for matrix-vector multiplication or convolution.The dense parameter W is masked before the layer operation.Source: §A.3 · p. 12
  • Back-propagation gives W a performance-gradient branch and a structure-gradient branch, while the threshold vector t receives a gradient.The passage specifies the gradient expressions for W and t.Source: §A.3 · p. 13
  • Because W, H(Q), and sgn(W) are dense under the stated value ranges, masked weights can receive the structure gradient.The passage attributes the value ranges to ℓ2 regularization and notes an exception for an unmasked last layer.Source: §A.3 · p. 13

A.4 CHANGE OF REMAINING RATIO IN OTHER NETWORK ARCHITECTURES

The remaining-ratio pattern is reported as similar across three WideResNet architectures and various α values, while sparse patterns are presented for two additional architecture settings.

A.4 CHANGE OF REMAINING RATIO IN OTHER NETWORK ARCHITECTURES

  • Similar remaining-ratio changes appear across WideResNet-16-8, WideResNet-16-10, and WideResNet-28-8 for various α.The figures show the first five convolutional layers and the final fully connected output layer.Source: §A.4 · p. 13, Figure 8 · p. 13, Figure 9 · p. 14
  • For the remaining-ratio plots, only the first five convolutional layers and final fully connected output layer are shown.This is the stated layer selection for WideResNet models, which have one fully connected output layer.Source: §A.4 · p. 13

A.5 CONSISTENT SPARSE PATTERN IN OTHER NETWORK ARCHITECTURES

  • Figures 11–12 compare sparse patterns and remaining-parameter percentages across α choices for WideResNet-16-8 and WideResNet-16-10.The accompanying text describes the sparse pattern as consistent across these tested architectures.Source: §A.5 · p. 13, Figure 11 · p. 14, Figure 12 · p. 14
Loading 2005.06870v1…