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 work proposes an improved method to find a performative subnetwork within a randomly weighted neural network. In contrast to prior work, a fraction of unused weights are periodically reset. This is a novel idea with empirical improvements and theoretical justification.