Soft Threshold Weight Reparameterization for Learnable Sparsity

Introduction

Deep Neural Networks (DNNs) are the state-of-the-art models for many important tasks in the domains of Computer Vision, Natural Language Processing, etc. To enable highly accurate solutions, DNNs require large model sizes resulting in huge inference costs, which many times become the main bottleneck in the real-world deployment of the solutions. During inference, a typical DNN model stresses the following aspects of the compute environment: 1) RAM - working memory, 2) Processor compute - Floating Point Operations (FLOPsOne Multiply-Add is counted as one FLOP), and 3) Flash - model size. Various techniques are proposed to make DNNs efficient including model pruning (sparsity) (Han et al., 2015), knowledge distillation (Buciluǎ et al., 2006), model architectures (Howard et al., 2017) and quantization (Rastegari et al., 2016).

Sparsity of the model, in particular, has potential for impact across a variety of inference settings as it reduces the model size and inference cost (FLOPs) without significant change in training pipelines. Naturally, several interesting projects address inference speed-ups via sparsity on existing frameworks (Liu et al., 2015; Elsen et al., 2019) and commodity hardware (Ashby et al., ). On-premise or Edge computing is another domain where sparse DNNs have potential for deep impact as it is governed by billions of battery limited devices with single-core CPUs. These devices, including mobile phones (Anguita et al., 2012) and IoT sensors (Patil et al., 2019; Roy et al., 2019), can benefit significantly from sparsity as it can enable real-time on-device solutions.

Sparsity in DNNs, surveyed extensively in Section 2, has been the subject of several papers where new algorithms are designed to obtain models with a given parameter budget. But state-of-the-art DNN models tend to have a large number of layers with highly non-uniform distribution both in terms of the number of parameters as well as FLOPs required per layer. Most existing methods rely either on uniform sparsity across all parameter tensors (layers) or on heuristic non-uniform sparsity budgets leading to a sub-optimal weight allocation across layers and can lead to a significant loss in accuracy. Furthermore, if the budget is set at a global level, some of the layers with a small number of parameters would be fully dense as their contribution to the budget is insignificant. However, those layers can have significant FLOPs, e.g., in an initial convolution layer, a simple tiny 3$\times$3 kernel would be applied to the entire image. Hence, while such models might decrease the number of non-zeroes significantly, their FLOPs could still be large.

Motivated by the above-mentioned challenges, this works addresses the following question: “Can we design a method to learn non-uniform sparsity budget across layers that is optimized per-layer, is stable, and is accurate?”.

Most existing methods for learning sparse DNNs have their roots in the long celebrated literature of high-dimension statistics and, in particular, sparse regression. These methods are mostly based on well-known Hard and Soft Thresholding techniques, which are essentially projected gradient methods with explicit projection onto the set of sparse parameters. However, these methods require a priori knowledge of sparsity, and as mentioned above, mostly heuristic methods are used to set the sparsity levels per layer.

Due to layer-specific thresholds and sparsity, ${\rm STR}$ is able to achieve state-of-the-art accuracy for unstructured sparsity in CNNs across various sparsity regimes. ${\rm STR}$ makes even small-parameter layers sparse resulting in models with significantly lower inference FLOPs than the baselines. For example, ${\rm STR}$ for 90% sparse MobileNetV1 on ImageNet-1K results in a 0.3% boost in accuracy with 50% fewer FLOPs. Empirically, ${\rm STR}$’s learnt non-uniform budget makes it a very effective choice for ultra (99%) sparse ResNet50 as well where it is $\sim$10% more accurate than baselines on ImageNet-1K. ${\rm STR}$ can also be trivially modified to induce structured sparsity, demonstrating its generalizability to a variety of DNN architectures across domains. Finally, ${\rm STR}$’s learnt non-uniform sparsity budget transfers across tasks thus discovering an efficient sparse backbone of the model.

The 3 major contributions of this paper are:

Soft Threshold Reparameterization (${\rm STR}$), for the weights in DNNs, to induce sparsity via learning the per-layer pruning thresholds thereby obtaining a better non-uniform sparsity budget across layers.

Extensive experimentation showing that ${\rm STR}$ achieves the state-of-the-art accuracy for sparse CNNs (ResNet50 and MobileNetV1 on ImageNet-1K) along with a significant reduction in inference FLOPs.

Extension of ${\rm STR}$ to structured sparsity, that is useful for the direct implementation of fast inference in practice.

Related Work

This section covers the spectrum of work on sparsity in DNNs. The sparsity in the discussion can be characterized as (a) unstructured and (b) structured while sparsification techniques can be (i) dense-to-sparse, and (ii) sparse-to-sparse. Finally, the sparsity budget in DNNs can either be (a) uniform, or (b) non-uniform across layers. This will be a key focus of this paper, as different budgets result in different inference compute costs as measured by FLOPs. This section also discusses the recent work on learnable sparsity.

Unstructured sparsity does not take the structure of the model (e.g. channels, rank, etc.,) into account. Typically, unstructured sparsity is induced in DNNs by making the parameter tensors sparse directly based on heuristics (e.g. weight magnitude) thereby creating sparse tensors that might not be capable of leveraging the speed-ups provided by commodity hardware during training and inference. Unstructured sparsity has been extensively studied and includes methods which use gradient, momentum, and Hessian based heuristics (Evci et al., 2020; Lee et al., 2019; LeCun et al., 1990; Hassibi & Stork, 1993; Dettmers & Zettlemoyer, 2019), and magnitude-based pruning (Han et al., 2015; Guo et al., 2016; Zhu & Gupta, 2017; Frankle & Carbin, 2019; Gale et al., 2019; Mostafa & Wang, 2019; Bellec et al., 2018; Mocanu et al., 2018; Narang et al., 2019; Kusupati et al., 2018; Wortsman et al., 2019). Unstructured sparsity can also be induced by $L_{0},L_{1}$ regularization (Louizos et al., 2018), and Variational Dropout (VD) (Molchanov et al., 2017).

Gradual Magnitude Pruning (GMP), proposed in (Zhu & Gupta, 2017), and studied further in (Gale et al., 2019), is a simple magnitude-based weight pruning applied gradually over the course of the training. Discovering Neural Wirings (DNW) (Wortsman et al., 2019) also relies on magnitude-based pruning while utilizing a straight-through estimator for the backward pass. GMP and DNW are the state-of-the-art for unstructured pruning in DNNs (especially in CNNs) demonstrating the effectiveness of magnitude pruning. VD gets accuracy comparable to GMP (Gale et al., 2019) for CNNs but at a cost of $2\times$ memory and $4\times$ compute during training making it hard to be used ubiquitously.

Structured sparsity takes structure into account making the models scalable on commodity hardware with the standard computation techniques/architectures. Structured sparsity includes methods which make parameter tensors low-rank (Jaderberg et al., 2014; Alizadeh et al., 2020; Lu et al., 2016), prune out channels, filters and induce block/group sparsity (Liu et al., 2019; Wen et al., 2016; Li et al., 2017; Luo et al., 2017; Gordon et al., 2018; Yu & Huang, 2019). Even though structured sparsity can leverage speed-ups provided by parallelization, the highest levels of model pruning are only possible with unstructured sparsity techniques.

2 Dense-to-sparse and Sparse-to-sparse Training

Until recently, most sparsification methods were dense-to-sparse i.e., the DNN starts fully dense and is made sparse by the end of the training. Dense-to-sparse training in DNNs encompasses the techniques presented in (Han et al., 2015; Zhu & Gupta, 2017; Molchanov et al., 2017; Frankle & Carbin, 2019; Renda et al., 2020).

The lottery ticket hypothesis (Frankle & Carbin, 2019) sparked an interest in training sparse neural networks end-to-end. This is referred to as sparse-to-sparse training and a lot of recent work (Mostafa & Wang, 2019; Bellec et al., 2018; Evci et al., 2020; Lee et al., 2019; Dettmers & Zettlemoyer, 2019) aims to do sparse-to-sparse training using techniques which include re-allocation of weights to improve accuracy.

Dynamic Sparse Reparameterization (DSR) (Mostafa & Wang, 2019) heuristically obtains a global magnitude threshold along with the re-allocation of the weights based on the non-zero weights present at every step. Sparse Networks From Scratch (SNFS) (Dettmers & Zettlemoyer, 2019) utilizes momentum of the weights to re-allocate weights across layers and the Rigged Lottery (RigL) (Evci et al., 2020) uses the magnitude to drop and the periodic dense gradients to regrow weights. SNFS and RigL are state-of-the-art in sparse-to-sparse training but fall short of GMP for the same experimental settings. It should be noted that, even though sparse-to-sparse can reduce the training cost, the existing frameworks (Paszke et al., 2019; Abadi et al., 2016) consider the models as dense resulting in minimal gains.

DNW (Wortsman et al., 2019) and Dynamic Pruning with Feedback (DPF) (Lin et al., 2020) fall between both as DNW uses a fully dense gradient in the backward pass and DPF maintains a copy of the dense model in parallel to optimize the sparse model through feedback. Note that DPF is complementary to most of the techniques discussed here.

3 Uniform and Non-uniform Sparsity

Uniform sparsity implies that all the layers in the DNN have the same amount of sparsity in proportion. Quite a few works have used uniform sparsity (Gale et al., 2019), given its ease and lack of hyperparameters. However, some works keep parts of the model dense, including the first or the last layers (Lin et al., 2020; Mostafa & Wang, 2019; Zhu & Gupta, 2017). In general, making the first or the last layers dense benefits all the methods. GMP typically uses uniform sparsity and achieves state-of-the-art results.

Non-uniform sparsity permits different layers to have different sparsity budgets. Weight re-allocation heuristics have been used for non-uniform sparsity in DSR and SNFS. It can be a fixed budget like the ERK (Erdos-Renyi-Kernel) heuristic described in RigL (Evci et al., 2020). A global pruning threshold (Han et al., 2015) can also induce non-uniform sparsity and has been leveraged in Iterative Magnitude Pruning (IMP) (Frankle & Carbin, 2019; Renda et al., 2020). A good non-uniform sparsity budget can help in maintaining accuracy while also reducing the FLOPs due to a better parameter distribution. The aforementioned methods with non-uniform sparsity do not reduce the FLOPs compared to uniform sparsity in practice. Very few techniques like AMC (He et al., 2018), using expensive reinforcement learning, minimize FLOPs with non-uniform sparsity.

Most of the discussed techniques rely on intelligent heuristics to obtain non-uniform sparsity. Learning the pruning thresholds and in-turn learning the non-uniform sparsity budget is the main contribution of this paper.

4 Learnable Sparsity

Concurrent to our work, (Savarese et al., 2019; Liu et al., 2020; Lee, 2019; Xiao et al., 2019; Azarian et al., 2020) have proposed learnable sparsity methods through training of the sparse masks and weights simultaneously with minimal heuristics. The reader is urged to review these works for a more complete picture of the field. Note that, while ${\rm STR}$ is proposed to induce layer-wise unstructured sparsity, it can be easily adapted for global, filter-wise, or per-weight sparsity as discussed in Appendix A.5.

Method - STRSTR{\rm STR}

Optimization under sparsity constraint on the parameter set is a well studied area spanning more than three decades (Donoho, 1995; Candes et al., 2007; Jain et al., 2014), and is modeled as:

Projected Gradient Descent (PGD) in particular has been popular for both the problems as the projection onto both $L_{0}$ as well as the $L_{1}$ ball is computable in almost closed form (Beck & Teboulle, 2009; Jain et al., 2014); $L_{0}$ ball projection is called Hard Thresholding while $L_{1}$ ball projection is known as Soft Thresholding. Further, these methods have been the guiding principle for many modern DNN model pruning (sparsity) techniques (Han et al., 2015; Zhu & Gupta, 2017; Narang et al., 2019).

However, projection-based methods suffer from the problem of dense gradient and intermediate parameter structure, as the gradient descent iterate can be arbitrarily out of the set and is then projected back onto $L_{0}$ or $L_{1}$ ball. At a scale of billions of parameters, computing such dense gradients and updates can be daunting. More critically, the budget parameter $k$ is set at the global level, so it is not clear how to partition the budget for each layer, as the importance of each layer can be significantly different.

In this work, we propose a reparameterization, Soft Threshold Reparameterization (${\rm STR}$) based on the soft threshold operator (Donoho, 1995), to alleviate both the above mentioned concerns. That is, instead of first updating $\mathcal{W}$ via gradient descent and then computing its projection, we directly optimize over projected $\mathcal{W}$. Let $\mathcal{S}_{g}(\mathcal{W};s)$ be the projection of $\mathcal{W}$ parameterized by $s$ and function $g$. $\mathcal{S}$ is applied to each element of $\mathcal{W}$ and is defined as:

Reparameterizing the optimization problem with $\mathcal{S}$ modifies (note that it is not equivalent) it to:

For $L$-layer DNN architectures, we divide $\mathcal{W}$ into: $\mathcal{W}=\left[\mathbf{W}_{l}\right]_{l=1}^{L}$ where $\mathbf{W}_{l}$ is the parameter tensor for the $l$-th layer. As mentioned earlier, different layers of DNNs are unique can have significantly different number of parameters. Similarly, different layers might need different sparsity budget for the best accuracy. So, we set the trainable pruning parameter for each layer as $s_{l}$. That is, $\bm{s}=[s_{1},\dots,s_{L}]$.

Now, using the above mentioned reparameterization for each $\mathbf{W}_{l}$ and adding a standard $L_{2}$ regularization per layer, we get the following Gradient Descent (GD) update equation at the $t$-th step for $\mathbf{W}_{l},\ \forall\ l\in\left[L\right]$:

where $\eta_{t}$ is the learning rate at the $t$-th step, and $\lambda$ is the $L_{2}$ regularization (weight-decay) hyper-parameter. $\nabla_{\mathbf{W}_{l}}\mathcal{S}_{g}(\mathbf{W}_{l},s_{l})$ is the gradient of $\mathcal{S}_{g}(\mathbf{W}_{l},s_{l})$ w.r.t. $\mathbf{W}_{l}$.

Now, $\mathcal{S}$ is non-differentiable, so we use sub-gradient which leads to the following update equation:

where $\mathbf{1}\left\{\cdot\right\}$ is the indicator function and $A\odot B$ denotes element-wise (Hadamard) product of tensors $A$ and $B$.

Now, if $g$ is a continuous function, then using the ${\rm STR}$ (2) and (1), it is clear that $\mathcal{L}(\mathcal{S}_{g}(\mathcal{W},\bm{s}),\mathcal{D})$ is a continuous function of $\bm{s}$. Further, sub-gradient of $\mathcal{L}$ w.r.t. $\bm{s}$, can be computed and uses for gradient descent on $\bm{s}$ as well; see Appendix A.2. Algorithm 1 in the Appendix shows the implementation of ${\rm STR}$ on 2D convolution along with extensions to global, per-filter & per-weight sparsity. ${\rm STR}$ can be modified and applied on the eigenvalues of a parameter tensor, instead of individual entries mentioned above, resulting in low-rank tensors; see Section 4.2.1 for further details. Note that $\bm{s}$ also has the same weight-decay parameter $\lambda$.

Naturally, $g$ plays a critical role here, as a sharp $g$ can lead to an arbitrary increase in threshold leading to poor accuracy while a flat $g$ can lead to slow learning. Practical considerations for choice of $g$ are discussed in Appendix A.1. For the experiments, $g$ is set as the Sigmoid function for unstructured sparsity and the exponential function for structured sparsity. Typically, $\left\{s_{l}\right\}_{l\in\left[L\right]}$ are initialized with $s_{\rm init}$ to ensure that the thresholds $\left\{\alpha_{l}=g(s_{l})\right\}_{l\in\left[L\right]}$ start close to . Figure 1 shows that the thresholds’ dynamics are guided by a combination of gradients from $\mathcal{L}$ and the weight-decay on $\bm{s}$. Further, the overall sparsity budget for ${\rm STR}$ is not set explicitly. Instead, it is controlled by the weight-decay parameter ($\lambda$), and can be further fine-tuned using $s_{\rm init}$. Interestingly, this curve is similar to the handcrafted heuristic for thresholds defined in (Narang et al., 2019). Figure 2 shows the overall learnt sparsity budget for ResNet50 during training. The curve looks similar to GMP (Zhu & Gupta, 2017) sparsification heuristic, however, ${\rm STR}$ learns it via backpropagation and SGD.

Finally, each parameter tensor learns a different threshold value, $\left\{\alpha_{l}\right\}_{l\in\left[L\right]}$, resulting in unique final thresholds across the layers, as shown in Figure 3 for ResNet50. This, in turn, results in the non-uniform sparsity budget (see Figure 6) which is empirically shown to be effective in increasing prediction accuracy while reducing FLOPs. Moreover, (4) shows that the gradient update itself is sparse as gradient of $\mathcal{L}$ is multiplied with an indicator function of $\mathcal{S}_{g}(\mathbf{W}_{l})\neq 0$ which gets sparser over iterations (Figure 2). So ${\rm STR}$ addresses both the issues with standard PGD methods (Hard/Soft Thresholding) that we mentioned above.

The reparameterization trick using the projection operator’s functional form can be used for standard constrained optimization problems as well (assuming the projection operator has a closed-form). However, it is easy to show that in general, such a method need not converge to the optimal solution even for convex functions over convex sets. This raises a natural question about the effectiveness of the technique for sparse weights learning problem. It turns out that for sparsity constrained problems, ${\rm STR}$ is very similar to backward pruning (Hastie et al., 2009) which is a well-known technique for sparse regression. Note that, similar to Hard/Soft Thresholding, standard backward pruning also does not support differentiable tuning thresholds which makes it challenging to apply it to DNNs.

Experiments

This section showcases the experimentation followed by the observations from applying ${\rm STR}$ for (a) unstructured sparsity in CNNs and (b) structured sparsity in RNNs.

ImageNet-1K (Deng et al., 2009) is a widely used large-scale image classification dataset with 1K classes. All the CNN experiments presented are on ImageNet-1K. ResNet50 (He et al., 2016) and MobileNetV1 (Howard et al., 2017) are two popular CNN architectures. ResNet50 is extensively used in literature to show the effectiveness of sparsity in CNNs. Experiments on MobileNetV1 argue for the generalizability of the proposed technique (${\rm STR}$). Dataset and models’ details can be found in Appendix A.7.

${\rm STR}$ was compared against strong state-of-the-art baselines in various sparsity regimes including GMP (Gale et al., 2019), DSR (Mostafa & Wang, 2019), DNW (Wortsman et al., 2019), SNFS (Dettmers & Zettlemoyer, 2019), RigL (Evci et al., 2020) and DPF (Lin et al., 2020). GMP and DNW always use a uniform sparsity budget. RigL, SNFS, DSR, and DPF were compared in their original form. Exceptions for the uniform sparsity are marked in Table 1. The “+ ERK” suffix implies the usage of ERK budget (Evci et al., 2020) instead of the original sparsity budget. Even though VD (Molchanov et al., 2017) achieves state-of-the-art results, it is omitted due to the 2$\times$ memory and 4$\times$ compute footprint during training. Typically VD and IMP use a global threshold for global sparsity (GS) (Han et al., 2015) which can also be learnt using ${\rm STR}$. The unstructured sparsity experiments presented compare the techniques which induce layer-wise sparsity. Note that ${\rm STR}$ is generalizable to other scenarios as well. Open-source implementations, pre-trained models, and reported numbers of the available techniques were used as the baselines. Experiments were run on a machine with 4 NVIDIA Titan X (Pascal) GPUs.

All baselines use the hyperparameter settings defined in their implementations/papers. The experiments for ${\rm STR}$ use a batch size of 256, cosine learning rate routine and are trained for 100 epochs following the hyperparameter settings in (Wortsman et al., 2019) using SGD + momentum. ${\rm STR}$ has weight-decay ($\lambda$) and $s_{\rm init}$ hyperparameters to control the overall sparsity in CNNs and can be found in Appendix A.6. $\text{GMP}_{1.5\times}$ (Gale et al., 2019) and $\text{RigL}_{5\times}$ (Evci et al., 2020) show that training the networks longer increases accuracy. However, due to the limited compute and environmental concerns (Schwartz et al., 2019), all the experiments were run only for around 100 epochs ($\sim$3 days each). Unstructured sparsity in CNNs with ${\rm STR}$ is enforced by learning one threshold per-layer as shown in Figure 3. PyTorch ${\rm STRConv}$ code can be found in Algorithm 1 of Appendix.

1.2 ResNet50 on ImageNet-1K

A fully dense ResNet50 trained on ImageNet-1K has 77.01% top-1 validation accuracy. ${\rm STR}$ is compared extensively to other baselines on ResNet50 in the sparsity ranges of 80%, 90%, 95%, 96.5%, 98%, and 99%. Table 1 shows that DNW and GMP are state-of-the-art among the baselines across all the aforementioned sparsity regimes. As ${\rm STR}$ might not be able to get exactly to the sparsity budget, numbers are reported for the models which nearby. Note that the 90.23% sparse ResNet50 on ImageNet-1K with ${\rm STR}$ is referred to as the 90% sparse ResNet50 model learnt with ${\rm STR}$.

${\rm STR}$ comfortably beats all the baselines across all the sparsity regimes as seen in Table 1 and is the state-of-the-art for unstructured sparsity. Figure 4 shows that ${\rm STR}$ forms a frontier curve encompassing all the baselines at all the levels of sparsity. Very few methods are stable in the ultra sparse regime of 98-99% sparsity and GMP can achieve 99% sparsity. ${\rm STR}$ is very stable even in the ultra sparse regime, as shown in Table 1 and Figure 4, while being up to 10% higher in accuracy than GMP at 99% sparsity.

${\rm STR}$ induces non-uniform sparsity across layers, Table 1 and Figure 5 show that ${\rm STR}$ produces models which have lower or similar inference FLOPs compared to the baselines while having better prediction accuracy in all the sparsity regimes. This hints at the fact that ${\rm STR}$ could be redistributing the parameters thereby reducing the FLOPs. In the 80% sparse models, ${\rm STR}$ is at least 0.19% better in accuracy than the baselines while having at least 60M (6.5%) lesser FLOPs. Similarly, ${\rm STR}$ has state-of-the-art accuracy in 90%, 95%, and 96.5% sparse regimes while having at least 68M (16.5%), 45M (22%) and 140M (54%) lesser FLOPs than the best baselines respectively. In the ultra sparse regime of 98% and 99% sparsity, ${\rm STR}$ has similar or slightly higher FLOPs compared to the baselines but is up to 4.6% and 10% better in accuracy respectively. Table 1 summarizes that the non-uniform sparsity baselines like SNFS, SNFS+ERK, and RigL+ERK can have up to 2-4$\times$ higher inference cost (FLOPs) due to non-optimal layer-wise distribution of the parameter weights.

Observations: ${\rm STR}$ on ResNet50 shows some interesting observations related to sparsity and inference cost (FLOPs). These observations will be further discussed in Section 5:

${\rm STR}$ is state-of-the-art for unstructured sparsity.

${\rm STR}$ minimizes inference cost (FLOPs) while maintaining accuracy in the 80-95% sparse regime.

${\rm STR}$ maximizes accuracy while maintaining inference cost (FLOPs) in 98-99% ultra sparse regime.

${\rm STR}$ learns a non-uniform layer-wise sparsity, shown in Figure 6, which shows that the initial layers of the CNN can be sparser than that of the existing non-uniform sparsity methods. All the learnt non-uniform budgets through ${\rm STR}$ can be found in Appendix A.3.

Figure 6 also shows that the last layers through ${\rm STR}$ are denser than that of the other methods which is contrary to the understanding in the literature of non-uniform sparsity (Mostafa & Wang, 2019; Dettmers & Zettlemoyer, 2019; Evci et al., 2020; Gale et al., 2019). This leads to a sparser backbone for transfer learning. The backbone sparsities can be found in Appendix A.3.

Figure 7 shows the layer-wise FLOPs distribution for the non-uniform sparsity methods. ${\rm STR}$ adjusts the FLOPs across layers such that it has lower FLOPs than the baselines. Note that the other non-uniform sparsity budgets lead to heavy compute overhead in the initial layers due to denser parameter tensors.

${\rm STR}$ can also induce global sparsity (GS) (Han et al., 2015) with similar accuracy at $\sim 2\times$ FLOPs compared to layer-wise for 90-98% sparsity (details in Appendix A.5.1).

1.3 MobileNetV1 on ImageNet-1K

MobileNetV1 was trained on ImageNet-1K for unstructured sparsity with ${\rm STR}$ to ensure generalizability. Since GMP is the state-of-the-art baseline as shown earlier, ${\rm STR}$ was only compared to GMP for 75% and 90% sparsity regimes. A fully dense MobileNetV1 has a top-1 accuracy of 71.95% on ImageNet-1K. GMP (Zhu & Gupta, 2017) has the first layer and depthwise convolution layers dense for MobileNetV1 to ensure training stability and maximize accuracy.

Table 2 shows the ${\rm STR}$ is at least 0.65% better than GMP for 75% sparsity, while having at least 62M (38%) lesser FLOPs. More interestingly, ${\rm STR}$ has state-of-the-art accuracy while having up to 50% (40M) lesser FLOPs than GMP in the 90% sparsity regime. All the observations made for ResNet50 hold for MobileNetV1 as well. The sparsity and FLOPs distribution across layers can be found in Appendix A.4.

2 Structured Sparsity in RNNs

Google-12 is a speech recognition dataset that has 12 classes made from the Google Speech Commands dataset (Warden, 2018). HAR-2 is a binarized version of the 6-class Human Activity Recognition dataset (Anguita et al., 2012). These two datasets stand as compelling cases for on-device resource-efficient machine learning at the edge. Details about the datasets can be found in Appendix A.7.

The baseline is low-rank FastGRNN where the ranks of the matrices are preset (Kusupati et al., 2018). EdgeML (Dennis et al., ) FastGRNN was used for the experiments with the hyperparameters suggested in the paper and is referred to as vanilla training. Hyperparameters for the models can be found in Appendix A.6.

2.2 FastGRNN on Google-12 and HAR-2

Table 3 presents the results for low-rank FastGRNN with vanilla training and ${\rm STR}$. Full-rank non-reparameterized FastGRNN has an accuracy of 92.60% and 96.10% on Google-12 and HAR-2 respectively.

${\rm STR}$ outperforms vanilla training by up to 1.67% in four different model-size reducing rank settings on Google-12. Similarly, on HAR-2, ${\rm STR}$ is better than vanilla training in all the rank settings by up to 2.47%. Note that the accuracy of the low-rank models obtained by ${\rm STR}$ is either better or on-par with the full rank models while being around 50% and 70% smaller in size (low-rank) for Google-12 and HAR-2 respectively.

These experiments for structured sparsity in RNNs show that ${\rm STR}$ can be applied to obtain low-rank parameter tensors. Similarly, ${\rm STR}$ can be extended for filter/channel pruning and block sparsity (He et al., 2017; Huang & Wang, 2018; Liu et al., 2019) and details for this adaptation can be found in Appendix A.5.2.

Discussion and Drawbacks

${\rm STR}$’s usage for unstructured sparsity leads to interesting observations as noted in Section 7. It is clear from Table 1 and Figures 4, 5 that ${\rm STR}$ achieves state-of-the-art accuracy for all the sparsity regimes and also reduces the FLOPs in doing so. ${\rm STR}$ helps in learning non-uniform sparsity budgets which are intriguing to study as an optimal non-uniform sparsity budget can ensure minimization of FLOPs while maintaining accuracy. Although it is not clear why ${\rm STR}$’s learning dynamics result in a non-uniform budget that minimizes FLOPs, the reduction in FLOPs is due to the better redistribution of parameters across layers.

Non-uniform sparsity budgets learnt by ${\rm STR}$ have the initial and middle layers to be sparser than the other methods while making the last layers denser. Conventional wisdom suggests that the initial layers should be denser as the early loss of information would be hard to recover, this drives the existing non-uniform sparsity heuristics. As most of the parameters are present in the deeper layers, the existing methods tend to make them sparser while not affecting the FLOPs by much. ${\rm STR}$, on the other hand, balances the FLOPs and sparsity across the layers as shown in Figures 6, 7 making it a lucrative and efficient choice. The denser final layers along with sparser initial and middle layers point to sparser CNN backbones obtained using ${\rm STR}$. These sparse backbones can be viable options for efficient representation/transfer learning for downstream tasks.

Table 4 shows the effectiveness/transferability of the learnt non-uniform budget through ${\rm STR}$ for 90% sparse ResNet50 on ImageNet-1K using DNW (Wortsman et al., 2019). DNW typically takes in a uniform sparsity budget and has an accuracy of 74% for a 90% sparse ResNet50. Using ERK non-uniform budget for 90% sparsity results in a 0.1% increase in accuracy at the cost 2.35$\times$ inference FLOPs. Training DNW with the learnt budget from ${\rm STR}$ results in a reduction of FLOPs by 66M (16%) while maintaining accuracy. In the 95% sparsity regime, the learnt budget can improve the accuracy of DNW by up to 1.42% over uniform along with a reduction in FLOPs by at least 22M (11%).

Similarly, these budgets can also be used for other methods like GMP (Zhu & Gupta, 2017). Table 5 shows that the learnt sparsity budgets can lead to an increase in accuracy by 0.22% and 1.57% in 90% and 98% sparsity regimes respectively when used with GMP. Accuracy gains over uniform sparsity are also accompanied by a significant reduction in inference FLOPs. Note that the learnt non-uniform sparsity budgets can also be obtained using smaller representative datasets instead of expensive large-scale experiments.

The major drawback of ${\rm STR}$ is the tuning of the weight-decay parameter, $\lambda$ and finer-tuning with $s_{\rm init}$ to obtain the targeted overall sparsity. One way to circumvent this issue is to freeze the non-uniform sparsity distribution in the middle of training when the overall sparsity constraints are met and train for the remaining epochs. This might not potentially give the best results but can give a similar budget which can be then transferred to methods like GMP or DNW. Another drawback of ${\rm STR}$ is the function $g$ for the threshold. The stability, expressivity, and sparsification capability of ${\rm STR}$ depends on $g$. However, it should be noted that sigmoid and exponential functions work just fine, as $g$, for ${\rm STR}$.

Conclusions

This paper proposed Soft Threshold Reparameterization (${\rm STR}$), a novel use of the soft-threshold operator, for the weights in DNN, to smoothly induce sparsity while learning layer-wise pruning thresholds thereby obtaining a non-uniform sparsity budget. Extensive experimentation showed that ${\rm STR}$ is state-of-the-art for unstructured sparsity in CNNs for ImageNet-1K while also being effective for structured sparsity in RNNs. Our method results in sparse models that have significantly lesser inference costs than the baselines. In particular, ${\rm STR}$ achieves the same accuracy as the baselines for 90% sparse MobileNetV1 with 50% lesser FLOPs. ${\rm STR}$ has $\sim$10% higher accuracy than the existing methods in ultra sparse (99%) regime for ResNet50 showing the effectiveness of the learnt non-uniform sparsity budgets. ${\rm STR}$ can also induce low-rank structure in RNNs while increasing the prediction accuracy showing the generalizability of the proposed reparameterization. Finally, ${\rm STR}$ is easy to adapt and the learnt budgets are transferable.

Acknowledgments

We are grateful to Keivan Alizadeh, Tapan Chugh, Tim Dettmers, Erich Elsen, Utku Evci, Daniel Gordon, Gabriel Ilharco, Sarah Pratt, James Park, Mohammad Rastegari and Matt Wallingford for helpful discussions and feedback. Mitchell Wortsman is in part supported by AI2 Fellowship in AI. Sham Kakade acknowledges funding from the Washington Research Foundation for Innovation in Data-intensive Discovery, and the NSF Awards CCF-1637360, CCF-1703574, and CCF-1740551. Ali Farhadi acknowledges funding from the NSF Awards IIS 1652052, IIS 17303166, DARPA N66001-19-2-4031, 67102239 and gifts from Allen Institute for Artificial Intelligence.

References

Appendix A Appendix

$0<g(s)$, $\lim\limits_{s\to-\infty}g(s)=0$, and $\lim\limits_{s\to\infty}g(s)=\infty$.

$g^{\prime}(s_{\rm init})<1$ providing us a handle on the dynamics of $s$.

The gradient of $s_{l}\ \forall\ l\in\left[L\right]$ takes an even interesting form

A.3 ResNet50 Learnt Budgets and Backbone Sparsities

Table 6 lists the non-uniform sparsity budgets learnt through ${\rm STR}$ across the sparsity regimes of 80%, 90%, 95%, 96.5%, 98% and 99% for ResNet50 on ImageNet-1K. The table also lists the backbone sparsities of every budget. It is clear that ${\rm STR}$ results in a higher than expected sparsity in the backbones of CNNs resulting in efficient backbones for transfer learning.

Table 7 summarizes all the sparsity budgets for 90% sparse ResNet50 on ImageNet-1K obtained using various methods. This table also shows that the backbone sparsities learnt through ${\rm STR}$ are considerably higher than that of the baselines.

One can use these budgets directly for techniques like GMP and DNW for a variety of datasets and have significant accuracy gains as shown in the Table 4.

A.4 MobileNetV1 Sparsity and FLOPs Budget Distributions

Table 8 summarizes all the sparsity budgets for 90% sparse MobileNetV1 on ImageNet-1K obtained using various methods. Note that GMP here makes the first and depthwise (dw) convolution layers dense, hence it is not the standard uniform sparsity. This table also shows that the backbone sparsities learnt through ${\rm STR}$ are considerably higher than that of GMP.

Figure 8 shows the sparsity distribution across layers when compared to GMP and Figure 9 shows the FLOPs distribution across layers when compared to GMP for 90% sparse MobileNetV1 models on ImageNet-1K.

It is interesting to notice that ${\rm STR}$ automatically keeps depthwise separable (the valleys in Figure 8) convolution layers less sparse than the rest to maximize accuracy which is the reason GMP keeps them fully dense.

A.5 STRSTR{\rm STR} Adaptations

Algorithm 1 has comments suggesting the simple modifications required for global and per-weight sparsity.

${\rm STR}$ can be trivially modified to learn the global threshold to induce global sparsity like in (Han et al., 2015; Frankle & Carbin, 2019). Instead of having an $s_{l}$ per layer $l$, share all the $s_{l}$ to create one single learnable global threshold $s_{g}$. This can be implemented by a simple modification in Algorithm 1. ${\rm STR}$’s capability to induce global sparsity was evaluated on ResNet50 for ImageNet-1K for 90-98% sparsity regimes.

Table 9 shows the performance of ${\rm STR}$-GS that learns the global threshold to induce global sparsity. While the accuracies are comparable to the state-of-the-art if not better, they do come at cost of $\sim 2\times$ inference cost compared to layer-wise sparsity due to poor non-uniform sparsity distribution which is a result of difference converged values of weights in each of the layers. ${\rm STR}$-GS has numbers similar to IMP (Frankle & Carbin, 2019) while being able to learn the threshold stablely.

A.5.2 STRSTR{\rm STR} for Filter/Channel Pruning

Let us assume there are $n_{out}$ filters of size $k\times k\times n_{in}$ in a given layer. Typically in channel/filter pruning techniques, each of these $n_{out}$ filters have an importance factor that represents the utility of the filter and is used to scale the corresponding filter. For a filter $f_{i}$ there exists a importance scalar $m_{i}$ learnt or obtained in some fashion and is used to get the effective filter in use $\hat{f}_{i}=m_{i}\cdot f_{i}$ where $m_{i}$ is broadcasted to scale $f_{i}$. In practice, $m_{i}$ is heuristically made to go to to induce structured sparsity through channel/filter pruning. Let us stack all the importance scalars of the filters in the layer, $\{m_{i}\}_{i\in[n_{out}]}$, as vector $\bm{m}_{l}$ where $l$ is the layer index. Now, this reduces to the same problem of inducing sparsity in a vector as in the learning of low-rank in RNN presented in Section 4.2.1. ${\rm STR}$ will be applied to each of the $\{m_{l}\}_{l\in[L]}$ where $L$ is the total number of layers in a deep neural network. The inference will use the importance scalars through ${\rm STR}$ ensuring channel/filter pruning due to the induced sparsity. This is very similar to the work-flow we used to induce low-rank in RNNs.

A.5.3 STRSTR{\rm STR} for Per-weight Pruning or Mask Learning

The adaptation of ${\rm STR}$ for per-weight pruning or mask learning is simple and is similar to layer-wise or global sparsity. Changing $s_{l}\to\mathbf{S}_{l}$ ie., changing the layer-wise thresholds from a scalar to a tensor of the size of $\mathbf{W}_{l}$ will hep STR adapt to do per-weight pruning or mask learning as discussed in the recent works (Zhou et al., 2019; Savarese et al., 2019; Ramanujan et al., 2020). We have explored this using a couple of experiments on CIFAR-10 (Krizhevsky et al., 2009) and ImageNet-1K. We observed that high amounts of sparsity were induced and the routine is very aggressive compared to other sparsification methods. For example, we were able to get 90% accuracy on CIFAR-10 using ResNet18 at a staggering 99.63% sparsity (270$\times$ lesser parameters than the dense model) which results in 41K parameters pushing it into very under parameterized regime. We suggest caution when running per-weight sparsity experiments with any method due to the high variance in the final accuracy.

A.6 Hyperparameters for Reproducibility

All the ResNet50 experiments use a batchsize of 256, cosine learning rate with warm-up as in (Wortsman et al., 2019) and trained for 100 epochs. $\lambda$ is the weight-decay hyperparameter. $s_{\text{init}}$ is the initial value of all $s_{i}$ where $i$ is the layer number. The hyper parameter setting for each of the sparse model can be found in Table 10.

All the MobileNetV1 experiments use a batchsize of 256, cosine learning rate with warm-up as in (Wortsman et al., 2019) and trained for 100 epochs. $\lambda$ is the weight-decay hyperparameter. $s_{\text{init}}$ is the initial value of all $s_{i}$ where $i$ is the layer number. The hyper parameter setting for each of the sparse model can be found in Table 11.

All the CNN experiments use $g(s)=\frac{1}{1+e^{-s}}$ for the ${\rm STR}$.

All the RNN experiments use $g(s)=e^{s}$ for the ${\rm STR}$.

A.7 Dataset and Model Details

ImageNet-1K: ImageNet-1K has RGB images with 224$\times$224 dimensions. The dataset has 1.3M training images, 50K validation images and 1000 classes. Images were transformed and augmented with the standard procedures as in (Wortsman et al., 2019).

Google-12: Google Speech Commands dataset (Warden, 2018) contains 1 second long utterances of 30 short words (30 classes) sampled at 16KHz. Standard log Mel-filter-bank featurization with 32 filters over a window size of 25ms and stride of 10ms gave 99 timesteps of 32 filter responses for a 1-second audio clip. For the 12 class version, 10 classes used in Kaggle’s Tensorflow Speech Recognition challenge were used and the remaining two classes were noise and background sounds (taken randomly from the remaining 20 short word utterances). The datasets were zero mean - unit variance normalized during training and prediction. Google-12 has 22,246 training points, 3,081 testing points. Each datapoint has 99 timesteps with each input being 32 dimensional making the datapoint 3,168 dimensional.

HAR-2: Human Activity Recognition (HAR) dataset was collected from an accelerometer and gyroscope on a Samsung Galaxy S3 smartphone. The features available on the repository were directly used for experiments. The 6 activities were merged to get the binarized version. The classes {Sitting, Laying, Walking_Upstairs} and {Standing, Walking, Walking_Downstairs} were merged to obtain the two classes. The dataset was zero mean - unit variance normalized during training and prediction. HAR-2 has 7,352 training points and 2,947 test points. Each datapoint has 1,152 dimensions, which will be split into 128 timesteps leading to dimensional per timestep inputs.

ResNet50: ResNet50 is a very popular CNN architecture and is widely used to showcase the effectiveness of sparsification techniques. ResNet50 has 54 parameter layers (including fc) and a couple of pooling layers (which contribute minimally to FLOPs). All the batchnorm parameters are left dense and are learnt during the training. ${\rm STR}$ can be applied per-layer, per-channel and even per-weight to obtain unstructured sparsity and the aggressiveness of sparsification increases in the same order. This paper only uses per-layer ${\rm STR}$ which makes it have 54 additional learnable scalars. The layer-wise parameters and FLOPs can be seen in Tables 7 and 6. All the layers had no bias terms.

MobileNetV1: MobileNetV1 is a popular efficient CNN architecture. It is used to showcase the generalizability of sparsification techniques. MobileNetV1 has 28 parameter layers (including fc) and a couple of pooling layers (which contribute minimally to FLOPs). All the batchnorm parameters are left dense and are learnt during the training. ${\rm STR}$ can be applied per-layer, per-channel and even per-weight to obtain unstructured sparsity and the aggressiveness of sparsification increases in the same order. This paper only uses per-layer ${\rm STR}$ which makes it have 28 additional learnable scalars. The layer-wise parameters and FLOPs can be seen in Tables 8. All the layers had no bias terms.

FastGRNN: FastGRNN’s update equations can be found in (Kusupati et al., 2018). FastGRNN, in general, benefits a lot from the low-rank reparameterization and this enables it to be deployed on tiny devices without losing any accuracy. FastGRNN’s biases and final classifier are left untouched in all the experiments and only the input and hidden projection matrices are made low-rank. All the hyperparameters were set specific to the datasets as in Kusupati et al. (2018).

A.8 Hard Threshold vs Soft Threshold

Figure 10 shows the difference between hard thresholding and soft thresholding for the same threshold value of $\alpha=2$. It is clear from Figure 10 that soft-threshold is a continuous function that is sub-differentiable. The abrupt change in hard-threshold leads to instability in training sometimes increasing dependence on fine tuning of the obtained sparse network. Soft-threshold is robust to such issues.