[论文解读] 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 Learning参考文献 48被引用 4

一句话总结

本文提出FISH Mask，一种预计算固定稀疏掩码的方法，基于经验Fisher信息量选择k个最重要参数用于训练过程中的梯度更新。通过在多个迭代中仅更新这些高重要性参数，该方法在迁移学习和分布式训练中显著降低了内存和通信开销，同时在10%稀疏度下仅带来极小的准确率下降，性能与标准训练相当或更优。

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.

研究动机与目标

• 解决大规模神经网络训练中的高存储与通信成本问题，特别是在迁移学习以及分布式/联邦学习场景中。
• 通过仅在每次迭代中更新一小部分固定参数子集，而非全部参数，降低模型更新开销。
• 开发一种与模型无关、高效的稀疏参数更新方法，避免引入新参数或动态改变更新掩码。
• 实现在带宽或存储受限场景下的实际部署，如模型检查点保存和联邦学习。
• 证明基于参数重要性（通过Fisher信息量衡量）的固定稀疏掩码可在性能下降极小的情况下维持高训练性能。

提出的方法

• 使用经验Fisher信息量估算参数重要性，其计算方式为在小批量数据上的样本梯度方差。
• 通过选择Fisher信息值最高的k个参数构建固定稀疏掩码，确保仅这些参数在训练过程中被更新。
• 利用单次前向和反向传播（N=1024，分布式设置中为N=256）的小样本集高效计算Fisher近似值。
• 应用标准随机梯度下降（SGD），但仅将权重更新限制在预计算的FISH掩码所选参数上。
• 在整个训练迭代过程中保持相同的掩码，避免重新计算或动态调整。
• 通过在优化器步骤中作为二值掩码应用，确保与现有训练流水线的兼容性，无需修改模型架构。

实验结果

研究问题

• RQ1 基于参数重要性的固定稀疏掩码是否能在降低通信与存储成本的同时维持高模型性能？
• RQ2 使用Fisher信息量作为参数重要性代理，相较于随机选择或启发式选择，在稀疏训练中表现如何？
• RQ3 掩码稀疏度与更新频率对参数高效迁移学习中模型准确率的影响是什么？
• RQ4 Fisher计算在数据与计算资源方面的效率如何？需要多少样本才能实现稳定性能？
• RQ5 固定FISH掩码是否可用于识别稀疏子网络，以支持类似彩票理论（Lottery Ticket Hypothesis）的训练方法？

主要发现

• 在10%稀疏度下，FISH Mask在GLUE基准测试中达到93.9%的准确率，与标准训练性能相当。
• 每轮训练更新一次FISH掩码时，性能与完整训练几乎完全一致，同时将存储需求降低5倍。
• 在所有测试的稀疏度水平下，FISH Mask的性能均显著优于随机掩码，即使在低稀疏度下也仅出现极小的性能下降。
• 仅使用32个样本计算Fisher近似即可获得接近最优的性能，表明重要性估计具有高度数据效率。
• 在10%稀疏度下，固定FISH掩码仅导致准确率下降几个百分点，表明其在识别彩票子网络方面具有潜力。
• 经验Fisher近似与真实Fisher信息量表现相当，实现了无需性能损失的更快、更实用的掩码计算。