SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Minsik Cho
Saurabh Adya
Devang Naik
{minsik, sadya, naik.d}@apple.com
DNN pruning is a popular way to reduce the size of a model, improve the inference latency, and minimize the power consumption on DNN accelerators. However, existing approaches might be too complex, expensive or ineffective to apply to a variety of vision/language tasks, DNN architectures and to honor structured pruning constraints. In this paper, we propose an efficient yet effective train-time pruning scheme, Parameter-free Differentiable Pruning (PDP), which offers state-of-the-art qualities in model size, accuracy, and training cost. PDP uses a dynamic function of weights during training to generate soft pruning masks for the weights in a parameter-free manner for a given pruning target. While differentiable, the simplicity and efficiency of PDP make it universal enough to deliver state-of-the-art random/structured/channel pruning results on various vision and natural language tasks. For example, for MobileNet-v1, PDP can achieve  $68.2\%$  top-1 ImageNet1k accuracy at  $86.6\%$  sparsity, which is  $1.7\%$  higher accuracy than those from the state-of-the-art algorithms. Also, PDP yields over  $83.1\%$  accuracy on Multi-Genre Natural Language Inference with  $90\%$  sparsity for BERT, while the next best from the existing techniques shows  $81.5\%$  accuracy. In addition, PDP can be applied to structured pruning, such as N:M pruning and channel pruning. For 1:4 structured pruning of ResNet18, PDP improved the top-1 ImageNet1k accuracy by over  $3.6\%$  over the state-of-the-art. For channel pruning of ResNet50, PDP reduced the top-1 ImageNet1k accuracy by  $0.6\%$  from the state-of-the-art.
Deep neural networks (DNN) have shown human performance on complex cognitive tasks [46], but their deployment onto mobile/edge devices for enhanced user experience (i.e., reduced latency and improved privacy) is still challenging. Most such on-device DNN systems are battery-powered and heavily resource-constrained, thus requiring high power/compute/storage efficiency [22, 48, 50, 53].
Such efficiency can be accomplished by mixing and matching various techniques, such as designing efficient DNN architectures like MobileNet/MobileViT/ MobileOne [36, 43, 48], distilling a complex DNN into a smaller architecture [41], quantizing/compressing the weights of DNNs [8, 18, 23, 31, 39, 58], and pruning near-zero weights [28, 34, 40, 44, 52, 56, 57, 61]. Also, pruning is known to be highly complementary to quantization/compression [51] when optimizing a DNN model. Training a larger model and then compressing it by pruning has been shown to be more effective in terms of model accuracy than training a smaller model from the beginning [32]. However, pruning comes at the cost of degraded model accuracy, and the trade-off is not straightforward [28].
Hence, a desirable pruning algorithm should achieve high accuracy and accelerate inference for various types of networks without significant training overheads in costs and complexity. In this
STRa	FTWTb	CSc	GraNetd	OptGe	ACDCFf	MVPg	POFAh	PDP
Conv-Net Accuracy	✓	✓✓	✓✓	✓	✓✓	✓✓	?	?	✓✓✓
Transformer Accuracy	✓	?	?	?	✓	?	✓✓	✓✓	✓✓✓
Inference speed	✓✓✓	✓✓✓	✓✓✓	✓	✓✓✓	✓	?	?	✓✓✓
Training speed	✓✓✓	?	✓	✓✓	✓✓	✓	?	?	✓✓✓
Training stability	✓	?	✓	✓✓	✓✓✓	✓✓✓	?	?	✓✓✓
Training flow	simple complex	complex complex	complex complex	complex complex	simple
Extra parameters	a few	many	many	none	many	none	many	none	none
a Soft threshold weight re-parameterization for learnable sparsity [28].
<sup>b</sup> Fire Together Wire Together: A Dynamic Pruning Approach with Self-Supervised Mask Prediction [11]
c Winning the Lottery with Continuous Sparsification [45]
d Sparse training via boosting pruning plasticity with neuro-regeneration [34].
e Optimizing gradient-driven criteria in network sparsity: Gradient is all you need [57].
f AC/DC: alternating compressed/decompressed training of deep neural networks [40].
$^{\mathrm{g}}$  Movement Pruning: Adaptive Sparsity by Fine-Tuning [44].
h Prune Once for All: Sparse Pre-Trained Language Models [56].
Table 1: Comparison of the state-of-the-art pruning schemes and PDP: PDP can explore a good trade-off bewteen accuracy and inference speed without introducing new learnable parameters and with a simple/fast training flow.
work, we propose a simple yet effective pruning technique, Parameter-free Differentiable Pruning or PDP, which uses a dynamic function of weights to generate soft pruning masks for the weights themselves. PDP requires neither additional learning parameters [57] nor complicated training flows [40], yet offers precise control on the target sparsity level [28], while pushing the state-of-the-arts in random/structured/channel pruning. The soft masks from PDP make pruning differentiable and let training loss decide whether/how each weight will be pruned [42]. Table 1 compares PDP with the latest state-of-the-art pruning schemes, and our major contributions include:
- PDP outperforms the state-of-the-art schemes on a variety of models and tasks. Being differentiable and parameter-free, PDP offers efficient and effective pruning without complex techniques.
- PDP offers a universal and holistic approach for efficient random/structured/channel pruning, while delivering a high-quality model optimization for a given pruning target.
- With a dynamic function of weights, PDP generates a soft pruning mask which does not need training, and thus does require neither gradient synchronization nor SGD-update.
Pruning in DNN incurs a complex trade-off between model accuracy and inference speed in terms of MAC (mult-add operations) [28]. A weight can contribute differently to the model accuracy, depending on the number of times it is used for prediction (i.e., a weight in convolution filter for a large input) and the criticality of the layer it belongs to (i.e., a weight in a bottleneck layer). Therefore, even if two models are pruned to the same level, the accuracy and inference speed of each can be vastly different, which makes exploring the best trade-off challenging yet crucial in DNN pruning. A body of work in random and structured pruning has been proposed to optimize such a trade-off. Please see Fig. 8 and Section B in Appendix for additional reviews and details.
Differentiable techniques gained popularity due to the advances in network architecture search methods as well as network compression techniques [6, 8, 33], and is a powerful technique for random/structured pruning as it drives pruning based on task loss. However, differentiable pruning does not necessarily deliver the best quality pruning (i.e., many state-of-the-art results are from non-differentiable methods), because the computational overhead and training complexity can limit its benefit. Existing differentiable pruning methods can be classified into a few approaches.
Learning pruning budget allocation is to determine target sparsity for each layer in a differentiable way, rather than computing pruning masks. Optimizing and training a per-layer pruning threshold with ReLU based on dynamic re-parameterization was proposed to allocate the sparsity across all layers in STR [28], which shows good model accuracy and low inference overhead. DSA [38] finds the layer-wise pruning ratios in a differentiable fashion by computing a channel-wise keep probability drawn from the Bernoulli distribution and keeping the expected sparsity ratio satisfied.
(a) Using a trainable threshold
(b) Using soft masks with  $t$  (see Section 3.1)
Figure 1: Unlike learning a threshold parameter for differentiable pruning [28, 38] in (a), PDP generates soft masks without extra trainable parameters to accomplish differentiable pruning as in (b). The moment a weight is pruned in (a), the weight will not get any update as the mask zeros out the gradient, depriving the opportunity for better model training. However, the scheme in (a) ensures the training and test behaviors are identical, yet makes the sparsity control difficult. PDP generates such soft masks without extra parameter/training overheads unlike existing schemes [34, 40, 44, 45, 57] (see Section 3.1). While PDP reveals different behaviors during training/test times because the test needs hard masks, our soft mask allows a weight to recover from being pruned over time, which yields higher model accuracies efficiently.
Learning pruning mask as an extra parameter is a popular way for differentiable pruning. In CS [45], a new trainable mask parameter is learned to continuously remove sub-networks based on the lottery hypothesis [13]. In detail, the new mask parameter is first multiplied by a scheduled temperature parameter, and then L0-regularized inside the Sigmoid function to continuously increase the sparsity. OptG [57] proposed learning a mask parameter based on the loss changes due to pruning a particular weight, which is proportional to the gradient and the magnitude of the corresponding masked weight when approximated with STE [5]. AutoPrune [54] also introduces extra learnable parameters (controlled by a regularizer) to generate a differentiable/approximated pruning mask using STE [5]. In [42], new learnable parameters,  $\alpha$  are introduced and  $\mathrm{L}_1$ -regularized to enforce sparsity. Then,  $\alpha$  is compared with a hyper-param  $t$  to derive the masks for the model parameters (i.e., prune if  $\alpha < t$ ). As such comparison is not differentiable, comparison is approximated into a foothill function [4]. Learning channel pruning mask is proposed in SCP [25] based on the operation-specific observation that a feature map with a large and negative batch mean will get zeroed by ReLU.
Generating pruning masks is proposed in FTWT [11] where new layers are attached to the exiting main architecture. The attached new layers are fed with activations from the main network, then trained to generate masks for the weights in the main network. Thus, although it does not learn masks directly, it still requires new trainable parameters to learn how to generate masks, not to mention the changes to the model network. Using knowledge distillation to generate pruning masks is explored in NISP [55] and DCP [62], where a fully trained teacher network guides the mask generation in a way that the distortion to the layer responses can be minimized. A differentiable Markov process is studied to generate a channel pruning mask probabilistically, where a state represents a channel and the transitions from states account for the probability of keeping one channel unpruned [16].
3 PDP: Parameter-free Differentiable Pruning
Complex pruning schemes do not always yield the best quality results, and their complexity and cost can make them impractical and difficult to use. The proposed Parameter-free Differentiable Pruning (PDP) is a highly effective and efficient scheme that generates soft pruning masks using a dynamic function of weights in a parameter-free and differentiable fashion. Since PDP is differentiable, the task loss can directly guide the pruning decision, offering an effective pruning solution. Simultaneously, being parameter-free, PDP can be fast and less intrusive to the existing training flow. Overall, PDP finds a weight distribution that is best for task loss and pruning. Instead of having extra parameters, PDP indirectly influences the weight distribution for high-quality pruning. For example, if a weight
(a) PDP training flow
(b) Probability plot by  $w \in R_{[0,1]}$  with  $t = 0.6$
Figure 2: Computing  $z(w), m(w)$  for the chances for  $Z$  and  $M$  with  $t$  for the equal chance to be in  $Z$  and  $M$ .
$w$  is destined to be pruned for some reason, instead of having a new parameter to denote "to-prune", PDP lets SGD gradually make  $w$  itself smaller relatively against other parameters in the same entity, increasing its chance to be pruned over time. We will first present PDP in Section 3.1, provide in-depth analysis in Section 3.2, discuss the benefits of PDP over existing differentiable pruning approaches in Section 3.3, and then show the extension to structured and channel pruning in Section 3.4.
3.1 PDP Algorithm
To address the drawbacks of existing differentiable pruning algorithms, we propose PDP. A soft mask should ideally represent the chance of a weight  $w$  being in one of two symbolic states, "to-prune" (noted as  $Z$ ) or "not-to-prune" (noted as  $M$ ), without requiring extra parameters or expensive bookkeeping. While the chance of  $w$  being in either state is not straightforward to compute, PDP generates a soft mask based on the fact that there exists an equal chance point for both states. Let us consider differentiable functions,  $z, m: \mathbb{R}[0, \infty] \mapsto \mathbb{R}[0, 1]$ , to compute the chances of being in  $Z$  and  $M$ , respectively, which must satisfy the following conditions as a soft mask for magnitude-based pruning:
-  $z(|a|) < z(|b|)$  for  $|a| > |b|$ : a weight with a smaller magnitude has a higher chance to be in  $Z$ .
-  $m(|a|) > m(|b|)$  for  $|a| > |b|$ : a weight with a larger magnitude has a higher chance to be in  $M$ .
-  $z(w) + m(w) = 1$  for any  $w$ : the total probability is 1.
$$
z (w) = \left\{ \begin{array}{l l} 1 & \text {i f} w = 0 \\ 0 & \text {i f} | w | = \infty \\ \frac {1}{2} & \text {i f} | w | = t \end{array} \right.
$$
Then, by the monotonicity and continuity, there exists  $t \in \mathbb{R}_{\geq 0}$  such that  $z(t) = m(t) = 0.5$  (the equal chance for  $Z$  and  $\bar{M}$ ), which leads to the following boundary conditions on the left. Any function that satisfies these conditions can be used to compute  $m(w)$  as a soft mask of  $w$  for train-time pruning. Let denote that
topK(X,k)/bottomK(X,k) is selecting the largest/smallest  $k$  elements from a matrix  $X$ ,  $\operatorname{abs}(X)$  is an element-wise absolute operation, and  $n(X)$  returns the element count. In PDP, we uniquely identify  $t$  for a given prune ratio  $r \in [0,1)$  for a layer with a weight matrix  $W$  as in Fig. 2 (a) based on the pruning context.
- The sparsity  $r$  for each  $W$  can be easily obtained by sorting the weights from the network per magnitude after a few epochs w.r.t the global target sparsity as a one-time task, by a pruning budget allocation algorithm [14, 28], or set manually by a user.
- Right after the SGD weight update,  $t$  is computed for the weights  $W$  in each layer or entity. The role of  $t$  is to abstract the current weight distribution of each layer/entity for pruning.
- During forward-pass, a soft mask,  $m(w)$  for the weight  $w$  is computed and then the masked weight  $\hat{w}$  is applied.  $\tau$  is the temperature parameter (see Section E in Appendix for details).
- Computing  $t$  and generating  $\hat{w}$  repeat iteratively to adapt to the updated weight distribution.
Figure 2 (b) shows how the value of  $t$  is obtained in PDP and a soft mask is computed. Specifically,  $t$  is set to the value that is halfway between the largest pruned weight and the smallest unpruned weight when a hard mask is applied for a given sparsity ratio  $r$ . This ensures that each weight has an equal probability of being pruned or kept. As a result, PDP satisfies all the constraints for  $z$  and  $m$ . More details on PDP training flow are in Section D in Appendix.
PDP uses a dynamic function of  $t$  to generate soft pruning masks of  $W$  without the need for any extra trainable parameters. Instead, PDP lets the weights of the network adjust themselves such that the information that would otherwise be learned by the extra trainable parameters is instead fused into the weights themselves and their distribution. This is possible because each weight  $w$  is not only a coefficient in a layer, but also an indicator of the relative chance of that weight being pruned against the other weights in  $W$ . This relative chance is determined by the value of  $t$ , as shown in Figure 2 (b). Overall, PDP looks simple but is shown to be quite effective, which opens up broad applicability.
3.2 PDP Gradient Analysis
To better understand how PDP helps pruning decision, we derive the gradient of the masked weight  $\hat{w}$ . As in Fig. 2 (a),  $\hat{w}$  is the following:
$$
\hat {w} = m (w) \cdot w = \frac {e ^ {\frac {w ^ {2}}{\tau}}}{e ^ {\frac {w ^ {2}}{\tau}} + e ^ {\frac {t ^ {2}}{\tau}}} \cdot w \tag {1}
$$
Then, the gradient of  $w$ ,  $\Delta w$  can be simplified as the next:
$$
\Delta w = m (w) \Delta \hat {w} + 2 \frac {w ^ {2}}{\tau} m (w) \{1 - m (w) \} \Delta \hat {w} \tag {2}
$$
where we can make the following observations:
- The 1st term is a typical gradient in mask-based pruning, and the 2nd term is an additional gradient with a positive factor (i.e.,  $\frac{w^2}{\tau} m(w)\{1 - m(w)\} \geq 0$ ) from PDP.
- If  $m(w) \approx 0$  or hard-prune,  $\Delta w \approx 0$ , which is true for any pruning algorithm.
- If  $m(w) \approx 1$  or hard-not-to-prune,  $\Delta w \approx \Delta \hat{w}$ , which is true for any pruning algorithm.
- When  $m(w) \approx 0.5$  (i.e., pruning decision is unclear),  $m_w(1 - m_w)$  is maximized and accordingly  $\Delta w$  is too, boosting the  $w$  movement.
-  $\tau$  serves as an inverse scaling factor for the boosted gradient.
Hence, PDP will accelerate the SGD updates for the weights near the pruning boundary (i.e.,  $t$ ) toward a loss-decreasing direction, which can encourage the weights to settle with the proper pruning decision at the end. Even if the current gradient is not globally beneficial for the task, many second chances will eventually help recover the damages in an accelerated manner. We conjecture these features enable PDP to make a better pruning decision for the weight on the boundary.
3.3 PDP vs. Existing Differentiable Pruning Strategies
Learning Pruning Budget Allocation focuses on obtaining a pruning threshold in a differentiable way based on marginal loss or  $\mathrm{L}_1$ -regularization [25, 38]. Meanwhile, PDP directly generates a soft pruning mask for each weight also in a differentiable way. Fig. 1 illustrates the differences with an example. Learning a pruning threshold is helpful for global pruning budget allocation, as the threshold gets adjusted per the task loss as in Fig. 1 (a), but has the following drawbacks:
- Sparsity level is hard to control, as the pruning threshold is not explicitly related to the sparsity, as masks and the thresholds are not directly co-optimized. In Fig. 1 (a), after one weight update, the threshold is reduced from 0.205 to 0.202, but the sparsity is increased from  $50\%$  to  $66\%$  as  $w_{4}$  becomes smaller than the threshold.
- Once a weight gets pruned during training, it does not get updated as the gradient is masked out. The bottom weights  $w_{\{0,1,2,3\}}$  get no weight update due to the zero masks, and remain pruned.
On the contrary, PDP allows all the weights to be updated through soft masks as in Fig. 1 (b), providing higher flexibility and recovery from undesirable pruning [17]. For example, consider the weight  $w_{2}$  of a value 0.19 which would have been permanently pruned in (a). PDP discourages a weight from being permanently pruned through the weight update,  $w_{2}$  still receives a scaled-down gradient. If a near-zero weight continues to get negative gradients over time (although scaled down by the soft mask), it can eventually get unpruned at the end of training. Similarly, if a very large weight gets positive gradients many times enough to be near zero, it will get pruned eventually, even
(a) PDP: epoch 20
(b) PDP: epoch 30
(c) PDP: epoch 40
(d) PDP: epoch 50
(e) PDP: epoch 60
(f) PDP: epoch 70
(g) PDP: epoch 80
(h) PDP: epoch 90
(i)STR:epoch 20
(j)STR:epoch 30
(k)STR:epoch 40
Figure 3: The effects of PDP and STR [28] (from Table 1) for the first 3x3 Conv2d layer with 32 filters in MobileNet-v1 on ImageNet1k where each small rectangle represents one 3x3 kernel: For PDP, during the entire end-2-end training, we temporarily round the soft mask value to make the pruning decisions. The white cell indicates such pruning decision for the corresponding weight has been flipped at least once during the particular epoch, and the black cell indicates the other case.
(1) STR: epoch 50,70,80,90
if it was not pruned in the early stage. Hence, such gradual pruning decision over time during training allows PDP to make better pruning decisions w.r.t. the task loss.
Such soft masking allows pruning decisions to be flipped multiple times during the entire training process, and such effects are compared between PDP and STR [28] (from Table 1) in Fig. 3 where the pruning decisions for the first Conv2d (3x3 kernels with 32 output channels) in MobileNet-v1 for the end-to-end ImageNet1k training are captured. The white cell means that the pruning decision has been flipped (from to-prune to not-to-prune or vice-versa), while the black cell means the decision has never changed, during the given epoch. Then, we can observe from Fig. 3 that while PDP in (a)-(h) continues to flip the decisions until the very late training stage, STR in (i)-(l) finalizes the pruning decisions early (i.e., in the first 30 epochs).
Learning/Generating Pruning Masks with Extra Parameters allows the pruning decision to be driven by a task loss through back-propagation rather than the weight value itself (i.e., a hard mask will zero out the gradient of a pruned weight), thus addressing the problems with differentiable pruning budget allocation, but comes with the following issues.
- A pruning mask is (or derives from) a learnable parameter. Hence the number of total trainable parameters increases significantly, making the training process slow and complex [11, 44, 45, 57].
- A hard mask is approximated into a soft mask using a differentiable function, without guaranteeing the key properties of a pruning mask, such as the [0,1] value range or monotonicity [42].
Table 2 compares parameter-free PDP with a differentiable mask pruning scheme, CS [45] from Table 1 (see Table 8 in Appendix for detail). The results show that PDP outperforms CS and is 2.8 faster, without adding extra trainable parameters or complicating the training flow. Table 3 also demonstrates being parameter-free can provide the substantial training efficiency gains for very large language models like GPT. When trained and sparsified on 32 GPUs in mixed-precision (to fit OptG into GPU memory) using the recipe in [1], PDP delivers the best perplexity at  $75\%$  sparsity with  $42\%$  lower cost than OptG [57].
Method	Top-1 (%)	Sparsity (%)	Model #param	Extra #param	Runtime min	Inference MAC(xe6)
Dense	91.4	0	273k	0	22.1	41.0
CS	89.1	86.3	273k	267k	84.7	5.80
PDP	90.4	86.3	273k	0	30.3	5.79
Table 2: Compared with CS for ResNet20/CIFAR10, PDP delivers  ${2.7}\mathrm{x}$  speed up due to no extra parameters.
Method	Perplexity	Model #param	Extra #param	GPU cost($)
Dense	22.4	163M	0
GMP [61]	37.7	163M	0	6997
OptG	33.7	163M	124M	11210
PDP	33.7	163M	0	7499
Table 3: With  ${75}\%$  sparsity,being parameter-free becomes more important for training GPT2 with OpenWebText.
(a) 1:4 pruning (75% sparsity)
(b) Channel pruning (50% sparsity)
Figure 4: PDP is simple and universal enough to be applied directly to structured and channel pruning.
3.4 PDP for Structured and Channel Pruning
The simplicity and non-intrusive nature of PDP make it readily applicable to differentiable structured and channel pruning. As an example of structured pruning, we consider N:M pruning, where only  $N$  weights are kept out of every  $M$  consecutive weights. N:M pruning is attracting high research and industrial attention because top-of-the-line GPUs support 2:4 configuration [24]. To apply PDP to N:M pruning, we apply it to every  $M$  consecutive weights of the layer, as if the layer were composed of many sub-layers, each with  $M$  weights, which is illustrated in Figure 4 (a). Since N:M dictates the target sparsity, we can easily find the threshold  $t$  and generate the soft mask, as shown in Figure 2 (a).