PaperlayerView exact revision on arXiv

Source-linked AI summary

Soft Threshold Weight Reparameterization for Learnable Sparsity

Aditya Kusupati, Vivek Ramanujan, Raghav Somani, Mitchell Wortsman, Prateek Jain, Sham Kakade, Ali Farhadi

arXiv:2002.03231v9cs.LGcs.CVstat.ML

TL;DR

Existing sparsity methods can allocate parameters poorly across layers, affecting accuracy or inference cost. STR uses soft-thresholded weights with learned layer-specific pruning thresholds to induce non-uniform sparsity. It achieves state-of-the-art CNN sparsity results, reduces FLOPs, improves ultra-sparse accuracy, and extends to structured sparsity, while requiring tuning of weight decay and initialization.

  • Problem

    Existing uniform or heuristic non-uniform sparsity budgets can produce sub-optimal layer-wise parameter allocation, limiting accuracy or increasing inference FLOPs.

  • Method

    STR soft-thresholds DNN weights and learns layer-specific pruning thresholds through backpropagation to smoothly induce non-uniform sparsity.

  • Results

    STR achieves state-of-the-art unstructured CNN sparsity results, with 50% fewer FLOPs for 90% sparse MobileNetV1 and approximately 10% higher accuracy for 99% sparse ResNet50.

  • Takeaways & Limitations

    Learned non-uniform budgets can redistribute parameters across layers while maintaining accuracy and reducing inference costs, and STR also supports structured sparsity.

  • Takeaways & Limitations

    STR requires tuning the weight-decay parameter and initialization to reach the targeted overall sparsity, and its behavior depends on the threshold function g.

Abstract

from arXiv · show

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.

1. Introduction

DNN deployment faces high memory, compute, and model-size costs, while existing sparsity methods often allocate parameters sub-optimally across layers. STR learns layer-specific pruning thresholds to obtain non-uniform sparsity, achieving strong accuracy and lower inference cost.

  • Large DNNs create substantial RAM, FLOPs, and flash costs during inference, making deployment a central efficiency challenge.
  • Sparsity can reduce model size and inference FLOPs without significantly changing training pipelines, benefiting on-device and edge solutions.
  • Uniform and heuristic non-uniform budgets can allocate weights sub-optimally across layers and significantly reduce accuracy.
  • The paper asks whether non-uniform sparsity budgets can be learned per layer while remaining stable and accurate.
  • STR applies soft thresholding to each weight tensor and backpropagation learns layer-specific pruning thresholds that smoothly induce sparsity.
  • 0.3% accuracy boost and 50% fewer FLOPs are reported for 90% sparse MobileNetV1, while 99% sparse ResNet50 is approximately 10% more accurate than baselines.

2. Related Work

Prior sparsity research spans structural choices, training regimes, and layer-wise budget allocation. STR addresses the limitation of heuristic non-uniform budgets by learning pruning thresholds directly.

  • DNN sparsity methods differ by structure, training direction, and whether sparsity budgets are uniform or non-uniform across layers.
  • Unstructured sparsity directly zeros tensor entries but may not exploit commodity-hardware speedups, whereas structured sparsity supports scalable standard computation.
  • GMP and DNW are magnitude-based unstructured pruning methods regarded as state-of-the-art for CNNs, while VD incurs 2× memory and 4× training compute.
  • Dense-to-sparse methods begin with dense models, whereas sparse-to-sparse methods train sparse networks end-to-end and may reallocate weights.
  • DSR, SNFS, and RigL use heuristic or fixed mechanisms to reallocate weights across layers during sparse training.
  • Non-uniform budgets can preserve accuracy and reduce FLOPs, but existing approaches do not reduce FLOPs relative to uniform sparsity in practice.
  • STR’s main contribution is learning pruning thresholds and thereby learning the non-uniform sparsity budget rather than relying on heuristics.

3. Method - STR

STR reparameterizes DNN weights with learnable, layer-specific soft thresholds, smoothly inducing sparsity while adapting each layer’s budget. This approach addresses limitations of projection-based methods and produces non-uniform sparsity that can improve accuracy and reduce FLOPs.

  • Projection-based sparsity methods can require a global budget whose layer-wise partition is unclear, while their gradient iterates may remain dense.The resulting dense gradients and updates can be difficult to compute for models with billions of parameters.
  • STR applies soft thresholding directly to weights, setting values with magnitude at most g(s) to zero.The operator is Sg(w, s) = sign(w) · ReLU(|w| − g(s)), where s is learnable and g(s) is the pruning threshold.
  • Each layer receives a trainable parameter sl, allowing backpropagation to learn layer-specific thresholds and induce sparsity smoothly.The loss becomes continuous in the threshold parameters when g is continuous, enabling gradient descent on s.
  • STR’s gradient update is itself sparse because gradients are multiplied by an indicator selecting nonzero soft-thresholded weights.The gradient becomes sparser over iterations as the learned overall budget progresses.
  • Learned thresholds differ across layers, producing a non-uniform sparsity budget that empirically improves prediction accuracy while reducing FLOPs.For 90% sparse ResNet50, the thresholds are shown across all layers; the resulting budget is linked to both accuracy and FLOP reduction.
  • The reparameterization is not generally guaranteed to converge to an optimum, but for sparse weights it is closely related to backward pruning.Unlike standard hard or soft thresholding and backward pruning, STR supports differentiable threshold tuning for DNNs.

4. Experiments

Experiments evaluate STR for unstructured sparsity in ResNet50 and MobileNetV1 on ImageNet-1K, and for low-rank structured sparsity in FastGRNN. Across these settings, STR achieves strong accuracy while learning layer-wise sparsity and reducing inference cost or model size.

  • ResNet50 on ImageNet-1K: STR achieves state-of-the-art unstructured-sparsity accuracy for ResNet50 across the evaluated sparsity regimes.Figure 4 and Table 1 compare STR against multiple strong baselines across 80%, 90%, 95%, 96.5%, 98%, and 99% sparsity.
  • ResNet50 on ImageNet-1K: Up to 10% higher accuracy than GMP at 99% sparsity demonstrates STR’s stability in the ultra-sparse ResNet50 regime.Very few methods are stable at 98–99% sparsity, while STR remains stable at 99%.
  • ResNet50 on ImageNet-1K: In the 80–96.5% sparse ResNet50 regimes, STR maintains state-of-the-art accuracy while reducing FLOPs relative to the best baselines.The reported reductions are at least 6.5%, 16.5%, 22%, and 54% for 80%, 90%, 95%, and 96.5% sparsity, respectively.
  • Layer-wise budgets: STR learns non-uniform layer-wise sparsity and adjusts FLOPs across layers, including sparser initial layers and denser final layers than other non-uniform methods.The resulting layer-wise allocation is associated with lower FLOPs than baselines and a sparser backbone for transfer learning.
  • MobileNetV1 on ImageNet-1K: For MobileNetV1, STR is at least 0.65% more accurate than GMP at 75% sparsity and uses up to 50% fewer FLOPs at 90% sparsity.The 75% result includes at least 38% fewer FLOPs, while the 90% result retains state-of-the-art accuracy with up to 40M fewer FLOPs.
  • FastGRNN on Google-12 and HAR-2: STR learns low-rank FastGRNN parameter tensors, outperforming vanilla training by up to 1.67% on Google-12 and 2.47% on HAR-2.The resulting low-rank models are about 50% and 70% smaller on Google-12 and HAR-2, respectively, while matching or exceeding full-rank accuracy.

5. Discussion and Drawbacks

STR learns non-uniform layer-wise sparsity budgets that achieve strong accuracy while reducing inference FLOPs, and these budgets transfer to other pruning methods. Its main drawbacks are tuning requirements and dependence on the threshold function.

  • Discussion: STR’s learned non-uniform budgets redistribute parameters across layers, achieving state-of-the-art accuracy while reducing inference FLOPs.The FLOP reduction is attributed to better parameter redistribution across layers, although the learning dynamics behind the budget are not fully explained.
  • Discussion: STR makes initial and middle layers sparser while keeping later layers denser than competing non-uniform methods.This allocation balances layer-wise FLOPs and sparsity differently from heuristics that typically sparsify deeper layers more aggressively.
  • Transferability: 66M (16%) fewer FLOPs are obtained when DNW uses STR’s learned budget for 90% sparse ResNet50 while maintaining accuracy.At 95% sparsity, the learned budget improves DNW accuracy by up to 1.42% over uniform sparsity while reducing FLOPs by at least 22M (11%).
  • Transferability: 0.22% and 1.57% accuracy gains are obtained with GMP at 90% and 98% sparsity, respectively, alongside reduced inference FLOPs.The learned budgets can also be obtained using smaller representative datasets rather than expensive large-scale experiments.
  • Drawbacks: STR requires tuning weight decay and sinit to reach target overall sparsity, and its stability and sparsification depend on the threshold function g.Sigmoid and exponential choices for g are reported to work adequately.

6. Conclusions

The paper introduces STR to learn layer-wise pruning thresholds and non-uniform sparsity budgets. It reports strong unstructured and structured-sparsity results, lower inference costs, and transferable learned budgets.

  • Conclusions: STR smoothly induces sparsity by learning layer-wise pruning thresholds, producing non-uniform sparsity budgets.The method applies a soft-threshold operator to DNN weights.
  • Conclusions: STR achieves state-of-the-art unstructured sparsity results in CNNs on ImageNet-1K while substantially reducing inference costs.It matches baseline accuracy for 90% sparse MobileNetV1 with 50% fewer FLOPs.
  • Conclusions: STR is about 10% more accurate than existing methods for ultra-sparse 99% sparse ResNet50 models.The result is presented as evidence of the effectiveness of learned non-uniform sparsity budgets.
  • Conclusions: STR also induces low-rank structure in RNNs while increasing prediction accuracy, and its learned budgets are transferable.The conclusion describes the reparameterization as generalizable and easy to adapt.

A.1. Characterization of g

The appendix specifies desired properties for the threshold function g and gives sigmoid and exponential functions as simple choices satisfying the training-dynamics requirements.

  • Characterization of g: g must have a positive derivative bounded above by a finite positive constant across all real inputs.The condition is expressed as 0 < g′(s) ≤ G for every s ∈ R.
  • Characterization of g: The derivative at sinit should be less than 1 to provide control over the dynamics of s.This condition constrains the initial threshold-parameter behavior.
  • Characterization of g: The logistic sigmoid and exponential functions are used as simple choices for g, with the CNN experiments using the logistic sigmoid with k = 1.The paper notes that experimental training typically has s < 0 and permits choosing k as a scaling factor based on the final weight distribution.

A.2. Gradient w.r.t. {sl}l∈[L]

The supplied passage identifies λ as the ℓ2 regularization hyperparameter in the gradient-related derivation.

  • Gradient w.r.t. {sl}l∈[L]: λ denotes the ℓ2 regularization hyperparameter in the gradient expression.The passage supplies the parameter’s meaning but not the full gradient formula or its role in the broader pipeline.

A.3. ResNet50 Learnt Budgets and Backbone Sparsities

STR learns non-uniform sparsity budgets across ResNet50 layers and produces higher-than-expected backbone sparsity across multiple ImageNet-1K sparsity regimes.

  • STR learns non-uniform sparsity budgets for ResNet50 across 80%–99% sparsity regimes on ImageNet-1K.The reported regimes are 80%, 90%, 95%, 96.5%, 98%, and 99%.
  • STR produces higher-than-expected sparsity in CNN backbones, yielding efficient backbones for transfer learning.
  • For 90% sparse ResNet50, STR learns considerably higher backbone sparsity than the compared methods.
  • The learned ResNet50 budgets can be used directly with GMP and DNW across datasets, with significant accuracy gains reported in Table 4.

A.4. MobileNetV1 Sparsity and FLOPs Budget Distributions

For 90% sparse MobileNetV1 on ImageNet-1K, STR learns a layer-wise sparsity distribution that preserves accuracy-sensitive depthwise separable layers while reducing inference FLOPs.

  • MobileNetV1 sparsity budgets: STR learns considerably higher backbone sparsity than GMP for 90% sparse MobileNetV1 on ImageNet-1K.
  • Budget and FLOPs views: The MobileNetV1 budget comparison covers STR and GMP across layer-wise sparsity and FLOPs distributions.Figures 8 and 9 show the corresponding sparsity and FLOPs distributions for 90% sparse models.
  • Layer-wise distribution: STR automatically keeps depthwise separable convolution layers less sparse than other layers to maximize accuracy.
  • Implementation: STRConv applies soft-thresholding to convolution weights before the convolution operation.The implementation replaces the convolution weights with a soft-thresholded sparseWeight tensor.

A.5. STR Adaptations

STR extends beyond layer-wise unstructured sparsity to global thresholds, structured filter pruning, and per-weight mask learning, while soft thresholding supports stable optimization.

  • Global sparsity: STR can share layer thresholds into one learnable global threshold to induce global sparsity.
  • Global sparsity: Global-threshold STR achieves comparable accuracy to layer-wise sparsity but incurs approximately 2× the inference cost.The higher cost is attributed to a poor non-uniform sparsity distribution caused by differing converged weight values across layers.
  • Structured sparsity: For structured sparsity, STR applies soft thresholding to learned filter-importance scalars, enabling channel or filter pruning.
  • Per-weight sparsity: Per-weight pruning or mask learning is obtained by replacing each layer’s scalar threshold with a tensor matching the weight tensor.
  • Optimization behavior: Soft thresholding is continuous and sub-differentiable, whereas hard thresholding can introduce training instability.
Loading 2002.03231v9…