[论文解读] Local Learning Matters: Rethinking Data Heterogeneity in Federated Learning

Matías Mendieta, Taojiannan Yang+4 more|arXiv (Cornell University)|Nov 28, 2021

Privacy-Preserving Technologies in Data参考文献 29被引用 10

一句话总结

本文提出 FedAlign，一种资源高效的联邦学习正则化方法，通过增强本地模型泛化能力而非依赖近端约束来提升性能。通过将知识蒸馏应用于学生模型以正则化最后一层的Lipschitz常数，FedAlign 在视觉基准测试中实现了最先进（SOTA）的准确率，同时相比 MOON 和 GradAug 等 SOTA 方法，计算和内存开销分别降低了 65% 以上和 70% 以上。

ABSTRACT

Federated learning (FL) is a promising strategy for performing privacy-preserving, distributed learning with a network of clients (i.e., edge devices). However, the data distribution among clients is often non-IID in nature, making efficient optimization difficult. To alleviate this issue, many FL algorithms focus on mitigating the effects of data heterogeneity across clients by introducing a variety of proximal terms, some incurring considerable compute and/or memory overheads, to restrain local updates with respect to the global model. Instead, we consider rethinking solutions to data heterogeneity in FL with a focus on local learning generality rather than proximal restriction. To this end, we first present a systematic study informed by second-order indicators to better understand algorithm effectiveness in FL. Interestingly, we find that standard regularization methods are surprisingly strong performers in mitigating data heterogeneity effects. Based on our findings, we further propose a simple and effective method, FedAlign, to overcome data heterogeneity and the pitfalls of previous methods. FedAlign achieves competitive accuracy with state-of-the-art FL methods across a variety of settings while minimizing computation and memory overhead. Code is available at https://github.com/mmendiet/FedAlign

研究动机与目标

• 将联邦学习中数据异质性的处理方式从近端正则化重新定位为本地学习泛化。
• 识别出在不产生高计算或内存成本的前提下提升模型泛化能力的有效正则化技术。
• 设计一种轻量化、高效的正则化方法，使其在多种非独立同分布（non-IID）数据设置下仍能保持高性能。
• 验证提升本地学习泛化能力可带来更好的全局模型收敛性和鲁棒性。
• 证明在联邦学习中，资源效率与高准确率并非相互排斥。

提出的方法

• FedAlign 使用知识蒸馏来正则化神经网络的最后模块，重点在于最后一层表示的Lipschitz常数。
• 引入一种学生-教师蒸馏损失，将客户端模型的特征表示与全局教师模型对齐。
• 该方法仅对最后一层应用正则化，从而最小化额外的计算和内存开销。
• 蒸馏损失在本地训练过程中进行优化，促进泛化而非限制对全局模型的更新。
• 通过调优超参数（如蒸馏权重 ω_S = 0.25 和温度 μ = 0.45）以实现最佳性能。
• 该方法与网络架构无关，适用于视觉任务，也可能适用于自然语言处理（NLP）任务，因其在特征空间中操作。

实验结果

研究问题

• RQ1 在联邦学习中，提升本地学习泛化能力是否能超越传统近端正则化方法，以更好地应对数据异质性？
• RQ2 在非独立同分布（non-IID）联邦学习设置下，哪些正则化技术在性能与计算效率之间提供了最佳权衡？
• RQ3 将正则化聚焦于网络最后一层是否能在极低资源成本下带来显著性能提升？
• RQ4 在多种数据集上，FedAlign 与 MOON 和 GradAug 等 SOTA 方法相比，在准确率和效率方面表现如何？
• RQ5 基于知识蒸馏的简单方法是否能实现 SOTA 性能，同时显著降低计算和内存开销？

主要发现

• FedAlign 在 CIFAR-100 和 ImageNet-200 上实现了最先进（SOTA）的准确率，相比 FedAvg 准确率提升约 4.0%，且仅增加 1.02 倍的 FLOPs。
• 与 MOON 相比，FedAlign 在准确率上提升约 1.9%，同时本地计算开销降低 65% 以上，内存使用降低 70% 以上。
• 与 GradAug 相比，FedAlign 在 n=2 和 n=1 的设置下，分别将计算需求降低约 47% 和 33%，且未牺牲准确率。
• 在 ImageNet-200 数据集上，FedAlign 的性能与 GradAug（n=1 和 n=2）相当或更优，同时资源消耗显著更低。
• StochDepth 在 ResNet18 上表现较差，原因在于网络深度减少，凸显了如 FedAlign 这类基于宽度的正则化方法在实际联邦学习部署中的优势。
• 消融实验表明，聚焦于最后一层的基于蒸馏的正则化方法既有效又高效，尤其在真实的客户端采样场景下表现突出。