[Paper Review] Training Neural Networks with Fixed Sparse Masks

Yi-Lin Sung, Varun Nair+1 more|arXiv (Cornell University)|Nov 18, 2021

Domain Adaptation and Few-Shot Learning48 references4 citations

TL;DR

This paper proposes FISH Mask, a method to pre-compute a fixed sparse mask selecting the k most important parameters—based on empirical Fisher information—for gradient updates during training. By updating only these high-importance parameters across many iterations, the approach reduces memory and communication costs in transfer learning and distributed training while matching or exceeding standard training performance with minimal accuracy drop at 10% sparsity.

ABSTRACT

During typical gradient-based training of deep neural networks, all of the model's parameters are updated at each iteration. Recent work has shown that it is possible to update only a small subset of the model's parameters during training, which can alleviate storage and communication requirements. In this paper, we show that it is possible to induce a fixed sparse mask on the model's parameters that selects a subset to update over many iterations. Our method constructs the mask out of the $k$ parameters with the largest Fisher information as a simple approximation as to which parameters are most important for the task at hand. In experiments on parameter-efficient transfer learning and distributed training, we show that our approach matches or exceeds the performance of other methods for training with sparse updates while being more efficient in terms of memory usage and communication costs. We release our code publicly to promote further applications of our approach.

Motivation & Objective

• Address high storage and communication costs in large-scale neural network training, especially in transfer learning and distributed/federated settings.
• Reduce model update overhead by updating only a small, fixed subset of parameters instead of all parameters at each iteration.
• Develop a model-agnostic, efficient method for sparse parameter updates that avoids adding new parameters or dynamically changing the update mask.
• Enable practical deployment in settings with limited bandwidth or storage, such as model checkpointing and federated learning.
• Demonstrate that fixed sparse masks based on parameter importance (via Fisher information) can maintain high performance with minimal degradation.

Proposed method

• Estimate parameter importance using empirical Fisher information, computed as the variance of per-sample gradients over a small data batch.
• Construct a fixed sparse mask by selecting the k parameters with the highest Fisher information values, ensuring only these are updated during training.
• Use a single forward and backward pass over a small sample set (N=1024, or N=256 in distributed settings) to compute the Fisher approximation efficiently.
• Apply standard stochastic gradient descent (SGD) but restrict weight updates to only the parameters selected by the pre-computed FISH mask.
• Maintain the same mask across all training iterations, avoiding re-computation or dynamic reconfiguration.
• Ensure compatibility with existing training pipelines by applying the mask as a binary mask during the optimizer step, without modifying model architecture.

Experimental results

Research questions

• RQ1 Can a fixed sparse mask based on parameter importance maintain high model performance while reducing communication and storage costs in training?
• RQ2 How does using Fisher information as a proxy for parameter importance compare to random or heuristic selection in sparse training?
• RQ3 What is the impact of mask sparsity and update frequency on model accuracy in parameter-efficient transfer learning?
• RQ4 How efficient is the Fisher computation in terms of data and computational resources, and how many samples are needed for stable performance?
• RQ5 Can a fixed FISH mask be used to identify sparse subnetworks for lottery ticket-style training, as suggested by the Lottery Ticket Hypothesis?

Key findings

• At 10% sparsity, FISH Mask training achieves 93.9% accuracy on the GLUE benchmark, matching standard training performance.
• Updating the FISH mask every epoch maintains near-identical performance to full training at 10% sparsity, reducing storage needs by a factor of 5.
• The performance of FISH Mask is consistently superior to a random mask across all sparsity levels tested, with minimal degradation even at low sparsity.
• Using only 32 samples to compute the Fisher approximation yields near-optimal performance, indicating data efficiency in importance estimation.
• A fixed FISH mask at 10% sparsity causes only a few percent drop in accuracy, suggesting potential for identifying lottery ticket subnetworks.
• The empirical Fisher approximation performs as well as the true Fisher, enabling faster and more practical mask computation without performance loss.