arXiv:2410.03281cs.LG2024-10被引 5

解决联邦学习中批量归一化统计量漂移问题,提升模型性能。

BN-SCAFFOLD: controlling the drift of Batch Normalization statistics in Federated Learning

  • 扩展SCAFFOLD算法,修正批量归一化统计量的客户端偏移
  • 理论证明收敛性,在MNIST和CIFAR-10上性能媲美FedTAN
  • 通信开销远低于FedTAN,适合深度网络的联邦学习场景

联邦学习(FL)作为分布式机器学习范式日益流行。批量归一化(BN)在深度神经网络(DNN)中广泛应用,能加速收敛并提升泛化能力。然而,当面对异构数据时,BN会降低模型性能。近期提出的FedTAN通过聚合所有客户端的BN统计量与梯度来缓解该问题,但通信开销随网络深度线性增长。SCAFFOLD是一种高效的方差缩减算法,可估计并校正客户端漂移,但在含BN的模型中表现不佳。本文提出BN-SCAFFOLD,将SCAFFOLD的客户端漂移校正机制扩展至BN统计量。我们建立统一的理论框架分析方差缩减算法在含BN的DNN中的收敛性,证明原SCAFFOLD无法消除由BN引入的偏差。实验在MNIST和CIFAR-10上验证了理论结果,表明BN-SCAFFOLD达到与FedTAN相当的性能,且通信开销显著更低,优于联邦平均(FedAvg)、SCAFFOLD及其他针对BN异构设计的算法。

原文摘要 · Abstract (English)

Federated Learning (FL) is gaining traction as a learning paradigm for training Machine Learning (ML) models in a decentralized way. Batch Normalization (BN) is ubiquitous in Deep Neural Networks (DNN), as it improves convergence and generalization. However, BN has been reported to hinder performance of DNNs in heterogeneous FL. Recently, the FedTAN algorithm has been proposed to mitigate the effect of heterogeneity on BN, by aggregating BN statistics and gradients from all the clients. However, it has a high communication cost, that increases linearly with the depth of the DNN. SCAFFOLD is a variance reduction algorithm, that estimates and corrects the client drift in a communication-efficient manner. Despite its promising results in heterogeneous FL settings, it has been reported to underperform for models with BN. In this work, we seek to revive SCAFFOLD, and more generally variance reduction, as an efficient way of training DNN with BN in heterogeneous FL. We introduce a unified theoretical framework for analyzing the convergence of variance reduction algorithms in the BN-DNN setting, inspired of by the work of Wang et al. 2023, and show that SCAFFOLD is unable to remove the bias introduced by BN. We thus propose the BN-SCAFFOLD algorithm, which extends the client drift correction of SCAFFOLD to BN statistics. We prove convergence using the aforementioned framework and validate the theoretical results with experiments on MNIST and CIFAR-10. BN-SCAFFOLD equals the performance of FedTAN, without its high communication cost, outperforming Federated Averaging (FedAvg), SCAFFOLD, and other FL algorithms designed to mitigate BN heterogeneity.

联邦学习批量归一化方差缩减通信效率

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。