Enhanced Sparsification via Stimulative Training

Existing sparsification-based pruning methods can be divided into two categories: static sparsification and dynamic sparsification. Static sparsification aims at imposing a global and unchanged penalty strength to all parameters of the model [43, 3, 38, 40, 60, 4, 19, 45, 14, 62]. The penalty terms in these works are designed to be a norm-based regularization item w.r.t. model parameters. Different from static sparsification methods, dynamic sparsification considers the individual and time-varying sparsification strength of different parameters [59, 58, 13, 59, 67, 65, 72, 9, 4]. AFP [13] utilizes auto-balanced regularization for different weights to transfer the representational capacity from the whole network to the kept part. Greg [58] merely applies sparsification on the dropped parameters and grows the penalty term large gradually to achieve higher sparsity. Previous works mostly follow the suppressed sparsification paradigm, which imposes constraints on dropped parameters, causing expressivity damage before pruning. Differently, our method achieves relative sparsity by enhancing kept parameters and maintaining the magnitude of dropped ones via stimulative training.

In general, pruning a network inevitably deteriorates its performance. To compensate for the performance degradation, conducting knowledge distillation (KD) after pruing has been proved feasible [1, 2, 63, 47, 48, 7, 37]. In [2], the authors firstly obtain a pruned network via APoZ [29] and then transfer knowledge from the unpruned network to the pruned one via a cosine-similarity-based loss. The authors in [47] apply a Kullback-Leibler-divergence-based loss for KD and verify its effectiveness in both one-shot and iterative magnitude pruning. Built upon iterative magnitude pruning, the authors in [48] propose a cross-correlation-based KD loss to better align the pruned network with the unpruned one. The aforementioned methods treat KD as a performance booster isolated from pruning. Differently, we treat KD itself as a pruning method, not merely a support technique. Our proposed method incorporates KD into pruning, utilizing self-distillation for sparsification and KD loss to obtain the pruning mask.

Stimulative training (ST) [68] aims to boost the main network by transferring the knowledge from the main network to each depth-wise subnet. In ST, the main network can be seen as an ensemble of shallow subnets, and the final performance is determined by each subnet under the unraveled view [57, 53]. Thus, while training the main network, ST randomly skips some layers to generate subnets and utilizes knowledge distillation to supervise each subnet. The experiments in ST show that ST can obtain performance improvement on the main network while bringing various subnets with marginal performance reduction compared with the main network. This additional benefit of subnets suggests that ST can achieve simultaneous enhancement at both local and global of the main network, which inspires us to study the distribution of parameters in ST and exploit ST as an enhanced sparsification paradigm for pruning.

Pruning aims to eliminate parameters from a given network $\mathcal{N}$ with parameters $\theta$ to obtain a high-performance compact subnet $\mathcal{N}_{s}$ with retaining parameters $\theta_{s}$ with specific resource constraint. Although $\mathcal{N}$ can be given with pretrained weights, our method concentrates on pruning from randomly initialized parameters, same as [4, 27]. The pruning process can be abstracted as a function $F(\cdot)$, which takes $\theta$ as input and output $\theta_{s}$, $\theta_{s}=F(\theta)$²²2Because the dependence of dataset is upon the pruning algorithm, for simplicity, we omit the given dataset $D$ as input.. General pruning mainly consists of two categories of transformations 1) $\mathcal{M}=F_{1}(\theta)$, which generates a binary mask to determine remaining and dropping parts 2) $\hat{\theta}=F_{2}(\theta)$, which adjusts the parameter distribution for eliminating with low damage. A typical pruning process can be denoted as:

The exact format depends on the specific pruning algorithm. For one-shot pruning, $F_{1}$ is applied once [36] while alternately for iterative pruning [17]. We denote the dropped parameters as $\theta_{d}=\theta\setminus\theta_{s}$. The specification process belongs to $F_{2}(\cdot)$ and a typical method [59, 58, 4] introduces $L_{2}$ norm penalty term for $\theta_{d}$, which optimizes the following objective:

where $\mathcal{L}_{task}$ is the task loss; $D$ is the given dataset;$\lambda$ is the coefficient of penalty term. The penalty term drives $||\theta_{d}||$ to converge to zero and eliminates the importance of $\theta_{d}$ for painless removal, which belongs to suppressed paradigm. Whereas $\theta_{d}$ is a part of $\theta$, suppression of $\theta_{d}$ reduces the capacity of $\theta$, which limits the upper bound of $\theta_{s}$ and compromises the final performance.

To boost the performance of the main network, stimulative training (ST) views a network as an ensemble of multiple shallow subnets and transfers the knowledge from the main network to each weight-shared subnet. Formally, ST optimizes the following loss:

where $x$ and $y$ are the input images and ground truth labels, respectively; $CE(\cdot)$ is the cross entropy; $KL(\cdot)$ is the Kullback-Leibler divergence; $\lambda$ is the balance coefficient; $p(x;\theta)$ is the output probability with $x$ and $\theta$; $\theta_{m}$ and $\theta_{s}$ are the parameters of the main network and subnet, respectively, and $\theta_{s}$ is the subset of $\theta_{m}$. According to the optimization of the second term of $\mathcal{L}_{ST}$, i.e., $KL(p(x;\theta_{m})||p(x;\theta_{s}))$, ST can gradually transfer the representative ability from the whole parameters ($\theta_{m}$) to the partial ones ($\theta_{s}$), causing the enhancement of subnet weights and implying the potential of relative sparsity.

To investigate the relative sparsity effect in ST, we conduct ST by fixing a specific subnet on ResNet-50 and visually show the magnitude of a convolutional layer’s weights. Notely, while the original ST only samples subnet in the depth dimension, we extend it into the width dimension. As shown in Fig. 8, compared with the random weights’ distribution in the baseline, i.e., standard training with stochastic gradient descent, ST can significantly enhance the chosen weights, causing the weights’ magnitude relatively concentrating on the chosen weights (the black rectangle). It is worth noting that Fig. 8 reports the relative comparison, see Fig. 11 for detailed absolute values. To further display the effect of different sparsification methods before and after pruning, we conduct experiments with different target sparsity. As shown in Fig. 2, the $L$ norm sparsifications harm the performance before pruning (before removing the dropped parameters), while ST maintains and even improves the performance. It suggests ST introduces a moderate and relative sparsity effect via KD to keep the capacity of the dropped parts, which encourages learning better sparsity distribution.

The proposed ST guided pruning (STP) framework is illustrated in Fig.9. In detail, given a dense network, i.e., the main network and a target FLOPs constraint, we initialize an architecture pool containing structured subnet architectures sampled from multi-dimension FLOPs-aware sampling space. The training starts with randomly initialized parameters. In each training iteration, the procedure can be divided into five steps: 1) sample a subnet architecture based on the score in the architecture pool; 2) forward the main network to generate the main output, compute and backward the cross entropy between ground truth and main output; 3) apply the sampled subnet architecture mask into the main network to construct a weight-shared subnet and supervise it using the main output by knowledge distillation. 4) mutate the sampled architecture mask multi-dimensionally into a larger architecture mask, which is utilized to generate a weight-shared mutating network and supervise it using the main output; 5) remove the subnet architecture with the highest score (KD loss) in the architecture pool every $k$ iterations. After $T$ training iterations, only one subnet architecture mask exists in the architecture pool. We apply it to the main network to obtain the final pruned network while the other parameters are deserted. A typical value of $k$ is the iteration number of an epoch. The pseudo-code is shown in the Appendix. Three critical designs of STP are introduced as follows.

Besides the multi-dimension design, we impose a FLOPs-aware constraint on the sampling space, only sampling the subnets with target FLOPs. It can guarantee a more accurate practical speedup ratio compared with sparsity constraint. Because the sampled subnets generally possess much lower capacity than the main network, we adopt the normalized Kullback-Leibler divergence (KL-) [66] to eliminate the capacity gap for higher performance.
Formally, given a structured mask $\mathcal{M}_{s}$, the parameters of subnets are $\theta_{s}=\mathcal{M}_{s}\odot\theta_{m}$. ST sparsification is optimized by:

$\hat{p}(x;\theta)$ is the normalized probability which is donated as:

where $Z(x;\theta)$ is the output logits. $L_{STS}$ is essentially a knowledge distillation loss, and illustrated as $KD_{sub}$ in Fig.9. The subnets will also be called target subnets in the following section.
Subnet Mutating Expansion. To mitigate the performance degradation resulting from focusing on training tiny subnets [55], we introduce additional larger subnets to provide support during training. As a distinction, we refer to them as support subnets. Considering that randomly sampling larger support subnets will disrupt the parameter magnitude concentration of target subnets, we propose to mutate and expand the width and depth of target subnets to generate support subnets.

Formally, given a target subnet with layer-wise sparsity ratio $S_{tar}=\{S_{tar}^{1},S_{tar}^{2}\\ ,...,S_{tar}^{L}\}$, for each $S_{tar}^{i}$ we uniformly sample a larger ratio $S_{sup}^{i}$ which satisfies:

where $U(\cdot)$ denotes the uniform distribution. The support subnet will be supervised by the main network similar to Formulation 4, denoted as:

where $\theta_{sup}$ is the parameters of the support subnet and $\mathcal{L}_{SME}$ is illustrated as $KD_{mut}$ in Fig. 9. Formulation 6 can guarantee that the support subnet contains the whole parameters of the sampled target subnet. Thus, training the support subnet also enhances the target subnet.

Specifically, suppose there is architecture pool with $N$ candidate masks, denoted as $\mathcal{P}=\{\mathcal{M}_{1},\mathcal{M}_{2},...,\mathcal{M}_{N}\}$. Each architecture has a corresponding score, denoted as $Sco=\{Sco_{1},Sco_{2},...,Sco_{N}\}$. While training a target subnet with $\mathcal{M}_{i}$, we record the $\mathcal{L}_{STS}$ and update the corresponding score via exponential moving average (EMA) [21], which can be denoted as:

where $\alpha\in(0,1)$ is the updating ratio. The pool is shrunk by removing the architectures with the highest scores every $k$ iterations. After $T$ rounds of shrinking, the pool contains a single architecture, which is chosen as the pruned network.

To further demonstrate the generality of STP, we apply STP on two typical Transformers, i.e., ViT [56] and Swin Transformer [41]. Similar to CNNs, we conduct experiments on CIFAR-100 w.r.t. three FLOPs targets, including 15%, 35%, and 55%. The results are shown in Table 3. It can be observed that STP is superior for both Transformers under all target FLOPs compared with RST-S. Moreover, under 55% FLOPs, the STP pruned Transformers can almost realize lossless compression, suffering merely 0.25 and 0.02 performance drop for ViT and Swin Transformer, respectively. Note that the proposed STP is not tailored for Transformers, however, it can still achieve competitive results under 55% FLOPs, implying its immense potential in pruning Transformers.

We eliminate the proposed components of STP, i.e., ST relative sparsification, subnet mutating expansion, KD guided architecture exploration and multi-dimension sampling, to validate the single and joint effect of them. The results are shown in Table 5. Firstly, we conduct experiments with the separate removal of each component except the ST relative sparsification and observe different degrees of performance degradation, which demonstrates their indispensability. Then, we remove these components step-by-step and observe consistent performance degradation, which demonstrates their joint effectiveness. Note that, to remove the ST relative sparsification, we replace it with the $L_{2}$ regularization, which causes significant performance damage (-2.7% Top-1 accuracy), demonstrating the superiority of the ST relative sparsification.

To validate the practical capacity of STP, we conduct experiments on downstream tasks, including object detection and instance segmentation. Specifically, we apply STP on ResNet-50 in ImageNet to generate pretrained pruned networks with three remaining FLOPs, i.e., 15%, 35%, 55%. These pretrained pruned networks are utilized as the backbone in Mask R-CNN [22], which is a multi-task framework. Then, the whole Mask R-CNN is fine-tuned and trained on COCO2017 [39] under the supervision of object detection and instance segmentation. For the fine-tuning settings, we follow the standard 1x training schedule [6] and adopt AdamW [44] to train the Mask R-CNN for 12 epochs with weight decay 0.1. The batch size is 16 and the initial learning rate is 0.0001. The learning rate decays 10% at the 8-th and the 12-th epoch. We report the mean average precision (mAP) under different intersection over union (IoU) of the bounding boxe and mask as the evaluation metrics of object detection and instance segmentation, respectively.

The results are shown in Table 6. It can be observed that STP consistently surpasses GReg [58] on both detection and segmentation. Under relatively low FLOPs (35%), STP causes marginal performance reduction (-0.1% on Det [email protected] and Seg [email protected]) on downstream tasks. In the case of extremely low FLOPs(15%), STP can maintain the basic performance and does not cause unbearable performance degradation. It is worth noting that STP can improve performance (+0.3% on Det [email protected] and +0.9% on Seg [email protected]) with 1.8x speed up (55% remaining FLOPs). We consider the reason is that the ST guided training framework can drive the pruned network to extract more representative features for downstream tasks. The above observations demonstrate that STP can be simply and practically deployed in different tasks and obtain competitive performance with various target FLOPs, which is friendly for resource-limited edge devices. Besides, the performance improvement suggests the potential of STP to improve pareto frontier between FLOPs and performance.

To verify that our architecture exploration strategy can gradually explore the subnets with outstanding performance, we record the pool in different training epochs. We evaluate the effectiveness from two aspects: 1) performance: the average performance of architectures in the pool 2) clustering: the average sum of squared error (SSE) distance between the clustering center and architectures in the pool, where the higher SSE means lower clustering. Notely, for estimating the clustering of architectures, we convert the normalized sparsity ratio of each layer and depth into a vector and regard it as a proxy of architectures. Due to the unbearable computational overhead of fully training and measuring each subnet in the initialized pool(containing more than 1000 subnets), we conduct another experiment that utilizes the original ST while uniformly sampling the whole pool. After ST, we evaluate all subnets in the initialized pool on the test dataset as the proxy (or a look table) of the architectures’ performance.

The results are shown in Fig. 11. At the early stage of STP, the randomly initialized pool contains various architectures with low average performance and markedly different structures. As training progresses, the average performance is gradually raised, and at the same time, the similarity of structures also increases. It demonstrates that under the constraint of FLOPs, our strategy can progressively drive the architecture exploration to converge to an optimal region and finally obtain a high-performance architecture for sparsification.

To verify the relative sparsification effect of STP, we record the average magnitude of the chosen (kept), unchosen (dropped), and baseline (standard training without sparsification) parameters in different convolutional layers of ResNet-50. As illustrated in Fig. 11, the magnitude of chosen parameters is remarkably and consistently higher than the unchosen part in the front layers, while the unchosen part has a similar magnitude as the baseline. It demonstrates that STP achieves the evident relative sparsity effect like the original ST without suppressing the dropped part. Due to the subnet mutating expansion, the unchosen part has small parability of being enhanced, thus, the magnitude of unchosen part is marginally higher than the baseline. There is another phenomenon that in the later layer of stages, the relative sparsification effect becomes less pronounced. That is because during training, for low FLOPs targets, e.g., 15%, the later layers are dropped with a greater probability and sparsified less. Combining with the empirical results in Table 1, we conclude that STP can obtain a sparsity parameter distribution for pruning with less performance degradation.

The CIFAR-100 [33] is a typical classification dataset with 100 categories, consisting of 50,000 training images and 10,000 testing images. For ResNet-50 [23] and MBV3 [28], we adopt the training settings of [68]. Specifically, the epoch number and batch size are 500 and 64, respectively. The SGD is chosen as the optimizer with a 0.05 initial learning rate and a 0.0003 weight decay. We use the cosine decay schedule to adjust the learning rate over the training process. For WRN28-10 [73], we adopt the training settings of [73]. Specifically, the epoch number and batch size are 200 and 128, respectively. The SGD is chosen as the optimizer with a 0.1 initial learning rate and a 0.0005 weight decay. The learning rate scheduler is also the cosine decay schedule. For ViT [56] and Swin Transformer [41], we use an image size of 32x32 and a patch size of 4. The epoch number and batch size are 200 and 128, respectively. The optimizer is AdamW [44] with an initial learning rate of 0.001/0.003 for Swin/ViT and a 0.05 weight decay. The learning rate is warmed up for 10 epochs. The data augmentations are the same as the ones in [35]. Different from CNNs, where we regard the channel numbers of convolutional and linear layers as the width dimension, to prune the width of Transformers, we take the head numbers of attention layers and the channel numbers of linear layers into account.

The Tiny Imagenet dataset is inherited from the ImageNet dataset [12], containing 200 categories, 100,000 training images, and 10,000 testing images. For ResNet-50 [23] and MBV3 [28], the epoch number and batch size are 500 and 64, respectively. The optimizer is SGD with a 0.1 initial learning rate and a 0.0003 weight decay. We utilize a step-wise learning rate scheduler, downscaling the learning rate to 0.1 and 0.01 of the original one at the 250-th and 375-th epoch, respectively. For WRN28-10 [73], we adopt the training settings in [51]. The epoch number and batch size are 200 and 128, respectively. The optimizer is SGD with a 0.2 initial learning rate and a 0.0001 weight decay. We utilize a step-wise learning rate scheduler, downscaling the learning rate to 0.1 and 0.01 of the original one at the 100-th and 150-th epoch, respectively.

The ImageNet dataset [12] is a widely used classification benchmark, containing 1000 categories, 1.2 million training images, and 50000 testing images. For the evaluated ResNet-50 [23], the epoch number and batch size are 200 and 512, respectively. We utilize SGD as the optimizer. The learning rate is initialized as 0.2 and is controlled by a cosine decay schedule. The weight decay is 0.0001. Besides, we apply the commonly used data augmentations according to [30, 54].

To build a relatively more general and user-friendly pruning framework, we resort to the symbolic tracing of Pytorch, called FX tracing [50]. Given any FX-compatible model, we first extract the FX graph from it. Based on the graph, we try to extract dependencies groups, which are similar to the concepts in [16, 10]. A dependency group represents a series of layers whose outputs are expected to be integrated, such as being added or concatenated in the dimension to prune. When pruned, all layers in the same group are expected to have the same number of remaining channels.

Given the inputs to the model, we run an interpreter [50] based on the extracted FX graph instead of using the traditional model forward. The interpreter propagates the inputs from the root node of the graph to the last node, executing all encountered nodes according to their recorded operations, promising the same result as the one produced by the traditional model forward while providing the freedom to intercept any intermediate operation for injecting necessary functionality of pruning. Benefiting from the interpreter, we can dynamically reduce the layer number or shrink channels based on the sampled architecture without altering the model structure or parameters. To realize the dynamic interpreter, we dispatch all commonly used operators, such as trivial binary operators (add, mul, div, etc.), Pytorch tensor operators, and “torch.nn" forward functions, in the FX graph to a pre-registered dynamic handler, which can be easily extended to support customized operations.

Besides dynamic forward, we closely integrate FLOPs/parameters estimation with the above pruning framework via running an interpreter with a proxy tensor. A proxy tensor does not require dense calculation, such as convolution or matrix multiplication, resulting in significant acceleration of the FLOPs/parameters estimation for dynamic models, which is the key to estimating the FLOPs of all architectures within the pool in a bearable time. During an interpreter run, a proxy tensor only contains the shape information of a tensor, and each encountered operation is required to infer the output shape of the tensor in addition to the FLOPs and parameters introduced by this operation. Similar to the dispatch of dynamic forward, we have registered the handlers of FLOPs/parameters estimation for all commonly used operators in the FX graph.