通过拼接预训练模块实现异构性感知的联邦学习协调
机器学习
2024-09-12 v1 人工智能
摘要
联邦学习(FL)协调多个设备以共同训练共享模型,同时保持数据隐私。然而,大型内存占用和训练过程中的高能耗会排除低端设备,它们无法贡献自己的数据以改善全局模型,这在实际场景中严重恶化了模型性能。本文提出FedStitch,一种用于异构性联邦学习的分层协调框架,支持预训练模块。与传统方法从头开始训练全局模型不同,对于新任务,FedStitch通过拼接预训练模块来构建全局模型。具体而言,每个参与客户端从候选池中(由预训练模型的模块组成)选择最适合自己本地数据的模块。服务器然后聚合用于拼接的最优模块。该过程迭代进行,直至生成新的拼接网络。除了新的训练范式,FedStitch包含以下三个核心组件:1)基于强化学习的加权聚合器,2)部署在服务器端的搜索空间优化器,3)部署在每个参与客户端的本地能源优化器。基于强化学习的加权聚合器有助于在非独立同分布情形下选择正确的模块,而搜索空间优化器在拼接过程中持续减少候选模块池的大小。同时,本地能源优化器旨在最小化每个客户端的能耗,同时保证整体训练进度。结果表明,与现有方法相比,FedStitch在模型准确率上提升了最高可达20.93%,同时实现了最高可达8.12%的加速,内存占用降低最高可达79.5%,在学习过程中实现最高可达89.41%的能源节省。
引用
@article{arxiv.2409.07202,
title = {Heterogeneity-Aware Coordination for Federated Learning via Stitching Pre-trained blocks},
author = {Shichen Zhan and Yebo Wu and Chunlin Tian and Yan Zhao and Li Li},
journal= {arXiv preprint arXiv:2409.07202},
year = {2024}
}