当远见剪枝遇上零阶优化：面向低内存设备的高效联邦学习
机器学习
2024-05-09 v1 人工智能
分布式、并行与集群计算
摘要
尽管联邦学习（FL）使得人工智能物联网（AIoT）设计中的协作学习成为可能，但由于其内存占用过高，无法在低内存AIoT设备上运行。为解决此问题，研究者提出了多种联邦剪枝方法以减少推理过程中的内存使用。然而，其中很少有方法能显著减轻剪枝和训练过程中的内存负担。作为替代方案，零阶或无反向传播（BP-Free）方法可以部分缓解内存消耗，但它们在扩展性和计算开销方面存在问题，因为梯度估计误差和浮点运算次数（FLOPs）会随着模型参数维度的增加而增加。本文提出了一种基于神经正切核（NTK）的联邦远见剪枝方法，该方法能够无缝集成到联邦BP-Free训练框架中。我们通过使用局部NTK矩阵给出了联邦NTK计算的近似。此外，我们证明了该方法无数据依赖的特性能够显著降低极端数据异质性场景下的近似误差。由于我们的方法以更少的FLOPs提升了原始BP-Free方法的性能，并真正减轻了训练和推理过程中的内存压力，因此它使FL对低内存设备更加友好。基于仿真和真实测试平台的大量实验结果表明，我们的联邦远见剪枝方法不仅保留了稠密模型的能力（内存减少高达9倍），而且以显著更少的FLOPs提升了原始BP-Free方法的性能。
引用
@article{arxiv.2405.04765,
title = {When Foresight Pruning Meets Zeroth-Order Optimization: Efficient Federated Learning for Low-Memory Devices},
author = {Pengyu Zhang and Yingjie Liu and Yingbo Zhou and Xiao Du and Xian Wei and Ting Wang and Mingsong Chen},
journal= {arXiv preprint arXiv:2405.04765},
year = {2024}
}