[论文解读] Test-Time Robust Personalization for Federated Learning

Liangze Jiang, Tao Lin|arXiv (Cornell University)|May 22, 2022

Privacy-Preserving Technologies in Data被引用 7

一句话总结

本文提出了一种用于联邦学习的测试时鲁棒个性化方法，通过测试时微调和可学习集成权重动态调整模型头。通过基于测试样本历史和特征表示来优化本地和全局头参数，该方法在分布偏移和非独立同分布（non-IID）数据条件下显著优于基线方法，性能提升明显。

ABSTRACT

Federated Learning (FL) is a machine learning paradigm where many clients collaboratively learn a shared global model with decentralized training data. Personalized FL additionally adapts the global model to different clients, achieving promising results on consistent local training and test distributions. However, for real-world personalized FL applications, it is crucial to go one step further: robustifying FL models under the evolving local test set during deployment, where various distribution shifts can arise. In this work, we identify the pitfalls of existing works under test-time distribution shifts and propose Federated Test-time Head Ensemble plus tuning(FedTHE+), which personalizes FL models with robustness to various test-time distribution shifts. We illustrate the advancement of FedTHE+ (and its computationally efficient variant FedTHE) over strong competitors, by training various neural architectures (CNN, ResNet, and Transformer) on CIFAR10 andImageNet with various test distributions. Along with this, we build a benchmark for assessing the performance and robustness of personalized FL methods during deployment. Code: https://github.com/LINs-lab/FedTHE.

研究动机与目标

• 为解决联邦学习中在数据分布偏移和非独立同分布客户端数据条件下的模型个性化挑战。
• 在不重新训练全局模型或访问客户端数据的前提下，提升测试时性能。
• 开发一种在推理过程中利用本地测试统计信息动态调整个性化模型头的方法。
• 通过在推理过程中同时利用全局和本地模型组件，增强对分布偏移的鲁棒性。
• 以可学习、自适应的方式在测试时优化全局与本地头之间的集成权重。

提出的方法

• 该方法使用全局模型头参数 $\mathbf{w}^g$ 和本地模型头参数 $\mathbf{w}^l$，二者均从全局模型初始化。
• 对本地头 $\mathbf{w}^l$ 应用测试时微调，使用 $\tau$ 个优化步骤在输入的测试样本 $\mathbf{x}_n$ 上进行，学习率 $\eta$ 可学习。
• 通过可学习权重 $e \in [0,1]$ 构成全局与本地头的集成，并在测试时适应过程中通过梯度优化更新该权重。
• 模型维护一个测试历史 $\mathbf{h}^{\text{history}}$，以指导自适应头加权和特征表示的调整。
• 特征表示被划分为全局 $\mathbf{h}^g$ 和本地 $\mathbf{h}^l$ 两部分，后者通过在测试样本上微调进行更新。
• 超参数 $\alpha = 0.1$ 和 $\beta = 0.3$ 控制集成权重和头更新优化目标中各项损失的平衡。

实验结果

研究问题

• RQ1 如何使联邦学习中的模型个性化在推理过程中对分布偏移具有鲁棒性？
• RQ2 在不重新训练全局模型的前提下，测试时对模型头进行动态调整是否能提升性能？
• RQ3 使用可学习的全局与本地头集成对测试时准确率有何影响？
• RQ4 对本地头进行测试时微调在非独立同分布数据下对泛化能力有何影响？
• RQ5 维护测试历史在多大程度上能提升个性化预测的鲁棒性？

主要发现

• 所提方法在分布偏移条件下显著优于标准个性化基线方法，准确率大幅提升。
• 采用自适应集成加权的测试时微调，相比静态头集成，误差率最高降低25%。
• 即使客户端数据分布偏离训练数据，该方法仍保持高性能，展现出强大的鲁棒性。
• 使用可学习集成权重 $e$ 能够自动平衡全局与本地头，提升泛化能力。
• 该方法在非独立同分布数据设置下表现有效，且在不同客户端数据分布下均保持一致的性能增益。
• 消融实验确认，测试时微调和基于历史的自适应调整均对性能提升有显著贡献。