Winning the Lottery with Continuous Sparsification

1 INTRODUCTION . Although deep neural networks have become ubiquitous in fields such as computer vision and natural language processing , extreme overparameterization is typically required to achieve state-ofthe-art results ( Xie et al. , 2017 ; Devlin et al. , 2018 ) , causing higher training costs and hindering applications where memory or inference time are constrained . Recent theoretical work suggest that overparameterization plays a key role in both the capacity and generalization of a network ( Neyshabur et al. , 2018 ) , and in training dynamics ( Allen-Zhu et al. , 2019 ) . However , it remains unclear whether overparameterization is truly necessary to train networks to state-of-the-art performance . At the same time , empirical approaches have been successful in finding less overparameterized neural networks , either by reducing the network after training ( Han et al. , 2015 ; 2016 ) or through more efficient architectures that can be trained from scratch ( Iandola et al. , 2016 ) . Recently , the combination of these two approaches lead to new methods which discover efficient architectures through optimization instead of design ( Liu et al. , 2019 ; Savarese & Maire , 2019 ) . Nonetheless , parameter efficiency is typically maximized by pruning an already trained network . The fact that pruned networks are hard to train from scratch ( Han et al. , 2015 ; 2016 ) suggests that , while overparameterization is not necessary for a model ’ s capacity , it might be required for successful network training . Recently , this idea has been put into question by Frankle & Carbin ( 2019 ) , where heavily pruned networks are trained faster than their original counterparts , often yielding superior performance . A key finding is that the same parameter initialization should be used when re-training the pruned network . A winning ticket , defined by a sub-network and a setting of randomly-initialized parameters , is quickly trainable and has already found applications in , for example , transfer learning ( Morcos et al. , 2019 ; Mehta , 2019 ; Soelen & Sheppard , 2019 ) , making the search for winning tickets a problem of independent interest . Currently , the standard algorithm to find winning tickets is Iterative Magnitude Pruning ( IMP ) ( Frankle & Carbin , 2019 ; Frankle et al. , 2019 ) , which consists of a repeating a 2-stage procedure that alternates between parameter optimization and pruning . As a result , IMP relies on a sensible choice for pruning strategy , and is time-consuming : finding a winning ticket with 1 % of the original parameters in a 6-layer CNN requires over 20 rounds of training followed by pruning , totalling over 1000 epochs ( Frankle & Carbin , 2019 ) . Choosing a parameter ’ s magnitude as pruning criterion has also shown to be sub-optimal in some settings ( Zhou et al. , 2019 ) , leading to the question of whether better winning tickets can be found by different pruning methods . Moreover , at each iteration , IMP resets the parameters of the network back to initialization , hence considerable time is spent on re-training similar networks with different sparsities . With the goal of speeding up the search for winning tickets in deep neural networks , we design a novel method , Continuous Sparsification , which continuously removes weights from a network during training , instead of following a strategy to prune parameters at discrete time intervals . Unlike IMP , our method approaches the search for sparse networks as a ` 0-regularized optimization problem ( Louizos et al. , 2017 ) , resulting in a method that can be fully described in the optimization framework . To approximate ` 0-regularization , we propose a smooth re-parameterization , allowing for the subnetwork ’ s structure to be directly learned with gradient-based methods . Unlike previous works , our re-parameterization is deterministic , proving more convenient for the tasks of pruning and ticket search , while also yielding faster training times . Experimentally , our method offers superior performance when pruning VGG to extreme regimes , and is capable of finding winning tickets in Residual Networks trained on CIFAR-10 at a fraction of time taken by Iterative Magnitude Pruning . In particular , Continuous Sparsification successfully finds tickets in under 5 iterations , compared to 20 iterations required by Iterative Magnitude Pruning in the same setting . To further speed up the search for sub-networks , our method abdicates parameter rewinding , a key ingredient of Iterative Magnitude Pruning . By showing superior results without rewinding , our experiments offer insights on how ticket search should be performed . 2 RELATED WORK . 2.1 LOTTERY TICKET HYPOTHESIS . The Lottery Ticket Hypothesis ( Frankle & Carbin , 2019 ) states that for a network f ( x ; w ) , w ∈ Rd , and randomly-initialized parameters w0 ∼ D , there exists a sparse sub-network , defined by a configuration m ∈ { 0 , 1 } d , ‖m‖0 d , that , when trained from scratch , achieves higher performance than f ( x ; w ) while requiring fewer training iterations . The authors support this conjecture experimentally , showing that such sub-networks indeed exist : in particular , they can be discovered by repeatedly training , pruning , and re-initializing the network , through a procedure named Iterative Magnitude Pruning ( IMP ; Algorithm 1 ) ( Frankle et al. , 2019 ) . More specifically , IMP alternates between : ( 1 ) training the weights w of a network , ( 2 ) removing a fixed fraction of the weights with the smallest magnitude ( pruning ) , and ( 3 ) rewinding : setting the remaining weights back to their original initialization w0 . The sub-networks found by IMP , which indeed train faster and outperform their original , dense networks , are called winning tickets , and can generalize across datasets ( Mehta , 2019 ; Soelen & Sheppard , 2019 ) and training methods ( Morcos et al. , 2019 ) . In this sense , IMP can be a promising tool in applications that involve knowledge transfer , such as transfer or meta learning . Zhou et al . ( 2019 ) perform extensive experiments to re-evaluate and better understand the Lottery Ticket Hypothesis . Relevant to this work is the fact that the authors propose a method to learn the binary maskm in an end-to-end manner through SGD , instead of relying on magnitude-based pruning . The authors show that learning only the binary mask and not the weights is sufficient to achieve competitive performance , confirming that the learned masks are highly dependent on the initialized values w0 , and are also capable of encoding substantial information about a problem ’ s solution . 2.2 SPARSE NETWORKS . The core aspect of searching for a winning ticket is finding a sparse sub-network that attains high performance relative to its dense counterpart . One way to achieve this is through pruning methods ( LeCun et al. , 1990 ) , which follow a strategy to remove weights from a trained network while minimizing negative impacts on its performance . In Han et al . ( 2015 ) , a network is iteratively trained and pruned using parameter magnitudes as criterion : this iterative , two-stage algorithm is shown to outperform “ one-shot pruning ” : training and pruning the network only once . Other methods attempt to approximate ` 0 regularization on the weights of a network , yielding onestage procedures that can be fully described in the optimization framework . In order to find a sparse Algorithm 1 Iterative Magnitude Pruning ( Frankle et al. , 2019 ) 1 : Initialize w ← w0 ∼ D and m← ~1d 2 : Minimize L ( f ( x ; m w ) ) until wT is produced 3 : Set mi = 0 for the active weights with smallest magnitudes ( |wT , i| ≤ τ and mi = 1 ) 4 : If satisfied , output ticket f ( x ; m wk ) 5 : Otherwise , set w ← wk and go back to step 2 Algorithm 2 Iterative Stochastic Sparsification ( inspired by Zhou et al . ( 2019 ) ) 1 : Initialize w ← w0 ∼ D , s← ~s0 2 : Minimize Em∼Ber ( σ ( s ) ) [ L ( f ( x ; m w ) ) ] + λ ‖σ ( s ) ‖1 until wT and sT are produced 3 : If satisfied , output ticket f ( x ; m wk ) , m ∼ Ber ( σ ( sT ) ) 4 : Otherwise , set w ← wk , si ← −∞ for si , T < si,0 , and go back to step 2 setting m ∈ { 0 , 1 } d of a network f ( x ; m w ) , Srinivas et al . ( 2016 ) and Louizos et al . ( 2017 ) use a stochastic re-parameterization m ∼ Bernoulli ( g ( s ) ) with s ∈ Rd and g : R → [ 0 , 1 ] applied element-wise . First-order methods , coupled with gradient estimators , are then used to train both w and s to minimize the expected loss . This approach performs continuous parameter removal during training in an automatic fashion : any component si of s that assumes a value during training where g ( si ) = 0 effectively removes wi from the network . Moreover , approximating ` 0 regularization has the advantage of not requiring a pruning strategy , which might be arbitrarily complex . 3 METHOD . Designing a method to quickly find winning tickets requires an efficient way to sparsify networks : ideally , sparsification should be done as early as possible in training , and the number of removed parameters should be maximized without harming the model ’ s performance . In other words , sparsification must be continuously maximized following a trade-off with the performance of the network . This goal is not met by Iterative Magnitude Pruning : sparsification is done at discrete time steps , only after fully training the network , and optimal pruning rates likely depend on the model ’ s performance and current sparsity : factors which are typically not accounted for – note that these are inherent characteristics of magnitude-based pruning . In light of this , we turn to ` 0-regularization methods for learning sparse networks , which consist of optimizing a clear trade-off between sparsity and performance . As we will see , performing sparsification continuously is not only straightforward , but done automatically by the optimizer . 3.1 CONTINUOUS SPARSIFICATION BY LEARNING DETERMINISTIC MASKS . We first frame the search for sparse networks as a loss minimization problem with ` 0 regularization ( Louizos et al. , 2017 ; Srinivas et al. , 2016 ) : min w∈Rd L ( f ( x ; w ) ) + λ · ‖w‖0 ( 1 ) where λ ≥ 0 controls the sparsity of the solution , and , with a slight abuse of notation , L ( f ( x ; w ) ) denotes the loss incurred by the network f ( x ; w ) ( e.g. , the cross-entropy loss over a training set ) . As ` 0 regularization is typically intractable , we re-state the above minimization problem as : min w∈Rd m∈ { 0,1 } d L ( f ( x ; m w ) ) + λ · ‖m‖1 ( 2 ) which uses the fact that , for m ∈ { 0 , 1 } d , ‖m‖0 = ‖m‖1 . The ` 1 term can be minimized with subgradient descent , however the m ∈ { 0 , 1 } d constraint makes the above problem combinatorial and poorly suited for local search methods like SGD . Algorithm 3 Continuous Sparsification 1 : Initialize w ← w0 ∼ D , s← s0 , β ← β0 2 : Minimize L ( f ( x ; σ ( β · s ) w ) ) + λ ‖σ ( βs ) ‖1 while increasing β , producing wT , sT , and βT 3 : If satisfied , output ticket f ( x ; b ( sT ) wk ) 4 : Otherwise , set s← min ( βT · sT , s0 ) , β ← β0 , ( optionally , w ← w0 ) , and go back to step 2 We can avoid the binary constraint m ∈ { 0 , 1 } d by re-parameterizing m as a function of a newlyintroduced variable s ∈ Rd . For example , Louizos et al . ( 2017 ) propose a stochastic mapping s 7→ m and use gradient methods to minimize the expected total loss , while using estimators for the gradients of s ( since m is still binary ) . Having a stochastic mask ( or , equivalently , a distribution over subnetworks ) poses an immediate challenge for the task of finding tickets , as it is not clear which ticket should be chosen once a distribution overm is learned . Moreover , relying on gradient estimators often causes gradients to have high variance , requiring longer training to reach optimality . Alternatively , we consider a deterministic parameterization m = b ( s ) , where s ∈ Rd6=0 and b : R 6=0 → { 0 , 1 } is applied element-wise : b ( z ) = { 1 , if z > 0 0 , if z < 0 ( 3 ) Applying this re-parameterization to Equation 2 yields : min w∈Rd s∈Rd6=0 L ( f ( x ; b ( s ) w ) ) + λ · ‖b ( s ) ‖1 ( 4 ) Clearly , the above problem is again intractable , as it is still equivalent to the original ` 0 problem in Equation 1 . More specifically , the step function b ( z ) is non-convex , and having zero gradients make gradient-based optimization ineffective . Instead , we consider the following smooth relaxation of b ( · ) : m : = σ ( β · s ) ( 5 ) where β ∈ R > 0 , and σ is the sigmoid function σ ( z ) = 11+e−z , applied element-wise . By controlling β , which acts as a temperature parameter , we effectively interpolate between σ ( s ) , a smooth function well-suited for SGD , and limβ→∞ σ ( β · s ) = b ( z ) , our original goal , which brings computational hardness to the problem . Figure 1 illustrates this behavior . Note that , if L ( f ( x ; w ) ) is continuous in w , then : min w∈Rd s∈Rd6=0 lim β→∞ L ( f ( x ; σ ( βs ) w ) ) + λ · ‖σ ( βs ) ‖1 = min w∈Rd s∈Rd6=0 L ( f ( x ; b ( s ) w ) ) + λ · ‖b ( s ) ‖1 ( 6 ) Although gradient methods will become ineffective as β →∞ due to vanishing gradients of s , we can increase β while optimizing s and w with gradient descent . That is , our loss at each iteration will be a function of β as follows : Lβ ( w , s ) = L ( f ( x ; σ ( βs ) w ) ) + λ · ‖σ ( β · s ) ‖1 ( 7 ) How does the soft mask m = σ ( β · s ) behave as we minimize Lβ ( w , s ) while increasing β ? As β →∞ , every negative component of s will be mapped to 0 , effectively removing its correspondent weight parameter from the network . While analytically the weights will never truly be zeroed-out , limited numerical precision has the fortunate side-effect of causing actual sparsification to the network during training , as long as β is increased to a large enough value . In a nutshell , we learn sparse networks by minimizing Lβ ( w , s ) for T parameter updates with gradient descent while jointly annealing β : producing wT , sT and βT , which is ideally large enough such that , numerically 1 , σ ( βT · sT ) = b ( sT ) . In case m is truly required to be binary ( as in the task of finding tickets ) , the dependence on numerical imprecision can be avoided by directly outputting m = b ( sT ) at the end of training . Finally , note that minimizing Lβ while increasing β is not generally equivalent to minimizing the original ` 0-regularized problem . Informally , the former aims to solve limβ→∞minw , s Lβ ( w , s ) , while the ` 0 problem is minw , s limβ→∞ Lβ ( w , s ) .

Summary: This paper proposes a novel objective function that can be used to jointly optimize a classification objective while at the same time encourage sparsification in a network. The lottery ticket hypothesis and associated work shows that the iterative pruning of a network can lead to a sparse network that performs with high accuracy. On the other hand, the work of Zhou et al. shows that sparse masks (dubbed "supermasks") may be learned without training the parameters of the network. In a sense, this paper tries to combine these ideas by simultaneously training a network while also optimizing the mask.