[论文解读] FedBABU: Towards Enhanced Representation for Federated Image Classification

Jaehoon Oh, Sangmook Kim+1 more|arXiv (Cornell University)|Jun 4, 2021

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

一句话总结

{

ABSTRACT

Federated learning has evolved to improve a single global model under data heterogeneity (as a curse) or to develop multiple personalized models using data heterogeneity (as a blessing). However, little research has considered both directions simultaneously. In this paper, we first investigate the relationship between them by analyzing Federated Averaging at the client level and determine that a better federated global model performance does not constantly improve personalization. To elucidate the cause of this personalization performance degradation problem, we decompose the entire network into the body (extractor), which is related to universality, and the head (classifier), which is related to personalization. We then point out that this problem stems from training the head. Based on this observation, we propose a novel federated learning algorithm, coined FedBABU, which only updates the body of the model during federated training (i.e., the head is randomly initialized and never updated), and the head is fine-tuned for personalization during the evaluation process. Extensive experiments show consistent performance improvements and an efficient personalization of FedBABU. The code is available at https://github.com/jhoon-oh/FedBABU.

研究动机与目标

• 激发对在数据异质性下单一全局模型如何与个性化模型交互的理解。
• 将神经网络解耦为通用的主体（body）和个性化的头部（head），以识别个性化退化的来源。
• 提出并验证 FedBABU，在联邦训练过程中仅更新主体，头部保持冻结。
• 证明在客户端之间固定头部可提供强表示能力，并加速设备端个性化。

提出的方法

• 将模型参数解耦为主体（提取器）和头部（分类器）。
• 在 FedBABU 期间，在每个客户端仅更新主体；头部使用固定的、全局随机初始化的头部。
• 仅聚合各客户端的主体参数；在联邦训练期间不更新或聚合头部。
• 通过比较初始全局模型性能与客户端特定微调性能来评估个性化。
• 显示固定头部在集中式性能方面可与共同训练的头部相媲美。
• 可选地用模板替换训练过的头部以评估无头的表示能力。

实验结果

研究问题

• RQ1 在具有异质性数据的联邦学习中，解耦主体和头部是否能缓解个性化退化？
• RQ2 在联邦训练中冻结头部是否能提高通用表示的质量，从而提高后续个性化微调？
• RQ3 在不同数据异质性下，FedBABU 与 FedAvg 及其他个性化方法在初始准确性和个性化准确性方面有何比较？
• RQ4 在不使用头部、仅用类别模板进行评估时，FedBABU 全局模型的表示能力是否稳健？
• RQ5 FedBABU 是否可以推广到如 FedProx 这样的正则化联邦学习框架？

主要发现

• FedBABU 仅训练主体在表示能力方面比 FedAvg 更强或相当，尤其在高度异质性下。
• 在客户端之间固定头部为表示学习提供了强劲的引导，提升个性化效率。
• FedBABU 能实现快速个性化，通常只需少量微调轮次，在若干场景中优于 FedAvg。
• 当头部未训练时，FedBABU 实现了更高或具有竞争力的初始精度，同时不牺牲个性化。
• 微调头部或整个模型可为 FedAvg 提供更好的个性化，而 FedBABU 只微调头部也能取得较强结果。