SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 1 Training flow for PDP
1: procedure  $\mathsf{TRAIN}(\epsilon ,s,r,W = [W_0,W_1,\ldots])$    
2: for epoch  $e$  in [0, 1, 2, s) do   
3: for each mini-batch do   
4: forward with  $[W_0,W_2,\dots ]$    
5: backward-pass and update W   
6: end for   
7: end for   
8:  $W_{p} = topK(-abs(W),r\cdot n(W))$    
9:  $[r_0,r_1,\dots] = [\frac{n(W_p\cap W_0)}{n(W_0)},\frac{n(W_p\cap W_1)}{n(W_1)},\dots]$    
10: for epoch  $e$  in  $[s,s + 1,s + 2,\dots ]$  do   
11:  $[\hat{r}_0,\hat{r}_1,\dots] = min(1,\epsilon \cdot (e - s))\cdot [r_0,r_1,\dots]$    
12: for each mini-batch do   
13: for  $i\in \{0,1,\ldots \}$  do   
14:  $W_{h} = topK(abs(W_{i}),(1 - \hat{r}_{i})\cdot n(W_{i}))$    
15:  $W_{l} = bottomK(abs(W_{i}),\hat{r}_{i}\cdot n(W)_{i})$    
16:  $t_i = 0.5\{\min (W_h) + max(W_l)\}$    
17: end for   
18: for  $i\in \{0,1,\ldots \}$  do   
19:  $[Z_i,M_i] = \mathrm{softmax}(\frac{[t_i^2\mathbb{I},W_i\circ W_i]}{\tau})//element-wise$    
20:  $\hat{W}_i = M_i\circ W_i//element-wise$    
21: forward-pass with  $\hat{W}_i$    
22: end for   
23: backward-pass and update W   
24: end for   
25: end for   
26:  $W_{i} = \lfloor M_{i}\rfloor \circ W_{i},\forall i\in \{0,1,\ldots \}$    
27: end procedure
network size, the target sparsity and so on. Hence,  $r$  is a user hyper-parameter. In our experiments, the following guidelines worked best (and this is how we selected it).
-  $r$  must be larger than the number of warm-up epochs.
- After  $r$  epochs, the validation accuracy hits consistently over the half of the accuracy upper-bound (which is  $50\%$  for classifications) for 5 epochs.
In lines 2-7, a normal training is performed for the first  $s$  epochs. Then, in lines 8 and 9, the per-layer target pruning ratio is computed by selecting the bottom  $r \cdot n(W)$  weights globally in terms of the magnitude. Then, in the remaining epochs, we use PDP to generate soft-masks as in the line 15, while gradually increasing the target ratio as in lines 10 and 11. The updated weight distribution is captured by updating  $t_i$  as in the line 16 for all weight matrices. Once the entire training is over, we binarize the last mask for each weight to output the fully pruned weight for inference as in the line 26. Overall, the average runtime complexity of PDP is  $O(W)$ , as we only need to exercise topK algorithm (i.e., sorintg  $W$  is not necessary).
While fewer learnable parameters will reduce the model footprint and the communication cost (thus speeding up the training) in a distributed training setup, it does not necessarily lead to smaller memory consumption as the soft mask from PDP needs to generate related activation tensors. To understand the memory consumption better, we made a small GPT2 case with block_size = 128 and  $n\_layer = 3$  (a downscale version of the one in Table 2). And, then GPU memory is measured at three different spots for PDP and OptG, and reported in Table 11.
- Spot0: right after the model/optimizer are created (the parameter and PyTorch overheads)
- Spot1: right after forward, before backward (the peak memory consumption)
- Spot2: right after backward, before weight update (other parameter-related overheads)
From Table 11, we can see the following:
Method	Spot0	Spot1	Spot2
OptG	2.93	18.3	3.58
PDP	1.30	16.3	2.93
Saving	1.63	2.0	0.65
- Not having extra mask parameters helps to reduce the model/optimizer overheads.
- The peak memory is dominated by the activations (both data and mask). Yet, about 10 saving in the peak memory is observed.
- After backward, the gradients for the learnable masks take up a large memory space, which have been all-reduced in the multi-GPU setting.
To isolate the benefit of the proposed soft mask in PDP from the gradual pruning in Algorithm 1, we ran PDP for ResNet50 and MobileNet-v1 training in the exactly same configuration as in Section 4 without the proposed soft mask, and obtained the following results. The result supports the efficacy of the proposed soft mask in PDP: without it, the ImageNet top1 accuracy drops by  $1.6\%$  and  $1.4\%$ , respectively.
Table 11: Memory footprint reduction from PDP.
Method	ResNet50	MobileNet-v1
PDP w/o softmax	73.1	66.8
PDP	74.7	68.2
Δ	1.6	1.4
Table 12: Top-1 accuracy with ImageNet1k witout the softmax in PDP
E Ablation Study: Hyper-Parameter  $\tau$  search
In the current PDP implementation, we use a global  $\tau$  to control the level of softness in the pruning mask. Therefore, the selection of  $\tau$  affects the model predictive power and should be carefully tuned. In order to explore the methodology for the  $\tau$  search, we tried various values for MobileNet-v2 training, and the results are plotted in Fig. 9. The selection of  $\tau$  affects the model predictive power as shown in Fig. 9 where there appears to be an optimal  $\tau$ . For examples of MobileNet-v2,  $\tau = 1e - 4$  is the best value and is used for all the experiments in Section 4.
Figure 9: MobileNet-v2 with varying  $\tau$  values.
Since, Fig. 9 shows a concave curve, one could use a binary search to find the best  $\tau$  values w.r.t. the top-1 accuracy. Also, it could be possible to cast  $\tau$  as a learnable parameter for each layer or apply some scheduling to improve the model accuracy further (as future work), but still both approaches need an excellent initial point which can be found using a binary search technique.
F Additional Result for Section 4.
Different approaches made different sparsity allocations per the characteristics of the algorithm for a given pruning target, which results in complex trade-offs between model accuracy and inference speed. We report the detailed sparsity and inference MAC break-down for each layer in Fig. 6 and Fig. 7 on ImageNet1k and summarize our observations as follows:
- OptG prunes the early convolution layers quite aggressively in ResNet18 and ResNet50, which leads to very low inference MACs as shown in Fig. 5 (a) and (b), yet at the cost of the worse Top-1 accuracy. For example, the inference MAC of ResNet18 from OptG is more than  $2\mathrm{x}$  less than that from ACDC,
- Interestingly, STR becomes aggressive in pruning the early convolution layers in MobileNet-v1/v2, while OptG does not expose such behavior to MobileNet-v1/v2 (unlike it did for ResNet18/50). Such characteristics also favor the low inference latency over the model accuracy. Also, STR tends not to prune the last linear layer much as discussed in [28].
- Unlike OptG and STR, ACDC does not prune the early convolution layers much for the tested networks, but prunes somewhat actively for the late convolution and linear layers, which leads to high model accuracies at the cost of higher inference latencies.
- PDP is somewhat between STR and ACDC and modest across all layers in pruning allocation for all the networks, leading to superior accuracy and inference trade-offs. For example, the layers model.13.3 of MobileNet-v1 and features.17 conv.1.0 of MobileNet-v2 have the most difference among algorithms, and PDP is modest in pruning these two layers.
- OptG has very low inference MACs on the earlier layers of ResNet18 and ResNet50 due to its aggressive pruning on these as seen in Fig. 6 (a) and (b), which leads to the extremely low inference latencies as shown in Fig. 5 (a) and (b).
- GraNet tends to prune the earlier layers less but the later layers more in general which explains why GraNet shows the highest inference MACs in Fig. 5.
Network	Method	Sparsity (%)
80	70	60	50
ResNet-18	PDP	69.8	70.8	71.0	71.3
ACDC	69.4	70.3	70.6	70.8
MobileNet-v1	PDP	69.5	71.0	71.6	71.9
OptG	68.1	69.1	69.6	69.7
ACDC	68.5	69.9	70.9	71.4
Table 13: Top-1 accuracy with ImageNet1k: PDP outperforms other schemes with various pruning rates.
Network	Method	Validation dataset	Sparsity (%)
80	70	60	50
Bert	PDP	matched mismatched	83.7	84.0	84.3	84.7
83.4	83.8	84.4	84.5
OptG	matched mismatched	80.3	80.7	81.3	81.2
80.1	80.7	80.5	81.0
Table 14: Accuracies with MNLI benchmark: PDP maintains the similar accuracy lead over other schemes.
Table. 9 shows the pruned weight histograms of MobileNet-v1 from Table 5. We can observe that each algorithm affects the weight distribution in a slightly different way.
- STR prefers to split the distribution more widely than others. For the example of the layer 5.0, STR clearly separated the positive and negative weights with a wide gap centered at the zero, while others sis not, except PDP created a slight dip around the zero to create mild separation.
- PDP tends to spread out the sparsified weight distributions more than others. For the example of the fc layer, the weights from PDP range from -0.5 to 2.0, while those from
others are from -0.5 to 1.5. On the other hand, GraNet tends to keep the weight distributions tight.
We also experimented with varying pruning rates for PDP, OptG and ACDC for MobileNet-v1 and ResNet-18 with ImageNet1k, and Bert with MNLI benchmark under the same configurations as in Section 4. Overall, as shown in Table 13 and Table 14, all tested schemes delivered higher accuracy with lower pruning rate, yet we can observe that PDP keeps its superiority to other schemes over all the tested pruning rates.
G Code References
- Dense https://pytorch.org/vision/stable/index.html
- GradNet https://github.com/VITA-Group/GraNet
- OptG https://github.com/zyxxmu/OptG
- ACDC https://github.com/IST-DASLab/ACDC
- STR https://github.com/RAIVNLab/STR
GMP https://github.com/RAIVNLab/STR
- DNW https://github.com/RAIVNLab/STR
- CS https://github.com/lolemacs/continuous-sparsification
- LNM https://github.com/NM-sparsity/NM-sparsity