[论文解读] Efficient Wireless Federated Learning with Partial Model Aggregation

Zhixiong Chen, Wenqiang Yi+2 more|arXiv (Cornell University)|Apr 20, 2022

Privacy-Preserving Technologies in Data被引用 5

一句话总结

本文提出部分模型聚合联邦学习（PMA-FL），通过在参数服务器仅聚合低层（特征提取层）而保留高层（预测层）在本地以实现个性化，从而提升无线联邦学习性能。通过拉姆齐优化联合优化设备调度、带宽分配与计算时间分配，PMA-FL降低了通信负载，并增强了对数据异构性的鲁棒性，在MNIST和CIFAR-10上分别实现了3.13%和11.8%的准确率提升，同时相比基准方案，通信能耗降低29%或通信时延减少20%。

ABSTRACT

The data heterogeneity across devices and the limited communication resources, e.g., bandwidth and energy, are two of the main bottlenecks for wireless federated learning (FL). To tackle these challenges, we first devise a novel FL framework with partial model aggregation (PMA). This approach aggregates the lower layers of neural networks, responsible for feature extraction, at the parameter server while keeping the upper layers, responsible for complex pattern recognition, at devices for personalization. The proposed PMA-FL is able to address the data heterogeneity and reduce the transmitted information in wireless channels. Then, we derive a convergence bound of the framework under a non-convex loss function setting to reveal the role of unbalanced data size in the learning performance. On this basis, we maximize the scheduled data size to minimize the global loss function through jointly optimize the device scheduling, bandwidth allocation, computation and communication time division policies with the assistance of Lyapunov optimization. Our analysis reveals that the optimal time division is achieved when the communication and computation parts of PMA-FL have the same power. We also develop a bisection method to solve the optimal bandwidth allocation policy and use the set expansion algorithm to address the device scheduling policy. Compared with the benchmark schemes, the proposed PMA-FL improves 3.13\% and 11.8\% accuracy on two typical datasets with heterogeneous data distribution settings, i.e., MINIST and CIFAR-10, respectively. In addition, the proposed joint dynamic device scheduling and resource management approach achieve slightly higher accuracy than the considered benchmarks, but they provide a satisfactory energy and time reduction: 29\% energy or 20\% time reduction on the MNIST; and 25\% energy or 12.5\% time reduction on the CIFAR-10.

研究动机与目标

• 解决联邦学习中的数据异构性与有限无线通信资源问题。
• 在非独立同分布（non-IID）数据及带宽与能量受限条件下，提升模型收敛性与准确率。
• 降低无线联邦学习系统中的通信开销与能耗。
• 构建设备调度、带宽分配与计算时间分配的联合优化框架。
• 通过计算与通信的平衡，实现更低能耗与更短时延下的更优学习性能。

提出的方法

• 提出一种新型PMA-FL框架，仅在参数服务器聚合神经网络的低层，保留高层用于本地个性化。
• 推导在非凸损失函数下PMA-FL的收敛边界，分析数据规模不平衡的影响。
• 采用拉姆齐优化方法构建联合优化问题，以平衡通信与计算能耗，并最大化调度数据量。
• 设计一种二分法求解资源约束下的最优带宽分配策略。
• 设计集合扩展算法以确定最优设备调度策略。
• 集成通信与计算阶段的动态时间分配机制，表明当两阶段功率均衡时性能最优。

实验结果

研究问题

• RQ1 在数据异构条件下，部分模型聚合如何影响无线联邦学习中的收敛性与性能？
• RQ2 PMA-FL中通信与计算能耗的最优平衡点为何？以最小化全局损失为目标。
• RQ3 联合优化设备调度、带宽分配与时间分配能否在降低能耗与时延的同时提升学习准确率？
• RQ4 在严格延迟与能耗约束下，PMA-FL相较于现有基准方案表现如何？
• RQ5 数据规模不平衡对PMA-FL收敛性与准确率的影响是什么？

主要发现

• 在非IID数据分布下，PMA-FL相较于基准方案在MNIST上准确率提升3.13%，在CIFAR-10上提升11.8%。
• 所提出的联合优化方法在MNIST上实现29%能耗降低或20%时延减少，在CIFAR-10上实现25%能耗降低或12.5%时延减少，同时保持或提升准确率。
• 在20%时延预算减少（如1.5s vs. 2.5s）条件下，所提算法在更严格的延迟约束下仍优于OCEAN。
• 在更高能耗约束（V=0.001）下，该算法性能与OCEAN相当，但更具能耗效率。
• 当通信与计算功率均衡时，时间分配策略达到最优，从而最小化整体系统成本。
• 调节拉姆齐权重参数V可实现准确率与能耗之间的权衡，当V被精细调优时可获得最优性能。