[论文解读] Multi-Level Branched Regularization for Federated Learning

Jinkyu Kim, Geeho Kim+1 more|arXiv (Cornell University)|Jul 14, 2022

Privacy-Preserving Technologies in Data被引用 11

一句话总结

本文提出 FedMLB，一种用于联邦学习的多级分支正则化技术，通过将本地和全局子网络嫁接到多个混合路径中，减少了非独立同分布（non-iid）设置下的模型漂移。利用在线知识蒸馏将本地主路径表示与混合分支的表示对齐，FedMLB 在不增加通信或内存开销的情况下提升了准确率和收敛速度，在多种骨干网络和数据异质性水平下均优于最先进方法。

ABSTRACT

A critical challenge of federated learning is data heterogeneity and imbalance across clients, which leads to inconsistency between local networks and unstable convergence of global models. To alleviate the limitations, we propose a novel architectural regularization technique that constructs multiple auxiliary branches in each local model by grafting local and global subnetworks at several different levels and that learns the representations of the main pathway in the local model congruent to the auxiliary hybrid pathways via online knowledge distillation. The proposed technique is effective to robustify the global model even in the non-iid setting and is applicable to various federated learning frameworks conveniently without incurring extra communication costs. We perform comprehensive empirical studies and demonstrate remarkable performance gains in terms of accuracy and efficiency compared to existing methods. The source code is available at our project page.

研究动机与目标

• 解决由于数据异质性和客户端数据分布非独立同分布（non-iid）带来的联邦学习中模型漂移问题。
• 在不增加通信或内存开销的前提下，减少训练过程中本地模型与全局模型之间的差异。
• 开发一种正则化技术，在保留全局表示知识的同时，支持客户端特定的适应性。
• 确保在多种联邦学习框架和骨干网络架构中的兼容性与有效性。
• 在低参与度和高度异质的客户端环境中，提升收敛速度和最终准确率。

提出的方法

• 通过在不同网络架构层级上嫁接来自全局模型的子网络，在每个本地模型中构建多个混合分支。
• 在训练本地模型主路径时，固定每个混合分支中全局子网络的参数。
• 在所有分支上应用交叉熵损失以保留分类能力，并通过主路径与混合分支输出之间的KL散度实现知识蒸馏。
• 在知识蒸馏中使用温度缩放的Softmax，以平滑输出分布并提升蒸馏稳定性。
• 反向传播过程中仅更新主路径参数，将混合分支参数作为固定的学生教师模型保留。
• 将该方法无缝集成到标准联邦学习流程（如FedAvg）中，无需修改通信或聚合机制。

实验结果

研究问题

• RQ1 多级架构正则化是否能在不增加通信或内存成本的前提下，减少非独立同分布联邦学习中的模型漂移？
• RQ2 与标准蒸馏或正则化相比，主路径与混合路径之间在线知识蒸馏在提升收敛性和准确率方面有何优势？
• RQ3 FedMLB在不同神经网络架构和数据分布设置下的泛化能力如何？
• RQ4 混合路径的数量及其在网络深度中的位置如何影响性能与效率？
• RQ5 在客户端参与率较低和数据异质性较高的情况下，FedMLB是否仍能保持性能？

主要发现

• 在CIFAR-100数据集上，使用ResNet-18和Dirichlet(0.3)非独立同分布数据时，FedMLB达到54.58%的top-1准确率，显著优于FedAvg（47.83%）、FedLS-NTD（49.29%）和FedGKD（47.96%）。
• 在VGG-9上，FedMLB在相同条件下实现54.54%的准确率，显著超过FedAvg（47.04%）和FedGKD（48.62%）。
• 该方法在所有测试的骨干网络中均一致提升准确率，包括轻量化模型如MobileNet（48.34% vs. 47.66% for FedLS-NTD）和ShuffleNet（42.29% vs. 39.27% for FedLS-NTD）。
• 增加混合路径数量可进一步提升性能，表8显示准确率随路径数量增加而呈现明显上升趋势。
• 即使在更高的本地计算量下（如100个本地训练轮次），FedMLB仍保持优异性能，而FedAVG和FedLS-NTD在迭代次数增加时表现出不一致的行为。
• 消融实验表明，混合路径中的交叉熵损失和KL散度损失均对性能有贡献，其中KL损失更为关键（见表9）。