Training Neural Networks with Fixed Sparse Masks

1 Introduction . Stochastic gradient descent ( SGD ) is a vital component of the modern pipeline for training deep neural networks . Along with the back-propagation algorithm , gradient descent allows for the efficient minimization of a loss function by gradually updating a model ’ s parameters . SGD minimizes the loss over a small random subset of the dataset at each training iteration , which allows training over large datasets . In practice , minimizing a large neural network ’ s training loss using SGD often results in models that generalize well to new data [ 3 , 17 , 27 ] , making SGD an invaluable tool . While effective , standard SGD requires that all model parameters are updated at every iteration of training . As a result , communicating changes to the model requires communicating the updated value of every parameter . Since modern neural networks often have millions or billions of parameters [ 9 , 7 , 40 ] , this communication can become excessively expensive . A concrete example of the negative impacts of these costs arises in the setting of transfer learning . In transfer learning , a model ’ s parameters are initialized from an existing pre-trained model before being fine-tuned ( i.e . trained ) on a task of interest . Pre-trained models can be fine-tuned a huge number of times – for example , the Hugging Face model repository3 has thousands of fine-tuned variants of the BERT model [ 12 ] . Each of these fine-tuned variants requires a unique copy of the model ’ s parameters , each of which takes up around 500MB of disk space . Relatedly , in distributed training [ 11 ] and federated learning [ 35 ] , workers compute updates for a centralized model in parallel on different subsets of data . After a certain number of updates , the workers each communicate the newly-computed parameter values back to the centralized model . The communication step can cause a significant amount of overhead ( particularly when the model is large ) since the workers must communicate the updated values of all parameters when using standard SGD . ∗Equal contribution . 2Code for our work can be found at https : //github.com/varunnair18/FISH . 3https : //huggingface.co/models 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . These issues could be mitigated if it was possible to only update a few parameters during training while still maintaining performance close to that of training all parameters . This has led to various work on parameter-efficient training of neural networks . For example , Adapters [ 19 , 41 , 5 ] introduce additional parameters into a pre-trained model in the form of small task-specific modules that are fine-tuned while the rest of the model ’ s parameters are kept fixed . Diff Pruning [ 16 ] and BitFit [ 6 ] demonstrate that it is possible to fine-tune a model while only updating a small subset of the existing parameters . In distributed and federated learning settings , Aji and Heafield [ 2 ] and Konečnỳ et al . [ 23 ] have shown that it is possible for each worker to only update a sparse subset of a model ’ s parameters , thereby reducing communication costs . Existing methods for training with sparse updates typically work in one of three ways : they either add parameters to the model ( Adapters ) , choose a hand-defined and heuristically-motivated subset of parameters ( BitFit ) , or allow the subset of parameters to change over the course of training ( Diff Pruning and methods in distributed and federated training ) . In this paper , we argue for pre-computing a sparse subset of existing parameters to update and keeping the subset fixed over many iterations of training . This approach yields various benefits : First , by updating a subset of existing parameters instead of adding parameters ( as is done in Adapters ) , we avoid any increase in the total size of the model . Second , by avoiding hand-defining the mask , we can ensure that our procedure is modelagnostic . Third , by pre-computing a mask , we avoid the computational and memory overhead that are apparent when updating the mask over the course of training . It also allows workers in the distributed training setup to operate on strictly complementary subsets of parameters . Finally , keeping the mask fixed over many iterations , we can ensure that only a specific fixed number of parameters are updated . We are not aware of any existing techniques that satisfy these desiderata . Motivated by these benefits , we introduce a new method for pre-computing fixed sparse masks . Our approach first estimates the importance of each parameter using an empirical approximation of its Fisher information . Then , we construct the mask by choosing the k parameters with the largest Fisher information . The resulting mask , which we deem a “ FISH ( Fisher-Induced Sparse uncHanging ) mask ” , can be re-used for many subsequent training iterations . We demonstrate the effectiveness of using a FISH Mask in a wide variety of settings , including parameter-efficient transfer learning , distributed training with long delays across workers , and reducing checkpoint size . Broadly speaking , FISH Mask training can dramatically reduce storage and communication requirements , while sacrificing minimal performance compared to standard gradient descent and outperforming relevant prior methods for training with sparse updates . 2 The Fisher-Induced Sparse uncHanging ( FISH ) Mask . The main contribution of this paper is a method for pre-computing a sparse subset of a model ’ s parameters to update over many subsequent training iterations . To construct such a subset , we use an approximation of each parameter ’ s Fisher information as a signal of how important the parameter is for a given task . We refer to the resulting mask ( i.e . binary array indicating which parameters are included in the subset ) as a FISH ( Fisher-Induced Sparse uncHanging ) mask . In this section , we provide the necessary background and detail the steps necessary for computing a FISH Mask . This process is diagrammed in fig . 1 . 2.1 Fisher Information . Our goal is to select the subset of parameters that are ( in some sense ) the most important to update . One way to measure a parameter ’ s importance is to consider how much changing the parameter will change the model ’ s output . We denote pθ ( y|x ) as the output distribution over y produced by a model with parameter vector θ ∈ R|θ| given input x . One way to measure how much a change in parameters would change a model ’ s prediction would be to compute DKL ( pθ ( y|x ) || pθ+δ ( y|x ) ) , where δ ∈ R|θ| is a small perturbation . It can be shown [ 34 , 37 ] that as δ → 0 , to second order , ExDKL ( pθ ( y|x ) || pθ+δ ( y|x ) ) = δTFθδ +O ( δ3 ) ( 1 ) where Fθ ∈ R|θ|×|θ| is the Fisher information matrix [ 14 , 4 ] , defined as Fθ = Ex∼p ( x ) [ Ey∼pθ ( y|x ) ∇θ log pθ ( y|x ) ∇θ log pθ ( y|x ) T ] ( 2 ) Given this relation , it can be seen that the Fisher information matrix is closely connected to how much each parameter affects the model ’ s predictions . Indeed , this has led the Fisher information matrix to be widely used in modern machine learning , e.g . as a measure of parameter importance [ 21 ] , as a preconditioner in gradient descent [ 4 , 37 , 34 ] , as a way to measure the amount of “ information ” in each parameter of a neural network [ 1 ] , or as a way to decide which parameters to prune when performing model compression [ 43 , 10 , 47 ] . When applied to large neural networks , the |θ| × |θ| size of the Fisher information matrix makes it intractable to compute . Prior work therefore frequently approximates Fθ as a diagonal matrix , or equivalently , as a vector in R|θ| . Separately , when training machine learning models we seldom have the ability to draw samples x ∼ p ( x ) ; instead , we are given a finite training set of such samples . Furthermore , it is rarely necessary to compute the expectation over x in eq . ( 2 ) over the full training set ; instead , it can often be well-approximated over N samples x1 , . . . , xN . These constraints result in the following common approximation : F̂θ = 1 N N∑ i=1 Ey∼pθ ( y|xi ) ( ∇θ log pθ ( y|xi ) ) 2 ( 3 ) where F̂θ ∈ R|θ| . This approximation also has an intuitive interpretation : A given entry in F̂θ relates to the average of the square gradient of the model ’ s output with respect to a given parameter . If a given parameter heavily affects the model ’ s output , then its corresponding entry in F̂θ will be large , so we can reasonably treat F̂θ as an approximation of the importance of each parameter . Note that both eq . ( 2 ) and eq . ( 3 ) include an expectation over y ∼ pθ ( y|x ) . When the number of classes is small , this expectation can be computed exactly . For tasks with many possible classes , it is common to approximate the expectation with a few samples from pθ ( y|x ) . In supervised learning settings , we have access to the ground-truth label yi for each sample xi in our training set . This leads to the possibility of replacing Ey∼pθ ( y|xi ) ( ∇θ log pθ ( y|xi ) ) 2 in eq . ( 3 ) with ( ∇θ log pθ ( yi|xi ) ) 2 . Performing this approximation is referred to as the “ empirical Fisher ” . It has been shown that using the empirical Fisher can lead to degenerate behavior when used as a preconditioner in an optimizer [ 34 , 26 ] . Since our use of the Fisher information is largely based on a heuristically-motivated notion of parameter importance , we experimented with both the empirical and standard ( eq . ( 3 ) ) Fisher approximations and found that they produced similar performance . Furthermore , the Empirical Fisher is faster to compute than the standard Fisher as long as more than one sample is used to approximate the expectation Ey∼pθ ( y|xi ) . We discuss this further in section 4.4.1 . 2.2 Computing Fixed Sparse Masks . Recall that our goal is to select a subset of parameters ( or , equivalently , a sparse mask over parameters ) to update over many iterations of training while keeping the remainder of the parameters fixed . Having established the Fisher information as a useful tool for estimating the importance of a given parameter , we therefore first compute the approximate Fisher information ( as described in the previous section ) for all of a model ’ s parameters . Then , to construct the FISH Mask , we simply choose the k parameters with the largest Fisher information , where k is set according to the desired mask sparsity level.4 Specifically , a FISH Mask comprises the parameters { θi | F̂θi ≥ sort ( F̂θ ) k } . Computing the FISH Mask is cheap because F̂θ can be computed efficiently using backpropagation , and ( as we will show in section 4.4.2 ) we can obtain a reliable mask for relatively small values of N . Further , the fact that we re-use the mask for many iterations prevents us from having to compute F̂θ frequently . As we will show in section 4 , we find that this simple procedure is sufficient to produce a mask that can be reused for many iterations ( over 100,000 iterations in some cases ) in a wide variety of settings without sacrificing substantial performance compared to standard gradient-based training . Note that in some applications of transfer learning , a new linear classifier layer must be added to the model to make it applicable to the downstream task . Since the FISH Mask depends on pθ ( y|x ) and is computed before training begins , this means that we must compute the FISH Mask using the randomly-initialized classifier before any training has begun . We find that computing the Fisher information through the randomly-initialized classifier layer still provides a good signal of parameter importance . When applying FISH Mask in transfer learning settings where a new classifier layer is added , we always include the parameters of the classifier in the mask .

Summary: In this paper, the authors propose a scheme to select and update only part of the weights of a neural network during training. The selection of the weights to update is done using an approximation of the Fisher information, and the mask is computed by keeping only the weights associated with the larger fisher information values. The authors then propose an extensive empirical evaluation of their method for transfer learning, distributed training and efficient checkpointing. Finally, the authors perform an ablation study to show the effect of their design choices.