[论文解读] Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning

Liangqiong Qu, Yuyin Zhou+6 more|arXiv (Cornell University)|Jun 10, 2021

Privacy-Preserving Technologies in Data参考文献 61被引用 12

一句话总结

本文提出将视觉变换器（ViTs）作为联邦学习中卷积神经网络（CNNs）的稳健架构替代方案，证明ViTs在异构数据下能显著减少灾难性遗忘并加速收敛。主要贡献在于实证证据表明，ViT-FL在多个基准测试和数据划分下均优于基于CNN的联邦学习，尤其在高数据异构性条件下表现更优。

ABSTRACT

Federated learning is an emerging research paradigm enabling collaborative training of machine learning models among different organizations while keeping data private at each institution. Despite recent progress, there remain fundamental challenges such as the lack of convergence and the potential for catastrophic forgetting across real-world heterogeneous devices. In this paper, we demonstrate that self-attention-based architectures (e.g., Transformers) are more robust to distribution shifts and hence improve federated learning over heterogeneous data. Concretely, we conduct the first rigorous empirical investigation of different neural architectures across a range of federated algorithms, real-world benchmarks, and heterogeneous data splits. Our experiments show that simply replacing convolutional networks with Transformers can greatly reduce catastrophic forgetting of previous devices, accelerate convergence, and reach a better global model, especially when dealing with heterogeneous data. We release our code and pretrained models at https://github.com/Liangqiong/ViT-FL-main to encourage future exploration in robust architectures as an alternative to current research efforts on the optimization front.

研究动机与目标

• 解决由于设备间数据异构性导致的联邦学习中非收敛和灾难性遗忘的持续挑战。
• 探究架构选择（特别是视觉变换器）是否能从根本上提升对联邦学习中分布偏移的鲁棒性。
• 提供系统性的实证基准，比较ViTs与CNNs在联邦学习算法、真实世界数据集和异构数据划分下的表现。
• 证明架构改进可立即带来性能提升，而无需额外的优化启发式方法或超参数调优。
• 将视觉变换器定位为未来联邦学习研究的自然起点，与现有基于优化的方法正交。

提出的方法

• 在多个联邦学习算法（FedAVG（并行）和CWT（串行））上，对视觉变换器（ViTs）和CNNs（ResNets、EfficientNets）进行大规模实证基准测试。
• 在图像分类任务中使用标准ViT实现，包括CIFAR-10和Retina，采用一致的数据增强和正则化策略。
• 在三种数据划分策略（Split-1至Split-3）下评估性能，逐步增加标签分布偏度，以模拟现实世界中的异构性。
• 应用UMAP可视化比较ViT和CNN模型学习到的特征嵌入，分析聚类分离度和表征鲁棒性。
• 将ViTs与现有基于优化的联邦学习方法（FedProx、FedAVG-Share）结合，测试兼容性和性能增益。
• 研究预训练以及训练超参数（如学习率、梯度裁剪、本地训练轮数）对ViT-FL性能的影响。

实验结果

研究问题

• RQ1 与CNN相比，视觉变换器是否能在高数据异构性条件下减轻联邦学习中的灾难性遗忘？
• RQ2 ViT-FL在不同联邦学习算法（FedAVG与CWT）和数据划分策略下的性能表现如何，相较于基于CNN的联邦学习？
• RQ3 在非独立同分布（non-IID）设置下，架构选择（ViT与CNN）在收敛速度和最终模型准确率方面的影响程度如何？
• RQ4 当缺乏大规模数据集时，预训练是否能显著提升ViT-FL的性能？
• RQ5 ViT-FL能否与现有基于优化的联邦学习方法有效结合，从而在异构数据上进一步提升性能？

主要发现

• 在CIFAR-10的高异构数据划分（Split-3）中，视觉变换器（ViTs）在测试准确率上显著优于ResNets，预训练下达到96.40%，而ResNet(50)-FedAVG仅为59.68%。
• ViT-FL减少了灾难性遗忘：与ResNet(50)-FedAVG相比，ViT(S)-FedAVG学习到的特征在UMAP可视化中表现出更清晰的类别分离。
• ViT-FL加速了收敛并达到了更优的全局模型，尤其在严重数据异构性下，性能增益随数据偏度增加而提升。
• 预训练对ViT-FL至关重要：当从零训练时，Swin(T)-FedAVG在Split-3上达到64.50%的准确率，优于未进行大规模预训练的ResNet(50)-FedAVG（59.68%）。
• ViT-FL与基于优化的方法正交：将ViTs与FedProx或FedAVG-Share结合，可在异构客户端上进一步提升性能。
• 为获得最佳性能，在高度异构设置下应使用较小的本地训练轮数（ViT-FedAVG中E ≤ 5，ViT-CWT中E = 1），而在同质情况下可使用更大的E以减少通信开销。