FedPAT: Federated Test-Time Adaptation via Prototype Affinity Topology
Abstract
Federated Learning (FL) enables privacy-preserving collaboration among distributed clients in open-world environments, but its performance often degrades under data heterogeneity and unpredictable distribution shifts. Test-Time Adaptation (TTA) has recently been introduced into FL to leverage unlabeled data from unseen clients for online adaptation. However, most existing federated TTA methods employ local feature statistics, which can be brittle under diverse and severe distribution shifts. In this work, we observe that despite significant variations in feature distributions, the relational structure among class prototypes—termed prototype affinity topology—remains remarkably stable across heterogeneous clients. Building on this insight, we propose FedPAT, a Federated TTA framework that leverages Prototype Affinity Topology as a cross-client structural prior. FedPAT learns a global PAT by aggregating class prototypes from source clients, capturing consensus inter-class relationships that are robust to local distribution variations. For unseen target clients, we design a topology-aware mechanism that enhances predictions via diffusion of the global PAT, fuses them with parametric outputs, and performs lightweight optimization for robust test-time adaptation. Extensive experiments demonstrate that FedPAT consistently outperforms advanced federated TTA and classical TTA methods across various distribution shifts.
Lay Summary
Federated learning enables many devices, such as smartphones, medical sensors, and edge computers, to collaboratively train a shared AI model without exposing their private data. In practice, however, this model is often deployed on entirely new devices operating under unfamiliar conditions, such as unexpected image corruptions or novel data distributions, causing significant performance drops. A promising solution is test-time adaptation, which allows the model to adjust itself on the fly using only the unlabeled data available at the new device. Yet existing approaches rely on local data statistics that become unreliable when distribution shifts are severe or highly varied. We identify a surprisingly stable property: even when raw data statistics differ substantially across devices, the relational structure among class categories remains consistent. A "cat" embedding stays closer to "dog" than to "airplane" regardless of corruption type or device characteristics. We call this structure the Prototype Affinity Topology (PAT). Building on this insight, we propose FedPAT, a federated test-time adaptation framework that extracts and aggregates inter-class relational structures from all participating devices during training, and uses the resulting global topology as a stable prior to guide adaptation on unseen devices, requiring no labeled data and no access to the original training set. Extensive experiments demonstrate that FedPAT consistently outperforms state-of-the-art methods across diverse and challenging distribution shifts, while also providing stronger resistance to privacy attacks compared to standard federated learning.