Learning Compact Representations of Neural Networks using DiscriminAtive Masking (DAM)

1 Introduction . 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 . 2 Related Work . 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 can not 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 .

Summary: 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.