[论文解读] On the effectiveness of partial variance reduction in federated learning with heterogeneous data

Bo Li, Mikkel N. Schmidt+2 more|arXiv (Cornell University)|Dec 5, 2022

Privacy-Preserving Technologies in Data被引用 4

一句话总结

本文提出 FedPVR，一种在联邦学习中应用的部分方差缩减方法，仅对深度神经网络的最终分类层应用方差缩减。通过在最后几层专门减少梯度方差，FedPVR 在非独立同分布（non-IID）数据上实现了比 FedAvg 及其他基线方法更快的收敛速度和更高的准确率，同时保持较低的通信开销，并在凸与非凸设置下证明了收敛性。

ABSTRACT

Data heterogeneity across clients is a key challenge in federated learning. Prior works address this by either aligning client and server models or using control variates to correct client model drift. Although these methods achieve fast convergence in convex or simple non-convex problems, the performance in over-parameterized models such as deep neural networks is lacking. In this paper, we first revisit the widely used FedAvg algorithm in a deep neural network to understand how data heterogeneity influences the gradient updates across the neural network layers. We observe that while the feature extraction layers are learned efficiently by FedAvg, the substantial diversity of the final classification layers across clients impedes the performance. Motivated by this, we propose to correct model drift by variance reduction only on the final layers. We demonstrate that this significantly outperforms existing benchmarks at a similar or lower communication cost. We furthermore provide proof for the convergence rate of our algorithm.

研究动机与目标

• 探究 FedAvg 在使用深度神经网络的非独立同分布（non-IID）异构数据设置下表现不佳的原因。
• 识别神经网络中受数据异构性和模型漂移影响最大的层。
• 设计一种通信高效的联邦学习方法，仅针对最易发生漂移的层进行修正。
• 在凸与非凸设置下证明所提方法的收敛性。
• 证明部分方差缩减在准确率与通信效率方面优于全模型方差缩减和基线方法。

提出的方法

• 该方法基于客户端与服务器分类器梯度差异引入控制变量，以减少最后全连接层的方差。
• 仅对最后几层（例如 VGG-11 中的第 27–28 层）应用方差缩减，其余层使用标准 FedAvg 更新。
• 引入漂移多样性度量，以量化每轮通信中各客户端之间梯度方向与尺度的变化。
• 在非方差缩减层中使用动量以稳定训练过程。
• 应用置信预测（conformal prediction）以提升高风险应用场景下的可靠性，利用一小部分验证集进行校准。
• 对收敛速率进行理论分析，表明其对数据异构性依赖较弱，且收敛速度可与集中式 SGD 相媲美。

实验结果

研究问题

• RQ1 在联邦学习中，深度神经网络的哪些层最易受数据异构性的影响？
• RQ2 仅对最终层应用部分方差缩减，是否能提升非独立同分布设置下的收敛速度与准确率？
• RQ3 在通信效率与最终模型性能方面，所提方法与 FedAvg、FedProx、FedDyn 和 SCAFFOLD 相比如何？
• RQ4 所提方法在非凸设置下是否可证明收敛？其收敛速率如何随数据异构性变化？
• RQ5 置信预测是否能在不增加通信开销的前提下进一步提升预测可靠性？

主要发现

• 在 CIFAR10 和 CIFAR100 上，FedPVR 在不同数据异构性水平（α=0.1 至 α=1.0）下，收敛速度优于 FedAvg、FedProx、FedDyn 和 SCAFFOLD。
• 在固定通信轮数下，FedPVR 在多个设置中达到的 top-1 准确率与集中式学习相当或更优。
• 仅对最终层应用方差缩减，可提升特征提取层的梯度多样性，从而增强模型表达能力。
• 该方法在达到相同准确率时所需通信轮数少于基线方法，且通信参数量与 FedAvg 相当或略多。
• 理论分析表明，FedPVR 的收敛速度与集中式 SGD 相当，且对数据异构性度量的依赖性极弱。
• 置信预测提升了预测可靠性，自适应预测集实现了高经验覆盖度，同时保持较小的平均预测集大小。