summaries: The authors propose a single stage pruning method (DAM) that jointly prunes and refines weights during training. The method uses a monotonically increasing gate function for the neurons/channels in each layer with one trainable parameter. The gate function only discriminates neurons based on the position of them in the layer. The proposed method achieves reasonable results without any fine-tuning. This paper proposes a novel structured pruning method called DiscriminAtive Masking (DAM) for neuron-level sparsity. A relu-tanh gating function is proposed to apply on the neuron activations, and gradually zeroing out the neurons at lower order by introducing a L0 norm. The proposed method is a single stage method, and does not require further after-prune finetuning. Extensive experiments demonstrate the effectiveness of the proposed method This paper proposes a gradual structured pruning method, DAM, that can achieve good performance with various applications. The paper conducts extensive experiments to support the effectiveness of the proposed methods. The method part is reasonable and makes sense. This paper presents a simple and novel in-training model pruning algorithm to learn sparse feature representation with reduced computational load w/o compromising accuracy. To do this, authors introduce a single learnable parameter to simplify the optimization involved in enforcing L0 sparsity. Authors further provide theoretic analysis on how reducing the single learnable parameter would enforce L0 sparsity. Finally, authors demonstrate the remarkable performance over various applications comparing with the established strong baselines.    
abstractText: A central goal in deep learning is to learn compact representations of features at every layer of a neural network, which is useful for both unsupervised representation learning and structured network pruning. While there is a growing body of work in structured pruning, current state-of-the-art methods suffer from two key limitations: (i) instability during training, and (ii) need for an additional step of fine-tuning, which is resource-intensive. At the core of these limitations is the lack of a systematic approach that jointly prunes and refines weights during training in a single stage, and does not require any fine-tuning upon convergence to achieve state-of-the-art performance. We present a novel single-stage structured pruning method termed DiscriminAtive Masking (DAM). The key intuition behind DAM is to discriminatively prefer some of the neurons to be refined during the training process, while gradually masking out other neurons. We show that our proposed DAM approach has remarkably good performance over a diverse range of applications in representation learning and structured pruning, including dimensionality reduction, recommendation system, graph representation learning, and structured pruning for image classification. We also theoretically show that the learning objective of DAM is directly related to minimizing the L0 norm of the masking layer. All of our codes and datasets are available https://github.com/jayroxis/ dam-pytorch.
authors.name: Jie Bu
authors.name: Arka Daw
references.authors: Sanjeev Arora Rong Ge Behnam Neyshabur Yi Zhang
references.title: Stronger generalization bounds for deep nets via a compression approach
references.venue: In International Conference on Machine Learning,
references.authors: Haoli Bai Jiaxiang Wu Irwin King Michael Lyu
references.title: Few shot network compression via cross distillation
references.venue: In AAAI,
references.authors: Yoshua Bengio Aaron Courville Pascal Vincent
references.title: Representation learning: A review and new perspectives
references.venue: IEEE transactions on pattern analysis and machine intelligence,
references.authors: Davis Blalock Jose Javier Gonzalez Ortiz Jonathan Frankle John Guttag
references.title: What is the state of neural network pruning
references.venue: arXiv preprint arXiv:2003.03033,
references.authors: Jie Bu Anuj Karpatne
references.title: Quadratic Residual Networks: A New Class of Neural Networks for Solving Forward and Inverse Problems in Physics Involving PDEs, pages 675–683
references.venue: URL https://epubs.siam.org/doi/abs/10
references.authors: Tim R Davidson Luca Falorsi Nicola De Cao Thomas Kipf Jakub M Tomczak
references.title: Hyperspherical variational auto-encoders
references.venue: arXiv preprint arXiv:1804.00891,
references.authors: Misha Denil Babak Shakibi Laurent Dinh Marc’Aurelio Ranzato Nando de Freitas
references.title: Predicting parameters in deep learning
references.venue: In Proceedings of the 26th International Conference on Neural Information Processing Systems-Volume
references.authors: Xinhan Di Pengqian Yu Rui Bu Mingchao Sun
references.title: Mutual information maximization in graph neural networks
references.venue: In 2020 International Joint Conference on Neural Networks (IJCNN),
references.authors: Utku Evci Trevor Gale Jacob Menick Pablo Samuel Castro Erich Elsen
references.title: Rigging the lottery: Making all tickets winners
references.venue: Proceedings of the 37th International Conference on Machine Learning,
references.authors: Jonathan Frankle Michael Carbin
references.title: The lottery ticket hypothesis: Finding sparse, trainable neural networks
references.venue: In ICLR,
references.authors: Jonathan Frankle Gintare Karolina Dziugaite Daniel M Roy Michael Carbin
references.title: Stabilizing the lottery ticket hypothesis
references.venue: arXiv preprint arXiv:1903.01611,
references.authors: Jonathan Frankle Gintare Karolina Dziugaite Daniel M. Roy Michael Carbin
references.title: Linear mode connectivity and the lottery ticket hypothesis, 2020
references.authors: Trevor Gale Erich Elsen Sara Hooker
references.title: The state of sparsity in deep neural networks, 2019
references.authors: Ariel Gordon Elad Eban Ofir Nachum Bo Chen Hao Wu Tien-Ju Yang Edward Choi
references.title: Morphnet: Fast & simple resource-constrained structure learning of deep networks
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
references.authors: Kathrin Grosse Michael Backes
references.title: Do winning tickets exist before DNN training?, pages 549–557
references.venue: URL https://epubs.siam
references.authors: Song Han Jeff Pool John Tran William Dally
references.title: Learning both weights and connections for efficient neural network
references.venue: Advances in Neural Information Processing Systems,
references.authors: I. Higgins Loïc Matthey A. Pal Christopher P. Burgess Xavier Glorot M. Botvinick S. Mohamed Alexander Lerchner
references.title: beta-vae: Learning basic visual concepts with a constrained variational framework
references.venue: In ICLR,
references.authors: G.E. Hinton R.R. Salakhutdinov
references.title: Reducing the dimensionality of data with neural networks
references.venue: Science, 313(5786):504–507,
references.authors: Diederik P Kingma Max Welling
references.title: Auto-encoding variational bayes
references.venue: arXiv preprint arXiv:1312.6114,
references.authors: Thomas N Kipf Max Welling
references.title: Semi-supervised classification with graph convolutional networks
references.venue: arXiv preprint arXiv:1609.02907,
references.authors: Thomas N Kipf Max Welling
references.title: Variational graph auto-encoders
references.venue: arXiv preprint arXiv:1611.07308,
references.authors: Simon Kornblith Mohammad Norouzi Honglak Lee Geoffrey Hinton
references.title: Similarity of neural network representations revisited
references.venue: In International Conference on Machine Learning,
references.authors: Duong Hoang Le Binh-Son Hua
references.title: Network pruning that matters: A case study on retraining variants
references.venue: In International Conference on Learning Representations,
references.authors: Quoc V Le Jiquan Ngiam Adam Coates Ahbik Lahiri Bobby Prochnow Andrew Y Ng
references.title: On optimization methods for deep learning
references.venue: In ICML,
references.authors: Namhoon Lee Thalaiyasingam Ajanthan Philip Torr
references.title: Snip: Single-shot network pruning based on connection sensitivity
references.venue: In International Conference on Learning Representations,
references.authors: Hao Li Asim Kadav Igor Durdanovic Hanan Samet Hans Peter Graf
references.title: Pruning filters for efficient convnets
references.venue: arXiv preprint arXiv:1608.08710,
references.authors: Shiwei Liu Lu Yin Decebal Constantin Mocanu Mykola Pechenizkiy
references.title: Do we actually need dense over-parameterization? in-time over-parameterization in sparse training, 2021
references.authors: Zhuang Liu Jianguo Li Zhiqiang Shen Gao Huang Shoumeng Yan Changshui Zhang
references.title: Learning efficient convolutional networks through network slimming
references.authors: Zhuang Liu Mingjie Sun Tinghui Zhou Gao Huang Trevor Darrell
references.title: Rethinking the value of network pruning
references.venue: In ICLR,
references.authors: Christos Louizos Max Welling Diederik P Kingma
references.title: Learning sparse neural networks through l_0 regularization
references.venue: In ICLR,
references.authors: Costas Mavromatis George Karypis
references.title: Graph infoclust: Leveraging cluster-level node information for unsupervised graph representation learning
references.venue: arXiv preprint arXiv:2009.06946,
references.authors: Decebal Constantin Mocanu Elena Mocanu Peter Stone Phuong H. Nguyen Madeleine Gibescu Antonio Liotta
references.title: Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science
references.venue: Nature Communications,
references.authors: Bruno A Olshausen David J Field
references.title: Emergence of simple-cell receptive field properties by learning a sparse code for natural images
references.authors: Shirui Pan Ruiqi Hu Guodong Long Jing Jiang Lina Yao Chengqi Zhang
references.title: Adversarially regularized graph autoencoder for graph embedding
references.venue: arXiv preprint arXiv:1802.04407,
references.authors: Marc’Aurelio Ranzato Y-Lan Boureau Yann LeCun
references.title: Sparse feature learning for deep belief networks
references.venue: Advances in neural information processing systems,
references.authors: Alex Renda Jonathan Frankle Michael Carbin
references.title: Comparing rewinding and fine-tuning in neural network pruning
references.venue: In International Conference on Learning Representations,
references.authors: Jorma Rissanen
references.title: Stochastic complexity and modeling
references.venue: The annals of statistics,
references.authors: Pedro Savarese Hugo Silva Michael Maire
references.title: Winning the lottery with continuous sparsification
references.authors: Skyler Seto Martin T Wells Wenyu Zhang
references.title: Halo: Learning to prune neural networks with shrinkage
references.venue: In Proceedings of the 2021 SIAM International Conference on Data Mining (SDM),
references.authors: Nitish Srivastava Geoffrey Hinton Alex Krizhevsky Ilya Sutskever Ruslan Salakhutdinov
references.title: Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning
references.authors: Ilya Sutskever James Martens George Dahl Geoffrey Hinton
references.title: On the importance of initialization and momentum in deep learning
references.venue: In International conference on machine learning,
references.authors: Robert Tibshirani
references.title: Regression shrinkage and selection via the lasso
references.venue: Journal of the Royal Statistical Society: Series B (Methodological),
references.authors: Rishabh Tiwari Udbhav Bamba Arnav Chavan Deepak Gupta
references.title: Chipnet: Budgetaware pruning with heaviside continuous approximations
references.venue: In ICLR,
references.authors: Pascal Vincent Hugo Larochelle Yoshua Bengio Pierre-Antoine Manzagol
references.title: Extracting and composing robust features with denoising autoencoders
references.venue: In Proceedings of the 25th International Conference on Machine Learning,
references.authors: Jiayun Wang Yubei Chen Rudrasis Chakraborty Stella X Yu
references.title: Orthogonal convolutional neural networks
references.venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,
references.authors: Hong Yang Shirui Pan Peng Zhang Ling Chen Defu Lian Chengqi Zhang
references.title: Binarized attributed network embedding
references.venue: IEEE International Conference on Data Mining (ICDM),
references.authors: Jihun Yun Peng Zheng Eunho Yang Aurelie Lozano Aleksandr Aravkin
references.title: Trimming the l1 regularizer: Statistical analysis, optimization, and applications to deep learning
references.venue: In International Conference on Machine Learning,
references.authors: Muhan Zhang Yixin Chen
references.title: Inductive matrix completion based on graph neural networks
references.venue: In ICLR,
references.authors: Wenda Zhou Victor Veitch Morgane Austern Ryan P. Adams Peter Orbanz
references.title: Compressibility and generalization in large-scale deep learning
references.venue: CoRR, abs/1804.05862,
references.authors: Michael Zhu Suyog Gupta
references.title: To prune, or not to prune: exploring the efficacy of pruning for model compression, 2017
references.authors: Xiaotian Zhu Wengang Zhou Houqiang Li
references.title: Improving deep neural network sparsity through decorrelation regularization
references.venue: In IJCAI,
sections.heading: 1 Introduction
sections.text: A central goal in deep learning is to learn compact (or sparse) representations of features at every layer of a neural network that are useful in a variety of machine learning tasks. For example, in unsupervised representation learning problems [3], there is a long-standing goal to learn lowdimensional embeddings of input features that are capable of reconstructing the original data [44, 35, 24, 33]. Similarly, in supervised learning problems, there is a growing body of work in the area of network pruning [4], where the goal is to reduce the size of modern-day neural networks (that are known to be heavily over-parameterized [7, 51, 22, 45]) so that they can be deployed in resource-constrained environments (e.g., over mobile devices) without compromising on their accuracy. From this unified view of representation learning and network pruning, the generic problem of “learning compact representations” has applications in several machine learning use-cases such as dimensionality reduction, graph representation learning, matrix factorization, and image classification.
  *These authors contributed equally to this work.
  35th Conference on Neural Information Processing Systems (NeurIPS 2021).
  A theoretically appealing approach for learning compact representations is to introduce regularization penalties in the learning objective of deep learning that enforce L0 sparsity of the network parameters, θ. However, directly minimizing the L0 norm requires performing a combinatorial search over all possible subsets of weights in θ, which is computationally intractable. In practice, a common approach for enforcing sparsity is to use a continuous approximation of the L0 penalty in the learning objective, e.g., L1-based regularization (or Lasso [42]) and its variants [39, 47]. While such techniques are capable of pruning individual weights and thus reducing storage requirements, they do not offer any direct gains in inference speed since the number of features generated at the hidden layers can still be large even though the network connectivity is sparse [4]. Instead, we are interested in the area of structured network pruning for learning compact representations, where the sparsity is induced at the level of neurons by pruning features (or channels) instead of individual weights.
  While there is a growing body of work in structured network pruning [26, 28, 43, 14], the basic structure of most state-of-theart (SOTA) methods in this area (e.g., ChipNet [43] and NetSlim [28]) can be described as training a learnable vector of mask parameters, g ∈ Rn,g = [g1, g2, ..., gn]⊺, which when multiplied with the features extracted at a hidden layer, h ∈ Rn, results in the pruned outputs of this layer, o = g ⊙ h. Sparsity in the mask parameters is generally enforced using different approximations of the L0 norm of g (e.g., use of Lasso in NetSlim and use of “crispness” loss in ChipNet). Despite recent progress in this area, current SOTA in structured pruning suffer from two key limitations. First, since most methods do not explicitly minimize the L0 norm of g during pruning, they often suffer from training instabilities. In particular, most SOTA methods [28, 43] involve thresholding techniques to set small non-zero weights to zero leading to large drops in accuracy during the training process, as
  evidenced by our results in this paper. Second, once the pruning process is complete and we have converged at a compact network, most SOTA methods still need an additional step of fine-tuning the network in order to achieve reasonable accuracy. This is not only resource-intensive but there is also an on-going discussion on whether and how we should fine-tune [36] or rewind to initial weights [23] or train from scratch [29], making it difficult to prefer one approach over another.
  At the core of these limitations is the lack of a systematic approach for structured network pruning that jointly prunes and refines weights during training in a single stage, and does not require any fine-tuning upon convergence to achieve SOTA performance. Notice that in existing methods for structured pruning, allowing every neuron j to be pruned differently using an independent mask parameter gj only increases the number of learnable (or free) parameters in the learning objective, increasing the complexity of the problem. Instead, we ask the question: “Can we leverage the intrinsic symmetry of neural network initialization to design a pruning mask using the least number of free parameters?”
  We present a simple solution to this question by proposing a new single-stage structured pruning method that learns compact representations while training and does not require fine-tuning, termed DiscriminAtive Masking (DAM). The basic idea of DAM is to use a monotonically increasing gate function G for masking every neuron in a layer that only depends on the index j = 1, 2, ..., n (or position) of the neuron in the layer and a scalar parameter β to be learned during training (see Figure 1). At the start of training, the gate function admits non-zero values for all neurons in the layer, allowing all neurons to be unmasked (or active). As the training progresses, the gate function gradually shifts from “left” to “right” as a result of updating β such that upon convergence, only a subset of neurons (on the extreme right) are active while all others are masked out. The key intuition behind DAM is to discriminatively prefer some of the neurons (on the right) to be refined (or re-adapted) during the training process for capturing useful features, while gradually masking out (or pruning) neurons on the left. This preferential pruning of neurons using a very simple gate function helps in regulating the number of features transmitted to the next layer.*
  Contributions: We show that our proposed DAM approach has remarkably good performance over various applications, including dimensionality reduction, recommendation system, graph represen-
  *This is somewhat analogous to how a physical dam regulates the flow of water by shifting a movable gate.
  tation learning, and structured pruning for image classification, achieving SOTA performance for structured pruning. This shows the versatility of DAM in learning compact representations on diverse problems and network choices, in contrast to baseline methods for structured pruning that are only developed and tested for the problem of image classification [28, 43]. Our approach is single-stage, does not require fine-tuning (and hence has lower training time), and does not suffer from training instabilities. We also theoretically show that the learning objective of DAM is directly related to minimizing the L0 norm of the discriminative mask, providing a new differentiable approximation for enforcing L0 sparsity. Our approach is also easy to implement and can be applied to any layer in a network architecture. The simplicity and effectiveness of DAM provides unique insights into the problem of learning compact representations and its relationship with preferential treatment of neurons for pruning and refining, opening novel avenues of systematic research in this rapidly growing area.
sections.heading: 2 Related Work
sections.text: Learning with Sparsity: There is a long history of methods for enforcing compactness (or sparsity) in learned parameters [33, 42], where it is understood that compact hypotheses can not only result in more efficient data representations but also promote better generalizability [37]. This connection between compactness and generalization has also been recently explored in the context of deep learning [49, 1], motivating our problem of learning compact representations.
  Unstructured Pruning: There are several methods that have been developed for unstructured network pruning [30, 10, 12, 38], where the goal is to prune individual weights of the network. A widely used technique in this field is referred to as the Lottery Ticket Hypothesis (LTH) [10], which suggests that a certain subset of weights in a network may be initialized in such a way that they are likely “winning tickets”, and the goal of network pruning is then to simply uncover such winning tickets from the initialization of an over-parameterized network using magnitude-based pruning. Similar to our work, there exists a line of work in the area of dynamic sparse training for unstructured pruning [50, 13, 32, 9, 27] that gradually prunes the model to the target sparsity during training. However, while unstructured pruning methods can show very large gains in terms of pruning ratio, they are not directly useful for learning compact representations, since the pruned weights, even though sparse, may not be arranged in a fashion conducive to the goal of reducing the number of features.
  Structured Network Pruning: The goal of structured pruning is to enforce sparsity over the network substructures (e.g., neurons or convolutional channels) that is directly relevant to the goal of learning compact representations. Most structured pruning SOTA methods need a three-stage process [16], i.e., training, pruning and finetuning, in order to get highly compressed models. Early exploration on the L1 norm based filter pruning can be traced back to [26], followed by Network Slimming (NetSlim) [28], which uses a sparsity-inducing regularization based on L1 norm. Following this line of work, a recent development in the area of network pruning is ChipNet [43], which introduces an additional crispness loss to the L1 regularization, and achieves SOTA performance on benchmark tasks. However, network slimming can only work if the underlying architecture has batch normalization layers. It also has training instabilities due to soft-thresholding inherent to Lasso. ChipNet also iterates between soft- and hard-pruning, which is done after the training phase. Further, both these approaches require finetuning or retraining, which can be resource intensive.
  Faster pruning schemes have also been explored by previous works. Bai et al [2] proposed a few-shot structured pruning algorithm utilizing filter cross distillation. The SNIP [25] took one step ahead by performing single-stage pruning. However, SNIP falls in the category of unstructured pruning. To the best of our knowledge, no existing structured pruning method has been demonstrated to achieve SOTA performance without finetuning. For example, while the method proposed in [30] can be potentially used as a single-stage structured pruning method with the help of group sparsity, the empirical analysis of the paper only focused on unstructured pruning.
  Deep Representation Learning: Leveraging the power of deep neural networks, high-dimensional data can be encoded into a low-dimensional representation [18]. Such process is known as deep representation learning. A good representation extracts useful information by capturing the underlying explanatory factors of the observed input. The existence of noises, spurious patterns and complicated correlations among features make it a challenging problem. Previous work [44] show that an
  Autoencoder architecture can be used for denoising images. To disentangle correlated features, VAE [19] and β-VAE [17] encourages the posterior distribution over the generative factors q(z|x) to be closer to the isotropic Gaussian N (0, I) that explicitly promotes disentanglement of the latent generative factors. Driven by the idea of the autoencoder, the deep representation learning has been found remarkably successful in various of applications, e.g., graph representation learning [21, 6, 8, 31, 46, 34] and recommendation system [48].
  However, most existing works in representation learning treat the embedding dimension as a hyperparameter, which can be crucial and difficult to choose properly. Small embedding dimensions cannot sufficiently represent the important information in the data that leads to bad representations. On the other hand, large embedding dimensions allow some level of redundancy in the learned representations, even picking up spurious patterns that can degrade model performances.
sections.heading: 3 Proposed Approach
sections.text: Problem Statement: Let us represent the general architecture of a neural network with l layers as F(x) = Ψl ◦ Ψl−1 ◦ ...Ψ2 ◦ Ψ1(x), where x are the input features to the network, Ψi is the learnable function at layer i that generates hidden features hi = Ψi(hi−1), and f ◦ g represents the composite function f(g(·)). Let us also consider the “masked” version of this network where the features learned at every hidden layer, hi are multiplied with a vector of mask values gi using a mask function Mi(hi) = gi ⊙ hi. The masked neural network can then be described as Fmasked(x) = Ψl ◦Ml−1 ◦Ψl−1 ◦ ...M2 ◦Ψ2 ◦M1 ◦Ψ1(x). The goal of learning compact representations can be framed as minimizing the following learning objective:
  min θ Ex∼Dtr
  [ L (Fmasked(x)) + λ
  l − 1 l−1∑ i=1 ∥gi∥0
  ] , (1)
  where θ are the learnable parameters of the network, Dtr is the training set, L(.) is the loss function, and λ is the trade-off parameter for enforcing L0 sparsity in the mask values gi. Note that this formulation provides a unified view of both representation learning and structured network pruning, which differ in the choice of loss function (typically, reconstruction error for representation learning and cross-entropy for structured pruning), and the number of layers that are pruned (in representation learning, we only aim to prune the bottleneck layer while in structured pruning, we aim to prune all of the layers in the network, see Figure 2).
  DiscriminAtive Masking (DAM): Our proposed approach presents a novel perspective of structured pruning by introducing a new strategy of “ordering-and-masking.” To understand the basic idea of this strategy, let us assume that we are given the ideal “ordering” of neurons in a layer that reflect their importance (i.e., neurons with higher order number are more useful for the learning task). Given this ordering, we can perform structured pruning by incrementally “masking” out neurons with lower order (i.e., setting their mask values to 0) before going to neurons with higher order during the training process. In other words, we discriminatively (or preferentially) focus on pruning neurons lower in the order, while retaining and refining the features learned at neurons higher in the order. This approach for masking can be realized using a monotonic gate function G that maps neuron order to mask values, i.e., lower order neurons are assigned smaller mask values than higher order neurons. By dynamically regulating the gate function during the course of training, we can control the percentage of neurons pruned at every training epoch.
  Specifically, let us assign every neuron j at layer i with an “order number,” µij . Before we describe how we arrive at these neuron orders, we first present the formulation of our proposed gate function G, whose value at neuron j monotonically increases with the order number of the neuron, µij , as follows:
  gij = G(µij , αi, βi) = ReLU[tanh (αi(µij + βi))], (2) where αi is a constant scalar parameter (termed as the steepness parameter) that is not optimized during training, while βi is a ‘learnable’ scalar parameter (termed the offset parameter) that is the only parameter optimized during training to control (or regulate) the process of masking in DAM. Figures 2(c) and (d) show examples of gate function values (in red) at two different values of βi. Note that using a single learnable parameter βi to implement the masking process introduces significant simplicity in the design of the learnable mask, in contrast to SOTA methods in structured
  pruning [28, 43] where a different masking parameter is trained for every neuron in the network. We hypothesize this simplicity of the gate function, in conjunction with the preferential pruning property of DAM, to result in remarkable gains in learning compact representations, as evidenced later in our experiments. Further, note that while we can use other monotonically increasing function to implement our gates, we found ReLU-tanh to perform best in practice.
  Assigning Neuron Orders µij : We now discuss our approach for assigning neuron order numbers µij . Note that the order numbers only have to be assigned during the initialization step, and do not change during the training process. Further, note that every neuron is initialized with weights that are independent and identically distributed (because of the i.i.d nature of commonly used initialization methods in deep learning [41]). Hence, all neurons are symmetrically invariant to each other at initialization and any neuron has an equally good chance of refining themselves to capture useful features later in the training process as any other neuron. Hence, any random ordering of the neurons in the DAM approach would result in a similar pruning process (we empirically demonstrate the permutation invariance of µij in Section 5). For simplicity, we consider a trivial choice of µij = kj/ni, where k is a constant parameter that determines the domain size of µij , and ni is the total number of neurons at layer i.
  How Does Reducing βi Enforce L0 Sparsity? Let us first understand how mask values gi respond to the only learnable parameter βi. Figure 2 (c) shows an example of gi when βi = −2. At this stage, we can see that only a small number of neurons have zero gate values (white). As we reduce βi to −6 in Figure 2 (d), we can see that the gate moves towards the ’right’, resulting in more excessive pruning of neurons. In general, the number of non-zero values of the gate function (and hence, its L0 norm) is directly related to βi as follows:
  ∥gi∥0 = ni∑ j=1 1(gij > 0) = ni∑ j=1 (1− 1(µij ≤ βi)) = ⌈ni(1 + βi/k)⌉, for βi ≥ −k, (3)
  where 1 indicates the identity operator and ⌈·⌉ is the ceiling operator. In order to back-propagate gradients, we use a continuous approximation of Equation (3) (by dropping the ceiling operator) for regularization as ∥gi∥0 ≈ ni(1 + βi/k). Hence, minimizing L0 norm of gi is directly proportional to minimizing βi with a scaling factor of ni/k.
  Learning Objective of DAM: While minimizing βi with scaling factors proportional to ni can explicitly minimize L0 norm of the gate values, we found in practice that layers with smaller number of neurons ni still require adequate pruning of features for learning compact representations proportional to layers with large ni. Hence, in our learning objective, we drop the scaling factor and directly minimize the sum of βi across all layers in the following objective function.
  min θ,β Ex∼Dtr
  [ L (Fmasked(x)) + λ
  l − 1 l−1∑ i=1 βi
  ] , (4)
  where, β is the set of all βi’s in the network, {β1, β2, . . . , βl−1}. Theoretical Analysis of DAM: Analysis of the dynamics of the gate functions and additional specifications for choosing αi and the gate function are provided in Appendix A.
  Parameter Specification and Implementation Details: In all of our experiments, we used αi = 1 and k = 5. In our structured network pruning experiments, we used a cold-start of 20 epochs (i.e., the βi’s were frozen for the duration of cold-start), so as to allow the leftmost neurons to undergo
  some epochs of refinement before beginning the pruning process. We also set the initial value of βi to 1, which can also be thought of as another form of cold-starting (since pruning of a layer only starts when βi becomes less than zero).
sections.heading: 4 Experiments on Representation Learning Problems
sections.heading: 4.1 Results on Dimensionality Reduction (DR) Problems
sections.text: Problem Setup: We evaluate the effectiveness of DAM in recovering the embeddings of synthetically generated dimensionality reduction problems. The general form of the problem is expressed by (5). Suppose Ω ∈ Rr and all of its r dimensions are independent from each other. A d-dimensional data X can be expressed using a transformation Ψ : Rr −→ Rd defined on Hilbert space, and X = Ψ(Ω), d > r. Let M be the learnable mask, we formulated the dimensionality reduction problem as:
  min Θ
  { ∥FD ◦M ◦ FE(X)−X∥2F + λβ } , Θ = {θE , θD, β} (5)
  The FE ,M,FD are trained end-to-end using gradient descent. In our experiments, Ω is sampled from a isotropic normal distribution, i.e., Ω ∼ Nr(0, Ir). We let Ψ to use the same structure as the decoder FD with randomly generated parameters θΨ, i.e., Ψ = FD|θD=θΨ , so that the minimum number of dimensions needed in the encoded representation for the decoder to reconstruct X is r.
  Linear DR: We test DAM for removing linear correlations in a given rank-deficient matrix. In this case, Ψ is a matrix representing a linear projection from Rr to Rd. We use two real matrices as the encoder and decoder and train the DAM to find the full-rank representation of X .
  Nonlinear DR: To test the ability of DAM for disentangling nonlinearly correlated dimensions, we present two cases: (i) Ψ is a polynomial kernel of degree 2 (we use a QRes layer from [5]); (ii) Ψ is a nonlinear transformation expressed by a neural network. We use a deep neural network as the encoder with sufficient complexity.
  Observation: Figure 3 shows the convergence of bottleneck dimensions over the training epochs for different synthetic data with varying sizes of underlying factors (r). The curves and shades show the mean and standard deviation of five different runs. We can see that DAM consistently uncovers the exact dimension of the underlying factors r. We provides details of in-depth sensitivity analysis of DAM results based on its hyper-parameters (e.g., learning rate and λ) in Appendix C.
  Theoretical Analysis of DAM for linear DR: We theoretically show that DAM converges to the optimal solution for the linear DR case (see Appendix B for details).
sections.heading: 4.2 Recommendation System Results
sections.text: Problem Setup: We consider the state-of-the-art method for recommendation system problems (IGMC [48]), which transforms the rating matrix completion task to a link prediction problem, where the ratings are interpreted as links between users and items.
  In particular, IGMC generates 128-dimensional embeddings to represent the enclosing subgraph and further avoids overfitting using a dropout layer with a dropout rate of 0.5. We replaced the dropout
  layer with a DAM layer in the IGMC to reduce the dimension of the learned representations. We train our IGMC-DAM model for 100 epochs under the same training configurations.
  Observation: Table 1 shows that DAM successfully reduces the dimensions of the learned representations without any increase in errors, demonstrating that DAM is able to boost the performance of IGMC by learning compact representations over the target graph.
sections.heading: 4.3 Graph Representation Learning Results
sections.text: Problem Setup: Following the previous experiment, we further explore the effectiveness of DAM in learning compact graph representations. Our goal is to learn low-dimensional embeddings for each node in a graph that captures the structure of interaction among nodes. A simple graph autoencoder, e.g., GAE [21], uses a graph convolutional network (GCN) [20] as an encoder and an inner product as a decoder. The encoder calculates the embedding matrix Z from the node feature matrix X with the adjacency matrix A, and the decoder reconstructs a adjacency matrix Â such that: Â = σ(ZZ⊤), with Z = GCN(X,A). The reconstruction loss, BCE(A, Â) is backpropagated to train the GAE model. To reduce the dimension of the learned representation, we add a DAM layer in a GAE after the encoder (GAE-DAM).
  Observation: Figure 4 shows that DAM based GAE method is able to learn meaningful structural information in compact latent embeddings. DAM improves the link prediction performance of the simple GAE for Cora, CiteSeer and PubMed dataset by learning compact representations.
sections.heading: 4.4 Representation Learning Results for MNIST Dataset
sections.text: Problem Setup: We further evaluate DAM on the problem of dimensionality reduction using simple auto-encoders on the MNIST dataset, and also compare the effectiveness of DAM (that directly enforces L0 sparsity) with an ablation of our approach that instead minimizes the L1 norm (similar to Lasso). We vary the trade-off parameter λ to compare the sparsity of the learned representations between DAM and L1-norm based ablation method.
  Observation: Figure 5 shows that DAM is able to achieve lower reconstruction error as well as higher F1 scores over the same bottleneck dimensions than L1 based method. In Figure 5 (c), we can see that DAM shows a near-linear descending trend of bottleneck dimensions as λ increases (further evidence presented in Appendix E1). In contrast, the L1 based method has a saturating effect at large λ values since the weights come close to zero but require some thresholding to be pruned. This shows that DAM is amenable to learning highly compact representations in contrast to L1 based methods.
sections.heading: 5 Experiments on Structured Network Pruning
sections.text: Evaluation Setup Here we demonstrate the effectiveness of DAM on the task of structured network pruning. We compared the performance of our proposed DAM with two SOTA structured pruning algorithms, Network Slimming (Net-Slim) [28] and ChipNet [43]. Net-Slim is well-established SOTA that is widely used as a strong baseline, while ChipNet is the latest SOTA representing a very recent development in the field. ChipNet and Net-Slim are both pretrained for 200 epochs, while the DAM performs single-stage training and pruning in the same number of epochs. Since, our DAM approach does not require additional fine-tuning, we imposed a limited budget of 50 epochs for fine-tuning the SOTA methods. Note that ChipNet also has a pruning stage which involves 20 additional epochs. We evaluate the performance of these methods on the PreResnet-164 architecture on benchmark computer vision datasets, CIFAR-10 and CIFAR-100. Additional evaluation setup details are in Appendix E2.
  Performance and Running time comparison: Figures 6a and 6b demonstrate the performance of different network pruning methods for various pruning ratios. We observe that for both datasets, DAM is able to outperform Net-Slim by a significant margin specially when the models are sparse. Additionally, DAM is able to achieve similar or slightly better performance than ChipNet with no additional fine-tuning. We also compare the total running time for pruning across the different models (divided into the three categories) in Figure 6c. We can see that the training time of the three methods are almost the same, with ChipNet being slightly small. However, the running time for the pruning and fine-tuning stages for DAM are both zero. Net-Slim also does not involve additional pruning after training. However, ChipNet involves 20 epochs of pruning which is significant and is almost comparable to its pretraining time. Finally, comparing the total running time taken by each of the structured pruning methods, we observe that DAM is significantly faster than the current SOTA counterparts owing to its single-stage pruning approach.
  Stability Analysis: We further analyse the stability of our proposed DAM approach for network pruning through visualization of the training dynamics. We observe that the training cross-entropy (CE) loss and the validation CE loss are very similar to what we expect from a Preresnet model trained on benchmark vision datasets with learning rate changes at 100 and 150 epochs, respectively. We further notice that the total training loss for the DAM is also very stable and does not show any sudden variations, leading to steady convergence. Finally, from the convergence plot for E[βi] (green-dashed line), we can see that our pruning does not involve any sudden steps and happens in a continuous and steady rate throughout the training stage.
  The training dynamics of the ChipNet (blue shaded region in the middle panel) is exactly the same as training vanilla vision models since it does not involve sparsity constraints. However, as we move to the pruning stage, we notice a sharp rise in both training and validation losses (white region). This is due to the use of a different optimizer AdamW as opposed to SGD which was used in the training stage. Further, a very interesting phenomenon happens in the fine-tuning stage (red shaded region) where the training and validation losses increases for the first 25 epochs and then takes a sharp decrease once the learning rate is adjusted. This suggests that after the pruning stage, finetuning the model with large learning rate forces the model to make large updates in the initial epochs (akin to training from scratch), which further drives the model out of the local minima. Once the learning rate is adjusted, the model is able to slowly progress towards a new local optima and ultimately converges to the fine-tuned minima.
  Finally, both the training and fine-tuning dynamics of the Net-Slim (rightmost panel) appears to be highly unstable. This suggests that the pre-training with sparsity loss is not stable and can lead to large fluctuations in the training loss. Also, note that Net-Slim does not have a validation set as it instead uses the test set to choose the best model. We have refrained from using the best model by looking at the performance on the test set, and instead evaluate the model from the last epoch.
  Permutation Invariance: To test the effect of permuting the neuron order on the performance of DAM, Table 2 provides results across five randomly permuted values of µij across every neuron in the network, while using the same initialization as previous experiments. We can see that DAM is invariant to permutations of neuron order during initialization, validating our simple choice of assigning neuron ordering based on their index values.
  Additional Results: To illustrate generalizibility we provide additional network pruning results for more networks such as VGG-19, PreResnet-20, PreResNet-56 and PreResnet-110 on CIFAR datasets and LeNet-5 on MNIST in Appendix E3. To demonstrate decorrelated kernels obtained after pruning using DAMs we use CKA similarity [22]. The results can be found in Appendix E3. For extreme values of λ, we demonstrate that DAM is able to pruning entire bottleneck layers for ResNet architectures by visualizing the pruned architectures in Appendix E4.
sections.heading: 6 Discussion, Open Questions, and Future Works
sections.text: The effectiveness of DAM in learning compact representations provides a novel perspective into structured network pruning and raises several open questions about its relationships with other concepts in this area. In the following, we list some of these connections and describe the limitations and future extensions of our work.
  Connections with Lottery Ticket Hypothesis: Our work builds on the intuition of discriminative masking, i.e., we preferentially favor some neurons to be refined while favoring some other neurons to be pruned. This has interesting connections with the lottery ticket hypothesis (LTH) [10], where some subsets of initialized weights in the network are hypothesized to contain “winning tickets” that are uncovered after a few epochs of training. While there is some open discussion on whether the winning tickets exist even before initialization [11, 15], there are several studies supporting this hypothesis [38, 12]. Our results attempts to throw light on the question: “can we find winning tickets if we discriminatively search for it in a certain subregion of the network?” Further, we are able to show that the results of DAM are invariant to random permutation of the neuron indices at initialization, since all neurons receive identically distributed weights. Along these lines, we can also explore if there are certain ordering of neurons (e.g., in accordance with their likelihood of containing winning tickets revealed through LTH) that can perform better than random ordering in DAM.
  Connections with Dropout: Another interesting connection of DAM is with the popular regularization technique of Dropout [40] that randomly drops out some neurons during training to break “co-adaptation” patterns in the learned features so as to avoid overfitting. Essentially, by randomly dropping a neuron at some epoch of training, the other neurons are forced to pick up the learned features of the dropped neuron, thus resulting in the learning of robust features. This is similar to the “re-adaptation” of weights at the active neurons at some epoch of DAM training, while the mask value of the neurons in the transitioning zone (when 0 < gj < 1) are gradually dropped by the shifting of the gate function. This motivates us to postulate the weight re-adaptation hypothesis as a potential reason behind the effectiveness of DAM in learning compact representations, which requires further theoretical justifications.
  Budget-aware Variants of DAM: One promising future extension of our current DAM formulation would be to make it budget-aware, i.e., to stop the training process once we arrive at a target level of sparsity. Note that there is a direct correspondence between the value of the learnable gate offset parameter βi (that is minimized at every epoch) and the resulting L0 sparsity (see Equation 3). Hence, a simple budget-aware extension of our current DAM formulation would be to start from a large value of lambda (to provide sufficient scope for aggressive pruning) and keep monitoring the level of sparsity at every layer during the training process. As soon as the target sparsity mark is achieved, we can freeze βi’s at every layer thus essentially halting the pruning process (note that the network architecture becomes immutable if βi’s is fixed).
  DAM Variants for Pre-trained Networks: One of the fundamental assumptions of our current DAM formulation is that all neurons are symmetrically invariant to each at initialization, and hence we can simply use random neuron ordering for discriminative masking. While this assumption is valid for “training from scratch”, this may not hold for neurons in a pre-trained network. Future extensions of DAM can include advanced ways of ordering neurons that do not rely on the above assumption and hence can even be used with pre-trained networks. For example, we can order neurons based on the magnitudes of weights of every neuron or the value of mutual information between neuron activations and outputs. This would open novel avenues of research in structured network pruning by finding effective neuron orderings for different initializations of network weights.
sections.heading: 7 Acknowledgement
sections.text: This work was supported by the NSF Eager Grant #2026710.
title: Learning Compact Representations of Neural Networks using DiscriminAtive Masking (DAM)