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

Wenxuan Bao, Tianxin Wei+2 more|arXiv (Cornell University)|Oct 28, 2023

Privacy-Preserving Technologies in Data被引用 5

一句话总结

该论文提出ATP，一种用于联邦学习的新型自适应测试时个性化框架，通过从源客户端数据中学习模块特定的适应率，以在无标签测试数据的情况下处理多样的分布偏移。通过根据客户端间分布偏移动态调整适应率，ATP在标签偏移、图像损坏和领域偏移等多种场景下均实现了优越的泛化性能，优于现有TTA方法在多个数据集和模型上的表现。

ABSTRACT

Personalized federated learning algorithms have shown promising results in adapting models to various distribution shifts. However, most of these methods require labeled data on testing clients for personalization, which is usually unavailable in real-world scenarios. In this paper, we introduce a novel setting called test-time personalized federated learning (TTPFL), where clients locally adapt a global model in an unsupervised way without relying on any labeled data during test-time. While traditional test-time adaptation (TTA) can be used in this scenario, most of them inherently assume training data come from a single domain, while they come from multiple clients (source domains) with different distributions. Overlooking these domain interrelationships can result in suboptimal generalization. Moreover, most TTA algorithms are designed for a specific kind of distribution shift and lack the flexibility to handle multiple kinds of distribution shifts in FL. In this paper, we find that this lack of flexibility partially results from their pre-defining which modules to adapt in the model. To tackle this challenge, we propose a novel algorithm called ATP to adaptively learns the adaptation rates for each module in the model from distribution shifts among source domains. Theoretical analysis proves the strong generalization of ATP. Extensive experiments demonstrate its superiority in handling various distribution shifts including label shift, image corruptions, and domain shift, outperforming existing TTA methods across multiple datasets and model architectures. Our code is available at https://github.com/baowenxuan/ATP .

研究动机与目标

• 为应对新客户端在面临复杂分布偏移时个性化联邦模型的挑战。
• 克服现有测试时自适应（TTA）方法的局限性，这些方法假设单一领域训练，缺乏在多客户端联邦学习设置中的灵活性。
• 通过学习基于源客户端间观察到的分布偏移来决定哪些模型组件需要自适应，实现灵活且模块无关的自适应。
• 为跨设备联邦学习场景中的无监督个性化提供理论基础坚实且鲁棒的方法。

提出的方法

• ATP在源客户端上进行训练，通过模拟无监督的测试时自适应，并优化模块特定的适应率以最大化性能。
• 适应率按模型模块（如批归一化统计量、权重）分别学习，并在中央聚合以提升泛化能力。
• 推理阶段，目标客户端使用学习到的适应率，仅依靠无标签测试数据对全局模型进行本地个性化。
• 该方法对批量适应的模型状态进行累积平均，实现在线自适应，从而提升稳定性和性能。
• 理论分析表明，由于利用多源数据和低维适应率学习，ATP具有良好的泛化能力。
• 允许负适应率，从而在标签偏移下实现反直觉但有效的“去对齐”特征操作。

实验结果

研究问题

• RQ1 我们能否在联邦学习中不依赖标签测试数据，学习到有效且灵活的模型模块适应率？
• RQ2 如何利用源域客户端间的客户端间分布偏移，来指导测试时联邦学习中的无监督个性化？
• RQ3 自适应的、模块特定的自适应是否优于在多种分布偏移下固定的模块TTA方法？
• RQ4 负适应率是否能提升标签偏移下的性能？如果是，原因是什么？

主要发现

• ATP在多个数据集（如CIFAR-10、CIFAR-100、Tiny ImageNet）和模型架构（ResNet、Vision Transformer）上持续优于现有TTA方法。
• 在标签偏移下，ATP为批归一化运行统计量学习到负适应率，通过反直觉的“去对齐”特征操作提升了准确率。
• 实验表明，即使在部分客户端参与（低至C=15）的情况下，ATP仍表现出稳健的泛化能力，收敛稳定且测试损失低。
• 在全参与和部分参与设置下，训练和测试损失曲线均稳定收敛，表明具有强大的泛化能力和鲁棒性。
• 在小样本示例中，ATP采用α = -0.5时在标签偏移下达到92%准确率，优于α = 0（89%）和α = 1（73%），验证了负适应率的有效性。
• ATP的理论分析证实了其因使用多个源域和低维适应率学习而具备强大的泛化能力。