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 structured pruning method called Only-Train-Once (OTO). OTO first partitions the trainable parameters into zero-invariant groups called ZIGs. This way allows one to prune without affecting the output and therefore it requires no additional training steps after pruning. Once they define ZIGs, they perform Half-space Stochastic Projected Gradient (HSPG), starting from a pre-trained estimate using SGD, it performs a form of projected gradient descent while gradually increasing group sparsity. The authors demonstrate the effectiveness of OTO on CIFAR, ImageNet, SQuAD by showing an impressive gain in terms of FLOPs while maintaining the performance of resulting slim network to the original dense counterparts.