Pruning-as-Search: Efficient Neural Architecture Search via Channel Pruning and Structural Reparameterization

1 Introduction

Recently, Deep Neural Networks (DNNs) have achieved great success in various applications. However, their tremendous memory and computation consumption lead to difficulties for the wide deployment of DNNs especially on resourcelimited edge devices. To mitigate this gap, Neural Architecture Search (NAS) and network pruning are proposed to lower DNN memory occupancy and improve inference speed.

Although NAS can automatically explore desired operator type, skip connections, as well as model depth and width, it suffers from significant searching overhead with exponentially growing search space. On the other hand, network pruning is designed to remove redundant weights from pre-defined network architectures. The fundamental challenge that impedes the pruning to serve as an efficient automatic width search method is that the layer-wise pruning ratios are hard to be determined automatically. It usually requires a trial-and-error process with domain expertise.

∗Work is done at Kwai Inc.

To overcome the disadvantages of NAS and pruning, we propose pruning-as-search (PaS), an efficient end-to-end channel pruning method to automatically search the desirable subnetworks with only a similar cost to one typical training process. Specifically, we introduce a depth-wise binary convolutional (DBC) layer as a pruning indicator after each convolutional (CONV) layer, as shown in Figure 1. We incorporate Straight Through Estimator (STE) technique to facilitate the trainability of DBC layers. Starting from a well-trained network, PaS only requires one fine-tuning process to learn layer-wise pruning policy and fine-tune the pruned DNN simultaneously. Moreover, another superiority of PaS is the ability to step out of the magnitude trap, which distinguishes PaS from prior arts that use the magnitude of weights or feature maps to select the pruning location (more details in Section 3). We demonstrate that PaS consistently outperforms prior arts by 0.3% to 2.7% on top-1 accuracy under the same computation constraint on ImageNet with ResNet-50.

Besides, we reveal a new dimension of flexibility in DNN width design inspired by structural reparameterization in [Ding et al., 2021]. The design of DNN width is rather rigid in state-of-the-arts. To enable residual connection, most SOTA backbone networks [He et al., 2016; Tan and Le, 2019] have to set block width identical within a stage, otherwise the dimension of output feature map mismatches those from identity path. However, do we really have to sacrifice the flexibility of layer width within a stage due to the identical width by stage design? In this work, we propose a feasible solution to this problem. By skipping only one convolution layer, identity path can be reparameterized into CONV layer during inference. We propose to search on a super-net with reparameterization design, and deploy compact sub-nets without residual connections, such that each layer can have arbitrary width and be free from the aforementioned constraint. We release a brand new family of VGG-like networks with stack of pure 3 × 3 CONVs as well as lightweight DNNs with depth-wise separable CONVs. Our searched network architecture can achieve 2.1× inference speedup and 0.3% higher accuracy compared with ResNet-50 on ImageNet.

We summarize our main contributions as following:

• We develop a PaS algorithm that directly learn pruning policy via DBC layers. Our method saves searching cost compared to brute search, enables exploration which is beyond the capabilities of traditional pruning, and is easily integratable to complex tasks.

• We are the first to investigate the special role of structural reparameterazation in width optimization. Compared to SOTA backbone DNNs with residual connections, we introduce a new flexibility in channel configuration, enabling arbitrary width by layer instead of by stage.
• We release a new family of backbone networks based on our width search and reparameterization. We also validate the effectiveness of our PaS method on complex tasks including segmentation and generative models.

2 RelatedWork

Neural Architecture Search (NAS). NAS aims to design high-performance and efficient network architecture by leveraging powerful computation resources to reduce the intervention of humans. Prior works [Zoph and Le, 2017] incorporate RL into searching process. An agent such as an RNN is trained during the search process to generate desired candidate network architectures. In general, RL-based and evolution-based NAS methods require long searching time, which can even take hundreds to thousands of GPU days. Gradient-based methods [Liu et al., 2018] relax the discrete architecture representation into a continuous and differentiable form, enabling more efficient architecture search by leveraging gradient descent. However, these methods either require conducting the search process on a proxy task (on the smaller dataset) or need a huge memory cost for a large number of candidate operators.

Network Pruning. DNN pruning removes redundant weights in DNNs, thus effectively reduces both storage and computation cost. According to the sparsity type, network pruning are generally categorized into (1) unstructured pruning [Han et al., 2015] and (2) structured pruning [Li et al., 2019]. Unstructured pruning usually achieves a high sparsity ratio, but it is hard to have a considerable acceleration due to its irregular sparsity. Structured pruning removes entire filters/channels of weights and has the potential to reconstruct the pruned model to a small dense model, enabling significant acceleration and storage reduction. Specifically, network pruning in this work refers to channel pruning [Hu et al., 2016; He et al., 2017; Liu et al., 2017; Guo et al., 2021; Su et al., 2021] as we aim to find compact models with width shrinking.

Structural Reparameterization. As a multi-branch strategy, structural reparameterization is proposed either to improve accuracy of origin architecture [Marnerides et al., 2018], or to deliver plain models for speedup purposes [Ding et al., 2021] by merging branch operator into mainstream operator (CONV in general) in a mathematically equivalent way. In this work, we integrate structural reparameterization to decouple the model architecture at training-time and inference-time. We take advantage of residual connection during training but deploy plain architecture with arbitrary widths.

3 Challenges andMotivations

We discuss the key challenges in width optimization. Specifically, (C1) Tremendous searching cost for NAS, (C2) magnitude trap in pruning, and (C3) rigid width constraint.

3.1 Tremendous Search Cost for NAS

In NAS, the exponentially growing search space is a major problem. Specifically, the RL-based NAS methods [Zoph and Le, 2017] typically need to train each candidate architecture with multiple epochs, incurring large searching cost. For a model with NL layers, each layer has NC candidate width, (e.g. {0, 16, 32, 48, 64, ...}), and each candidate needs to be trained by NE (e.g. epochs) computation force. Then the computation complexity is O((NC)NL · NE). Besides, the differentiable NAS methods [Liu et al., 2018] build a memoryintensive super-net to train multiple architectures simultaneously with memory complexity of O(NCNL) for a super-net of NL layers and NC candidates each layer, leading to limited discrete search space up-bounded by the available memory.

3.2 Magnitude Trap for Pruning

The key challenge of pruning is to decide the per-layer pruning ratio and pruning positions. Following [Han et al., 2016], magnitude-based method is widely employed which prunes channels smaller than a global threshold. It is based on the assumption that filters/channels with smaller magnitudes serve less important for final accuracy. However, this assumption is not necessarily true. As shown in Appendix A.1 , we demonstrate that simply penalizing small magnitude filters/channels leads to non-recoverable pruning which means that small channels do not have a chance to become large enough in contribution to the accuracy due to penalizing. It becomes pure exploitation without exploration, and we refer this as the magnitude trap, as illustrated in Tab. 1 We show the comparison of different types of methods including NAS [Zoph and Le, 2017; Zhong et al., 2018] Search-to-Prune [Liu et al., 2019; He et al., 2018] Uniform pruning [Luo et al., 2017], Handcrafted Pruning [Zhang et al., 2018b] Magnitude plus OneShot pruning [Han et al., 2016], Magnitude plus Iterative Update [Zhang et al., 2018a], Equally Regularized [Liu et al., 2017; Guan et al., 2022].

Till now, there are very few work recognizing this issue. [Liu et al., 2017] and [Guan et al., 2022] apply uniform regularization to all pruning indicators. All channels are forced to be close to zero, then the smallest are pruned. Though enabling policy updating thus exploration, the pruned channels are finally selected by magnitude, which falls into magnitude trap as well. Plus, this process usually destroys the performance of super-net, as shown in Appendix A.1.

3.3 RigidWidth Design

Inspired by ResNet, most state-of-the-art DNNs incorporate the residual connections in their building blocks. It requires that the input dimension matches the output dimension. This constraints the DNN design paradigm to divide DNNs into stages by downsampling positions. Within each stage, feature sizes are identical, and each block shares same width to ensure the residual connection can work without size mismatch issues, as demonstrated in Fig. 2 (a, b, c).

Current compression techniques, including NAS and pruning, still suffer from these constraints, significantly decreasing the design flexibility. In NAS, candidate blocks are still residual connected and should share the same width within a stage. Consequently, the freedom is to search expansion ratio by block, and width by stage, which is still at coarse level, as shown in Fig. 2 (b) and (c). Existing pruning methods struggle to satisfy the aforementioned design constraint as well. (i) [Wu et al., 2019] remove the identity path or replace them with a 1 × 1 CONV to match the channel number, sacrificing the advantages of identity path with considerable performance degradation. (ii) [Guo et al., 2020] incorporates constraints to force all blocks sharing the same output dimension within each stage. As observed in Fig. 3, this method is obviously non-optimal as the pruned channel width in a specific stage are not necessarily identical. (iii) As discussed in [Guan et al., 2022], the last CONV within a stage can be pruned freely by incorporating a reshaping operation. The output of the narrow & dense CONV is inflated with zero channels to match the dimension of identity path. But the input of next block is fixed to be the original width, limiting the computation reduction.

3.4 Motivation

We summarize model search methods by four metrics in Table 1. To overcome the aforementioned challenges, we propose PaS to automatically search the model width. Besides, we propose a new pruning-deployment paradigm based on structural reparameterization, unleashing a new dimension of flexibility in per-layer width design.

4 Pruning as Search Algorithm

We first introduce our method and then discuss how our method can deal with the three challenges.

4.1 Depth-Wise Binary Convolution Layers

To automate channel pruning as width search, we formulate the per-layer pruning policy to be trainable along with regular network training, by creating a differentiable depth-wise binary convolution (DBC) layer for each pruned layer. Specifically, we insert a depth-wise 1 × 1 CONV layer following a CONV layer that is supposed to be pruned, as below,

where ⊙ is the convolution operation. wl ∈ Ro×i×k×k is the weight parameters in l-th CONV layer, with o output channels, i input channels, and kernels of size k × k. al ∈ Rn×o×s×s′ represents the output features of l-th layer (with the DBC layer), with o channels and s×s′ feature size. n denotes batch size. vl ∈ Ro×1×1×1 is the DBC layer weights.

Each element in vl corresponds to an output channel of wl ⊙ al-1. Thus we use vl as the pruning indicator and pruning is performed with reference to the magnitude of vl elements. Then the problem of channel pruning is relaxed as binarization of vl. The binarization operations have zero derivatives, leading to difficulties for backpropagation. We propose to integrate Straight Through Estimator (STE) [Bengio et al., 2013] as shown below to train DBC along with original network parameters.

where bl ∈ {0, 1}o×1×1×1 is the binarized vl. The thres is an adjustable threshold, and in our case it is simply set as 0.5. With the DBC layers and the STE method, we can train the model parameters W = {wl} and the policy parameters V = {vl} simultaneously, and pruning policy is decoupled from weight or feature magnitudes.

STE method is originally proposed to avoid the nondifferentiable problem in quantization task. We highlight that the benefits of integrating STE with DBC layers in pruning task is two-fold. First, we decouple the pruning policy frommodel parameter magnitudes. Second, the information in pruned channels is preserved since zeros in DBC layers block gradient flow. As a result, pruned channels are free to recover and contribute to accuracy as they originally did. While in most aforementioned pruning methods, weights in pruned layers are destroyed as they are forced to get close to zeros. Thus our DBC with STE outperforms other complicated strategies which relax the non-differentiable binary masks into sigmoid-like function, such as [Guo et al., 2020; Guan et al., 2022]. Our comprehensive evaluation demonstrates the effectiveness of the STE method on large-scale datasets with different architectures in various applications.

In order to deploy pruned sparse models, the next step is to convert sparse models to dense models by squeezing/- grouping the DBC layer based on the binary values. Let b = {0}o0×1×1×1 ⊕ {1}o1×1×1×1 , where o0 and o1 denote the number of zeros and ones, respectively, with o0 + o1 = o, and ⊕ refers to channel-wise concatenation. Thus we have

where zero channels are squeezed in the last equality.

4.2 Training Loss Function

With differentiable DBC layers, we can train and prune the model via SGD simultaneously with the loss function,

where Lreg is the regularization term related to the computation complexity or on-chip latency. For simplicity, we take Multiply-Accumulate operations (MACs) as the constraint/- target rather than parameter number to estimate on-device execution cost more precisely. β can weight the loss and stabilize training. Lreg can be simply defined as squared ℓ2 norm between current MACs and target MACs C,

4.3 Simultaneous Pruning and Training for C1

By adopting DBC layers to introduce the pruning indicators and leveraging STE method to enable gradients backpropagation, our end-to-end channel pruning algorithm can train the model parameters and pruning indicators at the same time. The training is just like training a typical unpruned model until convergence. Compared with NAS to train each architecture with multiple epochs or other typical pruning methods to iteratively prune and train the model, our method can save tremendous training efforts.

4.4 DBC Layers as Indicators for C2

As discussed, we decouple pruning policy from parameter or feature magnitudes. With the DBC layers as pruning indicators, the benefits come two-fold. First, the indicators are trained with stochastic gradient descent which naturally enable exploration. Thus we jump out of one-shot magnitude-based decision and the pruning policy can be updated dynamically during training. Second, with STE, we zero out not only channels but also gradients during training. If a channel is determined to be pruned by DBC layer, its corresponding weights are not updated. As a result, the weights and performance of pruned layers are preserved thus are ready to recover anytime for re-evaluation. This distinguish our method from equal-penalty soft-mask ones which destroy origin weights.

4.5 Structural Reparameterization for C3

To deal with the residual constraints, RepVGG [Ding et al., 2021] takes the identity path to skip only one convolution at a time, as a typical reparameterization instance. Thus, the identity path can be merged into the convolution during inference, outperforming same-level ResNet by a considerable margin in terms of inference speed, while preserving benefits of residual connections in gradient flow. Deriving from Eq. (1), we show the reparameterization as follows,

where there is a residual connection besides wl ⊙ al-1. We merge the identity path into CONV weights and obtain the merged weights ml. The DBC layer is added after the merged CONVml ⊙al-1 rather than the original CONVwl ⊙al-1 to enable reparameterization, as shown in Fig. 1.

After training the pruning indicators and reparameterization, we zero out both pruned channel from CONV layer and its corresponding channel from identity path, so that the sparse block can be squeezed into a compact one for inference acceleration. The channel dimension of both the current layer's output channel and the next layer's input channel can be reduced, as shown in Fig. 1, Fig. 2 (d→f) and (e→g).

It is notable that these unique-width-per-layer architecture can not be created and trained directly. We demonstrate the searched architectures in section 5.1.

5 Experiment Results

All experiments are conducted on PyTorch 1.7 using NVIDIA RTX TITAN and GeForce RTX 2080Ti GPUs. To demonstrate the efficiency of prune as search method against brute NAS methods and sophisticated search to prune methods, we directly conducted experiments on large scale datasets or complex tasks including ImageNet for classification, MS COCO dataset for segmentation, Generative Adversarial Networks (GANs) for image generation.

5.1 PaS on ImageNet

ImageNet ILSVRC-2012 contains 1.2 million training images and 50k testing images. Note that all results are based on the standard resolution (224×224) for fair comparison. Following standard data augmentation, we prune from pretrained model with weight decay setting to 3.05×10-5 , momentum as 0.875. Learning rate is rewound to 0.4 for a total batch size of 1024 synchronized on 8 GPUs. We search for 10 epochs, which is enough for per-layer pruning policy convergence as shown in Fig. 4. Then we freeze the policy (parameters in DBC) and anneal learning rate by cosine schedule for 50 epochs to achieve final accuracy. Thus we use a total of 60 epochs to deliver the compact well-trained model.

PaS Networks.

As demonstrated in Tab. 2, our searched networks outperform handcrafted design in terms of accuracy and inference speed by large margin under similar GMACs or speed. We provide 3 searched architectures of heavy model under different computation constraints. PaS-A outperforms RepVGG-A2 from the origin paper [Ding et al., 2021] by 0.9% higher top-1 accuracy with less computation cost (0.8 GMACs), because of the optimized layer widths. PaS-B targets on RepVGG-B0 which is a proportionally narrower version of RepVGG-B1. With 0.2 GMACs smaller computation cost, PaS-B outperforms RepVGG-B0 by 0.7% top-1 accuracy. PaS-C achieves 1% higher top-1 accuracy than RepVGG-A1 with the same computation complexity.

Besides accuracy, our proposed PaS-A, PaS-B and PaS-C achieve faster inference speed (in terms of frame per second).

Moreover, compared with ResNet-50 which has residual connections at inference time, PaS-A achieves 0.3% higher top-1 accuracy and 114% faster inference speed, demonstrating the potential of reparameterization-based backbones.

PaS-A, B, C exhibit high flexibility in layer width as shown in Fig.6. We also compared the searched architectures of PaSB and RepVGG-B0 which share the same depth in Fig. 3. We can observe that the first few layers in PaS-B are narrower than those in RepVGG-B0, while the last ones are significantly wider. In classification task, doubling channel numbers when downsampling proves to be a good practice, as seen in the searched PaS-B. However, we should also gradually increase channel numbers within each stage, which is usually not supported in previous work due to residual constraints.

We also experiment over lightweight networks [Sandler et al., 2018; Tan and Le, 2019; Tan et al., 2019], which demonstrate significantly better parameter and computation efficiency and have become de facto choice on edge devices. We perform search on lightweight building blocks as shown in Tab. 2 and release a series of models in Appendix Fig.7, composed with 3 × 3 depth-wise convolutions and 1 × 1 regular convolutions. Our searched architecture exhibits high flexibility in layer width. Experiment results demonstrate that our searched network achieves best accuracy-complexity trade-off. Since we are investigating width optimization, in this work we do not employ sophisticated network elements including SE-block [Hu et al., 2018], or activation functions [Ramachandran et al., 2018]. Thus we only compare accuracy performance with plain-built baselines.

Prune Traditional Models.

Despite the newly proposed networks with arbitrary width, we also show the effectiveness of our PaS algorithm on common models. We validate our PaS on ResNet-50 [He et al., 2016]. As shown in Tab. 3, PaS consistently outperforms prior arts by 0.3% to 2.7% top-1 accuracy under same computation constraint. Moreover, PaS just need to train once to deliver final accuracy, saving search cost (in GPU-hours) compared with other baselines with sophisticated searching process.

By comparing the pruned architecture of PaS and other baselines, we notice an interesting phenomenon that the first layer is pruned 50% to 80% with our method, which is significantly different from prior arts that claim it to be the most sensitive. On the contrary, the last layers are pruned less. We conclude this phenomenon as the last several layers extract high-level features and contribute more per-MACs-information.

5.2 PaS on Instance Segmentation and GAN

We evaluate the proposed method on instance segmentation with MS COCO dataset . We select the GPU real-time YoLACT [Bolya et al., 2019] model as supernet. As shown in Tab. 4, the proposed PaS method can achieve 1.7× total compression rate without degradation in mask mAP. Compared with other pruning methods such as uniform pruning or magnitude pruning, PaS can achieve better mask mAP and box mAP by a large margin under the same computation constraint. We also show the performance of PaS on unpaired image style transfer [Zhu et al., 2017]. Specifically, we demonstrate horse2zebra dataset in Tab. 4. Though GAN suffers from unstable training, our PaS still achieves satisfactory compression results. Due to space limit, more detailed experiment settings and results on segmentation and GANs can be found in Appendix. Our PaS exhibits better generalization on complex tasks and datasets compared to prior arts.

5.3 Convergence of Pruning Policy

We demonstrate the robust convergence of PaS in Fig. 4. The pruning policy converges within a few epochs (less than 10). During PaS, we observe that pruned channels have the chance to recover and be re-evaluated, which proves our ability of exploration and exploitation. The policy converges to a stable point which proves the robustness of our PaS.

6 Conclusion

We propose PaS to optimize sub-network width with regular training, achieving high efficiency and robustness. Combining PaS with reparameterization, we generate a brand new series of reparameterization-based networks, which benefit from residual connections during training while enabling arbitrary width for each layer during inference, squeezing out an unrevealed dimension of width optimization.

Acknowledgments

The research reported here was funded in whole or in part by the Army Research Office/Army Research Laboratory via grant W911-NF-20-1-0167 to Northeastern University. Any errors and opinions are not those of the Army Research Office or Department of Defense and are attributable solely to the author(s). This research is also partially supported by National Science Foundation CCF-1937500 and CCF-1901378.