Only Train Once: A One-Shot Neural Network Training And Pruning Framework

1 Introduction . Deep neural networks ( DNNs ) have been shown to be effective in various real applications ( 51 ; 28 ) . It is widely acknowledged that large-scale DNN models not only learn faster but also outperform their slimmer counterparts . However , such heavy models pose a great challenge to the deployment stage due to their resource-consuming nature . In addressing this issue , many model compression ∗Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . techniques ( 5 ; 11 ) are proposed in the past decade that aim at compressing those large and complex models into slimmer and simpler ones while suffering negligible loss in performance . Pruning methods as one of the main categories of model compression , focus on identifying and pruning redundant structures via various mechanisms to achieve a slimmer architecture , and thus improve the interpretability of a DNN model ( 26 ; 11 ; 65 ) . For example , ( 32 ; 33 ) adopt fine-grained pruning via ` 1 or ` 2 regularization , which prune the small-weight connections based on some hard threshold . ( 36 ; 57 ; 60 ) measure the importance of filters to accelerate the networks by removing insignificant feature maps . ( 37 ; 7 ) utilize reinforcement learning agent to predict compression action . Nevertheless , many of the existing pruning methods ( i ) often rely on criteria based on heuristics or empirical cues , e.g. , magnitude of a connection weight and importance score of a filter , to identify redundant parameters , which may cause divergence during optimization ; ( ii ) thus require complex multi-stage pipelines that involve either a retraining or fine-tuning procedure to regain the accuracy during constructing a slimmer model , which is time-consuming ; and ( iii ) are specific to certain architectures or applications , and are consequently less applicable to various downstream scenarios . Recently , there have been a few efforts ( 14 ; 58 ; 8 ) to directly train the network with sparsity inducing regularizers , which provide generality and convergence guarantee . However , these approaches focus on either merely the individual sparsity of the parameters or the group sparsity of the filters , and thus can not directly remove those zero components ( still require subsequent fine-tuning ) since the zeros are entangled with other commonly used components , e.g. , bias , batch normalization or skip connection . Furthermore , the optimization algorithms used in ( 14 ; 58 ) lack sufficient capability to explore ( group ) sparsity in DNNs effectively and require a post-processing step to yield exact zeros . In this paper , we overcome the above limitations of existing pruning methods by proposing a one-shot neural network pruning framework , with which we are able to train a full heavy model from scratch only once , and obtain a slim architecture without fine-tuning while maintain high performance . As shown in Figure 1 , the key to its success is twofold : ( i ) we identify and partition the parameters of DNNs into zero-invariant groups ( ZIGs ) , enabling us to prune redundant structures according to zero groups without affecting the output of the network ; and ( ii ) to promote zero groups , we formulate the pruning task as a structured-sparsity optimization problem and propose a novel optimization method , Half-Space Stochastic Projected Gradient ( HSPG ) , to solve it , which outperforms the standard proximal methods on sparsity exploration and maintains comparable convergence . We highlight that both zero-invariant group partition and the novel optimization algorithm in promoting zero group lead to achieve one-shot neural network training and pruning regardless of its architecture . Our main contributions are summarized as follows . • One-Shot Training and Pruning . We propose OTO , a training and pruning framework that compresses a full neural network with competitive performance by Only-Train-Once , thereby one-shot . OTO dramatically simplifies the complex multi-stage training pipelines of the existing pruning approaches , fits various architectures and applications , and hence is generic and efficient . • Zero-Invariant Group . We define zero-invariant groups for neural networks . If a network is partitioned into ZIGs , it allows us to prune the zero groups without affecting the output , which results in one-shot pruning . Such property is applicable to various popular structures from plain fully connected layers to sophisticated ones such as residual blocks and multi-head attention . • Novel Structured-Sparsity Optimization Algorithm . We propose Half-Space Stochastic Projected Gradient ( HSPG ) , a method that solves structured-sparsity inducing regularization problem . We show and analyze the superiority of HSPG in promoting zero groups of networks than the standard proximal methods and the competitive objective convergence in practice . The fact that ZIG and HSPG are designed agnostic to networks makes OTO generic to various applications . • Experimental Results . We train and compress full models simultaneously from scratch without fine-tuning for inference speedup and parameter reduction , and achieve state-of-the-art results on compression benchmark VGG for CIFAR10 , ResNet50 for CIFAR10/ImageNet , Bert for SQuAD . 2 Related Work . Structured pruning focuses on identifying and pruning the redundant structures in a full model to achieve slimmer architectures for efficient model inference and storage ( 26 ; 32 ) , where there have been numerous efforts dedicated . For CNN compression , the general procedure can be largely summarized as : ( i ) train a full model ; ( ii ) identify and prune the redundant structures to build a slimmer model based on various criteria , including ( structured ) sparsity ( 58 ; 85 ; 14 ; 56 ; 102 ; 27 ; 102 ; 62 ; 91 ) , Bayesian pruning ( 101 ; 65 ; 59 ; 81 ) , ranking importance ( 54 ; 60 ; 41 ; 36 ; 57 ; 100 ) , reinforcement learning ( 37 ; 7 ) , adversarial robustness ( 76 ) , scientific control ( 79 ) , lottery ticket ( 23 ; 24 ; 72 ) , joint quantization learning ( 80 ; 90 ) , etc . ; ( iii ) retrain or iteratively fine-tune the slimmer model to regain the accuracy regression during pruning . These methods can not avoid the extra and usually timeconsuming fine-tuning step because the identified redundant structures , even parametrized with zeros , actually contribute to the model output , thereby additional fine-tuning step is an absolute necessity . For pruning Bert ( 82 ) , knowledge distillation ( 40 ) and LayerDropout ( 21 ) shorten Bert by reducing the number of layers directly . Other methods ( 29 ; 75 ; 30 ) build slimmer Berts in the manner of individual sparsity , but require specially designed data structure for storage and computing library to take advantage of sparse data ( 31 ; 10 ) , and typically can not achieve inference speedup against the highly optimized library ( 16 ) for dense model due to the discontiguous memory allocation ( 9 ) . The structured sparsity for weight pruning is the most relevant to the algorithm described in this paper . The existing structure learning works ( 58 ; 85 ; 14 ; 56 ; 102 ) have the respective disadvantages : ( i ) multiple trainings during the whole procedure since their group partition can not isolate the impact of pruned structures to the model output ; and ( ii ) heuristic post-processing to generate zero groups as the standard proximal methods ( 19 ; 87 ; 88 ; 12 ) and ADMM ( 100 ; 58 ; 4 ) defective on the sparsity exploration for deep learning ( 8 ) , which may deteriorate the performance of the model significantly . Avoiding fine-tuning step during the whole pruning procedure is receiving more and more attentions because of its efficiency . In particular , SNIP ( 52 ) and GraSP ( 83 ) identify redundancy via salience scores at the initialization stage to construct pruned structures , then train the pruned models by the standard optimizers . SCP ( 48 ) isolates the impact of batch normalization , while lacks the consideration of more general DNN architectures . 3 OTO . In essence , OTO frames the network training and pruning as a structure learning problem . Given a full modelM , OTO trains and compresses it simultaneously from scratch without fine-tuning , and achieves significant reduction in both FLOPs and parameters . Particularly , as stated in Algorithm 1 , the trainable parameters ofM are firstly partitioned into a ZIG set G ( Section 3.1 ) . We then construct and solve a structured-sparsity inducing optimization problem ( Section 3.2 ) by proposed stochastic optimizer ( HSPG ) to seek a highly group-sparse solution x∗HSPG ( Section 3.3 ) . Lastly , we obtain a compressed modelM∗ by directly pruning these zero groups ( Section 3.4 ) . Algorithm 1 Outline of OTO . 1 : Input : Full modelM ( no need to be pretrained ) . 2 : Construct G : Partition the trainable parameters ofM into a ZIG set G. 3 : Train : Train the modelM using HSPG ( Algorithm . 2 ) to obtain a group-sparse solution x∗HSPG . 4 : Prune : Construct a slimmer model architectureM∗ by directly pruning zero groups of x∗HSPG . 5 : Output : Compressed modelM∗ . 3.1 Zero-Invariant Group . The root cause of the existing methods having multi-stage training pipeline is that despite the pruned structure ( e.g. , 3D filter ) being zeros , its associated structures ( e.g. , non-zero bias ) still contribute to its corresponding output to the next layer ( e.g. , feature map ) . As a result , the model accuracy regresses , hence an extra step of fine-tuning is necessary . OTO avoids the necessity by partitioning the parameters of DNNs into a set of so-called zero-invariant groups ( ZIGs ) G defined as follows . Definition 1 ( Zero-Invariant Groups ( ZIGs ) ) . For a layer-wise DNN , we partition its entire trainable parameters into disjoint groups G = { g } . Then we call G zero-invariant groups ( ZIGs ) if each group g ∈ G is zero-invariant in the sense that all of the parameters in g being zeros results in its corresponding output to the next layer to be zeros as well . In effect , if and only if a DNN model is partitioned into a ZIG set G and one or more of its element g are parameterized by zeros , the entire corresponding structures contribute none to the model outputs and hence can be pruned directly . Such partition is applicable to various structures of DNN models . Without loss of generality , we define and describe ZIG partition for three most popular structures : ( i ) Conv-BN , ( ii ) Residual Block , and ( iii ) Fully Connected and Multi-Head Attention Layer . ZIG of Conv-BN . Convolutional layer ( Conv ) followed by batch-normalization layer ( BN ) is extensively used in DNN models . Figure 2a shows the ZIG partition for Conv-BN . The 4D filter tensor Kl is flattened into a filter matrix K̂l . During the forward pass , the input tensor I l is transformed into the output tensor Ol of Conv and then into the input tensor of the ( l + 1 ) th layer I l+1 by Ol ← I l ⊗ K̂ l + bl , I l+1 ← a ( O l ) − µl σl γl + βl , ( 1 ) where denoted by ⊗ the convolutional operation , the element-wise multiplication and a ( · ) the activation function . BN is parameterized by mean µl , standard deviation σl , weight γl and bias βl respectively . The activation function needs to be zero-invariant , i.e. , a ( 0 ) = 0 , where most instances satisfy , e.g. , ReLU ( 25 ) , PReLU ( 34 ) , GELU ( 39 ) and LeakyReLU ( 89 ) . Hence , each row of the flattened filter matrix K̂l and its bias bl belong to one ZIG because they being zeros results in their corresponding channel of Ol ( i.e. , feature map ) to be zeros as well . Subsequently , γl and βl of this corresponding channel in BN are also included into this ZIG to avoid the value shift ( zero to non-zero ) during normalization . Note that grouping these four sets of parameters channel-wisely makes Conv-BN zero-invariant regardless of the value of µl and σl , and hence they are excluded from the ZIG . For illustration , each ZIG is highlighted in the same color ( e.g. , gl1 in blue ) . ZIG of Residual Block . The residual block adds another layer of challenge because its output tensor is the summation of the outputs of two Conv-BNs . Figure 2b shows the ZIG partition for the residual block . As illustrated , before propagated to Conv3 , the outputs of Conv1-BN1 and Conv2-BN2 are summarized and hence share the same dimension . As such , to make residual block zero-invariant , we group the four sets of parameters channel-wisely of both Conv1-BN1 and Conv2-BN2 into ZIGs , i.e. , each row of K̂1 , b1 , γ1 , β1 of Conv1-BN1 and each row of K̂2 , b2 , γ2 , β2 of Conv2-BN2 . In Appendix A.1 , we describe the zero-invariant group partition of ResNet50 in greater detail . ZIG of Fully Connected and Multi-Head Attention Layer . Figure 2c shows the ZIG partition for fully connected and multi-head attention layer . Particularly , we partition each row of weight matrix and its associated bias into a ZIG , and therefore any input element is turned to zero if that ZIG is parameterized with zeros , making the fully connected layer zero-invariant . Multi-head attention layer is the key building block of the transformer architectures ( 82 ) . Its trainable parameters contain a weight matrix and bias vector , consisting of the sub-matrix and sub-vector of each head ( we use two heads as an example ) . We form ZIG by grouping each row of every sub-matrix and sub-vector , i.e. , each row of Wh1 , bh1 , Wh2 and bh2 of h1 and h2 , respectively . Automatic ZIG Partition . Based on the above illustrating examples , we provide prescribed ZIG partition for the tested DNNs in Section 4 . Furthermore , given an arbitrary DNN architecture , the procedure of partitioning variables into ZIGs could be automatically proceeded , wherein the key would be identifying the connections among various layers , then performing corresponding group partition . We will leave the automatic ZIG partition for arbitrary DNNs as future work .

Summary: This paper proposes a group sparsity optimisation approach for training pruned models. This is done by splitting the optimisation into an initial phase where the non-differentiable regularisation objective is optimised using subgradients, followed by a second stage where channels/units (plus their corresponding bias/BN/residual path) are mapped to exact zero by augmenting the sub-gradients with an indicator function.  For each parameter group, this function compares the direction of subgradient with the direction you'd get as a result of pruning that group, and if the angle between the two is larger than 90 it prunes them away.