Pruning Randomly Initialized Neural Networks with Iterative Randomization

1 Introduction . The lottery ticket hypothesis , which was originally proposed by Frankle and Carbin [ 5 ] , has been an important topic in the research of deep neural networks ( DNNs ) . The hypothesis claims that an over-parameterized DNN has a sparse subnetwork ( called a winning ticket ) that can achieve almost the same accuracy as the fully-trained entire network when trained independently . If the hypothesis holds for a given network , then we can reduce the computational cost by using the sparse subnetwork instead of the entire network while maintaining the accuracy [ 14 , 19 ] . In addition to the practical benefit , the hypothesis also suggests that the over-parametrization of DNNs is no longer necessary and their subnetworks alone are sufficient to achieve full accuracy . Ramanujan et al . [ 23 ] went one step further . They proposed and empirically demonstrated a conjecture related to the above hypothesis , called the strong lottery ticket hypothesis , which informally states that there exists a subnetwork in a randomly initialized neural network such that it already achieves almost the same accuracy as a fully trained network , without any optimization of the weights of the network . A remarkable consequence of this hypothesis is that neural networks could be trained by solving a discrete optimization problem . That is , we may train a randomly initialized neural network by finding an optimal subnetwork ( which we call weight-pruning optimization ) , instead of optimizing the network weights continuously ( which we call weight-optimization ) with stochastic gradient descent ( SGD ) . However , the weight-pruning optimization requires a problematic amount of parameters in the random network before pruning . Pensia et al . [ 22 ] theoretically showed that the required network width for the weight-pruning optimization needs to be logarithmically wider than the weight-optimization at ∗Corresponding author : daiki.chijiwa.mk @ hco.ntt.co.jp 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . least in the case of shallow networks . Therefore , the weight-pruning optimization requires more parameters , and thus more memory space , than the weight-optimization to achieve the same accuracy . In other words , under a given memory constraint , the weight-pruning optimization can have lower final accuracy than the weight-optimization in practice . In this paper , we propose a novel optimization method for neural networks called weight-pruning with iterative randomization ( IteRand ) , which extends the weight-pruning optimization to overcome the parameter inefficiency . The key idea is to virtually increase the network width by randomizing pruned weights at each iteration of the weight-pruning optimization , without any additional memory consumption . Indeed , we theoretically show that the required network width can be reduced by the randomizing operations . More precisely , our theoretical result indicates that , if the number of randomizing operations is large enough , we can reduce the required network width for weight-pruning to the same as that for a network fully trained by the weight-optimization up to constant factors , in contrast to the logarithmic factors of the previous results [ 22 , 21 ] . We also empirically demonstrate that , under a given amount of network parameters , IteRand boosts the accuracy of the weight-pruning optimization in multiple vision experiments . 2 Background . In this section , we review the prior works on pruning randomly initialized neural networks . Notation and setup . Let d , N ∈ N. Let f ( x ; θ ) be an l-layered ReLU neural network with an input x ∈ Rd and parameters θ = ( θi ) 1≤i≤n ∈ Rn , where each weight θi is randomly sampled from a distributionDparam over R. A subnetwork of f ( x ; θ ) is written as f ( x ; θ m ) where m ∈ { 0 , 1 } n is a binary mask and `` `` represents an element-wise multiplication . Ramanujan et al . [ 23 ] empirically observed that we can train the randomly initialized neural network f ( x ; θ ) by solving the following discrete optimization problem , which we call weight-pruning optimization : min m∈ { 0,1 } n E ( x , y ) ∼Dlabeled [ L ( f ( x ; θ m ) , y ) ] , ( 1 ) where Dlabeled is a distribution on a set of labeled data ( x , y ) and L is a loss function . To solve this optimization problem , Ramanujan et al . [ 23 ] proposed an optimization algorithm , called edge-popup ( Algorithm 1 ) . Algorithm 1 : Weight-pruning optimization by edge-popup [ 23 ] 1 Initialize θ ∼ Dnparam , s ∼ Dnscore ; // Dparam and Dscore are distributions over R 2 while k = 0 , · · · , N − 1 do 3 Sample a labeled data ( x , y ) ∼ Dlabeled ; 4 m , s← TrainMask ( θ , s ; ( x , y ) ) ; // optimize importance scores s and update m 5 end 6 return m , θ ; The TrainMask ( Algorithm 2 ) is the key process in Algorithm 1 . It has a latent variable s = ( si ) 1≤i≤n ∈ Rn , where each element si represents an importance score of the corresponding weight θi , and optimizes s instead of directly optimizing the discrete variable m. Given the score s , the corresponding m is computed by the function CalculateMask ( s ) , which returns m = ( mi ) 1≤i≤n defined as follows : mi = 1 if si is top 100 ( 1 − p ) % in { si } 1≤i≤n , otherwise mi = 0 , where p ∈ ( 0 , 1 ) is a hyperparameter representing a sparsity rate of the pruned network . In the line 3 of Algorithm 2 , SGDη , λ , µ ( s , g ) returns the updated value of s by stochastic gradient descent with a learning rate η , weight decay λ , momentum coefficient µ , and gradient vector g. On the theoretical side , Malach et al . [ 16 ] first provided a mathematical justification of the above empirical observation . They formulated it as an approximation theorem with some assumptions on the network width as follows . Theorem 2.1 ( informal statement of Theorem 2.1 in [ 16 ] ) Let ftarget ( x ) be an l-layered network with bounded weight matrices , and g ( x ) be a randomly initialized 2l-layered neural network . If the width of g ( x ) is larger than ftarget ( x ) by the factor of a polynomial term , then there probably exists a subnetwork of g ( x ) that approximates ftarget ( x ) . Algorithm 2 : Pseudo code of TrainMask 1 Input : θ , s ∈ Rn , ( x , y ) : a labeled data ; 2 m← CalculateMask ( s ) ; // calculate the mask m with the current scores 3 s← SGDη , λ , µ ( s , ∇s=mL ( f ( x ; θ s ) , y ) ) ; // update s by the gradient at s = m 4 m← CalculateMask ( s ) ; // calculate new mask m with the updated scores 5 return m , s ; By considering a well-trained network as ftarget , Theorem 2.1 indicates that pruning a sufficiently wide g ( x ) may reveal a subnetwork which achieves good test accuracy as ftarget , in principle . In the follow-up works [ 22 , 21 ] , the assumption on the network width was improved by reducing the factor of the required width to a logarithmic term . However , Pensia et al . [ 22 ] showed that the logarithmic order is unavoidable at least in the case of l = 1 . While their results imply the optimality of the logarithmic bound , it also means that we can not further relax the assumption on the network width as long as we work in the same setting . This indicates a limitation of the weight-pruning optimization , i.e . the weight-pruning optimization can train only less expressive models than ones trained with the conventional weight-optimization like SGD , under a given amount of memory or network parameters . 3 Method . In this section , we present a novel method called weight-pruning with iterative randomization ( IteRand ) for randomly initialized neural networks . As discussed in Section 2 , although the original weight-pruning optimization ( Algorithm 1 ) can achieve good accuracy , it still has a limitation in the expressive power under a fixed amount of memory or network parameters . Our method is designed to overcome this limitation . The main idea is to randomize pruned weights at each iteration of the weight-pruning optimization . As we prove in Section 4 , this reduces the required size of an entire network to be pruned . We use the same notation and setup as Section 2 . In addition , we assume that each weight θi of the network f ( x ; θ ) can be re-sampled from Dparam at each iteration of the weight-pruning optimization . 3.1 Algorithm . Algorithm 3 describes our proposed method , IteRand , which extends Algorithm 1 . The differences from Algorithm 1 are lines 5-7 in Algorithm 3 . IteRand has a hyperparameter Kper ∈ N≥1 ( line 5 ) . At the k-th iteration , whenever k can be divided by Kper , pruned weights are randomized by Randomize ( θ , m ) function ( line 6 ) . There are multiple possible designs for Randomize ( θ , m ) , which will be discussed in the next subsection . Algorithm 3 : Weight-pruning optimization with iterative randomization ( IteRand ) 1 Initialize θ ∼ Dparam , s ∼ Dscore ; 2 while k = 0 , · · · , N − 1 do 3 Sample a labeled data ( x , y ) ∼ Dlabeled ; 4 m , s← TrainMask ( θ , s ; ( x , y ) ) ; 5 if k + 1 can be divided by Kper then // this if-block is newly added to Algorithm 1 6 θ ← Randomize ( θ , m ) ; // randomize a subset of pruned weights 7 end 8 end 9 return m , θ ; Note that Kper controls how frequently the algorithm randomizes the pruned weights . Indeed the total number of the randomizing operations is bN/Kperc . If Kper is too small , the algorithm is likely to be unstable because it may randomize even the important weights before their scores are well-optimized , and also the overhead of the randomizing operations can not be ignored . In contrast , if Kper is too large , the algorithm becomes almost same as the original weight-pruning Algorithm 1 , and thus the effect of the randomization disappears . We fix Kper = 300 on CIFAR-10 ( about 1 epoch ) and Kper = 1000 on ImageNet ( about 1/10 epochs ) in our experiments ( Section 5 ) . 3.2 Designs of Randomize ( θ , m ) . Here , we discuss how to define Randomize ( θ , m ) function . There are several possible ways to randomize a subset of the parameters θ . Naive randomization . For any distribution Dparam , a naive definition of the randomization function ( which we call naive randomization ) can be given as follows . Randomize ( θ , m ) i : = { θi , ( if mi = 1 ) θ̃i , ( otherwise ) where we denote the i-th component of Randomize ( θ , m ) ∈ Rn as Randomize ( θ , m ) i , and each θ̃i ∈ R is a random variable with the distribution Dparam . Also this can be written in another form as Randomize ( θ , m ) : = θ m+ θ̃ ( 1−m ) , ( 2 ) where θ̃ = ( θ̃i ) 1≤i≤n ∈ Rn is a random variable with the distribution Dnparam . Partial randomization . The naive randomization ( Eq . ( 2 ) ) is likely to be unstable because it entirely replaces all pruned weights with random values every Kper iteration . To increase the stability , we modify the definition of the naive randomization as it replaces a randomly chosen subset of the pruned weights as follows ( which we call partial randomization ) : Randomize ( θ , m ) : = θ m+ ( θ ( 1− br ) + θ̃ br ) ( 1−m ) , ( 3 ) where θ̃ ∈ Rn is the same as in Eq . ( 2 ) , r ∈ [ 0 , 1 ] is a hyperparameter and br = ( br , i ) 1≤i≤n ∈ { 0 , 1 } n is a binary vector whose each element is sampled from the Bernoulli distribution Bernoulli ( r ) , i.e . br , i = 1 with probability r and br , i = 0 with probability 1− r. The partial randomization replaces randomly chosen 100r % of all pruned weights with random values . Note that , when r = 1 , the partial randomization is equivalent to the naive randomization ( Eq . ( 2 ) ) . In contrast , when r = 0 , it never randomizes any weights and thus is equivalent to Algorithm 1 . In Figure 1 , we observe that r = 0.1 works well with various network architectures on CIFAR-10 , where we use the Kaiming uniform distribution ( whose definition will be given in Section 5 ) for Dparam and Kper = 300 .

Summary: This paper proposes a memory-efficient approach for \textit{weight-pruning optimization}, which suffers from a much larger memory cost than standard \textit{weight optimization}.   The proposed approach, weight-pruning with iterative randomization (IteRand), decreases the required network width by introducing pruned weights randomization at each iteration. The author theoretically shows that, with sufficient randomizing operations, the width can be reduced to constant factors of that of a standard network trained by weight-optimization.   The empirical experiments demonstrate IteRand achieves better performance than existing weight-pruning optimization methods under the same memory budget.