[Paper Review] Soft Threshold Weight Reparameterization for Learnable Sparsity

Aditya Kusupati, Vivek Ramanujan+5 more|arXiv (Cornell University)|Feb 8, 2020

Advanced Neural Network ApplicationsComputer Science58 references82 citations

TL;DR

STR learns per-layer pruning thresholds via soft thresholding to induce non-uniform sparsity. It achieves state-of-the-art unstructured sparsity accuracy on CNNs (ResNet50, MobileNetV1 on ImageNet) while reducing FLOPs, and extends to structured sparsity in RNNs.

ABSTRACT

Sparsity in Deep Neural Networks (DNNs) is studied extensively with the focus of maximizing prediction accuracy given an overall parameter budget. Existing methods rely on uniform or heuristic non-uniform sparsity budgets which have sub-optimal layer-wise parameter allocation resulting in a) lower prediction accuracy or b) higher inference cost (FLOPs). This work proposes Soft Threshold Reparameterization (STR), a novel use of the soft-threshold operator on DNN weights. STR smoothly induces sparsity while learning pruning thresholds thereby obtaining a non-uniform sparsity budget. Our method achieves state-of-the-art accuracy for unstructured sparsity in CNNs (ResNet50 and MobileNetV1 on ImageNet-1K), and, additionally, learns non-uniform budgets that empirically reduce the FLOPs by up to 50%. Notably, STR boosts the accuracy over existing results by up to 10% in the ultra sparse (99%) regime and can also be used to induce low-rank (structured sparsity) in RNNs. In short, STR is a simple mechanism which learns effective sparsity budgets that contrast with popular heuristics. Code, pretrained models and sparsity budgets are at https://github.com/RAIVNLab/STR.

Motivation & Objective

• Motivate non-uniform sparsity budgets across layers to maximize accuracy under a parameter budget.
• Introduce Soft Threshold Reparameterization (STR) that learns per-layer pruning thresholds via backpropagation.
• Show STR achieves state-of-the-art unstructured sparsity performance on CNNs and reduces FLOPs.
• Demonstrate STR can induce structured sparsity and transfer learned sparsity budgets across tasks.

Proposed method

• Define STR as the soft-thresholded projection S_g(W_l, s_l) = sign(W_l) * ReLU(|W_l| - g(s_l)) with learnable per-layer s_l.
• Reparameterize the optimization to minimize L(S_g(W, s), D) instead of L(W, D), enabling backpropagation through the pruning thresholds.
• Use a continuous function g (sigmoid for unstructured, exponential for structured sparsity) to map s_l to a per-layer threshold alpha_l = g(s_l).
• Train W_l with SGD and weight decay while updating s_l via gradients through L with respect to STR projections.
• Demonstrate that STR induces sparse gradients and non-uniform budgets across layers, and discuss convergence intuition via connections to backward pruning.

Experimental results

Research questions

• RQ1 Can we learn per-layer sparsity budgets that optimize accuracy under a global parameter limit?
• RQ2 Does STR provide stable training and improved accuracy over existing uniform or heuristic non-uniform sparsity methods?
• RQ3 Can STR extend to structured sparsity and other architectures (e.g., RNNs) while maintaining gains?
• RQ4 Do learned sparsity budgets transfer across tasks or datasets?
• RQ5 How does STR affect FLOPs distribution across layers compared to baselines?

Key findings

• STR achieves state-of-the-art accuracy for unstructured sparsity in ResNet50 and MobileNetV1 on ImageNet-1K across multiple sparsity levels.
• STR learns non-uniform layer budgets that reduce FLOPs by up to about 50% in some cases.
• In ultra-sparse regimes (99%), STR outperforms baselines by up to ~10% accuracy.
• STR induces distinct final per-layer thresholds, leading to non-uniform sparsity budgets that can concentrate computation reductions in earlier layers and preserve density in others.
• The learned sparsity budgets are transferable across tasks and can be extended to structured sparsity (low-rank in RNNs).