Training Neural Networks with Fixed Sparse Masks

TL;DR

The paper tackles the high cost of updating all parameters during gradient-based training by introducing FISH Mask, a fixed sparse mask pre-computed using the Fisher information to identify the most influential parameters. The method computes an approximation of the Fisher information $\hat{F}_{\theta}$ (or the true $F_\theta$) and selects the top-$k$ parameters to update. Experiments show that across parameter-efficient transfer learning, distributed training, and efficient checkpointing, updating only 0.5–10% of parameters yields competitive performance with substantial reductions in memory and communication. The approach is model-agnostic and does not increase parameter count, and is complementary to existing techniques; code is released.

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.

Figures (3)

• Figure 1: Diagram comparing our proposed method to standard SGD. In traditional gradient-based training (left), all of a model's parameters are updated at every iteration. We propose FISH Mask, a method for precomputing a sparse subset of parameters to update over many subsequent training iterations. To construct the FISH Mask, we find the $k$ parameters with the largest Fisher information (right, top). Then, we train the model with traditional gradient descent, but only update those parameters chosen by the mask (right, bottom).
• Figure 2: (Left) GLUE validation performance of a randomly selected mask and the FISH Mask at varying levels of mask sparsity. Compared to the densely fine-tuned baseline score of 85%, training with the FISH Mask is competitive at 0.5% mask sparsity. (Right) GLUE validation performance at varying levels of dataset samples used to compute the FISH Mask. Few samples are needed to effectively compute the FISH Mask and obtain good performance. Results in both (Left) and (Right) are averaged over 5 seeds.
• Figure 3: CIFAR-10 validation set accuracy achieved by a ResNet-34 through distributed training at different communication costs. X-axis refers to the total number of model communications required for a single worker. Standard (non-distributed) training achieves an accuracy of 93.9%.