summaries: This paper proposes an efficient sparse training method to solves the problem of accelerating the training speed of deep neural networks and save memory usage.  The authors first formulate the training process as a continuous minimization problem under global sparsity constraints.  Then, separating the optimization process into two steps, corresponding to weight update and structure parameter update. A variance-reduced policy gradient estimator is proposed to update the structure parameter. Finally, the experimental results are solid and strong. The problem has lots of meanings for engineering implementation.   This paper addresses the topic of sparse training of neural networks. They target channel-level sparsity with the goal of realizing practical training speedups with existing software and hardware. The authors study the limitations of existing techniques in terms of their ability to exploit sparsity during training and show a wide array of empirical results with their proposed technique. Paper proposes a variational based pruning method where, with 2 structurally sparse forward passes and 1 sparse backward, the training can be significantly faster. Description is clear and theoretically sounds, equations are easy to follow, code is provided, results are encouraging. The paper might be a great benchmark and starting point for faster pruning-while-training with efficient inference.  The authors propose a channel-level sparse training algorithm that can effectively utilize the channel sparsity to accelerate both forward and backward propagation process. They divide the sparse training process into weight updating step and mask updating step. And solve the problem separately. More specifically, they develop a variance-reduced policy gradient estimator to update the trainable pruning mask without the need for backpropagation, and hence a sparse computation on both forward and backward propagation. Network sparsity has historically targeted only inference.  Recent explorations in training with sparsity fall short of their potential, though, by not fully exploiting sparsity in the forward and backward propagation steps.  The authors decompose training a channel-wise sparse network into two parts: training the weights themselves (as usual) and learning the sparse structure.  Rather than using a chain-rule based gradient step for the structure, a unique policy gradient estimator with reduced variance is used; it is composed of two forward passes, eliminating the need for the chain rule to be used to learn the sparse structure.  Experiments show competitive accuracy for saved parameters and FLOPs for a number of networks, as well as actual time saved during training with the proposed technique for VGG19 on CIFAR-10.
abstractText: Sparse training is a natural idea to accelerate the training speed of deep neural networks and save the memory usage, especially since large modern neural networks are significantly over-parameterized. However, most of the existing methods cannot achieve this goal in practice because the chain rule based gradient (w.r.t. structure parameters) estimators adopted by previous methods require dense computation at least in the backward propagation step. This paper solves this problem by proposing an efficient sparse training method with completely sparse forward and backward passes. We first formulate the training process as a continuous minimization problem under global sparsity constraint. We then separate the optimization process into two steps, corresponding to weight update and structure parameter update. For the former step, we use the conventional chain rule, which can be sparse via exploiting the sparse structure. For the latter step, instead of using the chain rule based gradient estimators as in existing methods, we propose a variance reduced policy gradient estimator, which only requires two forward passes without backward propagation, thus achieving completely sparse training. We prove that the variance of our gradient estimator is bounded. Extensive experimental results on real-world datasets demonstrate that compared to previous methods, our algorithm is much more effective in accelerating the training process, up to an order of magnitude faster.
authors.name: Xiao Zhou
authors.name: Weizhong Zhang
authors.name: Zonghao Chen
authors.name: Shizhe Diao
authors.name: Tong Zhang
references.authors: M. Abadi P. Barham J. Chen Z. Chen A. Davis J. Dean M. Devin S. Ghemawat G. Irving M. Isard
references.title: et al
references.venue: Tensorflow: A system for large-scale machine learning. In 12th {USENIX} symposium on operating systems design and implementation ({OSDI} 16), pages 265–283
references.authors: Y. Bengio N. Léonard A. Courville
references.title: Estimating or propagating gradients through stochastic neurons for conditional computation
references.venue: arXiv preprint arXiv:1308.3432
references.authors: J. Deng W. Dong R. Socher L.-J. Li K. Li L. Fei-Fei
references.title: Imagenet: A large-scale hierarchical image database
references.venue: 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee
references.authors: T. Dettmers L. Zettlemoyer
references.title: Sparse networks from scratch: Faster training without losing performance
references.venue: arXiv preprint arXiv:1907.04840
references.authors: U. Evci T. Gale J. Menick P.S. Castro E. Elsen
references.title: Rigging the lottery: Making all tickets winners
references.venue: International Conference on Machine Learning, pages 2943–2952. PMLR
references.authors: J. Frankle G.K. Dziugaite D.M. Roy M. Carbin
references.title: Stabilizing the lottery ticket hypothesis
references.venue: arXiv preprint arXiv:1903.01611
references.authors: Y. Gu W. Zhang C. Fang J.D. Lee T. Zhang
references.title: How to characterize the landscape of overparameterized convolutional neural networks
references.venue: Advances in Neural Information Processing Systems, volume 33, pages 3797–3807
references.authors: Y. Guo A. Yao Y. Chen
references.title: Dynamic network surgery for efficient dnns
references.venue: Advances in neural information processing systems, pages 1379–1387
references.authors: S. Han H. Mao W.J. Dally
references.title: Deep compression: Compressing deep neural networks with pruning
references.venue: trained quantization and huffman coding. International Conference on Learning Representations
references.authors: S. Han J. Pool J. Tran W. Dally
references.title: Learning both weights and connections for efficient neural network
references.venue: Advances in neural information processing systems, pages 1135–1143
references.authors: K. He X. Zhang S. Ren J. Sun
references.title: Deep residual learning for image recognition
references.venue: Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778
references.authors: Y. He G. Kang X. Dong Y. Fu Y. Yang
references.title: Soft filter pruning for accelerating deep convolutional neural networks
references.venue: IJCAI International Joint Conference on Artificial Intelligence
references.authors: Y. He X. Zhang J. Sun
references.title: Channel pruning for accelerating very deep neural networks
references.venue: Proceedings of the IEEE International Conference on Computer Vision, pages 1389–1397
references.authors: A.G. Howard M. Zhu B. Chen D. Kalenichenko W. Wang T. Weyand M. Andreetto H. Adam
references.title: Mobilenets: Efficient convolutional neural networks for mobile vision applications
references.venue: arXiv preprint arXiv:1704.04861
references.authors: E. Jang S. Gu B. Poole
references.title: Categorical reparameterization with gumbel-softmax
references.venue: International Conference on Learning Representations
references.authors: M. Kang B. Han
references.title: Operation-aware soft channel pruning using differentiable masks
references.venue: International Conference on Machine Learning, pages 5122–5131. PMLR
references.authors: A. Krizhevsky
references.title: Learning multiple layers of features from tiny images
references.venue: Master’s thesis, University of Tront
references.authors: A. Kusupati V. Ramanujan R. Somani M. Wortsman P. Jain S. Kakade A. Farhadi
references.title: Soft threshold weight reparameterization for learnable sparsity
references.venue: In Proceedings of the International Conference on Machine Learning,
references.authors: C. Lemaire A. Achkar P.-M. Jodoin
references.title: Structured pruning of neural networks with budget-aware regularization
references.venue: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9108–9116
references.authors: H. Li A. Kadav I. Durdanovic H. Samet H.P. Graf
references.title: Pruning filters for efficient convnets
references.venue: International Conference on Learning Representations
references.authors: L. Liebenwein C. Baykal H. Lang D. Feldman D. Rus
references.title: Provable filter pruning for efficient neural networks
references.venue: International Conference on Learning Representations
references.authors: B. Liu M. Wang H. Foroosh M. Tappen M. Pensky
references.title: Sparse convolutional neural networks
references.venue: Proceedings of the IEEE conference on computer vision and pattern recognition, pages 806–814
references.authors: S. Liu T. Chen X. Chen Z. Atashgahi L. Yin H. Kou L. Shen M. Pechenizkiy Z. Wang D.C. Mocanu
references.title: Sparse training via boosting pruning plasticity with neuroregeneration
references.venue: arXiv preprint arXiv:2106.10404
references.authors: Z. Liu J. Li Z. Shen G. Huang S. Yan C. Zhang
references.title: Learning efficient convolutional networks through network slimming
references.venue: Proceedings of the IEEE International Conference on Computer Vision, pages 2736–2744
references.authors: Z. Liu M. Sun T. Zhou G. Huang T. Darrell
references.title: Rethinking the value of network pruning
references.venue: International Conference on Learning Representations
references.authors: C. Louizos M. Welling D.P. Kingma
references.title: Learning sparse neural networks through l_0 regularization
references.venue: International Conference on Learning Representations
references.authors: J.-H. Luo J. Wu W. Lin
references.title: Thinet: A filter level pruning method for deep neural network compression
references.venue: Proceedings of the IEEE international conference on computer vision, pages 5058–5066
references.authors: S. Lym E. Choukse S. Zangeneh W. Wen S. Sanghavi M. Erez
references.title: Prunetrain: fast neural network training by dynamic sparse model reconfiguration
references.venue: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–13
references.authors: D.C. Mocanu E. Mocanu P. Stone P.H. Nguyen M. Gibescu A. Liotta
references.title: Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science
references.venue: Nature communications, 9(1):1–12
references.authors: H. Mostafa X. Wang
references.title: Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization
references.venue: International Conference on Machine Learning, pages 4646–4655. PMLR
references.authors: A. Paszke S. Gross F. Massa A. Lerer J. Bradbury G. Chanan T. Killeen Z. Lin N. Gimelshein L. Antiga
references.title: et al
references.venue: Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32:8026–8037
references.authors: A. Renda J. Frankle M. Carbin
references.title: Comparing rewinding and fine-tuning in neural network pruning
references.venue: International Conference on Learning Representations
references.authors: D.J. Rezende S. Mohamed D. Wierstra
references.title: Stochastic backpropagation and approximate inference in deep generative models
references.venue: International conference on machine learning, pages 1278–1286. PMLR
references.authors: P. Savarese H. Silva M. Maire
references.title: Winning the lottery with continuous sparsification
references.venue: Advances in Neural Information Processing Systems, 33
references.authors: K. Simonyan A. Zisserman
references.title: Very deep convolutional networks for large-scale image recognition
references.venue: International Conference on Learning Representations
references.authors: M. Song A. Montanari P. Nguyen
references.title: A mean field view of the landscape of two-layers neural networks
references.venue: Proceedings of the National Academy of Sciences, 115:E7665–E7671
references.authors: S. Srinivas A. Subramanya R. Venkatesh Babu
references.title: Training sparse neural networks
references.venue: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pages 138–145
references.authors: A. Vaswani N. Shazeer N. Parmar J. Uszkoreit L. Jones A.N. Gomez
references.title: L
references.venue: u. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems
references.authors: C. Wang G. Zhang R. Grosse
references.title: Picking winning tickets before training by preserving gradient flow
references.venue: International Conference on Learning Representations
references.authors: Z. Wang J. Wohlwend T. Lei
references.title: Structured pruning of large language models
references.venue: Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 6151–6162
references.authors: W. Wen C. Wu Y. Wang Y. Chen H. Li
references.title: Learning structured sparsity in deep neural networks
references.venue: Advances in neural information processing systems, pages 2074–2082
references.authors: X. Xiao Z. Wang S. Rajasekaran
references.title: Autoprune: Automatic network pruning by regularizing auxiliary parameters
references.venue: Advances in Neural Information Processing Systems, pages 13681–13691
references.authors: M. Ye C. Gong L. Nie D. Zhou A. Klivans Q. Liu
references.title: Good subnetworks provably exist: Pruning via greedy forward selection
references.venue: International Conference on Machine Learning, pages 10820–10830. PMLR
references.authors: X. Yuan P.H.P. Savarese M. Maire
references.title: Growing efficient deep networks by structured continuous sparsification
references.venue: International Conference on Learning Representations
references.authors: S. Zagoruyko N. Komodakis
references.title: Wide residual networks
references.venue: British Machine Vision Conference 2016. British Machine Vision Association
references.authors: W. Zeng R. Urtasun
references.title: MLPrune: Multi-layer pruning for automated neural network compression, 2019
references.authors: X. Zhou W. Zhang H. Xu T. Zhang
references.title: Effective sparsification of neural networks with global sparsity constraint
references.venue: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3599–3608
references.authors: M. Zhu S. Gupta
references.title: To prune
references.venue: or not to prune: exploring the efficacy of pruning for model compression. arXiv preprint arXiv:1710.01878
sections.heading: 1 Introduction
sections.text: In the last decade, deep neural networks (DNNs) [35, 11, 38] have proved their outstanding performance in various fields such as computer vision and natural language processing. However, training such large-sized networks is still very challenging, requiring huge computational power and storage. This hinders us from exploring larger networks, which are likely to have better performance. Moreover, it is a widely-recognized property that modern neural networks are significantly overparameterized, which means that a fully trained network can always be sparsified dramatically by network pruning techniques [9, 8, 25, 46, 20] into a small sub-network with negligible degradation in accuracy. After pruning, the inference efficiency can be greatly improved. Therefore, a natural question is can we exploit this sparsity to improve the training efficiency?
  The emerging technique called sparse network training [10] is closely related with our question, which can obtain sparse networks by training from scratch. We can divide existing methods into two categories, i.e., parametric and non-parametric, based on whether they explicitly parameterize network structures with trainable variables (termed structure parameters). Empirical results [24, 34, 44, 23] demonstrate that the sparse networks they obtain have comparable accuracy with those ∗Equal contribution †Jointly with Google Research
  35th Conference on Neural Information Processing Systems (NeurIPS 2021).
  obtained from network pruning. However, most of them narrowly aim at finding a sparse subnetwork instead of simultaneously sparsifying the computation of training by exploiting the sparse structure. As a consequence, it is hard for them to effectively accelerate the training process in practice on general platforms, e.g., Tensorflow [1] and Pytorch [31]. Detailed reasons are discussed below:
  • Non-parametric methods find the sparse network by repeating a two-stage procedure that alternates between weight optimization and pruning [10, 6], or by adding a proper sparsityinducing regularizer on the weights to the objective [22, 41]. The two-stage methods prune the networks in weight space and usually require retraining the obtained subnetwork from scratch every time when new weights are pruned, which makes training process even more time-consuming. Moreover, the computation of regularized methods is dense since the gradients of a zero-valued weights/filters are still nonzero.
  • All the parametric approaches estimate the gradients based on chain rule. The gradient w.r.t. the structure parameters can be nonzero even when the corresponding channel/weight is pruned. Thus, to calculate the gradient via backward propagation, the error has to be propagated through all the neurons/channels. This means that the computation of backward propagation has to be dense. Concrete analysis can be found in Section 3.
  We notice that some existing methods [4, 28] can achieve training speedup by careful implementation. For example, the dense to sparse algorithm [28] removes some channels if the corresponding weights are quite small for a long time. However, these methods always need to work with a large model at the beginning epochs and consume huge memory and heavy computation in the early stage. Therefore, even with such careful implementations, the speedups they can achieve are still limited.
  In this paper, we propose an efficient channel-level parametric sparse neural network training method, which is comprised of completely sparse (See Remark 1) forward and backward propagation. We adopt channel-level sparsity since such sparsity can be efficiently implemented on the current training platforms to save the computational cost. In our method, we first parameterize the network structure by associating each filter with a binary mask modeled as an independent Bernoulli random variable, which can be continuously parameterized by the probability. Next, inspired by the recent work [47], we globally control the network size during the whole training process by controlling the sum of the Bernoulli distribution parameters. Thus, we can formulate the sparse network training problem into a constrained minimization problem on both the weights and structure parameters (i.e., the probability). The main novelty and contribution of this paper lies in our efficient training method called completely sparse neural network training for solving the minimization problem. Specifically, to fully exploit the sparse structure, we separate training iteration into two parts, i.e., weight update and structure parameter update. For weight update, the conventional backward propagation is used to calculate the gradient, which can be sparsified completely because the gradients of the filters with zero valued masks are also zero. For structure parameter update, we develop a new variance reduced policy gradient estimator (VR-PGE). Unlike the conventional chain rule based gradient estimators (e.g., straight through[2]), VR-PGE estimates the gradient via two forward propagations, which is completely sparse because of the sparse subnetwork. Finally, extensive empirical results demonstrate that our method can significantly accelerate the training process of neural networks.
  The main contributions of this paper can be summarized as follows:
  • We develop an efficient sparse neural network training algorithm with the following three appealing features:
  – In our algorithm, the computation in both forward and backward propagations is completely sparse, i.e., they do not need to go through any pruned channels, making the computational complexity significantly lower than that in standard training.
  – During the whole training procedure, our algorithm works on small sub-networks with the target sparsity instead of follows a dense-to-sparse scheme.
  – Our algorithm can be implemented easily on widely-used platforms, e.g., Pytorch and Tensorflow, to achieve practical speedup.
  • We develop a variance reduced policy gradient estimator VR-PGE specifically for sparse neural network training, and prove that its variance is bounded.
  • Experimental results demonstrate that our methods can achieve significant speed-up in training sparse neural networks. This implies that our method can enable us to explore larger-sized neural networks in the future.
  Remark 1. We call a sparse training algorithm completely sparse if both its forward and backward propagation do not need to go through any pruned channels. For such algorithms, the computational cost in forward and backward propagation cost can be roughly reduced to ρ2 ∗ 100%, with ρ being the ratio of remaining unpruned channels.
sections.heading: 2 Related Work
sections.text: In this section, we briefly review the studies on neural network pruning, which refers to the algorithms that prune DNNs after fully trained, and the recent works on sparse neural network training.
sections.heading: 2.1 Neural Network Pruning
sections.text: Network Pruning [10] is a promising technique for reducing the model size and inference time of DNNs. The key idea of existing methods [10, 8, 46, 20, 27, 13, 48, 40, 43, 32, 16] is to develop effective criteria (e.g, weight magnitude) to identify and remove the massive unimportant weights contained in networks after training. To achieve practical speedup on general devices, some of them prune networks in a structured manner, i.e., remove the weights in a certain group (e.g., filter) together, while others prune the weights individually. It has been reported in the literature [8, 25, 46, 20] that they can improve inference efficiency and reduce memory usage of DNNs by orders of magnitudes with minor loss in accuracy, which enables the deployment of DNNs on low-power devices.
  We notice that although some pruning methods can be easily extended to train sparse networks, they cannot accelerate or could even slow down the training process. One reason is they are developed in the scenario that a fully trained dense network is given, and cannot work well on the models learned in the early stage of training. Another reason is after each pruning iteration, one has to fine tune or even retrain the network for lots of epoch to compensate the caused accuracy degradation.
sections.heading: 2.2 Sparse Neural Network Training
sections.text: The research on sparse neural network training has emerged in the recent years. Different from the pruning methods, they can find sparse networks without pre-training a dense one. Existing works can be divided into four categories based on their granularity in pruning and whether the network structures are explicitly parameterized. To the best of our knowledge, no significant training speedups achieved in practice are reported in the literature. Table 1 summarizes some representative works.
  Weight-level non-parametric methods, e.g., [6, 10, 48, 29, 30], always adopt a two-stage training procedure that alternates between weight optimization and pruning. They differ in the schedules of tuning the prune ratio over training and layers. [10] prunes the weights with the magnitude below a certain threshold and [48, 6] gradually increase the pruning rate during training. [30, 5] automatically reallocate parameters across layers during training via controlling the global sparsity.
  Channel-level non-parametric methods [12, 41] are proposed to achieve a practical acceleration in inference. [41] is a structured sparse learning method, which adds a group Lasso regularization into the objective function of DNNs with each group comprised of the weights in a filter. [12] proposes a soft filter pruning method. It zeroizes instead of hard pruning the filters with small `2 norm, after which these filters are treated the same with other filters in training. It is obvious that these methods cannot achieve significant speedup in training since they need to calculate the full gradient in backward propagation although the forward propagation could be sparsified if implemented carefully.
  Parametric methods multiply each weight/channel with a binary [47, 44, 37, 42] or continuous [24, 26, 19, 18] mask, which can be either deterministic [24, 42] or stochastic [47, 44, 26, 37, 19, 18]. The mask is always parameterized via a continuous trainable variable, i.e., structure parameter. The
  sparsity is achieved by adding sparsity-inducing regularizers on the masks. The novelties of these methods lie in estimating the gradients w.r.t structure parameters in training. To be precise,
  • Deterministic Binary Mask. [42] parameterizes its deterministic binary mask as a simple step function and estimates the gradients via sigmoid straight through estimator (STE) [2].
  • Deterministic Continuous Mask. [24] uses the linear coefficients of batch normalization (BN) as a continuous mask and enforces most of them to 0 by penalizing the objective with `1 norm of the coefficients. [18] defines the mask as a soft threshold function with learnable threshold. These methods can estimate the gradients via standard backward propagation.
  • Stochastic Binary Mask. [44, 37] model the mask as a bernoulli random variable and the gradients w.r.t. the parameters of bernoulli distributions are estimated via STE. [47] estimates the gradients via Gumbel-Softmax trick [15], which is more accurate than STE.
  • Stochastic Continuous Mask. [26, 19] parameterize the mask as a continuous function g(c, ), which is differentiable w.r.t. c, and is a parameter free noise, e.g., Gaussian noise N (0, 1). In this way, the gradients can be calcuated via conventional backward propagation.
  Therefore, we can see that all of these parametric methods estimate the gradients of the structure parameters based on the chain rule in backward propagation. This makes the training iteration cannot be sparsified by exploiting the sparse network structure. For the details, please refer to Section 3.
sections.heading: 3 Why Existing Parameteric Methods Cannot Achieve Practical Speedup?
sections.text: In this section, we reformulate existing parametric channel-level methods into a unified framework to explain why they cannot accelerate the training process in practice.
  Notice that convolutional layer can be viewed as a generalized fully connected layer, i.e., viewing the channels as neurons and convolution of two matrices as a generalized multiplication (see [7]). Hence, for simplicity, we consider the fully connected network in Figure 1. Moreover, since the channels in CNNs are corresponding to the neurons in fully connected networks, we consider neuron-level instead of weight-level sparse training in our example.
  As discussed in Section 2, existing methods parameterize the 4 kinds of mask in the following ways:
  (i): mi = φ(si); (ii): mi = ψ(si); (iii): mi = g(si, ), ∼ N (0, 1); (iv): m ∼ Bern(pi(s)), where the function φ(si) is binary, e.g., step function; ψ(si) is a continuous function; g(si, ) is differentiable w.r.t. si. All existing methods estimate the gradient of the loss `(ŷ, y) w.r.t. si based on chain rule, which can be formulated into a unified form below.
  Specifically, we take the pruned neuron x3 in Figure 1 as an example, the gradient is calculated as
  ∇s3`(ŷ, y) = ∂`(ŷ, y)
  ∂x̂3︸ ︷︷ ︸ a
  ( w>:,3xin )︸ ︷︷ ︸ forward ∂m3 ∂s3 . (1)
  Existing parametric methods developed different ways to estimate ∂m3∂s3 . Actually, for cases (ii) and (iii), the gradients are well-defined and thus can be calculated directly. STE is used to estimate the gradient in case (i) [42]. For cases (iv), [44, 37, 47] adopt STE and Gumbel-Softmax.
  In Eqn.(1), the term (a) is always nonzero especially when x̂3 is followed by BN. Hence, we can see that even for the pruned neuron x3, the gradient ∂m3∂s3 can be nonzero in all four cases. This means the backward propagation has to go though all the neurons/channels, leading to dense computation.
  At last, we can know from Eqn.(1) that forward propagation in existing methods cannot be completely sparse. Although w>:,3xin can be computed sparsely as in general models xin could be a sparse tensor of a layer with some channels being pruned, we need to calculate it for each neuron via forward propagation to calculate RHS of Eqn.(1). Thus, even if carefully implemented, the computational cost of forward propagation can only be reduced to ρ ∗ 100% instead of ρ2 ∗ 100% as in inference. That’s why we argue that existing methods need dense computation at least in backward propagation. So they cannot speed up the training process effectively in practice. Remark 2. The authors of GrowEfficient [44] confirmed that actually they also calculated the gradient of qc w.r.t, sc in their Eqn.(6) via STE even if qc = 0. Thus need dense backward propagation.
sections.heading: 4 Channel-level Completely Sparse Neural Network Training
sections.text: Below, we present our sparse neural network training framework and the efficient training algorithm.
sections.heading: 4.1 Framework of Channel-level Sparse Training
sections.text: Given a convolutional network f(x;w), let {Fc : c ∈ C} be the set of filters with C being the set of indices of all the channels. To parameterize the network structure, we associate each Fc with a binary mask mc, which is an independent Bernoulli random variable. Thus, each channel is computed as
  xout, c = xin ∗ (Fcmc) ,
  with ∗ being the convolution operation. Inspired by [47], to avoid the problems, e.g., gradient vanishing, we parameterize mc directly on the probability sc, i.e., mc equals to 1 and 0 with the probabilities sc and 1 − sc, respectively. Thus, we can control the channel size by the sum of sc. Following [47], we can formulate channel-level sparse network training into the following framework:
  min w,s
  Ep(m|s) L(w,m) := 1
  N N∑ i=1 ` (f (xi;w,m) ,yi) (2)
  s.t. w ∈ Rn, s ∈ S := {s ∈ [0, 1]|C| : 1>s ≤ K},
  where {(xi,yi)}Ni=1 is the training dataset, w is the weights of the original network, f (·; ·, ·) is the pruned network, and `(·, ·) is the loss function, e.g, cross entropy loss. K = ρ|C| controls the remaining channel size with ρ being the remaining ratio of the channels.
  Discussion. We’d like to point out that although our framework is inspired by [47], our main contribution is the efficient solver comprised of completely sparse forward/backward propagation for Problem (2). Moreover, our framework can prune the weights in fully connected layers together, since we can associate each weight with an independent mask.
sections.heading: 4.2 Completely Sparse Training with Variance Reduced Policy Gradient
sections.text: Now we present our completely sparse training method, which can solve Problem (2) via completely sparse forward and backward propagation. The key idea is to separate the training iteration into filter update and structure parameter update so that the sparsity can be fully exploited.
sections.heading: 4.2.1 Filter Update via Completely Sparse Computation
sections.text: It is easy to see that the computation of the gradient w.r.t. the filters can be sparsified completely. To prove this point, we just need to clarify the following two things:
  • We do not need to update the filters corresponding to the pruned channels. Consider a pruned channel c, i.e., mc = 0, then due to the chain rule, we can have
  ∂` (f (xi;w,m)) ∂Fc = ∂` (f (xi;w,m)) ∂xout,c ∂xout,c ∂Fc ≡ 0,
  the last equation holds since xout,c ≡ 0. This indicates that the gradient w.r.t the pruned filter Fc is always 0, and thus Fc does not need to be updated.
  • The error cannot pass the pruned channels via backward propagation. Consider a pruned channel c, we denote its output before masking as x̂out,c = xin ∗ Fc, then the error propagating through this channel can be computed as
  ∂` (f (xi;w,m)) ∂x̂out,c = ∂` (f (xi;w,m)) ∂xout,c ∂xout,c x̂out,c ≡ 0.
  This demonstrates that to calculate the gradient w.r.t. the unpruned filters, the backward propagation does not need to go through any pruned channels.
  Therefore, the filters can be updated via completely sparse backward propagation.
sections.heading: 4.2.2 Structure Parameter Update via Variance Reduced Policy Gradient
sections.text: We notice that policy gradient estimator (PGE) can estimate the gradient via forward propagation, avoiding the pathology of chain rule based estimators as dicussed in Section 3. For abbreviation, we denote L(w,m) as L(m) since w can be viewed as a constant here. The objective can be written as
  Φ(s) = Ep(m|s) L(m),
  which can be optimized using gradient descent:
  s← s− η∇Φ(s).
  with learning rate η. One can obtain a stochastic unbiased estimate of the gradient∇Φ(s) using PGE:
  ∇Φ(s) = Ep(m|s) L(m)∇s ln p(m|s), (PGE)
  leading to Policy Gradient method, which may be regarded as a stochastic gradient descent algorithm:
  s← s− ηL(m)∇s ln p(m|s). (3)
  In Eqn.(3), L(m) can be computed via completely sparse forward propagation and the computational cost of ∇s ln p(m|s) = m−ss(1−s) is negligible, therefore PGE is computationally efficient.
  However, in accordance with the empirical results reported in [33, 15], we found that standard PGE suffers from high variance and does not work in practice. Below we will develop a Variance Reduced Policy Gradient Estimator (VR-PGE) starting from theoretically analyzing the variance of PGE.
  Firstly, we know that this variance of PGE is
  Ep(m|s) L2(m)‖∇s ln p(m|s)‖22 − ‖∇Φ(s)‖22,
  which can be large because L(m) is large. Mean Field theory [36] indicates that, while L(m) can be large, the term L(m)− L(m′) is small when m and m′ are two independent masks sampled from a same distributionp(m|s) (see the appendix for the details). This means that we may consider the following variance reduced preconditioned policy gradient estimator:
  Em′∼p(m′|s)Em∼p(m|s) (L(m)− L(m′))Hα(s)∇s ln p(m|s), (VR-PGE)
  where Hα(s) is a specific diagonal preconditioning matrix
  Hα(s) = diag (s ◦ (1− s))α , (4)
  with α ∈ (0, 1) and ◦ being the element-wise product. It plays a role as adaptive step size and it is shown that this term can reduce the variance of the stochastic PGE term∇s ln p(m|s). The details can be found in the appendix. Thus Φ(s) can be optimized via:
  s← s− η (L (m)− L (m′))Hα(s)∇s ln p(m|s). (5)
  In our experiments, we set α to be 12 for our estimator VR-PGE. The theorem below demonstrates that VR-PGE can have bounded variance.
  Algorithm 1 Completely Sparse Neural Network Training Input: target remaining ratio ρ, a dense network w, the step size η, and parameter α in (4) . 1: Initialize w, let s = ρ1. 2: for training epoch t = 1, 2 . . . T do 3: for each training iteration do 4: Sample mini batch of data B = {(x1,y1) , . . . , (xB ,yB)}. 5: Sample m(i) from p(m|s), i = 1, 2. 6: Update s and w
  s← projS(z) with z = s− η ( LB(w,m(1))− LB(w,m(2)) ) Hα(s)m (1)−s s(1−s) ,
  w ← w − η∇wLB ( w,m(1) ) 7: end for 8: end for 9: return A pruned network w ◦m by sampling a mask m from the distribution p(m|s).
  Theorem 1. Suppose m and m′ are two independent masks sampled from the Bernoulli distribution p(m|s), then for any α ∈ [ 12 , 1) and s ∈ (0, 1) |C|, the variance is bounded for
  (L(m)− L(m′))Hα(s)∇s ln p(m|s)
  Finally, we provide a complete view of our sparse training algorithm in Algorithm 1, which is essentially a projected stochastic gradient descent equipped with our efficient gradient estimators above. The projection operator in Algorithm 1 can be computed efficiently using Theorem 1 of [47].
  Discussion. In our algorithm, benefited from our constraint on s, the channel size of the neural network during training can be strictly controlled. This is in contrast with GrowEfficient [44], which ultilizes regularizer term to control the model size and has situations where model size largely drift away from desired. This will have larger demand for the GPU memory storage and have more risk that memory usage may explode, especially when we utilize sparse learning to explore larger models. Moreover, our forward and backward propagations are completely sparse, i.e., they do not need to go through any pruned channels. Therefore, the computational cost of each training iteration can be roughly reduced to ρ2 ∗ 100% of the dense network.
sections.heading: 5 Experiments
sections.text: In this section, we conduct a series of experiments to demonstrate the outstanding performance of our method. We divide the experiments into five parts. In part one, we compare our method with several state-of-the-art methods on CIFAR-10 [17] using VGG-16 [35], ResNet-20 [11] and WideResNet-28-10 [45] to directly showcase the superiority of our method. In part two, we directly compare with state-of-the-art method GrowEfficient [44] especially on extremely sparse regions, and on two high capacity networks VGG-19 [35] and ResNet-32 [11] on CIFAR-10/100 [17]. In part three, we conduct experiments on a large-scale dataset ImageNet [3] with ResNet-50 [11] and MobileNetV1 [14] and compare with GrowEfficient [44] across a wide sparsity region. In part four, we present the train-computational time as a supplementary to the conceptual train-cost savings to justify the applicability of sparse training method into practice. In part five, we present further analysis on epoch-wise train-cost dynamics and experimental justification of variance reduction of VR-PGE. Due to the space limitation, we postpone the experimental configurations, calculation schemes on train-cost savings and train-computational time and additional experiments into appendix.
sections.heading: 5.1 VGG-16, ResNet-20 and WideResNet-28-10 on CIFAR-10
sections.text: Table 2 presents Top-1 validation accuracy, parameters, FLOPs and train-cost savings comparisons with channel pruning methods L1-Pruning [20], SoftNet [12], ThiNet [27], Provable [21] and sparse training method GrowEfficient [44]. SoftNet can train from scratch but requires completely dense computation. Other pruning methods all require pretraining of dense model and multiple rounds of pruning and finetuning, which makes them slower than vanilla dense model training. Therefore the train-cost savings of these methods are below 1× and thus shown as ("-") in Table 2.
  GrowEfficient [44] is a recently proposed state-of-the-art channel-level sparse training method showing train-cost savings compared with dense training. As described in Section 3, GrowEfficient features completely dense backward and partially sparse forward pass, making its train-cost saving limited by 32 . By contrast, the train-cost savings of our method is not limited by any constraint. The details of how train-cost savings are computed can be found in appendix.
  Table 2 shows that our method generally exhibits better performance in terms of validation accuracy, parameters and particularly FLOPs. In terms of train-cost savings, our method shows at least 1.85× speed-up against GrowEfficient [44] and up to 9.39× speed-up against dense training.
sections.heading: 5.2 Wider Range of Sparsity on CIFAR-10/100 on VGG-19 and ResNet-32
sections.text: In this section, we explore sparser regions of training efficiency to present a broader comparision with state-of-the-art channel sparse training method GrowEfficient [44].
  We plot eight figures demonstrating the relationships between the Top-1 validation accuracy, FLOPs and train-cost savings. We find that our method generally achieves higher accuracy under same FLOPs
  settings. To be noted, the train-cost savings of our method is drastically higher than GrowEfficient [44], reaching up to 58.8× when sparisty approches 1.56% on ResNet-32 on CIFAR-100, while the speed-up of GrowEfficient is limited by 32 .
sections.heading: 5.3 ResNet-50 and MobileNetV1 on ImageNet-1K
sections.text: In this section, we present the performance boost obtained by our method on ResNet-50 and MobileNetV1 on ImageNet-1K [3]. Our method searches a model with 76.0% Top-1 accuracy, 48.2% parameters and 46.8% FLOPs beating all compared state-of-the-art methods. The train-cost saving comes up to 1.60× and is not prominent due to the accuracy constraint to match up with compared methods. Therefore we give a harder limit to the channel size and present sparser results on the same Table 3, reaching up to 7.36× speed-up while still preserving 69.3% Top-1 accuracy. For the already compact model MobileNetV1, we plot two figures in Figure 3 comparing with GrowEfficient [44]. We find that our method is much stabler in sparse regions and obtains much higher train-cost savings.
sections.heading: 5.4 Actual Training Computational Time Testing
sections.text: In this section, we provide actual training computational time on VGG-19 and CIFAR-10. The GPU in test is RTX 2080 Ti and the deep learning framework is Pytorch [31]. The intent of this section is to justify the feasibility of our method in reducing actual computational time cost, rather than staying in conceptual training FLOPs reduction. The computational time cost is measured by wall clock time, focusing on forward and backward propagation. We present training computational time in Table 4 with varying sparsity as in Figure 2. It shows that the computational time savings increases steadily with the sparisty. We also notice the gap between the savings in FLOPS and computational time. The gap comes from the difference between FLOPs and actual forward/backward time. More specifically, forward/backward time is slowed down by data-loading processes and generally affected by hardware latency and throughput, network architecture, etc. At extremely sparse regions, the pure computational time of sparse networks only occupies little of the forward/backward time and the cost of data management and hardware latency dominates the wall-clock time. Despite this gap, it can be expected that our train-cost savings can be better translated into real speed-up in exploring large
  models where the pure computational time dominates the forward/backward time, which promises a bright future for making training infeasibly large models into practice.
sections.heading: 5.5 Further Analysis
sections.text: [Epoch-wise Train-cost Dynamics of Sparse Training Process] We plot the train-cost dynamics in Figure 3. The vertical label is the ratio of train-cost to dense training, the inverse of train-cost savings. This demonstrates huge difference between our method and GrowEfficient [44]. The model searched by our method exhibits 92.73% Top-1 accuracy, 16.68% parameters, 14.28% FLOPs with 5.28× train-cost savings, while the model searched by GrowEfficient exhibits 92.47% Top-1 accuracy, 18.08% parameters, 22.74% FLOPs with 1.21× train-cost savings. [Experimental Verification of Variance Reduction of VR-PGE against PGE] We plot the mean of variance of gradients of channels from different layers. The model checkpoint and input data are selected randomly. The gradients are calculated in two approaches, VR-PGE and PGE. From the rightmost graph of Figure 3, we find that the VR-PGE reduces variance significantly, up to 3 orders of magnitude.
sections.heading: 6 Conclusion
sections.text: This paper proposes an efficient sparse neural network training method with completely sparse forward and backward passes. A novel gradient estimator named VR-PGE is developed for updating structure parameters, which estimates the gradient via two sparse forward propagation. We theoretically proved that VR-PGE has bounded variance. In this way, we can separate the weight and structure update in training and making the whole training process completely sparse. Emprical results demonstrate that the proposed method can significantly accelerate the training process of DNNs in practice. This enables us to explore larger-sized neural networks in the future.
  Acknowledgments and Disclosure of Funding
  This work is supported by GRF 16201320.
title: Efficient Neural Network Training via Forward and Backward Propagation Sparsification