Source-linked AI summary
Training Neural Networks with Fixed Sparse Masks
Yi-Lin Sung, Varun Nair, Colin Raffel
TL;DR
Updating all neural-network parameters can make training communication and storage expensive. FISH Mask precomputes a fixed subset using the largest approximate Fisher-information values and reuses it during training. The method matches dense performance in a GLUE example while updating 0.5% of BERTLARGE parameters and reduces communication or storage requirements.
Problem
Standard training updates every parameter, creating expensive communication and storage demands for large neural networks.
Method
FISH Mask estimates parameter importance with approximate Fisher information, selects the k largest values, and reuses that fixed sparse subset across training iterations.
Results
FISH Mask achieved 82.6% GLUE performance versus 82.5% for dense fine-tuning while updating 0.5% of BERTLARGE parameters.
Takeaways & Limitations
FISH Mask can reduce storage and communication requirements while sacrificing minimal performance compared with standard gradient descent.
Takeaways & Limitations
The empirical Fisher approximation can lead to degenerate behavior when used as an optimizer preconditioner, and newly added classifiers require special handling.
Abstract
from arXiv · showhide
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.
1 Introduction
SGD trains neural networks by updating all parameters each iteration, but those updates create substantial communication costs for large models. FISH Mask instead precomputes and repeatedly updates a fixed, Fisher-informed subset of existing parameters.
- SGD minimizes loss on small random dataset subsets while updating model parameters and often producing models that generalize well.
- Updating every parameter makes communicating model changes expensive for neural networks with millions or billions of parameters.
- FISH Mask precomputes a fixed subset of existing parameters using the k parameters with the largest Fisher information, then updates only that subset.The mask is reused over many subsequent training iterations.
- Parameter-efficient methods reduce updates by adding modules, hand-selecting parameters, or changing the updated subset during training.
- The fixed-mask design avoids increasing model size, supports model-agnostic selection, and reduces computational and memory overhead from repeatedly changing masks.
2 The Fisher-Induced Sparse uncHanging (FISH) Mask
FISH Mask estimates parameter importance with an approximate Fisher information measure and selects the k largest entries for a reusable sparse update mask. The method includes practical approximations and a transfer-learning rule for newly added classifiers.
- FISH Mask represents the paper’s sparse subset of parameters to update over many iterations, selected using approximate Fisher information.
- 2.1 Fisher Information: Parameter importance is motivated by how much a small parameter change alters the model’s output distribution, a relationship captured locally by Fisher information.
- 2.1 Fisher Information: Because the full Fisher matrix is intractable for large models, the method commonly uses a diagonal vector approximation based on finite samples.
- 2.1 Fisher Information: The empirical Fisher replaces the model-output expectation over labels with the squared gradient for each ground-truth label in supervised learning.The paper notes that this approximation can lead to degenerate behavior when used as an optimizer preconditioner.
- 2.2 Computing Fixed Sparse Masks: The mask is constructed by computing approximate Fisher information for all parameters and selecting the k largest values, with k determined by the desired mask sparsity.Backpropagation computes the approximation efficiently, and mask reuse avoids frequent recomputation.
- 2.2 Computing Fixed Sparse Masks: When transfer learning adds a classifier, FISH Mask is computed before training through the random classifier, whose parameters are always included.
3 Related Work
Related work addresses parameter-efficient transfer learning, distributed communication, and network pruning through different forms of parameter or update sparsity. FISH Mask differs by selecting and updating a fixed subset of existing parameters rather than adding parameters, changing masks, or zeroing weights.
- 3.1 Parameter-Efficient Transfer Learning: Transfer learning initializes a model from a pretrained checkpoint before fine-tuning it on a related downstream task.
- 3.1 Parameter-Efficient Transfer Learning: Adapters add small subnetworks whose parameters are updated alongside the output layer while the pretrained model remains otherwise fixed.
- 3.1 Parameter-Efficient Transfer Learning: Diff Pruning learns an update mask during training and can outperform Adapters on BERT GLUE, but requires significantly more training memory to store and update the mask.
- 3.2 Distributed Training: Distributed training shares computation across workers, making communication costs important because workers regularly exchange parameter updates.
- 3.2 Distributed Training: Federated learning uses asynchronous workers that train on private data before sending changes to a centralized model, with stale gradients and communication costs as key issues.
- Network pruning sparsifies neural networks by zeroing parameters for compression, whereas FISH Mask trains with sparse updates without creating sparse weights.
4 Experiments
Experiments evaluate FISH Masks for parameter-efficient transfer learning, distributed training, checkpointing, and Fisher-estimation choices. Across these settings, task-informed sparse masks generally outperform random masks while preserving performance with substantially fewer updated parameters or communications.
- Parameter-Efficient Transfer Learning: 82.6% average GLUE performance matches dense fine-tuning at 82.5% while updating only 0.5% of BERTLARGE parameters.
- Parameter-Efficient Transfer Learning: The FISH Mask consistently outperforms random masks across GLUE mask-sparsity levels and remains strong at 0.1% sparsity.Saving updated parameters and their indices can reduce the storage required for many fine-tuned models.
- Distributed Training: FISH Masks provide a better distributed-training communication/performance trade-off than random masks, and 10% sparsity can match standard training when communicating two model copies.Using only 2% of parameters produced relatively poor performance, indicating a lower bound for reasonable results in this setting.
- Efficient Checkpointing: Sparse checkpointing stores updated values and their indices rather than full parameter copies, reducing storage costs when few parameters change between checkpoints.
- Ablations: True and empirical Fisher estimates both achieve an average GLUE validation score of 82.5 for 0.5%-sparse masks.The empirical Fisher is used because it avoids marginalization or sampling from the model output distribution and requires one ground-truth-label gradient.
- Ablations: Only 32 samples are needed to achieve the highest possible GLUE performance when computing the FISH Mask.Performance remains surprisingly stable across many sample counts.
5 Conclusion
The paper proposes FISH Mask training, which pre-computes fixed sparse parameter masks using Fisher information and applies them over many iterations. It demonstrates usefulness across transfer learning, distributed training, and checkpoint storage, while identifying lower-sparsity performance and real-world federated learning as future directions.
- FISH Mask training pre-computes fixed sparse parameter masks for updating over many subsequent iterations.The mask estimates each parameter’s Fisher information and selects the k parameters with the largest values.
- The method estimates parameter importance with approximated Fisher information, then selects the k parameters with the largest values.
- Experiments demonstrate FISH Mask training in parameter-efficient transfer learning, distributed training, and reducing model-checkpoint storage requirements.
- Future work targets better performance at lower mask sparsity levels, potentially using alternative measures of parameter importance.
- The paper also identifies real-world federated learning and sharing FISH Masks across tasks as future directions.
A Distributed BERT fine-tuning experiments
Distributed fine-tuning experiments on BERTLARGE evaluated 0.5%-sparse FISH Masks on GLUE, including worker-shared and nonoverlapping mask designs. Shared FISH Mask training matched dense baselines and outperformed random masks, whereas nonoverlapping masks were detrimental.
- The experiments fine-tuned BERTLARGE on GLUE for 7 epochs, with 3.5 epochs performed by each worker.
- Segmenting the FISH Mask into nonoverlapping worker-specific portions was detrimental compared with sharing the same FISH Mask across all workers.The segmented design aimed to give workers complementary parameter sets with roughly balanced Fisher information.
- At 0.5% mask sparsity, FISH Mask training achieved comparable performance to standard training and densely updated distributed training across varying worker updates.Standard non-distributed fine-tuning attained an average GLUE validation score of 85.0%.
- FISH Mask training performed significantly better than random-mask training in distributed BERTLARGE fine-tuning.